From a9cef4dcde6affd7f7f9bf6384cabe4e27e413e9 Mon Sep 17 00:00:00 2001 From: Madhavendra Rathore Date: Thu, 4 Jun 2026 01:52:15 +0530 Subject: [PATCH 1/3] Fix U2M/M2M OAuth for SPOG (unified) AWS hosts SPOG / unified hosts (a single *.databricks.com custom URL that fronts workspaces across accounts, addressed with ?o=) failed OAuth in two ways: 1. InferCloudFromHost returned Unknown for bare *.databricks.com hosts, so the U2M authenticator errored ("unhandled cloud type") before any OAuth. 2. For AWS/GCP, GetEndpoint discovered the OIDC endpoint by hitting https:///oidc directly. On a SPOG host that resolves to the account-agnostic account-console authorize endpoint, which mints a token for the caller's default account. The target workspace (owned by a different account) then rejects it with 400 "Invalid Token". Fixes: - InferCloudFromHost: classify bare *.databricks.com hosts as AWS (checked after Azure/GCP so those remain correctly classified). - GetEndpoint (AWS/GCP): resolve the OIDC issuer via /.well-known/databricks-config, substituting the {account_id} placeholder, so unified/SPOG hosts use their account-rooted endpoint. Normal workspace hosts advertise https:///oidc, so the issuer is identical to before. Any failure (absent endpoint, non-200, unparseable, timeout) falls back to the previous bare-host issuer, so existing behavior is preserved. This mirrors databricks-sdk-go's config host-metadata resolution. Adds oauth_test.go covering cloud inference (incl. SPOG and the GCP/Azure disambiguation), {account_id} substitution, and the databricks-config failure/fallback paths. Tested end-to-end against staging via the real driver auth paths: - AWS non-SPOG (e2-dogfood...cloud.databricks.com): U2M and M2M pass. - AWS SPOG (dogfood.staging.databricks.com): U2M and M2M pass. - Azure SPOG (dogfood-spog.staging.azuredatabricks.net): U2M passes (Azure path unchanged; account resolution handled by the host redirector). - Full unit suite passes with no regressions. Co-authored-by: Isaac --- auth/oauth/oauth.go | 89 ++++++++++++++++++++++++- auth/oauth/oauth_test.go | 137 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 225 insertions(+), 1 deletion(-) create mode 100644 auth/oauth/oauth_test.go diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index 0df9d5c4..8e0134f2 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -2,14 +2,21 @@ package oauth import ( "context" + "encoding/json" "errors" "fmt" + "net/http" "strings" + "time" "github.com/coreos/go-oidc/v3/oidc" "golang.org/x/oauth2" ) +// hostConfigTimeout bounds the /.well-known/databricks-config lookup so it cannot +// stall connection setup; on any failure we fall back to bare-host OIDC discovery. +const hostConfigTimeout = 30 * time.Second + var azureTenants = map[string]string{ ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68", @@ -35,7 +42,12 @@ func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error) return oauth2.Endpoint{AuthURL: authURL, TokenURL: tokenURL}, nil } - issuerURL := fmt.Sprintf("https://%s/oidc", hostName) + // AWS / GCP. Resolve the OIDC issuer via /.well-known/databricks-config so that + // unified / SPOG hosts (one host fronting workspaces across multiple accounts) + // use their account-rooted endpoint instead of the account-agnostic console login. + // For normal workspace hosts this resolves to https:///oidc, identical to + // the previous behavior. + issuerURL := resolveOIDCIssuer(ctx, hostName) ctx = oidc.InsecureIssuerURLContext(ctx, issuerURL) provider, err := oidc.NewProvider(ctx, issuerURL) if err != nil { @@ -47,6 +59,71 @@ func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error) return endpoint, err } +// hostMetadata is the subset of /.well-known/databricks-config we consume. +type hostMetadata struct { + OIDCEndpoint string `json:"oidc_endpoint"` + AccountID string `json:"account_id"` +} + +// resolveOIDCIssuer returns the OIDC issuer URL to use for AWS/GCP OAuth discovery. +// +// On a unified / SPOG host, the bare-host OIDC discovery doc points at the +// account-agnostic account-console login. That mints a token for the caller's +// default account, which the target workspace rejects ("Invalid Token") when the +// workspace belongs to a different account. Such hosts advertise the correct, +// account-rooted OIDC endpoint via /.well-known/databricks-config (with an +// {account_id} placeholder); we consult it and substitute the account id. +// +// For a normal workspace host the advertised endpoint is just https:///oidc, +// so the result is identical to the historical bare-host issuer. Any failure +// (endpoint absent, non-200, unparseable, missing field, timeout) falls back to +// the bare-host issuer, preserving existing behavior. +func resolveOIDCIssuer(ctx context.Context, hostName string) string { + fallback := fmt.Sprintf("https://%s/oidc", hostName) + + url := fmt.Sprintf("https://%s/.well-known/databricks-config", hostName) + client := &http.Client{Timeout: hostConfigTimeout} + + meta, ok := fetchHostMetadata(ctx, client, url) + if !ok || meta.OIDCEndpoint == "" { + return fallback + } + + return substituteAccountID(meta) +} + +// substituteAccountID resolves the {account_id} placeholder in the advertised +// oidc_endpoint. Workspace hosts have no placeholder and are returned unchanged. +func substituteAccountID(meta hostMetadata) string { + return strings.ReplaceAll(meta.OIDCEndpoint, "{account_id}", meta.AccountID) +} + +// fetchHostMetadata GETs /.well-known/databricks-config and decodes it. The bool +// is false on any failure (request error, non-200, unparseable body) so callers +// fall back to bare-host discovery. +func fetchHostMetadata(ctx context.Context, client *http.Client, url string) (hostMetadata, bool) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return hostMetadata{}, false + } + + resp, err := client.Do(req) + if err != nil { + return hostMetadata{}, false + } + defer resp.Body.Close() //nolint:errcheck + + if resp.StatusCode != http.StatusOK { + return hostMetadata{}, false + } + + var meta hostMetadata + if err := json.NewDecoder(resp.Body).Decode(&meta); err != nil { + return hostMetadata{}, false + } + return meta, true +} + func GetScopes(hostName string, scopes []string) []string { for _, s := range []string{oidc.ScopeOfflineAccess} { if !HasScope(scopes, s) { @@ -135,6 +212,16 @@ func InferCloudFromHost(hostname string) CloudType { return GCP } } + + // Unified / SPOG (Single Pane of Glass) AWS hosts use bare *.databricks.com + // custom URLs (e.g. .databricks.com, .staging.databricks.com) that + // match none of the lists above. Treat them as AWS. This is checked last so the + // more specific Azure (.azuredatabricks.net) and GCP (.gcp.databricks.com) hosts + // are classified first. + if strings.Contains(hostname, "databricks.com") { + return AWS + } + return Unknown } diff --git a/auth/oauth/oauth_test.go b/auth/oauth/oauth_test.go new file mode 100644 index 00000000..fe0bfd8e --- /dev/null +++ b/auth/oauth/oauth_test.go @@ -0,0 +1,137 @@ +package oauth + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +func TestInferCloudFromHost(t *testing.T) { + cases := []struct { + host string + want CloudType + }{ + // Standard per-workspace hosts. + {"dbc-1234.cloud.databricks.com", AWS}, + {"example.cloud.databricks.us", AWS}, + {"foo.dev.databricks.com", AWS}, + {"adb-123.azuredatabricks.net", Azure}, + {"x.databricks.azure.us", Azure}, + {"y.databricks.azure.cn", Azure}, + {"ws.gcp.databricks.com", GCP}, + // SPOG / unified custom-URL AWS hosts (the fix): must classify as AWS, + // not Unknown, and must NOT be swallowed by the GCP/Azure checks. + {"pecoaws.databricks.com", AWS}, + {"dogfood.staging.databricks.com", AWS}, + // Azure SPOG stays Azure. + {"dogfood-spog.staging.azuredatabricks.net", Azure}, + // GCP custom host must remain GCP even though it contains "databricks.com". + {"foo.gcp.databricks.com", GCP}, + // Truly unrelated host stays Unknown. + {"example.com", Unknown}, + } + + for _, tc := range cases { + t.Run(tc.host, func(t *testing.T) { + if got := InferCloudFromHost(tc.host); got != tc.want { + t.Fatalf("InferCloudFromHost(%q) = %v, want %v", tc.host, got, tc.want) + } + }) + } +} + +func TestGetAzureDnsZone(t *testing.T) { + // Documents current behavior: the generic suffix is matched first, so staging + // and dev Azure hosts resolve to the prod tenant. (Separate fix tracked.) + cases := []struct { + host string + want string + }{ + {"adb-123.azuredatabricks.net", ".azuredatabricks.net"}, + {"x.databricks.azure.us", ".databricks.azure.us"}, + {"nope.example.com", ""}, + } + for _, tc := range cases { + t.Run(tc.host, func(t *testing.T) { + if got := GetAzureDnsZone(tc.host); got != tc.want { + t.Fatalf("GetAzureDnsZone(%q) = %q, want %q", tc.host, got, tc.want) + } + }) + } +} + +func TestResolveOIDCIssuer_substitutesAccountID(t *testing.T) { + // Unified / SPOG host advertises an account-rooted endpoint with a placeholder. + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/.well-known/databricks-config" { + w.WriteHeader(http.StatusNotFound) + return + } + _, _ = w.Write([]byte(`{ + "oidc_endpoint": "https://spog.example.com/oidc/accounts/{account_id}", + "account_id": "7a99b43c-b46c-432b-b0a7-814217701909", + "host_type": "unified" + }`)) + })) + defer srv.Close() + + meta, ok := fetchHostMetadata(context.Background(), srv.Client(), srv.URL+"/.well-known/databricks-config") + if !ok { + t.Fatal("fetchHostMetadata returned ok=false, want true") + } + got := substituteAccountID(meta) + want := "https://spog.example.com/oidc/accounts/7a99b43c-b46c-432b-b0a7-814217701909" + if got != want { + t.Fatalf("issuer = %q, want %q", got, want) + } +} + +func TestResolveOIDCIssuer_workspaceHostUnchanged(t *testing.T) { + // Normal workspace host: endpoint has no placeholder, so it is returned as-is + // (equivalent to the historical https:///oidc). + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{ + "oidc_endpoint": "https://ws.cloud.databricks.com/oidc", + "account_id": "7a99b43c-b46c-432b-b0a7-814217701909", + "host_type": "workspace" + }`)) + })) + defer srv.Close() + + meta, ok := fetchHostMetadata(context.Background(), srv.Client(), srv.URL+"/.well-known/databricks-config") + if !ok { + t.Fatal("fetchHostMetadata returned ok=false, want true") + } + if got := substituteAccountID(meta); got != "https://ws.cloud.databricks.com/oidc" { + t.Fatalf("issuer = %q, want unchanged workspace endpoint", got) + } +} + +func TestFetchHostMetadata_failuresFallBack(t *testing.T) { + t.Run("404", func(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + if _, ok := fetchHostMetadata(context.Background(), srv.Client(), srv.URL); ok { + t.Fatal("ok=true on 404, want false (fallback)") + } + }) + + t.Run("garbage body", func(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("not json")) + })) + defer srv.Close() + if _, ok := fetchHostMetadata(context.Background(), srv.Client(), srv.URL); ok { + t.Fatal("ok=true on garbage body, want false (fallback)") + } + }) + + t.Run("unreachable", func(t *testing.T) { + if _, ok := fetchHostMetadata(context.Background(), &http.Client{}, "https://127.0.0.1:1/nope"); ok { + t.Fatal("ok=true on unreachable host, want false (fallback)") + } + }) +} From 548e4bfbd58c2543ae1e81875c5c859b2bfb3a65 Mon Sep 17 00:00:00 2001 From: Madhavendra Rathore Date: Thu, 4 Jun 2026 11:02:06 +0530 Subject: [PATCH 2/3] Address review: harden databricks-config issuer resolution - Fall back to the bare-host issuer when oidc_endpoint has an {account_id} placeholder but account_id is empty (previously produced a malformed ".../accounts/" issuer and a hard error, violating the fallback contract). - Validate the metadata-supplied issuer before use (isValidDatabricksIssuer): must be https, on a recognized Databricks domain, with no unresolved placeholder. Covers the http:// cleartext case and bounds the trust placed in the host-supplied document given InsecureIssuerURLContext disables the issuer-match check. - Log (debug/warn) on every fallback path so a misconfigured SPOG host leaves a breadcrumb instead of silently reproducing the opaque 400. - Inject the *http.Client into resolveOIDCIssuer and bound OIDC discovery with the same client (oidc.ClientContext); tighten the timeout 30s -> 10s. - Tests: drive resolveOIDCIssuer end-to-end via httptest (substitute / workspace-unchanged / empty-account_id / empty-endpoint / non-https / non-databricks-host / 404 / garbage / unreachable) and add isValidDatabricksIssuer coverage. Co-authored-by: Isaac --- auth/oauth/oauth.go | 72 ++++++++++++++---- auth/oauth/oauth_test.go | 156 +++++++++++++++++++++++++++++---------- 2 files changed, 176 insertions(+), 52 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index 8e0134f2..b47e9207 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -6,16 +6,23 @@ import ( "errors" "fmt" "net/http" + "net/url" "strings" "time" "github.com/coreos/go-oidc/v3/oidc" + "github.com/rs/zerolog/log" "golang.org/x/oauth2" ) -// hostConfigTimeout bounds the /.well-known/databricks-config lookup so it cannot -// stall connection setup; on any failure we fall back to bare-host OIDC discovery. -const hostConfigTimeout = 30 * time.Second +// hostConfigTimeout bounds the /.well-known/databricks-config lookup (and the +// subsequent OIDC discovery) so they cannot stall connection setup; on any +// failure we fall back to bare-host OIDC discovery. +const hostConfigTimeout = 10 * time.Second + +// accountIDPlaceholder is the token the host-metadata oidc_endpoint uses for the +// account id, e.g. "https:///oidc/accounts/{account_id}". +const accountIDPlaceholder = "{account_id}" var azureTenants = map[string]string{ ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", @@ -47,7 +54,14 @@ func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error) // use their account-rooted endpoint instead of the account-agnostic console login. // For normal workspace hosts this resolves to https:///oidc, identical to // the previous behavior. - issuerURL := resolveOIDCIssuer(ctx, hostName) + // + // NOTE: this client uses the default transport, matching the existing + // oidc.NewProvider discovery below. A connector-supplied transport / TLS config + // (WithTransport, WithSkipTLSHostVerify) is not yet threaded into the OAuth + // endpoint-resolution path; that is a pre-existing limitation, tracked separately. + client := &http.Client{Timeout: hostConfigTimeout} + issuerURL := resolveOIDCIssuer(ctx, client, hostName) + ctx = oidc.ClientContext(ctx, client) ctx = oidc.InsecureIssuerURLContext(ctx, issuerURL) provider, err := oidc.NewProvider(ctx, issuerURL) if err != nil { @@ -75,27 +89,57 @@ type hostMetadata struct { // {account_id} placeholder); we consult it and substitute the account id. // // For a normal workspace host the advertised endpoint is just https:///oidc, -// so the result is identical to the historical bare-host issuer. Any failure -// (endpoint absent, non-200, unparseable, missing field, timeout) falls back to -// the bare-host issuer, preserving existing behavior. -func resolveOIDCIssuer(ctx context.Context, hostName string) string { +// so the result is identical to the historical bare-host issuer. Any failure or +// unusable value (endpoint absent, non-200, unparseable, missing/empty account_id, +// non-https, non-Databricks host, timeout) falls back to the bare-host issuer, +// preserving existing behavior. +func resolveOIDCIssuer(ctx context.Context, client *http.Client, hostName string) string { fallback := fmt.Sprintf("https://%s/oidc", hostName) - url := fmt.Sprintf("https://%s/.well-known/databricks-config", hostName) - client := &http.Client{Timeout: hostConfigTimeout} - - meta, ok := fetchHostMetadata(ctx, client, url) + cfgURL := fmt.Sprintf("https://%s/.well-known/databricks-config", hostName) + meta, ok := fetchHostMetadata(ctx, client, cfgURL) if !ok || meta.OIDCEndpoint == "" { + log.Debug().Msgf("oauth: no usable databricks-config for %q; using bare-host OIDC issuer", hostName) + return fallback + } + + // An account-rooted endpoint needs a non-empty account_id; otherwise the + // placeholder would resolve to a malformed ".../accounts/" issuer. Fall back + // rather than emit it (the function's documented contract). + if strings.Contains(meta.OIDCEndpoint, accountIDPlaceholder) && meta.AccountID == "" { + log.Warn().Msgf("oauth: databricks-config for %q has an %s placeholder but empty account_id; using bare-host OIDC issuer", hostName, accountIDPlaceholder) return fallback } - return substituteAccountID(meta) + issuer := substituteAccountID(meta) + if !isValidDatabricksIssuer(issuer) { + log.Warn().Msgf("oauth: databricks-config for %q advertised an unusable oidc_endpoint %q; using bare-host OIDC issuer", hostName, issuer) + return fallback + } + return issuer } // substituteAccountID resolves the {account_id} placeholder in the advertised // oidc_endpoint. Workspace hosts have no placeholder and are returned unchanged. func substituteAccountID(meta hostMetadata) string { - return strings.ReplaceAll(meta.OIDCEndpoint, "{account_id}", meta.AccountID) + return strings.ReplaceAll(meta.OIDCEndpoint, accountIDPlaceholder, meta.AccountID) +} + +// isValidDatabricksIssuer guards the metadata-supplied OIDC issuer before it is +// passed to discovery: it must be an https URL on a recognized Databricks domain +// with the {account_id} placeholder fully resolved. This bounds the trust placed +// in the host-supplied document (the issuer-match check is disabled via +// InsecureIssuerURLContext because the discovered issuer is cross-host) and avoids +// cleartext OAuth from an http:// endpoint. +func isValidDatabricksIssuer(issuer string) bool { + if strings.Contains(issuer, accountIDPlaceholder) { + return false + } + u, err := url.Parse(issuer) + if err != nil || u.Scheme != "https" || u.Hostname() == "" { + return false + } + return InferCloudFromHost(u.Hostname()) != Unknown } // fetchHostMetadata GETs /.well-known/databricks-config and decodes it. The bool diff --git a/auth/oauth/oauth_test.go b/auth/oauth/oauth_test.go index fe0bfd8e..c6a7e99c 100644 --- a/auth/oauth/oauth_test.go +++ b/auth/oauth/oauth_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "strings" "testing" ) @@ -61,50 +62,129 @@ func TestGetAzureDnsZone(t *testing.T) { } } -func TestResolveOIDCIssuer_substitutesAccountID(t *testing.T) { - // Unified / SPOG host advertises an account-rooted endpoint with a placeholder. - srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/.well-known/databricks-config" { - w.WriteHeader(http.StatusNotFound) - return - } - _, _ = w.Write([]byte(`{ - "oidc_endpoint": "https://spog.example.com/oidc/accounts/{account_id}", - "account_id": "7a99b43c-b46c-432b-b0a7-814217701909", - "host_type": "unified" - }`)) - })) - defer srv.Close() +// TestResolveOIDCIssuer drives resolveOIDCIssuer end-to-end against an httptest +// server: the server stands in for the connection host, so a 200 with a given +// databricks-config body exercises the real fetch + substitution + validation + +// fallback wiring. +func TestResolveOIDCIssuer(t *testing.T) { + const fallbackForUnreachable = "https://no-such-host.invalid/oidc" - meta, ok := fetchHostMetadata(context.Background(), srv.Client(), srv.URL+"/.well-known/databricks-config") - if !ok { - t.Fatal("fetchHostMetadata returned ok=false, want true") + cases := []struct { + name string + // body served at /.well-known/databricks-config (status 200). Empty body + // with status!=200 simulates a missing endpoint. + status int + body string + // want is the exact resolved issuer. If wantFallback is true we instead + // assert the bare-host fallback (https:///oidc). + want string + wantFallback bool + }{ + { + name: "unified host substitutes account_id", + status: 200, + body: `{"oidc_endpoint":"https://spog.databricks.com/oidc/accounts/{account_id}","account_id":"acc-123","host_type":"unified"}`, + want: "https://spog.databricks.com/oidc/accounts/acc-123", + }, + { + name: "workspace host returned unchanged", + status: 200, + body: `{"oidc_endpoint":"https://ws.cloud.databricks.com/oidc","account_id":"acc-123","host_type":"workspace"}`, + want: "https://ws.cloud.databricks.com/oidc", + }, + { + name: "placeholder with empty account_id falls back", + status: 200, + body: `{"oidc_endpoint":"https://spog.databricks.com/oidc/accounts/{account_id}","account_id":""}`, + wantFallback: true, + }, + { + name: "empty oidc_endpoint falls back", + status: 200, + body: `{"account_id":"acc-123"}`, + wantFallback: true, + }, + { + name: "non-https endpoint falls back", + status: 200, + body: `{"oidc_endpoint":"http://spog.databricks.com/oidc","account_id":"acc-123"}`, + wantFallback: true, + }, + { + name: "non-databricks host falls back", + status: 200, + body: `{"oidc_endpoint":"https://evil.example.com/oidc","account_id":"acc-123"}`, + wantFallback: true, + }, + { + name: "404 falls back", + status: 404, + wantFallback: true, + }, + { + name: "garbage body falls back", + status: 200, + body: `not json`, + wantFallback: true, + }, } - got := substituteAccountID(meta) - want := "https://spog.example.com/oidc/accounts/7a99b43c-b46c-432b-b0a7-814217701909" - if got != want { - t.Fatalf("issuer = %q, want %q", got, want) + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/.well-known/databricks-config" { + w.WriteHeader(http.StatusNotFound) + return + } + if tc.status != 200 { + w.WriteHeader(tc.status) + return + } + _, _ = w.Write([]byte(tc.body)) + })) + defer srv.Close() + + host := strings.TrimPrefix(srv.URL, "https://") + got := resolveOIDCIssuer(context.Background(), srv.Client(), host) + + if tc.wantFallback { + if want := "https://" + host + "/oidc"; got != want { + t.Fatalf("issuer = %q, want fallback %q", got, want) + } + return + } + if got != tc.want { + t.Fatalf("issuer = %q, want %q", got, tc.want) + } + }) } -} -func TestResolveOIDCIssuer_workspaceHostUnchanged(t *testing.T) { - // Normal workspace host: endpoint has no placeholder, so it is returned as-is - // (equivalent to the historical https:///oidc). - srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte(`{ - "oidc_endpoint": "https://ws.cloud.databricks.com/oidc", - "account_id": "7a99b43c-b46c-432b-b0a7-814217701909", - "host_type": "workspace" - }`)) - })) - defer srv.Close() + // Sanity: an unreachable host falls back without panicking. + if got := resolveOIDCIssuer(context.Background(), &http.Client{}, "no-such-host.invalid"); got != fallbackForUnreachable { + t.Fatalf("unreachable host issuer = %q, want %q", got, fallbackForUnreachable) + } +} - meta, ok := fetchHostMetadata(context.Background(), srv.Client(), srv.URL+"/.well-known/databricks-config") - if !ok { - t.Fatal("fetchHostMetadata returned ok=false, want true") +func TestIsValidDatabricksIssuer(t *testing.T) { + cases := []struct { + issuer string + want bool + }{ + {"https://spog.databricks.com/oidc/accounts/acc-123", true}, + {"https://ws.cloud.databricks.com/oidc", true}, + {"https://adb-1.azuredatabricks.net/oidc", true}, + {"https://spog.databricks.com/oidc/accounts/{account_id}", false}, // unresolved placeholder + {"http://spog.databricks.com/oidc", false}, // not https + {"https://evil.example.com/oidc", false}, // not a databricks host + {"://bad", false}, // unparseable + {"", false}, } - if got := substituteAccountID(meta); got != "https://ws.cloud.databricks.com/oidc" { - t.Fatalf("issuer = %q, want unchanged workspace endpoint", got) + for _, tc := range cases { + t.Run(tc.issuer, func(t *testing.T) { + if got := isValidDatabricksIssuer(tc.issuer); got != tc.want { + t.Fatalf("isValidDatabricksIssuer(%q) = %v, want %v", tc.issuer, got, tc.want) + } + }) } } From ae8bf4de543d0355f456fa51ab8850f9a03d3566 Mon Sep 17 00:00:00 2001 From: Madhavendra Rathore Date: Thu, 4 Jun 2026 11:56:05 +0530 Subject: [PATCH 3/3] Address follow-up: suffix-match the issuer trust gate The hardening commit made InferCloudFromHost (substring matching) load-bearing for the cross-host issuer trust decision in isValidDatabricksIssuer, so a host like "databricks.com.evil.example" would have passed. Replace the InferCloudFromHost call in the trust gate with isDatabricksHost, which uses HasSuffix against an explicit list of Databricks DNS suffixes (case-insensitive). Cloud routing (InferCloudFromHost) is left unchanged. Tests: add TestIsDatabricksHost and extend TestIsValidDatabricksIssuer / TestResolveOIDCIssuer with suffix-spoof ("databricks.com.evil.example"), "-databricks.com", and bare-apex cases, plus the real SPOG/non-SPOG issuer hosts as positives. Verified end-to-end against staging over the network (opt-in test, not committed): GetEndpoint resolves the SPOG host to its account-rooted endpoint (accounts.staging.cloud.databricks.com/oidc/accounts//v1/{authorize,token}) and the non-SPOG host to its workspace endpoint, unchanged by the gate. Co-authored-by: Isaac --- auth/oauth/oauth.go | 27 +++++++++++++++++++++++++- auth/oauth/oauth_test.go | 41 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index b47e9207..0525e8cd 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -139,7 +139,32 @@ func isValidDatabricksIssuer(issuer string) bool { if err != nil || u.Scheme != "https" || u.Hostname() == "" { return false } - return InferCloudFromHost(u.Hostname()) != Unknown + return isDatabricksHost(u.Hostname()) +} + +// databricksHostSuffixes is the set of DNS suffixes treated as first-party +// Databricks hosts when deciding whether to trust a metadata-supplied OIDC issuer. +// +// Suffix matching (not the substring matching InferCloudFromHost uses for cloud +// routing) is deliberate here: this is a trust gate for a cross-host issuer, so +// "databricks.com.evil.example" must NOT pass. ".databricks.com" covers the +// .cloud/.dev/.gcp variants and the bare unified/SPOG custom URLs. +var databricksHostSuffixes = []string{ + ".databricks.com", + ".cloud.databricks.us", + ".azuredatabricks.net", + ".databricks.azure.us", + ".databricks.azure.cn", +} + +func isDatabricksHost(host string) bool { + host = strings.ToLower(host) + for _, suffix := range databricksHostSuffixes { + if strings.HasSuffix(host, suffix) { + return true + } + } + return false } // fetchHostMetadata GETs /.well-known/databricks-config and decodes it. The bool diff --git a/auth/oauth/oauth_test.go b/auth/oauth/oauth_test.go index c6a7e99c..e9e64291 100644 --- a/auth/oauth/oauth_test.go +++ b/auth/oauth/oauth_test.go @@ -116,6 +116,12 @@ func TestResolveOIDCIssuer(t *testing.T) { body: `{"oidc_endpoint":"https://evil.example.com/oidc","account_id":"acc-123"}`, wantFallback: true, }, + { + name: "suffix-spoof host falls back", + status: 200, + body: `{"oidc_endpoint":"https://databricks.com.evil.example/oidc/accounts/{account_id}","account_id":"acc-123"}`, + wantFallback: true, + }, { name: "404 falls back", status: 404, @@ -170,12 +176,21 @@ func TestIsValidDatabricksIssuer(t *testing.T) { issuer string want bool }{ + // Real hosts from the validated SPOG/non-SPOG flows must stay valid. + {"https://dogfood.staging.databricks.com/oidc/accounts/7a99b43c", true}, // SPOG (unified) issuer + {"https://e2-dogfood.staging.cloud.databricks.com/oidc", true}, // non-SPOG workspace + {"https://accounts.staging.cloud.databricks.com/oidc/accounts/x", true}, // account-rooted {"https://spog.databricks.com/oidc/accounts/acc-123", true}, {"https://ws.cloud.databricks.com/oidc", true}, {"https://adb-1.azuredatabricks.net/oidc", true}, + {"https://x.cloud.databricks.us/oidc", true}, + // Rejections. {"https://spog.databricks.com/oidc/accounts/{account_id}", false}, // unresolved placeholder {"http://spog.databricks.com/oidc", false}, // not https {"https://evil.example.com/oidc", false}, // not a databricks host + {"https://databricks.com.evil.example/oidc", false}, // suffix-spoof: must NOT pass + {"https://notdatabricks.com/oidc", false}, // "-databricks.com" is not ".databricks.com" + {"https://databricks.com/oidc", false}, // bare apex, no subdomain {"://bad", false}, // unparseable {"", false}, } @@ -188,6 +203,32 @@ func TestIsValidDatabricksIssuer(t *testing.T) { } } +func TestIsDatabricksHost(t *testing.T) { + cases := []struct { + host string + want bool + }{ + {"dogfood.staging.databricks.com", true}, + {"e2-dogfood.staging.cloud.databricks.com", true}, + {"foo.gcp.databricks.com", true}, + {"adb-1.azuredatabricks.net", true}, + {"x.cloud.databricks.us", true}, + {"DOGFOOD.STAGING.DATABRICKS.COM", true}, // case-insensitive + {"databricks.com.evil.example", false}, + {"notdatabricks.com", false}, + {"databricks.com", false}, + {"evil.example", false}, + {"", false}, + } + for _, tc := range cases { + t.Run(tc.host, func(t *testing.T) { + if got := isDatabricksHost(tc.host); got != tc.want { + t.Fatalf("isDatabricksHost(%q) = %v, want %v", tc.host, got, tc.want) + } + }) + } +} + func TestFetchHostMetadata_failuresFallBack(t *testing.T) { t.Run("404", func(t *testing.T) { srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {