fix(sdk): use bearer auth for non-dpop Connect tokens

sdk/auth/access_token_source.go

@@ -3,15 +3,49 @@ package auth
 import (
 	"context"
 	"net/http"
+	"strings"
 
 	"github.com/lestrrat-go/jwx/v2/jwk"
 )
 
 type AccessToken string
 
+type TokenType string
+
+const (
+	TokenTypeBearer TokenType = "Bearer"
+	TokenTypeDPoP   TokenType = "DPoP"
+)
+
+// TokenTypeFromOAuthTokenType returns the authentication scheme required for
+// an OAuth access token. A DPoP token type requires a DPoP proof; all other
+// values use the Bearer scheme.
+func TokenTypeFromOAuthTokenType(tokenType string) TokenType {
+	if strings.EqualFold(tokenType, string(TokenTypeDPoP)) {
+		return TokenTypeDPoP
+	}
+
+	return TokenTypeBearer
+}
+
 type AccessTokenSource interface {
 	AccessToken(ctx context.Context, client *http.Client) (AccessToken, error)
 	// MakeToken probably better to use `crypto.AsymDecryption` here than roll our own since this should be
 	// more closely linked to what happens in KAS in terms of crypto params
 	MakeToken(func(jwk.Key) ([]byte, error)) ([]byte, error)
 }
+
+// AccessTokenCredential binds an access token to its authentication scheme.
+type AccessTokenCredential struct {
+	Token AccessToken
+	Type  TokenType
+}
+
+// AccessTokenCredentialSource is implemented by token sources that can return
+// an access token and its authentication scheme atomically.
+//
+// Token sources without this optional interface retain the SDK's existing DPoP
+// behavior for backwards compatibility.
+type AccessTokenCredentialSource interface {
+	AccessTokenCredential(ctx context.Context, client *http.Client) (AccessTokenCredential, error)
+}

sdk/auth/access_token_source_test.go

@@ -0,0 +1,22 @@
+package auth
+
+import "testing"
+
+func TestTokenTypeFromOAuthTokenType(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		tokenType string
+		want      TokenType
+	}{
+		{name: "dpop", tokenType: "DPoP", want: TokenTypeDPoP},
+		{name: "bearer", tokenType: "Bearer", want: TokenTypeBearer},
+		{name: "lowercase dpop", tokenType: "dpop", want: TokenTypeDPoP},
+		{name: "missing token type", want: TokenTypeBearer},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := TokenTypeFromOAuthTokenType(tc.tokenType); got != tc.want {
+				t.Errorf("TokenTypeFromOAuthTokenType(%q) = %q, want %q", tc.tokenType, got, tc.want)
+			}
+		})
+	}
+}

sdk/auth/token_adding_interceptor.go

