-
Notifications
You must be signed in to change notification settings - Fork 156
Fix auth profiles misclassifying SPOG hosts as workspace configs
#4929
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 4 commits
32e780a
7b8d122
456280f
a7198ac
c680516
05c4f39
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,16 @@ | ||
|
|
||
| === SPOG account profile should be valid | ||
| >>> [CLI] auth profiles --output json | ||
| { | ||
| "profiles": [ | ||
| { | ||
| "name":"spog-account", | ||
| "host":"[DATABRICKS_URL]", | ||
| "account_id":"spog-acct-123", | ||
| "workspace_id":"none", | ||
| "cloud":"aws", | ||
| "auth_type":"pat", | ||
| "valid":true | ||
| } | ||
| ] | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,15 @@ | ||
| sethome "./home" | ||
|
|
||
| # Create a SPOG account profile: non-accounts.* host with account_id, no workspace_id. | ||
| # Before the fix, this was misclassified as WorkspaceConfig and validated with | ||
| # CurrentUser.Me, which fails on account-scoped SPOG hosts. | ||
| cat > "./home/.databrickscfg" <<EOF | ||
| [spog-account] | ||
| host = ${DATABRICKS_HOST} | ||
| account_id = spog-acct-123 | ||
| workspace_id = none | ||
| token = test-token | ||
| EOF | ||
|
|
||
| title "SPOG account profile should be valid" | ||
| trace $CLI auth profiles --output json |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| Ignore = [ | ||
| "home" | ||
| ] | ||
|
|
||
| # Override .well-known to return account-scoped OIDC, simulating a SPOG host. | ||
| # The oidc_endpoint contains /oidc/accounts/ which signals account-scoped OIDC. | ||
| [[Server]] | ||
| Pattern = "GET /.well-known/databricks-config" | ||
| Response.Body = ''' | ||
| { | ||
| "account_id": "spog-acct-123", | ||
| "oidc_endpoint": "localhost/oidc/accounts/spog-acct-123" | ||
| } | ||
| ''' | ||
|
|
||
| # Provide a handler for the account workspaces list endpoint used by | ||
| # auth profiles validation for account-type configs. | ||
| [[Server]] | ||
| Pattern = "GET /api/2.0/accounts/spog-acct-123/workspaces" | ||
| Response.Body = '[]' |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,9 +5,11 @@ import ( | |
| "errors" | ||
| "fmt" | ||
| "io/fs" | ||
| "strings" | ||
| "sync" | ||
| "time" | ||
|
|
||
| "github.com/databricks/cli/libs/auth" | ||
| "github.com/databricks/cli/libs/cmdio" | ||
| "github.com/databricks/cli/libs/databrickscfg" | ||
| "github.com/databricks/cli/libs/databrickscfg/profile" | ||
|
|
@@ -56,7 +58,30 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV | |
| return | ||
| } | ||
|
|
||
| switch cfg.ConfigType() { | ||
| // ConfigType() classifies based on the host URL prefix (accounts.* → | ||
| // AccountConfig, everything else → WorkspaceConfig). SPOG hosts don't | ||
| // match the accounts.* prefix so they're misclassified as WorkspaceConfig. | ||
| // Use the resolved DiscoveryURL (from .well-known/databricks-config) to | ||
| // detect SPOG hosts with account-scoped OIDC, matching the routing logic | ||
| // in auth.AuthArguments.ToOAuthArgument(). | ||
| configType := cfg.ConfigType() | ||
| isAccountScopedOIDC := cfg.DiscoveryURL != "" && strings.Contains(cfg.DiscoveryURL, "/oidc/accounts/") | ||
| if configType != config.AccountConfig && cfg.AccountID != "" && isAccountScopedOIDC { | ||
| if cfg.WorkspaceID != "" && cfg.WorkspaceID != auth.WorkspaceIDNone { | ||
| configType = config.WorkspaceConfig | ||
| } else { | ||
| configType = config.AccountConfig | ||
| } | ||
| } | ||
|
|
||
| // Legacy backward compat: profiles with Experimental_IsUnifiedHost where | ||
| // .well-known is unreachable (so DiscoveryURL is empty). Matches the | ||
| // fallback in auth.AuthArguments.ToOAuthArgument(). | ||
| if configType == config.InvalidConfig && cfg.Experimental_IsUnifiedHost && cfg.AccountID != "" { | ||
| configType = config.AccountConfig | ||
| } | ||
|
||
|
|
||
| switch configType { | ||
| case config.AccountConfig: | ||
| a, err := databricks.NewAccountClient((*databricks.Config)(cfg)) | ||
| if err != nil { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,6 +1,10 @@ | ||
| package auth | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "net/http" | ||
| "net/http/httptest" | ||
| "os" | ||
| "path/filepath" | ||
| "runtime" | ||
| "testing" | ||
|
|
@@ -74,3 +78,185 @@ func TestProfilesDefaultMarker(t *testing.T) { | |
| require.NoError(t, err) | ||
| assert.Equal(t, "profile-a", defaultProfile) | ||
| } | ||
|
|
||
| // newProfileTestServer creates a mock server for profile validation tests. | ||
| // It serves /.well-known/databricks-config with the given OIDC shape and | ||
| // responds to the workspace/account validation API endpoints. | ||
| func newProfileTestServer(t *testing.T, accountScoped bool, accountID string) *httptest.Server { | ||
|
||
| t.Helper() | ||
| 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": | ||
| oidcEndpoint := r.Host + "/oidc" | ||
| if accountScoped { | ||
| oidcEndpoint = r.Host + "/oidc/accounts/" + accountID | ||
| } | ||
| _ = json.NewEncoder(w).Encode(map[string]any{ | ||
| "account_id": accountID, | ||
| "oidc_endpoint": oidcEndpoint, | ||
| }) | ||
| 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": | ||
| _ = json.NewEncoder(w).Encode([]map[string]any{}) | ||
| default: | ||
| w.WriteHeader(http.StatusNotFound) | ||
| } | ||
| })) | ||
| t.Cleanup(server.Close) | ||
| return server | ||
| } | ||
|
|
||
| func TestProfileLoadSPOGConfigType(t *testing.T) { | ||
| spogServer := newProfileTestServer(t, true, "spog-acct") | ||
| wsServer := newProfileTestServer(t, false, "ws-acct") | ||
|
|
||
| cases := []struct { | ||
| name string | ||
| host string | ||
| accountID string | ||
| workspaceID string | ||
| wantValid bool | ||
| }{ | ||
mihaimitrea-db marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| { | ||
| name: "SPOG account profile validated as account", | ||
| host: spogServer.URL, | ||
| accountID: "spog-acct", | ||
| wantValid: true, | ||
| }, | ||
| { | ||
| name: "SPOG workspace profile validated as workspace", | ||
| host: spogServer.URL, | ||
| accountID: "spog-acct", | ||
| workspaceID: "ws-123", | ||
| wantValid: true, | ||
| }, | ||
| { | ||
| name: "SPOG profile with workspace_id=none validated as account", | ||
| host: spogServer.URL, | ||
| accountID: "spog-acct", | ||
| workspaceID: "none", | ||
| wantValid: true, | ||
| }, | ||
| { | ||
| name: "classic workspace with account_id from discovery stays workspace", | ||
| host: wsServer.URL, | ||
| accountID: "ws-acct", | ||
| wantValid: true, | ||
| }, | ||
| } | ||
|
|
||
| for _, tc := range cases { | ||
| t.Run(tc.name, func(t *testing.T) { | ||
| dir := t.TempDir() | ||
| configFile := filepath.Join(dir, ".databrickscfg") | ||
| t.Setenv("HOME", dir) | ||
| if runtime.GOOS == "windows" { | ||
| t.Setenv("USERPROFILE", dir) | ||
| } | ||
|
|
||
| content := "[test-profile]\nhost = " + tc.host + "\ntoken = test-token\n" | ||
| if tc.accountID != "" { | ||
| content += "account_id = " + tc.accountID + "\n" | ||
| } | ||
| if tc.workspaceID != "" { | ||
| content += "workspace_id = " + tc.workspaceID + "\n" | ||
| } | ||
| require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600)) | ||
|
|
||
| p := &profileMetadata{ | ||
| Name: "test-profile", | ||
| Host: tc.host, | ||
| AccountID: tc.accountID, | ||
| } | ||
| p.Load(t.Context(), configFile, false) | ||
|
|
||
| assert.Equal(t, tc.wantValid, p.Valid, "Valid mismatch") | ||
| assert.NotEmpty(t, p.Host, "Host should be set") | ||
| assert.NotEmpty(t, p.AuthType, "AuthType should be set") | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| func TestProfileLoadUnifiedHostFallback(t *testing.T) { | ||
| // When Experimental_IsUnifiedHost is set but .well-known is unreachable, | ||
| // ConfigType() returns InvalidConfig. The fallback should reclassify as | ||
| // AccountConfig so the profile is still validated. | ||
| 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": | ||
| w.WriteHeader(http.StatusNotFound) | ||
| case "/api/2.0/accounts/unified-acct/workspaces": | ||
| _ = json.NewEncoder(w).Encode([]map[string]any{}) | ||
| 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 := "[unified-profile]\nhost = " + server.URL + "\ntoken = test-token\naccount_id = unified-acct\nexperimental_is_unified_host = true\n" | ||
| require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600)) | ||
|
|
||
| p := &profileMetadata{ | ||
| Name: "unified-profile", | ||
| Host: server.URL, | ||
| AccountID: "unified-acct", | ||
| } | ||
| p.Load(t.Context(), configFile, false) | ||
|
|
||
| assert.True(t, p.Valid, "unified host profile should be valid via fallback") | ||
| assert.NotEmpty(t, p.Host) | ||
| assert.NotEmpty(t, p.AuthType) | ||
| } | ||
|
|
||
| func TestProfileLoadClassicAccountHost(t *testing.T) { | ||
| // Verify that a host with account-scoped OIDC from discovery is validated | ||
|
||
| // as an account config (via Workspaces.List, not CurrentUser.Me). | ||
| 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": "classic-acct", | ||
| "oidc_endpoint": r.Host + "/oidc/accounts/classic-acct", | ||
| }) | ||
| case "/api/2.0/accounts/classic-acct/workspaces": | ||
| _ = json.NewEncoder(w).Encode([]map[string]any{}) | ||
| 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 := "[acct-profile]\nhost = " + server.URL + "\ntoken = test-token\naccount_id = classic-acct\n" | ||
| require.NoError(t, os.WriteFile(configFile, []byte(content), 0o600)) | ||
|
|
||
| p := &profileMetadata{ | ||
| Name: "acct-profile", | ||
| Host: server.URL, | ||
| AccountID: "classic-acct", | ||
| } | ||
| p.Load(t.Context(), configFile, false) | ||
|
|
||
| assert.True(t, p.Valid, "classic account profile should be valid") | ||
| assert.NotEmpty(t, p.Host) | ||
| assert.NotEmpty(t, p.AuthType) | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
suggestion (follow-up): The SPOG detection heuristic (DiscoveryURL contains
/oidc/accounts/+ IsUnifiedHost fallback) is now duplicated between here andlibs/auth/arguments.go:69-82. The cross-referencing comments help, but if the routing logic changes in one place the other can drift.Consider extracting a shared helper in
libs/auth/in a follow-up, e.g.:Not blocking for this PR since the logic is simple and well-documented.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This would be good to have. Bonus points that it moves some of the logic form the SDK to the CLI.