Move all auth code to a common utility package (#897)

Video description:
https://drive.google.com/file/d/1C-jso8ZOhHZJ5Qw7C1NbRJTx__7BO9RY/view?usp=sharing

GitOrigin-RevId: 1f6bf0cd0537d00a147e47a7c9098bd0f33f103b

Author: Dylan

sdp-go/auth/auth.go auth/auth.gosdp-go/auth/auth.go renamed to auth/auth.go

@@ -90,8 +90,14 @@ func WithImpersonateAccount(account string) TokenSourceOptionsFunc {
 // Cache this between invocations to avoid additional charges by Auth0 for M2M
 // tokens. The oAuthTokenURL looks like this:
 // https://somedomain.auth0.com/oauth/token
-func (flowConfig ClientCredentialsConfig) TokenSource(oAuthTokenURL, oAuthAudience string, opts ...TokenSourceOptionsFunc) oauth2.TokenSource {
-	ctx := context.Background()
+//
+// The context that is passed to this function is used when getting new tokens,
+// which will happen initially, and then subsequently when the token expires.
+// This means that if this token source is going to be stored and used for many
+// requests, it should not use the context of the request that created it, as
+// this will be cancelled. Instead it should probably use `context.Background()`
+// or similar.
+func (flowConfig ClientCredentialsConfig) TokenSource(ctx context.Context, oAuthTokenURL, oAuthAudience string, opts ...TokenSourceOptionsFunc) oauth2.TokenSource {
 	// inject otel into oauth2
 	ctx = context.WithValue(ctx, oauth2.HTTPClient, otelhttp.DefaultClient)
 
@@ -116,47 +122,6 @@ func (flowConfig ClientCredentialsConfig) TokenSource(oAuthTokenURL, oAuthAudien
 	return conf.TokenSource(ctx)
 }
 
-// NewOAuthTokenClient creates a token client that uses the provided TokenSource
-// to get a NATS token. `overmindAPIURL` is the root URL of the NATS token
-// exchange API that will be used e.g. https://api.server.test/v1
-//
-// Tokens will be minted under the specified account as long as the client has
-// admin permissions, if not, the account that is attached to the client via
-// Auth0 metadata will be used
-func NewOAuthTokenClient(overmindAPIURL string, account string, ts oauth2.TokenSource) *natsTokenClient {
-	return NewOAuthTokenClientWithContext(context.Background(), overmindAPIURL, account, ts)
-}
-
-// NewOAuthTokenClientWithContext creates a token client that uses the provided
-// TokenSource to get a NATS token. `overmindAPIURL` is the root URL of the NATS
-// token exchange API that will be used e.g. https://api.server.test/v1
-//
-// Tokens will be minted under the specified account as long as the client has
-// admin permissions, if not, the account that is attached to the client via
-// Auth0 metadata will be used
-//
-// The provided context is used for cancellation and to lookup the HTTP client
-// used by oauth2. See the oauth2.HTTPClient variable.
-//
-// Provide an account name and an admin token to create a token client for a
-// foreign account.
-func NewOAuthTokenClientWithContext(ctx context.Context, overmindAPIURL string, account string, ts oauth2.TokenSource) *natsTokenClient {
-	authenticatedClient := oauth2.NewClient(ctx, ts)
-
-	// backwards compatibility: remove previously existing "/api" suffix from URL for connect
-	apiUrl, err := url.Parse(overmindAPIURL)
-	if err == nil {
-		apiUrl.Path = ""
-		overmindAPIURL = apiUrl.String()
-	}
-
-	return &natsTokenClient{
-		Account:     account,
-		adminClient: sdpconnect.NewAdminServiceClient(authenticatedClient, overmindAPIURL),
-		mgmtClient:  sdpconnect.NewManagementServiceClient(authenticatedClient, overmindAPIURL),
-	}
-}
-
 // natsTokenClient A client that is capable of getting NATS JWTs and signing the
 // required nonce to prove ownership of the NKeys. Satisfies the `TokenClient`
 // interface
@@ -414,3 +379,44 @@ func NewStaticTokenClient(overmindAPIURL, token, tokenType string) (*natsTokenCl
 		mgmtClient:  sdpconnect.NewManagementServiceClient(&httpClient, overmindAPIURL),
 	}, nil
 }
+
+// NewOAuthTokenClient creates a token client that uses the provided TokenSource
+// to get a NATS token. `overmindAPIURL` is the root URL of the NATS token
+// exchange API that will be used e.g. https://api.server.test/v1
+//
+// Tokens will be minted under the specified account as long as the client has
+// admin permissions, if not, the account that is attached to the client via
+// Auth0 metadata will be used
+func NewOAuthTokenClient(overmindAPIURL string, account string, ts oauth2.TokenSource) *natsTokenClient {
+	return NewOAuthTokenClientWithContext(context.Background(), overmindAPIURL, account, ts)
+}
+
+// NewOAuthTokenClientWithContext creates a token client that uses the provided
+// TokenSource to get a NATS token. `overmindAPIURL` is the root URL of the NATS
+// token exchange API that will be used e.g. https://api.server.test/v1
+//
+// Tokens will be minted under the specified account as long as the client has
+// admin permissions, if not, the account that is attached to the client via
+// Auth0 metadata will be used
+//
+// The provided context is used for cancellation and to lookup the HTTP client
+// used by oauth2. See the oauth2.HTTPClient variable.
+//
+// Provide an account name and an admin token to create a token client for a
+// foreign account.
+func NewOAuthTokenClientWithContext(ctx context.Context, overmindAPIURL string, account string, ts oauth2.TokenSource) *natsTokenClient {
+	authenticatedClient := oauth2.NewClient(ctx, ts)
+
+	// backwards compatibility: remove previously existing "/api" suffix from URL for connect
+	apiUrl, err := url.Parse(overmindAPIURL)
+	if err == nil {
+		apiUrl.Path = ""
+		overmindAPIURL = apiUrl.String()
+	}
+
+	return &natsTokenClient{
+		Account:     account,
+		adminClient: sdpconnect.NewAdminServiceClient(authenticatedClient, overmindAPIURL),
+		mgmtClient:  sdpconnect.NewManagementServiceClient(authenticatedClient, overmindAPIURL),
+	}
+}

sdp-go/auth/auth_client.go auth/auth_client.gosdp-go/auth/auth_client.go renamed to auth/auth_client.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/overmindtech/cli/sdp-go"
 	"github.com/overmindtech/cli/sdp-go/sdpconnect"
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
@@ -37,7 +36,7 @@ func (y *AuthenticatedTransport) RoundTrip(req *http.Request) (*http.Response, e
 // NewAuthenticatedClient creates a new AuthenticatedClient from the given
 // context and http.Client.
 func NewAuthenticatedClient(ctx context.Context, from *http.Client) *http.Client {
-	token, ok := ctx.Value(sdp.UserTokenContextKey{}).(string)
+	token, ok := ctx.Value(UserTokenContextKey{}).(string)
 	if !ok {
 		token = ""
 	}

sdp-go/auth/auth_test.go auth/auth_test.gosdp-go/auth/auth_test.go renamed to auth/auth_test.go

@@ -97,7 +97,7 @@ func GetTestOAuthTokenClient(t *testing.T) *natsTokenClient {
 	return NewOAuthTokenClient(
 		exchangeURL,
 		"overmind-development",
-		flowConfig.TokenSource(fmt.Sprintf("https://%v/oauth/token", domain), os.Getenv("API_SERVER_AUDIENCE")),
+		flowConfig.TokenSource(t.Context(), fmt.Sprintf("https://%v/oauth/token", domain), os.Getenv("API_SERVER_AUDIENCE")),
 	)
 }
 

sdp-go/middleware.go auth/middleware.gosdp-go/middleware.go renamed to auth/middleware.go

@@ -1,4 +1,4 @@
-package sdp
+package auth
 
 import (
 	"context"
@@ -18,11 +18,11 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
-// AuthBypassedContextKey is a key that is stored in the request context when auth is
-// actively being bypassed, e.g. in development. When this is set the
-// `HasScopes()` function will always return true, and can be set using the
-// `BypassAuth()` middleware.
-type AuthBypassedContextKey struct{}
+// ScopeCheckBypassedContextKey is a key that is stored in the request context
+// when scope checking is actively being bypassed, e.g. in development. When
+// this is set the `HasScopes()` function will always return true, and can be
+// set using the `WithBypassScopeCheck()` middleware.
+type ScopeCheckBypassedContextKey struct{}
 
 // CustomClaimsContextKey is the key that is used to store the custom claims
 // from the JWT
@@ -84,7 +84,7 @@ func HasAllScopes(ctx context.Context, requiredScopes ...string) bool {
 		attribute.StringSlice("ovm.auth.requiredScopes.all", requiredScopes),
 	)
 
-	if ctx.Value(AuthBypassedContextKey{}) == true {
+	if ctx.Value(ScopeCheckBypassedContextKey{}) == true {
 		// this is always set when auth is bypassed
 		// set it here again to capture non-standard auth configs
 		span.SetAttributes(attribute.Bool("ovm.auth.bypass", true))
@@ -116,7 +116,7 @@ func HasAnyScopes(ctx context.Context, requiredScopes ...string) bool {
 		attribute.StringSlice("ovm.auth.requiredScopes.any", requiredScopes),
 	)
 
-	if ctx.Value(AuthBypassedContextKey{}) == true {
+	if ctx.Value(ScopeCheckBypassedContextKey{}) == true {
 		// this is always set when auth is bypassed
 		// set it here again to capture non-standard auth configs
 		span.SetAttributes(attribute.Bool("ovm.auth.bypass", true))
@@ -169,7 +169,20 @@ func ExtractAccount(ctx context.Context) (string, error) {
 // must also be set.
 func NewAuthMiddleware(config AuthConfig, next http.Handler) http.Handler {
 	processOverrides := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := OverrideCustomClaims(r.Context(), config.ScopeOverride, config.AccountOverride)
+		options := []OverrideAuthOptionFunc{}
+
+		if config.ScopeOverride != nil {
+			options = append(options, WithScope(*config.ScopeOverride))
+		}
+
+		if config.AccountOverride != nil {
+			options = append(options, WithAccount(*config.AccountOverride))
+		}
+
+		ctx := r.Context()
+		if len(options) > 0 {
+			ctx = OverrideAuth(r.Context(), options...)
+		}
 
 		r = r.Clone(ctx)
 
@@ -179,51 +192,83 @@ func NewAuthMiddleware(config AuthConfig, next http.Handler) http.Handler {
 	return ensureValidTokenHandler(config, processOverrides)
 }
 
-// AddBypassAuthConfig Adds the requires keys to the context so that
-// authentication is bypassed. This is intended to be used in tests
-func AddBypassAuthConfig(ctx context.Context) context.Context {
-	return context.WithValue(ctx, AuthBypassedContextKey{}, true)
-}
+type OverrideAuthOptionFunc func(ctx context.Context) context.Context
 
-// OverrideAuthContext overrides the authentication data and token stored in the context.
-// This is mostly useful for testing or delegating access locally into a protected API.
-func OverrideAuthContext(ctx context.Context, claims *validator.ValidatedClaims) context.Context {
-	customClaims := claims.CustomClaims.(*CustomClaims)
-	ctx = context.WithValue(ctx, jwtmiddleware.ContextKey{}, claims)
-	ctx = context.WithValue(ctx, CustomClaimsContextKey{}, customClaims)
-	ctx = context.WithValue(ctx, CurrentSubjectContextKey{}, claims.RegisteredClaims.Subject)
-	ctx = context.WithValue(ctx, AccountNameContextKey{}, customClaims.AccountName)
-	return ctx
+// Sets the scope in the context to the given value. This should be the value
+// that would be embedded directly in the token, with each scope being separated
+// by a space.
+func WithScope(scope string) OverrideAuthOptionFunc {
+	return withCustomClaims(func(claims *CustomClaims) {
+		claims.Scope = scope
+	})
 }
 
-// OverrideCustomClaims Overrides the custom claims in the context that have
-// been set at CustomClaimsContextKey
-func OverrideCustomClaims(ctx context.Context, scope *string, account *string) context.Context {
-	// Read existing claims from the context
-	i := ctx.Value(CustomClaimsContextKey{})
-
-	var claims *CustomClaims
-	var newClaims CustomClaims
-	var ok bool
+// Sets the account in the context to the given value.
+func WithAccount(account string) OverrideAuthOptionFunc {
+	return withCustomClaims(func(claims *CustomClaims) {
+		claims.AccountName = account
+	})
+}
 
-	if claims, ok = i.(*CustomClaims); ok {
-		// clone out the values to avoid false sharing
-		newClaims = *claims
+// Sets the auth info in the context directly from the validated claims produced
+// by the `github.com/auth0/go-jwt-middleware/v2/validator` package. This is
+// essentially what the middleware already does when receiving a request, and
+// therefore should only be used in exceptional circumstances, like testing, when the
+// middleware is not being used.
+//
+// If this is being used, there is no need to use the `WithScope` or `WithAccount`
+// options as the claims will be extracted directly from the validated claims.
+func WithValidatedClaims(claims *validator.ValidatedClaims) OverrideAuthOptionFunc {
+	return func(ctx context.Context) context.Context {
+		customClaims := claims.CustomClaims.(*CustomClaims)
+		ctx = context.WithValue(ctx, jwtmiddleware.ContextKey{}, claims)
+		ctx = context.WithValue(ctx, CustomClaimsContextKey{}, customClaims)
+		ctx = context.WithValue(ctx, CurrentSubjectContextKey{}, claims.RegisteredClaims.Subject)
+		ctx = context.WithValue(ctx, AccountNameContextKey{}, customClaims.AccountName)
+		return ctx
 	}
+}
 
-	if scope != nil {
-		newClaims.Scope = *scope
+// Bypasses the scope check, meaning that `HasScopes()` and `HasAllScopes` will
+// always return true. This is useful for testing.
+func WithBypassScopeCheck() OverrideAuthOptionFunc {
+	return func(ctx context.Context) context.Context {
+		return context.WithValue(ctx, ScopeCheckBypassedContextKey{}, true)
 	}
+}
 
-	if account != nil {
-		newClaims.AccountName = *account
+// Overrides the authentication that is currently stored in the context. This
+// can only be used within a single process, and doesn't mean that the overrides
+// set here will be passed on if you are using `NewAuthenticatedClient` to pass
+// through auth. It is however useful for testing, or for calling other handlers
+// within the same process.
+func OverrideAuth(ctx context.Context, opts ...OverrideAuthOptionFunc) context.Context {
+	for _, opt := range opts {
+		ctx = opt(ctx)
 	}
+	return ctx
+}
 
-	// Store the new claims in the context
-	ctx = context.WithValue(ctx, CustomClaimsContextKey{}, &newClaims)
-	ctx = context.WithValue(ctx, AccountNameContextKey{}, newClaims.AccountName)
+func withCustomClaims(modify func(*CustomClaims)) OverrideAuthOptionFunc {
+	return func(ctx context.Context) context.Context {
+		i := ctx.Value(CustomClaimsContextKey{})
+		var claims *CustomClaims
+		var newClaims CustomClaims
+		var ok bool
 
-	return ctx
+		if claims, ok = i.(*CustomClaims); ok {
+			// clone out the values to avoid sharing
+			newClaims = *claims
+		}
+
+		modify(&newClaims)
+
+		// Store the new claims in the context
+		ctx = context.WithValue(ctx, CustomClaimsContextKey{}, &newClaims)
+		ctx = context.WithValue(ctx, AccountNameContextKey{}, newClaims.AccountName)
+
+		return ctx
+	}
 }
 
 // ensureValidTokenHandler is a middleware that will check the validity of our
@@ -375,7 +420,7 @@ func ensureValidTokenHandler(config AuthConfig, next http.Handler) http.Handler
 		span.SetAttributes(attribute.Bool("ovm.auth.bypass", shouldBypass))
 
 		if shouldBypass {
-			ctx = AddBypassAuthConfig(ctx)
+			ctx = OverrideAuth(ctx, WithBypassScopeCheck())
 
 			r = r.Clone(ctx)