Merge pull request #669 from ably/fix/jwt-authentication

[ECO-4550] Fix JWT authentication

Author: sacOO7

README.md

@@ -326,6 +326,56 @@ if err != nil {
 fmt.Print(status, status.ChannelId)
 ```
 
+### Authentication
+
+Initialize `ably.NewREST` using `ABLY_KEY`. Check [Authentication Doc](https://ably.com/docs/auth) for more information types of auth and it's server/client-side usage.
+
+```go
+restClient, err := ably.NewREST(ably.WithKey("API_KEY"))
+```
+
+Token requests are signed using provided `API_KEY` and issued by your servers.
+
+```go
+// e.g. Gin server endpoint
+router.GET("/token", getToken)
+func getToken(c *gin.Context) {
+    token, err := restClient.Auth.CreateTokenRequest(nil)
+    c.IndentedJSON(http.StatusOK, token)
+}
+```
+
+- When using `WithAuthURL` clientOption at client side, for [JWT token](https://ably.com/tutorials/jwt-authentication) response, contentType header should be set to  `text/plain` or `application/jwt`. For `ably.TokenRequest`/ `ably.TokenDetails`, set it as  `application/json`.
+
+### Using the Token auth at client side
+
+`WithAuthUrl` clientOption automatically decodes response based on the response contentType, `WithAuthCallback` needs manual decoding based on the response. See [official token auth documentation](https://ably.com/docs/auth/token?lang=go) for more information.
+
+```go
+// Return token of type ably.TokenRequest, ably.TokenDetails or ably.TokenString
+authCallback := ably.WithAuthCallback(func(ctx context.Context, tp ably.TokenParams) (ably.Tokener, error) {
+        // HTTP client impl. to fetch token, you can pass tokenParams based on your requirement
+        tokenReqJsonString, err := requestTokenFrom(ctx, "/token");
+        if err != nil {
+                return nil, err
+        }
+        var req ably.TokenRequest
+        err := json.Unmarshal(tokenReqJsonString, &req)
+        return req, err
+})
+
+```
+If [JWT token](https://ably.com/tutorials/jwt-authentication) is returned by server
+```go
+authCallback := ably.WithAuthCallback(func(ctx context.Context, tp ably.TokenParams) (ably.Tokener, error) {
+        jwtTokenString, err := requestTokenFrom(ctx, "/jwtToken"); // jwtTokenString starts with "ey"
+        if err != nil {
+                return nil, err
+        }
+        return ably.TokenString(jwtTokenString), err
+})
+```
+
 ### Configure logging
 - By default, internal logger prints output to `stdout` with default logging level of `warning`.
 - You need to create a custom Logger that implements `ably.Logger` interface.
@@ -442,9 +492,6 @@ As of release 1.2.0, the following are not implemented and will be covered in fu
 
 - [Push notifications admin API](https://ably.com/docs/api/rest-sdk/push-admin) is not implemented.
 
-- [JWT authentication](https://ably.com/docs/auth/token?lang=javascript#jwt) using `auth-url` is not implemented.
-See [jwt auth issue](https://github.com/ably/ably-go/issues/569) for more details.
-
 ### Realtime API
 
 - Channel suspended state is partially implemented. See [suspended channel state](https://github.com/ably/ably-go/issues/568).

ably/auth.go

@@ -462,7 +462,7 @@ func (a *Auth) requestAuthURL(ctx context.Context, params *TokenParams, opts *au
 		return nil, a.newError(40004, err)
 	}
 	switch typ {
-	case "text/plain":
+	case "text/plain", "application/jwt":
 		token, err := io.ReadAll(resp.Body)
 		if err != nil {
 			return nil, a.newError(40000, err)

ably/auth_integration_test.go

@@ -4,10 +4,13 @@
 package ably_test
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"net/url"
 	"strings"
@@ -342,6 +345,164 @@ func TestAuth_RequestToken(t *testing.T) {
 	}
 }
 
+func TestAuth_JWT_Token_RSA8c(t *testing.T) {
+
+	t.Run("Get JWT from echo server", func(t *testing.T) {
+		app := ablytest.MustSandbox(nil)
+		defer safeclose(t, app)
+		jwt, err := app.CreateJwt(3*time.Second, false)
+		assert.NoError(t, err)
+		assert.True(t, strings.HasPrefix(jwt, "ey"))
+
+		// JWT header assertions
+		header := strings.Split(jwt, ".")[0]
+		jwtToken, err := base64.RawStdEncoding.DecodeString(header)
+		assert.NoError(t, err)
+		var result map[string]interface{}
+		err = json.Unmarshal(jwtToken, &result)
+		assert.NoError(t, err)
+		assert.Equal(t, "HS256", result["alg"])
+		assert.Equal(t, "JWT", result["typ"])
+		assert.Contains(t, result, "kid")
+		assert.NoError(t, err)
+
+		// JWT payload assertions
+		payload := strings.Split(jwt, ".")[1]
+		jwtToken, err = base64.RawStdEncoding.DecodeString(payload)
+		assert.NoError(t, err)
+		err = json.Unmarshal(jwtToken, &result)
+		assert.NoError(t, err)
+		assert.Contains(t, result, "iat")
+		assert.Contains(t, result, "exp")
+
+		// check expiry of 3 seconds
+		assert.Equal(t, float64(3), result["exp"].(float64)-result["iat"].(float64))
+	})
+
+	t.Run("Should be able to use it as a token", func(t *testing.T) {
+		app := ablytest.MustSandbox(nil)
+		defer safeclose(t, app)
+		jwt, err := app.CreateJwt(3*time.Second, false)
+		assert.NoError(t, err)
+		assert.True(t, strings.HasPrefix(jwt, "ey"))
+
+		rec, optn := ablytest.NewHttpRecorder()
+		rest, err := ably.NewREST(
+			ably.WithToken(jwt),
+			ably.WithEnvironment(app.Environment),
+			optn[0],
+		)
+		assert.NoError(t, err, "rest()=%v", err)
+
+		_, err = rest.Stats().Pages(context.Background())
+		assert.NoError(t, err, "Stats()=%v", err)
+
+		assert.Len(t, rec.Requests(), 1)
+		assert.Len(t, rec.Responses(), 1)
+
+		statsRequest := rec.Request(0)
+		assert.Equal(t, "/stats", statsRequest.URL.Path)
+		encodedToken := base64.StdEncoding.EncodeToString([]byte(jwt))
+		assert.Equal(t, "Bearer "+encodedToken, statsRequest.Header.Get("Authorization"))
+	})
+
+	t.Run("RSA8g, RSA3d: Should be able to authenticate using authURL", func(t *testing.T) {
+		app := ablytest.MustSandbox(nil)
+		defer safeclose(t, app)
+
+		rec, optn := ablytest.NewHttpRecorder()
+		rest, err := ably.NewREST(
+			ably.WithAuthURL(ablytest.CREATE_JWT_URL),
+			ably.WithAuthParams(app.GetJwtAuthParams(30*time.Second, false)),
+			ably.WithEnvironment(app.Environment),
+			optn[0],
+		)
+		assert.NoError(t, err, "rest()=%v", err)
+
+		_, err = rest.Stats().Pages(context.Background())
+		assert.NoError(t, err, "Stats()=%v", err)
+
+		assert.Len(t, rec.Requests(), 2)
+		assert.Len(t, rec.Responses(), 2)
+
+		// first request is jwt request
+		jwtRequest := rec.Request(0).URL
+		assert.Equal(t, ablytest.CREATE_JWT_URL, "https://"+jwtRequest.Host+jwtRequest.Path)
+		// response is jwt token
+		jwtResponse, err := io.ReadAll(rec.Response(0).Body)
+		assert.NoError(t, err)
+		assert.True(t, bytes.HasPrefix(jwtResponse, []byte("ey")))
+
+		// Second request is made to stats with given jwt token (base64 encoded)
+		statsRequest := rec.Request(1)
+		assert.Equal(t, "/stats", statsRequest.URL.Path)
+		encodedToken := base64.StdEncoding.EncodeToString(jwtResponse)
+		assert.Equal(t, "Bearer "+encodedToken, statsRequest.Header.Get("Authorization"))
+	})
+
+	t.Run("RSA8g, RSA3d: Should be able to authenticate using authCallback", func(t *testing.T) {
+		app := ablytest.MustSandbox(nil)
+		defer safeclose(t, app)
+
+		jwtToken := ""
+		authCallback := ably.WithAuthCallback(func(ctx context.Context, tp ably.TokenParams) (ably.Tokener, error) {
+			jwtTokenString, err := app.CreateJwt(time.Second*30, false)
+			jwtToken = jwtTokenString
+			if err != nil {
+				t.Fatalf("Error creating JWT: %v", err)
+				return nil, err
+			}
+			return ably.TokenString(jwtTokenString), nil
+		})
+
+		rec, optn := ablytest.NewHttpRecorder()
+		rest, err := ably.NewREST(
+			ably.WithEnvironment(app.Environment),
+			authCallback,
+			optn[0],
+		)
+		assert.NoError(t, err)
+
+		_, err = rest.Stats().Pages(context.Background())
+		assert.NoError(t, err, "Stats()=%v", err)
+
+		assert.Len(t, rec.Requests(), 1)
+		assert.Len(t, rec.Responses(), 1)
+
+		assert.True(t, strings.HasPrefix(jwtToken, "ey"))
+		// Second request is made to stats with given jwt token (base64 encoded)
+		statsRequest := rec.Request(0)
+		assert.Equal(t, "/stats", statsRequest.URL.Path)
+		encodedToken := base64.StdEncoding.EncodeToString([]byte(jwtToken))
+		assert.Equal(t, "Bearer "+encodedToken, statsRequest.Header.Get("Authorization"))
+	})
+
+	t.Run("RSA4e, RSA4b: Should return error when JWT is invalid", func(t *testing.T) {
+		app := ablytest.MustSandbox(nil)
+		defer safeclose(t, app)
+
+		rec, optn := ablytest.NewHttpRecorder()
+		rest, err := ably.NewREST(
+			ably.WithAuthURL(ablytest.CREATE_JWT_URL),
+			ably.WithAuthParams(app.GetJwtAuthParams(30*time.Second, true)),
+			ably.WithEnvironment(app.Environment),
+			optn[0],
+		)
+		assert.NoError(t, err, "rest()=%v", err)
+
+		_, err = rest.Stats().Pages(context.Background())
+		var errorInfo *ably.ErrorInfo
+		assert.Error(t, err)
+		assert.ErrorAs(t, err, &errorInfo)
+		assert.Equal(t, 40144, int(errorInfo.Code))
+		assert.Equal(t, 401, errorInfo.StatusCode)
+		assert.Contains(t, err.Error(), "invalid JWT format")
+
+		assert.Len(t, rec.Requests(), 2)
+		assert.Len(t, rec.Responses(), 2)
+	})
+}
+
 func TestAuth_ReuseClientID(t *testing.T) {
 	opts := []ably.ClientOption{ably.WithUseTokenAuth(true)}
 	app, client := ablytest.NewREST(opts...)

ably/realtime_conn_spec_integration_test.go

@@ -3022,8 +3022,63 @@ func TestRealtimeConn_RTC8a_ExplicitAuthorizeWhileConnected(t *testing.T) {
 		ablytest.Instantly.Recv(t, nil, authorizeDone, t.Fatalf)
 	})
 
-	t.Run("RTC8a4: reauthorize with JWT token", func(t *testing.T) {
-		t.Skip("not implemented")
+	t.Run("RTC8a4, RSA3d: reauthorize with JWT token", func(t *testing.T) {
+		t.Parallel()
+		app := ablytest.MustSandbox(nil)
+		defer safeclose(t, app)
+
+		authCallbackTokens := []string{}
+		tokenExpiry := 3 * time.Second
+		// Returns token that expires after 3 seconds causing disconnect every 3 seconds
+		authCallback := func(ctx context.Context, tp ably.TokenParams) (ably.Tokener, error) {
+			jwtTokenString, err := app.CreateJwt(tokenExpiry, false)
+			if err != nil {
+				return nil, err
+			}
+			authCallbackTokens = append(authCallbackTokens, jwtTokenString)
+			return ably.TokenString(jwtTokenString), nil
+		}
+
+		realtimeMsgRecorder := NewMessageRecorder()
+		realtime, err := ably.NewRealtime(
+			ably.WithAutoConnect(false),
+			ably.WithEnvironment(ablytest.Environment),
+			ably.WithDial(realtimeMsgRecorder.Dial),
+			ably.WithAuthCallback(authCallback))
+
+		assert.NoError(t, err)
+		defer realtime.Close()
+
+		err = ablytest.Wait(ablytest.ConnWaiter(realtime, realtime.Connect, ably.ConnectionEventConnected), nil)
+		assert.NoError(t, err)
+
+		changes := make(ably.ConnStateChanges, 2)
+		off := realtime.Connection.OnAll(changes.Receive)
+		defer off()
+		var state ably.ConnectionStateChange
+
+		// Disconnects due to timeout
+		ablytest.Soon.Recv(t, &state, changes, t.Fatalf)
+		assert.Equal(t, ably.ConnectionEventDisconnected, state.Event)
+		// Reconnect again using new JWT token
+		ablytest.Soon.Recv(t, &state, changes, t.Fatalf)
+		assert.Equal(t, ably.ConnectionEventConnecting, state.Event)
+		ablytest.Soon.Recv(t, &state, changes, t.Fatalf)
+		assert.Equal(t, ably.ConnectionEventConnected, state.Event)
+		assert.Nil(t, state.Reason)
+		assert.Equal(t, ably.ConnectionStateConnected, realtime.Connection.State())
+
+		ablytest.Instantly.NoRecv(t, nil, changes, t.Fatalf)
+
+		// Make sure requested tokens are JWT tokens
+		assert.Len(t, authCallbackTokens, 2)
+		assert.True(t, strings.HasPrefix(authCallbackTokens[0], "ey"))
+		assert.True(t, strings.HasPrefix(authCallbackTokens[1], "ey"))
+		assertUnique(t, authCallbackTokens)
+		// 2 Dial attempts made
+		assert.Len(t, realtimeMsgRecorder.URLs(), 2)
+		assert.Equal(t, authCallbackTokens[0], realtimeMsgRecorder.URLs()[0].Query().Get("access_token"))
+		assert.Equal(t, authCallbackTokens[1], realtimeMsgRecorder.URLs()[1].Query().Get("access_token"))
 	})
 
 	t.Run("RTC8a2: Failed reauth moves connection to FAILED", func(t *testing.T) {

ablytest/sandbox.go

@@ -256,6 +256,67 @@ func (app *Sandbox) URL(paths ...string) string {
 	return "https://" + app.Environment + "-rest.ably.io/" + path.Join(paths...)
 }
 
+// Source code for the same => https://github.com/ably/echoserver/blob/main/app.js
+var CREATE_JWT_URL string = "https://echo.ably.io/createJWT"
+
+// GetJwtAuthParams constructs the authentication parameters required for JWT creation.
+// Required when authUrl is chosen as a mode of auth
+//
+// Parameters:
+// - expiresIn: The duration until the JWT expires.
+// - invalid: A boolean flag indicating whether to use an invalid key secret.
+//
+// Returns: A url.Values object containing the authentication parameters.
+func (app *Sandbox) GetJwtAuthParams(expiresIn time.Duration, invalid bool) url.Values {
+	key, secret := app.KeyParts()
+	authParams := url.Values{}
+	authParams.Add("environment", app.Environment)
+	authParams.Add("returnType", "jwt")
+	authParams.Add("keyName", key)
+	if invalid {
+		authParams.Add("keySecret", "invalid")
+	} else {
+		authParams.Add("keySecret", secret)
+	}
+	authParams.Add("expiresIn", fmt.Sprint(expiresIn.Seconds()))
+	return authParams
+}
+
+// CreateJwt generates a JWT with the specified expiration time.
+//
+// Parameters:
+// - expiresIn: The duration until the JWT expires.
+// - invalid: A boolean flag indicating whether to use an invalid key secret.
+//
+// Returns:
+// - A string containing the generated JWT.
+// - An error if the JWT creation fails.
+func (app *Sandbox) CreateJwt(expiresIn time.Duration, invalid bool) (string, error) {
+	u, err := url.Parse(CREATE_JWT_URL)
+	if err != nil {
+		return "", err
+	}
+	u.RawQuery = app.GetJwtAuthParams(expiresIn, invalid).Encode()
+	req, err := http.NewRequest(http.MethodGet, u.String(), nil)
+	if err != nil {
+		return "", fmt.Errorf("client: could not create request: %s", err)
+	}
+	res, err := app.client.Do(req)
+	if err != nil {
+		res.Body.Close()
+		return "", fmt.Errorf("client: error making http request: %s", err)
+	}
+	defer res.Body.Close()
+	resBody, err := io.ReadAll(res.Body)
+	if err != nil {
+		return "", fmt.Errorf("client: could not read response body: %s", err)
+	}
+	if res.StatusCode != 200 {
+		return "", fmt.Errorf("non-success response received: %v:%s", res.StatusCode, resBody)
+	}
+	return string(resBody), nil
+}
+
 func NewHTTPClient() *http.Client {
 	const timeout = time.Minute
 	return &http.Client{