diff --git a/postgres/auth.go b/postgres/auth.go index b2004e7..9e4ccec 100644 --- a/postgres/auth.go +++ b/postgres/auth.go @@ -31,10 +31,19 @@ func NewAuthToken(ctx context.Context, cfg *Config, user string) (string, error) if cfg.DynamicAuth == nil { return "", nil } - if cfg.DynamicAuth.AWSRDSIAM != nil { + if err := singleDynamicAuthBackend(cfg.DynamicAuth); err != nil { + return "", err + } + switch { + case cfg.DynamicAuth.AWSRDSIAM != nil: return awsRDSIAMToken(ctx, cfg, user) + case cfg.DynamicAuth.AzureAD != nil: + return azureADToken(ctx) + case cfg.DynamicAuth.GCPCloudSQLIAM != nil: + return gcpCloudSQLIAMToken(ctx) + default: + return "", errors.New("unreachable: singleDynamicAuthBackend guarantees exactly one backend is set") } - return "", errors.New("dynamicAuth is set but no supported auth method (e.g., awsRdsIam) is configured") } // NewDynamicAuthFunc returns a BeforeConnect hook that resolves a fresh @@ -52,10 +61,19 @@ func NewDynamicAuthFunc(ctx context.Context, cfg *Config, user string) (BeforeCo if cfg.DynamicAuth == nil { return nil, errors.New("dynamic authentication is not configured") } - if cfg.DynamicAuth.AWSRDSIAM != nil { + if err := singleDynamicAuthBackend(cfg.DynamicAuth); err != nil { + return nil, err + } + switch { + case cfg.DynamicAuth.AWSRDSIAM != nil: return awsRDSIAMBeforeConnect(ctx, cfg, user) + case cfg.DynamicAuth.AzureAD != nil: + return azureADBeforeConnect() + case cfg.DynamicAuth.GCPCloudSQLIAM != nil: + return gcpCloudSQLIAMBeforeConnect() + default: + return nil, errors.New("unreachable: singleDynamicAuthBackend guarantees exactly one backend is set") } - return nil, errors.New("dynamicAuth is set but no supported auth method (e.g., awsRdsIam) is configured") } // wrapAuthError prefixes dynamic-auth errors with a consistent label so they diff --git a/postgres/auth_test.go b/postgres/auth_test.go index d5a047e..c568587 100644 --- a/postgres/auth_test.go +++ b/postgres/auth_test.go @@ -48,6 +48,17 @@ func TestNewAuthToken(t *testing.T) { }, wantErr: testErrRegionMissing, }, + { + name: testCaseMultipleBackend, + cfg: &Config{ + Host: "h", Port: 5432, User: "u", Database: "d", + DynamicAuth: &DynamicAuthConfig{ + AWSRDSIAM: &DynamicAuthAWSRDSIAM{Region: testRegion}, + GCPCloudSQLIAM: &DynamicAuthGCPCloudSQLIAM{}, + }, + }, + wantErr: testErrMultipleBackend, + }, } for _, tt := range tests { @@ -102,6 +113,17 @@ func TestNewDynamicAuthFunc(t *testing.T) { }, wantErr: testErrRegionMissing, }, + { + name: testCaseMultipleBackend, + cfg: &Config{ + Host: "h", Port: 5432, User: "u", Database: "d", + DynamicAuth: &DynamicAuthConfig{ + AzureAD: &DynamicAuthAzureAD{}, + GCPCloudSQLIAM: &DynamicAuthGCPCloudSQLIAM{}, + }, + }, + wantErr: testErrMultipleBackend, + }, } for _, tt := range tests { diff --git a/postgres/azuread.go b/postgres/azuread.go new file mode 100644 index 0000000..9ec36bf --- /dev/null +++ b/postgres/azuread.go @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/jackc/pgx/v5" +) + +// azureADScope is the OAuth2 scope Azure Database for PostgreSQL Flexible +// Server (and Single Server) requires for Entra ID (formerly Azure AD) +// token-based authentication. +const azureADScope = "https://ossrdbms-aad.database.windows.net/.default" + +// azureADToken returns a single Entra ID access token usable as a PostgreSQL +// password for Azure Database for PostgreSQL. +func azureADToken(ctx context.Context) (string, error) { + cred, err := newAzureCredential() + if err != nil { + return "", wrapAuthError("azureAd", err) + } + token, err := cred.GetToken(ctx, policy.TokenRequestOptions{Scopes: []string{azureADScope}}) + if err != nil { + return "", wrapAuthError("azureAd", fmt.Errorf("failed to acquire Entra ID token: %w", err)) + } + return token.Token, nil +} + +// azureADBeforeConnect returns a BeforeConnect hook that generates a fresh +// Entra ID token before each connection attempt. The credential is +// constructed once, at hook-construction time; azidentity caches and +// refreshes the underlying token internally, so per-connection cost is a +// single (usually cached) token acquisition. +func azureADBeforeConnect() (BeforeConnectFn, error) { + cred, err := newAzureCredential() + if err != nil { + return nil, wrapAuthError("azureAd", err) + } + return func(ctx context.Context, conn *pgx.ConnConfig) error { + token, err := cred.GetToken(ctx, policy.TokenRequestOptions{Scopes: []string{azureADScope}}) + if err != nil { + return wrapAuthError("azureAd", fmt.Errorf("failed to acquire Entra ID token: %w", err)) + } + conn.Password = token.Token + return nil + }, nil +} + +// newAzureCredential builds the credential used to acquire Entra ID tokens. +// Construction never contacts Azure — DefaultAzureCredential resolves +// lazily, on the first GetToken call, following its normal chain +// (environment variables — including AZURE_CLIENT_ID to select a +// user-assigned managed identity — workload identity, managed identity, +// Azure CLI, ...). +func newAzureCredential() (*azidentity.DefaultAzureCredential, error) { + cred, err := azidentity.NewDefaultAzureCredential(nil) + if err != nil { + return nil, fmt.Errorf("failed to construct Azure credential: %w", err) + } + return cred, nil +} diff --git a/postgres/azuread_test.go b/postgres/azuread_test.go new file mode 100644 index 0000000..fa41696 --- /dev/null +++ b/postgres/azuread_test.go @@ -0,0 +1,36 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestAzureADBeforeConnect_ReturnsHookWithoutContactingAzure verifies the +// constructor returns a non-nil hook with no error, deterministically and +// without any ambient Azure credentials: unlike GCP's +// google.DefaultTokenSource (see gcpiam.go's doc comment), azidentity's +// DefaultAzureCredential resolves lazily, on the first GetToken call, so +// construction alone never contacts Azure or requires credentials to be +// present. Actually invoking the returned hook would require real Azure +// credentials and is out of scope for unit tests — same posture as +// TestAwsRDSIAMBeforeConnect_ReturnsHookForStaticRegion. +func TestAzureADBeforeConnect_ReturnsHookWithoutContactingAzure(t *testing.T) { + t.Parallel() + fn, err := azureADBeforeConnect() + require.NoError(t, err) + assert.NotNil(t, fn) +} + +// No equivalent GCP constructor test exists: unlike AWS and Azure, +// google.DefaultTokenSource (gcpCloudSQLIAMBeforeConnect's first call) +// resolves Application Default Credentials eagerly and returns an error +// immediately when none are found, rather than deferring resolution to the +// first token fetch. That makes even "does the constructor return a non-nil +// hook" environment-dependent — a machine with real GCP credentials +// configured would get a different, equally valid result. See gcpiam.go's +// newGCPTokenSource doc comment. diff --git a/postgres/config.go b/postgres/config.go index 6ab571a..a82445d 100644 --- a/postgres/config.go +++ b/postgres/config.go @@ -95,6 +95,19 @@ type Config struct { type DynamicAuthConfig struct { // AWSRDSIAM enables AWS RDS IAM authentication tokens. AWSRDSIAM *DynamicAuthAWSRDSIAM + + // AzureAD enables Azure Entra ID (formerly Azure AD) authentication + // tokens for Azure Database for PostgreSQL. + AzureAD *DynamicAuthAzureAD + + // GCPCloudSQLIAM enables GCP Cloud SQL IAM database authentication + // tokens over a direct TCP+TLS connection. This is NOT the Cloud SQL Go + // connector (cloudsqlconn): it requires the instance to have a reachable + // IP (public, or private with direct network routing) and does not get + // Cloud SQL's automatic mTLS tunnel. Instances reachable only through + // the connector need a different integration this package does not + // provide. + GCPCloudSQLIAM *DynamicAuthGCPCloudSQLIAM } // DynamicAuthAWSRDSIAM configures AWS RDS IAM dynamic authentication. @@ -104,6 +117,51 @@ type DynamicAuthAWSRDSIAM struct { Region string } +// DynamicAuthAzureAD configures Azure Entra ID (formerly Azure AD) +// authentication for Azure Database for PostgreSQL. It has no fields: the +// token is minted from DefaultAzureCredential's normal resolution order +// (environment variables — including AZURE_CLIENT_ID to select a +// user-assigned managed identity — workload identity, system-assigned +// managed identity, Azure CLI, ...). +type DynamicAuthAzureAD struct{} + +// DynamicAuthGCPCloudSQLIAM configures GCP Cloud SQL IAM database +// authentication. It has no fields: the token is minted from ambient +// Application Default Credentials, scoped for Cloud SQL login. +type DynamicAuthGCPCloudSQLIAM struct{} + +// countDynamicAuthBackends returns how many backend fields on da are set. +func countDynamicAuthBackends(da *DynamicAuthConfig) int { + n := 0 + if da.AWSRDSIAM != nil { + n++ + } + if da.AzureAD != nil { + n++ + } + if da.GCPCloudSQLIAM != nil { + n++ + } + return n +} + +// singleDynamicAuthBackend rejects a DynamicAuthConfig with zero or more than +// one backend configured. Config.Validate calls this, and so do NewAuthToken +// and NewDynamicAuthFunc: a caller may build a Config and mint a one-shot +// token (for example, a migration tool) without going through Validate +// first, so the dispatchers re-check rather than assuming an ambiguous +// configuration was already caught upstream. +func singleDynamicAuthBackend(da *DynamicAuthConfig) error { + switch n := countDynamicAuthBackends(da); { + case n == 0: + return errors.New("dynamicAuth is set but no supported auth method (e.g., awsRdsIam, azureAd, gcpCloudSqlIam) is configured") + case n > 1: + return errors.New("dynamicAuth must configure exactly one auth method, but more than one is set") + default: + return nil + } +} + // Validate checks Config for required-field and consistency errors and // returns the first violation encountered. func (c *Config) Validate() error { @@ -129,10 +187,10 @@ func (c *Config) Validate() error { return fmt.Errorf("database must not contain any of %q or whitespace", databaseForbiddenChars) } if c.DynamicAuth != nil { - if c.DynamicAuth.AWSRDSIAM == nil { - return errors.New("dynamicAuth is set but no supported auth method (e.g., awsRdsIam) is configured") + if err := singleDynamicAuthBackend(c.DynamicAuth); err != nil { + return err } - if c.DynamicAuth.AWSRDSIAM.Region == "" { + if c.DynamicAuth.AWSRDSIAM != nil && c.DynamicAuth.AWSRDSIAM.Region == "" { return errors.New("dynamicAuth.awsRdsIam.region is required") } } diff --git a/postgres/config_test.go b/postgres/config_test.go index f294c45..df3714a 100644 --- a/postgres/config_test.go +++ b/postgres/config_test.go @@ -95,6 +95,17 @@ func TestConfig_Validate(t *testing.T) { }, wantErr: testErrRegionConfigured, }, + { + name: testCaseMultipleBackend, + cfg: &Config{ + Host: "h", Port: 5432, User: "u", Database: "d", + DynamicAuth: &DynamicAuthConfig{ + AWSRDSIAM: &DynamicAuthAWSRDSIAM{Region: testRegion}, + AzureAD: &DynamicAuthAzureAD{}, + }, + }, + wantErr: testErrMultipleBackend, + }, { name: "valid minimal config", cfg: validConfig(), @@ -108,6 +119,24 @@ func TestConfig_Validate(t *testing.T) { }, }, }, + { + name: "valid with Azure AD", + cfg: &Config{ + Host: "h", Port: 5432, User: "u", Database: "d", + DynamicAuth: &DynamicAuthConfig{ + AzureAD: &DynamicAuthAzureAD{}, + }, + }, + }, + { + name: "valid with GCP Cloud SQL IAM", + cfg: &Config{ + Host: "h", Port: 5432, User: "u", Database: "d", + DynamicAuth: &DynamicAuthConfig{ + GCPCloudSQLIAM: &DynamicAuthGCPCloudSQLIAM{}, + }, + }, + }, } for _, tt := range tests { diff --git a/postgres/doc.go b/postgres/doc.go index b73ff04..41a4e33 100644 --- a/postgres/doc.go +++ b/postgres/doc.go @@ -29,7 +29,8 @@ remains the caller's responsibility. # Dynamic Authentication Setting Config.DynamicAuth causes NewPool to install a BeforeConnect hook -that resolves a fresh credential before every connection attempt. +that resolves a fresh credential before every connection attempt. Exactly +one backend must be configured. Currently supported backends: @@ -37,12 +38,31 @@ Currently supported backends: AWS credentials (env vars, EC2 instance profile, EKS web identity, …). Region "detect" auto-discovers the region via IMDS. -Example: + - Azure Entra ID (formerly Azure AD) — OAuth2 tokens for Azure Database + for PostgreSQL, acquired via DefaultAzureCredential's normal chain + (environment variables, workload identity, managed identity, Azure + CLI, …). AZURE_CLIENT_ID selects a user-assigned managed identity. + + - GCP Cloud SQL IAM — OAuth2 tokens minted from ambient Application + Default Credentials, connecting over a direct TCP+TLS connection. This + is NOT the Cloud SQL Go connector (cloudsqlconn): it requires the + instance to have a reachable IP (public, or private with direct + network routing) and does not get Cloud SQL's automatic mTLS tunnel. + +Examples: cfg.DynamicAuth = &postgres.DynamicAuthConfig{ AWSRDSIAM: &postgres.DynamicAuthAWSRDSIAM{Region: "us-east-1"}, } + cfg.DynamicAuth = &postgres.DynamicAuthConfig{ + AzureAD: &postgres.DynamicAuthAzureAD{}, + } + + cfg.DynamicAuth = &postgres.DynamicAuthConfig{ + GCPCloudSQLIAM: &postgres.DynamicAuthGCPCloudSQLIAM{}, + } + For short-lived connections that cannot use a pool hook (for example golang-migrate's one-shot migration connection), call NewAuthToken to materialize a single token, then embed it via BuildConnectionStringWithAuth: diff --git a/postgres/gcpiam.go b/postgres/gcpiam.go new file mode 100644 index 0000000..79a3ad6 --- /dev/null +++ b/postgres/gcpiam.go @@ -0,0 +1,101 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + + "github.com/jackc/pgx/v5" + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" +) + +// gcpCloudSQLIAMScope is the OAuth2 scope required for GCP Cloud SQL IAM +// database authentication — distinct from the broader sqlservice.admin scope +// used to manage Cloud SQL instances themselves. +const gcpCloudSQLIAMScope = "https://www.googleapis.com/auth/sqlservice.login" + +// gcpCloudSQLIAMToken returns a single GCP OAuth2 access token usable as a +// PostgreSQL password for Cloud SQL IAM database authentication. This is the +// direct-TCP token-swap path, not the Cloud SQL Go connector (cloudsqlconn): +// it requires the instance to have a reachable IP and does not get Cloud +// SQL's automatic mTLS tunnel — see DynamicAuthGCPCloudSQLIAM's doc comment. +func gcpCloudSQLIAMToken(ctx context.Context) (string, error) { + ts, err := newGCPTokenSource(ctx) + if err != nil { + return "", wrapAuthError("gcpCloudSqlIam", err) + } + token, err := tokenWithContext(ctx, ts) + if err != nil { + return "", wrapAuthError("gcpCloudSqlIam", fmt.Errorf("failed to acquire GCP access token: %w", err)) + } + return token.AccessToken, nil +} + +// gcpCloudSQLIAMBeforeConnect returns a BeforeConnect hook that generates a +// fresh GCP access token before each connection attempt. The token source is +// constructed once, at hook-construction time, using context.Background() +// rather than the ctx passed in here: google.DefaultTokenSource captures +// whatever context it's given for the lifetime of the returned TokenSource +// (used to build its internal token-refresh HTTP client), and this +// constructor's ctx is often request- or pool-construction-scoped and may be +// canceled shortly after NewPool returns — which would poison every later +// token refresh through this source. oauth2.TokenSource caches and refreshes +// the underlying token internally. +func gcpCloudSQLIAMBeforeConnect() (BeforeConnectFn, error) { + ts, err := newGCPTokenSource(context.Background()) + if err != nil { + return nil, wrapAuthError("gcpCloudSqlIam", err) + } + return func(ctx context.Context, conn *pgx.ConnConfig) error { + token, err := tokenWithContext(ctx, ts) + if err != nil { + return wrapAuthError("gcpCloudSqlIam", fmt.Errorf("failed to acquire GCP access token: %w", err)) + } + conn.Password = token.AccessToken + return nil + }, nil +} + +// tokenWithContext calls ts.Token(), honoring ctx for cancellation even +// though oauth2.TokenSource's Token method takes no context of its own. +// This returns as soon as ctx is done, but — since the underlying call +// cannot itself be aborted through this interface — the goroutine calling +// Token() keeps running in the background until it completes on its own. +func tokenWithContext(ctx context.Context, ts oauth2.TokenSource) (*oauth2.Token, error) { + type result struct { + token *oauth2.Token + err error + } + ch := make(chan result, 1) + go func() { + token, err := ts.Token() + ch <- result{token, err} + }() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case r := <-ch: + return r.token, r.err + } +} + +// newGCPTokenSource builds the token source used to acquire GCP access +// tokens. Unlike the AWS and Azure backends, this call is NOT +// credential-free: google.DefaultTokenSource resolves Application Default +// Credentials eagerly and returns an error immediately when none are found, +// rather than deferring resolution to the first Token() call. That makes its +// outcome environment-dependent in a way the other two backends' +// constructors are not, which is why this package has no equivalent to +// TestAwsRDSIAMBeforeConnect_ReturnsHookForStaticRegion / +// TestAzureADBeforeConnect_ReturnsHookWithoutContactingAzure for this +// backend — see the explanatory comment at the bottom of azuread_test.go. +func newGCPTokenSource(ctx context.Context) (oauth2.TokenSource, error) { + ts, err := google.DefaultTokenSource(ctx, gcpCloudSQLIAMScope) + if err != nil { + return nil, fmt.Errorf("failed to load GCP default credentials: %w", err) + } + return ts, nil +} diff --git a/postgres/gcpiam_test.go b/postgres/gcpiam_test.go new file mode 100644 index 0000000..7889f76 --- /dev/null +++ b/postgres/gcpiam_test.go @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// blockingTokenSource is a stub oauth2.TokenSource that blocks until unblock +// is closed, then returns token/err. It stands in for a stalled GCP token +// refresh so tokenWithContext's cancellation behavior can be tested without +// real network calls or credentials. +type blockingTokenSource struct { + unblock chan struct{} + token *oauth2.Token + err error +} + +func (b *blockingTokenSource) Token() (*oauth2.Token, error) { + <-b.unblock + return b.token, b.err +} + +func TestTokenWithContext_ReturnsPromptlyOnCancellation(t *testing.T) { + t.Parallel() + + ts := &blockingTokenSource{unblock: make(chan struct{})} // never closed: Token() blocks forever + ctx, cancel := context.WithCancel(t.Context()) + + done := make(chan struct{}) + var gotErr error + go func() { + _, gotErr = tokenWithContext(ctx, ts) + close(done) + }() + + cancel() + select { + case <-done: + require.Error(t, gotErr) + assert.ErrorIs(t, gotErr, context.Canceled) + case <-time.After(2 * time.Second): + t.Fatal("tokenWithContext did not return promptly after ctx cancellation") + } +} + +func TestTokenWithContext_ReturnsTokenOnSuccess(t *testing.T) { + t.Parallel() + + want := &oauth2.Token{AccessToken: "test-access-token"} + ts := &blockingTokenSource{unblock: make(chan struct{}), token: want} + close(ts.unblock) + + got, err := tokenWithContext(t.Context(), ts) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +func TestTokenWithContext_PropagatesUnderlyingError(t *testing.T) { + t.Parallel() + + wantErr := errors.New("token refresh failed") + ts := &blockingTokenSource{unblock: make(chan struct{}), err: wantErr} + close(ts.unblock) + + _, err := tokenWithContext(t.Context(), ts) + require.Error(t, err) + assert.ErrorIs(t, err, wantErr) +} diff --git a/postgres/testdata_test.go b/postgres/testdata_test.go index 1f1b3bd..8482250 100644 --- a/postgres/testdata_test.go +++ b/postgres/testdata_test.go @@ -9,6 +9,8 @@ const ( testErrConfigNil = "config is nil" testCaseNoBackend = "dynamic auth without backend" testErrNoSupportedAuth = "no supported auth method" + testCaseMultipleBackend = "dynamic auth with more than one backend" + testErrMultipleBackend = "more than one is set" testSSLModeDisable = "disable" testErrRegionMissing = "AWS RDS IAM region is not configured" testErrRegionConfigured = "dynamicAuth.awsRdsIam.region is required" diff --git a/redis/gcpiam.go b/redis/gcpiam.go index 57d3b7c..09b64a0 100644 --- a/redis/gcpiam.go +++ b/redis/gcpiam.go @@ -49,8 +49,8 @@ func gcpMemorystoreIAMCredentialsFunc() (CredentialsFunc, error) { if err != nil { return nil, wrapAuthError("gcpMemorystoreIam", err) } - return func(context.Context) (string, string, error) { - token, err := ts.Token() + return func(ctx context.Context) (string, string, error) { + token, err := tokenWithContext(ctx, ts) if err != nil { return "", "", wrapAuthError("gcpMemorystoreIam", fmt.Errorf("failed to acquire GCP access token: %w", err)) } @@ -58,6 +58,29 @@ func gcpMemorystoreIAMCredentialsFunc() (CredentialsFunc, error) { }, nil } +// tokenWithContext calls ts.Token(), honoring ctx for cancellation even +// though oauth2.TokenSource's Token method takes no context of its own. +// This returns as soon as ctx is done, but — since the underlying call +// cannot itself be aborted through this interface — the goroutine calling +// Token() keeps running in the background until it completes on its own. +func tokenWithContext(ctx context.Context, ts oauth2.TokenSource) (*oauth2.Token, error) { + type result struct { + token *oauth2.Token + err error + } + ch := make(chan result, 1) + go func() { + token, err := ts.Token() + ch <- result{token, err} + }() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case r := <-ch: + return r.token, r.err + } +} + // newGCPTokenSource builds the token source used to acquire GCP access // tokens, scoped for Memorystore IAM authentication. func newGCPTokenSource(ctx context.Context) (oauth2.TokenSource, error) { diff --git a/redis/gcpiam_test.go b/redis/gcpiam_test.go new file mode 100644 index 0000000..742dec4 --- /dev/null +++ b/redis/gcpiam_test.go @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package redis + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// blockingTokenSource is a stub oauth2.TokenSource that blocks until unblock +// is closed, then returns token/err. It stands in for a stalled GCP token +// refresh so tokenWithContext's cancellation behavior can be tested without +// real network calls or credentials. +type blockingTokenSource struct { + unblock chan struct{} + token *oauth2.Token + err error +} + +func (b *blockingTokenSource) Token() (*oauth2.Token, error) { + <-b.unblock + return b.token, b.err +} + +func TestTokenWithContext_ReturnsPromptlyOnCancellation(t *testing.T) { + t.Parallel() + + ts := &blockingTokenSource{unblock: make(chan struct{})} // never closed: Token() blocks forever + ctx, cancel := context.WithCancel(t.Context()) + + done := make(chan struct{}) + var gotErr error + go func() { + _, gotErr = tokenWithContext(ctx, ts) + close(done) + }() + + cancel() + select { + case <-done: + require.Error(t, gotErr) + assert.ErrorIs(t, gotErr, context.Canceled) + case <-time.After(2 * time.Second): + t.Fatal("tokenWithContext did not return promptly after ctx cancellation") + } +} + +func TestTokenWithContext_ReturnsTokenOnSuccess(t *testing.T) { + t.Parallel() + + want := &oauth2.Token{AccessToken: "test-access-token"} + ts := &blockingTokenSource{unblock: make(chan struct{}), token: want} + close(ts.unblock) + + got, err := tokenWithContext(t.Context(), ts) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +func TestTokenWithContext_PropagatesUnderlyingError(t *testing.T) { + t.Parallel() + + wantErr := errors.New("token refresh failed") + ts := &blockingTokenSource{unblock: make(chan struct{}), err: wantErr} + close(ts.unblock) + + _, err := tokenWithContext(t.Context(), ts) + require.Error(t, err) + assert.ErrorIs(t, err, wantErr) +}