Repository: supabase/gotrue
Pull request #1332 (2023-12-04 05:50:47)

feat: add custom access token hook
## What kind of change does this PR introduce?

Allow developers to customize access token by adding claims or altering metadata

Other affected changes:
 - `generateAccessToken` is moved to a method on `API` and now takes in a `context`. In tests we use a default `context.Background()` 



Patch:
diff --git a/go.mod b/go.mod
index c723a6a3..18cdc5d4 100644
--- a/go.mod
+++ b/go.mod
@@ -39,10 +39,10 @@ require (
 
 require (
 	github.com/bits-and-blooms/bitset v1.10.0 // indirect
-	github.com/bits-and-blooms/bloom/v3 v3.6.0 // indirect
 	github.com/go-jose/go-jose/v3 v3.0.1 // indirect
 	github.com/gobuffalo/nulls v0.4.2 // indirect
-	github.com/supabase/hibp v0.0.0-20231124125943-d225752ae869 // indirect
+	github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect
+	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
 )
 
 require (
@@ -66,12 +66,15 @@ require (
 )
 
 require (
+	github.com/bits-and-blooms/bloom/v3 v3.6.0
 	github.com/crewjam/saml v0.4.14
 	github.com/deepmap/oapi-codegen v1.12.4
 	github.com/fatih/structs v1.1.0
 	github.com/gobuffalo/pop/v6 v6.1.1
 	github.com/jackc/pgx/v4 v4.17.2
+	github.com/supabase/hibp v0.0.0-20231124125943-d225752ae869
 	github.com/supabase/mailme v0.0.0-20230628061017-01f68480c747
+	github.com/xeipuuv/gojsonschema v1.2.0
 )
 
 require (
diff --git a/go.sum b/go.sum
index c77fb09c..3db4b75b 100644
--- a/go.sum
+++ b/go.sum
@@ -491,7 +491,14 @@ github.com/supabase/hibp v0.0.0-20231124125943-d225752ae869 h1:VDuRtwen5Z7QQ5ctu
 github.com/supabase/hibp v0.0.0-20231124125943-d225752ae869/go.mod h1:eHX5nlSMSnyPjUrbYzeqrA8snCe2SKyfizKjU3dkfOw=
 github.com/supabase/mailme v0.0.0-20230628061017-01f68480c747 h1:FIUdLV4o5JLsJno4Poum157kAxDKINeJo6liBfauLrI=
 github.com/supabase/mailme v0.0.0-20230628061017-01f68480c747/go.mod h1:kWsnmPfUBZTavlXYkfJrE9unzmmRAIi/kqsxXfEWEY8=
+github.com/twmb/murmur3 v1.1.6 h1:mqrRot1BRxm+Yct+vavLMou2/iJt0tNVTTC0QoIjaZg=
 github.com/twmb/murmur3 v1.1.6/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ=
+github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f h1:J9EGpcZtP0E/raorCMxlFGSTBrsSlaDGf3jU/qvAE2c=
+github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=
+github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHovont7NscjpAxXsDA8S8BMYve8Y5+7cuRE7R0=
+github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ=
+github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74=
+github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
diff --git a/internal/api/mfa.go b/internal/api/mfa.go
index 0ce5a2e5..dece1671 100644
--- a/internal/api/mfa.go
+++ b/internal/api/mfa.go
@@ -289,6 +289,45 @@ func (a *API) invokeHook(ctx context.Context, input, output any) error {
 		}
 
 		return nil
+	case *hooks.CustomAccessTokenInput:
+		hookOutput, ok := output.(*hooks.CustomAccessTokenOutput)
+		if !ok {
+			panic("output should be *hooks.CustomAccessTokenOutput")
+		}
+
+		if _, err := a.runHook(ctx, config.Hook.CustomAccessToken.HookName, input, output); err != nil {
+			return internalServerError("Error invoking access token hook.").WithInternalError(err)
+		}
+
+		if hookOutput.IsError() {
+			httpCode := hookOutput.HookError.HTTPCode
+
+			if httpCode == 0 {
+				httpCode = http.StatusInternalServerError
+			}
+
+			httpError := &HTTPError{
+				Code:    httpCode,
+				Message: hookOutput.HookError.Message,
+			}
+
+			return httpError.WithInternalError(&hookOutput.HookError)
+		}
+		if err := validateTokenClaims(hookOutput.Claims); err != nil {
+			httpCode := hookOutput.HookError.HTTPCode
+
+			if httpCode == 0 {
+				httpCode = http.StatusInternalServerError
+			}
+
+			httpError := &HTTPError{
+				Code:    httpCode,
+				Message: err.Error(),
+			}
+
+			return httpError
+		}
+		return nil
 
 	default:
 		panic("unknown hook input type")
diff --git a/internal/api/token.go b/internal/api/token.go
index 8deaf98f..85bd3912 100644
--- a/internal/api/token.go
+++ b/internal/api/token.go
@@ -9,8 +9,10 @@ import (
 	"strconv"
 	"time"
 
+	"fmt"
 	"github.com/gofrs/uuid"
 	"github.com/golang-jwt/jwt"
+	"github.com/xeipuuv/gojsonschema"
 
 	"github.com/supabase/gotrue/internal/conf"
 	"github.com/supabase/gotrue/internal/hooks"
@@ -285,7 +287,8 @@ func (a *API) PKCE(ctx context.Context, w http.ResponseWriter, r *http.Request)
 	return sendJSON(w, http.StatusOK, token)
 }
 
-func generateAccessToken(tx *storage.Connection, user *models.User, sessionId *uuid.UUID, config *conf.JWTConfiguration) (string, int64, error) {
+func (a *API) generateAccessToken(ctx context.Context, tx *storage.Connection, user *models.User, sessionId *uuid.UUID, authenticationMethod models.AuthenticationMethod) (string, int64, error) {
+	config := a.config
 	aal, amr := models.AAL1.String(), []models.AMREntry{}
 	sid := ""
 	if sessionId != nil {
@@ -301,15 +304,15 @@ func generateAccessToken(tx *storage.Connection, user *models.User, sessionId *u
 	}
 
 	issuedAt := time.Now().UTC()
-	expiresAt := issuedAt.Add(time.Second * time.Duration(config.Exp)).Unix()
+	expiresAt := issuedAt.Add(time.Second * time.Duration(config.JWT.Exp)).Unix()
 
-	claims := &GoTrueClaims{
+	claims := &hooks.GoTrueClaims{
 		StandardClaims: jwt.StandardClaims{
 			Subject:   user.ID.String(),
 			Audience:  user.Aud,
 			IssuedAt:  issuedAt.Unix(),
 			ExpiresAt: expiresAt,
-			Issuer:    config.Issuer,
+			Issuer:    config.JWT.Issuer,
 		},
 		Email:                         user.GetEmail(),
 		Phone:                         user.GetPhone(),
@@ -321,17 +324,37 @@ func generateAccessToken(tx *storage.Connection, user *models.User, sessionId *u
 		AuthenticationMethodReference: amr,
 	}
 
-	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	var token *jwt.Token
+	if config.Hook.CustomAccessToken.Enabled {
+		input := hooks.CustomAccessTokenInput{
+			UserID:               user.ID,
+			Claims:               claims,
+			AuthenticationMethod: authenticationMethod.String(),
+		}
+
+		output := hooks.CustomAccessTokenOutput{}
+
+		err := a.invokeHook(ctx, &input, &output)
+		if err != nil {
+			return "", 0, err
+		}
+		goTrueClaims := jwt.MapClaims(output.Claims)
 
-	if config.KeyID != "" {
+		token = jwt.NewWithClaims(jwt.SigningMethodHS256, goTrueClaims)
+
+	} else {
+		token = jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	}
+
+	if config.JWT.KeyID != "" {
 		if token.Header == nil {
 			token.Header = make(map[string]interface{})
 		}
 
-		token.Header["kid"] = config.KeyID
+		token.Header["kid"] = config.JWT.KeyID
 	}
 
-	signed, err := token.SignedString([]byte(config.Secret))
+	signed, err := token.SignedString([]byte(config.JWT.Secret))
 	if err != nil {
 		return "", 0, err
 	}
@@ -362,7 +385,7 @@ func (a *API) issueRefreshToken(ctx context.Context, conn *storage.Connection, u
 			return terr
 		}
 
-		tokenString, expiresAt, terr = generateAccessToken(tx, user, refreshToken.SessionId, &config.JWT)
+		tokenString, expiresAt, terr = a.generateAccessToken(ctx, tx, user, refreshToken.SessionId, authenticationMethod)
 		if terr != nil {
 			return internalServerError("error generating jwt token").WithInternalError(terr)
 		}
@@ -422,7 +445,7 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 			return err
 		}
 
-		tokenString, expiresAt, terr = generateAccessToken(tx, user, &sessionId, &config.JWT)
+		tokenString, expiresAt, terr = a.generateAccessToken(ctx, tx, user, &sessionId, models.TOTPSignIn)
 
 		if terr != nil {
 			return internalServerError("error generating jwt token").WithInternalError(terr)
@@ -494,3 +517,27 @@ func (a *API) clearCookieToken(config *conf.GlobalConfiguration, name string, w
 		Domain:   config.Cookie.Domain,
 	})
 }
+
+func validateTokenClaims(outputClaims map[string]interface{}) error {
+	schemaLoader := gojsonschema.NewStringLoader(hooks.MinimumViableTokenSchema)
+
+	documentLoader := gojsonschema.NewGoLoader(outputClaims)
+
+	result, err := gojsonschema.Validate(schemaLoader, documentLoader)
+	if err != nil {
+		return err
+	}
+
+	if !result.Valid() {
+		var errorMessages string
+
+		for _, desc := range result.Errors() {
+			errorMessages += fmt.Sprintf("- %s\n", desc)
+			fmt.Printf("- %s\n", desc)
+		}
+		return fmt.Errorf("output claims do not conform to the expected schema: \n%s", errorMessages)
+
+	}
+
+	return nil
+}
diff --git a/internal/api/token_refresh.go b/internal/api/token_refresh.go
index c492ccbe..b1bfb6d3 100644
--- a/internal/api/token_refresh.go
+++ b/internal/api/token_refresh.go
@@ -217,7 +217,7 @@ func (a *API) RefreshTokenGrant(ctx context.Context, w http.ResponseWriter, r *h
 				issuedToken = newToken
 			}
 
-			tokenString, expiresAt, terr = generateAccessToken(tx, user, issuedToken.SessionId, &config.JWT)
+			tokenString, expiresAt, terr = a.generateAccessToken(ctx, tx, user, issuedToken.SessionId, models.TokenRefresh)
 			if terr != nil {
 				return internalServerError("error generating jwt token").WithInternalError(terr)
 			}
diff --git a/internal/conf/configuration.go b/internal/conf/configuration.go
index 162e233d..c758142d 100644
--- a/internal/conf/configuration.go
+++ b/internal/conf/configuration.go
@@ -443,6 +443,7 @@ type WebhookConfig struct {
 type HookConfiguration struct {
 	MFAVerificationAttempt      ExtensibilityPointConfiguration `json:"mfa_verification_attempt" split_words:"true"`
 	PasswordVerificationAttempt ExtensibilityPointConfiguration `json:"password_verification_attempt" split_words:"true"`
+	CustomAccessToken           ExtensibilityPointConfiguration `json:"custom_access_token" split_words:"true"`
 }
 
 type ExtensibilityPointConfiguration struct {
@@ -455,6 +456,7 @@ func (h *HookConfiguration) Validate() error {
 	points := []ExtensibilityPointConfiguration{
 		h.MFAVerificationAttempt,
 		h.PasswordVerificationAttempt,
+		h.CustomAccessToken,
 	}
 	for _, point := range points {
 		if err := point.ValidateAndPopulateExtensibilityPoint(); err != nil {
diff --git a/internal/hooks/auth_hooks.go b/internal/hooks/auth_hooks.go
index a2fc7018..0a1711da 100644
--- a/internal/hooks/auth_hooks.go
+++ b/internal/hooks/auth_hooks.go
@@ -2,6 +2,8 @@ package hooks
 
 import (
 	"github.com/gofrs/uuid"
+	"github.com/golang-jwt/jwt"
+	"github.com/supabase/gotrue/internal/models"
 )
 
 type HookType string
@@ -25,6 +27,78 @@ type HookOutput interface {
 	Error() string
 }
 
+// #nosec
+const MinimumViableTokenSchema = `{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "object",
+  "properties": {
+    "aud": {
+      "type": "string"
+    },
+    "exp": {
+      "type": "integer"
+    },
+    "jti": {
+      "type": "string"
+    },
+    "iat": {
+      "type": "integer"
+    },
+    "iss": {
+      "type": "string"
+    },
+    "nbf": {
+      "type": "integer"
+    },
+    "sub": {
+      "type": "string"
+    },
+    "email": {
+      "type": "string"
+    },
+    "phone": {
+      "type": "string"
+    },
+    "app_metadata": {
+      "type": "object",
+      "additionalProperties": true
+    },
+    "user_metadata": {
+      "type": "object",
+      "additionalProperties": true
+    },
+    "role": {
+      "type": "string"
+    },
+    "aal": {
+      "type": "string"
+    },
+    "amr": {
+      "type": "array",
+      "items": {
+        "type": "object"
+      }
+    },
+    "session_id": {
+      "type": "string"
+    }
+  },
+  "required": ["aud", "exp", "iat", "sub", "email", "phone", "role", "aal"]
+}`
+
+// GoTrueClaims is a struct thats used for JWT claims
+type GoTrueClaims struct {
+	jwt.StandardClaims
+	Email                         string                 `json:"email"`
+	Phone                         string                 `json:"phone"`
+	AppMetaData                   map[string]interface{} `json:"app_metadata"`
+	UserMetaData                  map[string]interface{} `json:"user_metadata"`
+	Role                          string                 `json:"role"`
+	AuthenticatorAssuranceLevel   string                 `json:"aal,omitempty"`
+	AuthenticationMethodReference []models.AMREntry      `json:"amr,omitempty"`
+	SessionId                     string                 `json:"session_id,omitempty"`
+}
+
 type MFAVerificationAttemptInput struct {
 	UserID   uuid.UUID `json:"user_id"`
 	FactorID uuid.UUID `json:"factor_id"`
@@ -49,6 +123,17 @@ type PasswordVerificationAttemptOutput struct {
 	HookError        AuthHookError `json:"error"`
 }
 
+type CustomAccessTokenInput struct {
+	UserID               uuid.UUID     `json:"user_id"`
+	Claims               *GoTrueClaims `json:"claims"`
+	AuthenticationMethod string        `json:"authentication_method"`
+}
+
+type CustomAccessTokenOutput struct {
+	Claims    map[string]interface{} `json:"claims"`
+	HookError AuthHookError          `json:"error,omitempty"`
+}
+
 func (mf *MFAVerificationAttemptOutput) IsError() bool {
 	return mf.HookError.Message != ""
 }
@@ -65,6 +150,14 @@ func (p *PasswordVerificationAttemptOutput) Error() string {
 	return p.HookError.Message
 }
 
+func (ca *CustomAccessTokenOutput) IsError() bool {
+	return ca.HookError.Message != ""
+}
+
+func (ca *CustomAccessTokenOutput) Error() string {
+	return ca.HookError.Message
+}
+
 type AuthHookError struct {
 	HTTPCode int    `json:"http_code,omitempty"`
 	Message  string `json:"message,omitempty"`
diff --git a/internal/models/factor.go b/internal/models/factor.go
index 64a7f9ff..57fbdd63 100644
--- a/internal/models/factor.go
+++ b/internal/models/factor.go
@@ -44,6 +44,7 @@ const (
 	MagicLink
 	EmailSignup
 	EmailChange
+	TokenRefresh
 )
 
 func (authMethod AuthenticationMethod) String() string {
@@ -68,6 +69,8 @@ func (authMethod AuthenticationMethod) String() string {
 		return "email/signup"
 	case EmailChange:
 		return "email_change"
+	case TokenRefresh:
+		return "token_refresh"
 	}
 	return ""
 }
@@ -97,6 +100,8 @@ func ParseAuthenticationMethod(authMethod string) (AuthenticationMethod, error)
 		return EmailSignup, nil
 	case "email_change":
 		return EmailChange, nil
+	case "token_refresh":
+		return TokenRefresh, nil
 	}
 	return 0, fmt.Errorf("unsupported authentication method %q", authMethod)
 }


Test patch:
diff --git a/internal/api/audit_test.go b/internal/api/audit_test.go
index 2e54db29..aa80dbae 100644
--- a/internal/api/audit_test.go
+++ b/internal/api/audit_test.go
@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -48,7 +49,7 @@ func (ts *AuditTestSuite) makeSuperAdmin(email string) string {
 	u.Role = "supabase_admin"
 
 	var token string
-	token, _, err = generateAccessToken(ts.API.db, u, nil, &ts.Config.JWT)
+	token, _, err = ts.API.generateAccessToken(context.Background(), ts.API.db, u, nil, models.PasswordGrant)
 	require.NoError(ts.T(), err, "Error generating access token")
 
 	p := jwt.Parser{ValidMethods: []string{jwt.SigningMethodHS256.Name}}
diff --git a/internal/api/invite_test.go b/internal/api/invite_test.go
index ba4d2f1d..774aa370 100644
--- a/internal/api/invite_test.go
+++ b/internal/api/invite_test.go
@@ -2,6 +2,7 @@ package api
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -59,7 +60,7 @@ func (ts *InviteTestSuite) makeSuperAdmin(email string) string {
 	u.Role = "supabase_admin"
 
 	var token string
-	token, _, err = generateAccessToken(ts.API.db, u, nil, &ts.Config.JWT)
+	token, _, err = ts.API.generateAccessToken(context.Background(), ts.API.db, u, nil, models.Invite)
 
 	require.NoError(ts.T(), err, "Error generating access token")
 
diff --git a/internal/api/logout_test.go b/internal/api/logout_test.go
index 7a9b3ebb..2c44efbb 100644
--- a/internal/api/logout_test.go
+++ b/internal/api/logout_test.go
@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -42,7 +43,7 @@ func (ts *LogoutTestSuite) SetupTest() {
 
 	// generate access token to use for logout
 	var t string
-	t, _, err = generateAccessToken(ts.API.db, u, nil, &ts.Config.JWT)
+	t, _, err = ts.API.generateAccessToken(context.Background(), ts.API.db, u, nil, models.PasswordGrant)
 	require.NoError(ts.T(), err)
 	ts.token = t
 }
diff --git a/internal/api/mfa_test.go b/internal/api/mfa_test.go
index 014736e1..16449b7b 100644
--- a/internal/api/mfa_test.go
+++ b/internal/api/mfa_test.go
@@ -87,7 +87,7 @@ func (ts *MFATestSuite) SetupTest() {
 }
 
 func (ts *MFATestSuite) generateToken(user *models.User, sessionId *uuid.UUID) string {
-	token, _, err := generateAccessToken(ts.API.db, user, sessionId, &ts.Config.JWT)
+	token, _, err := ts.API.generateAccessToken(context.Background(), ts.API.db, user, sessionId, models.TOTPSignIn)
 	require.NoError(ts.T(), err, "Error generating access token")
 	return token
 }
@@ -96,7 +96,7 @@ func (ts *MFATestSuite) TestEnrollFactor() {
 	testFriendlyName := "bob"
 	alternativeFriendlyName := "john"
 
-	token, _, err := generateAccessToken(ts.API.db, ts.TestUser, nil, &ts.Config.JWT)
+	token, _, err := ts.API.generateAccessToken(context.Background(), ts.API.db, ts.TestUser, nil, models.TOTPSignIn)
 	require.NoError(ts.T(), err)
 
 	var cases = []struct {
diff --git a/internal/api/phone_test.go b/internal/api/phone_test.go
index 9243b32a..396b1366 100644
--- a/internal/api/phone_test.go
+++ b/internal/api/phone_test.go
@@ -157,7 +157,7 @@ func (ts *PhoneTestSuite) TestMissingSmsProviderConfig() {
 	require.NoError(ts.T(), ts.API.db.Update(u), "Error updating new test user")
 
 	var token string
-	token, _, err = generateAccessToken(ts.API.db, u, nil, &ts.Config.JWT)
+	token, _, err = ts.API.generateAccessToken(context.Background(), ts.API.db, u, nil, models.OTP)
 	require.NoError(ts.T(), err)
 
 	cases := []struct {
diff --git a/internal/api/token_test.go b/internal/api/token_test.go
index 179d39b8..a6dfad57 100644
--- a/internal/api/token_test.go
+++ b/internal/api/token_test.go
@@ -10,6 +10,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -58,6 +59,8 @@ func (ts *TokenTestSuite) SetupTest() {
 	ts.User = u
 	ts.RefreshToken, err = models.GrantAuthenticatedUser(ts.API.db, u, models.GrantParams{})
 	require.NoError(ts.T(), err, "Error creating refresh token")
+	ts.Config.Hook.CustomAccessToken.Enabled = false
+
 }
 
 func (ts *TokenTestSuite) TestSessionTimebox() {
@@ -649,3 +652,111 @@ func (ts *TokenTestSuite) TestPasswordVerificationHook() {
 	}
 
 }
+func (ts *TokenTestSuite) TestCustomAccessToken() {
+	type customAccessTokenTestcase struct {
+		desc            string
+		uri             string
+		hookFunctionSQL string
+		expectedClaims  map[string]interface{}
+		shouldError     bool
+	}
+	cases := []customAccessTokenTestcase{
+		{
+			desc: "Add a new claim",
+			uri:  "pg-functions://postgres/auth/custom_access_token_add_claim",
+			hookFunctionSQL: ` create or replace function custom_access_token_add_claim(input jsonb) returns jsonb as $$ declare result jsonb; begin if jsonb_typeof(jsonb_object_field(input, 'claims')) is null then result := jsonb_build_object('error', jsonb_build_object('http_code', 400, 'message', 'Input does not contain claims field')); return result; end if;
+    input := jsonb_set(input, '{claims,newclaim}', '"newvalue"', true);
+    result := jsonb_build_object('claims', input->'claims');
+    return result;
+end; $$ language plpgsql;`,
+			expectedClaims: map[string]interface{}{
+				"newclaim": "newvalue",
+			},
+		}, {
+			desc: "Delete the Role claim",
+			uri:  "pg-functions://postgres/auth/custom_access_token_delete_claim",
+			hookFunctionSQL: `
+create or replace function custom_access_token_delete_claim(input jsonb)
+returns jsonb as $$
+declare
+    result jsonb;
+begin
+    input := jsonb_set(input, '{claims}', (input->'claims') - 'role');
+    result := jsonb_build_object('claims', input->'claims');
+    return result;
+end; $$ language plpgsql;`,
+			expectedClaims: map[string]interface{}{},
+			shouldError:    true,
+		}, {
+			desc: "Delete a non-required claim (UserMetadata)",
+			uri:  "pg-functions://postgres/auth/custom_access_token_delete_usermetadata",
+			hookFunctionSQL: `
+create or replace function custom_access_token_delete_usermetadata(input jsonb)
+returns jsonb as $$
+declare
+    result jsonb;
+begin
+    input := jsonb_set(input, '{claims}', (input->'claims') - 'user_metadata');
+    result := jsonb_build_object('claims', input->'claims');
+    return result;
+end; $$ language plpgsql;`,
+			// Not used
+			expectedClaims: map[string]interface{}{
+				"user_metadata": nil,
+			},
+			shouldError: false,
+		},
+	}
+	for _, c := range cases {
+		ts.T().Run(c.desc, func(t *testing.T) {
+			ts.Config.Hook.CustomAccessToken.Enabled = true
+			ts.Config.Hook.CustomAccessToken.URI = c.uri
+			require.NoError(t, ts.Config.Hook.CustomAccessToken.ValidateAndPopulateExtensibilityPoint())
+
+			err := ts.API.db.RawQuery(c.hookFunctionSQL).Exec()
+			require.NoError(t, err)
+
+			var buffer bytes.Buffer
+			require.NoError(t, json.NewEncoder(&buffer).Encode(map[string]interface{}{
+				"refresh_token": ts.RefreshToken.Token,
+			}))
+
+			req := httptest.NewRequest(http.MethodPost, "http://localhost/token?grant_type=refresh_token", &buffer)
+			req.Header.Set("Content-Type", "application/json")
+
+			w := httptest.NewRecorder()
+			ts.API.handler.ServeHTTP(w, req)
+
+			var tokenResponse struct {
+				AccessToken string `json:"access_token"`
+			}
+			require.NoError(t, json.NewDecoder(w.Result().Body).Decode(&tokenResponse))
+			if c.shouldError {
+				require.Equal(t, http.StatusInternalServerError, w.Code)
+			} else {
+				parts := strings.Split(tokenResponse.AccessToken, ".")
+				require.Equal(t, 3, len(parts), "Token should have 3 parts")
+
+				payload, err := base64.RawURLEncoding.DecodeString(parts[1])
+				require.NoError(t, err)
+
+				var responseClaims map[string]interface{}
+				require.NoError(t, json.Unmarshal(payload, &responseClaims))
+
+				for key, expectedValue := range c.expectedClaims {
+					if expectedValue == nil {
+						// Since c.shouldError is false here, we only need to check if the claim should be removed
+						_, exists := responseClaims[key]
+						assert.False(t, exists, "Claim should be removed")
+					} else {
+						assert.Equal(t, expectedValue, responseClaims[key])
+					}
+				}
+			}
+
+			cleanupHookSQL := fmt.Sprintf("drop function if exists %s", ts.Config.Hook.CustomAccessToken.HookName)
+			require.NoError(t, ts.API.db.RawQuery(cleanupHookSQL).Exec())
+			ts.Config.Hook.CustomAccessToken.Enabled = false
+		})
+	}
+}
diff --git a/internal/api/user_test.go b/internal/api/user_test.go
index 16230c4b..f6af83b9 100644
--- a/internal/api/user_test.go
+++ b/internal/api/user_test.go
@@ -47,7 +47,7 @@ func (ts *UserTestSuite) SetupTest() {
 }
 
 func (ts *UserTestSuite) generateToken(user *models.User, sessionId *uuid.UUID) string {
-	token, _, err := generateAccessToken(ts.API.db, user, sessionId, &ts.Config.JWT)
+	token, _, err := ts.API.generateAccessToken(context.Background(), ts.API.db, user, sessionId, models.PasswordGrant)
 	require.NoError(ts.T(), err, "Error generating access token")
 	return token
 }
diff --git a/internal/api/verify_test.go b/internal/api/verify_test.go
index 6c93e07c..4640cfdd 100644
--- a/internal/api/verify_test.go
+++ b/internal/api/verify_test.go
@@ -2,6 +2,7 @@ package api
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -178,7 +179,7 @@ func (ts *VerifyTestSuite) TestVerifySecureEmailChange() {
 
 		// Generate access token for request
 		var token string
-		token, _, err = generateAccessToken(ts.API.db, u, nil, &ts.Config.JWT)
+		token, _, err = ts.API.generateAccessToken(context.Background(), ts.API.db, u, nil, models.MagicLink)
 		require.NoError(ts.T(), err)
 		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))