@@ -94,14 +94,25 @@ func (i TokenAddingInterceptor) AddCredentialsConnect() connect.UnaryInterceptor
 			ctx context.Context,
 			req connect.AnyRequest,
 		) (connect.AnyResponse, error) {
-			accessToken, err := i.tokenSource.AccessToken(ctx, i.httpClient)
+			credential := AccessTokenCredential{Type: TokenTypeDPoP}
+			var err error
+			if credentialSource, ok := i.tokenSource.(AccessTokenCredentialSource); ok {
+				credential, err = credentialSource.AccessTokenCredential(ctx, i.httpClient)
+			} else {
+				credential.Token, err = i.tokenSource.AccessToken(ctx, i.httpClient)
+			}
 			if err != nil {
 				slog.ErrorContext(ctx, "error getting access token", slog.Any("error", err))
 				return nil, connect.NewError(connect.CodeUnauthenticated, err)
 			}
 
+			if credential.Type != TokenTypeDPoP {
+				req.Header().Set("Authorization", fmt.Sprintf("Bearer %s", credential.Token))
+				return next(ctx, req)
+			}
+
 			// Add Authorization header
-			req.Header().Set("Authorization", fmt.Sprintf("DPoP %s", accessToken))
+			req.Header().Set("Authorization", fmt.Sprintf("DPoP %s", credential.Token))
 
 			// Add DPoP header if possible
 			slog.DebugContext(
@@ -110,7 +121,7 @@ func (i TokenAddingInterceptor) AddCredentialsConnect() connect.UnaryInterceptor
 				slog.String("dpop_htm", http.MethodPost),
 				slog.Any("stream_type", req.Spec().StreamType),
 			)
-			dpopTok, err := i.GetDPoPToken(req.Spec().Procedure, http.MethodPost, string(accessToken))
+			dpopTok, err := i.GetDPoPToken(req.Spec().Procedure, http.MethodPost, string(credential.Token))
 			if err == nil {
 				req.Header().Set("DPoP", dpopTok)
 			} else {

sdk/auth/token_adding_interceptor_test.go

@@ -41,9 +41,12 @@ func setupTokenAddingInterceptor(t *testing.T) (TokenAddingInterceptor, jwk.Key)
 	err = key.Set(jwk.AlgorithmKey, jwa.RS256)
 	require.NoError(t, err, "error setting the algorithm on the JWK")
 
-	ts := FakeTokenSource{
-		key:         key,
-		accessToken: "thisisafakeaccesstoken",
+	ts := DPoPFakeTokenSource{
+		FakeTokenSource: FakeTokenSource{
+			key:         key,
+			accessToken: "thisisafakeaccesstoken",
+		},
+		tokenType: TokenTypeDPoP,
 	}
 
 	oo := NewTokenAddingInterceptorWithClient(&ts, httputil.SafeHTTPClientWithTLSConfig(&tls.Config{
@@ -119,6 +122,82 @@ func TestAddingTokensToOutgoingRequest_Connect(t *testing.T) {
 	checkAccessAndDpopTokens(t, serverConnect.accessToken, serverConnect.dpopToken, key)
 }
 
+func TestAddingTokensToOutgoingRequest_Connect_BearerToken(t *testing.T) {
+	ts := DPoPFakeTokenSource{
+		FakeTokenSource: FakeTokenSource{
+			accessToken: "thisisafakeaccesstoken",
+		},
+		tokenType: TokenTypeBearer,
+	}
+	interceptor := NewTokenAddingInterceptorWithClient(&ts, httputil.SafeHTTPClient())
+	serverConnect := FakeAccessServiceServerConnect{}
+	clientConnect, stopC := runConnectServer(&serverConnect, interceptor)
+	defer stopC()
+
+	_, err := clientConnect.PublicKey(t.Context(), connect.NewRequest(&kas.PublicKeyRequest{}))
+	require.NoError(t, err, "error making call")
+
+	assert.Equal(t, []string{"Bearer thisisafakeaccesstoken"}, serverConnect.accessToken)
+	assert.Equal(t, []string{""}, serverConnect.dpopToken)
+	assert.False(t, ts.makeTokenCalled)
+}
+
+func TestAddingTokensToOutgoingRequest_Connect_UsesAtomicCredential(t *testing.T) {
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+	key, err := jwk.FromRaw(privateKey)
+	require.NoError(t, err)
+	require.NoError(t, key.Set(jwk.AlgorithmKey, jwa.RS256))
+
+	for _, tc := range []struct {
+		name              string
+		credential        AccessTokenCredential
+		wantAuthorization string
+		wantDPoP          bool
+	}{
+		{
+			name: "dpop",
+			credential: AccessTokenCredential{
+				Token: "dpop-access-token",
+				Type:  TokenTypeDPoP,
+			},
+			wantAuthorization: "DPoP dpop-access-token",
+			wantDPoP:          true,
+		},
+		{
+			name: "bearer",
+			credential: AccessTokenCredential{
+				Token: "bearer-access-token",
+				Type:  TokenTypeBearer,
+			},
+			wantAuthorization: "Bearer bearer-access-token",
+			wantDPoP:          false,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			ts := AtomicCredentialTokenSource{
+				FakeTokenSource: FakeTokenSource{key: key},
+				credential:      tc.credential,
+			}
+			interceptor := NewTokenAddingInterceptorWithClient(&ts, httputil.SafeHTTPClient())
+			serverConnect := FakeAccessServiceServerConnect{}
+			clientConnect, stopC := runConnectServer(&serverConnect, interceptor)
+			defer stopC()
+
+			_, err := clientConnect.PublicKey(t.Context(), connect.NewRequest(&kas.PublicKeyRequest{}))
+			require.NoError(t, err)
+
+			assert.Equal(t, []string{tc.wantAuthorization}, serverConnect.accessToken)
+			if tc.wantDPoP {
+				assert.NotEmpty(t, serverConnect.dpopToken[0])
+			} else {
+				assert.Equal(t, []string{""}, serverConnect.dpopToken)
+				assert.False(t, ts.makeTokenCalled)
+			}
+		})
+	}
+}
+
 func Test_InvalidCredentials_DoesNotSendMessage(t *testing.T) {
 	ts := FakeTokenSource{key: nil, accessToken: ""}
 	serverGrpc := FakeAccessServiceServer{}
@@ -194,8 +273,36 @@ func (f *FakeAccessServiceServer) Rewrap(context.Context, *kas.RewrapRequest) (*
 }
 
 type FakeTokenSource struct {
-	key         jwk.Key
-	accessToken string
+	key             jwk.Key
+	accessToken     string
+	makeTokenCalled bool
+}
+
+type DPoPFakeTokenSource struct {
+	FakeTokenSource
+	tokenType TokenType
+}
+
+func (fts *DPoPFakeTokenSource) AccessTokenCredential(ctx context.Context, client *http.Client) (AccessTokenCredential, error) {
+	token, err := fts.AccessToken(ctx, client)
+	if err != nil {
+		return AccessTokenCredential{}, err
+	}
+
+	return AccessTokenCredential{Token: token, Type: fts.tokenType}, nil
+}
+
+type AtomicCredentialTokenSource struct {
+	FakeTokenSource
+	credential AccessTokenCredential
+}
+
+func (fts *AtomicCredentialTokenSource) AccessToken(context.Context, *http.Client) (AccessToken, error) {
+	return "", errors.New("AccessToken must not be called when AccessTokenCredential is available")
+}
+
+func (fts *AtomicCredentialTokenSource) AccessTokenCredential(context.Context, *http.Client) (AccessTokenCredential, error) {
+	return fts.credential, nil
 }
 
 func (fts *FakeTokenSource) AccessToken(context.Context, *http.Client) (AccessToken, error) {
@@ -206,6 +313,7 @@ func (fts *FakeTokenSource) AccessToken(context.Context, *http.Client) (AccessTo
 }
 
 func (fts *FakeTokenSource) MakeToken(f func(jwk.Key) ([]byte, error)) ([]byte, error) {
+	fts.makeTokenCalled = true
 	if fts.key == nil {
 		return nil, errors.New("no such key")
 	}

sdk/idp_access_token_source.go

@@ -94,19 +94,28 @@ func NewIDPAccessTokenSource(
 }
 
 // AccessToken use a pointer receiver so that the token state is shared
-func (t *IDPAccessTokenSource) AccessToken(_ context.Context, client *http.Client) (auth.AccessToken, error) {
+func (t *IDPAccessTokenSource) AccessToken(ctx context.Context, client *http.Client) (auth.AccessToken, error) {
+	credential, err := t.AccessTokenCredential(ctx, client)
+	return credential.Token, err
+}
+
+// AccessTokenCredential returns the cached access token and its DPoP status atomically.
+func (t *IDPAccessTokenSource) AccessTokenCredential(_ context.Context, client *http.Client) (auth.AccessTokenCredential, error) {
 	t.tokenMutex.Lock()
 	defer t.tokenMutex.Unlock()
 
 	if t.token == nil || t.token.Expired() {
 		tok, err := oauth.GetAccessToken(client, t.idpTokenEndpoint.String(), t.scopes, t.credentials, t.dpopKey)
 		if err != nil {
-			return "", fmt.Errorf("error getting access token: %w", err)
+			return auth.AccessTokenCredential{}, fmt.Errorf("error getting access token: %w", err)
 		}
 		t.token = tok
 	}
 
-	return auth.AccessToken(t.token.AccessToken), nil
+	return auth.AccessTokenCredential{
+		Token: auth.AccessToken(t.token.AccessToken),
+		Type:  auth.TokenTypeFromOAuthTokenType(t.token.TokenType),
+	}, nil
 }
 
 func (t *IDPAccessTokenSource) MakeToken(tokenMaker func(jwk.Key) ([]byte, error)) ([]byte, error) {

sdk/idp_cert_exchange.go

@@ -42,19 +42,28 @@ func NewCertExchangeTokenSource(logger *slog.Logger, info oauth.CertExchangeInfo
 }
 
 func (c *CertExchangeTokenSource) AccessToken(ctx context.Context, _ *http.Client) (auth.AccessToken, error) {
+	credential, err := c.AccessTokenCredential(ctx, nil)
+	return credential.Token, err
+}
+
+// AccessTokenCredential returns the cached certificate-exchange token and its DPoP status atomically.
+func (c *CertExchangeTokenSource) AccessTokenCredential(ctx context.Context, _ *http.Client) (auth.AccessTokenCredential, error) {
 	c.tokenMutex.Lock()
 	defer c.tokenMutex.Unlock()
 
 	if c.token == nil || c.token.Expired() {
 		tok, err := oauth.DoCertExchange(ctx, c.IdpEndpoint, c.info, c.credentials, c.key)
 		if err != nil {
-			return "", err
+			return auth.AccessTokenCredential{}, err
 		}
 
 		c.token = tok
 	}
 
-	return auth.AccessToken(c.token.AccessToken), nil
+	return auth.AccessTokenCredential{
+		Token: auth.AccessToken(c.token.AccessToken),
+		Type:  auth.TokenTypeFromOAuthTokenType(c.token.TokenType),
+	}, nil
 }
 
 func (c *CertExchangeTokenSource) MakeToken(tokenMaker func(jwk.Key) ([]byte, error)) ([]byte, error) {

sdk/idp_oauth_access_token_source.go

@@ -11,7 +11,7 @@ import (
 	"golang.org/x/oauth2"
 )
 
-// OAuthAccessTokenSource allow connecting to an IDP and obtain a DPoP bound access token
+// OAuthAccessTokenSource allows connecting to an IDP and obtaining an access token.
 type OAuthAccessTokenSource struct {
 	source         oauth2.TokenSource
 	scopes         []string
@@ -40,19 +40,28 @@ func NewOAuthAccessTokenSource(
 }
 
 // AccessToken use a pointer receiver so that the token state is shared
-func (t *OAuthAccessTokenSource) AccessToken(_ context.Context, _ *http.Client) (auth.AccessToken, error) { // must satisfy auth.AccessTokenSource interface
+func (t *OAuthAccessTokenSource) AccessToken(ctx context.Context, client *http.Client) (auth.AccessToken, error) {
+	credential, err := t.AccessTokenCredential(ctx, client)
+	return credential.Token, err
+}
+
+// AccessTokenCredential returns an access token and its DPoP status from the same token response.
+func (t *OAuthAccessTokenSource) AccessTokenCredential(_ context.Context, _ *http.Client) (auth.AccessTokenCredential, error) {
 	tok, err := t.source.Token()
 	if err != nil {
-		return "", fmt.Errorf("error getting access token: %w", err)
+		return auth.AccessTokenCredential{}, fmt.Errorf("error getting access token: %w", err)
 	}
 
 	// Non-nil with AccessToken and not Expired
 	if !tok.Valid() {
-		return "", ErrAccessTokenInvalid
+		return auth.AccessTokenCredential{}, ErrAccessTokenInvalid
 		// TODO: refresh tokens if expired?
 	}
 
-	return auth.AccessToken(tok.AccessToken), nil
+	return auth.AccessTokenCredential{
+		Token: auth.AccessToken(tok.AccessToken),
+		Type:  auth.TokenTypeFromOAuthTokenType(tok.Type()),
+	}, nil
 }
 
 func (t *OAuthAccessTokenSource) MakeToken(tokenMaker func(jwk.Key) ([]byte, error)) ([]byte, error) {