diff --git a/README.md b/README.md index 371bbe5e..06ef90af 100644 --- a/README.md +++ b/README.md @@ -18,9 +18,10 @@ right top corner that will open a dropdown menu with the account ID along other options. Another requirement is to have valid credentials to run the connector with. The -connector authenticates with the OAuth client credentials flow, using an -account-level service principal that works across the account and every -workspace it has access to. +connector authenticates as an account-level service principal that works across +the account and every workspace it has access to. It can prove it is that +service principal in two ways: with a client secret, which is the default, or +with workload identity federation. See [Authentication](#authentication). To set this up, create a service principal and add an OAuth secret (client ID and secret) to it. You can do that by going to the user management tab and @@ -84,13 +85,53 @@ deployment name, or numeric workspace ID, matched case-insensitively. ## Authentication -OAuth is the only authentication method: an account-level service principal's -client ID and secret, supplied through `--databricks-client-id` and -`--databricks-client-secret`. There is no personal access token (PAT) or +The connector always authenticates as an account-level service principal, named +by `--databricks-client-id`. There is no personal access token (PAT) or workspace-token option, and no username/password option. -OAuth requires a reachable account API. If the account API check fails at -startup, the connector fails validation rather than falling back to a +Exactly one of the following credentials must be supplied. Passing more than one +is a configuration error. + +### Client secret (default) + +Pass the service principal's OAuth secret to `--databricks-client-secret` (or +`BATON_DATABRICKS_CLIENT_SECRET`). The connector runs the OAuth client +credentials flow against the account token endpoint. + +### Workload identity federation + +Instead of a long-lived secret, the connector can present an externally issued +JWT that identifies the workload, such as a SPIFFE JWT-SVID or a Kubernetes +projected service account token, and exchange it for a short-lived Databricks +OAuth token ([RFC 8693](https://datatracker.ietf.org/doc/html/rfc8693)). + +Create a [service principal federation policy](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-exchange) +on the same service principal, trusting the issuer, subject and audience of that +JWT. Then select the method and point the connector at the token: + +``` +--auth-method workload-identity-file \ + --databricks-oidc-token-file /var/run/secrets/token +``` + +`--auth-method` (or `BATON_AUTH_METHOD`) is required here. Without it the +connector validates against the default method and asks for a client secret. + +The file is read again on every token refresh, so whatever rotates it controls +the lifetime and the connector needs no restart. + +This method runs the connector yourself: on the command line or in service mode. +A ConductorOne-hosted connector has no file system you control, so the option is +not offered there. Hosted support needs ConductorOne to issue the JWT, which is +tracked in CXE-1557. + +Nothing else changes: the exchange happens at the same account token endpoint, +and the resulting token is used for both account and workspace APIs. + +### Validation + +Either method requires a reachable account API. If the account API check fails +at startup, the connector fails validation rather than falling back to a workspace-only sync, even when `--workspaces` is set. To instead exclude specific workspaces from the sync, pass them to the @@ -101,7 +142,7 @@ ID. Excluded workspaces and their roles are skipped entirely. ## Group provisioning -Account groups are provisioned through the OAuth client ID and secret flow. +Account groups are provisioned as the authenticated service principal. [Here](https://docs.databricks.com/aws/en/admin/users-groups/groups#:~:text=Types%20of%20groups%20in%20Databricks,permissions%20to%20identity%20federated%20workspaces.) are the different types of groups in Databricks. @@ -156,6 +197,7 @@ Flags: --databricks-client-id string required: The Databricks service principal's client ID used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_ID) --databricks-client-secret string required: The Databricks service principal's client secret used to connect to the Databricks Account and Workspace API ($BATON_DATABRICKS_CLIENT_SECRET) --databricks-exclude-workspaces strings Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces. ($BATON_DATABRICKS_EXCLUDE_WORKSPACES) + --databricks-oidc-token-file string required: Path to a file holding an externally issued JWT identifying the workload. The file is read again on every token refresh, so an external rotation is picked up without a restart. ($BATON_DATABRICKS_OIDC_TOKEN_FILE) --external-resource-c1z string The path to the c1z file to sync external baton resources with ($BATON_EXTERNAL_RESOURCE_C1Z) --external-resource-entitlement-id-filter string The entitlement that external users, groups must have access to sync external baton resources ($BATON_EXTERNAL_RESOURCE_ENTITLEMENT_ID_FILTER) --external-resource-traits strings Resource type traits (e.g. "user", "group", "app") to sync and match from the external resource c1z. When unset the matcher falls back to user and group; passing this flag replaces the full set rather than adding to it. ($BATON_EXTERNAL_RESOURCE_TRAITS) diff --git a/config_schema.json b/config_schema.json index e387c1ed..8a5c9fa7 100644 --- a/config_schema.json +++ b/config_schema.json @@ -164,5 +164,23 @@ ], "displayName": "Databricks", "helpUrl": "/docs/baton/databricks", - "iconUrl": "/static/app-icons/databricks.svg" + "iconUrl": "/static/app-icons/databricks.svg", + "fieldGroups": [ + { + "name": "client-secret", + "displayName": "Client Secret", + "helpText": "Authenticate with the Databricks service principal's OAuth client ID and client secret.", + "fields": [ + "account-hostname", + "account-id", + "databricks-client-id", + "hostname", + "workspaces", + "base-url", + "databricks-exclude-workspaces", + "databricks-client-secret" + ], + "default": true + } + ] } \ No newline at end of file diff --git a/docs/connector.mdx b/docs/connector.mdx index f57273e3..7de21995 100644 --- a/docs/connector.mdx +++ b/docs/connector.mdx @@ -29,7 +29,16 @@ Revoking the workspace **member** entitlement removes the principal's entire wor ## Authentication methods -The connector authenticates with **OAuth** — an account-level service principal's client ID and secret. This is the only authentication method the connector supports. +The connector authenticates as an account-level service principal, named by its client ID. It proves it is that service principal in one of two ways: + +- **OAuth client secret** (default): the service principal's client ID and secret. +- **Workload identity federation**: the connector presents an externally issued JWT that identifies the workload, such as a SPIFFE JWT-SVID or a Kubernetes projected service account token, and Databricks exchanges it for a short-lived token. No long-lived secret is stored, and the Databricks administrator controls access through a service principal federation policy. This fits a deployment where the platform running the connector already issues it a workload identity token. + +The authentication method selects the credential. There is no personal access token (PAT), workspace token, or username and password option. + + +Workload identity federation is available when you run the connector yourself, on the command line or in service mode. The connector reads the JWT from a file, and a ConductorOne-hosted connector has no file system you control, so the option is not offered there. Select it with `--auth-method workload-identity-file`. + ## Gather Databricks credentials @@ -55,7 +64,7 @@ A user with the **Account admin** role in each Databricks workspace you want to ### Generate Databricks credentials -The Databricks connector authenticates with OAuth, which syncs info from all Databricks workspaces the service principal can access. +The Databricks connector authenticates as a service principal, which syncs info from all Databricks workspaces that service principal can access. @@ -72,6 +81,24 @@ The Databricks connector authenticates with OAuth, which syncs info from all Dat - OAuth client ID - OAuth client secret +### Use workload identity federation instead of a secret + +You can skip the OAuth secret and have the connector present an externally issued JWT instead. The connector must be able to read that JWT where it runs, so this suits a deployment whose platform injects a workload identity token. + + + + Create the service principal as described above, but do not generate an OAuth secret. + + + Follow the [Databricks OAuth token federation documentation](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-exchange) to add a service principal federation policy that trusts the issuer, subject and audience of the JWT your workload receives. + + + Make the JWT available to the connector, ideally as a file that your platform rotates. + + + +**Done.** In this case the credentials are the account ID, the OAuth client ID, and the JWT, in place of the OAuth client secret. + Next, move on to the instructions for your chosen setup method. ## Configure the Databricks connector @@ -204,6 +231,18 @@ stringData: BATON_DATABRICKS_CLIENT_ID: BATON_DATABRICKS_CLIENT_SECRET: + # Alternative to BATON_DATABRICKS_CLIENT_SECRET: workload identity federation. + # Both variables below are required together. Without BATON_AUTH_METHOD the + # connector validates against the client secret method and fails on the + # missing secret. + # Point the file at a projected service account token, or any other externally + # issued JWT your federation policy trusts. It is read again on every token + # refresh, so a rotated token needs no restart. + # Nothing creates this file for you: uncomment the matching volume and + # volumeMount in the Deployment below, or mount the JWT there yourself. + # BATON_AUTH_METHOD: workload-identity-file + # BATON_DATABRICKS_OIDC_TOKEN_FILE: /var/run/secrets/databricks/token + # Optional: limit the sync to specific workspaces # (workspace name, deployment name, or numeric ID). # Mutually exclusive with BATON_DATABRICKS_EXCLUDE_WORKSPACES — set one or the other, never both. @@ -223,7 +262,7 @@ stringData: -OAuth requires a reachable account API. If the account API check fails at startup, the connector fails validation instead of falling back to a workspace-only sync, even when `BATON_WORKSPACES` is set. +Both credential options require a reachable account API. If the account API check fails at startup, the connector fails validation instead of falling back to a workspace-only sync, even when `BATON_WORKSPACES` is set. See the connector's README or run `--help` to see all available configuration flags and environment variables. @@ -259,6 +298,21 @@ spec: envFrom: - secretRef: name: baton-databricks-secrets + # Only for BATON_DATABRICKS_OIDC_TOKEN_FILE: mount the JWT at the path the + # secret points to. Leave this out when you use the client secret. + # volumeMounts: + # - name: databricks-oidc-token + # mountPath: /var/run/secrets/databricks + # readOnly: true + # volumes: + # - name: databricks-oidc-token + # projected: + # sources: + # - serviceAccountToken: + # path: token + # # Must match the audience of the Databricks federation policy. + # audience: + # expirationSeconds: 3600 ``` ### Step 3: Deploy the connector diff --git a/docs/docs-info.md b/docs/docs-info.md index fcebe05c..d0f198d1 100644 --- a/docs/docs-info.md +++ b/docs/docs-info.md @@ -17,28 +17,31 @@ While developing the connector, please fill out this form. This information is n 1. What credentials or information are needed to set up the connector? (For example, API key, client ID and secret, domain, etc.) - > The connector requires a Databricks account ID and OAuth credentials: + > The connector requires a Databricks account ID and a service principal OAuth client ID. It syncs the account and every workspace that service principal can access. The service principal is then authenticated in one of two ways: > - > - **OAuth**: a service principal OAuth client ID and client secret. Syncs the account and every workspace the service principal can access. This is the connector's only authentication method. + > - **OAuth client secret** (default): the service principal's OAuth client secret. + > - **Workload identity federation**: an externally issued JWT identifying the workload, such as a SPIFFE JWT-SVID or a Kubernetes projected service account token, read from a file. The connector exchanges it for a short-lived Databricks token, so no long-lived secret is stored. This requires a service principal federation policy in Databricks. It is available only when you run the connector yourself, on the command line or in service mode, and is selected with `--auth-method workload-identity-file`. + > + > The authentication method selects the credential. > > Google Cloud Platform and Azure Databricks customers also provide the account hostname and hostname. 2. For each item in the list above: - - * How does a user create or look up that credential or info? Please include links to (non-gated) documentation, screenshots (of the UI or of gated docs), or a video of the process. + - How does a user create or look up that credential or info? Please include links to (non-gated) documentation, screenshots (of the UI or of gated docs), or a video of the process. > - **Account ID**: in the Databricks account console, open the menu next to your username in the upper-right corner; the account ID is shown there. > - **OAuth client ID and secret**: follow the [Databricks OAuth (M2M) documentation](https://docs.databricks.com/en/dev-tools/auth/oauth-m2m.html) to create a service principal and generate an OAuth secret. + > - **Workload identity federation**: follow the [Databricks OAuth token federation documentation](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-exchange) to add a federation policy to that same service principal. The JWT itself comes from the workload's own platform, not from Databricks. > - **Deployment name**: the subdomain in the workspace URL (not the numeric workspace ID). - * Does the credential need any specific scopes or permissions? If so, list them here. + - Does the credential need any specific scopes or permissions? If so, list them here. > The credential must have admin access to each resource it reads or writes: account-admin on the Databricks account (for account-level sync and provisioning) and workspace-admin on each workspace being synced. - * If applicable: Is the list of scopes or permissions different to sync (read) versus provision (read-write)? If so, list the difference here. + - If applicable: Is the list of scopes or permissions different to sync (read) versus provision (read-write)? If so, list the difference here. > No separate scopes: Databricks admin access covers both read (sync) and read-write (provision). - * What level of access or permissions does the user need in order to create the credentials? (For example, must be a super administrator, must have access to the admin console, etc.) + - What level of access or permissions does the user need in order to create the credentials? (For example, must be a super administrator, must have access to the admin console, etc.) > Account admin access to the Databricks account console (to create the service principal and OAuth secret), and workspace admin on each workspace being synced. diff --git a/pkg/config/conf.gen.go b/pkg/config/conf.gen.go index dd2a3bba..5550c7e5 100644 --- a/pkg/config/conf.gen.go +++ b/pkg/config/conf.gen.go @@ -8,6 +8,7 @@ type Databricks struct { AccountId string `mapstructure:"account-id"` DatabricksClientId string `mapstructure:"databricks-client-id"` DatabricksClientSecret string `mapstructure:"databricks-client-secret"` + DatabricksOidcTokenFile string `mapstructure:"databricks-oidc-token-file"` Hostname string `mapstructure:"hostname"` Workspaces []string `mapstructure:"workspaces"` BaseUrl string `mapstructure:"base-url"` diff --git a/pkg/config/config.go b/pkg/config/config.go index 35afe969..1e0aa11d 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -4,6 +4,13 @@ import ( "github.com/conductorone/baton-sdk/pkg/field" ) +// Auth method names. The connector branches on the selected one, so the group +// definitions and that branch read the same constant. +const ( + ClientSecretAuthMethod = "client-secret" + WorkloadIdentityFileAuthMethod = "workload-identity-file" +) + var ( AccountIdField = field.StringField( "account-id", @@ -24,6 +31,14 @@ var ( field.WithRequired(true), field.WithDisplayName("OAuth2 Client Secret"), ) + DatabricksOIDCTokenFileField = field.StringField( + "databricks-oidc-token-file", + field.WithDescription("Path to a file holding an externally issued JWT identifying the workload. "+ + "The file is read again on every token refresh, so an external rotation is picked up without a restart."), + field.WithDisplayName("Workload Identity Token File"), + field.WithRequired(true), + field.WithExportTarget(field.ExportTargetCLIOnly), + ) WorkspacesField = field.StringSliceField( "workspaces", field.WithDescription( @@ -54,11 +69,46 @@ var ( field.WithDescription("Workspaces to exclude from sync, identified by workspace name, deployment name, or numeric workspace ID. Mutually exclusive with workspaces."), field.WithDisplayName("Exclude Workspaces"), ) + // commonFields apply whichever way the connector authenticates. Validation + // skips any field the selected group does not list, so every group repeats + // them: account-id is required and would stop being checked otherwise. + commonFields = []field.SchemaField{ + AccountHostnameField, + AccountIdField, + DatabricksClientIdField, + HostnameField, + WorkspacesField, + BaseURLField, + ExcludeWorkspacesField, + } + + // The client secret is the default. The token file exchanges an externally + // issued JWT, and is read again on every exchange, so a rotated file keeps + // working without a restart. + fieldGroups = []field.SchemaFieldGroup{ + { + Name: ClientSecretAuthMethod, + DisplayName: "Client Secret", + HelpText: "Authenticate with the Databricks service principal's OAuth client ID and client secret.", + Default: true, + Fields: append(append([]field.SchemaField{}, commonFields...), DatabricksClientSecretField), + }, + { + Name: WorkloadIdentityFileAuthMethod, + DisplayName: "Workload Identity (token file)", + HelpText: "Exchange an externally issued JWT for a Databricks token. " + + "The JWT is read from a file on every exchange, so whatever writes that file controls the lifetime.", + ExportTarget: field.ExportTargetCLIOnly, + Fields: append(append([]field.SchemaField{}, commonFields...), DatabricksOIDCTokenFileField), + }, + } + configFields = []field.SchemaField{ AccountHostnameField, AccountIdField, DatabricksClientIdField, DatabricksClientSecretField, + DatabricksOIDCTokenFileField, HostnameField, WorkspacesField, BaseURLField, @@ -72,6 +122,7 @@ var Config = field.NewConfiguration( field.WithConnectorDisplayName("Databricks"), field.WithHelpUrl("/docs/baton/databricks"), field.WithIconUrl("/static/app-icons/databricks.svg"), + field.WithFieldGroups(fieldGroups), field.WithConstraints( field.FieldsMutuallyExclusive(WorkspacesField, ExcludeWorkspacesField), ), diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go new file mode 100644 index 00000000..9cbc4969 --- /dev/null +++ b/pkg/config/config_test.go @@ -0,0 +1,143 @@ +package config + +import ( + "strings" + "testing" + + "github.com/conductorone/baton-sdk/pkg/field" +) + +// Validation skips every field the selected group does not list, so a group +// that forgets a required field silently stops checking it. +func TestEveryAuthGroupCarriesTheRequiredCommonFields(t *testing.T) { + for _, group := range Config.FieldGroups { + fields := group.FieldMap() + for _, required := range []field.SchemaField{AccountIdField, DatabricksClientIdField} { + if _, ok := fields[required.FieldName]; !ok { + t.Errorf("group %q omits %q, so it would no longer be validated", group.Name, required.FieldName) + } + } + } +} + +func TestAuthGroupsOfferOneCredentialEach(t *testing.T) { + var defaults []string + for _, group := range Config.FieldGroups { + if group.Default { + defaults = append(defaults, group.Name) + } + } + if len(defaults) != 1 || defaults[0] != "client-secret" { + t.Errorf("default groups = %v, want exactly [client-secret]", defaults) + } + + for _, tt := range []struct { + group, holds, omits string + }{ + {"client-secret", DatabricksClientSecretField.FieldName, DatabricksOIDCTokenFileField.FieldName}, + {"workload-identity-file", DatabricksOIDCTokenFileField.FieldName, DatabricksClientSecretField.FieldName}, + } { + fields := Config.FieldGroupFields(tt.group) + if _, ok := fields[tt.holds]; !ok { + t.Errorf("group %q does not offer %q", tt.group, tt.holds) + } + if _, ok := fields[tt.omits]; ok { + t.Errorf("group %q also offers %q, so the two credentials are not separated", tt.group, tt.omits) + } + } +} + +// A hosted connector has no file system the tenant admin controls, so the token +// file must never reach the web form. +func TestTheTokenFileIsNotOfferedToHostedTenants(t *testing.T) { + if DatabricksOIDCTokenFileField.ExportTarget != field.ExportTargetCLIOnly { + t.Errorf("databricks-oidc-token-file ExportTarget = %q, want %q", + DatabricksOIDCTokenFileField.ExportTarget, field.ExportTargetCLIOnly) + } + for _, group := range Config.FieldGroups { + if group.Name != "workload-identity-file" { + continue + } + if group.ExportTarget != field.ExportTargetCLIOnly { + t.Errorf("group %q ExportTarget = %q, want %q", group.Name, group.ExportTarget, field.ExportTargetCLIOnly) + } + } +} + +func TestCredentialValidation(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + authMethod string + cfg Databricks + wantErr string + }{ + { + name: "client secret under its own method", + authMethod: "client-secret", + cfg: Databricks{DatabricksClientSecret: "secret"}, + }, + { + name: "client secret missing under its own method", + authMethod: "client-secret", + cfg: Databricks{}, + wantErr: "databricks-client-secret", + }, + { + name: "token file under its own method", + authMethod: "workload-identity-file", + cfg: Databricks{DatabricksOidcTokenFile: "/var/run/secrets/token"}, //nolint:gosec // G101: a test fixture path, not a credential. + }, + { + name: "token file missing under its own method", + authMethod: "workload-identity-file", + cfg: Databricks{}, + wantErr: "databricks-oidc-token-file", + }, + { + // The secret belongs to the other group, so it is not asked for here. + name: "token file alone does not require the secret", + authMethod: "workload-identity-file", + cfg: Databricks{DatabricksOidcTokenFile: "/var/run/secrets/token"}, //nolint:gosec // G101: a test fixture path, not a credential. + }, + { + name: "no method selected falls back to the default group", + cfg: Databricks{}, + wantErr: "databricks-client-secret", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + cfg := tt.cfg + cfg.AccountId = "acc-1" + cfg.DatabricksClientId = "sp-client-id" + + err := field.Validate(Config, &cfg, field.WithAuthMethod(tt.authMethod)) + + switch { + case tt.wantErr == "" && err != nil: + t.Errorf("Validate returned %v, want no error", err) + case tt.wantErr != "" && err == nil: + t.Errorf("Validate succeeded, want an error naming %q", tt.wantErr) + case tt.wantErr != "" && !strings.Contains(err.Error(), tt.wantErr): + t.Errorf("Validate returned %q, want it to name %q", err, tt.wantErr) + } + }) + } +} + +func TestCredentialFieldRequirements(t *testing.T) { + if !DatabricksClientSecretField.Required { + t.Error("databricks-client-secret must be required: its group offers no other credential") + } + if !DatabricksOIDCTokenFileField.Required { + t.Error("databricks-oidc-token-file must be required: its group offers no other credential") + } + if !DatabricksClientIdField.Required { + t.Error("databricks-client-id must stay required: federation is per service principal") + } +} diff --git a/pkg/connector/auth_selection_test.go b/pkg/connector/auth_selection_test.go new file mode 100644 index 00000000..dc8bfd25 --- /dev/null +++ b/pkg/connector/auth_selection_test.go @@ -0,0 +1,69 @@ +package connector + +import ( + "testing" + + "github.com/conductorone/baton-databricks/pkg/config" + "github.com/conductorone/baton-databricks/pkg/databricks" +) + +// The selected auth method decides the credential. A field left over from +// another deployment must not change it. +func TestPrepareClientAuthSelectsTheCredential(t *testing.T) { + t.Parallel() + + const bothSet = "both fields set" + + tests := []struct { + name string + authMethod string + cfg config.Databricks + federated bool + }{ + { + name: "client secret", + authMethod: config.ClientSecretAuthMethod, + cfg: config.Databricks{DatabricksClientSecret: "secret"}, + }, + { + name: "token file", + authMethod: config.WorkloadIdentityFileAuthMethod, + cfg: config.Databricks{DatabricksOidcTokenFile: "/var/run/secrets/token"}, //nolint:gosec // G101: a fixture path, not a credential. + federated: true, + }, + { + name: bothSet + " under the client secret method", + authMethod: config.ClientSecretAuthMethod, + cfg: config.Databricks{DatabricksClientSecret: "secret", DatabricksOidcTokenFile: "/leftover/token"}, + }, + { + name: bothSet + " under the token file method", + authMethod: config.WorkloadIdentityFileAuthMethod, + cfg: config.Databricks{DatabricksClientSecret: "secret", DatabricksOidcTokenFile: "/var/run/secrets/token"}, //nolint:gosec // G101: a fixture path, not a credential. + federated: true, + }, + { + // No method selected falls back to the default group, which + // validation has already checked for a client secret. + name: "no method selected", + cfg: config.Databricks{DatabricksClientSecret: "secret", DatabricksOidcTokenFile: "/leftover/token"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + cfg := tt.cfg + cfg.AccountId = "acc-1" + cfg.DatabricksClientId = "sp-client-id" + cfg.Hostname = "cloud.databricks.com" + + auth := prepareClientAuth(&cfg, tt.authMethod) + _, federated := auth.(*databricks.TokenFederation) + if federated != tt.federated { + t.Errorf("prepareClientAuth returned %T, federated = %v, want %v", auth, federated, tt.federated) + } + }) + } +} diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index f97021e6..4aea4a9e 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -11,6 +11,7 @@ import ( "github.com/conductorone/baton-sdk/pkg/annotations" "github.com/conductorone/baton-sdk/pkg/cli" "github.com/conductorone/baton-sdk/pkg/connectorbuilder" + "github.com/conductorone/baton-sdk/pkg/field" ) type Databricks struct { @@ -147,9 +148,19 @@ func New( } // NewConnector returns a new connector builder from a configuration struct. -func NewConnector(ctx context.Context, cfg *config.Databricks, _ *cli.ConnectorOpts) (connectorbuilder.ConnectorBuilderV2, []connectorbuilder.Opt, error) { +func NewConnector(ctx context.Context, cfg *config.Databricks, opts *cli.ConnectorOpts) (connectorbuilder.ConnectorBuilderV2, []connectorbuilder.Opt, error) { + // The auth method decides which fields are required. Validating without it + // would check the default group against a configuration built for another. + authMethod := "" + if opts != nil { + authMethod = opts.SelectedAuthMethod + } + if err := field.Validate(config.Config, cfg, field.WithAuthMethod(authMethod)); err != nil { + return nil, nil, err + } + accountHostname := getAccountHostname(cfg, cfg.Hostname) - auth := prepareClientAuth(cfg) + auth := prepareClientAuth(cfg, authMethod) cb, err := New( ctx, @@ -168,16 +179,27 @@ func NewConnector(ctx context.Context, cfg *config.Databricks, _ *cli.ConnectorO return cb, nil, nil } -func prepareClientAuth(cfg *config.Databricks) databricks.Auth { +// prepareClientAuth follows the selected auth method, not whichever field +// happens to be set: a token file left over from another deployment must not +// silently replace the credential the operator chose. +func prepareClientAuth(cfg *config.Databricks, authMethod string) databricks.Auth { accountID := cfg.AccountId databricksClientId := cfg.DatabricksClientId - databricksClientSecret := cfg.DatabricksClientSecret accountHostname := getAccountHostname(cfg, cfg.Hostname) + if authMethod == config.WorkloadIdentityFileAuthMethod { + return databricks.NewTokenFederation( + accountID, + databricksClientId, + accountHostname, + databricks.FileSubjectTokenSource{Path: cfg.DatabricksOidcTokenFile}, + ) + } + return databricks.NewOAuth2( accountID, databricksClientId, - databricksClientSecret, + cfg.DatabricksClientSecret, accountHostname, ) } diff --git a/pkg/databricks/auth.go b/pkg/databricks/auth.go index 63cda5bc..8d4105b9 100644 --- a/pkg/databricks/auth.go +++ b/pkg/databricks/auth.go @@ -2,15 +2,52 @@ package databricks import ( "context" + "encoding/base64" + "encoding/json" "fmt" "net/http" + "net/url" + "os" + "strings" + "time" "github.com/conductorone/baton-sdk/pkg/uhttp" "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" + "google.golang.org/grpc/codes" ) +const ( + scopeAllAPIs = "all-apis" + + tokenPathOIDC = "oidc" + tokenPathAccounts = "accounts" + tokenPathVersion = "v1" + tokenPathToken = "token" + + //nolint:gosec // G101: these are the public RFC 8693 URNs, not credentials. + grantTypeTokenExchange = "urn:ietf:params:oauth:grant-type:token-exchange" + //nolint:gosec // G101: these are the public RFC 8693 URNs, not credentials. + subjectTokenTypeJWT = "urn:ietf:params:oauth:token-type:jwt" + + tokenExchangeTimeout = 30 * time.Second + + fallbackTokenLifetime = 5 * time.Minute +) + +type tokenExchangeError struct { + Code string `json:"error"` + Description string `json:"error_description"` + RequestID string `json:"request_id"` +} + +type tokenExchangeResponse struct { + AccessToken string `json:"access_token"` + TokenType string `json:"token_type"` + ExpiresIn int64 `json:"expires_in"` +} + type Auth interface { Apply(req *http.Request) GetClient(ctx context.Context) (*http.Client, error) @@ -18,6 +55,115 @@ type Auth interface { type NoAuth struct{} +type OAuth2 struct { + cfg *clientcredentials.Config +} + +// SubjectTokenSource yields the externally issued JWT that identifies the +// workload. Token() asks for one on every exchange, so a source that cannot +// produce a fresh token does not belong behind this interface. +type SubjectTokenSource interface { + SubjectToken() (string, error) +} + +// FileSubjectTokenSource reads the JWT from disk on every call, so whatever +// rotates the file controls the lifetime and the connector needs no restart. +type FileSubjectTokenSource struct { + Path string +} + +func (f FileSubjectTokenSource) SubjectToken() (string, error) { + raw, err := os.ReadFile(f.Path) + if err != nil { + return "", uhttp.WrapErrors(codes.FailedPrecondition, fmt.Sprintf("failed to read workload identity token file %q", f.Path), err) + } + + token := strings.TrimSpace(string(raw)) + if token == "" { + return "", uhttp.WrapErrors(codes.FailedPrecondition, fmt.Sprintf("workload identity token file %q is empty", f.Path)) + } + + return token, nil +} + +type TokenFederation struct { + clientID string + tokenURL string + source SubjectTokenSource +} + +type federationTokenSource struct { + ctx context.Context + auth *TokenFederation + httpClient *uhttp.BaseHttpClient +} + +func (e *tokenExchangeError) Message() string { + msg := e.Description + if msg == "" { + msg = e.Code + } + if e.RequestID != "" { + msg = fmt.Sprintf("%s (request id %s)", msg, e.RequestID) + } + + return msg +} + +// accountTokenURL is the account-level token endpoint, shared by both auth methods. +func accountTokenURL(accountHostname, accountID string) string { + u := &url.URL{Scheme: "https", Host: accountHostname} + + return u.JoinPath(tokenPathOIDC, tokenPathAccounts, accountID, tokenPathVersion, tokenPathToken).String() +} + +// subjectToken returns the JWT. A file is read on every call, so a rotated token is +// picked up without restarting the connector. +func (t *TokenFederation) subjectToken() (string, error) { + if t.source == nil { + return "", uhttp.WrapErrors(codes.FailedPrecondition, "no workload identity token source configured") + } + + return t.source.SubjectToken() +} + +// tokenExpiry never returns the zero time, which oauth2 reads as a token that +// never expires. +func tokenExpiry(now time.Time, expiresIn int64, accessToken string) time.Time { + if expiresIn > 0 { + return now.Add(time.Duration(expiresIn) * time.Second) + } + + if exp, ok := jwtExpiry(accessToken); ok && exp.After(now) { + return exp + } + + return now.Add(fallbackTokenLifetime) +} + +// jwtExpiry reads the exp claim without verifying the signature. The claim only +// decides when to exchange again. +func jwtExpiry(token string) (time.Time, bool) { + parts := strings.Split(token, ".") + if len(parts) != 3 { + return time.Time{}, false + } + + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return time.Time{}, false + } + + var claims struct { + Exp int64 `json:"exp"` + } + if err := json.Unmarshal(payload, &claims); err != nil || claims.Exp <= 0 { + return time.Time{}, false + } + + return time.Unix(claims.Exp, 0), true +} + func (n *NoAuth) Apply(req *http.Request) {} func (n *NoAuth) GetClient(ctx context.Context) (*http.Client, error) { @@ -29,17 +175,13 @@ func (n *NoAuth) GetClient(ctx context.Context) (*http.Client, error) { return httpClient, nil } -type OAuth2 struct { - cfg *clientcredentials.Config -} - func NewOAuth2(accId, clientId, clientSecret, accountHostname string) *OAuth2 { return &OAuth2{ cfg: &clientcredentials.Config{ ClientID: clientId, ClientSecret: clientSecret, - TokenURL: fmt.Sprintf("https://%s/oidc/accounts/%s/v1/token", accountHostname, accId), - Scopes: []string{"all-apis"}, + TokenURL: accountTokenURL(accountHostname, accId), + Scopes: []string{scopeAllAPIs}, }, } } @@ -54,3 +196,93 @@ func (o *OAuth2) GetClient(ctx context.Context) (*http.Client, error) { func (o *OAuth2) Apply(req *http.Request) { // No need to set the Authorization header here, the oauth2 client does it automatically } + +// NewTokenFederation builds the federation auth. Exactly one of token or tokenFile +// carries the JWT. +func NewTokenFederation(accID, clientID, accountHostname string, source SubjectTokenSource) *TokenFederation { + return &TokenFederation{ + clientID: clientID, + tokenURL: accountTokenURL(accountHostname, accID), + source: source, + } +} + +func (t *TokenFederation) GetClient(ctx context.Context) (*http.Client, error) { + httpClient, err := uhttp.NewClient(ctx, uhttp.WithLogger(true, ctxzap.Extract(ctx))) + if err != nil { + return nil, err + } + + exchangeClient, err := uhttp.NewBaseHttpClientWithContext(ctx, httpClient) + if err != nil { + return nil, err + } + + source := &federationTokenSource{ctx: ctx, auth: t, httpClient: exchangeClient} + + return oauth2.NewClient(ctx, oauth2.ReuseTokenSource(nil, source)), nil +} + +func (t *TokenFederation) Apply(req *http.Request) {} + +// Token exchanges the JWT for a short-lived Databricks OAuth token. +// POST https://{accountHostname}/oidc/accounts/{accountId}/v1/token, RFC 8693. +// It needs a service principal federation policy trusting the JWT's issuer, subject and audience. +// The context is the one the client was built with, because oauth2.TokenSource takes none. +func (s *federationTokenSource) Token() (*oauth2.Token, error) { + subjectToken, err := s.auth.subjectToken() + if err != nil { + return nil, err + } + + form := url.Values{ + "grant_type": {grantTypeTokenExchange}, + "subject_token": {subjectToken}, + "subject_token_type": {subjectTokenTypeJWT}, + "client_id": {s.auth.clientID}, + "scope": {scopeAllAPIs}, + } + + ctx, cancel := context.WithTimeout(s.ctx, tokenExchangeTimeout) + defer cancel() + + endpoint, err := url.Parse(s.auth.tokenURL) + if err != nil { + return nil, fmt.Errorf("failed to parse the token endpoint: %w", err) + } + + req, err := s.httpClient.NewRequest(ctx, http.MethodPost, endpoint, + uhttp.WithFormBody(form.Encode()), + uhttp.WithAcceptJSONHeader(), + ) + if err != nil { + return nil, fmt.Errorf("failed to build the token exchange request: %w", err) + } + + var exchanged tokenExchangeResponse + var exchangeErr tokenExchangeError + + resp, err := s.httpClient.Do(req, + uhttp.WithJSONResponse(&exchanged), + uhttp.WithErrorResponse(&exchangeErr), + ) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if exchanged.AccessToken == "" { + return nil, uhttp.WrapErrors(codes.FailedPrecondition, "token exchange returned no access token") + } + + tokenType := exchanged.TokenType + if tokenType == "" { + tokenType = "Bearer" + } + + return &oauth2.Token{ + AccessToken: exchanged.AccessToken, + TokenType: tokenType, + Expiry: tokenExpiry(time.Now(), exchanged.ExpiresIn, exchanged.AccessToken), + }, nil +} diff --git a/pkg/databricks/auth_test.go b/pkg/databricks/auth_test.go new file mode 100644 index 00000000..5f0cdb0e --- /dev/null +++ b/pkg/databricks/auth_test.go @@ -0,0 +1,385 @@ +package databricks + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/conductorone/baton-sdk/pkg/uhttp" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// exchangeServer serves the token endpoint and records the last form it received. +func exchangeServer(t *testing.T, handler func(form url.Values) (int, string)) (*httptest.Server, *url.Values) { + t.Helper() + + received := &url.Values{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err != nil { + t.Errorf("ParseForm: %v", err) + w.WriteHeader(http.StatusBadRequest) + return + } + *received = r.PostForm + + status, body := handler(r.PostForm) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if _, err := w.Write([]byte(body)); err != nil { + t.Errorf("write response: %v", err) + } + })) + t.Cleanup(server.Close) + + return server, received +} + +// stubSubjectTokenSource stands in for a real source so the exchange itself can +// be tested without a file. Production offers no static source: a fixed JWT +// cannot be refreshed. +type stubSubjectTokenSource string + +func (s stubSubjectTokenSource) SubjectToken() (string, error) { return string(s), nil } + +func newTestTokenSource(t *testing.T, tokenURL string, source SubjectTokenSource) *federationTokenSource { + t.Helper() + + ctx := context.Background() + httpClient, err := uhttp.NewBaseHttpClientWithContext(ctx, &http.Client{}) + if err != nil { + t.Fatalf("NewBaseHttpClientWithContext: %v", err) + } + + return &federationTokenSource{ + ctx: ctx, + auth: &TokenFederation{clientID: "sp-client-id", tokenURL: tokenURL, source: source}, + httpClient: httpClient, + } +} + +func writeTokenFile(t *testing.T, contents string) string { + t.Helper() + + path := filepath.Join(t.TempDir(), "token.jwt") + if err := os.WriteFile(path, []byte(contents), 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + return path +} + +func TestTokenFederationExchange(t *testing.T) { + server, received := exchangeServer(t, func(url.Values) (int, string) { + return http.StatusOK, `{"access_token":"databricks-token","token_type":"Bearer","expires_in":3600}` + }) + + source := newTestTokenSource(t, server.URL, stubSubjectTokenSource("subject-jwt")) + + token, err := source.Token() + if err != nil { + t.Fatalf("Token: %v", err) + } + + if token.AccessToken != "databricks-token" { + t.Errorf("AccessToken = %q, want %q", token.AccessToken, "databricks-token") + } + if token.TokenType != "Bearer" { + t.Errorf("TokenType = %q, want %q", token.TokenType, "Bearer") + } + // A zero Expiry passes token.Valid(), so assert the value itself. + if d := time.Until(token.Expiry); d < 55*time.Minute || d > time.Hour { + t.Errorf("token expires in %v, want about an hour from now", d) + } + + want := map[string]string{ + "grant_type": grantTypeTokenExchange, + "subject_token": "subject-jwt", + "subject_token_type": subjectTokenTypeJWT, + "client_id": "sp-client-id", + "scope": scopeAllAPIs, + } + for key, value := range want { + if got := received.Get(key); got != value { + t.Errorf("form %q = %q, want %q", key, got, value) + } + } +} + +// A rotated token file must be picked up without restarting the connector. +func TestTokenFederationRereadsTokenFile(t *testing.T) { + server, received := exchangeServer(t, func(url.Values) (int, string) { + return http.StatusOK, `{"access_token":"databricks-token","expires_in":3600}` + }) + + path := writeTokenFile(t, "first-jwt\n") + source := newTestTokenSource(t, server.URL, FileSubjectTokenSource{Path: path}) + + if _, err := source.Token(); err != nil { + t.Fatalf("first Token: %v", err) + } + if got := received.Get("subject_token"); got != "first-jwt" { + t.Fatalf("first subject_token = %q, want %q (trailing newline must be trimmed)", got, "first-jwt") + } + + if err := os.WriteFile(path, []byte("rotated-jwt"), 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + if _, err := source.Token(); err != nil { + t.Fatalf("second Token: %v", err) + } + if got := received.Get("subject_token"); got != "rotated-jwt" { + t.Errorf("second subject_token = %q, want %q", got, "rotated-jwt") + } +} + +func TestTokenFederationExchangeErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + status int + body string + wantErr string + }{ + { + name: "rejected by the federation policy", + status: http.StatusUnauthorized, + body: `{"error":"invalid_request","error_description":"no matching federation policy"}`, + wantErr: "no matching federation policy", + }, + { + name: "response without an access token", + status: http.StatusOK, + body: `{"token_type":"Bearer","expires_in":3600}`, + wantErr: "no access token", + }, + { + name: "response that is not JSON", + status: http.StatusOK, + body: `gateway`, + wantErr: "unmarshal json response", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + server, _ := exchangeServer(t, func(url.Values) (int, string) { + return tt.status, tt.body + }) + + _, err := newTestTokenSource(t, server.URL, stubSubjectTokenSource("subject-jwt")).Token() + if err == nil { + t.Fatal("Token succeeded, want an error") + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Errorf("error = %q, want it to contain %q", err, tt.wantErr) + } + }) + } +} + +// A failed exchange must reach the syncer with a gRPC code. An unclassified error +// reads as Unknown, so a throttled token endpoint would end the sync instead of retrying. +func TestTokenFederationExchangeErrorsAreClassified(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + status int + want codes.Code + }{ + {"rejected by the federation policy", http.StatusBadRequest, codes.InvalidArgument}, + {"bad credentials", http.StatusUnauthorized, codes.Unauthenticated}, + {"forbidden", http.StatusForbidden, codes.PermissionDenied}, + {"throttled stays retryable", http.StatusTooManyRequests, codes.Unavailable}, + {"server error stays retryable", http.StatusServiceUnavailable, codes.Unavailable}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + server, _ := exchangeServer(t, func(url.Values) (int, string) { + return tt.status, `{"error":"invalid_grant","error_description":"no matching federation policy"}` + }) + + _, err := newTestTokenSource(t, server.URL, stubSubjectTokenSource("subject-jwt")).Token() + if err == nil { + t.Fatal("Token succeeded, want an error") + } + + if got := status.Code(err); got != tt.want { + t.Errorf("grpc code = %v, want %v", got, tt.want) + } + + // wrapTransportAuthError runs on this error at the API call sites and must + // not downgrade the code. + if got := status.Code(wrapTransportAuthError(err)); got != tt.want { + t.Errorf("grpc code after wrapTransportAuthError = %v, want %v", got, tt.want) + } + + if !strings.Contains(err.Error(), "no matching federation policy") { + t.Errorf("error = %q, want it to keep the vendor's reason", err) + } + }) + } +} + +func TestTokenFederationSubjectToken(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + file string + wantErr string + }{ + {"missing file", filepath.Join(t.TempDir(), "absent.jwt"), "failed to read"}, + {"empty file", writeTokenFile(t, " \n"), "is empty"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + _, err := newTestTokenSource(t, "", FileSubjectTokenSource{Path: tt.file}).Token() + if err == nil { + t.Fatal("Token succeeded, want an error") + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Errorf("error = %q, want it to contain %q", err, tt.wantErr) + } + if got := status.Code(err); got != codes.FailedPrecondition { + t.Errorf("grpc code = %v, want %v", got, codes.FailedPrecondition) + } + }) + } +} + +// A zero Expiry reads as "never expires" in oauth2, pinning the first exchange. +func TestTokenExpiry(t *testing.T) { + t.Parallel() + + // A token that carries no exp claim to fall back to. + const notAJWT = "an-opaque-string" + + now := time.Date(2026, 10, 26, 12, 0, 0, 0, time.UTC) + + tests := []struct { + name string + expiresIn int64 + accessToken string + want time.Time + }{ + { + name: "expires_in wins", + expiresIn: 3600, + want: now.Add(time.Hour), + }, + { + name: "falls back to the exp claim", + accessToken: testJWT(t, now.Add(30*time.Minute)), + want: now.Add(30 * time.Minute), + }, + { + name: "falls back to a short lifetime for an opaque token", + accessToken: notAJWT, + want: now.Add(fallbackTokenLifetime), + }, + { + name: "falls back to a short lifetime when the exp claim is in the past", + accessToken: testJWT(t, now.Add(-time.Hour)), + want: now.Add(fallbackTokenLifetime), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + if got := tokenExpiry(now, tt.expiresIn, tt.accessToken); !got.Equal(tt.want) { + t.Errorf("tokenExpiry = %v, want %v", got, tt.want) + } + }) + } +} + +func testJWT(t *testing.T, expiry time.Time) string { + t.Helper() + + claims, err := json.Marshal(map[string]int64{"exp": expiry.Unix()}) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + + encode := base64.RawURLEncoding.EncodeToString + + return fmt.Sprintf("%s.%s.%s", encode([]byte(`{"alg":"RS256"}`)), encode(claims), encode([]byte("signature"))) +} + +func TestAccountTokenURL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + accountHostname string + accountID string + want string + }{ + { + name: "account id", + accountHostname: "accounts.cloud.databricks.com", + accountID: "8c6f99ec-78ce-4654-8f92-e716b3dd67a7", + want: "https://accounts.cloud.databricks.com/oidc/accounts/8c6f99ec-78ce-4654-8f92-e716b3dd67a7/v1/token", + }, + { + name: "azure account host", + accountHostname: "accounts.azuredatabricks.net", + accountID: "acc-1", + want: "https://accounts.azuredatabricks.net/oidc/accounts/acc-1/v1/token", + }, + { + // A '?' must be escaped, or it truncates the path before /v1/token. + name: "account id is escaped", + accountHostname: "accounts.cloud.databricks.com", + accountID: "acc?x=1", + want: "https://accounts.cloud.databricks.com/oidc/accounts/acc%3Fx=1/v1/token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + if got := accountTokenURL(tt.accountHostname, tt.accountID); got != tt.want { + t.Errorf("accountTokenURL = %q, want %q", got, tt.want) + } + }) + } +} + +// A federation built without a source cannot answer an exchange, and it must +// say so instead of reading an empty path. +func TestTokenFederationWithoutASourceIsRefused(t *testing.T) { + t.Parallel() + + _, err := (&TokenFederation{clientID: "sp-client-id"}).subjectToken() + if err == nil { + t.Fatal("subjectToken succeeded with no source, want an error") + } + if !strings.Contains(err.Error(), "no workload identity token source configured") { + t.Errorf("error = %q, want it to name the missing source", err) + } +}