From 73a927beca9e7d213773ca3366bfe45c1260dc43 Mon Sep 17 00:00:00 2001 From: Chen Shou Date: Thu, 20 Aug 2026 17:26:21 +0000 Subject: [PATCH 1/2] 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 | 36 ++ cmd/auth/docker/docker.go | 44 +++ cmd/auth/docker/docker_profile.go | 33 ++ cmd/auth/docker/docker_profile_test.go | 79 ++++ cmd/auth/docker/docker_token.go | 107 +++++ cmd/auth/docker/docker_token_test.go | 373 ++++++++++++++++++ libs/dockercredentials/registry.go | 103 +++++ libs/dockercredentials/registry_test.go | 89 +++++ 19 files changed, 987 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/docker.go create mode 100644 cmd/auth/docker/docker_profile.go create mode 100644 cmd/auth/docker/docker_profile_test.go create mode 100644 cmd/auth/docker/docker_token.go create mode 100644 cmd/auth/docker/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..6bbee3730b3 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()) 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..37f80d1a9d1 --- /dev/null +++ b/cmd/auth/docker.go @@ -0,0 +1,36 @@ +package auth + +import ( + "context" + + authdocker "github.com/databricks/cli/cmd/auth/docker" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/spf13/cobra" + "golang.org/x/oauth2" +) + +func newDockerCommand() *cobra.Command { + return authdocker.New(loadDockerToken) +} + +func loadDockerToken(ctx context.Context, req authdocker.TokenRequest) (*oauth2.Token, error) { + tokenStore, mode, err := storage.ResolveStore(ctx, storage.StorageModeUnknown) + if err != nil { + return nil, err + } + return loadToken(ctx, loadTokenArgs{ + authArguments: &auth.AuthArguments{ + Host: req.Profile.Host, + AccountID: req.Profile.AccountID, + WorkspaceID: req.Profile.WorkspaceID, + }, + profileName: req.Profile.Name, + tokenTimeout: req.Timeout, + forceRefresh: req.ForceRefresh, + profiler: profile.DefaultProfiler, + tokenStore: tokenStore, + mode: mode, + }) +} diff --git a/cmd/auth/docker/docker.go b/cmd/auth/docker/docker.go new file mode 100644 index 00000000000..1cd3a46d08d --- /dev/null +++ b/cmd/auth/docker/docker.go @@ -0,0 +1,44 @@ +package docker + +import ( + "time" + + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/spf13/cobra" +) + +const defaultTokenTimeout = time.Hour + +// New returns the Docker authentication command. +func New(load TokenLoader) *cobra.Command { + cmd := &cobra.Command{ + Use: "docker", + Short: "(Experimental) Manage Docker authentication for Databricks Artifact Registry", + } + cmd.AddCommand(newDockerTokenCommand(load)) + return cmd +} + +func newDockerTokenCommand(load TokenLoader) *cobra.Command { + cmd := &cobra.Command{ + Use: "token", + Short: "(Experimental) Generate a Docker credential", + } + + var tokenTimeout time.Duration + cmd.Flags().DurationVar(&tokenTimeout, "timeout", defaultTokenTimeout, "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() + return runDockerToken(ctx, cmd, tokenOptions{ + 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, + }, load) + } + return cmd +} diff --git a/cmd/auth/docker/docker_profile.go b/cmd/auth/docker/docker_profile.go new file mode 100644 index 00000000000..af056110240 --- /dev/null +++ b/cmd/auth/docker/docker_profile.go @@ -0,0 +1,33 @@ +package docker + +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 != authlib.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/docker_profile_test.go b/cmd/auth/docker/docker_profile_test.go new file mode 100644 index 00000000000..6b7253b9cd9 --- /dev/null +++ b/cmd/auth/docker/docker_profile_test.go @@ -0,0 +1,79 @@ +package docker + +import ( + "testing" + + "github.com/databricks/cli/libs/auth" + "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: auth.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: auth.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: auth.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/docker_token.go b/cmd/auth/docker/docker_token.go new file mode 100644 index 00000000000..75fa44df716 --- /dev/null +++ b/cmd/auth/docker/docker_token.go @@ -0,0 +1,107 @@ +package docker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "time" + + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/dockercredentials" + "github.com/spf13/cobra" + "golang.org/x/oauth2" +) + +// TokenRequest contains the resolved profile and token acquisition options. +type TokenRequest struct { + Profile profile.Profile + Timeout time.Duration + ForceRefresh bool +} + +// TokenLoader acquires an OAuth token for a Docker credential request. +type TokenLoader func(context.Context, TokenRequest) (*oauth2.Token, error) + +type tokenOptions struct { + tokenTimeout time.Duration + forceRefresh bool + profiler profile.Profiler +} + +type dockerGetResponse struct { + Username string `json:"Username"` + Secret string `json:"Secret"` +} + +func runDockerToken(ctx context.Context, cmd *cobra.Command, opts tokenOptions, 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, opts.profiler) + if err != nil { + return err + } + + t, err := load(ctx, TokenRequest{ + Profile: selectedProfile, + Timeout: opts.tokenTimeout, + ForceRefresh: opts.forceRefresh, + }) + 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/docker_token_test.go b/cmd/auth/docker/docker_token_test.go new file mode 100644 index 00000000000..36ff35c15ce --- /dev/null +++ b/cmd/auth/docker/docker_token_test.go @@ -0,0 +1,373 @@ +package docker + +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: auth.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 := newDockerTokenCommand(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 TokenRequest) (*oauth2.Token, error) { + gotProfile = args.Profile.Name + 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 TokenRequest) (*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: auth.AuthTypeDatabricksCli, + }, + }, + } + + var gotProfile string + loadToken := func(_ context.Context, args TokenRequest) (*oauth2.Token, error) { + gotProfile = args.Profile.Name + 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, tokenOptions{ + 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: auth.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, tokenOptions{ + profiler: profiler, + }, func(_ context.Context, args TokenRequest) (*oauth2.Token, error) { + assert.Equal(t, "profile-account", args.Profile.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: auth.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: auth.AuthTypeDatabricksCli, + }, + { + Name: "dev", + Host: "https://workspace.dev.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + AuthType: auth.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: auth.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 := newDockerTokenCommand(func(context.Context, TokenRequest) (*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 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: auth.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(newDockerTokenCommand(func(context.Context, TokenRequest) (*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 := newDockerTokenCommand(func(context.Context, TokenRequest) (*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 := newDockerTokenCommand(func(context.Context, TokenRequest) (*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: auth.AuthTypeDatabricksCli, + })) + } + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + t.Setenv("USERPROFILE", dir) + + cmd := newDockerTokenCommand(func(context.Context, TokenRequest) (*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") +} From 1e1a4acad48f6d3d5e62c68638b15c135e637db2 Mon Sep 17 00:00:00 2001 From: Chen Shou Date: Thu, 20 Aug 2026 17:29:19 +0000 Subject: [PATCH 2/2] Configure Docker authentication for Artifact Registry --- .nextchanges/cli/configure-docker.md | 1 + .../auth/docker/configure/docker-config.json | 7 + .../cmd/auth/docker/configure/out.test.toml | 2 + .../cmd/auth/docker/configure/output.txt | 23 + .../docker/configure/profile.databrickscfg | 5 + acceptance/cmd/auth/docker/configure/script | 44 ++ .../cmd/auth/docker/configure/test.toml | 7 + acceptance/cmd/auth/docker/help/output.txt | 1 + cmd/auth/docker/docker.go | 1 + cmd/auth/docker/docker_configure.go | 267 +++++++++ cmd/auth/docker/docker_configure_test.go | 521 ++++++++++++++++++ libs/dockercredentials/docker_config.go | 122 ++++ libs/dockercredentials/docker_config_test.go | 145 +++++ libs/dockercredentials/registry.go | 35 ++ libs/dockercredentials/registry_test.go | 71 +++ libs/dockercredentials/shim.go | 142 +++++ libs/dockercredentials/shim_test.go | 280 ++++++++++ 17 files changed, 1674 insertions(+) create mode 100644 .nextchanges/cli/configure-docker.md create mode 100644 acceptance/cmd/auth/docker/configure/docker-config.json create mode 100644 acceptance/cmd/auth/docker/configure/out.test.toml create mode 100644 acceptance/cmd/auth/docker/configure/output.txt create mode 100644 acceptance/cmd/auth/docker/configure/profile.databrickscfg create mode 100644 acceptance/cmd/auth/docker/configure/script create mode 100644 acceptance/cmd/auth/docker/configure/test.toml create mode 100644 cmd/auth/docker/docker_configure.go create mode 100644 cmd/auth/docker/docker_configure_test.go create mode 100644 libs/dockercredentials/docker_config.go create mode 100644 libs/dockercredentials/docker_config_test.go create mode 100644 libs/dockercredentials/shim.go create mode 100644 libs/dockercredentials/shim_test.go diff --git a/.nextchanges/cli/configure-docker.md b/.nextchanges/cli/configure-docker.md new file mode 100644 index 00000000000..1b2fd1b5709 --- /dev/null +++ b/.nextchanges/cli/configure-docker.md @@ -0,0 +1 @@ +* Add experimental `databricks auth docker configure` to configure Docker credential helper access for Databricks Artifact Registry. ([#6336](https://github.com/databricks/cli/pull/6336)) diff --git a/acceptance/cmd/auth/docker/configure/docker-config.json b/acceptance/cmd/auth/docker/configure/docker-config.json new file mode 100644 index 00000000000..3e346bda29b --- /dev/null +++ b/acceptance/cmd/auth/docker/configure/docker-config.json @@ -0,0 +1,7 @@ +{ + "auths": { + "registry.example.test": { + "auth": "preserved" + } + } +} diff --git a/acceptance/cmd/auth/docker/configure/out.test.toml b/acceptance/cmd/auth/docker/configure/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/cmd/auth/docker/configure/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/cmd/auth/docker/configure/output.txt b/acceptance/cmd/auth/docker/configure/output.txt new file mode 100644 index 00000000000..c594514cd9e --- /dev/null +++ b/acceptance/cmd/auth/docker/configure/output.txt @@ -0,0 +1,23 @@ + +=== Docker config +{ + "auths": { + "registry.example.test": { + "auth": "preserved" + } + }, + "credHelpers": { + "[NUMID].container.us-west-2.cloud.databricks.com": "databricks" + } +} + +=== Installed shim +Platform shim verified + +=== Credential from shim + +>>> docker-credential-databricks get +{ + "Secret": "oauth-token", + "Username": "oauthtoken" +} diff --git a/acceptance/cmd/auth/docker/configure/profile.databrickscfg b/acceptance/cmd/auth/docker/configure/profile.databrickscfg new file mode 100644 index 00000000000..385fdc8c49b --- /dev/null +++ b/acceptance/cmd/auth/docker/configure/profile.databrickscfg @@ -0,0 +1,5 @@ +[docker-test] +# This host is used only to derive the registry DNS zone and is replaced before the helper performs authentication. +host = https://workspace.cloud.databricks.com +auth_type = databricks-cli +workspace_id = 123456789 diff --git a/acceptance/cmd/auth/docker/configure/script b/acceptance/cmd/auth/docker/configure/script new file mode 100644 index 00000000000..4b99ff6299a --- /dev/null +++ b/acceptance/cmd/auth/docker/configure/script @@ -0,0 +1,44 @@ +setup_docker_auth + +mkdir -p "$HOME/.databricks" +cp "$TESTDIR/../token/token-cache.json" "$HOME/.databricks/token-cache.json" +cp profile.databrickscfg "$HOME/.databrickscfg" + +export DOCKER_CONFIG="$HOME/docker" +mkdir -p "$DOCKER_CONFIG" +cp docker-config.json "$DOCKER_CONFIG/config.json" + +mkdir -p "$HOME/bin" +if [[ "$OSTYPE" == "msys" || "$OSTYPE" == "cygwin" || "$OSTYPE" == "win32" ]]; then + cli_path="$HOME/bin/databricks.exe" + cp "$(cygpath -u "$CLI")" "$cli_path" + helper="$HOME/bin/docker-credential-databricks.cmd" +else + cli_path="$HOME/bin/databricks" + cp "$CLI" "$cli_path" + chmod +x "$cli_path" + helper="$HOME/bin/docker-credential-databricks" +fi +export PATH="$(dirname "$helper"):$PATH" + +"$cli_path" auth docker configure docker-test --region us-west-2 >LOG.configure 2>&1 + +title "Docker config\n" +jq -S . "$DOCKER_CONFIG/config.json" + +title "Installed shim\n" +if [[ "$OSTYPE" == "msys" || "$OSTYPE" == "cygwin" || "$OSTYPE" == "win32" ]]; then + test -f "$helper" +else + test -x "$helper" +fi +printf 'Platform shim verified\n' + +title "Credential from shim\n" +cat > "$HOME/.databrickscfg" < to specify which profile to use") + } + + profiles, err := profiler.LoadProfiles(ctx, profile.MatchWorkspaceProfiles) + if err != nil { + return "", err + } + currentDefault, _ := databrickscfg.GetDefaultProfile(ctx, env.Get(ctx, "DATABRICKS_CONFIG_FILE")) + return profile.SelectProfile(ctx, profile.SelectConfig{ + Label: "Select a workspace profile", + Profiles: profiles, + StartInSearchMode: true, + Default: currentDefault, + ActiveTemplate: `▸ {{.Name | bold}}{{if .IsDefault}} {{ "[default]" | green }}{{end}}{{if .AccountID}} (account: {{.AccountID|faint}}){{else if .Host}} ({{.Host|faint}}){{end}}`, + InactiveTemplate: ` {{.Name}}{{if .IsDefault}} [default]{{end}}{{if .AccountID}} (account: {{.AccountID|faint}}){{else if .Host}} ({{.Host|faint}}){{end}}`, + SelectedTemplate: `{{ "Using profile" | faint }}: {{ .Name | bold }}`, + }) +} + +func loadAndValidateConfigureDockerProfile(ctx context.Context, profileName string, profiler profile.Profiler) (profile.Profile, error) { + profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName)) + if err != nil { + return profile.Profile{}, err + } + if len(profiles) == 0 { + return profile.Profile{}, fmt.Errorf("profile %q not found", profileName) + } + if err := validateDockerCredentialProfile(profiles[0]); err != nil { + return profile.Profile{}, err + } + return profiles[0], nil +} + +func resolveConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, executable string, deps configureDockerDeps) (string, error) { + if p.WorkspaceID != "" && p.WorkspaceID != authlib.WorkspaceIDNone { + return p.WorkspaceID, nil + } + + cfg := &databricks.Config{ + Profile: p.Name, + Host: p.Host, + AccountID: p.AccountID, + AuthType: p.AuthType, + ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), + Loaders: databrickscfg.ProfileAuthLoaders, + DatabricksCliPath: executable, + } + w, err := deps.newWorkspaceClient(cfg) + if err != nil { + return "", fmt.Errorf("load workspace profile %q: %w. Run databricks auth login --host and retry with that profile", p.Name, err) + } + // The selected profile may contain the CLI-only "none" sentinel, which the SDK would send as a routing header. + w.Config.WorkspaceID = "" + workspaceID, err := deps.resolveWorkspaceID(ctx, w) + if err != nil { + return "", fmt.Errorf("resolve workspace ID for profile %q: %w. Run databricks auth login --host and retry with that profile", p.Name, err) + } + return workspaceID, nil +} + +// ensureConfigureDockerUniqueProfile rejects registry-to-profile mappings that would be ambiguous at credential lookup time. +func ensureConfigureDockerUniqueProfile(ctx context.Context, profiler profile.Profiler, p profile.Profile, workspaceID string) error { + matches, err := profiler.LoadProfiles(ctx, func(candidate profile.Profile) bool { + return candidate.WorkspaceID == workspaceID + }) + if err != nil { + return err + } + + if p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone { + matches = append(matches, p) + } + var validProfiles profile.Profiles + for _, candidate := range matches { + if validateDockerCredentialProfile(candidate) == nil { + validProfiles = append(validProfiles, candidate) + } + } + names := validProfiles.Names() + if len(names) <= 1 { + return nil + } + + return fmt.Errorf("multiple Databricks profiles match workspace ID %s: %s. Remove duplicate workspace_id entries before using Docker credential helper", workspaceID, strings.Join(names, " and ")) +} + +func persistConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, workspaceID string) error { + return databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), + Profile: p.Name, + WorkspaceID: workspaceID, + }) +} + +func configureDockerConfigPath(ctx context.Context) (string, error) { + if dockerConfig := env.Get(ctx, "DOCKER_CONFIG"); dockerConfig != "" { + return filepath.Join(dockerConfig, "config.json"), nil + } + home, err := env.UserHomeDir(ctx) + if err != nil { + return "", err + } + return filepath.Join(home, ".docker", "config.json"), nil +} diff --git a/cmd/auth/docker/docker_configure_test.go b/cmd/auth/docker/docker_configure_test.go new file mode 100644 index 00000000000..403522e6886 --- /dev/null +++ b/cmd/auth/docker/docker_configure_test.go @@ -0,0 +1,521 @@ +package docker + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "os" + "path/filepath" + "runtime" + "testing" + + authlib "github.com/databricks/cli/libs/auth" + "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/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newDockerConfigureTestCommand(ctx context.Context, args ...string) *cobra.Command { + return newDockerConfigureTestCommandWithDeps(ctx, defaultConfigureDockerDeps(), args...) +} + +func newDockerConfigureTestCommandWithDeps(ctx context.Context, deps configureDockerDeps, args ...string) *cobra.Command { + cmd := &cobra.Command{Use: "auth"} + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.PersistentFlags().String("host", "", "Databricks Host") + cmd.PersistentFlags().String("account-id", "", "Databricks Account ID") + cmd.PersistentFlags().String("workspace-id", "", "Databricks Workspace ID") + dockerCmd := &cobra.Command{Use: "docker"} + dockerCmd.AddCommand(newDockerConfigureCommandWithDeps(deps)) + cmd.AddCommand(dockerCmd) + cmd.SetContext(ctx) + cmd.SetArgs(args) + return cmd +} + +func writeConfigureDockerProfile(t *testing.T, ctx context.Context, configFile string, cfg *config.Config) { + t.Helper() + cfg.ConfigFile = configFile + require.NoError(t, databrickscfg.SaveToProfile(ctx, cfg)) +} + +func readCredentialHelpers(t *testing.T, path string) map[string]string { + t.Helper() + raw, err := os.ReadFile(path) + require.NoError(t, err) + + var cfg struct { + CredHelpers map[string]string `json:"credHelpers"` + } + require.NoError(t, json.Unmarshal(raw, &cfg)) + return cfg.CredHelpers +} + +func configureDockerRegistryHostStub(t *testing.T, wantWorkspaceID, wantRegion, wantWorkspaceHost, registryHost string) func(string, string, string) (string, error) { + t.Helper() + return func(workspaceID, region, workspaceHost string) (string, error) { + assert.Equal(t, wantWorkspaceID, workspaceID) + assert.Equal(t, wantRegion, region) + assert.Equal(t, wantWorkspaceHost, workspaceHost) + return registryHost, nil + } +} + +func writeConfigureDockerExecutable(t *testing.T, dir string) string { + t.Helper() + name := "databricks" + if runtime.GOOS == "windows" { + name += ".exe" + } + path := filepath.Join(dir, name) + require.NoError(t, os.MkdirAll(dir, 0o755)) + require.NoError(t, os.WriteFile(path, []byte("databricks executable"), 0o755)) + return path +} + +func TestConfigureDockerCommandWritesDockerConfigAndShim(t *testing.T) { + ctx, stderr := cmdio.NewTestContextWithStderr(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + binDir := filepath.Join(dir, "bin") + workspaceHost := "https://workspace.staging.cloud.databricks.test" + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "DEFAULT", + Host: workspaceHost, + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("PATH", binDir) + + registryHost := "123456789.container.us-west-2.staging.cloud.databricks.test" + deps := defaultConfigureDockerDeps() + databricksPath := writeConfigureDockerExecutable(t, binDir) + deps.executable = func() (string, error) { + return databricksPath, nil + } + deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, registryHost) + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "--profile", "DEFAULT", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) + + helpers := readCredentialHelpers(t, filepath.Join(dockerDir, "config.json")) + assert.Equal(t, dockercredentials.HelperName, helpers[registryHost]) + + helperName := "docker-credential-databricks" + if runtime.GOOS == "windows" { + helperName += ".cmd" + } + _, err := os.Stat(filepath.Join(binDir, helperName)) + require.NoError(t, err) + assert.Contains(t, stderr.String(), registryHost) + assert.Contains(t, stderr.String(), filepath.ToSlash(filepath.Join(dockerDir, "config.json"))) +} + +func TestConfigureDockerCommandDocumentsRegionRequirement(t *testing.T) { + cmd := newDockerConfigureCommandWithDeps(defaultConfigureDockerDeps()) + + assert.Equal(t, "configure [PROFILE] --region REGION", cmd.Use) + assert.Contains(t, cmd.Flag("region").Usage, "workspace home region") +} + +func TestConfigureDockerCommandRequiresRegion(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "DEFAULT", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", "DEFAULT") + err := cmd.Execute() + assert.ErrorContains(t, err, "--region is required because workspace region cannot be inferred from this profile; it must match the workspace home region") +} + +func TestConfigureDockerCommandRejectsAccountOnlyProfile(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "account", + Host: "https://accounts.cloud.databricks.test", + AccountID: "acc", + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + + cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", "account", "--region", "us-west-2") + err := cmd.Execute() + assert.ErrorContains(t, err, "databricks auth login --host ") + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandPersistsResolvedWorkspaceID(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + homeDir := filepath.Join(dir, "home") + workspaceHost := "https://workspace.gcp.databricks.test" + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "workspace", + Host: workspaceHost, + WorkspaceID: authlib.WorkspaceIDNone, + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DATABRICKS_WORKSPACE_ID", "ambient-workspace") + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", homeDir) + + server := testserver.New(t) + server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { + assert.Empty(t, req.Headers.Get(authlib.WorkspaceIDHeader)) + return testserver.Response{ + Headers: http.Header{"X-Databricks-Org-Id": {"999999"}}, + Body: map[string]any{}, + } + }) + testserver.AddDefaultHandlers(server) + + deps := defaultConfigureDockerDeps() + databricksPath := writeConfigureDockerExecutable(t, filepath.Join(dir, "bin")) + deps.executable = func() (string, error) { + return databricksPath, nil + } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + assert.Equal(t, databricksPath, cfg.DatabricksCliPath) + require.NoError(t, (*config.Config)(cfg).EnsureResolved()) + assert.Equal(t, authlib.WorkspaceIDNone, cfg.WorkspaceID) + cfg.Host = server.URL + cfg.Token = "test-token" + cfg.AuthType = "pat" + cfg.Profile = "" + return databricks.NewWorkspaceClient(cfg) + } + deps.registryHost = configureDockerRegistryHostStub(t, "999999", "us-west-2", workspaceHost, "999999.container.us-west-2.gcp.databricks.test") + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "workspace", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) + + raw, err := os.ReadFile(configFile) + require.NoError(t, err) + assert.Contains(t, string(raw), "workspace_id = 999999") + + helpers := readCredentialHelpers(t, filepath.Join(dockerDir, "config.json")) + assert.Equal(t, dockercredentials.HelperName, helpers["999999.container.us-west-2.gcp.databricks.test"]) +} + +func TestConfigureDockerCommandRejectsUnsupportedWorkspaceHostBeforeProfileAndDockerConfigMutation(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "DEFAULT", + Host: "https://workspace.example.test", + AuthType: authlib.AuthTypeDatabricksCli, + }) + before, err := os.ReadFile(configFile) + require.NoError(t, err) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", filepath.Join(dir, "home")) + + deps := defaultConfigureDockerDeps() + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + return &databricks.WorkspaceClient{Config: (*config.Config)(cfg)}, nil + } + deps.resolveWorkspaceID = func(context.Context, *databricks.WorkspaceClient) (string, error) { + return "123456789", nil + } + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + t.Fatal("installShim should not be called") + return dockercredentials.ShimInstallResult{}, nil + } + deps.setCredentialHelper = func(string, string) error { + t.Fatal("setCredentialHelper should not be called") + return nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + err = cmd.Execute() + assert.ErrorContains(t, err, `"workspace.example.test" is not a supported Databricks workspace host`) + after, err := os.ReadFile(configFile) + require.NoError(t, err) + assert.Equal(t, string(before), string(after)) + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandRejectsUnsupportedAuthProfiles(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + homeDir := filepath.Join(dir, "home") + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "pat", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: "pat", + }) + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "m2m", + Host: "https://m2m.cloud.databricks.test", + WorkspaceID: "987654321", + ClientID: "client-id", + ClientSecret: "client-secret", + }) + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "blank-auth", + Host: "https://blank-auth.cloud.databricks.test", + WorkspaceID: "111222333", + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", homeDir) + + for _, profileName := range []string{"pat", "m2m", "blank-auth"} { + t.Run(profileName, func(t *testing.T) { + cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", profileName, "--region", "us-west-2") + err := cmd.Execute() + assert.ErrorContains(t, err, "requires a profile created by databricks auth login") + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) + }) + } +} + +func TestConfigureDockerCommandRejectsExplicitInheritedFlags(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "DEFAULT", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", filepath.Join(dir, "docker")) + t.Setenv("HOME", filepath.Join(dir, "home")) + + cases := [][]string{ + {"docker", "configure", "DEFAULT", "--region", "us-west-2", "--host", "https://other.cloud.databricks.test"}, + {"docker", "configure", "DEFAULT", "--region", "us-west-2", "--account-id", "abc"}, + {"docker", "configure", "DEFAULT", "--region", "us-west-2", "--workspace-id", "987654321"}, + } + + for _, args := range cases { + t.Run(args[len(args)-2], func(t *testing.T) { + cmd := newDockerConfigureTestCommand(ctx, args...) + err := cmd.Execute() + assert.ErrorContains(t, err, "is not supported for auth docker configure") + }) + } +} + +func TestConfigureDockerCommandRejectsAmbiguousWorkspaceIDBeforeDockerConfig(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + + for _, name := range []string{"one", "two"} { + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: name, + Host: "https://" + name + ".cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + } + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", filepath.Join(dir, "home")) + + deps := defaultConfigureDockerDeps() + deps.registryHost = func(workspaceID, region, _ string) (string, error) { + return workspaceID + ".container." + region + ".cloud.databricks.test", nil + } + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + t.Fatal("installShim should not be called") + return dockercredentials.ShimInstallResult{}, nil + } + deps.setCredentialHelper = func(string, string) error { + t.Fatal("setCredentialHelper should not be called") + return nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "one", "--region", "us-west-2") + err := cmd.Execute() + assert.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") + assert.ErrorContains(t, err, "Remove duplicate workspace_id entries") + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandRejectsSameWorkspaceIDInDifferentEnvironment(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "prod", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "dev", + Host: "https://workspace.dev.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + deps := defaultConfigureDockerDeps() + deps.executable = func() (string, error) { + return filepath.Join(dir, "databricks"), nil + } + deps.registryHost = func(workspaceID, region, workspaceHost string) (string, error) { + zone := ".cloud.databricks.test" + if workspaceHost == "https://workspace.dev.cloud.databricks.test" { + zone = ".dev.cloud.databricks.test" + } + return workspaceID + ".container." + region + zone, nil + } + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + t.Fatal("installShim should not be called") + return dockercredentials.ShimInstallResult{}, nil + } + deps.setCredentialHelper = func(string, string) error { + t.Fatal("setCredentialHelper should not be called") + return nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "prod", "--region", "us-west-2") + err := cmd.Execute() + assert.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") +} + +func TestConfigureDockerAllowsUnsupportedDuplicateProfile(t *testing.T) { + p := profile.Profile{ + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + } + profiler := profile.InMemoryProfiler{Profiles: profile.Profiles{ + p, + { + Name: "m2m", + Host: p.Host, + WorkspaceID: p.WorkspaceID, + HasClientCredentials: true, + }, + }} + err := ensureConfigureDockerUniqueProfile(t.Context(), profiler, p, p.WorkspaceID) + require.NoError(t, err) +} + +func TestConfigureDockerCommandInstallsShimBeforeDockerConfig(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + workspaceHost := "https://workspace.cloud.databricks.test" + + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "DEFAULT", + Host: workspaceHost, + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", filepath.Join(dir, "home")) + + deps := defaultConfigureDockerDeps() + deps.executable = func() (string, error) { + return "/usr/local/bin/databricks", nil + } + deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, "123456789.container.us-west-2.cloud.databricks.test") + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + return dockercredentials.ShimInstallResult{}, errors.New("install failed") + } + deps.setCredentialHelper = func(string, string) error { + t.Fatal("setCredentialHelper should not be called after install failure") + return nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + err := cmd.Execute() + assert.ErrorContains(t, err, "install failed") + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandWarnsAboutPATHAndPATHEXT(t *testing.T) { + ctx, stderr := cmdio.NewTestContextWithStderr(t.Context()) + workspaceHost := "https://workspace.cloud.databricks.test" + registryHost := "123456789.container.us-west-2.cloud.databricks.test" + t.Setenv("DOCKER_CONFIG", t.TempDir()) + + deps := defaultConfigureDockerDeps() + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{ + { + Name: "DEFAULT", + Host: workspaceHost, + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }, + }} + deps.executable = func() (string, error) { + return "/usr/local/bin/databricks", nil + } + deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, registryHost) + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + return dockercredentials.ShimInstallResult{ + Path: "/usr/local/bin/docker-credential-databricks", + OnPath: false, + }, nil + } + deps.setCredentialHelper = func(string, string) error { + return nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) + assert.Contains(t, stderr.String(), "PATH") + assert.Contains(t, stderr.String(), ".CMD is in PATHEXT on Windows") +} diff --git a/libs/dockercredentials/docker_config.go b/libs/dockercredentials/docker_config.go new file mode 100644 index 00000000000..5f95a2cc6f6 --- /dev/null +++ b/libs/dockercredentials/docker_config.go @@ -0,0 +1,122 @@ +package dockercredentials + +import ( + "encoding/json" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" +) + +// HelperName is the suffix Docker uses to resolve docker-credential-databricks. +const HelperName = "databricks" + +// SetCredentialHelper assigns docker-credential-databricks to registryHost without changing other Docker configuration. +// See https://docs.docker.com/reference/cli/docker/login/#credential-helpers. +func SetCredentialHelper(path, registryHost string) error { + path, err := resolveDockerConfigPath(path) + if err != nil { + return err + } + config, err := readDockerConfig(path) + if err != nil { + return err + } + + helpers := map[string]string{} + if raw, ok := config["credHelpers"]; ok { + if err := json.Unmarshal(raw, &helpers); err != nil { + return fmt.Errorf("read Docker config %s: %w", path, err) + } + } + if helpers == nil { + helpers = map[string]string{} + } + + if helpers[registryHost] == HelperName { + return nil + } + + helpers[registryHost] = HelperName + rawHelpers, err := json.Marshal(helpers) + if err != nil { + return err + } + config["credHelpers"] = rawHelpers + + return writeDockerConfig(path, config) +} + +// resolveDockerConfigPath follows a config symlink so replacement does not remove the link itself. +func resolveDockerConfigPath(path string) (string, error) { + info, err := os.Lstat(path) + if errors.Is(err, fs.ErrNotExist) { + return path, nil + } + if err != nil { + return "", fmt.Errorf("inspect Docker config %s: %w", path, err) + } + if info.Mode()&os.ModeSymlink == 0 { + return path, nil + } + + resolved, err := filepath.EvalSymlinks(path) + if err != nil { + return "", fmt.Errorf("resolve Docker config symlink %s: %w", path, err) + } + return resolved, nil +} + +func readDockerConfig(path string) (map[string]json.RawMessage, error) { + raw, err := os.ReadFile(path) + if errors.Is(err, fs.ErrNotExist) { + return map[string]json.RawMessage{}, nil + } + if err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + + var config map[string]json.RawMessage + if err := json.Unmarshal(raw, &config); err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + if config == nil { + config = map[string]json.RawMessage{} + } + return config, nil +} + +func writeDockerConfig(path string, config map[string]json.RawMessage) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("create Docker config directory %s: %w", dir, err) + } + + raw, err := json.MarshalIndent(config, "", " ") + if err != nil { + return err + } + raw = append(raw, '\n') + + tmp, err := os.CreateTemp(dir, ".config.json.*") + if err != nil { + return fmt.Errorf("create temporary Docker config in %s: %w", dir, err) + } + tmpPath := tmp.Name() + defer func() { + _ = os.Remove(tmpPath) + }() + + if _, err := tmp.Write(raw); err != nil { + _ = tmp.Close() + return fmt.Errorf("write temporary Docker config %s: %w", tmpPath, err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("close temporary Docker config %s: %w", tmpPath, err) + } + if err := os.Rename(tmpPath, path); err != nil { + return fmt.Errorf("write Docker config %s: %w", path, err) + } + return nil +} diff --git a/libs/dockercredentials/docker_config_test.go b/libs/dockercredentials/docker_config_test.go new file mode 100644 index 00000000000..1949339d4cc --- /dev/null +++ b/libs/dockercredentials/docker_config_test.go @@ -0,0 +1,145 @@ +package dockercredentials + +import ( + "encoding/json" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testRegistryHost = "123.container.us-west-2.cloud.databricks.test" + +func readDockerConfigForTest(t *testing.T, path string) map[string]any { + t.Helper() + + raw, err := os.ReadFile(path) + require.NoError(t, err) + + var got map[string]any + require.NoError(t, json.Unmarshal(raw, &got)) + return got +} + +func TestConfigureDockerCredentialHelperCreatesConfig(t *testing.T) { + path := filepath.Join(t.TempDir(), "docker", "config.json") + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + got := readDockerConfigForTest(t, path) + assert.Equal(t, map[string]any{ + testRegistryHost: HelperName, + }, got["credHelpers"]) + + info, err := os.Stat(path) + require.NoError(t, err) + if runtime.GOOS != "windows" { + assert.Equal(t, os.FileMode(0o600), info.Mode().Perm()) + } +} + +func TestConfigureDockerCredentialHelperPreservesExistingConfig(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + require.NoError(t, os.WriteFile(path, []byte(`{ + "auths": { + "registry.example.test": {"auth": "abc"} + }, + "credsStore": "desktop", + "credHelpers": { + "registry.example.test": "desktop" + }, + "experimental": "enabled" +}`), 0o600)) + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + got := readDockerConfigForTest(t, path) + assert.Equal(t, "desktop", got["credsStore"]) + assert.Equal(t, "enabled", got["experimental"]) + assert.Equal(t, map[string]any{ + "registry.example.test": "desktop", + testRegistryHost: HelperName, + }, got["credHelpers"]) + assert.Contains(t, got, "auths") +} + +func TestConfigureDockerCredentialHelperPreservesConfigSymlink(t *testing.T) { + dir := t.TempDir() + target := filepath.Join(dir, "target.json") + path := filepath.Join(dir, "config.json") + require.NoError(t, os.WriteFile(target, []byte(`{"credsStore":"desktop"}`), 0o600)) + if err := os.Symlink(target, path); err != nil { + t.Skipf("symlinks are unavailable: %v", err) + } + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + info, err := os.Lstat(path) + require.NoError(t, err) + assert.NotZero(t, info.Mode()&os.ModeSymlink) + got := readDockerConfigForTest(t, target) + assert.Equal(t, "desktop", got["credsStore"]) + assert.Equal(t, map[string]any{testRegistryHost: HelperName}, got["credHelpers"]) +} + +func TestConfigureDockerCredentialHelperIsIdempotent(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + require.NoError(t, os.WriteFile(path, []byte(`{ + "credHelpers": { + "123.container.us-west-2.cloud.databricks.test": "databricks" + } +}`), 0o600)) + + before, err := os.ReadFile(path) + require.NoError(t, err) + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + after, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, before, after) +} + +func TestConfigureDockerCredentialHelperReplacesExistingHelper(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + require.NoError(t, os.WriteFile(path, []byte(`{ + "credHelpers": { + "123.container.us-west-2.cloud.databricks.test": "desktop" + } +}`), 0o600)) + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + got := readDockerConfigForTest(t, path) + assert.Equal(t, map[string]any{ + testRegistryHost: HelperName, + }, got["credHelpers"]) +} + +func TestConfigureDockerCredentialHelperTreatsNullCredHelpersAsEmpty(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + require.NoError(t, os.WriteFile(path, []byte(`{"credHelpers": null}`), 0o600)) + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + got := readDockerConfigForTest(t, path) + assert.Equal(t, map[string]any{ + testRegistryHost: HelperName, + }, got["credHelpers"]) +} + +func TestConfigureDockerCredentialHelperRejectsInvalidJSON(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + require.NoError(t, os.WriteFile(path, []byte("{not valid json"), 0o600)) + + err := SetCredentialHelper(path, testRegistryHost) + assert.ErrorContains(t, err, "read Docker config") +} diff --git a/libs/dockercredentials/registry.go b/libs/dockercredentials/registry.go index cc3c6e7b864..34d69a2324b 100644 --- a/libs/dockercredentials/registry.go +++ b/libs/dockercredentials/registry.go @@ -22,6 +22,29 @@ type Registry struct { Host string } +// RegistryHost builds a registry host in the workspace's cloud and environment DNS zone. +func RegistryHost(workspaceID, region, workspaceHost string) (string, error) { + workspaceID = strings.TrimSpace(workspaceID) + region = strings.TrimSpace(region) + if workspaceID == "" { + return "", errors.New("workspace ID is required") + } + if region == "" { + return "", errors.New("region is required") + } + if !isDNSLabel(workspaceID) { + return "", fmt.Errorf("invalid workspace ID %q", workspaceID) + } + if !isDNSLabel(region) { + return "", fmt.Errorf("invalid region %q", region) + } + dnsZone, err := registryDNSZoneForWorkspaceHost(workspaceHost) + if err != nil { + return "", err + } + return fmt.Sprintf("%s.container.%s%s", workspaceID, region, dnsZone), nil +} + // 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) { @@ -75,6 +98,18 @@ func ParseRegistryHost(raw string) (Registry, error) { }, nil } +func registryDNSZoneForWorkspaceHost(raw string) (string, error) { + host, err := normalizeServerAddress(raw) + if err != nil { + return "", fmt.Errorf("parse workspace host: %w", err) + } + dnsZone, ok := matchingDatabricksDNSZone(host) + if !ok { + return "", fmt.Errorf("%q is not a supported Databricks workspace host", host) + } + return dnsZone, 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. diff --git a/libs/dockercredentials/registry_test.go b/libs/dockercredentials/registry_test.go index 071f93c6adf..acc216ea2c4 100644 --- a/libs/dockercredentials/registry_test.go +++ b/libs/dockercredentials/registry_test.go @@ -8,6 +8,73 @@ import ( "github.com/stretchr/testify/require" ) +func TestRegistryHost(t *testing.T) { + tests := []struct { + name string + workspaceHost string + region string + want string + }{ + { + name: "AWS production", + workspaceHost: "https://adb-123.456.cloud.databricks.com", + region: "us-west-2", + want: "123456789.container.us-west-2.cloud.databricks.com", + }, + { + name: "AWS staging", + workspaceHost: "https://workspace.staging.cloud.databricks.com", + region: "us-west-2", + want: "123456789.container.us-west-2.staging.cloud.databricks.com", + }, + { + name: "Azure production", + workspaceHost: "https://adb-123.456.azuredatabricks.net", + region: "eastus", + want: "123456789.container.eastus.azuredatabricks.net", + }, + { + name: "Azure development", + workspaceHost: "https://workspace.dev.azuredatabricks.net", + region: "eastus", + want: "123456789.container.eastus.dev.azuredatabricks.net", + }, + { + name: "GCP production", + workspaceHost: "https://workspace.gcp.databricks.com", + region: "us-central1", + want: "123456789.container.us-central1.gcp.databricks.com", + }, + { + name: "GCP development", + workspaceHost: "https://workspace.dev.gcp.databricks.com", + region: "us-central1", + want: "123456789.container.us-central1.dev.gcp.databricks.com", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := RegistryHost("123456789", tt.region, tt.workspaceHost) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestRegistryHostRejectsEmptyParts(t *testing.T) { + _, err := RegistryHost("", "us-west-2", "https://workspace.cloud.databricks.test") + assert.ErrorContains(t, err, "workspace ID is required") + + _, err = RegistryHost("123456789", "", "https://workspace.cloud.databricks.test") + assert.ErrorContains(t, err, "region is required") +} + +func TestRegistryHostRejectsUnsupportedWorkspaceHost(t *testing.T) { + _, err := RegistryHost("123456789", "us-west-2", "https://workspace.example.test") + assert.ErrorContains(t, err, `"workspace.example.test" is not a supported Databricks workspace host`) +} + func TestParseRegistryHost(t *testing.T) { cases := []string{ "123456789.container.us-west-2.cloud.databricks.com", @@ -35,6 +102,10 @@ func TestParseRegistryHostSupportsAllDatabricksEnvironmentZones(t *testing.T) { } t.Run(dnsZone, func(t *testing.T) { wantHost := "123456789.container.test-region" + dnsZone + got, err := RegistryHost("123456789", "test-region", "https://workspace"+dnsZone) + require.NoError(t, err) + assert.Equal(t, wantHost, got) + registry, err := ParseRegistryHost("https://" + wantHost + "/v2/") require.NoError(t, err) assert.Equal(t, Registry{ diff --git a/libs/dockercredentials/shim.go b/libs/dockercredentials/shim.go new file mode 100644 index 00000000000..d11e0aa1ed5 --- /dev/null +++ b/libs/dockercredentials/shim.go @@ -0,0 +1,142 @@ +package dockercredentials + +import ( + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" +) + +// ShimInstallResult reports where the helper was installed and whether Docker can resolve it from PATH. +type ShimInstallResult struct { + Path string + // OnPath reports whether Path is the first matching helper in the current PATH, using PATHEXT on Windows. + OnPath bool +} + +// InstallShim installs the Docker credential helper wrapper next to the Databricks CLI. +func InstallShim(databricksPath string) (ShimInstallResult, error) { + return installShimForGOOS(databricksPath, runtime.GOOS) +} + +func installShimForGOOS(databricksPath, goos string) (ShimInstallResult, error) { + if strings.TrimSpace(databricksPath) == "" { + return ShimInstallResult{}, errors.New("databricks executable path is required") + } + installDir := filepath.Dir(databricksPath) + + if err := os.MkdirAll(installDir, 0o755); err != nil { + return ShimInstallResult{}, fmt.Errorf("create Docker credential helper directory %s: %w", installDir, err) + } + + path := filepath.Join(installDir, shimFilename(goos)) + mode := os.FileMode(0o755) + script := shimScript(databricksPath) + if goos == "windows" { + mode = 0o644 + script = cmdShimScript(databricksPath) + } + if err := writeShimFile(path, []byte(script), mode); err != nil { + return ShimInstallResult{}, fmt.Errorf("write Docker credential helper %s: %w", path, err) + } + + return ShimInstallResult{ + Path: path, + OnPath: helperOnPathForGOOS(path, goos, exec.LookPath), + }, nil +} + +func shimFilename(goos string) string { + if goos == "windows" { + return "docker-credential-" + HelperName + ".cmd" + } + return "docker-credential-" + HelperName +} + +// shimScript accepts only get and sends CLI logs to stderr because Docker parses stdout as credential-helper JSON. +func shimScript(databricksPath string) string { + return fmt.Sprintf(`#!/bin/sh +if [ "$#" -ne 1 ] || [ "$1" != "get" ]; then + echo "docker-credential-databricks only supports get" >&2 + exit 1 +fi +shift +export DATABRICKS_LOG_FILE=stderr +exec %s auth docker token +`, posixShellQuote(databricksPath)) +} + +func cmdShimScript(databricksPath string) string { + // Percent signs must be doubled so cmd.exe treats them literally in executable paths. + script := fmt.Sprintf(`@echo off +setlocal EnableExtensions DisableDelayedExpansion +if not "%%~1"=="get" goto unsupported +if not "%%~2"=="" goto unsupported +set DATABRICKS_LOG_FILE=stderr +"%s" auth docker token +exit /b + +:unsupported +echo docker-credential-databricks only supports get 1>&2 +exit /b 1 +`, strings.ReplaceAll(databricksPath, "%", "%%")) + return strings.ReplaceAll(script, "\n", "\r\n") +} + +// writeShimFile preserves an existing helper if writing its replacement fails. +func writeShimFile(path string, script []byte, mode os.FileMode) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*") + if err != nil { + return err + } + tmpPath := tmp.Name() + defer func() { + _ = os.Remove(tmpPath) + }() + + if _, err := tmp.Write(script); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Chmod(mode); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + return os.Rename(tmpPath, path) +} + +func posixShellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +// helperOnPathForGOOS mirrors Docker's extensionless helper lookup through PATHEXT on Windows. +func helperOnPathForGOOS(helperPath, goos string, lookPath func(string) (string, error)) bool { + name := filepath.Base(helperPath) + if goos == "windows" { + name = strings.TrimSuffix(name, filepath.Ext(name)) + } + candidate, err := lookPath(name) + if err != nil { + return false + } + return samePath(candidate, helperPath) +} + +func samePath(a, b string) bool { + aInfo, aErr := os.Stat(a) + if aErr != nil { + return false + } + bInfo, bErr := os.Stat(b) + if bErr != nil { + return false + } + return os.SameFile(aInfo, bInfo) +} diff --git a/libs/dockercredentials/shim_test.go b/libs/dockercredentials/shim_test.go new file mode 100644 index 00000000000..aa828f8b559 --- /dev/null +++ b/libs/dockercredentials/shim_test.go @@ -0,0 +1,280 @@ +package dockercredentials + +import ( + "bytes" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func writeTestDatabricksExecutable(t *testing.T, dir string) string { + t.Helper() + name := "databricks" + if runtime.GOOS == "windows" { + name += ".exe" + } + path := filepath.Join(dir, name) + require.NoError(t, os.WriteFile(path, []byte("databricks executable"), 0o755)) + return path +} + +func TestShimFilename(t *testing.T) { + assert.Equal(t, "docker-credential-databricks", shimFilename("linux")) + assert.Equal(t, "docker-credential-databricks.cmd", shimFilename("windows")) +} + +func TestInstallWindowsShimWritesCommandScript(t *testing.T) { + installDir := t.TempDir() + databricksPath := filepath.Join(installDir, "databricks.exe") + require.NoError(t, os.WriteFile(databricksPath, []byte("databricks executable"), 0o755)) + t.Setenv("PATH", installDir) + + got, err := installShimForGOOS(databricksPath, "windows") + require.NoError(t, err) + assert.Equal(t, filepath.Join(installDir, "docker-credential-databricks.cmd"), got.Path) + + raw, err := os.ReadFile(got.Path) + require.NoError(t, err) + assert.NotEqual(t, "databricks executable", string(raw)) + assert.Contains(t, string(raw), "auth docker token") +} + +func TestUnixShimExecutesOnlyGetAndForcesLogsToStderr(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix shell shim test") + } + + dir := t.TempDir() + argsPath := filepath.Join(dir, "args") + envPath := filepath.Join(dir, "env") + stdinPath := filepath.Join(dir, "stdin") + fakeDir := filepath.Join(dir, "bin$DATABRICKS_LOG_FILE") + require.NoError(t, os.MkdirAll(fakeDir, 0o755)) + fakeDatabricks := filepath.Join(fakeDir, "data'bricks") + require.NoError(t, os.WriteFile(fakeDatabricks, []byte(`#!/bin/sh +printf '%s' "$*" > "$FAKE_ARGS_FILE" +printf '%s' "$DATABRICKS_LOG_FILE" > "$FAKE_ENV_FILE" +cat > "$FAKE_STDIN_FILE" +printf '{"Username":"oauthtoken","Secret":"secret"}\n' +`), 0o755)) + require.NoError(t, os.Chmod(fakeDatabricks, 0o755)) + + shim := filepath.Join(dir, "docker-credential-databricks") + require.NoError(t, os.WriteFile(shim, []byte(shimScript(fakeDatabricks)), 0o755)) + require.NoError(t, os.Chmod(shim, 0o755)) + + cmd := exec.Command(shim, "get") + cmd.Stdin = bytes.NewBufferString("registry-host") + cmd.Env = append(os.Environ(), + "DATABRICKS_LOG_FILE=stdout", + "FAKE_ARGS_FILE="+argsPath, + "FAKE_ENV_FILE="+envPath, + "FAKE_STDIN_FILE="+stdinPath, + ) + out, err := cmd.Output() + require.NoError(t, err) + assert.JSONEq(t, `{"Username":"oauthtoken","Secret":"secret"}`, string(out)) + + rawArgs, err := os.ReadFile(argsPath) + require.NoError(t, err) + assert.Equal(t, "auth docker token", string(rawArgs)) + + rawEnv, err := os.ReadFile(envPath) + require.NoError(t, err) + assert.Equal(t, "stderr", string(rawEnv)) + + rawStdin, err := os.ReadFile(stdinPath) + require.NoError(t, err) + assert.Equal(t, "registry-host", string(rawStdin)) + + err = exec.Command(shim, "store").Run() + assert.Error(t, err) + err = exec.Command(shim, "get", "store").Run() + assert.Error(t, err) +} + +func TestWindowsShimExecutesOnlyGetAndForcesLogsToStderr(t *testing.T) { + if runtime.GOOS != "windows" { + t.Skip("Windows command shim test") + } + + dir := t.TempDir() + argsPath := filepath.Join(dir, "args") + envPath := filepath.Join(dir, "env") + stdinPath := filepath.Join(dir, "stdin") + fakeDir := filepath.Join(dir, "bin%DOCKER_SHIM_UNSET%") + require.NoError(t, os.Mkdir(fakeDir, 0o755)) + fakeDatabricks := filepath.Join(fakeDir, "databricks.cmd") + require.NoError(t, os.WriteFile(fakeDatabricks, []byte(`@echo off +set /p registry= +> "%FAKE_ARGS_FILE%" echo %* +> "%FAKE_ENV_FILE%" echo %DATABRICKS_LOG_FILE% +> "%FAKE_STDIN_FILE%" echo %registry% +echo {"Username":"oauthtoken","Secret":"secret"} +`), 0o644)) + + shim := filepath.Join(dir, "docker-credential-databricks.cmd") + require.NoError(t, os.WriteFile(shim, []byte(cmdShimScript(fakeDatabricks)), 0o644)) + + cmd := exec.Command(shim, "get") + cmd.Stdin = bytes.NewBufferString("registry-host\n") + cmd.Env = append(os.Environ(), + "DATABRICKS_LOG_FILE=stdout", + "FAKE_ARGS_FILE="+argsPath, + "FAKE_ENV_FILE="+envPath, + "FAKE_STDIN_FILE="+stdinPath, + ) + out, err := cmd.Output() + require.NoError(t, err) + assert.JSONEq(t, `{"Username":"oauthtoken","Secret":"secret"}`, string(out)) + + rawArgs, err := os.ReadFile(argsPath) + require.NoError(t, err) + assert.Equal(t, "auth docker token", strings.TrimSpace(string(rawArgs))) + + rawEnv, err := os.ReadFile(envPath) + require.NoError(t, err) + assert.Equal(t, "stderr", strings.TrimSpace(string(rawEnv))) + + rawStdin, err := os.ReadFile(stdinPath) + require.NoError(t, err) + assert.Equal(t, "registry-host", strings.TrimSpace(string(rawStdin))) + + assert.Error(t, exec.Command(shim, "store").Run()) + assert.Error(t, exec.Command(shim, "get", "store").Run()) +} + +func TestWindowsShimDisablesInheritedDelayedExpansion(t *testing.T) { + if runtime.GOOS != "windows" { + t.Skip("Windows command shim test") + } + + dir := t.TempDir() + fakeDir := filepath.Join(dir, "bin!DOCKER_SHIM_UNSET!") + require.NoError(t, os.Mkdir(fakeDir, 0o755)) + fakeDatabricks := filepath.Join(fakeDir, "databricks.cmd") + require.NoError(t, os.WriteFile(fakeDatabricks, []byte("@echo {\"Username\":\"oauthtoken\",\"Secret\":\"secret\"}\r\n"), 0o644)) + + shim := filepath.Join(dir, "docker-credential-databricks.cmd") + require.NoError(t, os.WriteFile(shim, []byte(cmdShimScript(fakeDatabricks)), 0o644)) + + cmd := exec.Command("cmd.exe", "/D", "/V:ON", "/C", "call", shim, "get") + cmd.Stdin = bytes.NewBufferString("registry-host\n") + out, err := cmd.Output() + require.NoError(t, err) + assert.JSONEq(t, `{"Username":"oauthtoken","Secret":"secret"}`, string(out)) +} + +func TestInstallShimReportsPathStatus(t *testing.T) { + dir := t.TempDir() + databricksPath := writeTestDatabricksExecutable(t, dir) + t.Setenv("PATH", dir) + + got, err := InstallShim(databricksPath) + require.NoError(t, err) + assert.Equal(t, filepath.Join(dir, shimFilename(runtime.GOOS)), got.Path) + assert.True(t, got.OnPath) + + info, err := os.Stat(got.Path) + require.NoError(t, err) + if runtime.GOOS != "windows" { + assert.Equal(t, os.FileMode(0o755), info.Mode().Perm()) + } +} + +func TestInstallShimReportsNotOnPath(t *testing.T) { + dir := t.TempDir() + databricksPath := writeTestDatabricksExecutable(t, dir) + t.Setenv("PATH", t.TempDir()) + + got, err := InstallShim(databricksPath) + require.NoError(t, err) + assert.Equal(t, filepath.Join(dir, shimFilename(runtime.GOOS)), got.Path) + assert.False(t, got.OnPath) +} + +func TestInstallShimReportsNotOnPathWhenHelperIsShadowed(t *testing.T) { + installDir := t.TempDir() + databricksPath := writeTestDatabricksExecutable(t, installDir) + shadowDir := t.TempDir() + shadowPath := filepath.Join(shadowDir, shimFilename(runtime.GOOS)) + require.NoError(t, os.WriteFile(shadowPath, []byte("shadow"), 0o755)) + require.NoError(t, os.Chmod(shadowPath, 0o755)) + t.Setenv("PATH", shadowDir+string(os.PathListSeparator)+installDir) + + got, err := InstallShim(databricksPath) + require.NoError(t, err) + assert.Equal(t, filepath.Join(installDir, shimFilename(runtime.GOOS)), got.Path) + assert.False(t, got.OnPath) +} + +func TestHelperOnPathUsesDockerLookupName(t *testing.T) { + dir := t.TempDir() + helperPath := filepath.Join(dir, "docker-credential-databricks.cmd") + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o755)) + + var gotName string + found := helperOnPathForGOOS(helperPath, "windows", func(name string) (string, error) { + gotName = name + return helperPath, nil + }) + + assert.True(t, found) + assert.Equal(t, "docker-credential-databricks", gotName) +} + +func TestHelperOnPathRejectsEmptyUnixPathEntry(t *testing.T) { + dir := t.TempDir() + helperPath := filepath.Join(dir, shimFilename(runtime.GOOS)) + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o755)) + + t.Chdir(dir) + t.Setenv("PATH", "") + assert.False(t, helperOnPathForGOOS(helperPath, runtime.GOOS, exec.LookPath)) +} + +func TestSamePathUsesFileIdentity(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink test") + } + + dir := t.TempDir() + target := filepath.Join(dir, "docker-credential-databricks") + link := filepath.Join(dir, "helper-link") + require.NoError(t, os.WriteFile(target, []byte("helper"), 0o755)) + require.NoError(t, os.Symlink(target, link)) + + assert.True(t, samePath(target, link)) +} + +func TestInstallShimDoesNotTruncateExistingHelperWhenTempCreateFails(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("permission-forced failure test") + } + if os.Geteuid() == 0 { + t.Skip("permission-forced failure test requires a non-root user") + } + + installDir := filepath.Join(t.TempDir(), "missing") + require.NoError(t, os.MkdirAll(installDir, 0o755)) + databricksPath := writeTestDatabricksExecutable(t, installDir) + shimPath := filepath.Join(installDir, shimFilename(runtime.GOOS)) + require.NoError(t, os.WriteFile(shimPath, []byte("existing helper"), 0o755)) + require.NoError(t, os.Chmod(installDir, 0o500)) + t.Cleanup(func() { + _ = os.Chmod(installDir, 0o755) + }) + + _, err := InstallShim(databricksPath) + assert.Error(t, err) + + raw, readErr := os.ReadFile(shimPath) + require.NoError(t, readErr) + assert.Equal(t, "existing helper", string(raw)) +}