From 7f52acb7c77b38d1a66c427d7766e76cd45ae980 Mon Sep 17 00:00:00 2001 From: Chen Shou Date: Thu, 20 Aug 2026 17:26:21 +0000 Subject: [PATCH] Add Docker credential token command --- .nextchanges/cli/docker-token.md | 1 + acceptance/cmd/auth/docker/help/out.test.toml | 2 + acceptance/cmd/auth/docker/help/output.txt | 58 +++ acceptance/cmd/auth/docker/help/script | 2 + acceptance/cmd/auth/docker/script.prepare | 27 ++ .../cmd/auth/docker/token/out.test.toml | 2 + acceptance/cmd/auth/docker/token/output.txt | 6 + acceptance/cmd/auth/docker/token/script | 6 + acceptance/cmd/auth/docker/token/test.toml | 7 + .../cmd/auth/docker/token/token-cache.json | 11 + cmd/auth/auth.go | 1 + cmd/auth/docker.go | 54 +++ cmd/auth/docker_profile.go | 33 ++ cmd/auth/docker_profile_test.go | 78 ++++ cmd/auth/docker_token.go | 97 +++++ cmd/auth/docker_token_test.go | 391 ++++++++++++++++++ libs/dockercredentials/registry.go | 103 +++++ libs/dockercredentials/registry_test.go | 89 ++++ 18 files changed, 968 insertions(+) create mode 100644 .nextchanges/cli/docker-token.md create mode 100644 acceptance/cmd/auth/docker/help/out.test.toml create mode 100644 acceptance/cmd/auth/docker/help/output.txt create mode 100644 acceptance/cmd/auth/docker/help/script create mode 100644 acceptance/cmd/auth/docker/script.prepare create mode 100644 acceptance/cmd/auth/docker/token/out.test.toml create mode 100644 acceptance/cmd/auth/docker/token/output.txt create mode 100644 acceptance/cmd/auth/docker/token/script create mode 100644 acceptance/cmd/auth/docker/token/test.toml create mode 100644 acceptance/cmd/auth/docker/token/token-cache.json create mode 100644 cmd/auth/docker.go create mode 100644 cmd/auth/docker_profile.go create mode 100644 cmd/auth/docker_profile_test.go create mode 100644 cmd/auth/docker_token.go create mode 100644 cmd/auth/docker_token_test.go create mode 100644 libs/dockercredentials/registry.go create mode 100644 libs/dockercredentials/registry_test.go diff --git a/.nextchanges/cli/docker-token.md b/.nextchanges/cli/docker-token.md new file mode 100644 index 00000000000..d17bd20901d --- /dev/null +++ b/.nextchanges/cli/docker-token.md @@ -0,0 +1 @@ +* Add experimental `databricks auth docker token` to generate Docker credentials for Databricks Artifact Registry. ([#6194](https://github.com/databricks/cli/pull/6194)) diff --git a/acceptance/cmd/auth/docker/help/out.test.toml b/acceptance/cmd/auth/docker/help/out.test.toml new file mode 100644 index 00000000000..98ea5040486 --- /dev/null +++ b/acceptance/cmd/auth/docker/help/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/docker/help/output.txt b/acceptance/cmd/auth/docker/help/output.txt new file mode 100644 index 00000000000..530233eba3a --- /dev/null +++ b/acceptance/cmd/auth/docker/help/output.txt @@ -0,0 +1,58 @@ + +>>> [CLI] auth --help +Authentication related commands. For more information regarding how +authentication for the Databricks CLI and SDKs work please refer to the documentation +linked below. + +AWS: https://docs.databricks.com/dev-tools/auth/index.html +Azure: https://learn.microsoft.com/azure/databricks/dev-tools/auth +GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html + +Usage: + databricks auth [command] + +Available Commands: + describe Describes the credentials and the source of those credentials, being used by the CLI to authenticate + docker (Experimental) Manage Docker authentication for Databricks Artifact Registry + login Log into a Databricks workspace or account + logout Log out of a Databricks profile + profiles Lists profiles from ~/.databrickscfg + switch Set the default profile + token Get authentication token + +Flags: + --account-id string Databricks Account ID + -h, --help help for auth + --host string Databricks Host + --workspace-id string Databricks Workspace ID + +Global Flags: + --debug enable debug logging + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) + +Use "databricks auth [command] --help" for more information about a command. + +>>> [CLI] auth docker --help +(Experimental) Manage Docker authentication for Databricks Artifact Registry + +Usage: + databricks auth docker [command] + +Available Commands: + token (Experimental) Generate a Docker credential + +Flags: + -h, --help help for docker + +Global Flags: + --account-id string Databricks Account ID + --debug enable debug logging + --host string Databricks Host + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) + --workspace-id string Databricks Workspace ID + +Use "databricks auth docker [command] --help" for more information about a command. diff --git a/acceptance/cmd/auth/docker/help/script b/acceptance/cmd/auth/docker/help/script new file mode 100644 index 00000000000..fe76e441c59 --- /dev/null +++ b/acceptance/cmd/auth/docker/help/script @@ -0,0 +1,2 @@ +trace "$CLI" auth --help +trace "$CLI" auth docker --help diff --git a/acceptance/cmd/auth/docker/script.prepare b/acceptance/cmd/auth/docker/script.prepare new file mode 100644 index 00000000000..6731358bbcb --- /dev/null +++ b/acceptance/cmd/auth/docker/script.prepare @@ -0,0 +1,27 @@ +setup_docker_auth() { + export DATABRICKS_HOST_ORIG="$DATABRICKS_HOST" + sethome "./home" + + unset DATABRICKS_ACCOUNT_ID + unset DATABRICKS_AUTH_TYPE + unset DATABRICKS_CLIENT_ID + unset DATABRICKS_CLIENT_SECRET + unset DATABRICKS_CONFIG_FILE + unset DATABRICKS_CONFIG_PROFILE + unset DATABRICKS_DISCOVERY_URL + unset DATABRICKS_HOST + unset DATABRICKS_TOKEN + unset DATABRICKS_WORKSPACE_ID + + export DATABRICKS_AUTH_STORAGE=plaintext + export DATABRICKS_LOG_LEVEL=error + export TEST_DAR_REGISTRY_HOST=123456789.container.us-west-2.cloud.databricks.com + + cat > "$HOME/.databrickscfg" <>> [CLI] auth docker token +{ + "Secret": "oauth-token", + "Username": "oauthtoken" +} diff --git a/acceptance/cmd/auth/docker/token/script b/acceptance/cmd/auth/docker/token/script new file mode 100644 index 00000000000..2681f35740e --- /dev/null +++ b/acceptance/cmd/auth/docker/token/script @@ -0,0 +1,6 @@ +setup_docker_auth + +mkdir -p "$HOME/.databricks" +cp token-cache.json "$HOME/.databricks/token-cache.json" + +printf '%s\n' "$TEST_DAR_REGISTRY_HOST" | trace "$CLI" auth docker token | jq -S . diff --git a/acceptance/cmd/auth/docker/token/test.toml b/acceptance/cmd/auth/docker/token/test.toml new file mode 100644 index 00000000000..c009f8c8e93 --- /dev/null +++ b/acceptance/cmd/auth/docker/token/test.toml @@ -0,0 +1,7 @@ +Ignore = [ + "home", +] + +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +Timeout = "15s" diff --git a/acceptance/cmd/auth/docker/token/token-cache.json b/acceptance/cmd/auth/docker/token/token-cache.json new file mode 100644 index 00000000000..92aff621332 --- /dev/null +++ b/acceptance/cmd/auth/docker/token/token-cache.json @@ -0,0 +1,11 @@ +{ + "version": 1, + "tokens": { + "docker-test": { + "access_token": "cached-access-token", + "token_type": "Bearer", + "refresh_token": "test-refresh-token", + "expiry": "2099-01-01T00:00:00Z" + } + } +} diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 7ef3a9f72ac..be8d7bc50fd 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -35,6 +35,7 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`, cmd.AddCommand(newLogoutCommand()) cmd.AddCommand(newProfilesCommand()) cmd.AddCommand(newTokenCommand(&authArguments)) + cmd.AddCommand(newDockerCommand(&authArguments)) cmd.AddCommand(newDescribeCommand()) cmd.AddCommand(newSwitchCommand()) return cmd diff --git a/cmd/auth/docker.go b/cmd/auth/docker.go new file mode 100644 index 00000000000..3ae32388ec6 --- /dev/null +++ b/cmd/auth/docker.go @@ -0,0 +1,54 @@ +package auth + +import ( + "time" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/spf13/cobra" +) + +func newDockerCommand(authArguments *auth.AuthArguments) *cobra.Command { + cmd := &cobra.Command{ + Use: "docker", + Short: "(Experimental) Manage Docker authentication for Databricks Artifact Registry", + } + cmd.AddCommand(newDockerTokenCommand(authArguments)) + return cmd +} + +func newDockerTokenCommand(authArguments *auth.AuthArguments) *cobra.Command { + return newDockerTokenCommandWithTokenLoader(authArguments, loadToken) +} + +func newDockerTokenCommandWithTokenLoader(authArguments *auth.AuthArguments, load tokenLoader) *cobra.Command { + cmd := &cobra.Command{ + Use: "token", + Short: "(Experimental) Generate a Docker credential", + } + + var tokenTimeout time.Duration + cmd.Flags().DurationVar(&tokenTimeout, "timeout", defaultTimeout, "Timeout for acquiring a token.") + var noForceRefresh bool + cmd.Flags().BoolVar(&noForceRefresh, "no-force-refresh", false, "Use a valid cached token instead of forcing a refresh.") + + cmd.PreRunE = validateDockerTokenRequest + cmd.RunE = func(cmd *cobra.Command, _ []string) error { + ctx := cmd.Context() + tokenStore, mode, err := storage.ResolveStore(ctx, storage.StorageModeUnknown) + if err != nil { + return err + } + return runDockerToken(ctx, cmd, loadTokenArgs{ + authArguments: authArguments, + tokenTimeout: tokenTimeout, + // Docker may reuse one credential for a long upload, so maximize its lifetime by refreshing it by default. + forceRefresh: !noForceRefresh, + profiler: profile.DefaultProfiler, + tokenStore: tokenStore, + mode: mode, + }, load) + } + return cmd +} diff --git a/cmd/auth/docker_profile.go b/cmd/auth/docker_profile.go new file mode 100644 index 00000000000..dcaf979801b --- /dev/null +++ b/cmd/auth/docker_profile.go @@ -0,0 +1,33 @@ +package auth + +import ( + "fmt" + + authlib "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/databricks-sdk-go/config" +) + +func validateDockerCredentialProfile(p profile.Profile) error { + if p.HasClientCredentials { + return fmt.Errorf("profile %q uses client credentials. Docker credential helper requires a profile created by databricks auth login", p.Name) + } + if p.AuthType != authTypeDatabricksCLI { + return fmt.Errorf("profile %q uses auth_type %q. Docker credential helper requires a profile created by databricks auth login", p.Name, p.AuthType) + } + if isDockerCredentialAccountOnlyProfile(p) { + return fmt.Errorf("profile %q does not target a workspace. Run databricks auth login --host and retry with that profile", p.Name) + } + return nil +} + +func isDockerCredentialAccountOnlyProfile(p profile.Profile) bool { + if p.Host == "" { + return true + } + cfg := &config.Config{Host: p.Host, AccountID: p.AccountID, WorkspaceID: p.WorkspaceID} + if authlib.IsClassicAccountHost(cfg.CanonicalHostName()) { + return true + } + return p.AccountID != "" && (p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone) +} diff --git a/cmd/auth/docker_profile_test.go b/cmd/auth/docker_profile_test.go new file mode 100644 index 00000000000..17db05d0503 --- /dev/null +++ b/cmd/auth/docker_profile_test.go @@ -0,0 +1,78 @@ +package auth + +import ( + "testing" + + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/stretchr/testify/assert" +) + +func TestValidateDockerCredentialProfile(t *testing.T) { + tests := []struct { + name string + profile profile.Profile + wantError string + }{ + { + name: "workspace", + profile: profile.Profile{ + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + }, + }, + { + name: "client credentials", + profile: profile.Profile{ + Name: "m2m", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + HasClientCredentials: true, + }, + wantError: "requires a profile created by databricks auth login", + }, + { + name: "unsupported auth type", + profile: profile.Profile{ + Name: "pat", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: "pat", + }, + wantError: "requires a profile created by databricks auth login", + }, + { + name: "classic account", + profile: profile.Profile{ + Name: "account", + Host: "https://accounts.cloud.databricks.test", + AccountID: "account-id", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + }, + wantError: "does not target a workspace", + }, + { + name: "unified account", + profile: profile.Profile{ + Name: "account", + Host: "https://workspace.cloud.databricks.test", + AccountID: "account-id", + AuthType: authTypeDatabricksCLI, + }, + wantError: "does not target a workspace", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateDockerCredentialProfile(tt.profile) + if tt.wantError == "" { + assert.NoError(t, err) + return + } + assert.ErrorContains(t, err, tt.wantError) + }) + } +} diff --git a/cmd/auth/docker_token.go b/cmd/auth/docker_token.go new file mode 100644 index 00000000000..b17a7ac7648 --- /dev/null +++ b/cmd/auth/docker_token.go @@ -0,0 +1,97 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/dockercredentials" + "github.com/spf13/cobra" + "golang.org/x/oauth2" +) + +type tokenLoader func(context.Context, loadTokenArgs) (*oauth2.Token, error) + +type dockerGetResponse struct { + Username string `json:"Username"` + Secret string `json:"Secret"` +} + +func runDockerToken(ctx context.Context, cmd *cobra.Command, args loadTokenArgs, load tokenLoader) error { + rawServer, err := io.ReadAll(cmd.InOrStdin()) + if err != nil { + return fmt.Errorf("read Docker credential request: %w", err) + } + registry, err := dockercredentials.ParseRegistryHost(string(rawServer)) + if err != nil { + return err + } + + selectedProfile, err := dockerTokenProfile(ctx, registry, args.profiler) + if err != nil { + return err + } + + args.authArguments = &auth.AuthArguments{ + Host: selectedProfile.Host, + AccountID: selectedProfile.AccountID, + WorkspaceID: selectedProfile.WorkspaceID, + } + args.profileName = selectedProfile.Name + args.args = nil + + t, err := load(ctx, args) + if err != nil { + return err + } + + return json.NewEncoder(cmd.OutOrStdout()).Encode(dockerGetResponse{ + Username: dockercredentials.OAuthTokenUsername, + Secret: t.AccessToken, + }) +} + +func validateDockerTokenRequest(cmd *cobra.Command, args []string) error { + if len(args) > 0 { + return errors.New("auth docker token does not accept positional arguments") + } + for _, name := range []string{"profile", "host", "account-id", "workspace-id"} { + flag := cmd.Flag(name) + if flag != nil && flag.Changed { + return fmt.Errorf("auth docker token does not support --%s", name) + } + } + return nil +} + +func dockerTokenProfile(ctx context.Context, registry dockercredentials.Registry, profiler profile.Profiler) (profile.Profile, error) { + workspaceProfiles, err := profiler.LoadProfiles(ctx, func(p profile.Profile) bool { + return p.WorkspaceID == registry.WorkspaceID + }) + if err != nil { + return profile.Profile{}, err + } + if len(workspaceProfiles) == 0 { + return profile.Profile{}, fmt.Errorf("no Databricks profile found for workspace ID %s from registry host %s. Run databricks auth login --host and set workspace_id for that profile", registry.WorkspaceID, registry.Host) + } + + var matchingProfiles profile.Profiles + for _, p := range workspaceProfiles { + if validateDockerCredentialProfile(p) == nil { + matchingProfiles = append(matchingProfiles, p) + } + } + switch len(matchingProfiles) { + case 0: + return profile.Profile{}, validateDockerCredentialProfile(workspaceProfiles[0]) + case 1: + return matchingProfiles[0], nil + default: + return profile.Profile{}, fmt.Errorf("multiple Databricks profiles match workspace ID %s: %s. Remove duplicate workspace_id entries before using Docker credential helper", registry.WorkspaceID, strings.Join(matchingProfiles.Names(), " and ")) + } +} diff --git a/cmd/auth/docker_token_test.go b/cmd/auth/docker_token_test.go new file mode 100644 index 00000000000..d7992459ba1 --- /dev/null +++ b/cmd/auth/docker_token_test.go @@ -0,0 +1,391 @@ +package auth + +import ( + "bytes" + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/dockercredentials" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +func newTestDockerTokenCommand(t *testing.T, load tokenLoader) (*cobra.Command, *bytes.Buffer) { + t.Helper() + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + require.NoError(t, databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + })) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + stdout := &bytes.Buffer{} + cmd := newDockerTokenCommandWithTokenLoader(&auth.AuthArguments{}, load) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetOut(stdout) + return cmd, stdout +} + +func TestDockerTokenEmitsGetResponse(t *testing.T) { + var gotProfile string + loadToken := func(_ context.Context, args loadTokenArgs) (*oauth2.Token, error) { + gotProfile = args.profileName + return &oauth2.Token{AccessToken: "access-token"}, nil + } + + cmd, stdout := newTestDockerTokenCommand(t, loadToken) + + require.NoError(t, cmd.Execute()) + assert.Equal(t, "workspace", gotProfile) + + var got map[string]string + require.NoError(t, json.Unmarshal(stdout.Bytes(), &got)) + assert.Equal(t, map[string]string{ + "Username": "oauthtoken", + "Secret": "access-token", + }, got) +} + +func TestDockerTokenForceRefresh(t *testing.T) { + tests := []struct { + name string + args []string + want bool + }{ + {name: "default", want: true}, + {name: "disabled", args: []string{"--no-force-refresh"}, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got bool + loadToken := func(_ context.Context, args loadTokenArgs) (*oauth2.Token, error) { + got = args.forceRefresh + return &oauth2.Token{AccessToken: "access-token"}, nil + } + cmd, _ := newTestDockerTokenCommand(t, loadToken) + cmd.SetArgs(tt.args) + + require.NoError(t, cmd.Execute()) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestRunDockerTokenUsesConfiguredProfiler(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + }, + }, + } + + var gotProfile string + loadToken := func(_ context.Context, args loadTokenArgs) (*oauth2.Token, error) { + gotProfile = args.profileName + return &oauth2.Token{AccessToken: "access-token"}, nil + } + + cmd := &cobra.Command{Use: "token"} + var stdout bytes.Buffer + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetOut(&stdout) + + err := runDockerToken(ctx, cmd, loadTokenArgs{ + authArguments: &auth.AuthArguments{}, + profiler: profiler, + }, loadToken) + require.NoError(t, err) + assert.Equal(t, "workspace", gotProfile) + + var got dockerGetResponse + require.NoError(t, json.Unmarshal(stdout.Bytes(), &got)) +} + +func TestRunDockerTokenUsesMatchedProfileAccountID(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + AccountID: "profile-account", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + }, + }, + } + + cmd := &cobra.Command{Use: "token"} + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetOut(&bytes.Buffer{}) + + err := runDockerToken(ctx, cmd, loadTokenArgs{ + authArguments: &auth.AuthArguments{}, + profiler: profiler, + }, func(_ context.Context, args loadTokenArgs) (*oauth2.Token, error) { + assert.Equal(t, "profile-account", args.authArguments.AccountID) + return &oauth2.Token{AccessToken: "access-token"}, nil + }) + require.NoError(t, err) +} + +func TestDockerTokenProfileSelectsByWorkspaceID(t *testing.T) { + registry := dockercredentials.Registry{ + WorkspaceID: "123456789", + Host: "123456789.container.us-west-2.cloud.databricks.test", + } + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{{ + Name: "workspace", + Host: "https://workspace.dev.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }}, + } + selectedProfile, err := dockerTokenProfile(t.Context(), registry, profiler) + require.NoError(t, err) + assert.Equal(t, "workspace", selectedProfile.Name) +} + +func TestDockerTokenProfileRejectsDuplicateWorkspaceID(t *testing.T) { + registry := dockercredentials.Registry{ + WorkspaceID: "123456789", + Host: "123456789.container.us-west-2.cloud.databricks.test", + } + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "prod", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }, + { + Name: "dev", + Host: "https://workspace.dev.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }, + }, + } + _, err := dockerTokenProfile(t.Context(), registry, profiler) + assert.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") + assert.ErrorContains(t, err, "prod and dev") +} + +func TestDockerTokenProfileIgnoresUnsupportedDuplicateProfile(t *testing.T) { + registry := dockercredentials.Registry{ + WorkspaceID: "123456789", + Host: "123456789.container.us-west-2.cloud.databricks.test", + } + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }, + { + Name: "m2m", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + HasClientCredentials: true, + }, + }, + } + selectedProfile, err := dockerTokenProfile(t.Context(), registry, profiler) + require.NoError(t, err) + assert.Equal(t, "workspace", selectedProfile.Name) +} + +func TestDockerTokenRejectsPositionalArgs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + cmd := newDockerTokenCommandWithTokenLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetArgs([]string{"DEFAULT"}) + + err := cmd.Execute() + assert.ErrorContains(t, err, "auth docker token does not accept positional arguments") +} + +func TestDockerTokenValidatesBeforeResolvingTokenStore(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + t.Setenv(storage.EnvVar, "invalid") + + cmd := newDockerTokenCommandWithTokenLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"DEFAULT"}) + + err := cmd.Execute() + assert.ErrorContains(t, err, "auth docker token does not accept positional arguments") +} + +func TestDockerTokenRejectsAuthSelectionFlags(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + require.NoError(t, databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: "DEFAULT", + Host: "https://profile.cloud.databricks.test", + AuthType: authTypeDatabricksCLI, + })) + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + cases := [][]string{ + {"--profile", "DEFAULT"}, + {"--host", "https://workspace.cloud.databricks.test"}, + {"--profile", "DEFAULT", "--host", "https://workspace.cloud.databricks.test"}, + {"--account-id", "abc"}, + {"--workspace-id", "123456789"}, + } + + for _, args := range cases { + t.Run(strings.Join(args, " "), func(t *testing.T) { + var authArgs auth.AuthArguments + cmd := &cobra.Command{Use: "auth"} + cmd.PersistentFlags().StringVar(&authArgs.Host, "host", "", "Databricks Host") + cmd.PersistentFlags().StringVar(&authArgs.AccountID, "account-id", "", "Databricks Account ID") + cmd.PersistentFlags().StringVar(&authArgs.WorkspaceID, "workspace-id", "", "Databricks Workspace ID") + dockerCmd := &cobra.Command{Use: "docker"} + dockerCmd.AddCommand(newDockerTokenCommandWithTokenLoader(&authArgs, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + })) + cmd.AddCommand(dockerCmd) + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetArgs(append([]string{"docker", "token"}, args...)) + + err := cmd.Execute() + assert.ErrorContains(t, err, "auth docker token does not support") + }) + } +} + +func TestDockerTokenRejectsNonDARHost(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + cmd := newDockerTokenCommandWithTokenLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("registry.example.com\n")) + + err := cmd.Execute() + assert.ErrorContains(t, err, "is not a Databricks Artifact Registry host") +} + +func TestDockerTokenErrorsWithoutMatchingProfile(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + require.NoError(t, os.WriteFile(configFile, []byte(""), 0o600)) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + cmd := newDockerTokenCommandWithTokenLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + + err := cmd.Execute() + assert.ErrorContains(t, err, "no Databricks profile found for workspace ID 123456789") + assert.ErrorContains(t, err, "databricks auth login --host ") + assert.ErrorContains(t, err, "workspace_id") +} + +func TestDockerTokenErrorsWithMultipleMatchingProfiles(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + for _, name := range []string{"one", "two"} { + require.NoError(t, databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: name, + Host: "https://" + name + ".cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + })) + } + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + cmd := newDockerTokenCommandWithTokenLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + + err := cmd.Execute() + assert.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") + assert.ErrorContains(t, err, "one and two") + assert.ErrorContains(t, err, "Remove duplicate workspace_id entries") +} diff --git a/libs/dockercredentials/registry.go b/libs/dockercredentials/registry.go new file mode 100644 index 00000000000..cc3c6e7b864 --- /dev/null +++ b/libs/dockercredentials/registry.go @@ -0,0 +1,103 @@ +package dockercredentials + +import ( + "errors" + "fmt" + "net/url" + "strings" + "unicode" + + "github.com/databricks/databricks-sdk-go/common/environment" +) + +const ( + // OAuthTokenUsername is the username returned to Docker with an OAuth access token. + OAuthTokenUsername = "oauthtoken" + registryHostInfix = ".container." +) + +// Registry identifies the workspace and canonical host of a Databricks Artifact Registry endpoint. +type Registry struct { + WorkspaceID string + Host string +} + +// normalizeServerAddress accepts the bare host or HTTPS URL forms allowed by Docker's credential-helper protocol. +// Databricks Artifact Registry uses standard HTTPS and has no configurable port. +func normalizeServerAddress(raw string) (string, error) { + value := strings.TrimSpace(raw) + if value == "" { + return "", errors.New("server address is required") + } + + if !strings.Contains(value, "://") { + value = "https://" + value + } + u, err := url.Parse(value) + if err != nil { + return "", fmt.Errorf("parse server address %q: %w", raw, err) + } + if !strings.EqualFold(u.Scheme, "https") { + return "", fmt.Errorf("unsupported registry URL scheme %q", u.Scheme) + } + if u.Port() != "" { + return "", errors.New("registry address must not include a port") + } + + host := strings.TrimSuffix(strings.ToLower(u.Hostname()), ".") + if host == "" { + return "", errors.New("server address is required") + } + return host, nil +} + +// ParseRegistryHost normalizes a Databricks Artifact Registry address and extracts its workspace ID. +func ParseRegistryHost(raw string) (Registry, error) { + host, err := normalizeServerAddress(raw) + if err != nil { + return Registry{}, err + } + + dnsZone, ok := matchingDatabricksDNSZone(host) + if !ok { + return Registry{}, fmt.Errorf("%q is not a Databricks Artifact Registry host", host) + } + + trimmed := strings.TrimSuffix(host, dnsZone) + workspaceID, region, ok := strings.Cut(trimmed, registryHostInfix) + if !ok || !isDNSLabel(workspaceID) || !isDNSLabel(region) { + return Registry{}, fmt.Errorf("%q is not a Databricks Artifact Registry host", host) + } + + return Registry{ + WorkspaceID: workspaceID, + Host: host, + }, nil +} + +func matchingDatabricksDNSZone(host string) (string, bool) { + dnsZone := strings.ToLower(environment.GetEnvironmentForHostname(host).DnsZone) + // The SDK defaults unknown hosts to AWS production, so verify that the returned zone actually matched. + return dnsZone, dnsZone != "" && strings.HasSuffix(host, dnsZone) +} + +// isDNSLabel accepts one lowercase ASCII label: letters, digits, and interior hyphens, up to 63 bytes. +// Workspace IDs, regions, and Docker input are not uniformly prevalidated before becoming hostname components. +func isDNSLabel(label string) bool { + if label == "" || len(label) > 63 { + return false + } + for i, r := range label { + if r > unicode.MaxASCII { + return false + } + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') { + continue + } + if r == '-' && i > 0 && i < len(label)-1 { + continue + } + return false + } + return true +} diff --git a/libs/dockercredentials/registry_test.go b/libs/dockercredentials/registry_test.go new file mode 100644 index 00000000000..071f93c6adf --- /dev/null +++ b/libs/dockercredentials/registry_test.go @@ -0,0 +1,89 @@ +package dockercredentials + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/common/environment" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseRegistryHost(t *testing.T) { + cases := []string{ + "123456789.container.us-west-2.cloud.databricks.com", + "https://123456789.container.us-west-2.cloud.databricks.com", + "123456789.container.us-west-2.cloud.databricks.com/v2/", + } + + for _, input := range cases { + t.Run(input, func(t *testing.T) { + got, err := ParseRegistryHost(input) + require.NoError(t, err) + assert.Equal(t, Registry{ + WorkspaceID: "123456789", + Host: "123456789.container.us-west-2.cloud.databricks.com", + }, got) + }) + } +} + +func TestParseRegistryHostSupportsAllDatabricksEnvironmentZones(t *testing.T) { + for _, env := range environment.AllEnvironments() { + dnsZone := env.DnsZone + if dnsZone == "" { + continue + } + t.Run(dnsZone, func(t *testing.T) { + wantHost := "123456789.container.test-region" + dnsZone + registry, err := ParseRegistryHost("https://" + wantHost + "/v2/") + require.NoError(t, err) + assert.Equal(t, Registry{ + WorkspaceID: "123456789", + Host: wantHost, + }, registry) + }) + } +} + +func TestParseRegistryHostUsesLongestDNSZoneSuffix(t *testing.T) { + got, err := ParseRegistryHost("123456789.container.us-west-2.staging.cloud.databricks.com") + require.NoError(t, err) + assert.Equal(t, Registry{ + WorkspaceID: "123456789", + Host: "123456789.container.us-west-2.staging.cloud.databricks.com", + }, got) +} + +func TestParseRegistryHostRejectsNonDARHost(t *testing.T) { + _, err := ParseRegistryHost("registry.example.com") + assert.ErrorContains(t, err, `"registry.example.com" is not a Databricks Artifact Registry host`) +} + +func TestParseRegistryHostRejectsPluralContainersInfix(t *testing.T) { + _, err := ParseRegistryHost("123.containers.us-west-2.cloud.databricks.com") + assert.ErrorContains(t, err, `"123.containers.us-west-2.cloud.databricks.com" is not a Databricks Artifact Registry host`) +} + +func TestParseRegistryHostRejectsInvalidLabels(t *testing.T) { + _, err := ParseRegistryHost("-123.container.us-west-2.cloud.databricks.com") + assert.ErrorContains(t, err, `"-123.container.us-west-2.cloud.databricks.com" is not a Databricks Artifact Registry host`) + + _, err = ParseRegistryHost("123.container.-us-west-2.cloud.databricks.com") + assert.ErrorContains(t, err, `"123.container.-us-west-2.cloud.databricks.com" is not a Databricks Artifact Registry host`) +} + +func TestNormalizeServerAddress(t *testing.T) { + got, err := normalizeServerAddress("HTTPS://123.container.US-WEST-2.cloud.databricks.com/v2/") + require.NoError(t, err) + assert.Equal(t, "123.container.us-west-2.cloud.databricks.com", got) +} + +func TestNormalizeServerAddressRejectsNonHTTPSURL(t *testing.T) { + _, err := normalizeServerAddress("http://123.container.us-west-2.cloud.databricks.com") + assert.ErrorContains(t, err, "unsupported registry URL scheme") +} + +func TestNormalizeServerAddressRejectsPort(t *testing.T) { + _, err := normalizeServerAddress("https://123.container.us-west-2.cloud.databricks.com:443") + assert.ErrorContains(t, err, "registry address must not include a port") +}