From 13bc4ab41de7b4ccee20269fdc606e04ac982557 Mon Sep 17 00:00:00 2001 From: "c1-dev-bot[bot]" <2740113+c1-dev-bot[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:55:20 +0000 Subject: [PATCH] feat: add workload identity federation (OAuth token exchange) support Add a second authentication method so the connector can authenticate to Databricks using workload identity federation (RFC 8693 token exchange) instead of a static client secret. The connector presents an externally-issued JWT (e.g. SPIFFE JWT-SVID, Kubernetes projected ServiceAccount token, or any OIDC workload-identity token) and exchanges it for a short-lived Databricks OAuth token. New CLI flags / env vars: - --databricks-token-file / BATON_DATABRICKS_TOKEN_FILE Path to a file containing the external JWT. Re-read on each token refresh to support credential rotation by sidecars. - --databricks-token / BATON_DATABRICKS_TOKEN Inline external JWT for short-lived / CI scenarios. The existing client-secret method remains and is the default. Exactly one of client-secret, token-file, or token must be provided (enforced via mutual-exclusion and at-least-one config constraints). Fixes: CXH-2028 --- config_schema.json | 38 ++++++++++++--- pkg/config/conf.gen.go | 2 + pkg/config/config.go | 18 ++++++- pkg/connector/connector.go | 14 +++++- pkg/databricks/auth.go | 99 +++++++++++++++++++++++++++++++++++++- 5 files changed, 160 insertions(+), 11 deletions(-) diff --git a/config_schema.json b/config_schema.json index 672f3a60..17e81637 100644 --- a/config_schema.json +++ b/config_schema.json @@ -126,13 +126,21 @@ "name": "databricks-client-secret", "displayName": "OAuth2 Client Secret", "description": "The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API", - "isRequired": true, "isSecret": true, - "stringField": { - "rules": { - "isRequired": true - } - } + "stringField": {} + }, + { + "name": "databricks-token-file", + "displayName": "Federation Token File", + "description": "Path to a file containing an external JWT for workload identity federation (e.g. a SPIFFE JWT-SVID or Kubernetes projected ServiceAccount token). The file is re-read on each token refresh to support credential rotation.", + "stringField": {} + }, + { + "name": "databricks-token", + "displayName": "Federation Token", + "description": "An external JWT for workload identity federation (RFC 8693 token exchange). Use --databricks-token-file for automatic credential rotation.", + "isSecret": true, + "stringField": {} }, { "name": "hostname", @@ -143,6 +151,24 @@ } } ], + "constraints": [ + { + "kind": "CONSTRAINT_KIND_MUTUALLY_EXCLUSIVE", + "fieldNames": [ + "databricks-client-secret", + "databricks-token-file", + "databricks-token" + ] + }, + { + "kind": "CONSTRAINT_KIND_AT_LEAST_ONE", + "fieldNames": [ + "databricks-client-secret", + "databricks-token-file", + "databricks-token" + ] + } + ], "displayName": "Databricks", "helpUrl": "/docs/baton/databricks", "iconUrl": "/static/app-icons/databricks.svg" diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index 80fb65de..03d8860a 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -8,6 +8,8 @@ type Databricks struct { AccountId string `mapstructure:"account-id"` DatabricksClientId string `mapstructure:"databricks-client-id"` DatabricksClientSecret string `mapstructure:"databricks-client-secret"` + DatabricksTokenFile string `mapstructure:"databricks-token-file"` + DatabricksToken string `mapstructure:"databricks-token"` Hostname string `mapstructure:"hostname"` BaseUrl string `mapstructure:"base-url"` } diff --git a/pkg/config/config.go b/pkg/config/config.go index 3d0cc691..338a685b 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -21,9 +21,19 @@ var ( "databricks-client-secret", field.WithDescription("The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API"), field.WithIsSecret(true), - field.WithRequired(true), field.WithDisplayName("OAuth2 Client Secret"), ) + DatabricksTokenFileField = field.StringField( + "databricks-token-file", + field.WithDescription("Path to a file containing an external JWT for workload identity federation (e.g. a SPIFFE JWT-SVID or Kubernetes projected ServiceAccount token). The file is re-read on each token refresh to support credential rotation."), + field.WithDisplayName("Federation Token File"), + ) + DatabricksTokenField = field.StringField( + "databricks-token", + field.WithDescription("An external JWT for workload identity federation (RFC 8693 token exchange). Use --databricks-token-file for automatic credential rotation."), + field.WithIsSecret(true), + field.WithDisplayName("Federation Token"), + ) AccountHostnameField = field.StringField( "account-hostname", field.WithDefaultValue("accounts.cloud.databricks.com"), @@ -47,6 +57,8 @@ var ( AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, + DatabricksTokenFileField, + DatabricksTokenField, HostnameField, BaseURLField, } @@ -58,4 +70,8 @@ var Config = field.NewConfiguration( field.WithConnectorDisplayName("Databricks"), field.WithHelpUrl("/docs/baton/databricks"), field.WithIconUrl("/static/app-icons/databricks.svg"), + field.WithConstraints( + field.FieldsMutuallyExclusive(DatabricksClientSecretField, DatabricksTokenFileField, DatabricksTokenField), + field.FieldsAtLeastOneUsed(DatabricksClientSecretField, DatabricksTokenFileField, DatabricksTokenField), + ), ) diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 31d679bf..15177e3a 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -189,13 +189,23 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.Connect func prepareClientAuth(_ context.Context, cfg *config.Databricks, l *zap.Logger) databricks.Auth { accountID := cfg.AccountId databricksClientId := cfg.DatabricksClientId - databricksClientSecret := cfg.DatabricksClientSecret accountHostname := getAccountHostname(cfg, cfg.Hostname) + if cfg.DatabricksTokenFile != "" || cfg.DatabricksToken != "" { + l.Info("using workload identity federation (token exchange)") + return databricks.NewTokenFederation( + accountID, + databricksClientId, + cfg.DatabricksTokenFile, + cfg.DatabricksToken, + accountHostname, + ) + } + return databricks.NewOAuth2( accountID, databricksClientId, - databricksClientSecret, + cfg.DatabricksClientSecret, accountHostname, ) } diff --git a/pkg/databricks/auth.go b/pkg/databricks/auth.go index 63cda5bc..bd4e38f0 100644 --- a/pkg/databricks/auth.go +++ b/pkg/databricks/auth.go @@ -2,8 +2,14 @@ package databricks import ( "context" + "encoding/json" "fmt" + "io" "net/http" + "net/url" + "os" + "strings" + "time" "github.com/conductorone/baton-sdk/pkg/uhttp" "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" @@ -51,6 +57,95 @@ func (o *OAuth2) GetClient(ctx context.Context) (*http.Client, error) { return httpClient, nil } -func (o *OAuth2) Apply(req *http.Request) { - // No need to set the Authorization header here, the oauth2 client does it automatically +func (o *OAuth2) Apply(req *http.Request) {} + +// TokenFederation authenticates via RFC 8693 token exchange. +// It presents an externally-issued JWT (e.g. SPIFFE JWT-SVID, Kubernetes SA token) +// and exchanges it for a short-lived Databricks OAuth token. +type TokenFederation struct { + clientID string + tokenURL string + tokenFile string + staticToken string } + +func NewTokenFederation(accId, clientId, tokenFile, token, accountHostname string) *TokenFederation { + return &TokenFederation{ + clientID: clientId, + tokenURL: fmt.Sprintf("https://%s/oidc/accounts/%s/v1/token", accountHostname, accId), + tokenFile: tokenFile, + staticToken: token, + } +} + +func (tf *TokenFederation) readSubjectToken() (string, error) { + if tf.tokenFile != "" { + data, err := os.ReadFile(tf.tokenFile) + if err != nil { + return "", fmt.Errorf("failed to read token file %s: %w", tf.tokenFile, err) + } + return strings.TrimSpace(string(data)), nil + } + return tf.staticToken, nil +} + +// Token performs an RFC 8693 token exchange and returns the resulting OAuth2 token. +func (tf *TokenFederation) Token() (*oauth2.Token, error) { + subjectToken, err := tf.readSubjectToken() + if err != nil { + return nil, err + } + if subjectToken == "" { + return nil, fmt.Errorf("empty subject token for workload identity federation") + } + + form := url.Values{ + "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, + "subject_token": {subjectToken}, + "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, + "client_id": {tf.clientID}, + "scope": {"all-apis"}, + } + + resp, err := http.PostForm(tf.tokenURL, form) + if err != nil { + return nil, fmt.Errorf("token exchange request failed: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read token exchange response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("token exchange failed with status %d: %s", resp.StatusCode, string(body)) + } + + var tokenResp struct { + AccessToken string `json:"access_token"` + TokenType string `json:"token_type"` + ExpiresIn int64 `json:"expires_in"` + } + if err := json.Unmarshal(body, &tokenResp); err != nil { + return nil, fmt.Errorf("failed to decode token exchange response: %w", err) + } + + token := &oauth2.Token{ + AccessToken: tokenResp.AccessToken, + TokenType: tokenResp.TokenType, + } + if tokenResp.ExpiresIn > 0 { + token.Expiry = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second) + } + + return token, nil +} + +func (tf *TokenFederation) GetClient(ctx context.Context) (*http.Client, error) { + ts := oauth2.ReuseTokenSource(nil, tf) + httpClient := oauth2.NewClient(ctx, ts) + return httpClient, nil +} + +func (tf *TokenFederation) Apply(req *http.Request) {}