feat(auth): verifiers on one axis, and a reusable authenticator constructor

packages/auth/pkg/auth/internal/middleware/middleware.go

@@ -148,6 +148,59 @@ func adminValidationFunction(adminToken string) func(ctx context.Context, ginCtx
 	}
 }
 
+// AuthenticatorConfig describes a header-token security scheme.
+//
+// The constructors below cover the schemes this package defines, each with
+// its scheme name and context key fixed. This exists for a service that
+// defines its own — one that verifies a token under a scheme of its own name,
+// or records something other than a user or a team.
+//
+// Without it such a service reimplements the header handling and the 401
+// stamping below, which is how a scheme ends up subtly different from every
+// other one rather than merely differently named.
+type AuthenticatorConfig[T any] struct {
+	// SchemeName must match the securityScheme in the service's OpenAPI
+	// document; the validator dispatches on it.
+	SchemeName string
+
+	// Header is the header carrying the token.
+	Header string
+
+	// RequiredPrefix, when set, is a prefix the token must carry for this
+	// scheme to apply. A token without it is left to another authenticator.
+	RequiredPrefix string
+
+	// StrippedPrefix is removed before validation, e.g. "Bearer ".
+	StrippedPrefix string
+
+	// Validate turns a token into whatever the scheme establishes, or an
+	// APIError carrying the status the caller should see.
+	Validate func(ctx context.Context, ginCtx *gin.Context, token string) (T, *APIError)
+
+	// SetContext records the result for handlers. Optional: a scheme that
+	// only proves the caller may proceed has nothing to record.
+	SetContext func(ginCtx *gin.Context, value T)
+
+	// ErrorMessage prefixes the failure returned to the validator.
+	ErrorMessage string
+}
+
+// NewAuthenticator builds an Authenticator for a scheme this package does not
+// name itself.
+func NewAuthenticator[T any](config AuthenticatorConfig[T]) Authenticator {
+	return &commonAuthenticator[T]{
+		schemeName: config.SchemeName,
+		header: headerKey{
+			name:         config.Header,
+			prefix:       config.RequiredPrefix,
+			removePrefix: config.StrippedPrefix,
+		},
+		validationFunc: config.Validate,
+		setContextFunc: config.SetContext,
+		errorMessage:   config.ErrorMessage,
+	}
+}
+
 // NewApiKeyAuthenticator creates an authenticator for the ApiKeyAuth security scheme (X-API-Key header, e2b_ prefix).
 func NewApiKeyAuthenticator(validationFunc func(ctx context.Context, ginCtx *gin.Context, token string) (*types.Team, *APIError)) Authenticator {
 	return &commonAuthenticator[*types.Team]{
@@ -205,7 +258,7 @@ func NewAuthProviderTeamAuthenticator(validationFunc func(ctx context.Context, g
 }
 
 // NewAdminJWTAuthenticator creates an authenticator for the AdminJWTAuth security scheme.
-func NewAdminJWTAuthenticator(verifier *token.AdminVerifier) Authenticator {
+func NewAdminJWTAuthenticator(verifier *token.JWKSVerifier) Authenticator {
 	return &commonAuthenticator[struct{}]{
 		schemeName: "AdminJWTAuth",
 		header: headerKey{

packages/auth/pkg/auth/internal/middleware/middleware_test.go

@@ -79,3 +79,115 @@ func TestAdminTeamAuthenticatorSetsTeamContext(t *testing.T) {
 		t.Errorf("authcontext.GetTeamInfo(ginCtx).Team.ID = %s, want %s", got.Team.ID, teamID)
 	}
 }
+
+// A service defining its own scheme gets the same header handling and the
+// same 401 stamping as the named ones, and can record something other than a
+// user or a team.
+func TestNewAuthenticatorAppliesTheConfiguredScheme(t *testing.T) {
+	t.Parallel()
+
+	const subjectKey = "provider_subject"
+
+	authenticator := NewAuthenticator(AuthenticatorConfig[string]{
+		SchemeName:     "CustomBearerAuth",
+		Header:         HeaderAuthorization,
+		StrippedPrefix: PrefixBearer,
+		Validate: func(_ context.Context, _ *gin.Context, token string) (string, *APIError) {
+			if token != "good-token" {
+				return "", &APIError{Err: ErrInvalidAuthHeader, ClientMsg: "nope", Code: http.StatusUnauthorized}
+			}
+
+			return "subject-1", nil
+		},
+		SetContext:   func(c *gin.Context, subject string) { c.Set(subjectKey, subject) },
+		ErrorMessage: "Invalid custom token.",
+	})
+
+	require.Equal(t, "CustomBearerAuth", authenticator.SecuritySchemeName())
+
+	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", nil)
+	req.Header.Set(HeaderAuthorization, PrefixBearer+"good-token")
+	ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
+
+	require.NoError(t, authenticator.Authenticate(t.Context(), ginCtx, &openapi3filter.AuthenticationInput{
+		RequestValidationInput: &openapi3filter.RequestValidationInput{Request: req},
+	}))
+
+	subject, ok := ginCtx.Get(subjectKey)
+	require.True(t, ok)
+	require.Equal(t, "subject-1", subject)
+}
+
+// The prefix is stripped before validation, so a scheme sharing the
+// Authorization header does not have to strip it again.
+func TestNewAuthenticatorStripsTheConfiguredPrefix(t *testing.T) {
+	t.Parallel()
+
+	var seen string
+	authenticator := NewAuthenticator(AuthenticatorConfig[struct{}]{
+		SchemeName:     "CustomBearerAuth",
+		Header:         HeaderAuthorization,
+		StrippedPrefix: PrefixBearer,
+		Validate: func(_ context.Context, _ *gin.Context, token string) (struct{}, *APIError) {
+			seen = token
+
+			return struct{}{}, nil
+		},
+	})
+
+	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", nil)
+	req.Header.Set(HeaderAuthorization, PrefixBearer+"raw-token")
+	ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
+
+	require.NoError(t, authenticator.Authenticate(t.Context(), ginCtx, &openapi3filter.AuthenticationInput{
+		RequestValidationInput: &openapi3filter.RequestValidationInput{Request: req},
+	}))
+	require.Equal(t, "raw-token", seen)
+}
+
+// A missing header stamps 401 rather than leaving the validator's 400
+// fallback to win, which is what makes an auth failure look like one.
+func TestNewAuthenticatorStamps401OnAMissingHeader(t *testing.T) {
+	t.Parallel()
+
+	authenticator := NewAuthenticator(AuthenticatorConfig[struct{}]{
+		SchemeName: "CustomBearerAuth",
+		Header:     HeaderAuthorization,
+		Validate: func(context.Context, *gin.Context, string) (struct{}, *APIError) {
+			return struct{}{}, nil
+		},
+	})
+
+	recorder := httptest.NewRecorder()
+	ginCtx, _ := gin.CreateTestContext(recorder)
+
+	err := authenticator.Authenticate(t.Context(), ginCtx, &openapi3filter.AuthenticationInput{
+		RequestValidationInput: &openapi3filter.RequestValidationInput{
+			Request: httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", nil),
+		},
+	})
+	require.Error(t, err)
+	require.Equal(t, http.StatusUnauthorized, ginCtx.Writer.Status())
+}
+
+// SetContext is optional: a scheme that only proves the caller may proceed
+// has nothing to record, and must not require a setter to say so.
+func TestNewAuthenticatorAllowsNoContextSetter(t *testing.T) {
+	t.Parallel()
+
+	authenticator := NewAuthenticator(AuthenticatorConfig[struct{}]{
+		SchemeName: "CustomBearerAuth",
+		Header:     HeaderAdminToken,
+		Validate: func(context.Context, *gin.Context, string) (struct{}, *APIError) {
+			return struct{}{}, nil
+		},
+	})
+
+	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", nil)
+	req.Header.Set(HeaderAdminToken, "anything")
+	ginCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
+
+	require.NoError(t, authenticator.Authenticate(t.Context(), ginCtx, &openapi3filter.AuthenticationInput{
+		RequestValidationInput: &openapi3filter.RequestValidationInput{Request: req},
+	}))
+}

packages/auth/pkg/auth/internal/service/service.go

@@ -51,7 +51,7 @@ type Service interface {
 type AuthService struct {
 	store                authStore
 	teamCache            *authCache
-	authProviderVerifier *token.ProviderVerifier
+	authProviderVerifier *token.LinkedOIDCVerifier
 }
 
 // Compile-time assertion that *AuthService satisfies the Service interface.
@@ -82,7 +82,7 @@ func NewAuthService(
 	// OIDC bootstrap writes identity rows on the primary immediately before the
 	// next authenticated request; using the read replica here races replication lag.
 	identityLookup := newAuthIdentityLookup(authDB.Queries)
-	v, err := token.NewProviderVerifier(ctx, providerConfig, httpClient, identityLookup)
+	v, err := token.NewLinkedOIDCVerifier(ctx, providerConfig, httpClient, identityLookup)
 	if err != nil {
 		return nil, fmt.Errorf("initializing auth provider JWT verifier: %w", err)
 	}
@@ -188,7 +188,7 @@ func (s *AuthService) ValidateAuthProviderToken(ctx context.Context, ginCtx *gin
 	return s.validateJWTWithProvider(ctx, ginCtx, s.authProviderVerifier, token, "auth provider")
 }
 
-func (s *AuthService) validateJWTWithProvider(ctx context.Context, ginCtx *gin.Context, v *token.ProviderVerifier, token string, tokenSource string) (uuid.UUID, *APIError) {
+func (s *AuthService) validateJWTWithProvider(ctx context.Context, ginCtx *gin.Context, v *token.LinkedOIDCVerifier, token string, tokenSource string) (uuid.UUID, *APIError) {
 	userID, _, err := v.Verify(ctx, token)
 	if err != nil {
 		return uuid.UUID{}, &APIError{

packages/auth/pkg/auth/internal/token/jwks_verifier.go

@@ -0,0 +1,75 @@
+package token
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/golang-jwt/jwt/v5"
+
+	"github.com/e2b-dev/infra/packages/auth/pkg/auth/internal/token/jwks"
+)
+
+// jwksClockSkew is the leeway applied to time-based claims. Service
+// tokens are short-lived, so a clock a little out of step would otherwise
+// reject one that is legitimately current.
+const jwksClockSkew = 30 * time.Second
+
+// JWKSVerifier verifies JWTs against one or more configured issuers
+// and returns the first successful verification.
+//
+// Keys come from each issuer's conventional JWKS path rather than an OIDC
+// discovery document, which suits a token minted by a peer service: there is
+// no discovery to perform, and no separate issuer declaration to cross-check.
+type JWKSVerifier struct {
+	verifiers []*jwks.Verifier
+}
+
+// NewJWKSVerifier builds a verifier from the same ProviderConfig shape
+// used for AUTH_PROVIDER_CONFIG. It returns nil when the config declares no
+// issuers, leaving whichever scheme uses it unconfigured.
+//
+// Backs AdminJWTAuth today; nothing about it is specific to that scheme.
+func NewJWKSVerifier(ctx context.Context, config ProviderConfig, httpClient *http.Client) (*JWKSVerifier, error) {
+	normalized := config.normalize()
+	if !normalized.enabled() {
+		return nil, nil
+	}
+
+	verifiers := make([]*jwks.Verifier, 0, len(normalized.JWT))
+	for i, entry := range normalized.JWT {
+		verifier, err := jwks.NewVerifierFromIssuerJWKS(ctx, entry, httpClient,
+			jwks.WithParserOptions(jwt.WithLeeway(jwksClockSkew)),
+		)
+		if err != nil {
+			return nil, fmt.Errorf("service token jwt[%d]: %w", i, err)
+		}
+		verifiers = append(verifiers, verifier)
+	}
+
+	return &JWKSVerifier{verifiers: verifiers}, nil
+}
+
+// Verify iterates over the configured issuers and returns the claims of the
+// first successful verification.
+func (v *JWKSVerifier) Verify(ctx context.Context, tokenString string) (jwt.MapClaims, error) {
+	if v == nil || len(v.verifiers) == 0 {
+		return nil, errors.New("service token verifier is not configured")
+	}
+
+	errs := make([]error, 0, len(v.verifiers))
+	for _, verifier := range v.verifiers {
+		claims, err := verifier.Verify(ctx, tokenString)
+		if err != nil {
+			errs = append(errs, err)
+
+			continue
+		}
+
+		return claims, nil
+	}
+
+	return nil, fmt.Errorf("failed to verify service token: %w", errors.Join(errs...))
+}