Fix U2M/M2M OAuth for SPOG (unified) AWS hosts

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://<host>/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://<host>/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. <name>.databricks.com, <name>.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
 }
 

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://<host>/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)")
+		}
+	})
+}