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

auth/oauth/oauth.go

@@ -2,14 +2,28 @@ package oauth
 
 import (
 	"context"
+	"encoding/json"
 	"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 (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://<host>/oidc/accounts/{account_id}".
+const accountIDPlaceholder = "{account_id}"
+
 var azureTenants = map[string]string{
 	".dev.azuredatabricks.net":     "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc",
 	".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68",
@@ -35,7 +49,19 @@ 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.
+	//
+	// 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 {
@@ -47,6 +73,126 @@ 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 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)
+
+	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
+	}
+
+	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, 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 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
+// 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 +281,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
 }