From 2afbb90d2e1ee86074298e275bcd5bd220fab7a4 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sun, 6 Sep 2026 14:52:18 +0000 Subject: [PATCH] auth: move U2M authentication into the CLI --- cmd/auth/login.go | 8 +- cmd/auth/login_test.go | 2 +- cmd/auth/token.go | 4 +- cmd/auth/token_test.go | 2 +- libs/auth/arguments.go | 2 +- libs/auth/arguments_test.go | 2 +- libs/auth/credentials.go | 2 +- libs/auth/credentials_test.go | 2 +- libs/auth/error.go | 2 +- libs/auth/storage/cache.go | 12 +- libs/auth/storage/cache_test.go | 2 +- libs/auth/storage/dual_writing_cache.go | 6 +- libs/auth/storage/dual_writing_cache_test.go | 4 +- libs/auth/storage/keyring.go | 2 +- libs/auth/storage/not_found_hint.go | 6 +- libs/auth/storage/not_found_hint_test.go | 2 +- libs/auth/storage/sdku2m.go | 45 - libs/auth/storage/storage.go | 6 +- libs/auth/storage/u2m.go | 43 + libs/auth/u2m/account_oauth_argument.go | 74 + libs/auth/u2m/account_oauth_argument_test.go | 81 + libs/auth/u2m/cache/cache.go | 32 + libs/auth/u2m/cache/memory.go | 53 + libs/auth/u2m/cache/memory_test.go | 103 ++ libs/auth/u2m/callback.go | 156 ++ libs/auth/u2m/callback_test.go | 244 +++ libs/auth/u2m/discovery_oauth_argument.go | 63 + .../auth/u2m/discovery_oauth_argument_test.go | 106 ++ libs/auth/u2m/discovery_token_source.go | 191 +++ libs/auth/u2m/discovery_token_source_test.go | 394 +++++ libs/auth/u2m/doc.go | 47 + libs/auth/u2m/endpoint_supplier.go | 81 + libs/auth/u2m/endpoint_supplier_test.go | 115 ++ libs/auth/u2m/error.go | 14 + libs/auth/u2m/oauth_argument.go | 22 + libs/auth/u2m/page.tmpl | 104 ++ libs/auth/u2m/persistent_auth.go | 661 ++++++++ libs/auth/u2m/persistent_auth_test.go | 1473 +++++++++++++++++ libs/auth/u2m/unified_oauth_argument.go | 74 + libs/auth/u2m/unified_oauth_argument_test.go | 106 ++ libs/auth/u2m/workspace_oauth_argument.go | 87 + .../auth/u2m/workspace_oauth_argument_test.go | 119 ++ 42 files changed, 4476 insertions(+), 78 deletions(-) delete mode 100644 libs/auth/storage/sdku2m.go create mode 100644 libs/auth/storage/u2m.go create mode 100644 libs/auth/u2m/account_oauth_argument.go create mode 100644 libs/auth/u2m/account_oauth_argument_test.go create mode 100644 libs/auth/u2m/cache/cache.go create mode 100644 libs/auth/u2m/cache/memory.go create mode 100644 libs/auth/u2m/cache/memory_test.go create mode 100644 libs/auth/u2m/callback.go create mode 100644 libs/auth/u2m/callback_test.go create mode 100644 libs/auth/u2m/discovery_oauth_argument.go create mode 100644 libs/auth/u2m/discovery_oauth_argument_test.go create mode 100644 libs/auth/u2m/discovery_token_source.go create mode 100644 libs/auth/u2m/discovery_token_source_test.go create mode 100644 libs/auth/u2m/doc.go create mode 100644 libs/auth/u2m/endpoint_supplier.go create mode 100644 libs/auth/u2m/endpoint_supplier_test.go create mode 100644 libs/auth/u2m/error.go create mode 100644 libs/auth/u2m/oauth_argument.go create mode 100644 libs/auth/u2m/page.tmpl create mode 100644 libs/auth/u2m/persistent_auth.go create mode 100644 libs/auth/u2m/persistent_auth_test.go create mode 100644 libs/auth/u2m/unified_oauth_argument.go create mode 100644 libs/auth/u2m/unified_oauth_argument_test.go create mode 100644 libs/auth/u2m/workspace_oauth_argument.go create mode 100644 libs/auth/u2m/workspace_oauth_argument_test.go diff --git a/cmd/auth/login.go b/cmd/auth/login.go index f9f2531ac74..5835297b335 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/cli/libs/browser" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" @@ -21,7 +22,6 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" - "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/spf13/cobra" "golang.org/x/oauth2" ) @@ -47,8 +47,7 @@ const ( discoveryFallbackTip = "\n\nTip: you can specify a workspace directly with: databricks auth login --host " // discoveryHostEnvVar overrides the default https://login.databricks.com // host used by the discovery login flow. Intended for testing and - // development against non-production environments. See WithDiscoveryHost - // in github.com/databricks/databricks-sdk-go/credentials/u2m. + // development against non-production environments. discoveryHostEnvVar = "DATABRICKS_DISCOVERY_HOST" ) @@ -735,7 +734,8 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { // cluster_id, serverless_compute_id) from a prior login to a different host // type must be cleared so they don't leak into the new profile. account_id // and workspace_id are re-added from discovery/introspection results. - clearKeys = append(clearKeys, + clearKeys = append( + clearKeys, "account_id", "workspace_id", databrickscfg.ExperimentalIsUnifiedHostKey, diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index a8eafb4be43..cd9e20bc234 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -16,11 +16,11 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" diff --git a/cmd/auth/token.go b/cmd/auth/token.go index d5e88e64d72..dc8f54ecfc3 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -12,6 +12,8 @@ import ( "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/cli/libs/auth/u2m/cache" "github.com/databricks/cli/libs/browser" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" @@ -20,8 +22,6 @@ import ( "github.com/databricks/cli/libs/flags" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/credentials/u2m" - "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" "github.com/spf13/cobra" "golang.org/x/oauth2" ) diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index adda6888a40..11a5a937bd9 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -11,10 +11,10 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" - "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "github.com/stretchr/testify/assert" "golang.org/x/oauth2" diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go index d18d7058cd8..deac0b5b1cc 100644 --- a/libs/auth/arguments.go +++ b/libs/auth/arguments.go @@ -1,8 +1,8 @@ package auth import ( + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/credentials/u2m" ) // WorkspaceIDNone is a sentinel value persisted to .databrickscfg when the diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go index 6aeb16e22be..873ce557dfe 100644 --- a/libs/auth/arguments_test.go +++ b/libs/auth/arguments_test.go @@ -6,7 +6,7 @@ import ( "net/http/httptest" "testing" - "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/databricks/cli/libs/auth/u2m" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index 04f4a0f7e7b..10736a19bd1 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -5,11 +5,11 @@ import ( "errors" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/credentials" "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" - "github.com/databricks/databricks-sdk-go/credentials/u2m" ) // The credentials chain used by the CLI. It is a custom implementation diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 291b501d52e..1c6fceac2d7 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -10,9 +10,9 @@ import ( "testing" "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth" - "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" diff --git a/libs/auth/error.go b/libs/auth/error.go index 5bd5d3d43c6..d942a893ad4 100644 --- a/libs/auth/error.go +++ b/libs/auth/error.go @@ -7,9 +7,9 @@ import ( "net/http" "strings" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/credentials/u2m" ) // Auth type names returned by credential providers. diff --git a/libs/auth/storage/cache.go b/libs/auth/storage/cache.go index 4c874452cae..3754a97db96 100644 --- a/libs/auth/storage/cache.go +++ b/libs/auth/storage/cache.go @@ -5,11 +5,11 @@ import ( "errors" "fmt" + "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/cli/libs/auth/u2m/cache" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go/credentials/u2m" - "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" ) // storeFactories bundles the constructors ResolveStore depends on. Extracted @@ -46,9 +46,9 @@ func defaultStoreFactories() storeFactories { // fallback does not persist auth_storage = plaintext to [__settings__]; // pinning happens only on successful login. // -// Every CLI code path that calls u2m.NewPersistentAuth must route the result -// through u2m.WithTokenCache, otherwise the SDK defaults to the file cache -// and splits the user's tokens across two backends. +// Every CLI code path that calls u2m.NewPersistentAuth must supply the cache +// returned by this package, otherwise U2M uses an in-memory cache and bypasses +// the user's configured storage backend. func ResolveStore(ctx context.Context, override StorageMode) (Store, StorageMode, error) { return resolveStoreForReadWith(ctx, override, defaultStoreFactories()) } @@ -75,7 +75,7 @@ func ResolveStoreForLogin(ctx context.Context, override StorageMode) (Store, Sto return resolveStoreForLoginWith(ctx, override, defaultStoreFactories()) } -// OAuthTokenCache adapts a CLI Store to the SDK's u2m_cache.TokenCache for the +// OAuthTokenCache adapts a CLI Store to the U2M cache.TokenCache for the // U2M PersistentAuth flow, applying the not-found hint so a cache miss carries // actionable "run databricks auth login" guidance. Use on read and credential // paths. M2M/OIDC callers use the CLI Store directly and must not route through diff --git a/libs/auth/storage/cache_test.go b/libs/auth/storage/cache_test.go index ad0279526e1..9cbcf50e17c 100644 --- a/libs/auth/storage/cache_test.go +++ b/libs/auth/storage/cache_test.go @@ -9,9 +9,9 @@ import ( "path/filepath" "testing" + "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/env" - "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" diff --git a/libs/auth/storage/dual_writing_cache.go b/libs/auth/storage/dual_writing_cache.go index 4e12c4b97b2..bb5a4198c9a 100644 --- a/libs/auth/storage/dual_writing_cache.go +++ b/libs/auth/storage/dual_writing_cache.go @@ -1,8 +1,8 @@ package storage import ( - "github.com/databricks/databricks-sdk-go/credentials/u2m" - u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/databricks/cli/libs/auth/u2m" + u2m_cache "github.com/databricks/cli/libs/auth/u2m/cache" "golang.org/x/oauth2" ) @@ -12,7 +12,7 @@ import ( // implemented inside PersistentAuth.dualWrite in the SDK, now moved // caller-side per the cache-ownership split between SDK and CLI. // -// Mirroring happens inside Store, so every SDK-internal write (Challenge, +// Mirroring happens inside Store, so every U2M-internal write (Challenge, // refresh, discovery) dual-writes without requiring each call site to invoke // a helper explicitly. type DualWritingTokenCache struct { diff --git a/libs/auth/storage/dual_writing_cache_test.go b/libs/auth/storage/dual_writing_cache_test.go index c08433c03ac..86b7317ae59 100644 --- a/libs/auth/storage/dual_writing_cache_test.go +++ b/libs/auth/storage/dual_writing_cache_test.go @@ -5,8 +5,8 @@ import ( "sync" "testing" - "github.com/databricks/databricks-sdk-go/credentials/u2m" - u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/databricks/cli/libs/auth/u2m" + u2m_cache "github.com/databricks/cli/libs/auth/u2m/cache" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" diff --git a/libs/auth/storage/keyring.go b/libs/auth/storage/keyring.go index abfedf46d3d..ddc41648abc 100644 --- a/libs/auth/storage/keyring.go +++ b/libs/auth/storage/keyring.go @@ -14,7 +14,7 @@ import ( // keyringServiceName is the service name used for every entry the CLI writes // to the OS-native secure store. The account field carries the per-entry -// cache key the SDK passes through TokenCache.Store / Lookup. +// cache key the U2M manager passes through TokenCache.Store / Lookup. const keyringServiceName = "databricks-cli" // keyringProbeAccountPrefix is prefixed onto a per-call random suffix to form diff --git a/libs/auth/storage/not_found_hint.go b/libs/auth/storage/not_found_hint.go index 5b31db4bfd1..e531eeb8d12 100644 --- a/libs/auth/storage/not_found_hint.go +++ b/libs/auth/storage/not_found_hint.go @@ -7,8 +7,8 @@ import ( "os" "path/filepath" + "github.com/databricks/cli/libs/auth/u2m/cache" "github.com/databricks/cli/libs/env" - "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" "golang.org/x/oauth2" ) @@ -19,7 +19,7 @@ import ( // their cached credentials are no longer being read. // // errors.Is(err, cache.ErrNotFound) continues to return true because the -// wrap uses %w; the SDK's branches on ErrNotFound still fire. +// wrap uses %w; PersistentAuth's branches on ErrNotFound still fire. // // Store is delegated unchanged; only Lookup needs the message polish. type notFoundHintCache struct { @@ -45,7 +45,7 @@ func (c *notFoundHintCache) Lookup(key string) (*oauth2.Token, error) { // notFoundHint replaces cache.ErrNotFound's terse "token not found" string // with an actionable message while still satisfying errors.Is(err, -// cache.ErrNotFound). The SDK's loadToken wraps every cache error with +// cache.ErrNotFound). PersistentAuth.loadToken wraps every cache error with // "cache: %w", and fmt.Errorf("...: %w", ErrNotFound) would tack the // original "token not found" onto the end of our hint, producing // "cache: : token not found". A custom type lets us own the diff --git a/libs/auth/storage/not_found_hint_test.go b/libs/auth/storage/not_found_hint_test.go index cc7505a85bc..d61282b1009 100644 --- a/libs/auth/storage/not_found_hint_test.go +++ b/libs/auth/storage/not_found_hint_test.go @@ -8,7 +8,7 @@ import ( "path/filepath" "testing" - "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/databricks/cli/libs/auth/u2m/cache" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" diff --git a/libs/auth/storage/sdku2m.go b/libs/auth/storage/sdku2m.go deleted file mode 100644 index e9a779e390a..00000000000 --- a/libs/auth/storage/sdku2m.go +++ /dev/null @@ -1,45 +0,0 @@ -package storage - -import ( - "errors" - - "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" - "golang.org/x/oauth2" -) - -// ToU2MTokenCache adapts a CLI Store to the SDK's u2m_cache.TokenCache so it -// can be passed to u2m.WithTokenCache for the U2M PersistentAuth flow, the one -// place the SDK requires that interface. The SDK's Store(key, nil) "delete" -// convention maps to Store.Delete. -func ToU2MTokenCache(s Store) cache.TokenCache { - return &sdkTokenCache{store: s} -} - -// sdkTokenCache is the ToU2MTokenCache adapter. -type sdkTokenCache struct { - store Store -} - -// Store implements u2m_cache.TokenCache. A nil token is the SDK's delete -// signal; everything else is a plain put with no metadata. -func (sdktc *sdkTokenCache) Store(key string, t *oauth2.Token) error { - if t == nil { - return sdktc.store.Delete(key) - } - return sdktc.store.Put(key, Entry{Token: t}) -} - -// Lookup implements cache.TokenCache, translating the CLI miss sentinel to -// the SDK's so the SDK's errors.Is(err, cache.ErrNotFound) branches fire. -func (sdktc *sdkTokenCache) Lookup(key string) (*oauth2.Token, error) { - e, err := sdktc.store.Lookup(key) - if err != nil { - if errors.Is(err, ErrNotFound) { - return nil, cache.ErrNotFound - } - return nil, err - } - return e.Token, nil -} - -var _ cache.TokenCache = (*sdkTokenCache)(nil) diff --git a/libs/auth/storage/storage.go b/libs/auth/storage/storage.go index 785c716d224..dd208a3819e 100644 --- a/libs/auth/storage/storage.go +++ b/libs/auth/storage/storage.go @@ -16,8 +16,8 @@ import ( // ErrNotFound is returned by Store.Lookup when no entry exists for the key, or // when a stored entry cannot be decoded by this CLI version (an unknown format // is treated as a miss so the caller re-mints rather than failing). It is the -// CLI-owned counterpart to the SDK's u2m_cache.ErrNotFound; the adapter in -// ToU2MTokenCache translates between the two. +// CLI-owned counterpart to the U2M cache.ErrNotFound; the adapter in +// ToU2MTokenCache translates between the storage and OAuth layers. var ErrNotFound = errors.New("token not found") // Entry is the value held in the CLI token store. It wraps the credential so @@ -33,7 +33,7 @@ type Entry struct { // policy of its own. Implementations are the plaintext file cache and the OS // keyring cache. The interface is owned by the CLI rather than the SDK so the // entry schema can evolve (metadata, per-entry resilience) without being -// constrained by the SDK's U2M-internal u2m_cache.TokenCache, which only +// constrained by the U2M-internal cache.TokenCache, which only // carries a bare *oauth2.Token. type Store interface { // Put writes e under key, replacing any existing entry. diff --git a/libs/auth/storage/u2m.go b/libs/auth/storage/u2m.go new file mode 100644 index 00000000000..65b5f23fcae --- /dev/null +++ b/libs/auth/storage/u2m.go @@ -0,0 +1,43 @@ +package storage + +import ( + "errors" + + "github.com/databricks/cli/libs/auth/u2m/cache" + "golang.org/x/oauth2" +) + +// ToU2MTokenCache adapts a CLI Store to the U2M cache.TokenCache interface. +// The Store(key, nil) delete convention maps to Store.Delete. +func ToU2MTokenCache(s Store) cache.TokenCache { + return &u2mTokenCache{store: s} +} + +// u2mTokenCache is the ToU2MTokenCache adapter. +type u2mTokenCache struct { + store Store +} + +// Store implements cache.TokenCache. A nil token is the U2M delete +// signal; everything else is a plain put with no metadata. +func (tc *u2mTokenCache) Store(key string, t *oauth2.Token) error { + if t == nil { + return tc.store.Delete(key) + } + return tc.store.Put(key, Entry{Token: t}) +} + +// Lookup implements cache.TokenCache, translating the storage miss sentinel +// to the U2M package's sentinel. +func (tc *u2mTokenCache) Lookup(key string) (*oauth2.Token, error) { + e, err := tc.store.Lookup(key) + if err != nil { + if errors.Is(err, ErrNotFound) { + return nil, cache.ErrNotFound + } + return nil, err + } + return e.Token, nil +} + +var _ cache.TokenCache = (*u2mTokenCache)(nil) diff --git a/libs/auth/u2m/account_oauth_argument.go b/libs/auth/u2m/account_oauth_argument.go new file mode 100644 index 00000000000..0771c9ede0d --- /dev/null +++ b/libs/auth/u2m/account_oauth_argument.go @@ -0,0 +1,74 @@ +package u2m + +import ( + "fmt" +) + +// AccountOAuthArgument is an interface that provides the necessary information +// to authenticate using OAuth to a specific account. +type AccountOAuthArgument interface { + OAuthArgument + + // GetAccountHost returns the host of the account to authenticate to. + GetAccountHost() string + + // GetAccountId returns the account ID of the account to authenticate to. + GetAccountId() string +} + +// BasicAccountOAuthArgument is a basic implementation of the AccountOAuthArgument +// interface that links each account with exactly one OAuth token. +type BasicAccountOAuthArgument struct { + accountHost string + accountID string + + // profile is the optional profile name. When set, GetCacheKey() returns + // the profile name instead of the host-based key. + profile string +} + +var ( + _ AccountOAuthArgument = BasicAccountOAuthArgument{} + _ HostCacheKeyProvider = BasicAccountOAuthArgument{} +) + +// NewBasicAccountOAuthArgument creates a new BasicAccountOAuthArgument. +func NewBasicAccountOAuthArgument(accountsHost, accountID string) (BasicAccountOAuthArgument, error) { + return NewProfileAccountOAuthArgument(accountsHost, accountID, "") +} + +// NewProfileAccountOAuthArgument creates a new BasicAccountOAuthArgument with a +// profile name. When a profile is set, GetCacheKey() returns the profile name +// instead of the host-based key. +func NewProfileAccountOAuthArgument(accountsHost, accountID, profile string) (BasicAccountOAuthArgument, error) { + if err := validateHost(accountsHost); err != nil { + return BasicAccountOAuthArgument{}, err + } + return BasicAccountOAuthArgument{accountHost: accountsHost, accountID: accountID, profile: profile}, nil +} + +// GetAccountHost returns the host of the account to authenticate to. +func (a BasicAccountOAuthArgument) GetAccountHost() string { + return a.accountHost +} + +// GetAccountId returns the account ID of the account to authenticate to. +func (a BasicAccountOAuthArgument) GetAccountId() string { + return a.accountID +} + +// GetCacheKey returns a unique key for caching the OAuth token for the account. +// If a profile is set, the profile name is returned as the cache key. +// Otherwise, the key is in the format "/oidc/accounts/". +func (a BasicAccountOAuthArgument) GetCacheKey() string { + if a.profile != "" { + return a.profile + } + return a.GetHostCacheKey() +} + +// GetHostCacheKey returns the host-based cache key regardless of whether a +// profile is set. The key is in the format "/oidc/accounts/". +func (a BasicAccountOAuthArgument) GetHostCacheKey() string { + return fmt.Sprintf("%s/oidc/accounts/%s", a.accountHost, a.accountID) +} diff --git a/libs/auth/u2m/account_oauth_argument_test.go b/libs/auth/u2m/account_oauth_argument_test.go new file mode 100644 index 00000000000..55b6b74cb28 --- /dev/null +++ b/libs/auth/u2m/account_oauth_argument_test.go @@ -0,0 +1,81 @@ +package u2m + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNewProfileAccountOAuthArgument_ValidatesHost(t *testing.T) { + _, err := NewProfileAccountOAuthArgument("http://insecure.com", "abc", "my-profile") + assert.Error(t, err) + assert.Contains(t, err.Error(), "host must start with 'https://'") + + _, err = NewProfileAccountOAuthArgument("https://accounts.cloud.databricks.test/", "abc", "my-profile") + assert.Error(t, err) + assert.Contains(t, err.Error(), "host must not have a trailing slash") +} + +func TestBasicAccountOAuthArgument_ProfileCacheKeys(t *testing.T) { + tests := []struct { + name string + host string + accountID string + profile string + wantKey string + wantHostKey string + }{ + { + name: "without profile returns host-based key", + host: "https://accounts.cloud.databricks.test", + accountID: "abc", + wantKey: "https://accounts.cloud.databricks.test/oidc/accounts/abc", + wantHostKey: "https://accounts.cloud.databricks.test/oidc/accounts/abc", + }, + { + name: "with profile returns profile name", + host: "https://accounts.cloud.databricks.test", + accountID: "abc", + profile: "my-profile", + wantKey: "my-profile", + wantHostKey: "https://accounts.cloud.databricks.test/oidc/accounts/abc", + }, + { + name: "empty profile returns host-based key", + host: "https://accounts.cloud.databricks.test", + accountID: "abc", + profile: "", + wantKey: "https://accounts.cloud.databricks.test/oidc/accounts/abc", + wantHostKey: "https://accounts.cloud.databricks.test/oidc/accounts/abc", + }, + { + name: "different profiles on same host get different keys", + host: "https://accounts.cloud.databricks.test", + accountID: "abc", + profile: "profile-a", + wantKey: "profile-a", + wantHostKey: "https://accounts.cloud.databricks.test/oidc/accounts/abc", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var arg BasicAccountOAuthArgument + var err error + if tt.profile != "" { + arg, err = NewProfileAccountOAuthArgument(tt.host, tt.accountID, tt.profile) + } else { + arg, err = NewBasicAccountOAuthArgument(tt.host, tt.accountID) + } + assert.NoError(t, err) + assert.Equal(t, tt.wantKey, arg.GetCacheKey()) + assert.Equal(t, tt.wantHostKey, arg.GetHostCacheKey()) + }) + } +} + +func TestBasicAccountOAuthArgument_ImplementsHostCacheKeyProvider(t *testing.T) { + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "abc") + assert.NoError(t, err) + var _ HostCacheKeyProvider = arg +} diff --git a/libs/auth/u2m/cache/cache.go b/libs/auth/u2m/cache/cache.go new file mode 100644 index 00000000000..7dd1bb166ae --- /dev/null +++ b/libs/auth/u2m/cache/cache.go @@ -0,0 +1,32 @@ +/* +Package cache provides an interface for storing and looking up OAuth tokens. + +The cache should be primarily used for user-to-machine (U2M) OAuth flows. In U2M +OAuth flows, the application needs to store the token for later use, such as in +a separate process, and the cache provides a way to do so without requiring the +user to follow the OAuth flow again. + +In machine-to-machine (M2M) OAuth flows, the application is configured with a +secret and can fetch a new token on demand without user interaction, so the +token cache is not necessary. +*/ +package cache + +import ( + "errors" + + "golang.org/x/oauth2" +) + +var ErrNotFound = errors.New("token not found") + +// TokenCache is an interface for storing and looking up OAuth tokens. +type TokenCache interface { + // Store stores the token with the given key, replacing any existing token. + // If t is nil, it deletes the token. + Store(key string, t *oauth2.Token) error + + // Lookup looks up the token with the given key. If the token is not found, it + // returns ErrNotFound. + Lookup(key string) (*oauth2.Token, error) +} diff --git a/libs/auth/u2m/cache/memory.go b/libs/auth/u2m/cache/memory.go new file mode 100644 index 00000000000..401533fd3c7 --- /dev/null +++ b/libs/auth/u2m/cache/memory.go @@ -0,0 +1,53 @@ +package cache + +import ( + "sync" + + "golang.org/x/oauth2" +) + +// NewInMemoryTokenCache returns a TokenCache that stores tokens in process +// memory only. Tokens do not persist across process restarts. This is the +// default cache used by PersistentAuth when no WithTokenCache option is +// provided. Most production consumers should supply a persistent cache +// implementation (for example, the file-based cache in +// github.com/databricks/cli/libs/auth/storage). +func NewInMemoryTokenCache() TokenCache { + return &inMemoryTokenCache{tokens: map[string]*oauth2.Token{}} +} + +type inMemoryTokenCache struct { + mu sync.Mutex + tokens map[string]*oauth2.Token +} + +func cloneToken(t *oauth2.Token) *oauth2.Token { + if t == nil { + return nil + } + clone := *t + return &clone +} + +// Store implements TokenCache. A nil token deletes the key. +func (c *inMemoryTokenCache) Store(key string, t *oauth2.Token) error { + c.mu.Lock() + defer c.mu.Unlock() + if t == nil { + delete(c.tokens, key) + return nil + } + c.tokens[key] = cloneToken(t) + return nil +} + +// Lookup implements TokenCache. Returns ErrNotFound when the key is absent. +func (c *inMemoryTokenCache) Lookup(key string) (*oauth2.Token, error) { + c.mu.Lock() + defer c.mu.Unlock() + t, ok := c.tokens[key] + if !ok { + return nil, ErrNotFound + } + return cloneToken(t), nil +} diff --git a/libs/auth/u2m/cache/memory_test.go b/libs/auth/u2m/cache/memory_test.go new file mode 100644 index 00000000000..0d8a8c44b56 --- /dev/null +++ b/libs/auth/u2m/cache/memory_test.go @@ -0,0 +1,103 @@ +package cache + +import ( + "errors" + "fmt" + "sync" + "testing" + + "golang.org/x/oauth2" +) + +func TestInMemoryTokenCache_StoreAndLookup(t *testing.T) { + c := NewInMemoryTokenCache() + tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "def"} + if err := c.Store("key1", tok); err != nil { + t.Fatalf("Store: %v", err) + } + got, err := c.Lookup("key1") + if err != nil { + t.Fatalf("Lookup: %v", err) + } + if got.AccessToken != "abc" { + t.Errorf("AccessToken: want %q, got %q", "abc", got.AccessToken) + } + if got.RefreshToken != "def" { + t.Errorf("RefreshToken: want %q, got %q", "def", got.RefreshToken) + } +} + +func TestInMemoryTokenCache_StoreAndLookupUseCopies(t *testing.T) { + c := NewInMemoryTokenCache() + tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "def"} + if err := c.Store("key1", tok); err != nil { + t.Fatalf("Store: %v", err) + } + + tok.RefreshToken = "mutated-after-store" + + got, err := c.Lookup("key1") + if err != nil { + t.Fatalf("Lookup: %v", err) + } + if got.RefreshToken != "def" { + t.Fatalf("RefreshToken after store mutation: want %q, got %q", "def", got.RefreshToken) + } + + got.RefreshToken = "mutated-after-lookup" + + gotAgain, err := c.Lookup("key1") + if err != nil { + t.Fatalf("Lookup after lookup mutation: %v", err) + } + if gotAgain.RefreshToken != "def" { + t.Fatalf("RefreshToken after lookup mutation: want %q, got %q", "def", gotAgain.RefreshToken) + } +} + +func TestInMemoryTokenCache_StoreNilDeletesKey(t *testing.T) { + c := NewInMemoryTokenCache() + if err := c.Store("key1", &oauth2.Token{AccessToken: "abc"}); err != nil { + t.Fatalf("Store: %v", err) + } + if err := c.Store("key1", nil); err != nil { + t.Fatalf("Store(nil): %v", err) + } + _, err := c.Lookup("key1") + if !errors.Is(err, ErrNotFound) { + t.Errorf("Lookup after nil Store: want ErrNotFound, got %v", err) + } +} + +func TestInMemoryTokenCache_LookupUnsetKey(t *testing.T) { + c := NewInMemoryTokenCache() + _, err := c.Lookup("missing") + if !errors.Is(err, ErrNotFound) { + t.Errorf("Lookup(missing): want ErrNotFound, got %v", err) + } +} + +func TestInMemoryTokenCache_ConcurrentStoreAndLookup(t *testing.T) { + c := NewInMemoryTokenCache() + const goroutines = 20 + const iterations = 100 + var wg sync.WaitGroup + wg.Add(goroutines * 2) + for i := range goroutines { + prefix := fmt.Sprintf("writer-%d", i) + go func() { + defer wg.Done() + for j := range iterations { + key := fmt.Sprintf("%s-%d", prefix, j) + _ = c.Store(key, &oauth2.Token{AccessToken: key}) + } + }() + go func() { + defer wg.Done() + for j := range iterations { + _, _ = c.Lookup(fmt.Sprintf("%s-%d", prefix, j)) + } + }() + } + wg.Wait() +} diff --git a/libs/auth/u2m/callback.go b/libs/auth/u2m/callback.go new file mode 100644 index 00000000000..5db0eeeac6a --- /dev/null +++ b/libs/auth/u2m/callback.go @@ -0,0 +1,156 @@ +package u2m + +import ( + "context" + _ "embed" + "fmt" + "html/template" + "net/http" + "strings" + + "golang.org/x/text/cases" + "golang.org/x/text/language" +) + +//go:embed page.tmpl +var pageTmpl string + +type oauthResult struct { + Error string + ErrorDescription string + State string + Code string + Issuer string + Host string +} + +// callbackServer is a server that listens for the redirect from the Databricks +// identity provider. It renders a page.html template that shows the result of +// the authentication attempt. +type callbackServer struct { + // ctx is the context used when waiting for the redirect from the identity + // provider. This is needed because the Handler() method from the oauth2 + // library does not accept a context. + ctx context.Context + + // srv is the server that listens for the redirect from the identity provider. + srv http.Server + + // browser is a function that opens a browser to the given URL. + browser func(string) error + + // arg is the OAuth argument used to authenticate. + arg OAuthArgument + + // renderErrCh is a channel that receives an error if there is an error + // rendering the page.html template. + renderErrCh chan error + + // feedbackCh is a channel that receives the result of the authentication + // attempt. + feedbackCh chan oauthResult + + // tmpl is the template used to render the response page after the user is + // redirected back to the callback server. + tmpl *template.Template +} + +// newCallbackServer creates a new callback server that listens for the redirect +// from the Databricks identity provider. +func (a *PersistentAuth) newCallbackServer() (*callbackServer, error) { + tmpl, err := template.New("page").Funcs(template.FuncMap{ + "title": func(in string) string { + title := cases.Title(language.English) + return title.String(strings.ReplaceAll(in, "_", " ")) + }, + }).Parse(pageTmpl) + if err != nil { + return nil, err + } + cb := &callbackServer{ + feedbackCh: make(chan oauthResult), + renderErrCh: make(chan error), + tmpl: tmpl, + ctx: a.ctx, + browser: a.browser, + arg: a.oAuthArgument, + } + cb.srv.Handler = cb + go func() { + _ = cb.srv.Serve(a.ln) + }() + return cb, nil +} + +// Close closes the callback server. +func (cb *callbackServer) Close() error { + return cb.srv.Close() +} + +// ServeHTTP renders the page.html template. +func (cb *callbackServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + res := oauthResult{ + Error: r.FormValue("error"), + ErrorDescription: r.FormValue("error_description"), + Code: r.FormValue("code"), + State: r.FormValue("state"), + Issuer: r.FormValue("iss"), + Host: cb.getHost(), + } + if res.Error != "" { + w.WriteHeader(http.StatusBadRequest) + } else { + w.WriteHeader(http.StatusOK) + } + err := cb.tmpl.Execute(w, res) + if err != nil { + cb.renderErrCh <- err + } + cb.feedbackCh <- res +} + +func (cb *callbackServer) getHost() string { + switch a := cb.arg.(type) { + case AccountOAuthArgument: + return a.GetAccountHost() + case WorkspaceOAuthArgument: + return a.GetWorkspaceHost() + default: + return "" + } +} + +func (cb *callbackServer) awaitResult(authCodeURL string) (oauthResult, error) { + err := cb.browser(authCodeURL) + if err != nil { + fmt.Printf("Please continue the authentication process in your browser:\n%s\n", authCodeURL) + } + select { + case <-cb.ctx.Done(): + return oauthResult{}, cb.ctx.Err() + case renderErr := <-cb.renderErrCh: + return oauthResult{}, renderErr + case res := <-cb.feedbackCh: + if res.Error != "" { + return oauthResult{}, fmt.Errorf("%s: %s", res.Error, res.ErrorDescription) + } + return res, nil + } +} + +// Handler opens up a browser waits for redirect to come back from the identity provider +func (cb *callbackServer) Handler(authCodeURL string) (string, string, error) { + res, err := cb.awaitResult(authCodeURL) + if err != nil { + return "", "", err + } + return res.Code, res.State, nil +} + +func (cb *callbackServer) handlerWithIssuer(authCodeURL string) (code, state, issuer string, err error) { + res, err := cb.awaitResult(authCodeURL) + if err != nil { + return "", "", "", err + } + return res.Code, res.State, res.Issuer, nil +} diff --git a/libs/auth/u2m/callback_test.go b/libs/auth/u2m/callback_test.go new file mode 100644 index 00000000000..147d16cec4f --- /dev/null +++ b/libs/auth/u2m/callback_test.go @@ -0,0 +1,244 @@ +package u2m + +import ( + "fmt" + "html/template" + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "golang.org/x/oauth2" +) + +func TestCallbackServer_HandlerWithIssuerBindsIssuerToResult(t *testing.T) { + browserOpened := make(chan string, 1) + cb := &callbackServer{ + ctx: t.Context(), + browser: func(redirect string) error { + browserOpened <- redirect + return nil + }, + renderErrCh: make(chan error), + feedbackCh: make(chan oauthResult, 2), + tmpl: template.Must(template.New("page").Parse("")), + } + + const authCodeURL = "https://login.databricks.test/?destination_url=%2Foidc%2Fv1%2Fauthorize" + legitimate := struct { + code string + state string + issuer string + }{ + code: "legit-code", + state: "legit-state", + issuer: "https://adb-123.azuredatabricks.test/oidc", + } + + serveCallback := func(code, state, issuer string) { + callbackURL := fmt.Sprintf("/?code=%s&state=%s&iss=%s", + url.QueryEscape(code), url.QueryEscape(state), url.QueryEscape(issuer)) + req := httptest.NewRequest(http.MethodGet, callbackURL, nil) + rec := httptest.NewRecorder() + cb.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("ServeHTTP status = %d, want %d", rec.Code, http.StatusOK) + } + } + + // Queue two callbacks with different code/state/iss; the handler must + // return the values from the first result and not interleave. + serveCallback(legitimate.code, legitimate.state, legitimate.issuer) + serveCallback("second-code", "second-state", "https://other.example/oidc") + + code, state, issuer, err := cb.handlerWithIssuer(authCodeURL) + if err != nil { + t.Fatalf("handlerWithIssuer(): %v", err) + } + if code != legitimate.code { + t.Errorf("code = %q, want %q", code, legitimate.code) + } + if state != legitimate.state { + t.Errorf("state = %q, want %q", state, legitimate.state) + } + if issuer != legitimate.issuer { + t.Errorf("issuer = %q, want %q", issuer, legitimate.issuer) + } + + select { + case got := <-browserOpened: + if got != authCodeURL { + t.Errorf("browser opened %q, want %q", got, authCodeURL) + } + default: + t.Fatal("browser was not opened") + } +} + +func TestCallbackServer_ExtractsIssuer(t *testing.T) { + ctx := t.Context() + + browserOpened := make(chan string) + browser := func(redirect string) error { + browserOpened <- redirect + return nil + } + tokenCache := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + + p, err := NewPersistentAuth( + ctx, + WithTokenCache(tokenCache), + WithBrowser(browser), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=__ACCESS__&refresh_token=__REFRESH__`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + // Start a listener so we can create the callback server. + err = p.startListener(ctx) + if err != nil { + t.Fatalf("startListener(): %v", err) + } + + cb, err := p.newCallbackServer() + if err != nil { + t.Fatalf("newCallbackServer(): %v", err) + } + defer cb.Close() + + // Fire a callback with iss parameter. + issuerURL := "https://adb-123.azuredatabricks.test/oidc" + resp, err := http.Get(fmt.Sprintf("http://%s?code=xxx&state=yyy&iss=%s", p.redirectAddr, issuerURL)) + if err != nil { + t.Fatalf("http.Get(): %v", err) + } + defer resp.Body.Close() + + res := <-cb.feedbackCh + if got := res.Issuer; got != issuerURL { + t.Fatalf("result Issuer: want %q, got %q", issuerURL, got) + } +} + +func TestCallbackServer_NoIssuer(t *testing.T) { + ctx := t.Context() + + tokenCache := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + + p, err := NewPersistentAuth( + ctx, + WithTokenCache(tokenCache), + WithBrowser(func(string) error { return nil }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + err = p.startListener(ctx) + if err != nil { + t.Fatalf("startListener(): %v", err) + } + + cb, err := p.newCallbackServer() + if err != nil { + t.Fatalf("newCallbackServer(): %v", err) + } + defer cb.Close() + + // Fire a callback without iss parameter. + resp, err := http.Get(fmt.Sprintf("http://%s?code=xxx&state=yyy", p.redirectAddr)) + if err != nil { + t.Fatalf("http.Get(): %v", err) + } + defer resp.Body.Close() + + res := <-cb.feedbackCh + if got := res.Issuer; got != "" { + t.Fatalf("result Issuer: want %q, got %q", "", got) + } +} + +func TestCallbackServer_IssuerWithAccountPath(t *testing.T) { + ctx := t.Context() + + tokenCache := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + + p, err := NewPersistentAuth( + ctx, + WithTokenCache(tokenCache), + WithBrowser(func(string) error { return nil }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + err = p.startListener(ctx) + if err != nil { + t.Fatalf("startListener(): %v", err) + } + + cb, err := p.newCallbackServer() + if err != nil { + t.Fatalf("newCallbackServer(): %v", err) + } + defer cb.Close() + + // Fire a callback with iss containing an account path. + issuerURL := "https://nike.databricks.test/oidc/accounts/abc-123-def" + resp, err := http.Get(fmt.Sprintf("http://%s?code=xxx&state=yyy&iss=%s", p.redirectAddr, issuerURL)) + if err != nil { + t.Fatalf("http.Get(): %v", err) + } + defer resp.Body.Close() + + res := <-cb.feedbackCh + if got := res.Issuer; got != issuerURL { + t.Fatalf("result Issuer: want %q, got %q", issuerURL, got) + } +} diff --git a/libs/auth/u2m/discovery_oauth_argument.go b/libs/auth/u2m/discovery_oauth_argument.go new file mode 100644 index 00000000000..7fbe4a04862 --- /dev/null +++ b/libs/auth/u2m/discovery_oauth_argument.go @@ -0,0 +1,63 @@ +package u2m + +import "errors" + +// DiscoveryOAuthArgument is an OAuthArgument for bootstrapping the +// login.databricks.com discovery flow. Unlike other OAuthArgument types, it +// has no host at construction time. The host is discovered from the iss +// parameter in the OAuth callback after the user authenticates and selects a +// workspace. +// +// DiscoveryOAuthArgument is only intended for use with [WithDiscoveryLogin] +// during [PersistentAuth.Challenge]. Once the host has been discovered, +// callers should construct the usual host-based OAuthArgument +// (for example [WorkspaceOAuthArgument]) for future PersistentAuth instances. +type DiscoveryOAuthArgument interface { + OAuthArgument + + // SetDiscoveredHost stores the workspace host derived from the iss + // callback parameter. + SetDiscoveredHost(host string) + + // GetDiscoveredHost returns the workspace host discovered from the + // callback, or empty string if not yet discovered. + GetDiscoveredHost() string +} + +// BasicDiscoveryOAuthArgument is a basic implementation of +// [DiscoveryOAuthArgument] that uses the profile name as the cache key during +// discovery login bootstrap. +type BasicDiscoveryOAuthArgument struct { + profile string + discoveredHost string +} + +// NewBasicDiscoveryOAuthArgument creates a new [BasicDiscoveryOAuthArgument]. +// The profile name is required and used as the cache key during discovery +// login. Use it with [WithDiscoveryLogin]. After discovery, construct the +// usual host-based OAuthArgument with the discovered host. +func NewBasicDiscoveryOAuthArgument(profile string) (*BasicDiscoveryOAuthArgument, error) { + if profile == "" { + return nil, errors.New("profile name must not be empty for discovery login") + } + return &BasicDiscoveryOAuthArgument{profile: profile}, nil +} + +// GetCacheKey returns the profile name as the cache key. +func (a *BasicDiscoveryOAuthArgument) GetCacheKey() string { + return a.profile +} + +// SetDiscoveredHost stores the workspace host derived from the iss callback +// parameter. +func (a *BasicDiscoveryOAuthArgument) SetDiscoveredHost(host string) { + a.discoveredHost = host +} + +// GetDiscoveredHost returns the workspace host discovered from the callback, +// or empty string if not yet discovered. +func (a *BasicDiscoveryOAuthArgument) GetDiscoveredHost() string { + return a.discoveredHost +} + +var _ DiscoveryOAuthArgument = &BasicDiscoveryOAuthArgument{} diff --git a/libs/auth/u2m/discovery_oauth_argument_test.go b/libs/auth/u2m/discovery_oauth_argument_test.go new file mode 100644 index 00000000000..1eba2bd6c95 --- /dev/null +++ b/libs/auth/u2m/discovery_oauth_argument_test.go @@ -0,0 +1,106 @@ +package u2m + +import ( + "strings" + "testing" +) + +func TestDiscoveryOAuthArgument_GetCacheKey(t *testing.T) { + arg, err := NewBasicDiscoveryOAuthArgument("my-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): want no error, got %v", err) + } + got := arg.GetCacheKey() + if got != "my-profile" { + t.Errorf("GetCacheKey(): want %q, got %q", "my-profile", got) + } +} + +func TestDiscoveryOAuthArgument_SetAndGetDiscoveredHost(t *testing.T) { + arg, err := NewBasicDiscoveryOAuthArgument("test-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): want no error, got %v", err) + } + arg.SetDiscoveredHost("https://adb-123.azuredatabricks.test") + got := arg.GetDiscoveredHost() + if got != "https://adb-123.azuredatabricks.test" { + t.Errorf("GetDiscoveredHost(): want %q, got %q", "https://adb-123.azuredatabricks.test", got) + } +} + +func TestDiscoveryOAuthArgument_ZeroValueHasEmptyDiscoveredHost(t *testing.T) { + arg, err := NewBasicDiscoveryOAuthArgument("test-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): want no error, got %v", err) + } + got := arg.GetDiscoveredHost() + if got != "" { + t.Errorf("GetDiscoveredHost(): want empty string, got %q", got) + } +} + +func TestDiscoveryOAuthArgument_ConstructorRejectsEmptyProfile(t *testing.T) { + _, err := NewBasicDiscoveryOAuthArgument("") + if err == nil { + t.Fatal("NewBasicDiscoveryOAuthArgument(): want error for empty profile, got nil") + } +} + +func TestDiscoveryOAuthArgument_NewPersistentAuthRequiresDiscoveryLogin(t *testing.T) { + arg, err := NewBasicDiscoveryOAuthArgument("discovery-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): want no error, got %v", err) + } + _, err = NewPersistentAuth(t.Context(), WithOAuthArgument(arg)) + if err == nil { + t.Fatal("NewPersistentAuth(): want error for DiscoveryOAuthArgument without WithDiscoveryLogin, got nil") + } + if !strings.Contains(err.Error(), "requires WithDiscoveryLogin") { + t.Fatalf("NewPersistentAuth(): want discovery login requirement error, got %v", err) + } +} + +func TestDiscoveryOAuthArgument_NewPersistentAuthAcceptsItWithDiscoveryLogin(t *testing.T) { + arg, err := NewBasicDiscoveryOAuthArgument("discovery-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithOAuthArgument(arg), + WithDiscoveryLogin(), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() +} + +func TestDiscoveryOAuthArgument_DiscoveryLoginRejectsNonDiscoveryArgument(t *testing.T) { + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + _, err = NewPersistentAuth( + t.Context(), + WithOAuthArgument(arg), + WithDiscoveryLogin(), + ) + if err == nil { + t.Fatal("NewPersistentAuth(): want error for discovery login without DiscoveryOAuthArgument, got nil") + } + if !strings.Contains(err.Error(), "discovery login requires DiscoveryOAuthArgument") { + t.Fatalf("NewPersistentAuth(): want discovery login validation error, got %v", err) + } +} + +type unsupportedOAuthArgument struct{} + +func (u unsupportedOAuthArgument) GetCacheKey() string { return "unsupported" } + +func TestDiscoveryOAuthArgument_ValidateArgRejectsUnknownTypes(t *testing.T) { + _, err := NewPersistentAuth(t.Context(), WithOAuthArgument(unsupportedOAuthArgument{})) + if err == nil { + t.Fatal("NewPersistentAuth(): want error for unsupported OAuthArgument type, got nil") + } +} diff --git a/libs/auth/u2m/discovery_token_source.go b/libs/auth/u2m/discovery_token_source.go new file mode 100644 index 00000000000..daa0ee118e4 --- /dev/null +++ b/libs/auth/u2m/discovery_token_source.go @@ -0,0 +1,191 @@ +package u2m + +import ( + "errors" + "fmt" + "net/url" + "strings" + + "golang.org/x/oauth2" +) + +// defaultLoginDatabricksHost is the production host used for discovery login +// when no override is configured via WithDiscoveryHost. +const defaultLoginDatabricksHost = "https://login.databricks.com" + +// DeriveHostFromIssuer extracts the workspace host (scheme + host) from an +// issuer URL returned in the OAuth callback iss parameter. +// +// Examples: +// +// "https://adb-xxx.azuredatabricks.net/oidc" -> "https://adb-xxx.azuredatabricks.net" +// "https://nike.databricks.com/oidc/accounts/xxx" -> "https://nike.databricks.com" +func DeriveHostFromIssuer(issuer string) (string, error) { + if issuer == "" { + return "", errors.New("issuer must not be empty") + } + u, err := url.Parse(issuer) + if err != nil { + return "", fmt.Errorf("parsing issuer URL %q: %w", issuer, err) + } + // Allow http for localhost (consistent with validateHost in workspace_oauth_argument.go). + local := u.Scheme == "http" && u.Hostname() == "127.0.0.1" + if u.Scheme != "https" && !local { + return "", fmt.Errorf("issuer must use https scheme: %q", issuer) + } + if u.Host == "" { + return "", fmt.Errorf("issuer must have a non-empty host: %q", issuer) + } + return fmt.Sprintf("%s://%s", u.Scheme, u.Host), nil +} + +// DeriveTokenEndpoint derives the token endpoint from an issuer URL by +// appending /v1/token to the issuer path. +// +// Example: +// +// "https://adb-xxx.net/oidc" -> "https://adb-xxx.net/oidc/v1/token" +func DeriveTokenEndpoint(issuer string) string { + return strings.TrimRight(issuer, "/") + "/v1/token" +} + +// discoveryTargetAccount is the value of the `target` query parameter that +// tells login.databricks.com to land the user on the account selector instead +// of the workspace selector. Used when the caller has signalled (e.g. via +// WithDiscoveryAccountTarget) that they only want account-level access. +const discoveryTargetAccount = "ACCOUNT" + +// BuildDiscoveryAuthorizeURL builds the login.databricks.com URL that initiates +// the discovery OAuth flow. The OIDC authorize path with all OAuth query params +// is URL-encoded as the destination_url parameter. +func BuildDiscoveryAuthorizeURL(redirectAddr, state string, pkce PKCEParams, scopes []string) string { + return buildDiscoveryAuthorizeURL(defaultLoginDatabricksHost, redirectAddr, state, pkce, scopes, "") +} + +// buildDiscoveryAuthorizeURL builds the discovery authorize URL against the +// given host. Trailing slashes on host are trimmed so the result is +// well-formed regardless of how an override is written. When target is +// non-empty it is set as the top-level `target` query parameter, which +// login.databricks.com uses to route the user to a specific selector page +// (e.g. "ACCOUNT" for the account selector). +func buildDiscoveryAuthorizeURL(host, redirectAddr, state string, pkce PKCEParams, scopes []string, target string) string { + // Build the nested OIDC authorize path with query parameters. + authParams := url.Values{} + authParams.Set("client_id", appClientID) + authParams.Set("redirect_uri", "http://"+redirectAddr) + authParams.Set("response_type", "code") + authParams.Set("scope", strings.Join(scopes, " ")) + authParams.Set("state", state) + authParams.Set("code_challenge", pkce.Challenge) + authParams.Set("code_challenge_method", pkce.ChallengeMethod) + destinationURL := "/oidc/v1/authorize?" + authParams.Encode() + + // Wrap the authorize path as the destination_url query parameter on the + // discovery host. + topParams := url.Values{} + if target != "" { + topParams.Set("target", target) + } + topParams.Set("destination_url", destinationURL) + return strings.TrimRight(host, "/") + "/?" + topParams.Encode() +} + +// PKCEParams holds the PKCE challenge parameters used to build the discovery +// authorize URL. This mirrors authhandler.PKCEParams but is used directly so +// the caller does not need to import authhandler. +type PKCEParams struct { + Challenge string + ChallengeMethod string + Verifier string +} + +// discoveryTokenSource handles the OAuth PKCE flow for login.databricks.com +// discovery login. Unlike the standard flow, the token endpoint is not known +// until the callback provides the iss parameter identifying the workspace. +type discoveryTokenSource struct { + pa *PersistentAuth + // host overrides defaultLoginDatabricksHost when non-empty. + host string + // target is the value of the top-level `target` query parameter on the + // authorize URL. When non-empty (e.g. "ACCOUNT"), login.databricks.com + // routes the user directly to the corresponding selector. + target string +} + +// challenge initiates the discovery OAuth flow through login.databricks.com. +// It builds a custom authorize URL, opens the browser, waits for the callback, +// derives the workspace host and token endpoint from the iss parameter, and +// exchanges the authorization code for tokens. +func (d *discoveryTokenSource) challenge() error { + cb, err := d.pa.newCallbackServer() + if err != nil { + return fmt.Errorf("callback server: %w", err) + } + defer cb.Close() + + state, authPKCE, err := d.pa.stateAndPKCE() + if err != nil { + return fmt.Errorf("state and pkce: %w", err) + } + + scopes := d.pa.resolveScopes() + + pkce := PKCEParams{ + Challenge: authPKCE.Challenge, + ChallengeMethod: authPKCE.ChallengeMethod, + Verifier: authPKCE.Verifier, + } + host := d.host + if host == "" { + host = defaultLoginDatabricksHost + } + authorizeURL := buildDiscoveryAuthorizeURL(host, d.pa.redirectAddr, state, pkce, scopes, d.target) + + code, returnedState, issuer, err := cb.handlerWithIssuer(authorizeURL) + if err != nil { + return fmt.Errorf("authorize: %w", err) + } + + // Validate state matches what we generated before consuming callback data. + if returnedState != state { + return fmt.Errorf("state mismatch: expected %q, got %q", state, returnedState) + } + + if issuer == "" { + return errors.New("discovery login failed: callback did not include an issuer (iss) parameter") + } + + // Derive host and token endpoint from the issuer. + discoveredHost, err := DeriveHostFromIssuer(issuer) + if err != nil { + return fmt.Errorf("deriving host from issuer: %w", err) + } + tokenEndpoint := DeriveTokenEndpoint(issuer) + + // Exchange authorization code for tokens. + cfg := &oauth2.Config{ + ClientID: appClientID, + Endpoint: oauth2.Endpoint{ + TokenURL: tokenEndpoint, + AuthStyle: oauth2.AuthStyleInParams, + }, + RedirectURL: "http://" + d.pa.redirectAddr, + } + ctx := d.pa.setOAuthContext(d.pa.ctx) + token, err := cfg.Exchange(ctx, code, + oauth2.SetAuthURLParam("code_verifier", pkce.Verifier)) + if err != nil { + return fmt.Errorf("token exchange: %w", err) + } + + discoveryArg, ok := d.pa.oAuthArgument.(DiscoveryOAuthArgument) + if !ok { + return fmt.Errorf("discovery login requires DiscoveryOAuthArgument, got %T", d.pa.oAuthArgument) + } + discoveryArg.SetDiscoveredHost(discoveredHost) + + if err := d.pa.cache.Store(d.pa.oAuthArgument.GetCacheKey(), token); err != nil { + return fmt.Errorf("storing token: %w", err) + } + return nil +} diff --git a/libs/auth/u2m/discovery_token_source_test.go b/libs/auth/u2m/discovery_token_source_test.go new file mode 100644 index 00000000000..8941f909e82 --- /dev/null +++ b/libs/auth/u2m/discovery_token_source_test.go @@ -0,0 +1,394 @@ +package u2m + +import ( + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "testing" + "time" + + "golang.org/x/oauth2" +) + +func TestDeriveHostFromIssuer(t *testing.T) { + tests := []struct { + name string + issuer string + want string + wantErr bool + }{ + { + name: "workspace oidc", + issuer: "https://adb-xxx.azuredatabricks.test/oidc", + want: "https://adb-xxx.azuredatabricks.test", + }, + { + name: "workspace cloud", + issuer: "https://workspace.cloud.databricks.test/oidc", + want: "https://workspace.cloud.databricks.test", + }, + { + name: "spog with account path", + issuer: "https://nike.databricks.test/oidc/accounts/xxx", + want: "https://nike.databricks.test", + }, + { + name: "empty issuer", + issuer: "", + wantErr: true, + }, + { + name: "http scheme", + issuer: "http://insecure.net/oidc", + wantErr: true, + }, + { + name: "no host", + issuer: "https:///oidc", + wantErr: true, + }, + { + name: "http localhost allowed", + issuer: "http://127.0.0.1:12345/oidc", + want: "http://127.0.0.1:12345", + }, + { + name: "http non-local host rejected", + issuer: "http://127.0.0.1.attacker.tld/oidc", + wantErr: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := DeriveHostFromIssuer(tc.issuer) + if tc.wantErr { + if err == nil { + t.Fatalf("DeriveHostFromIssuer(%q): want error, got nil", tc.issuer) + } + return + } + if err != nil { + t.Fatalf("DeriveHostFromIssuer(%q): unexpected error: %v", tc.issuer, err) + } + if got != tc.want { + t.Errorf("DeriveHostFromIssuer(%q) = %q, want %q", tc.issuer, got, tc.want) + } + }) + } +} + +func TestDeriveTokenEndpoint(t *testing.T) { + tests := []struct { + name string + issuer string + want string + }{ + { + name: "standard", + issuer: "https://adb-xxx.net/oidc", + want: "https://adb-xxx.net/oidc/v1/token", + }, + { + name: "trailing slash", + issuer: "https://adb-xxx.net/oidc/", + want: "https://adb-xxx.net/oidc/v1/token", + }, + { + name: "account path", + issuer: "https://nike.databricks.test/oidc/accounts/abc123", + want: "https://nike.databricks.test/oidc/accounts/abc123/v1/token", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := DeriveTokenEndpoint(tc.issuer) + if got != tc.want { + t.Errorf("DeriveTokenEndpoint(%q) = %q, want %q", tc.issuer, got, tc.want) + } + }) + } +} + +func TestBuildDiscoveryAuthorizeURL(t *testing.T) { + pkce := PKCEParams{ + Challenge: "test-challenge", + ChallengeMethod: "S256", + Verifier: "test-verifier", + } + scopes := []string{"offline_access", "all-apis"} + redirectAddr := "localhost:8020" + state := "test-state" + + got := BuildDiscoveryAuthorizeURL(redirectAddr, state, pkce, scopes) + + // Parse the top-level URL. + u, err := url.Parse(got) + if err != nil { + t.Fatalf("parsing URL: %v", err) + } + if u.Scheme != "https" || u.Host != "login.databricks.com" { + t.Errorf("want host https://login.databricks.com, got %s://%s", u.Scheme, u.Host) + } + + // Extract and decode destination_url. + destURL := u.Query().Get("destination_url") + if destURL == "" { + t.Fatal("destination_url query param is empty") + } + + // Parse the destination_url as a relative URL with query params. + destParsed, err := url.Parse(destURL) + if err != nil { + t.Fatalf("parsing destination_url: %v", err) + } + if destParsed.Path != "/oidc/v1/authorize" { + t.Errorf("destination_url path = %q, want %q", destParsed.Path, "/oidc/v1/authorize") + } + + // Verify all expected OAuth params. + q := destParsed.Query() + expectations := map[string]string{ + "client_id": appClientID, + "redirect_uri": "http://localhost:8020", + "response_type": "code", + "scope": "offline_access all-apis", + "state": "test-state", + "code_challenge": "test-challenge", + "code_challenge_method": "S256", + } + for key, want := range expectations { + if got := q.Get(key); got != want { + t.Errorf("destination_url param %q = %q, want %q", key, got, want) + } + } +} + +func TestBuildDiscoveryAuthorizeURL_HostOverride(t *testing.T) { + pkce := PKCEParams{ + Challenge: "c", + ChallengeMethod: "S256", + Verifier: "v", + } + scopes := []string{"offline_access", "all-apis"} + tests := []struct { + name string + host string + }{ + { + name: "custom host", + host: "https://login.dev.databricks.test", + }, + { + name: "custom host with trailing slash", + host: "https://login.dev.databricks.test/", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := buildDiscoveryAuthorizeURL(tc.host, "localhost:8020", "s", pkce, scopes, "") + u, err := url.Parse(got) + if err != nil { + t.Fatalf("parsing URL: %v", err) + } + if u.Host != "login.dev.databricks.test" { + t.Errorf("host = %q, want login.dev.databricks.test", u.Host) + } + }) + } +} + +func TestBuildDiscoveryAuthorizeURL_Target(t *testing.T) { + pkce := PKCEParams{ + Challenge: "c", + ChallengeMethod: "S256", + Verifier: "v", + } + scopes := []string{"offline_access", "all-apis"} + tests := []struct { + name string + target string + wantTarget string + }{ + {name: "no target", target: "", wantTarget: ""}, + {name: "account target", target: "ACCOUNT", wantTarget: "ACCOUNT"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := buildDiscoveryAuthorizeURL(defaultLoginDatabricksHost, "localhost:8020", "s", pkce, scopes, tc.target) + u, err := url.Parse(got) + if err != nil { + t.Fatalf("parsing URL: %v", err) + } + if g := u.Query().Get("target"); g != tc.wantTarget { + t.Errorf("target = %q, want %q", g, tc.wantTarget) + } + // destination_url must still be present in every variant. + if u.Query().Get("destination_url") == "" { + t.Error("destination_url should be set regardless of target") + } + }) + } +} + +func TestWithDiscoveryAccountTarget(t *testing.T) { + var a PersistentAuth + if a.discoveryAccountTarget { + t.Fatal("discoveryAccountTarget should default to false") + } + WithDiscoveryAccountTarget()(&a) + if !a.discoveryAccountTarget { + t.Error("WithDiscoveryAccountTarget did not set discoveryAccountTarget") + } +} + +func TestWithDiscoveryHost_NormalizesScheme(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + {name: "empty stays empty", input: "", want: ""}, + {name: "https preserved", input: "https://login.dev.databricks.test", want: "https://login.dev.databricks.test"}, + {name: "http preserved", input: "http://localhost:8080", want: "http://localhost:8080"}, + {name: "no scheme gets https", input: "login.dev.databricks.test", want: "https://login.dev.databricks.test"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var a PersistentAuth + WithDiscoveryHost(tc.input)(&a) + if a.discoveryHost != tc.want { + t.Errorf("discoveryHost = %q, want %q", a.discoveryHost, tc.want) + } + }) + } +} + +func TestDiscoveryTokenSource_Challenge(t *testing.T) { + // Create a mock token server that responds to POST /oidc/v1/token. + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("token server: want POST, got %s", r.Method) + } + if r.URL.Path != "/oidc/v1/token" { + t.Errorf("token server: want path /oidc/v1/token, got %s", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"access_token":"test-access-token","refresh_token":"test-refresh-token","token_type":"Bearer","expires_in":3600}`) + })) + defer tokenServer.Close() + + // The issuer will be the mock server URL + /oidc. + issuer := tokenServer.URL + "/oidc" + + browserOpened := make(chan string, 1) + browserMock := func(u string) error { + browserOpened <- u + return nil + } + + storedTokens := map[string]*oauth2.Token{} + cacheMock := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + storedTokens[key] = tok + return nil + }, + } + + arg, err := NewBasicDiscoveryOAuthArgument("test-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): %v", err) + } + + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(cacheMock), + WithBrowser(browserMock), + WithHttpClient(tokenServer.Client()), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + WithDiscoveryLogin(), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + + // Start the listener so redirectAddr is set. + err = p.startListener(t.Context()) + if err != nil { + t.Fatalf("startListener(): %v", err) + } + defer p.Close() + + dts := &discoveryTokenSource{pa: p} + + errc := make(chan error, 1) + go func() { + errc <- dts.challenge() + }() + + // Wait for browser to be called and extract state from the URL. + var state string + select { + case authURL := <-browserOpened: + u, err := url.Parse(authURL) + if err != nil { + t.Fatalf("parsing auth URL: %v", err) + } + destURL := u.Query().Get("destination_url") + dest, err := url.Parse(destURL) + if err != nil { + t.Fatalf("parsing destination_url: %v", err) + } + state = dest.Query().Get("state") + if state == "" { + t.Fatal("state is empty in authorize URL") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for browser to be called") + } + + // Fire the callback with code, state, and iss. + callbackURL := fmt.Sprintf("http://%s?code=test-code&state=%s&iss=%s", + p.redirectAddr, url.QueryEscape(state), url.QueryEscape(issuer)) + resp, err := http.Get(callbackURL) + if err != nil { + t.Fatalf("callback GET: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("callback: want status 200, got %d", resp.StatusCode) + } + + // Wait for challenge to complete. + select { + case err := <-errc: + if err != nil { + t.Fatalf("challenge(): %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for challenge to complete") + } + + // Verify discovered host was set. + expectedHost, err := DeriveHostFromIssuer(issuer) + if err != nil { + t.Fatalf("DeriveHostFromIssuer(%q): %v", issuer, err) + } + if arg.GetDiscoveredHost() != expectedHost { + t.Errorf("discovered host = %q, want %q", arg.GetDiscoveredHost(), expectedHost) + } + if len(storedTokens) != 1 { + t.Fatalf("store count: want 1 key (profile), got %d", len(storedTokens)) + } + storedToken := storedTokens["test-profile"] + if storedToken == nil { + t.Fatalf("stored token for profile key is nil") + } + if storedToken.AccessToken != "test-access-token" { + t.Errorf("access token = %q, want %q", storedToken.AccessToken, "test-access-token") + } + if storedToken.RefreshToken != "test-refresh-token" { + t.Errorf("refresh token = %q, want %q", storedToken.RefreshToken, "test-refresh-token") + } +} diff --git a/libs/auth/u2m/doc.go b/libs/auth/u2m/doc.go new file mode 100644 index 00000000000..fa8882d95a1 --- /dev/null +++ b/libs/auth/u2m/doc.go @@ -0,0 +1,47 @@ +/* +Package u2m supports the user-to-machine (U2M) OAuth flow for authenticating with Databricks. + +Databricks uses the authorization code flow from OAuth 2.0 to authenticate users. This flow +consists of four steps: + 1. Retrieve an authorization code for a user by opening a browser and directing them to the + Databricks authorization URL. + 2. Exchange the authorization code for an access token. + 3. Use the access token to authenticate with Databricks. + 4. When the access token expires, use the refresh token to get a new access token. + +The token and authorization endpoints for Databricks vary depending on whether the host is +an account- or workspace-level host. Account-level endpoints are fixed based on the account +ID and host, while workspace-level endpoints are discovered using the OIDC discovery endpoint +at /oidc/.well-known/oauth-authorization-server. + +For host-agnostic login through login.databricks.com, use WithDiscoveryLogin together with a +DiscoveryOAuthArgument during Challenge(). This is a bootstrap flow only. Once the callback +reveals the workspace host, construct the usual host-based OAuthArgument for future +PersistentAuth instances. + +To trigger the authorization flow, construct a PersistentAuth object with an +OAuthArgument and call Challenge: + + arg, err := NewProfileWorkspaceOAuthArgument(host, profile) + if err != nil { + return err + } + auth, err := NewPersistentAuth(ctx, + WithOAuthArgument(arg), + WithTokenCache(tokenCache), + ) + if err != nil { + return err + } + defer auth.Close() + if err := auth.Challenge(); err != nil { + return err + } + token, err := auth.Token() + +Because the U2M flow requires user interaction, callers should provide a +persistent cache to avoid prompting the user on every invocation. Without +WithTokenCache, PersistentAuth uses an in-memory cache. See the cache package +for the cache contract. +*/ +package u2m diff --git a/libs/auth/u2m/endpoint_supplier.go b/libs/auth/u2m/endpoint_supplier.go new file mode 100644 index 00000000000..f5065aa1c61 --- /dev/null +++ b/libs/auth/u2m/endpoint_supplier.go @@ -0,0 +1,81 @@ +package u2m + +import ( + "context" + "errors" + "fmt" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/httpclient" +) + +var ErrOAuthNotSupported = errors.New("databricks OAuth is not supported for this host") + +// OAuthEndpointSupplier provides the http functionality needed for interacting with the +// Databricks OAuth APIs. +type OAuthEndpointSupplier interface { + // GetWorkspaceOAuthEndpoints returns the OAuth2 endpoints for the workspace. + GetWorkspaceOAuthEndpoints(ctx context.Context, workspaceHost string) (*OAuthAuthorizationServer, error) + + // GetAccountOAuthEndpoints returns the OAuth2 endpoints for the account. + GetAccountOAuthEndpoints(ctx context.Context, accountHost, accountId string) (*OAuthAuthorizationServer, error) + + // GetUnifiedOAuthEndpoints returns the OAuth2 endpoints for the unified host. + GetUnifiedOAuthEndpoints(ctx context.Context, host, accountId string) (*OAuthAuthorizationServer, error) + + // GetEndpointsFromURL fetches OAuth2 endpoints directly from an authorization + // server metadata URL, bypassing host-type detection. + GetEndpointsFromURL(ctx context.Context, rawURL string) (*OAuthAuthorizationServer, error) +} + +// BasicOAuthEndpointSupplier is an implementation of the OAuthEndpointSupplier interface. +type BasicOAuthEndpointSupplier struct { + // Client is the ApiClient to use for making HTTP requests. + Client *httpclient.ApiClient +} + +// GetEndpointsFromURL fetches OAuth2 endpoints directly from an authorization +// server metadata URL, bypassing host-type detection. +func (c *BasicOAuthEndpointSupplier) GetEndpointsFromURL(ctx context.Context, rawURL string) (*OAuthAuthorizationServer, error) { + var oauthEndpoints OAuthAuthorizationServer + if err := c.Client.Do(ctx, "GET", rawURL, httpclient.WithResponseUnmarshal(&oauthEndpoints)); err != nil { + if errors.Is(err, apierr.ErrNotFound) { + return nil, ErrOAuthNotSupported + } + return nil, fmt.Errorf("failed to get OAuth endpoints: %w", err) + } + return &oauthEndpoints, nil +} + +// GetWorkspaceOAuthEndpoints returns the OAuth endpoints for the given workspace. +func (c *BasicOAuthEndpointSupplier) GetWorkspaceOAuthEndpoints(ctx context.Context, workspaceHost string) (*OAuthAuthorizationServer, error) { + oidc := workspaceHost + "/oidc/.well-known/oauth-authorization-server" + return c.GetEndpointsFromURL(ctx, oidc) +} + +// GetAccountOAuthEndpoints returns the OAuth2 endpoints for the account. The +// account-level OAuth endpoints are fixed based on the account ID and host. +func (c *BasicOAuthEndpointSupplier) GetAccountOAuthEndpoints(ctx context.Context, accountHost, accountId string) (*OAuthAuthorizationServer, error) { + return &OAuthAuthorizationServer{ + AuthorizationEndpoint: fmt.Sprintf("%s/oidc/accounts/%s/v1/authorize", accountHost, accountId), + TokenEndpoint: fmt.Sprintf("%s/oidc/accounts/%s/v1/token", accountHost, accountId), + }, nil +} + +// GetUnifiedOAuthEndpoints returns the OAuth2 endpoints for the unified host +func (c *BasicOAuthEndpointSupplier) GetUnifiedOAuthEndpoints(ctx context.Context, host, accountId string) (*OAuthAuthorizationServer, error) { + oidc := fmt.Sprintf("%s/oidc/accounts/%s/.well-known/oauth-authorization-server", host, accountId) + return c.GetEndpointsFromURL(ctx, oidc) +} + +// OAuthAuthorizationServer contains the OAuth endpoints for a Databricks account +// or workspace. +type OAuthAuthorizationServer struct { + // AuthorizationEndpoint is the URL to redirect users to for authorization. + // It typically ends with /v1/authroize. + AuthorizationEndpoint string `json:"authorization_endpoint"` + + // TokenEndpoint is the URL to exchange an authorization code for an access token. + // It typically ends with /v1/token. + TokenEndpoint string `json:"token_endpoint"` +} diff --git a/libs/auth/u2m/endpoint_supplier_test.go b/libs/auth/u2m/endpoint_supplier_test.go new file mode 100644 index 00000000000..0d995cec2f7 --- /dev/null +++ b/libs/auth/u2m/endpoint_supplier_test.go @@ -0,0 +1,115 @@ +package u2m + +import ( + "context" + "errors" + "testing" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/common" + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/stretchr/testify/assert" +) + +func TestBasicOAuthClient_GetAccountOAuthEndpoints(t *testing.T) { + c := &BasicOAuthEndpointSupplier{} + s, err := c.GetAccountOAuthEndpoints(t.Context(), "https://abc", "xyz") + assert.NoError(t, err) + assert.Equal(t, "https://abc/oidc/accounts/xyz/v1/authorize", s.AuthorizationEndpoint) + assert.Equal(t, "https://abc/oidc/accounts/xyz/v1/token", s.TokenEndpoint) +} + +func TestGetWorkspaceOAuthEndpoints(t *testing.T) { + p := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Status: 200, + Response: map[string]string{ + "authorization_endpoint": "a", + "token_endpoint": "b", + }, + }, + }, + }) + c := &BasicOAuthEndpointSupplier{Client: p} + endpoints, err := c.GetWorkspaceOAuthEndpoints(t.Context(), "https://abc") + assert.NoError(t, err) + assert.Equal(t, "a", endpoints.AuthorizationEndpoint) + assert.Equal(t, "b", endpoints.TokenEndpoint) +} + +func TestGetEndpointsFromURL(t *testing.T) { + p := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.MappingTransport{ + "GET /oidc/v1/authorize-server": { + Status: 200, + Response: map[string]string{ + "authorization_endpoint": "https://abc/oidc/v1/authorize", + "token_endpoint": "https://abc/oidc/v1/token", + }, + }, + }, + }) + c := &BasicOAuthEndpointSupplier{Client: p} + endpoints, err := c.GetEndpointsFromURL(t.Context(), "https://abc/oidc/v1/authorize-server") + if err != nil { + t.Fatal(err) + } + if endpoints.AuthorizationEndpoint != "https://abc/oidc/v1/authorize" { + t.Errorf("unexpected AuthorizationEndpoint: %q", endpoints.AuthorizationEndpoint) + } + if endpoints.TokenEndpoint != "https://abc/oidc/v1/token" { + t.Errorf("unexpected TokenEndpoint: %q", endpoints.TokenEndpoint) + } +} + +func TestGetEndpointsFromURL_NotFound(t *testing.T) { + p := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.MappingTransport{ + "GET /oidc/v1/authorize-server": { + Status: 404, + }, + }, + // Mirror the error mapping done by Config.refreshTokenErrorMapper: map + // HTTP 404 to apierr.ErrNotFound so that errors.Is(err, apierr.ErrNotFound) + // returns true, matching the production code path. + ErrorMapper: func(ctx context.Context, resp common.ResponseWrapper) error { + err := httpclient.DefaultErrorMapper(ctx, resp) + if err == nil { + return nil + } + if httpErr, ok := errors.AsType[*httpclient.HttpError](err); ok { + if sdkErr, ok := apierr.ByStatusCode(httpErr.StatusCode); ok { + return sdkErr + } + } + return err + }, + }) + c := &BasicOAuthEndpointSupplier{Client: p} + _, err := c.GetEndpointsFromURL(t.Context(), "https://abc/oidc/v1/authorize-server") + if !errors.Is(err, ErrOAuthNotSupported) { + t.Errorf("expected ErrOAuthNotSupported, got %v", err) + } +} + +func TestGetUnifiedOAuthEndpoints(t *testing.T) { + p := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.MappingTransport{ + "GET /oidc/accounts/xyz/.well-known/oauth-authorization-server": { + Status: 200, + Response: map[string]string{ + "authorization_endpoint": "https://abc/oidc/accounts/xyz/v1/authorize", + "token_endpoint": "https://abc/oidc/accounts/xyz/v1/token", + }, + }, + }, + }) + c := &BasicOAuthEndpointSupplier{Client: p} + endpoints, err := c.GetUnifiedOAuthEndpoints(t.Context(), "https://abc", "xyz") + + assert.NoError(t, err) + assert.Equal(t, "https://abc/oidc/accounts/xyz/v1/authorize", endpoints.AuthorizationEndpoint) + assert.Equal(t, "https://abc/oidc/accounts/xyz/v1/token", endpoints.TokenEndpoint) +} diff --git a/libs/auth/u2m/error.go b/libs/auth/u2m/error.go new file mode 100644 index 00000000000..51ba9036551 --- /dev/null +++ b/libs/auth/u2m/error.go @@ -0,0 +1,14 @@ +package u2m + +import "errors" + +// ErrMissingRefreshToken is returned when a token refresh is requested but the +// cached OAuth token does not include a refresh token. +var ErrMissingRefreshToken = errors.New("cached token has no refresh token") + +// InvalidRefreshTokenError is returned from PersistentAuth's Token() and +// ForceRefreshToken() methods when a token refresh is attempted and the cached +// refresh token is invalid. +type InvalidRefreshTokenError struct { + error +} diff --git a/libs/auth/u2m/oauth_argument.go b/libs/auth/u2m/oauth_argument.go new file mode 100644 index 00000000000..84788312453 --- /dev/null +++ b/libs/auth/u2m/oauth_argument.go @@ -0,0 +1,22 @@ +package u2m + +// OAuthArgument is an interface that provides the necessary information to +// authenticate with PersistentAuth. Implementations of this interface must +// implement either the WorkspaceOAuthArgument or AccountOAuthArgument +// interface. +type OAuthArgument interface { + // GetCacheKey returns a unique key for the OAuthArgument. This key is used + // to store and retrieve the token from the token cache. + GetCacheKey() string +} + +// HostCacheKeyProvider is an interface for OAuthArgument implementations that +// can return a host-based cache key regardless of whether a profile is set. +// +// PersistentAuth itself no longer uses this key; it is exported so that +// external token cache implementations (for example, the CLI's file-based +// cache) can type-assert on it to mirror tokens under the host key for +// cross-SDK compatibility with older SDKs that only know host keys. +type HostCacheKeyProvider interface { + GetHostCacheKey() string +} diff --git a/libs/auth/u2m/page.tmpl b/libs/auth/u2m/page.tmpl new file mode 100644 index 00000000000..1540222dbec --- /dev/null +++ b/libs/auth/u2m/page.tmpl @@ -0,0 +1,104 @@ + + + + + {{if .Error }}{{ .Error | title }}{{ else }}Success{{end}} + + + + + + + +
+
+ + +
{{ .Error | title }}
+
{{ .ErrorDescription }}
+ +
Authenticated
+ {{- if .Host }} +
Go to {{.Host}}
+ {{- end}} + +
+ You can close this tab. Or go to documentation +
+
+
+ + diff --git a/libs/auth/u2m/persistent_auth.go b/libs/auth/u2m/persistent_auth.go new file mode 100644 index 00000000000..8a6756ee342 --- /dev/null +++ b/libs/auth/u2m/persistent_auth.go @@ -0,0 +1,661 @@ +package u2m + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "strings" + "time" + + cache "github.com/databricks/cli/libs/auth/u2m/cache" + "github.com/databricks/cli/libs/browser" + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/logger" + "golang.org/x/oauth2" + "golang.org/x/oauth2/authhandler" +) + +const ( + // appClientID is the public OAuth client ID assigned to the Databricks CLI. + appClientID = "databricks-cli" + + // defaultPort is the default port for the OAuth2 callback server. If the + // port is already in use, the next port is tried (8021, 8022, etc.). + defaultPort = 8020 + + // maxPortFallback is the maximum port to try when using the fallback + // mechanism. + maxPortFallback = 8040 + + // listenerTimeout is the maximum duration spent trying to acquire a + // listener (including port selection). + listenerTimeout = 45 * time.Second + + // tokenRefreshBuffer is the duration before token expiry at which the + // token is proactively refreshed. This prevents callers from receiving + // near-expired tokens that may expire before the next request. + tokenRefreshBuffer = 5 * time.Minute + + // Cache update recovery checks immediately and then waits for 25, 50, 100, + // and 200 milliseconds. This gives a concurrent cache writer time to finish + // while bounding the delay for persistent storage failures to 375 milliseconds. + cacheUpdateRecoveryAttempts = 5 + + cacheUpdateRecoveryInitialDelay = 25 * time.Millisecond + cacheUpdateRecoveryDelayFactor = 2 + + // Concurrent refreshes can finish at slightly different times, so their + // expiration times need not be identical. + cacheUpdateRecoveryExpiryDelta = time.Minute +) + +var ( + // Internal errors used for testing. + errListenerTimeout = errors.New("failed to listen on any port: timeout") + errNoPortAvailable = errors.New("no port available to listen on") +) + +// PersistentAuth is an OAuth manager that handles the U2M OAuth flow. Tokens +// are stored in and looked up from the provided cache. Tokens include the +// refresh token. On load, if the access token is expired or close to expiry, +// it is refreshed using the refresh token. +// +// The PersistentAuth is safe for concurrent use. The token cache is locked +// during token retrieval, refresh and storage. +type PersistentAuth struct { + // cache is the token cache to store and lookup tokens. + cache cache.TokenCache + + // client is the HTTP client to use for OAuth2 requests. + client *http.Client + + // endpointSupplier is the HTTP endpointSupplier to use for OAuth2 requests. + endpointSupplier OAuthEndpointSupplier + + // oAuthArgument defines the workspace or account to authenticate to and the + // cache key for the token. + oAuthArgument OAuthArgument + + // browser is the function to open a URL in the default browser. + browser func(url string) error + + // ln is the listener for the OAuth2 callback server. + ln net.Listener + + // ctx is the context to use for underlying operations. This is needed in + // order to implement the oauth2.TokenSource interface. + ctx context.Context + + // redirectAddr is the redirect address for OAuth2 callbacks. The value is + // set to localhost:PORT by startListener which will dynamically assign a + // random port. If a value is already provided, it will be used instead + // (e.g. for testing). + redirectAddr string + + // Optional port to use for the OAuth2 callback server. If set to 0, the + // default port with fallback is used. This means that setting a port will + // disable the fallback mechanism. + port int + + // netListen is an optional function to listen on a TCP address. If not set, + // it will use net.Listen by default. This is useful for testing. + netListen func(network, address string) (net.Listener, error) + + // scopes is the list of OAuth scopes to request. + scopes []string + + // disableOfflineAccess controls whether offline_access scope is requested. + // When true, offline_access will NOT be automatically added to scopes, + // meaning the token will not include a refresh token. + disableOfflineAccess bool + + // discoveryMode enables the login.databricks.com discovery flow. + // When true, Challenge() uses the discovery token source instead of + // the standard authhandler flow. + discoveryMode bool + + // discoveryHost overrides the default login.databricks.com host used by + // the discovery flow. Empty means the production host. + discoveryHost string + + // discoveryAccountTarget, when true, instructs the discovery flow to set + // the top-level `target=ACCOUNT` query parameter on the authorize URL so + // login.databricks.com lands the user on the account selector instead of + // the workspace selector. Use for account-only logins. + discoveryAccountTarget bool +} + +type PersistentAuthOption func(*PersistentAuth) + +// WithTokenCache sets the token cache for the PersistentAuth. +func WithTokenCache(c cache.TokenCache) PersistentAuthOption { + return func(a *PersistentAuth) { + a.cache = c + } +} + +// WithHttpClient sets the HTTP client for the PersistentAuth. +func WithHttpClient(c *http.Client) PersistentAuthOption { + return func(a *PersistentAuth) { + a.client = c + } +} + +// WithOAuthEndpointSupplier sets the OAuth endpoint supplier for the +// PersistentAuth. +func WithOAuthEndpointSupplier(c OAuthEndpointSupplier) PersistentAuthOption { + return func(a *PersistentAuth) { + a.endpointSupplier = c + } +} + +// WithOAuthArgument sets the OAuthArgument for the PersistentAuth. +func WithOAuthArgument(arg OAuthArgument) PersistentAuthOption { + return func(a *PersistentAuth) { + a.oAuthArgument = arg + } +} + +// WithBrowser sets the browser function for the PersistentAuth. +func WithBrowser(b func(url string) error) PersistentAuthOption { + return func(a *PersistentAuth) { + a.browser = b + } +} + +// WithPort sets the port for the PersistentAuth. +// +//deadcode:allow retained for parity with the deprecated SDK compatibility API +func WithPort(port int) PersistentAuthOption { + return func(a *PersistentAuth) { + a.port = port + } +} + +// WithScopes sets the OAuth scopes for the PersistentAuth. +func WithScopes(scopes []string) PersistentAuthOption { + return func(a *PersistentAuth) { + a.scopes = scopes + } +} + +// WithDisableOfflineAccess controls whether offline_access scope is requested. +func WithDisableOfflineAccess(disable bool) PersistentAuthOption { + return func(a *PersistentAuth) { + a.disableOfflineAccess = disable + } +} + +// WithDiscoveryLogin enables the login.databricks.com discovery flow. +// When enabled, Challenge() routes through login.databricks.com instead +// of directly to a workspace OIDC endpoint. +// +// This option is only valid with [DiscoveryOAuthArgument], which is a +// bootstrap-only argument type for discovery login. Once the workspace host +// has been discovered, callers should construct the usual host-based +// OAuthArgument for future PersistentAuth instances. +func WithDiscoveryLogin() PersistentAuthOption { + return func(a *PersistentAuth) { + a.discoveryMode = true + } +} + +// WithDiscoveryHost overrides the default https://login.databricks.com host +// used by the discovery login flow. Intended for testing and development +// against non-production environments; has no effect unless WithDiscoveryLogin +// is also set. If host has no scheme, https:// is prepended. Trailing slashes +// are trimmed. +func WithDiscoveryHost(host string) PersistentAuthOption { + return func(a *PersistentAuth) { + if host != "" && !strings.Contains(host, "://") { + host = "https://" + host + } + a.discoveryHost = host + } +} + +// WithDiscoveryAccountTarget sets the top-level `target=ACCOUNT` query +// parameter on the discovery authorize URL so login.databricks.com lands the +// user on the account selector instead of the workspace selector. Use for +// account-only logins where workspace selection would be a wasted step. +// +// Has no effect unless WithDiscoveryLogin is also set. +func WithDiscoveryAccountTarget() PersistentAuthOption { + return func(a *PersistentAuth) { + a.discoveryAccountTarget = true + } +} + +// NewPersistentAuth creates a new PersistentAuth with the provided options. +func NewPersistentAuth(ctx context.Context, opts ...PersistentAuthOption) (*PersistentAuth, error) { + p := &PersistentAuth{} + for _, opt := range opts { + opt(p) + } + // By default, PersistentAuth uses the default ApiClient to make HTTP + // requests. Furthermore, if the endpointSupplier is not provided, it uses + // this same client to fetch the OAuth endpoints. If the HTTP client is + // provided but the endpointSupplier is not, we construct a default + // ApiClient for use with BasicOAuthClient. + apiClient := httpclient.NewApiClient(httpclient.ClientConfig{}) + if p.client == nil { + p.client = &http.Client{ + Transport: apiClient, + // 30 seconds matches the default timeout of the ApiClient + Timeout: 30 * time.Second, + } + } + if p.endpointSupplier == nil { + p.endpointSupplier = &BasicOAuthEndpointSupplier{ + Client: apiClient, + } + } + if p.cache == nil { + p.cache = cache.NewInMemoryTokenCache() + } + if err := p.validateArg(); err != nil { + return nil, err + } + if p.browser == nil { + p.browser = func(url string) error { return browser.Open(ctx, url) } + } + p.ctx = ctx + return p, nil +} + +// loadToken loads the cached OAuth2 token for the configured OAuthArgument +// using GetCacheKey(). The returned token may be expired; callers are +// responsible for deciding whether and how to refresh it. +func (a *PersistentAuth) loadToken() (*oauth2.Token, error) { + t, err := a.cache.Lookup(a.oAuthArgument.GetCacheKey()) + if err != nil { + return nil, fmt.Errorf("cache: %w", err) + } + return t, nil +} + +// Token loads the OAuth2 token for the given OAuthArgument from the cache. If +// the token is expired or close to expiry, it is refreshed using the refresh +// token. When a proactive refresh (token still valid but near expiry) fails, +// the existing token is returned so the caller is not blocked. +func (a *PersistentAuth) Token() (*oauth2.Token, error) { + t, err := a.loadToken() + if err != nil { + return nil, err + } + if needsRefresh(t) { + if refreshedToken, err := a.refresh(t); err == nil { + t = refreshedToken + } else if !t.Valid() { + return nil, fmt.Errorf("token refresh: %w", err) + } else { + logger.Debugf(a.ctx, "proactive token refresh failed, returning existing token: %v", err) + } + } + t.RefreshToken = "" + return t, nil +} + +// ForceRefreshToken loads the OAuth2 token from the cache by GetCacheKey(), +// refreshes it unconditionally, and stores the refreshed token back under the +// same key. Unlike Token(), if the refresh fails the error is always returned +// -- the caller explicitly asked for a fresh token, so silently falling back +// to a stale one would be incorrect. +func (a *PersistentAuth) ForceRefreshToken() (*oauth2.Token, error) { + t, err := a.loadToken() + if err != nil { + return nil, err + } + t, err = a.refresh(t) + if err != nil { + return nil, fmt.Errorf("forced token refresh: %w", err) + } + t.RefreshToken = "" + return t, nil +} + +// needsRefresh returns true when the token should be refreshed, either because +// it is no longer valid or because it will expire within the +// tokenRefreshBuffer window. +func needsRefresh(t *oauth2.Token) bool { + if !t.Valid() { + return true + } + return !t.Expiry.IsZero() && time.Until(t.Expiry) < tokenRefreshBuffer +} + +// isFreshReplacement reports whether cached changed from old, is valid, and +// does not expire significantly sooner than candidate. +func isFreshReplacement(old, candidate, cached *oauth2.Token) bool { + if cached.AccessToken == old.AccessToken || !cached.Valid() { + return false + } + if cached.Expiry.IsZero() { + return true + } + if candidate.Expiry.IsZero() { + return false + } + return !cached.Expiry.Before(candidate.Expiry.Add(-cacheUpdateRecoveryExpiryDelta)) +} + +// recoverCacheUpdate checks whether a concurrent cache update completed. +// Retrying reads instead of writes avoids recreating the write race. +func (a *PersistentAuth) recoverCacheUpdate(old, candidate *oauth2.Token) *oauth2.Token { + delay := cacheUpdateRecoveryInitialDelay + for attempt := range cacheUpdateRecoveryAttempts { + if attempt > 0 { + timer := time.NewTimer(delay) + select { + case <-a.ctx.Done(): + timer.Stop() + return nil + case <-timer.C: + } + delay *= cacheUpdateRecoveryDelayFactor + } + + cached, err := a.cache.Lookup(a.oAuthArgument.GetCacheKey()) + if err == nil && isFreshReplacement(old, candidate, cached) { + return cached + } + } + return nil +} + +// refresh refreshes the token for the given OAuthArgument, storing the new +// token in the cache. +// +// This read-refresh-write sequence is not coordinated across processes. +// Because the CLI is stateless, two separate CLI invocations can load the same +// cached refresh token, both attempt a refresh, and race to update the cache. +// This should be fixed in a follow-up by adding cross-process coordination +// around refresh and cache writes. +func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) { + // Fail fast with ErrMissingRefreshToken instead of letting the oauth2 + // library attempt to refresh and return a misleading error (e.g. "token + // expired" when the real problem is that the cached token is not + // refresh-capable). + if oldToken.RefreshToken == "" { + return nil, ErrMissingRefreshToken + } + cfg, err := a.oauth2Config() + if err != nil { + return nil, err + } + ctx := a.setOAuthContext(a.ctx) + // Force the oauth2 library to refresh by ensuring the token appears + // expired. PersistentAuth owns the refresh decision (including the + // proactive buffer), so the oauth2 library should always perform the + // refresh when asked. + expired := *oldToken + expired.Expiry = time.Now().Add(-time.Minute) + t, err := cfg.TokenSource(ctx, &expired).Token() + if err != nil { + // The default RoundTripper of our httpclient.ApiClient returns an error + // if the response status code is not 2xx. This isn't compliant with the + // RoundTripper interface, so this error isn't handled by the oauth2 + // library. We need to handle it here. + if internalHttpError, ok := errors.AsType[*httpclient.HttpError](err); ok { + // error fields + // https://datatracker.ietf.org/doc/html/rfc6749#section-5.2 + var errResponse struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description"` + } + if unmarshalErr := json.Unmarshal([]byte(internalHttpError.Message), &errResponse); unmarshalErr != nil { + return nil, fmt.Errorf("unmarshal: %w", unmarshalErr) + } + // Invalid refresh tokens get their own error type so they can be + // better presented to users. + if errResponse.ErrorDescription == "Refresh token is invalid" { + return nil, &InvalidRefreshTokenError{err} + } + return nil, fmt.Errorf("%s (error code: %s)", errResponse.ErrorDescription, errResponse.Error) + } + + // Handle responses from well-behaved *http.Client implementations. + if httpErr, ok := errors.AsType[*oauth2.RetrieveError](err); ok { + // Invalid refresh tokens get their own error type so they can be + // better presented to users. + if httpErr.ErrorDescription == "Refresh token is invalid" { + return nil, &InvalidRefreshTokenError{err} + } + return nil, fmt.Errorf("%s (error code: %s)", httpErr.ErrorDescription, httpErr.ErrorCode) + } + return nil, err + } + err = a.cache.Store(a.oAuthArgument.GetCacheKey(), t) + if err != nil { + if cached := a.recoverCacheUpdate(oldToken, t); cached != nil { + return cached, nil + } + return nil, fmt.Errorf("cache update: %w", err) + } + return t, nil +} + +// Challenge initiates the OAuth2 login flow for the given OAuthArgument. The +// OAuth2 flow is started by opening the browser to the OAuth2 authorization +// URL. The user is redirected to the callback server on appRedirectAddr. The +// callback server listens for the redirect from the identity provider and +// exchanges the authorization code for an access token. +func (a *PersistentAuth) Challenge() error { + if a.discoveryMode { + return a.discoveryChallenge() + } + err := a.startListener(a.ctx) + if err != nil { + return fmt.Errorf("starting listener: %w", err) + } + // The listener will be closed by the callback server automatically, but if + // the callback server is not created, we need to close the listener manually. + defer a.Close() + + cfg, err := a.oauth2Config() + if err != nil { + return fmt.Errorf("fetching oauth config: %w", err) + } + cb, err := a.newCallbackServer() + if err != nil { + return fmt.Errorf("callback server: %w", err) + } + defer cb.Close() + + state, pkce, err := a.stateAndPKCE() + if err != nil { + return fmt.Errorf("state and pkce: %w", err) + } + // make OAuth2 library use our client + ctx := a.setOAuthContext(a.ctx) + ts := authhandler.TokenSourceWithPKCE(ctx, cfg, state, cb.Handler, pkce) + t, err := ts.Token() + if err != nil { + return fmt.Errorf("authorize: %w", err) + } + err = a.cache.Store(a.oAuthArgument.GetCacheKey(), t) + if err != nil { + return fmt.Errorf("store: %w", err) + } + return nil +} + +// discoveryChallenge handles the login.databricks.com discovery flow. +// The listener must be started before the discovery token source is invoked +// because the challenge needs the redirect address to build the authorize URL. +func (a *PersistentAuth) discoveryChallenge() error { + err := a.startListener(a.ctx) + if err != nil { + return fmt.Errorf("starting listener: %w", err) + } + defer a.Close() + ds := &discoveryTokenSource{pa: a, host: a.discoveryHost} + if a.discoveryAccountTarget { + ds.target = discoveryTargetAccount + } + return ds.challenge() +} + +// startListener starts a listener on appRedirectAddr, retrying if the address +// is already in use. +func (a *PersistentAuth) startListener(ctx context.Context) error { + if a.port != 0 { // if port is set, use it + return a.startListenerWithPort(a.port) + } + return a.startListenerWithFallback(ctx) +} + +// startListenerWithFallback starts a listener that will try to find a free +// port to listen on starting from the default port and incrementing by 1 until +// a free port is found. +func (a *PersistentAuth) startListenerWithFallback(ctx context.Context) error { + startTime := time.Now() + for port := defaultPort; port <= maxPortFallback; port++ { + if time.Since(startTime) > listenerTimeout { + return errListenerTimeout + } + if err := a.startListenerWithPort(port); err != nil { + logger.Debugf(ctx, "failed to listen on %d: %v, retrying", port, err) + continue + } + logger.Debugf(ctx, "OAuth callback server listening on %s", a.redirectAddr) + return nil + } + return errNoPortAvailable +} + +func (a *PersistentAuth) startListenerWithPort(port int) error { + addr := fmt.Sprintf("localhost:%d", port) + listener, err := a.listen("tcp", addr) + if err != nil { + return fmt.Errorf("failed to listen on %s: %w", addr, err) + } + a.ln = listener + a.redirectAddr = addr + return nil +} + +func (a *PersistentAuth) listen(network, addr string) (net.Listener, error) { + if a.netListen != nil { + return a.netListen(network, addr) + } + return net.Listen(network, addr) +} + +func (a *PersistentAuth) Close() error { + if a.ln == nil { + return nil + } + return a.ln.Close() +} + +// validateArg ensures that the OAuthArgument is either a WorkspaceOAuthArgument, +// AccountOAuthArgument, UnifiedOAuthArgument, or a bootstrap-only +// DiscoveryOAuthArgument paired with WithDiscoveryLogin. +func (a *PersistentAuth) validateArg() error { + if a.oAuthArgument == nil { + return errors.New("missing OAuthArgument") + } + _, isWorkspaceArg := a.oAuthArgument.(WorkspaceOAuthArgument) + _, isAccountArg := a.oAuthArgument.(AccountOAuthArgument) + _, isUnifiedArg := a.oAuthArgument.(UnifiedOAuthArgument) + _, isDiscoveryArg := a.oAuthArgument.(DiscoveryOAuthArgument) + if !isWorkspaceArg && !isAccountArg && !isUnifiedArg && !isDiscoveryArg { + return fmt.Errorf("unsupported OAuthArgument type: %T, must implement WorkspaceOAuthArgument, AccountOAuthArgument, UnifiedOAuthArgument, or DiscoveryOAuthArgument", a.oAuthArgument) + } + if isDiscoveryArg && !a.discoveryMode { + return fmt.Errorf("discovery OAuthArgument %T requires WithDiscoveryLogin; after discovery, construct a WorkspaceOAuthArgument with the discovered host", a.oAuthArgument) + } + if a.discoveryMode && !isDiscoveryArg { + return fmt.Errorf("discovery login requires DiscoveryOAuthArgument, got %T", a.oAuthArgument) + } + return nil +} + +// resolveScopes returns the effective OAuth scopes for this PersistentAuth. +// It defaults to "all-apis" for backwards compatibility, and prepends +// "offline_access" unless disabled. +func (a *PersistentAuth) resolveScopes() []string { + scopes := a.scopes + if len(scopes) == 0 { + scopes = []string{"all-apis"} + } + if !a.disableOfflineAccess { + scopes = append([]string{"offline_access"}, scopes...) + } + return scopes +} + +// oauth2Config returns the OAuth2 configuration for the given OAuthArgument. +func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) { + scopes := a.resolveScopes() + + var endpoints *OAuthAuthorizationServer + var err error + switch argg := a.oAuthArgument.(type) { + case WorkspaceOAuthArgument: + endpoints, err = a.endpointSupplier.GetWorkspaceOAuthEndpoints(a.ctx, argg.GetWorkspaceHost()) + case AccountOAuthArgument: + endpoints, err = a.endpointSupplier.GetAccountOAuthEndpoints( + a.ctx, argg.GetAccountHost(), argg.GetAccountId()) + case UnifiedOAuthArgument: + endpoints, err = a.endpointSupplier.GetUnifiedOAuthEndpoints(a.ctx, argg.GetHost(), argg.GetAccountId()) + case DiscoveryOAuthArgument: + return nil, fmt.Errorf("discovery OAuthArgument %T is only supported with WithDiscoveryLogin during Challenge; after discovery, construct a WorkspaceOAuthArgument with the discovered host", a.oAuthArgument) + default: + return nil, fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument, AccountOAuthArgument or UnifiedOAuthArgument interface", a.oAuthArgument) + } + if err != nil { + return nil, fmt.Errorf("fetching OAuth endpoints: %w", err) + } + return &oauth2.Config{ + ClientID: appClientID, + Endpoint: oauth2.Endpoint{ + AuthURL: endpoints.AuthorizationEndpoint, + TokenURL: endpoints.TokenEndpoint, + AuthStyle: oauth2.AuthStyleInParams, + }, + RedirectURL: "http://" + a.redirectAddr, + Scopes: scopes, + }, nil +} + +func (a *PersistentAuth) stateAndPKCE() (string, *authhandler.PKCEParams, error) { + verifier, err := a.randomString(64) + if err != nil { + return "", nil, fmt.Errorf("verifier: %w", err) + } + verifierSha256 := sha256.Sum256([]byte(verifier)) + challenge := base64.RawURLEncoding.EncodeToString(verifierSha256[:]) + state, err := a.randomString(16) + if err != nil { + return "", nil, fmt.Errorf("state: %w", err) + } + return state, &authhandler.PKCEParams{ + Challenge: challenge, + ChallengeMethod: "S256", + Verifier: verifier, + }, nil +} + +func (a *PersistentAuth) randomString(size int) (string, error) { + raw := make([]byte, size) + // ignore error as rand.Reader never returns an error + _, err := rand.Read(raw) + if err != nil { + return "", fmt.Errorf("rand.Read: %w", err) + } + return base64.RawURLEncoding.EncodeToString(raw), nil +} + +func (a *PersistentAuth) setOAuthContext(ctx context.Context) context.Context { + return context.WithValue(ctx, oauth2.HTTPClient, a.client) +} + +var _ oauth2.TokenSource = (*PersistentAuth)(nil) diff --git a/libs/auth/u2m/persistent_auth_test.go b/libs/auth/u2m/persistent_auth_test.go new file mode 100644 index 00000000000..31c95913869 --- /dev/null +++ b/libs/auth/u2m/persistent_auth_test.go @@ -0,0 +1,1473 @@ +package u2m + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/u2m/cache" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "golang.org/x/oauth2" +) + +type tokenCacheMock struct { + store func(key string, t *oauth2.Token) error + lookup func(key string) (*oauth2.Token, error) +} + +func (m *tokenCacheMock) Store(key string, t *oauth2.Token) error { + return m.store(key, t) +} + +func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) { + return m.lookup(key) +} + +func TestToken(t *testing.T) { + cache := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != "https://abc/oidc/accounts/xyz" { + t.Fatalf("lookup(): want key 'https://abc/oidc/accounts/xyz', got %s", key) + } + return &oauth2.Token{ + AccessToken: "bcd", + Expiry: time.Now().Add(1 * time.Hour), + }, nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://abc", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth(t.Context(), WithTokenCache(cache), WithOAuthArgument(arg)) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "bcd" { + t.Errorf("p.Token(): want access token 'bcd', got %s", tok.AccessToken) + } + if tok.RefreshToken != "" { + t.Errorf("p.Token(): want refresh token '', got %s", tok.RefreshToken) + } +} + +func TestToken_WithProfile(t *testing.T) { + profileKey := "my-profile" + cache := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != profileKey { + t.Fatalf("lookup(): want key %q, got %q", profileKey, key) + } + return &oauth2.Token{ + AccessToken: "profile-token", + Expiry: time.Now().Add(1 * time.Hour), + }, nil + }, + } + arg, err := NewProfileAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz", profileKey) + if err != nil { + t.Fatalf("NewProfileAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth(t.Context(), WithTokenCache(cache), WithOAuthArgument(arg)) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "profile-token" { + t.Errorf("p.Token(): want access token 'profile-token', got %s", tok.AccessToken) + } +} + +type MockOAuthEndpointSupplier struct{} + +func (m MockOAuthEndpointSupplier) GetAccountOAuthEndpoints(ctx context.Context, accountHost, accountId string) (*OAuthAuthorizationServer, error) { + return &OAuthAuthorizationServer{ + AuthorizationEndpoint: fmt.Sprintf("%s/oidc/accounts/%s/v1/authorize", accountHost, accountId), + TokenEndpoint: fmt.Sprintf("%s/oidc/accounts/%s/v1/token", accountHost, accountId), + }, nil +} + +func (m MockOAuthEndpointSupplier) GetWorkspaceOAuthEndpoints(ctx context.Context, workspaceHost string) (*OAuthAuthorizationServer, error) { + return &OAuthAuthorizationServer{ + AuthorizationEndpoint: workspaceHost + "/oidc/v1/authorize", + TokenEndpoint: workspaceHost + "/oidc/v1/token", + }, nil +} + +func (m MockOAuthEndpointSupplier) GetUnifiedOAuthEndpoints(ctx context.Context, host, accountId string) (*OAuthAuthorizationServer, error) { + return &OAuthAuthorizationServer{ + AuthorizationEndpoint: fmt.Sprintf("%s/oidc/accounts/%s/v1/authorize", host, accountId), + TokenEndpoint: fmt.Sprintf("%s/oidc/accounts/%s/v1/token", host, accountId), + }, nil +} + +func (m MockOAuthEndpointSupplier) GetEndpointsFromURL(_ context.Context, _ string) (*OAuthAuthorizationServer, error) { + return nil, ErrOAuthNotSupported +} + +func TestToken_RefreshesExpiredAccessToken(t *testing.T) { + ctx := t.Context() + expectedKey := "https://accounts.cloud.databricks.test/oidc/accounts/xyz" + cache := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != expectedKey { + t.Fatalf("lookup(): want key %s, got %s", expectedKey, key) + } + return &oauth2.Token{ + AccessToken: "expired", + RefreshToken: "cde", + Expiry: time.Now().Add(-1 * time.Minute), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + if key != expectedKey { + t.Fatalf("store(): want key %s, got %s", expectedKey, key) + } + if tok.RefreshToken != "def" { + t.Fatalf("store(): want refresh token 'def', got %s", tok.RefreshToken) + } + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + ctx, + WithTokenCache(cache), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=refreshed&refresh_token=def`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Errorf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "refreshed" { + t.Errorf("p.Token(): want access token 'refreshed', got %s", tok.AccessToken) + } + if tok.RefreshToken != "" { + t.Errorf("p.Token(): want refresh token '', got %s", tok.RefreshToken) + } +} + +func TestToken_RefreshesTokenExpiringSoon(t *testing.T) { + ctx := t.Context() + expectedKey := "https://accounts.cloud.databricks.test/oidc/accounts/xyz" + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != expectedKey { + t.Fatalf("lookup(): want key %s, got %s", expectedKey, key) + } + return &oauth2.Token{ + AccessToken: "expiring-soon", + RefreshToken: "cde", + Expiry: time.Now().Add(4 * time.Minute), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + ctx, + WithTokenCache(c), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=refreshed&refresh_token=def`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "refreshed" { + t.Errorf("p.Token(): want access token 'refreshed', got %s", tok.AccessToken) + } +} + +func TestToken_ReturnsStillValidTokenWhenProactiveRefreshFails(t *testing.T) { + ctx := t.Context() + expectedKey := "https://accounts.cloud.databricks.test/oidc/accounts/xyz" + transport := fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `{"error": "temporarily_unavailable", "error_description": "temporarily unavailable"}`, + Status: 503, + }, + } + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != expectedKey { + t.Fatalf("lookup(): want key %s, got %s", expectedKey, key) + } + return &oauth2.Token{ + AccessToken: "still-valid", + RefreshToken: "cde", + Expiry: time.Now().Add(4 * time.Minute), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call - refresh should not succeed") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + ctx, + WithTokenCache(c), + WithHttpClient(&http.Client{Transport: transport}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "still-valid" { + t.Errorf("p.Token(): want access token 'still-valid', got %s", tok.AccessToken) + } + if tok.RefreshToken != "" { + t.Errorf("p.Token(): want refresh token '', got %s", tok.RefreshToken) + } + if transport[0].Method != "" { + t.Errorf("refresh(): want proactive refresh attempt, but request was not sent") + } +} + +func TestToken_DoesNotRefreshTokenNotExpiringSoon(t *testing.T) { + transport := fixtures.SliceTransport{} + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "still-valid", + Expiry: time.Now().Add(6 * time.Minute), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — token should not be refreshed") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{Transport: transport}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "still-valid" { + t.Errorf("p.Token(): want access token 'still-valid', got %s", tok.AccessToken) + } +} + +func TestToken_ZeroExpiryDoesNotTriggerRefresh(t *testing.T) { + transport := fixtures.SliceTransport{} + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "no-expiry", + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — zero-expiry token should not be refreshed") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{Transport: transport}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("p.Token(): want no error, got %v", err) + } + if tok.AccessToken != "no-expiry" { + t.Errorf("p.Token(): want access token 'no-expiry', got %s", tok.AccessToken) + } +} + +func TestToken_ReturnsError(t *testing.T) { + ctx := t.Context() + cache := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != "https://accounts.cloud.databricks.test/oidc/accounts/xyz" { + t.Fatalf("lookup(): want key 'https://accounts.cloud.databricks.test/oidc/accounts/xyz', got %s", key) + } + return &oauth2.Token{ + AccessToken: "expired", + RefreshToken: "cde", + Expiry: time.Now().Add(-1 * time.Minute), + }, nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + ctx, + WithTokenCache(cache), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `{"error": "invalid_grant", "error_description": "Invalid Client"}`, + Status: 401, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Errorf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + tok, err := p.Token() + + if tok != nil { + t.Errorf("p.Token(): want nil, got %v", tok) + } + if !strings.Contains(err.Error(), "Invalid Client (error code: invalid_grant)") { + t.Errorf("p.Token(): want error containing 'Invalid Client (error code: invalid_grant)', got %v", err) + } +} + +func TestToken_ReturnsInvalidRefreshTokenError(t *testing.T) { + ctx := t.Context() + cache := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + if key != "https://accounts.cloud.databricks.test/oidc/accounts/xyz" { + t.Fatalf("lookup(): want key 'https://accounts.cloud.databricks.test/oidc/accounts/xyz', got %s", key) + } + return &oauth2.Token{ + AccessToken: "expired", + RefreshToken: "cde", + Expiry: time.Now().Add(-1 * time.Minute), + }, nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + p, err := NewPersistentAuth( + ctx, + WithTokenCache(cache), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `{"error": "invalid_grant", "error_description": "Refresh token is invalid"}`, + Status: 401, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Errorf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + tok, err := p.Token() + if tok != nil { + t.Fatalf("p.Token(): want nil, got %v", tok) + } + if _, ok := errors.AsType[*InvalidRefreshTokenError](err); !ok { + t.Fatalf("p.Token(): want error of type InvalidRefreshTokenError, got %v", err) + } +} + +func TestToken_ReturnsMissingRefreshTokenErrorWhenExpired(t *testing.T) { + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "expired", + Expiry: time.Now().Add(-1 * time.Minute), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — Token() should fail before writing when refresh token is missing") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{}}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err == nil { + t.Fatal("Token(): want error when refresh token is missing, got nil") + } + if tok != nil { + t.Errorf("Token(): want nil token on error, got %v", tok) + } + if !errors.Is(err, ErrMissingRefreshToken) { + t.Fatalf("Token(): want ErrMissingRefreshToken, got %v", err) + } + want := "token refresh: cached token has no refresh token" + if got := err.Error(); got != want { + t.Errorf("Token(): want error %q, got %q", want, got) + } +} + +func TestToken_ReturnsExistingTokenWhenNearExpiryAndNoRefreshToken(t *testing.T) { + // Token is still valid but within the proactive refresh window (5 min). + // Token() attempts refresh, refresh() returns ErrMissingRefreshToken. + // Token() falls back to the existing token since it is still valid. + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "near-expiry", + Expiry: time.Now().Add(3 * time.Minute), // within tokenRefreshBuffer + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — Token() should return existing token without migrating") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{}}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.Token() + if err != nil { + t.Fatalf("Token(): want no error when falling back to valid token, got %v", err) + } + if tok == nil { + t.Fatal("Token(): want token, got nil") + } + if tok.AccessToken != "near-expiry" { + t.Errorf("Token(): want access token 'near-expiry', got %s", tok.AccessToken) + } +} + +func TestForceRefreshToken_RefreshesValidToken(t *testing.T) { + refreshCalled := false + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "still-valid", + RefreshToken: "refresh-me", + Expiry: time.Now().Add(1 * time.Hour), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + refreshCalled = true + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=force-refreshed&refresh_token=new-refresh`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err != nil { + t.Fatalf("ForceRefreshToken(): want no error, got %v", err) + } + if tok.AccessToken != "force-refreshed" { + t.Errorf("ForceRefreshToken(): want access token 'force-refreshed', got %s", tok.AccessToken) + } + if tok.RefreshToken != "" { + t.Errorf("ForceRefreshToken(): want refresh token redacted, got %s", tok.RefreshToken) + } + if !refreshCalled { + t.Error("ForceRefreshToken(): want refresh to be called for valid token, but it was not") + } +} + +func TestForceRefreshToken_RecoversConcurrentCacheUpdate(t *testing.T) { + now := time.Date(2100, time.January, 1, 0, 0, 0, 0, time.UTC) + old := &oauth2.Token{ + AccessToken: "old-access", + RefreshToken: "old-refresh", + Expiry: now.Add(time.Hour), + } + winner := &oauth2.Token{ + AccessToken: "winner-access", + RefreshToken: "winner-refresh", + Expiry: now.Add(time.Hour - 30*time.Second), + } + lookupCalls := 0 + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + lookupCalls++ + if lookupCalls <= 2 { + return old, nil + } + return winner, nil + }, + store: func(key string, tok *oauth2.Token) error { + if tok.AccessToken != "candidate-access" { + t.Fatalf("store(): want candidate access token, got %q", tok.AccessToken) + } + return errors.New("concurrent cache update") + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=candidate-access&refresh_token=candidate-refresh&expires_in=3600`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err != nil { + t.Fatalf("ForceRefreshToken(): want no error, got %v", err) + } + if tok.AccessToken != winner.AccessToken { + t.Errorf("ForceRefreshToken(): want winner access token %q, got %q", winner.AccessToken, tok.AccessToken) + } + if tok.RefreshToken != "" { + t.Errorf("ForceRefreshToken(): want refresh token redacted, got %q", tok.RefreshToken) + } + if lookupCalls != 3 { + t.Errorf("Lookup(): want 3 calls, got %d", lookupCalls) + } +} + +func TestIsFreshReplacement(t *testing.T) { + now := time.Date(2100, time.January, 1, 0, 0, 0, 0, time.UTC) + old := &oauth2.Token{AccessToken: "old", Expiry: now.Add(time.Hour)} + candidate := &oauth2.Token{AccessToken: "candidate", Expiry: now.Add(time.Hour)} + + tests := []struct { + name string + candidate *oauth2.Token + cached *oauth2.Token + want bool + }{ + { + name: "same token", + candidate: candidate, + cached: &oauth2.Token{AccessToken: "old", Expiry: now.Add(time.Hour)}, + want: false, + }, + { + name: "expired replacement", + candidate: candidate, + cached: &oauth2.Token{AccessToken: "winner", Expiry: time.Date(2000, time.January, 1, 0, 0, 0, 0, time.UTC)}, + want: false, + }, + { + name: "replacement expires too soon", + candidate: candidate, + cached: &oauth2.Token{AccessToken: "winner", Expiry: candidate.Expiry.Add(-time.Minute - time.Second)}, + want: false, + }, + { + name: "replacement expiry within tolerance", + candidate: candidate, + cached: &oauth2.Token{AccessToken: "winner", Expiry: candidate.Expiry.Add(-time.Minute)}, + want: true, + }, + { + name: "replacement expires later", + candidate: candidate, + cached: &oauth2.Token{AccessToken: "winner", Expiry: candidate.Expiry.Add(time.Minute)}, + want: true, + }, + { + name: "expiring candidate and non-expiring replacement", + candidate: candidate, + cached: &oauth2.Token{AccessToken: "winner"}, + want: true, + }, + { + name: "non-expiring candidate and replacement inside refresh buffer", + candidate: &oauth2.Token{AccessToken: "candidate"}, + cached: &oauth2.Token{AccessToken: "winner", Expiry: now.Add(time.Minute)}, + want: false, + }, + { + name: "non-expiring candidate and replacement", + candidate: &oauth2.Token{AccessToken: "candidate"}, + cached: &oauth2.Token{AccessToken: "winner"}, + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isFreshReplacement(old, tt.candidate, tt.cached); got != tt.want { + t.Errorf("isFreshReplacement(): want %t, got %t", tt.want, got) + } + }) + } +} + +func TestForceRefreshToken_WithInMemoryCachePreservesCachedRefreshToken(t *testing.T) { + tokenCache := cache.NewInMemoryTokenCache() + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + if err := tokenCache.Store(arg.GetCacheKey(), &oauth2.Token{ + AccessToken: "still-valid", + RefreshToken: "refresh-me", + Expiry: time.Now().Add(1 * time.Hour), + }); err != nil { + t.Fatalf("Store(): %v", err) + } + + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(tokenCache), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=force-refreshed&refresh_token=new-refresh`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err != nil { + t.Fatalf("ForceRefreshToken(): want no error, got %v", err) + } + if tok.RefreshToken != "" { + t.Fatalf("ForceRefreshToken(): want refresh token redacted, got %q", tok.RefreshToken) + } + + cached, err := tokenCache.Lookup(arg.GetCacheKey()) + if err != nil { + t.Fatalf("Lookup(): want cached token, got %v", err) + } + if cached.RefreshToken != "new-refresh" { + t.Fatalf("Lookup(): want cached refresh token %q, got %q", "new-refresh", cached.RefreshToken) + } +} + +func TestForceRefreshToken_FailsWithoutRefreshToken(t *testing.T) { + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "still-valid", + Expiry: time.Now().Add(1 * time.Hour), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — ForceRefreshToken() should fail before writing when refresh token is missing") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{}}), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err == nil { + t.Fatal("ForceRefreshToken(): want error when refresh token is missing, got nil") + } + if tok != nil { + t.Errorf("ForceRefreshToken(): want nil token on error, got %v", tok) + } + if !errors.Is(err, ErrMissingRefreshToken) { + t.Fatalf("ForceRefreshToken(): want ErrMissingRefreshToken, got %v", err) + } + want := "forced token refresh: cached token has no refresh token" + if got := err.Error(); got != want { + t.Errorf("ForceRefreshToken(): want error %q, got %q", want, got) + } +} + +func TestForceRefreshToken_FailsOnRefreshError(t *testing.T) { + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "still-valid", + RefreshToken: "bad-refresh", + Expiry: time.Now().Add(1 * time.Hour), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — refresh should not succeed") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `{"error": "temporarily_unavailable", "error_description": "temporarily unavailable"}`, + Status: 503, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if err == nil { + t.Fatal("ForceRefreshToken(): want error when refresh fails, got nil") + } + if tok != nil { + t.Errorf("ForceRefreshToken(): want nil token on error, got %v", tok) + } + want := "forced token refresh: temporarily unavailable (error code: temporarily_unavailable)" + if got := err.Error(); got != want { + t.Errorf("ForceRefreshToken(): want error %q, got %q", want, got) + } +} + +func TestForceRefreshToken_InvalidRefreshTokenError(t *testing.T) { + c := &tokenCacheMock{ + lookup: func(key string) (*oauth2.Token, error) { + return &oauth2.Token{ + AccessToken: "expired", + RefreshToken: "bad-refresh", + Expiry: time.Now().Add(-1 * time.Minute), + }, nil + }, + store: func(key string, tok *oauth2.Token) error { + t.Fatalf("store(): unexpected call — refresh should not succeed") + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): %v", err) + } + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(c), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `{"error": "invalid_grant", "error_description": "Refresh token is invalid"}`, + Status: 401, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + tok, err := p.ForceRefreshToken() + if tok != nil { + t.Fatalf("ForceRefreshToken(): want nil token, got %v", tok) + } + if _, ok := errors.AsType[*InvalidRefreshTokenError](err); !ok { + t.Fatalf("ForceRefreshToken(): want InvalidRefreshTokenError, got %v", err) + } +} + +func TestChallenge(t *testing.T) { + ctx := t.Context() + + browserOpened := make(chan string) + browser := func(redirect string) error { + u, err := url.ParseRequestURI(redirect) + if err != nil { + return err + } + if u.Path != "/oidc/accounts/xyz/v1/authorize" { + t.Fatalf("browser(): want path '/oidc/accounts/xyz/v1/authorize', got %s", u.Path) + } + // for now we're ignoring asserting the fields of the redirect + query := u.Query() + browserOpened <- query.Get("state") + return nil + } + cache := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + if key != "https://accounts.cloud.databricks.test/oidc/accounts/xyz" { + t.Fatalf("store(): want key 'https://accounts.cloud.databricks.test/oidc/accounts/xyz', got %s", key) + } + if tok.AccessToken != "__THAT__" { + t.Fatalf("store(): want access token '__THAT__', got %s", tok.AccessToken) + } + if tok.RefreshToken != "__SOMETHING__" { + t.Fatalf("store(): want refresh token '__SOMETHING__', got %s", tok.RefreshToken) + } + return nil + }, + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + + p, err := NewPersistentAuth( + ctx, + WithTokenCache(cache), + WithBrowser(browser), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: "/oidc/accounts/xyz/v1/token", + Response: `access_token=__THAT__&refresh_token=__SOMETHING__`, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + ) + if err != nil { + t.Errorf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + errc := make(chan error) + go func() { + err := p.Challenge() + errc <- err + close(errc) + }() + + state := <-browserOpened + resp, err := http.Get("http://localhost:8020?code=__THIS__&state=" + state) + if err != nil { + t.Fatalf("http.Get(): want no error, got %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("http.Get(): want status code 200, got %d", resp.StatusCode) + } + + err = <-errc + if err != nil { + t.Fatalf("p.Challenge(): want no error, got %v", err) + } +} + +func TestChallenge_ReturnsErrorOnFailure(t *testing.T) { + ctx := t.Context() + browserOpened := make(chan string) + browser := func(redirect string) error { + u, err := url.ParseRequestURI(redirect) + if err != nil { + return err + } + if u.Path != "/oidc/accounts/xyz/v1/authorize" { + t.Fatalf("browser(): want path '/oidc/accounts/xyz/v1/authorize', got %s", u.Path) + } + // for now we're ignoring asserting the fields of the redirect + query := u.Query() + browserOpened <- query.Get("state") + return nil + } + arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") + if err != nil { + t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err) + } + + p, err := NewPersistentAuth(ctx, WithBrowser(browser), WithOAuthArgument(arg)) + if err != nil { + t.Errorf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + errc := make(chan error) + go func() { + err := p.Challenge() + errc <- err + close(errc) + }() + + <-browserOpened + resp, err := http.Get("http://localhost:8020?error=access_denied&error_description=Policy%20evaluation%20failed%20for%20this%20request") + if err != nil { + t.Fatalf("http.Get(): want no error, got %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("http.Get(): want status code 400, got %d", resp.StatusCode) + } + + err = <-errc + if err == nil { + t.Fatalf("p.Challenge(): want error, got nil") + } + if !strings.Contains(err.Error(), "authorize: access_denied: Policy evaluation failed for this request") { + t.Fatalf("p.Challenge(): want error containing 'authorize: access_denied: Policy evaluation failed for this request', got %v", err) + } +} + +func TestPersistentAuth_startListener_startFrom8020(t *testing.T) { + pa := &PersistentAuth{} + pa.netListen = func(_, address string) (net.Listener, error) { + return nil, nil + } + + gotErr := pa.startListener(t.Context()) + + if gotErr != nil { + t.Fatalf("pa.startListener(): want no error, got %v", gotErr) + } + if pa.redirectAddr != "localhost:8020" { + t.Errorf("pa.redirectAddr should be localhost:8020, got %s", pa.redirectAddr) + } +} + +func TestPersistentAuth_startListener_incrementalFallBack(t *testing.T) { + pa := &PersistentAuth{} + pa.netListen = func(_, address string) (net.Listener, error) { + if address == "localhost:8020" { + return nil, errors.New("address already in use") + } + if address == "localhost:8021" { + return nil, errors.New("address already in use") + } + return nil, nil + } + + gotErr := pa.startListener(t.Context()) + + if gotErr != nil { + t.Fatalf("pa.startListener(): want no error, got %v", gotErr) + } + if pa.redirectAddr != "localhost:8022" { + t.Errorf("pa.redirectAddr should be localhost:8022, got %s", pa.redirectAddr) + } +} + +func TestPersistentAuth_startListener_noAvailablePort(t *testing.T) { + pa := &PersistentAuth{} + pa.netListen = func(_, address string) (net.Listener, error) { + return nil, errors.New("address already in use") + } + + gotErr := pa.startListener(t.Context()) + + if !errors.Is(gotErr, errNoPortAvailable) { + t.Fatalf("pa.startListener(): want error %v, got %v", errNoPortAvailable, gotErr) + } +} + +func TestPersistentAuth_startListener_maxPortFallbackIncluded(t *testing.T) { + maxAddress := fmt.Sprintf("localhost:%d", maxPortFallback) + pa := &PersistentAuth{} + pa.netListen = func(_, address string) (net.Listener, error) { + if address == maxAddress { + return nil, nil + } + return nil, errors.New("address already in use") + } + + gotErr := pa.startListener(t.Context()) + + if gotErr != nil { + t.Fatalf("pa.startListener(): want no error, got %v", gotErr) + } + if pa.redirectAddr != maxAddress { + t.Errorf("pa.redirectAddr should be %s, got %s", maxAddress, pa.redirectAddr) + } +} + +func TestPersistentAuth_startListener_explicitPort(t *testing.T) { + explicitPort := 1337 + pa := &PersistentAuth{port: explicitPort} + pa.netListen = func(_, address string) (net.Listener, error) { + return nil, nil + } + + gotErr := pa.startListener(t.Context()) + + if gotErr != nil { + t.Fatalf("pa.startListener(): want no error, got %v", gotErr) + } + if pa.redirectAddr != "localhost:1337" { + t.Errorf("pa.redirectAddr should be localhost:1337, got %s", pa.redirectAddr) + } +} + +func TestPersistentAuth_startListener_explicitPortNoFallBack(t *testing.T) { + testError := errors.New("test error") + explicitPort := 1337 + pa := &PersistentAuth{port: explicitPort} + pa.netListen = func(_, address string) (net.Listener, error) { + if address == "localhost:1337" { + return nil, testError + } + return nil, nil + } + + gotErr := pa.startListener(t.Context()) + + if !errors.Is(gotErr, testError) { + t.Fatalf("pa.startListener(): want error %v, got %v", testError, gotErr) + } +} + +// TestU2M_ScopesAndOfflineAccess verifies that OAuth scopes are correctly configured +// and sent during the authorization flow, and that the disableOfflineAccess flag +// correctly controls whether offline_access is added to the scope. +func TestU2M_ScopesAndOfflineAccess(t *testing.T) { + const ( + testWorkspaceHost = "https://workspace.cloud.databricks.test" + testTokenEndpoint = "/oidc/v1/token" + testCallbackURL = "http://localhost:8020" + ) + + tests := []struct { + name string + scopes []string + disableOffline bool + want string + }{ + { + name: "nil scopes uses default with offline_access", + scopes: nil, + disableOffline: false, + want: "offline_access all-apis", + }, + { + name: "empty scopes uses default with offline_access", + scopes: []string{}, + disableOffline: false, + want: "offline_access all-apis", + }, + { + name: "single scope with offline_access", + scopes: []string{"dashboards"}, + disableOffline: false, + want: "offline_access dashboards", + }, + { + name: "multiple scopes with offline_access", + scopes: []string{"files", "jobs", "mlflow:read"}, + disableOffline: false, + want: "offline_access files jobs mlflow:read", + }, + { + name: "disable offline_access", + scopes: []string{"files", "jobs"}, + disableOffline: true, + want: "files jobs", + }, + { + name: "nil scopes with disable offline_access", + scopes: nil, + disableOffline: true, + want: "all-apis", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := t.Context() + + var scopeReceived, stateReceived string + browserCalled := make(chan struct{}) + defer close(browserCalled) + browser := func(redirect string) error { + u, err := url.ParseRequestURI(redirect) + if err != nil { + return err + } + query := u.Query() + scopeReceived = query.Get("scope") + stateReceived = query.Get("state") + browserCalled <- struct{}{} + return nil + } + + cache := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + return nil + }, + } + + arg, err := NewBasicWorkspaceOAuthArgument(testWorkspaceHost) + if err != nil { + t.Fatalf("NewBasicWorkspaceOAuthArgument(): want no error, got %v", err) + } + + var tokenResponse string + if tt.disableOffline { + tokenResponse = `access_token=token` + } else { + tokenResponse = `access_token=token&refresh_token=refresh` + } + + opts := []PersistentAuthOption{ + WithTokenCache(cache), + WithBrowser(browser), + WithHttpClient(&http.Client{ + Transport: fixtures.SliceTransport{ + { + Method: "POST", + Resource: testTokenEndpoint, + Response: tokenResponse, + ResponseHeaders: map[string][]string{ + "Content-Type": {"application/x-www-form-urlencoded"}, + }, + }, + }, + }), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + WithDisableOfflineAccess(tt.disableOffline), + WithScopes(tt.scopes), + } + + p, err := NewPersistentAuth(ctx, opts...) + if err != nil { + t.Fatalf("NewPersistentAuth(): want no error, got %v", err) + } + defer p.Close() + + errc := make(chan error) + defer close(errc) + go func() { + err := p.Challenge() + errc <- err + }() + + select { + case <-browserCalled: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for browser to be called") + } + + if scopeReceived != tt.want { + t.Errorf("scope: want %q, got %q", tt.want, scopeReceived) + } + + resp, err := http.Get(fmt.Sprintf("%s?code=__CODE__&state=%s", testCallbackURL, stateReceived)) + if err != nil { + t.Fatalf("http.Get(): want no error, got %v", err) + } + defer resp.Body.Close() + + select { + case err = <-errc: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for Challenge() to complete") + } + if err != nil { + t.Fatalf("p.Challenge(): want no error, got %v", err) + } + }) + } +} + +func TestChallenge_Discovery(t *testing.T) { + // Mock token server that responds to POST /oidc/v1/token. + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("token server: want POST, got %s", r.Method) + } + if r.URL.Path != "/oidc/v1/token" { + t.Errorf("token server: want path /oidc/v1/token, got %s", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"access_token":"discovery-access-token","refresh_token":"discovery-refresh-token","token_type":"Bearer","expires_in":3600}`) + })) + defer tokenServer.Close() + + // The issuer is the mock server URL + /oidc. + issuer := tokenServer.URL + "/oidc" + + browserOpened := make(chan string, 1) + browserMock := func(u string) error { + browserOpened <- u + return nil + } + + storedTokens := map[string]*oauth2.Token{} + cacheMock := &tokenCacheMock{ + store: func(key string, tok *oauth2.Token) error { + storedTokens[key] = tok + return nil + }, + } + + arg, err := NewBasicDiscoveryOAuthArgument("discovery-profile") + if err != nil { + t.Fatalf("NewBasicDiscoveryOAuthArgument(): %v", err) + } + + p, err := NewPersistentAuth( + t.Context(), + WithTokenCache(cacheMock), + WithBrowser(browserMock), + WithHttpClient(tokenServer.Client()), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + WithOAuthArgument(arg), + WithDiscoveryLogin(), + ) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + defer p.Close() + + errc := make(chan error, 1) + go func() { + errc <- p.Challenge() + }() + + // Wait for browser to be called and extract state from the authorize URL. + var state string + select { + case authURL := <-browserOpened: + u, err := url.Parse(authURL) + if err != nil { + t.Fatalf("parsing auth URL: %v", err) + } + destURL := u.Query().Get("destination_url") + dest, err := url.Parse(destURL) + if err != nil { + t.Fatalf("parsing destination_url: %v", err) + } + state = dest.Query().Get("state") + if state == "" { + t.Fatal("state is empty in authorize URL") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for browser to be called") + } + + // Fire the callback with code, state, and iss. + callbackURL := fmt.Sprintf("http://%s?code=test-code&state=%s&iss=%s", + p.redirectAddr, url.QueryEscape(state), url.QueryEscape(issuer)) + resp, err := http.Get(callbackURL) + if err != nil { + t.Fatalf("callback GET: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("callback: want status 200, got %d", resp.StatusCode) + } + + // Wait for Challenge to complete. + select { + case err := <-errc: + if err != nil { + t.Fatalf("Challenge(): %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for Challenge to complete") + } + + // Verify discovered host was set on the argument. + expectedHost, err := DeriveHostFromIssuer(issuer) + if err != nil { + t.Fatalf("DeriveHostFromIssuer(%q): %v", issuer, err) + } + if arg.GetDiscoveredHost() != expectedHost { + t.Errorf("discovered host = %q, want %q", arg.GetDiscoveredHost(), expectedHost) + } + if len(storedTokens) != 1 { + t.Fatalf("store count: want 1 key (profile), got %d", len(storedTokens)) + } + storedToken := storedTokens["discovery-profile"] + if storedToken == nil { + t.Fatalf("stored token for profile key is nil") + } + if storedToken.AccessToken != "discovery-access-token" { + t.Errorf("access token = %q, want %q", storedToken.AccessToken, "discovery-access-token") + } + if storedToken.RefreshToken != "discovery-refresh-token" { + t.Errorf("refresh token = %q, want %q", storedToken.RefreshToken, "discovery-refresh-token") + } +} diff --git a/libs/auth/u2m/unified_oauth_argument.go b/libs/auth/u2m/unified_oauth_argument.go new file mode 100644 index 00000000000..733127b1421 --- /dev/null +++ b/libs/auth/u2m/unified_oauth_argument.go @@ -0,0 +1,74 @@ +package u2m + +import ( + "fmt" +) + +// UnifiedOAuthArgument is an interface that provides the necessary information +// to authenticate using OAuth to a host that supports both account and workspace APIs. +type UnifiedOAuthArgument interface { + OAuthArgument + + // GetHost returns the host to authenticate to. + GetHost() string + + // GetAccountId returns the account ID of the account to authenticate to. + GetAccountId() string +} + +// BasicUnifiedOAuthArgument is a basic implementation of the UnifiedOAuthArgument +// interface that links each account with exactly one OAuth token. +type BasicUnifiedOAuthArgument struct { + host string + accountID string + + // profile is the optional profile name. When set, GetCacheKey() returns + // the profile name instead of the host-based key. + profile string +} + +var ( + _ UnifiedOAuthArgument = BasicUnifiedOAuthArgument{} + _ HostCacheKeyProvider = BasicUnifiedOAuthArgument{} +) + +// NewBasicUnifiedOAuthArgument creates a new BasicUnifiedOAuthArgument. +func NewBasicUnifiedOAuthArgument(accountsHost, accountID string) (BasicUnifiedOAuthArgument, error) { + return NewProfileUnifiedOAuthArgument(accountsHost, accountID, "") +} + +// NewProfileUnifiedOAuthArgument creates a new BasicUnifiedOAuthArgument with a +// profile name. When a profile is set, GetCacheKey() returns the profile name +// instead of the host-based key. +func NewProfileUnifiedOAuthArgument(accountsHost, accountID, profile string) (BasicUnifiedOAuthArgument, error) { + if err := validateHost(accountsHost); err != nil { + return BasicUnifiedOAuthArgument{}, err + } + return BasicUnifiedOAuthArgument{host: accountsHost, accountID: accountID, profile: profile}, nil +} + +// GetHost returns the host to authenticate to. +func (a BasicUnifiedOAuthArgument) GetHost() string { + return a.host +} + +// GetAccountId returns the account ID of the account to authenticate to. +func (a BasicUnifiedOAuthArgument) GetAccountId() string { + return a.accountID +} + +// GetCacheKey returns a unique key for caching the OAuth token for the account. +// If a profile is set, the profile name is returned as the cache key. +// Otherwise, the key is in the format "/oidc/accounts/". +func (a BasicUnifiedOAuthArgument) GetCacheKey() string { + if a.profile != "" { + return a.profile + } + return a.GetHostCacheKey() +} + +// GetHostCacheKey returns the host-based cache key regardless of whether a +// profile is set. The key is in the format "/oidc/accounts/". +func (a BasicUnifiedOAuthArgument) GetHostCacheKey() string { + return fmt.Sprintf("%s/oidc/accounts/%s", a.host, a.accountID) +} diff --git a/libs/auth/u2m/unified_oauth_argument_test.go b/libs/auth/u2m/unified_oauth_argument_test.go new file mode 100644 index 00000000000..3ddef30c043 --- /dev/null +++ b/libs/auth/u2m/unified_oauth_argument_test.go @@ -0,0 +1,106 @@ +package u2m + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNewBasicUnifiedOAuthArgument_ValidatesHost(t *testing.T) { + tests := []struct { + name string + host string + accountID string + wantErr string + }{ + { + name: "invalid http protocol", + host: "http://insecure.com", + accountID: "account-123", + wantErr: "host must start with 'https://': http://insecure.com", + }, + { + name: "trailing slash", + host: "https://unified.databricks.test/", + accountID: "account-123", + wantErr: "host must not have a trailing slash: https://unified.databricks.test/", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := NewBasicUnifiedOAuthArgument(tt.host, tt.accountID) + assert.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestNewProfileUnifiedOAuthArgument_ValidatesHost(t *testing.T) { + _, err := NewProfileUnifiedOAuthArgument("http://insecure.com", "account-123", "my-profile") + assert.Error(t, err) + assert.Contains(t, err.Error(), "host must start with 'https://'") +} + +func TestBasicUnifiedOAuthArgument_ProfileCacheKeys(t *testing.T) { + tests := []struct { + name string + host string + accountID string + profile string + wantKey string + wantHostKey string + }{ + { + name: "without profile returns host-based key", + host: "https://unified.databricks.test", + accountID: "account-123", + wantKey: "https://unified.databricks.test/oidc/accounts/account-123", + wantHostKey: "https://unified.databricks.test/oidc/accounts/account-123", + }, + { + name: "with profile returns profile name", + host: "https://unified.databricks.test", + accountID: "account-123", + profile: "my-profile", + wantKey: "my-profile", + wantHostKey: "https://unified.databricks.test/oidc/accounts/account-123", + }, + { + name: "empty profile returns host-based key", + host: "https://unified.databricks.test", + accountID: "account-123", + profile: "", + wantKey: "https://unified.databricks.test/oidc/accounts/account-123", + wantHostKey: "https://unified.databricks.test/oidc/accounts/account-123", + }, + { + name: "different host and account", + host: "https://other-unified.databricks.test", + accountID: "account-456", + wantKey: "https://other-unified.databricks.test/oidc/accounts/account-456", + wantHostKey: "https://other-unified.databricks.test/oidc/accounts/account-456", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var arg BasicUnifiedOAuthArgument + var err error + if tt.profile != "" { + arg, err = NewProfileUnifiedOAuthArgument(tt.host, tt.accountID, tt.profile) + } else { + arg, err = NewBasicUnifiedOAuthArgument(tt.host, tt.accountID) + } + assert.NoError(t, err) + assert.Equal(t, tt.wantKey, arg.GetCacheKey()) + assert.Equal(t, tt.wantHostKey, arg.GetHostCacheKey()) + }) + } +} + +func TestBasicUnifiedOAuthArgument_ImplementsHostCacheKeyProvider(t *testing.T) { + arg, err := NewBasicUnifiedOAuthArgument("https://unified.databricks.test", "account-123") + assert.NoError(t, err) + var _ HostCacheKeyProvider = arg +} diff --git a/libs/auth/u2m/workspace_oauth_argument.go b/libs/auth/u2m/workspace_oauth_argument.go new file mode 100644 index 00000000000..66daba94ebe --- /dev/null +++ b/libs/auth/u2m/workspace_oauth_argument.go @@ -0,0 +1,87 @@ +package u2m + +import ( + "fmt" + "strings" +) + +// WorkspaceOAuthArgument is an interface that provides the necessary information +// to authenticate using OAuth to a specific workspace. +type WorkspaceOAuthArgument interface { + OAuthArgument + + // GetWorkspaceHost returns the host of the workspace to authenticate to. + GetWorkspaceHost() string +} + +// BasicWorkspaceOAuthArgument is a basic implementation of the WorkspaceOAuthArgument +// interface that links each host with exactly one OAuth token. +type BasicWorkspaceOAuthArgument struct { + // host is the host of the workspace to authenticate to. This must start + // with "https://" and must not have a trailing slash. + host string + + // profile is the optional profile name. When set, GetCacheKey() returns + // the profile name instead of the host-based key. + profile string +} + +func validateHost(host string) error { + // Allow http for localhost. This is necessary for local end to end testing + // of the `databricks auth login` command using a test server on localhost. + if strings.HasPrefix(host, "http://127.0.0.1") { + return nil + } + if !strings.HasPrefix(host, "https://") { + return fmt.Errorf("host must start with 'https://': %s", host) + } + if strings.HasSuffix(host, "/") { + return fmt.Errorf("host must not have a trailing slash: %s", host) + } + return nil +} + +// NewBasicWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument. +func NewBasicWorkspaceOAuthArgument(host string) (BasicWorkspaceOAuthArgument, error) { + return NewProfileWorkspaceOAuthArgument(host, "") +} + +// NewProfileWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument +// with a profile name. When a profile is set, GetCacheKey() returns the profile +// name instead of the host-based key. +func NewProfileWorkspaceOAuthArgument(host, profile string) (BasicWorkspaceOAuthArgument, error) { + if err := validateHost(host); err != nil { + return BasicWorkspaceOAuthArgument{}, err + } + return BasicWorkspaceOAuthArgument{host: host, profile: profile}, nil +} + +// GetWorkspaceHost returns the host of the workspace to authenticate to. +func (a BasicWorkspaceOAuthArgument) GetWorkspaceHost() string { + return a.host +} + +// GetCacheKey returns a unique key for caching the OAuth token for the workspace. +// If a profile is set, the profile name is returned as the cache key. +// Otherwise, the key is in the format "". +func (a BasicWorkspaceOAuthArgument) GetCacheKey() string { + if a.profile != "" { + return a.profile + } + return a.GetHostCacheKey() +} + +// GetHostCacheKey returns the host-based cache key regardless of whether a +// profile is set. The key is in the format "". +func (a BasicWorkspaceOAuthArgument) GetHostCacheKey() string { + host := strings.TrimSuffix(a.host, "/") + if !strings.HasPrefix(host, "http") { + host = "https://" + host + } + return host +} + +var ( + _ WorkspaceOAuthArgument = BasicWorkspaceOAuthArgument{} + _ HostCacheKeyProvider = BasicWorkspaceOAuthArgument{} +) diff --git a/libs/auth/u2m/workspace_oauth_argument_test.go b/libs/auth/u2m/workspace_oauth_argument_test.go new file mode 100644 index 00000000000..652f06206e1 --- /dev/null +++ b/libs/auth/u2m/workspace_oauth_argument_test.go @@ -0,0 +1,119 @@ +package u2m + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBasicWorkspaceOAuthArgument_GetCacheKey(t *testing.T) { + tests := []struct { + name string + host string + profile string + wantKey string + wantHostKey string + }{ + { + name: "without profile returns host-based key", + host: "https://myworkspace.cloud.databricks.test", + wantKey: "https://myworkspace.cloud.databricks.test", + wantHostKey: "https://myworkspace.cloud.databricks.test", + }, + { + name: "with profile returns profile name", + host: "https://myworkspace.cloud.databricks.test", + profile: "my-profile", + wantKey: "my-profile", + wantHostKey: "https://myworkspace.cloud.databricks.test", + }, + { + name: "empty profile returns host-based key", + host: "https://myworkspace.cloud.databricks.test", + profile: "", + wantKey: "https://myworkspace.cloud.databricks.test", + wantHostKey: "https://myworkspace.cloud.databricks.test", + }, + { + name: "localhost without profile", + host: "http://127.0.0.1:5656", + wantKey: "http://127.0.0.1:5656", + wantHostKey: "http://127.0.0.1:5656", + }, + { + name: "localhost with profile", + host: "http://127.0.0.1:5656", + profile: "local-profile", + wantKey: "local-profile", + wantHostKey: "http://127.0.0.1:5656", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var arg BasicWorkspaceOAuthArgument + var err error + if tt.profile != "" { + arg, err = NewProfileWorkspaceOAuthArgument(tt.host, tt.profile) + } else { + arg, err = NewBasicWorkspaceOAuthArgument(tt.host) + } + assert.NoError(t, err) + assert.Equal(t, tt.wantKey, arg.GetCacheKey()) + assert.Equal(t, tt.wantHostKey, arg.GetHostCacheKey()) + }) + } +} + +func TestNewProfileWorkspaceOAuthArgument_ValidatesHost(t *testing.T) { + _, err := NewProfileWorkspaceOAuthArgument("http://some-host.com", "my-profile") + assert.EqualError(t, err, "host must start with 'https://': http://some-host.com") + + _, err = NewProfileWorkspaceOAuthArgument("https://some-host.com/", "my-profile") + assert.EqualError(t, err, "host must not have a trailing slash: https://some-host.com/") +} + +func TestBasicWorkspaceOAuthArgument_ImplementsHostCacheKeyProvider(t *testing.T) { + arg, err := NewBasicWorkspaceOAuthArgument("https://myworkspace.cloud.databricks.test") + assert.NoError(t, err) + var _ HostCacheKeyProvider = arg +} + +func TestValidateHost(t *testing.T) { + tests := []struct { + host string + want string + }{ + // Valid hosts + {"https://some-host.com", ""}, + {"http://127.0.0.1", ""}, + {"http://127.0.0.1:5656", ""}, + + // Invalid hosts + {"http://some-host.com", "host must start with 'https://': http://some-host.com"}, + {"https://some-host.com/", "host must not have a trailing slash: https://some-host.com/"}, + } + + for _, test := range tests { + err := validateHost(test.host) + if test.want == "" { + assert.NoError(t, err) + } else { + assert.EqualError(t, err, test.want) + } + + _, err = NewBasicWorkspaceOAuthArgument(test.host) + if test.want == "" { + assert.NoError(t, err) + } else { + assert.EqualError(t, err, test.want) + } + + _, err = NewBasicAccountOAuthArgument(test.host, "123") + if test.want == "" { + assert.NoError(t, err) + } else { + assert.EqualError(t, err, test.want) + } + } +}