From 3cf9670cb61d01957fae58ec4029a402f815f76b Mon Sep 17 00:00:00 2001 From: mo khan Date: Tue, 4 Aug 2026 15:35:02 -0600 Subject: [PATCH] chore: introduce typed context keys * Add `xcontext.Key[T]`, a generic context key that carries its own With/From accessors. * Drops functionHooksKey, which had no readers or writers. * Where a wrongly typed value once panicked it now returns the zero value. --- internal/api/apitask/apitask.go | 7 +- internal/api/context.go | 169 +++++++++----------------------- internal/api/shared/context.go | 50 +++------- internal/ctxkey/key.go | 36 +++++++ internal/ctxkey/key_test.go | 70 +++++++++++++ internal/sbff/sbff.go | 9 +- internal/utilities/context.go | 21 +--- 7 files changed, 179 insertions(+), 183 deletions(-) create mode 100644 internal/ctxkey/key.go create mode 100644 internal/ctxkey/key_test.go diff --git a/internal/api/apitask/apitask.go b/internal/api/apitask/apitask.go index b5e2b5731e..99e5f1cc40 100644 --- a/internal/api/apitask/apitask.go +++ b/internal/api/apitask/apitask.go @@ -11,6 +11,7 @@ import ( "github.com/sirupsen/logrus" "github.com/supabase/auth/internal/api/apierrors" + "github.com/supabase/auth/internal/ctxkey" "github.com/supabase/auth/internal/observability" ) @@ -81,13 +82,13 @@ func With(ctx context.Context) context.Context { if !ok { wrk = &requestWorker{} } - return context.WithValue(ctx, ctxKey, wrk) + return ctxKey.WithValue(ctx, wrk) } -var ctxKey = new(int) +var ctxKey = ctxkey.New[*requestWorker]("apitask_worker") func from(ctx context.Context) (*requestWorker, bool) { - if st, ok := ctx.Value(ctxKey).(*requestWorker); ok && st != nil { + if st, ok := ctxKey.Lookup(ctx); ok && st != nil { return st, true } return nil, false diff --git a/internal/api/context.go b/internal/api/context.go index f8367a4abc..de87ba009b 100644 --- a/internal/api/context.go +++ b/internal/api/context.go @@ -7,49 +7,38 @@ import ( "github.com/gofrs/uuid" jwt "github.com/golang-jwt/jwt/v5" "github.com/supabase/auth/internal/api/shared" + "github.com/supabase/auth/internal/ctxkey" "github.com/supabase/auth/internal/models" ) -type contextKey string - -func (c contextKey) String() string { - return "gotrue api context key " + string(c) -} - -const ( - externalProviderTypeKey = contextKey("external_provider_type") - externalProviderEmailOptionalKey = contextKey("external_provider_allow_no_email") - - tokenKey = contextKey("jwt") - inviteTokenKey = contextKey("invite_token") - signatureKey = contextKey("signature") - targetUserKey = contextKey("target_user") - factorKey = contextKey("factor") - sessionKey = contextKey("session") - externalReferrerKey = contextKey("external_referrer") - functionHooksKey = contextKey("function_hooks") - adminUserKey = contextKey("admin_user") - oauthTokenKey = contextKey("oauth_token") // for OAuth1.0, also known as request token - oauthVerifierKey = contextKey("oauth_verifier") - ssoProviderKey = contextKey("sso_provider") - externalHostKey = contextKey("external_host") - oauthClientStateKey = contextKey("oauth_client_state_id") - flowStateContextKey = contextKey("flow_state") +var ( + externalProviderTypeKey = ctxkey.New[string]("external_provider_type") + externalProviderEmailOptionalKey = ctxkey.New[bool]("external_provider_allow_no_email") + + tokenKey = ctxkey.New[*jwt.Token]("jwt") + inviteTokenKey = ctxkey.New[string]("invite_token") + signatureKey = ctxkey.New[string]("signature") + targetUserKey = ctxkey.New[*models.User]("target_user") + factorKey = ctxkey.New[*models.Factor]("factor") + sessionKey = ctxkey.New[*models.Session]("session") + externalReferrerKey = ctxkey.New[string]("external_referrer") + adminUserKey = ctxkey.New[*models.User]("admin_user") + oauthTokenKey = ctxkey.New[string]("oauth_token") // for OAuth1.0, also known as request token + oauthVerifierKey = ctxkey.New[string]("oauth_verifier") + ssoProviderKey = ctxkey.New[*models.SSOProvider]("sso_provider") + externalHostKey = ctxkey.New[*url.URL]("external_host") + oauthClientStateKey = ctxkey.New[uuid.UUID]("oauth_client_state_id") + flowStateContextKey = ctxkey.New[*models.FlowState]("flow_state") ) // withToken adds the JWT token to the context. func withToken(ctx context.Context, token *jwt.Token) context.Context { - return context.WithValue(ctx, tokenKey, token) + return tokenKey.WithValue(ctx, token) } // getToken reads the JWT token from the context. func getToken(ctx context.Context) *jwt.Token { - obj := ctx.Value(tokenKey) - if obj == nil { - return nil - } - - return obj.(*jwt.Token) + return tokenKey.Value(ctx) } func getClaims(ctx context.Context) *AccessTokenClaims { @@ -67,12 +56,12 @@ func withUser(ctx context.Context, u *models.User) context.Context { // withTargetUser adds the target user for linking to the context. func withTargetUser(ctx context.Context, u *models.User) context.Context { - return context.WithValue(ctx, targetUserKey, u) + return targetUserKey.WithValue(ctx, u) } // with Factor adds the factor id to the context. func withFactor(ctx context.Context, f *models.Factor) context.Context { - return context.WithValue(ctx, factorKey, f) + return factorKey.WithValue(ctx, f) } // getUser reads the user from the context. @@ -82,103 +71,68 @@ func getUser(ctx context.Context) *models.User { // getTargetUser reads the user from the context. func getTargetUser(ctx context.Context) *models.User { - if ctx == nil { - return nil - } - obj := ctx.Value(targetUserKey) - if obj == nil { - return nil - } - return obj.(*models.User) + return targetUserKey.Value(ctx) } // getFactor reads the factor id from the context func getFactor(ctx context.Context) *models.Factor { - obj := ctx.Value(factorKey) - if obj == nil { - return nil - } - return obj.(*models.Factor) + return factorKey.Value(ctx) } // withSession adds the session to the context. func withSession(ctx context.Context, s *models.Session) context.Context { - return context.WithValue(ctx, sessionKey, s) + return sessionKey.WithValue(ctx, s) } // getSession reads the session from the context. func getSession(ctx context.Context) *models.Session { - if ctx == nil { - return nil - } - obj := ctx.Value(sessionKey) - if obj == nil { - return nil - } - return obj.(*models.Session) + return sessionKey.Value(ctx) } // withSignature adds the provided request ID to the context. func withSignature(ctx context.Context, id string) context.Context { - return context.WithValue(ctx, signatureKey, id) + return signatureKey.WithValue(ctx, id) } func withInviteToken(ctx context.Context, token string) context.Context { - return context.WithValue(ctx, inviteTokenKey, token) + return inviteTokenKey.WithValue(ctx, token) } func withOAuthClientStateID(ctx context.Context, oauthClientStateID uuid.UUID) context.Context { - return context.WithValue(ctx, oauthClientStateKey, oauthClientStateID) + return oauthClientStateKey.WithValue(ctx, oauthClientStateID) } func getOAuthClientStateID(ctx context.Context) uuid.UUID { - obj := ctx.Value(oauthClientStateKey) - if obj == nil { - return uuid.Nil - } - - return obj.(uuid.UUID) + return oauthClientStateKey.Value(ctx) } // withFlowState stores the entire FlowState object in the context func withFlowState(ctx context.Context, flowState *models.FlowState) context.Context { - return context.WithValue(ctx, flowStateContextKey, flowState) + return flowStateContextKey.WithValue(ctx, flowState) } // getFlowState retrieves the FlowState object from the context func getFlowState(ctx context.Context) *models.FlowState { - obj := ctx.Value(flowStateContextKey) - if obj == nil { - return nil - } - return obj.(*models.FlowState) + return flowStateContextKey.Value(ctx) } func getInviteToken(ctx context.Context) string { - obj := ctx.Value(inviteTokenKey) - if obj == nil { - return "" - } - - return obj.(string) + return inviteTokenKey.Value(ctx) } // withExternalProviderType adds the provided request ID to the context. func withExternalProviderType(ctx context.Context, id string, emailOptional bool) context.Context { - return context.WithValue(context.WithValue(ctx, externalProviderTypeKey, id), externalProviderEmailOptionalKey, emailOptional) + return externalProviderEmailOptionalKey.WithValue(externalProviderTypeKey.WithValue(ctx, id), emailOptional) } // getExternalProviderType returns the provider type and whether user data without email address should be allowed. func getExternalProviderType(ctx context.Context) (string, bool) { - idValue := ctx.Value(externalProviderTypeKey) - emailOptionalValue := ctx.Value(externalProviderEmailOptionalKey) - - id, okID := idValue.(string) + id, okID := externalProviderTypeKey.Lookup(ctx) if !okID { return "", false } - emailOptional, okEmailOptional := emailOptionalValue.(bool) + emailOptional, okEmailOptional := externalProviderEmailOptionalKey.Lookup(ctx) if !okEmailOptional { return "", false } @@ -187,77 +141,52 @@ func getExternalProviderType(ctx context.Context) (string, bool) { } func withExternalReferrer(ctx context.Context, token string) context.Context { - return context.WithValue(ctx, externalReferrerKey, token) + return externalReferrerKey.WithValue(ctx, token) } func getExternalReferrer(ctx context.Context) string { - obj := ctx.Value(externalReferrerKey) - if obj == nil { - return "" - } - - return obj.(string) + return externalReferrerKey.Value(ctx) } // withAdminUser adds the admin user to the context. func withAdminUser(ctx context.Context, u *models.User) context.Context { - return context.WithValue(ctx, adminUserKey, u) + return adminUserKey.WithValue(ctx, u) } // getAdminUser reads the admin user from the context. func getAdminUser(ctx context.Context) *models.User { - obj := ctx.Value(adminUserKey) - if obj == nil { - return nil - } - return obj.(*models.User) + return adminUserKey.Value(ctx) } // withRequestToken adds the request token to the context func withRequestToken(ctx context.Context, token string) context.Context { - return context.WithValue(ctx, oauthTokenKey, token) + return oauthTokenKey.WithValue(ctx, token) } func getRequestToken(ctx context.Context) string { - obj := ctx.Value(oauthTokenKey) - if obj == nil { - return "" - } - return obj.(string) + return oauthTokenKey.Value(ctx) } func withOAuthVerifier(ctx context.Context, token string) context.Context { - return context.WithValue(ctx, oauthVerifierKey, token) + return oauthVerifierKey.WithValue(ctx, token) } func getOAuthVerifier(ctx context.Context) string { - obj := ctx.Value(oauthVerifierKey) - if obj == nil { - return "" - } - return obj.(string) + return oauthVerifierKey.Value(ctx) } func withSSOProvider(ctx context.Context, provider *models.SSOProvider) context.Context { - return context.WithValue(ctx, ssoProviderKey, provider) + return ssoProviderKey.WithValue(ctx, provider) } func getSSOProvider(ctx context.Context) *models.SSOProvider { - obj := ctx.Value(ssoProviderKey) - if obj == nil { - return nil - } - return obj.(*models.SSOProvider) + return ssoProviderKey.Value(ctx) } func withExternalHost(ctx context.Context, u *url.URL) context.Context { - return context.WithValue(ctx, externalHostKey, u) + return externalHostKey.WithValue(ctx, u) } func getExternalHost(ctx context.Context) *url.URL { - obj := ctx.Value(externalHostKey) - if obj == nil { - return nil - } - return obj.(*url.URL) + return externalHostKey.Value(ctx) } diff --git a/internal/api/shared/context.go b/internal/api/shared/context.go index 81bfd47521..57dba243e6 100644 --- a/internal/api/shared/context.go +++ b/internal/api/shared/context.go @@ -3,70 +3,42 @@ package shared import ( "context" + "github.com/supabase/auth/internal/ctxkey" "github.com/supabase/auth/internal/models" ) -// ContextKey is the type for context keys to avoid collisions -type ContextKey string - -func (c ContextKey) String() string { - return "gotrue api context key " + string(c) -} - -// Context keys used across packages -const ( - UserKey ContextKey = "user" - SessionKey ContextKey = "session" - OAuthServerClientKey ContextKey = "oauth_server_client" +var ( + UserKey = ctxkey.New[*models.User]("user") + SessionKey = ctxkey.New[*models.Session]("session") + OAuthServerClientKey = ctxkey.New[*models.OAuthServerClient]("oauth_server_client") ) // GetUser reads the user from the context - shared implementation func GetUser(ctx context.Context) *models.User { - if ctx == nil { - return nil - } - obj := ctx.Value(UserKey) - if obj == nil { - return nil - } - return obj.(*models.User) + return UserKey.Value(ctx) } // WithUser adds the user to the context - shared implementation func WithUser(ctx context.Context, u *models.User) context.Context { - return context.WithValue(ctx, UserKey, u) + return UserKey.WithValue(ctx, u) } // GetSession reads the session from the context - shared implementation func GetSession(ctx context.Context) *models.Session { - if ctx == nil { - return nil - } - obj := ctx.Value(SessionKey) - if obj == nil { - return nil - } - return obj.(*models.Session) + return SessionKey.Value(ctx) } // WithSession adds the session to the context - shared implementation func WithSession(ctx context.Context, s *models.Session) context.Context { - return context.WithValue(ctx, SessionKey, s) + return SessionKey.WithValue(ctx, s) } // WithOAuthServerClient adds an OAuth server client to the context func WithOAuthServerClient(ctx context.Context, client *models.OAuthServerClient) context.Context { - return context.WithValue(ctx, OAuthServerClientKey, client) + return OAuthServerClientKey.WithValue(ctx, client) } // GetOAuthServerClient retrieves an OAuth server client from the context func GetOAuthServerClient(ctx context.Context) *models.OAuthServerClient { - if ctx == nil { - return nil - } - obj := ctx.Value(OAuthServerClientKey) - if obj == nil { - return nil - } - return obj.(*models.OAuthServerClient) + return OAuthServerClientKey.Value(ctx) } diff --git a/internal/ctxkey/key.go b/internal/ctxkey/key.go new file mode 100644 index 0000000000..7c54d57740 --- /dev/null +++ b/internal/ctxkey/key.go @@ -0,0 +1,36 @@ +// Package ctxkey provides a generic, strongly typed context key. Each key +// owns its accessors, so the type stored under a key and the type read back +// out are checked by the compiler rather than by hand at every call site. +package ctxkey + +import "context" + +type Key[T any] struct { + name string +} + +func New[T any](name string) *Key[T] { + return &Key[T]{name: name} +} + +func (k *Key[T]) String() string { + return k.name +} + +func (k *Key[T]) WithValue(ctx context.Context, value T) context.Context { + return context.WithValue(ctx, k, value) +} + +func (k *Key[T]) Lookup(ctx context.Context) (T, bool) { + var zero T + if ctx == nil { + return zero, false + } + value, ok := ctx.Value(k).(T) + return value, ok +} + +func (k *Key[T]) Value(ctx context.Context) T { + value, _ := k.Lookup(ctx) + return value +} diff --git a/internal/ctxkey/key_test.go b/internal/ctxkey/key_test.go new file mode 100644 index 0000000000..be55c5f6c0 --- /dev/null +++ b/internal/ctxkey/key_test.go @@ -0,0 +1,70 @@ +package ctxkey + +import ( + "context" + "fmt" + "testing" + + "github.com/gofrs/uuid" + "github.com/stretchr/testify/require" +) + +func TestKey(t *testing.T) { + type example struct { + ID uuid.UUID `db:"id" json:"id"` + } + key := New[*example]("example") + + t.Run("round trips a typed value", func(t *testing.T) { + item := &example{ID: uuid.Must(uuid.NewV4())} + ctx := key.WithValue(context.Background(), item) + require.Equal(t, item, key.Value(ctx)) + }) + + t.Run("returns the zero value when absent", func(t *testing.T) { + require.Nil(t, key.Value(context.Background())) + }) + + t.Run("returns the zero value for a nil context", func(t *testing.T) { + var ctx context.Context + require.Nil(t, key.Value(ctx)) + }) + + t.Run("Lookup distinguishes an absent value from a zero value", func(t *testing.T) { + flag := New[bool]("email_optional") + + value, ok := flag.Lookup(context.Background()) + require.False(t, value) + require.False(t, ok) + + value, ok = flag.Lookup(flag.WithValue(context.Background(), false)) + require.False(t, value) + require.True(t, ok) + }) + + t.Run("keys of different types do not collide", func(t *testing.T) { + other := New[string]("example") + ctx := other.WithValue(context.Background(), "not an item") + require.Nil(t, key.Value(ctx)) + require.Equal(t, "not an item", other.Value(ctx)) + }) + + t.Run("independently constructed keys never share a slot, even with the same type and name", func(t *testing.T) { + item := &example{ID: uuid.Must(uuid.NewV4())} + ctx := New[*example]("example").WithValue(context.Background(), item) + require.Nil(t, key.Value(ctx)) + }) + + t.Run("keys of the same type but a different name do not collide", func(t *testing.T) { + first, second := New[string]("first"), New[string]("second") + ctx := first.WithValue(context.Background(), "one") + ctx = second.WithValue(ctx, "two") + require.Equal(t, "one", first.Value(ctx)) + require.Equal(t, "two", second.Value(ctx)) + }) + + t.Run("names the key when a context is printed", func(t *testing.T) { + ctx := New[string]("first").WithValue(context.Background(), "one") + require.Contains(t, fmt.Sprint(ctx), "first") + }) +} diff --git a/internal/sbff/sbff.go b/internal/sbff/sbff.go index 33a5126643..032b89ae67 100644 --- a/internal/sbff/sbff.go +++ b/internal/sbff/sbff.go @@ -1,13 +1,13 @@ package sbff import ( - "context" "errors" "net" "net/http" "strings" "github.com/supabase/auth/internal/conf" + "github.com/supabase/auth/internal/ctxkey" ) // HeaderName is the Sb-Forwarded-For header name. It is all lowercase here as HTTP header names @@ -15,7 +15,7 @@ import ( const HeaderName = "sb-forwarded-for" var ( - ctxKeySBFF = &struct{}{} + ctxKeySBFF = ctxkey.New[string]("sbff_ip_address") ErrHeaderNotFound = errors.New("Sb-Forwarded-For header not found") ErrHeaderInvalid = errors.New("invalid Sb-Forwarded-For header value") @@ -35,7 +35,7 @@ func parseSBFFHeader(headerVal string) (string, error) { // SBForwardedForMiddleware. If no value is present in the request context, this function will // return ("", false). func GetIPAddress(r *http.Request) (addr string, found bool) { - if ipAddr, ok := r.Context().Value(ctxKeySBFF).(string); ok && ipAddr != "" { + if ipAddr := ctxKeySBFF.Value(r.Context()); ipAddr != "" { return ipAddr, true } @@ -57,8 +57,7 @@ func withIPAddress(r *http.Request) (*http.Request, error) { return nil, err } - ctx := r.Context() - newCtx := context.WithValue(ctx, ctxKeySBFF, parsedIPAddr) + newCtx := ctxKeySBFF.WithValue(r.Context(), parsedIPAddr) out := r.WithContext(newCtx) return out, nil diff --git a/internal/utilities/context.go b/internal/utilities/context.go index 06aa74a396..54248a0ed1 100644 --- a/internal/utilities/context.go +++ b/internal/utilities/context.go @@ -3,31 +3,20 @@ package utilities import ( "context" "sync" -) - -type contextKey string - -func (c contextKey) String() string { - return "gotrue api context key " + string(c) -} -const ( - requestIDKey = contextKey("request_id") + "github.com/supabase/auth/internal/ctxkey" ) +var requestIDKey = ctxkey.New[string]("request_id") + // WithRequestID adds the provided request ID to the context. func WithRequestID(ctx context.Context, id string) context.Context { - return context.WithValue(ctx, requestIDKey, id) + return requestIDKey.WithValue(ctx, id) } // GetRequestID reads the request ID from the context. func GetRequestID(ctx context.Context) string { - obj := ctx.Value(requestIDKey) - if obj == nil { - return "" - } - - return obj.(string) + return requestIDKey.Value(ctx) } // WaitForCleanup waits until all long-running goroutines shut