diff --git a/cmd/auth/profiles.go b/cmd/auth/profiles.go index 51c397a9ea9..86260b59b58 100644 --- a/cmd/auth/profiles.go +++ b/cmd/auth/profiles.go @@ -57,40 +57,53 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV return } - configType := auth.ResolveConfigType(cfg) - if configType != cfg.ConfigType() { - log.Debugf(ctx, "Profile %q: overrode config type from %s to %s (SPOG host)", c.Name, cfg.ConfigType(), configType) - } + // Validate by probing the API surfaces this profile has a signal for. + // Each signal — host shape or field presence — enables its corresponding + // probe, and the OR of the probe results is the verdict. + + // Host signals. + // isAccountHost: classic accounts.* host. + // isSPOGHost: unified host with account-scoped OIDC discovery. + // isWorkspaceHost: classic workspace host (neither of the above). + isAccountHost := auth.IsClassicAccountHost(cfg.CanonicalHostName()) + isSPOGHost := auth.IsSPOGHost(cfg) + isWorkspaceHost := auth.IsClassicWorkspaceHost(cfg) + + // Field signals. + // hasRealWorkspaceID: workspace_id is set to a real value. + // workspace_id is "" when not present in the profile, "none" when the user picked Skip during SPOG login. + hasRealWorkspaceID := cfg.WorkspaceID != "" && cfg.WorkspaceID != auth.WorkspaceIDNone + + // account_id isn't a probe trigger on its own: it gets back-filled by + // discovery on workspace profiles and can linger from a prior login on + // the same profile name, so its presence on a non-account host doesn't + // imply the user has account access. + tryAccount := isAccountHost || isSPOGHost + tryWorkspace := isWorkspaceHost || hasRealWorkspaceID - switch configType { - case config.AccountConfig: + var accountOK, workspaceOK bool + if tryAccount { a, err := databricks.NewAccountClient((*databricks.Config)(cfg)) - if err != nil { - return - } - _, err = a.Workspaces.List(ctx) - c.Host = cfg.Host - c.AuthType = cfg.AuthType - if err != nil { - return + if err == nil { + if _, err := a.Workspaces.List(ctx); err == nil { + accountOK = true + } } - c.Valid = true - case config.WorkspaceConfig: + } + if tryWorkspace { w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg)) - if err != nil { - return - } - _, err = w.CurrentUser.Me(ctx) - c.Host = cfg.Host - c.AuthType = cfg.AuthType - if err != nil { - return + if err == nil { + if _, err := w.CurrentUser.Me(ctx); err == nil { + workspaceOK = true + } } - c.Valid = true - case config.InvalidConfig: - // Invalid configuration, skip validation - return } + + // Capture Host/AuthType after the probes run: SDK Authenticate() sets + // cfg.AuthType lazily based on the credentials it actually exercised. + c.Host = cfg.Host + c.AuthType = cfg.AuthType + c.Valid = accountOK || workspaceOK } func newProfilesCommand() *cobra.Command { diff --git a/cmd/auth/profiles_test.go b/cmd/auth/profiles_test.go index 59803e210cf..7adb2c735f1 100644 --- a/cmd/auth/profiles_test.go +++ b/cmd/auth/profiles_test.go @@ -80,10 +80,11 @@ func TestProfilesDefaultMarker(t *testing.T) { } // newSPOGServer creates a mock SPOG server that returns account-scoped OIDC. -// It serves both validation endpoints since SPOG workspace profiles (with a -// real workspace_id) need CurrentUser.Me, while account profiles need -// Workspaces.List. The workspace-only newWorkspaceServer omits the account -// endpoint to prove routing correctness for non-SPOG hosts. +// The workspace endpoint deliberately returns 500 to mirror real SPOG hosts +// where account-audience OAuth tokens can't load workspace OAuth config. +// auth profiles probes both surfaces and accepts either success, so the test +// passes when Workspaces.List succeeds even though CurrentUser.Me fails — +// and a regression that drops the account probe surfaces as Valid=false. func newSPOGServer(t *testing.T, accountID string) *httptest.Server { t.Helper() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -97,8 +98,7 @@ func newSPOGServer(t *testing.T, accountID string) *httptest.Server { case "/api/2.0/accounts/" + accountID + "/workspaces": _ = json.NewEncoder(w).Encode([]map[string]any{}) case "/api/2.0/preview/scim/v2/Me": - // SPOG workspace profiles also need CurrentUser.Me to succeed. - _ = json.NewEncoder(w).Encode(map[string]any{"userName": "test-user"}) + http.Error(w, "SPOG profiles must validate via Workspaces.List, not CurrentUser.Me", http.StatusInternalServerError) default: w.WriteHeader(http.StatusNotFound) } @@ -108,8 +108,9 @@ func newSPOGServer(t *testing.T, accountID string) *httptest.Server { } // newWorkspaceServer creates a mock workspace server that returns workspace-scoped -// OIDC and only serves the workspace validation endpoint. The account validation -// endpoint returns 404 to prove the workspace path was taken. +// OIDC and a workspace_id in discovery (mirroring real workspace hosts since +// PR #4809). It serves CurrentUser.Me; the account endpoint returns 404 so a +// workspace probe is the only path that produces Valid=true. func newWorkspaceServer(t *testing.T, accountID string) *httptest.Server { t.Helper() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -118,6 +119,7 @@ func newWorkspaceServer(t *testing.T, accountID string) *httptest.Server { case "/.well-known/databricks-config": _ = json.NewEncoder(w).Encode(map[string]any{ "account_id": accountID, + "workspace_id": "ws-from-discovery", "oidc_endpoint": r.Host + "/oidc", }) case "/api/2.0/preview/scim/v2/Me": @@ -148,7 +150,10 @@ func TestProfileLoadSPOGConfigType(t *testing.T) { wantValid: true, }, { - name: "SPOG workspace profile validated as workspace", + // SPOG with a real workspace_id: workspace probe (CurrentUser.Me) + // fails on the mock, account probe (Workspaces.List) succeeds — + // the OR makes the profile valid. + name: "SPOG workspace profile valid when account probe succeeds", host: spogServer.URL, accountID: "spog-acct", workspaceID: "ws-123", @@ -201,6 +206,52 @@ func TestProfileLoadSPOGConfigType(t *testing.T) { } } +// TestProfileLoadSPOGWorkspaceCredential covers the inverse of the +// account-OAuth case: a workspace-scoped credential (e.g. a PAT) against a +// SPOG host. CurrentUser.Me succeeds, Workspaces.List fails (no account-level +// access). The OR of the two probes must still mark the profile Valid=true. +func TestProfileLoadSPOGWorkspaceCredential(t *testing.T) { + const accountID = "spog-acct" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/.well-known/databricks-config": + _ = json.NewEncoder(w).Encode(map[string]any{ + "account_id": accountID, + "oidc_endpoint": r.Host + "/oidc/accounts/" + accountID, + }) + case "/api/2.0/preview/scim/v2/Me": + _ = json.NewEncoder(w).Encode(map[string]any{"userName": "test-user"}) + case "/api/2.0/accounts/" + accountID + "/workspaces": + http.Error(w, "user lacks account-level access", http.StatusForbidden) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + t.Cleanup(server.Close) + + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + t.Setenv("HOME", dir) + if runtime.GOOS == "windows" { + t.Setenv("USERPROFILE", dir) + } + + content := "[ws-cred-on-spog]\nhost = " + server.URL + "\ntoken = test-token\naccount_id = " + accountID + "\nworkspace_id = ws-123\n" + require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600)) + + p := &profileMetadata{ + Name: "ws-cred-on-spog", + Host: server.URL, + AccountID: accountID, + } + p.Load(t.Context(), configFile, false) + + assert.True(t, p.Valid, "workspace probe alone should make the profile valid") + assert.NotEmpty(t, p.Host) + assert.NotEmpty(t, p.AuthType) +} + func TestClassicAccountsHostConfigType(t *testing.T) { // Classic accounts.* hosts can't be tested through Load() because httptest // generates 127.0.0.1 URLs. Verify directly that ConfigType() classifies @@ -217,9 +268,12 @@ func TestClassicAccountsHostConfigType(t *testing.T) { } func TestProfileLoadNoDiscoveryStaysWorkspace(t *testing.T) { - // When .well-known returns 404 and the unified-host fallback is false, - // the SPOG override should NOT trigger even if account_id is set. The - // profile should stay WorkspaceConfig and validate via CurrentUser.Me. + // account_id can linger in a profile from a prior account login on the + // same profile name (e.g. user logged into accounts.cloud.databricks.com, + // logged out, then re-used the profile name for a workspace login). A + // stale account_id must not promote the profile to account validation + // when the host itself isn't an account/SPOG surface — the workspace + // probe is still the right signal. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") switch r.URL.Path { diff --git a/libs/auth/config_type.go b/libs/auth/config_type.go index 0d93b1bf075..956da780581 100644 --- a/libs/auth/config_type.go +++ b/libs/auth/config_type.go @@ -27,10 +27,10 @@ func HasUnifiedHostSignal(discoveryURL string) bool { // of an accountID guard: SPOG routing requires an account ID to construct the // OAuth URL, so a nil or empty accountID always returns false. // -// The accountID parameter is separate from cfg.AccountID so that callers can -// control the source: ResolveConfigType passes cfg.AccountID (from config file), -// while ToOAuthArgument passes the caller-provided value to avoid env var -// contamination (DATABRICKS_ACCOUNT_ID or .well-known back-fill). +// The accountID parameter is separate from cfg.AccountID so that callers +// (currently ToOAuthArgument) can pass the caller-provided value to avoid +// env-var contamination (DATABRICKS_ACCOUNT_ID or .well-known back-fill) +// that would otherwise misroute plain workspace hosts as SPOG. func IsSPOG(cfg *config.Config, accountID string) bool { if accountID == "" { return false @@ -38,24 +38,19 @@ func IsSPOG(cfg *config.Config, accountID string) bool { return HasUnifiedHostSignal(cfg.DiscoveryURL) } -// ResolveConfigType determines the effective ConfigType for a resolved config. -// The SDK's ConfigType() classifies based on the host URL prefix alone, which -// misclassifies SPOG hosts (they don't match the accounts.* prefix). This -// function additionally uses IsSPOG to detect SPOG hosts. -// -// The cfg must already be resolved (via EnsureResolved) before calling this. -func ResolveConfigType(cfg *config.Config) config.ConfigType { - configType := cfg.ConfigType() - if configType == config.AccountConfig { - return configType - } - - if !IsSPOG(cfg, cfg.AccountID) { - return configType +// IsSPOGHost reports whether cfg points at a unified SPOG host: account-scoped +// OIDC discovery and NOT a classic accounts.* host. Classic accounts.* hosts +// share the same OIDC shape, so IsSPOG alone can't tell them apart; layer +// IsClassicAccountHost on top to disambiguate. +func IsSPOGHost(cfg *config.Config) bool { + if IsClassicAccountHost(cfg.CanonicalHostName()) { + return false } + return IsSPOG(cfg, cfg.AccountID) +} - if cfg.WorkspaceID != "" && cfg.WorkspaceID != WorkspaceIDNone { - return config.WorkspaceConfig - } - return config.AccountConfig +// IsClassicWorkspaceHost reports whether cfg points at a classic workspace +// host: neither a classic accounts.* host nor a SPOG host. +func IsClassicWorkspaceHost(cfg *config.Config) bool { + return !IsClassicAccountHost(cfg.CanonicalHostName()) && !IsSPOG(cfg, cfg.AccountID) } diff --git a/libs/auth/config_type_test.go b/libs/auth/config_type_test.go index 8ebe8ff7d68..7b35328934e 100644 --- a/libs/auth/config_type_test.go +++ b/libs/auth/config_type_test.go @@ -24,79 +24,77 @@ func TestHasUnifiedHostSignal(t *testing.T) { } } -func TestResolveConfigType(t *testing.T) { +func TestIsSPOG(t *testing.T) { cases := []struct { - name string - cfg *config.Config - want config.ConfigType + name string + cfg *config.Config + accountID string + want bool }{ { - name: "classic accounts host stays AccountConfig", - cfg: &config.Config{ - Host: "https://accounts.cloud.databricks.com", - AccountID: "acct-123", - }, - want: config.AccountConfig, + name: "account-scoped OIDC with account_id", + cfg: &config.Config{DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server"}, + accountID: "acct-123", + want: true, }, { - name: "SPOG account-scoped OIDC without workspace routes to AccountConfig", - cfg: &config.Config{ - Host: "https://spog.databricks.com", - AccountID: "acct-123", - DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server", - }, - want: config.AccountConfig, + name: "account-scoped OIDC without account_id", + cfg: &config.Config{DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server"}, + accountID: "", + want: false, }, { - name: "SPOG account-scoped OIDC with workspace routes to WorkspaceConfig", - cfg: &config.Config{ - Host: "https://spog.databricks.com", - AccountID: "acct-123", - WorkspaceID: "ws-456", - DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server", - }, - want: config.WorkspaceConfig, - }, - { - name: "SPOG account-scoped OIDC with workspace_id=none routes to AccountConfig", - cfg: &config.Config{ - Host: "https://spog.databricks.com", - AccountID: "acct-123", - WorkspaceID: "none", - DiscoveryURL: "https://spog.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server", - }, - want: config.AccountConfig, - }, - { - name: "workspace-scoped OIDC with account_id stays WorkspaceConfig", - cfg: &config.Config{ - Host: "https://workspace.databricks.com", - AccountID: "acct-123", - DiscoveryURL: "https://workspace.databricks.com/oidc/.well-known/oauth-authorization-server", - }, - want: config.WorkspaceConfig, - }, - { - name: "no discovery stays WorkspaceConfig", - cfg: &config.Config{ - Host: "https://workspace.databricks.com", - AccountID: "acct-123", - }, - want: config.WorkspaceConfig, - }, - { - name: "plain workspace without account_id", - cfg: &config.Config{ - Host: "https://workspace.databricks.com", - }, - want: config.WorkspaceConfig, + name: "workspace-scoped OIDC with account_id back-filled", + cfg: &config.Config{DiscoveryURL: "https://workspace.databricks.com/oidc/.well-known/oauth-authorization-server"}, + accountID: "acct-123", + want: false, }, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - got := ResolveConfigType(tc.cfg) - assert.Equal(t, tc.want, got) + assert.Equal(t, tc.want, IsSPOG(tc.cfg, tc.accountID)) }) } } + +// Configs used across the host-classification tests below. The three host +// shapes are mutually exclusive: exactly one helper returns true per cfg. +// accounts-dod.* is a second classic accounts variant — same OIDC shape +// and classification as accounts.*, different URL prefix. +var ( + classicAccountCfg = &config.Config{ + Host: "https://accounts.cloud.databricks.com", + AccountID: "acct-123", + DiscoveryURL: "https://accounts.cloud.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server", + } + classicAccountDodCfg = &config.Config{ + Host: "https://accounts-dod.cloud.databricks.us", + AccountID: "acct-123", + DiscoveryURL: "https://accounts-dod.cloud.databricks.us/oidc/accounts/acct-123/.well-known/oauth-authorization-server", + } + spogCfg = &config.Config{ + Host: "https://spog.gcp.databricks.com", + AccountID: "acct-123", + DiscoveryURL: "https://spog.gcp.databricks.com/oidc/accounts/acct-123/.well-known/oauth-authorization-server", + } + classicWorkspaceCfg = &config.Config{ + Host: "https://dbc-xxxx.cloud.databricks.com", + AccountID: "acct-123", + DiscoveryURL: "https://dbc-xxxx.cloud.databricks.com/oidc/.well-known/oauth-authorization-server", + } +) + +func TestIsSPOGHost(t *testing.T) { + assert.False(t, IsSPOGHost(classicAccountCfg), "classic accounts.* shares the SPOG OIDC shape but is not SPOG") + assert.False(t, IsSPOGHost(classicAccountDodCfg), "classic accounts-dod.* shares the SPOG OIDC shape but is not SPOG") + assert.True(t, IsSPOGHost(spogCfg)) + assert.False(t, IsSPOGHost(classicWorkspaceCfg)) +} + +func TestIsClassicWorkspaceHost(t *testing.T) { + assert.False(t, IsClassicWorkspaceHost(classicAccountCfg)) + assert.False(t, IsClassicWorkspaceHost(classicAccountDodCfg)) + assert.False(t, IsClassicWorkspaceHost(spogCfg)) + assert.True(t, IsClassicWorkspaceHost(classicWorkspaceCfg)) +}