diff --git a/.nextchanges/cli/configure-docker.md b/.nextchanges/cli/configure-docker.md new file mode 100644 index 00000000000..80b374df4de --- /dev/null +++ b/.nextchanges/cli/configure-docker.md @@ -0,0 +1 @@ +Added `databricks auth configure-docker` to configure Docker credential helper access for Databricks Artifact Registry. diff --git a/acceptance/cmd/auth/configure-docker-help/out.test.toml b/acceptance/cmd/auth/configure-docker-help/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/cmd/auth/configure-docker-help/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/cmd/auth/configure-docker-help/output.txt b/acceptance/cmd/auth/configure-docker-help/output.txt new file mode 100644 index 00000000000..6c1a8e581f0 --- /dev/null +++ b/acceptance/cmd/auth/configure-docker-help/output.txt @@ -0,0 +1,27 @@ + +>>> [CLI] auth configure-docker --help +Configure Docker authentication for Databricks Artifact Registry. + +This command installs docker-credential-databricks and configures Docker to use +it for the selected workspace's Artifact Registry host. If the selected profile +does not already include a workspace_id, the command resolves and saves it so +the Docker helper can map the registry host back to the profile. The required +region must match the workspace home region because it cannot be inferred from +the profile. Select the workspace with [PROFILE] or --profile; --host, +--account-id, and --workspace-id are not supported. + +Usage: + databricks auth configure-docker [PROFILE] --region REGION [flags] + +Flags: + -h, --help help for configure-docker + --region string Cloud region for the Databricks Artifact Registry host; must match the workspace home region + +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 diff --git a/acceptance/cmd/auth/configure-docker-help/script b/acceptance/cmd/auth/configure-docker-help/script new file mode 100644 index 00000000000..866b6a740cb --- /dev/null +++ b/acceptance/cmd/auth/configure-docker-help/script @@ -0,0 +1 @@ +trace $CLI auth configure-docker --help diff --git a/acceptance/cmd/auth/configure-docker-help/test.toml b/acceptance/cmd/auth/configure-docker-help/test.toml new file mode 100644 index 00000000000..9609e1af299 --- /dev/null +++ b/acceptance/cmd/auth/configure-docker-help/test.toml @@ -0,0 +1 @@ +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 7ef3a9f72ac..6d803b321db 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(newConfigureDockerCommand()) cmd.AddCommand(newDescribeCommand()) cmd.AddCommand(newSwitchCommand()) return cmd diff --git a/cmd/auth/configure_docker.go b/cmd/auth/configure_docker.go new file mode 100644 index 00000000000..315aade8faf --- /dev/null +++ b/cmd/auth/configure_docker.go @@ -0,0 +1,286 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + 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/env" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" +) + +// configureDockerDeps groups injectable profile reads, workspace resolution, executable discovery, and Docker operations. +type configureDockerDeps struct { + profiler profile.Profiler + newWorkspaceClient func(*databricks.Config) (*databricks.WorkspaceClient, error) + resolveWorkspaceID func(context.Context, *databricks.WorkspaceClient) (string, error) + executable func() (string, error) + registryHost func(string, string, string) (string, error) + installShim func(string, string) (dockercredentials.ShimInstallResult, error) + setCredentialHelper func(string, string) error +} + +// defaultConfigureDockerDeps provides production implementations for the command's injectable dependencies. +func defaultConfigureDockerDeps() configureDockerDeps { + return configureDockerDeps{ + profiler: profile.DefaultProfiler, + newWorkspaceClient: func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + return databricks.NewWorkspaceClient(cfg) + }, + resolveWorkspaceID: authlib.ResolveWorkspaceID, + executable: os.Executable, + registryHost: dockercredentials.RegistryHost, + installShim: dockercredentials.InstallShim, + setCredentialHelper: dockercredentials.SetCredentialHelper, + } +} + +// newConfigureDockerCommand is the production entry point; tests use the dependency-injected constructor. +func newConfigureDockerCommand() *cobra.Command { + return newConfigureDockerCommandWithDeps(defaultConfigureDockerDeps()) +} + +// newConfigureDockerCommandWithDeps accepts replacements for profile reads, workspace resolution, and Docker operations. +func newConfigureDockerCommandWithDeps(deps configureDockerDeps) *cobra.Command { + cmd := &cobra.Command{ + Use: "configure-docker [PROFILE] --region REGION", + Short: "Configure Docker authentication for Databricks Artifact Registry", + Long: `Configure Docker authentication for Databricks Artifact Registry. + +This command installs docker-credential-databricks and configures Docker to use +it for the selected workspace's Artifact Registry host. If the selected profile +does not already include a workspace_id, the command resolves and saves it so +the Docker helper can map the registry host back to the profile. The required +region must match the workspace home region because it cannot be inferred from +the profile. Select the workspace with [PROFILE] or --profile; --host, +--account-id, and --workspace-id are not supported.`, + Args: cobra.MaximumNArgs(1), + } + var region string + cmd.Flags().StringVar(®ion, "region", "", "Cloud region for the Databricks Artifact Registry host; must match the workspace home region") + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + if err := errorOnUnsupportedConfigureDockerFlags(cmd); err != nil { + return err + } + // Workspace profiles do not expose the home region needed for the registry hostname. + if region == "" { + return errors.New("--region is required because workspace region cannot be inferred from this profile; it must match the workspace home region") + } + + profileName, err := configureDockerProfileName(ctx, cmd, args, deps.profiler) + if err != nil { + return err + } + + p, err := loadAndValidateConfigureDockerProfile(ctx, profileName, deps.profiler) + if err != nil { + return err + } + + executable, err := deps.executable() + if err != nil { + return fmt.Errorf("locate databricks executable: %w", err) + } + workspaceID, err := resolveConfigureDockerWorkspaceID(ctx, p, executable, deps) + if err != nil { + return err + } + // The workspace host supplies the cloud and environment DNS zone for the registry hostname. + registryHost, err := deps.registryHost(workspaceID, region, p.Host) + if err != nil { + return err + } + if err := ensureConfigureDockerUniqueProfile(ctx, deps.profiler, p, workspaceID, region, registryHost, deps.registryHost); err != nil { + return err + } + if p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone { + if err := persistConfigureDockerWorkspaceID(ctx, p, workspaceID); err != nil { + return fmt.Errorf("save workspace ID to profile %q: %w", p.Name, err) + } + } + + // Installing beside this CLI lets an existing PATH entry discover both executables. + installDir := filepath.Dir(executable) + shim, err := deps.installShim(executable, installDir) + if err != nil { + return fmt.Errorf("install Docker credential helper: %w", err) + } + dockerConfigPath, err := configureDockerConfigPath(ctx) + if err != nil { + return err + } + if err := deps.setCredentialHelper(dockerConfigPath, registryHost); err != nil { + return fmt.Errorf("update Docker config %s: %w", dockerConfigPath, err) + } + + cmdio.LogString(ctx, "Configured Docker credential helper for "+registryHost) + cmdio.LogString(ctx, "Updated Docker config: "+dockerConfigPath) + cmdio.LogString(ctx, "Installed Docker credential helper: "+shim.Path) + if !shim.OnPath { + cmdio.LogString(ctx, fmt.Sprintf("Warning: ensure %s is on PATH before any other docker-credential-databricks helper, and that .EXE is in PATHEXT on Windows", installDir)) + } + return nil + } + + return cmd +} + +// errorOnUnsupportedConfigureDockerFlags rejects inherited selectors that bypass the durable profile-to-registry mapping. +func errorOnUnsupportedConfigureDockerFlags(cmd *cobra.Command) error { + for _, name := range []string{"host", "account-id", "workspace-id"} { + flag := cmd.Flag(name) + if flag != nil && flag.Changed { + return fmt.Errorf("--%s is not supported for configure-docker. Select the workspace with [PROFILE] or --profile instead", name) + } + } + return nil +} + +// configureDockerProfileName resolves an explicit profile before environment, default-profile, and interactive selection. +func configureDockerProfileName(ctx context.Context, cmd *cobra.Command, args []string, profiler profile.Profiler) (string, error) { + profileFlag := cmd.Flag("profile") + profileName := "" + if profileFlag != nil { + profileName = profileFlag.Value.String() + } + if len(args) == 1 { + if profileName != "" { + return "", fmt.Errorf("argument %q cannot be combined with --profile. Use --profile instead", args[0]) + } + return args[0], nil + } + if profileName != "" { + return profileName, nil + } + if profileName = env.Get(ctx, "DATABRICKS_CONFIG_PROFILE"); profileName != "" { + return profileName, nil + } + if profileName = databrickscfg.ResolveDefaultProfile(ctx); profileName != "" { + return profileName, nil + } + if !cmdio.IsPromptSupported(ctx) { + return "", errors.New("no profile specified. Use --profile 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")) + result, selected, err := pickAuthProfile(ctx, profiles, profilePickerOptions{ + Label: "Select a workspace profile", + Default: currentDefault, + }) + if err != nil { + return "", err + } + if result != profilePickerProfile { + return "", errors.New("no profile selected") + } + return selected, nil +} + +// loadAndValidateConfigureDockerProfile loads one named profile and checks that its metadata is eligible for workspace U2M authentication. +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 +} + +// resolveConfigureDockerWorkspaceID queries /Me when the profile has no usable ID without allowing ambient routing or the "none" sentinel into the request. +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 profiles that resolve to the same registry host because workspace IDs can repeat across environments. +func ensureConfigureDockerUniqueProfile(ctx context.Context, profiler profile.Profiler, p profile.Profile, workspaceID, region, selectedRegistryHost string, registryHost registryHostResolver) error { + matches, err := profiler.LoadProfiles(ctx, func(candidate profile.Profile) bool { + return candidate.WorkspaceID == workspaceID + }) + if err != nil { + return err + } + + var names []string + for _, candidate := range matches { + if validateDockerCredentialProfile(candidate) != nil { + continue + } + candidateRegistryHost, err := registryHost(workspaceID, region, candidate.Host) + if err == nil && candidateRegistryHost == selectedRegistryHost { + names = append(names, candidate.Name) + } + } + if p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone { + names = append(names, p.Name) + } + 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 ")) +} + +// persistConfigureDockerWorkspaceID adds the resolved ID to the selected profile without replacing its other settings. +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, + }) +} + +// configureDockerConfigPath honors Docker's DOCKER_CONFIG override before the per-user default. +// See https://docs.docker.com/reference/cli/docker/#configuration-files. +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/configure_docker_test.go b/cmd/auth/configure_docker_test.go new file mode 100644 index 00000000000..c38a7887e62 --- /dev/null +++ b/cmd/auth/configure_docker_test.go @@ -0,0 +1,533 @@ +package auth + +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" +) + +// newConfigureDockerTestCommand uses the production auth tree to cover command registration and inherited flags. +func newConfigureDockerTestCommand(ctx context.Context, args ...string) *cobra.Command { + cmd := New() + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs(args) + return cmd +} + +// newConfigureDockerTestCommandWithDeps retains inherited auth flags while replacing injectable command dependencies. +func newConfigureDockerTestCommandWithDeps(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") + cmd.AddCommand(newConfigureDockerCommandWithDeps(deps)) + cmd.SetContext(ctx) + cmd.SetArgs(args) + return cmd +} + +// writeConfigureDockerProfile uses production profile merging so fixtures match persisted config behavior. +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)) +} + +// readCredentialHelpers decodes only the mapping relevant to command-level assertions. +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 +} + +// configureDockerRegistryHostStub asserts every derivation input before returning a deterministic host. +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) { + require.Equal(t, wantWorkspaceID, workspaceID) + require.Equal(t, wantRegion, region) + require.Equal(t, wantWorkspaceHost, workspaceHost) + return registryHost, nil + } +} + +// writeConfigureDockerExecutable uses the platform suffix so shim installation exercises Windows path handling. +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: 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 := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "--profile", "DEFAULT", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) + + helpers := readCredentialHelpers(t, filepath.Join(dockerDir, "config.json")) + require.Equal(t, dockercredentials.HelperName, helpers[registryHost]) + + helperName := "docker-credential-databricks" + if runtime.GOOS == "windows" { + helperName += ".exe" + } + _, err := os.Stat(filepath.Join(binDir, helperName)) + require.NoError(t, err) + assert.Contains(t, stderr.String(), registryHost) + assert.Contains(t, stderr.String(), filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandDocumentsRegionRequirement(t *testing.T) { + cmd := newConfigureDockerCommandWithDeps(defaultConfigureDockerDeps()) + + require.Equal(t, "configure-docker [PROFILE] --region REGION", cmd.Use) + require.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: authTypeDatabricksCLI, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + cmd := newConfigureDockerTestCommand(ctx, "configure-docker", "DEFAULT") + err := cmd.Execute() + require.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: authTypeDatabricksCLI, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + + cmd := newConfigureDockerTestCommand(ctx, "configure-docker", "account", "--region", "us-west-2") + err := cmd.Execute() + require.ErrorContains(t, err, "databricks auth login --host ") + require.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: 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 { + require.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) { + require.Equal(t, databricksPath, cfg.DatabricksCliPath) + require.NoError(t, (*config.Config)(cfg).EnsureResolved()) + require.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 := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "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")) + require.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: 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, 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 := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "DEFAULT", "--region", "us-west-2") + err = cmd.Execute() + require.ErrorContains(t, err, `"workspace.example.test" is not a supported Databricks workspace host`) + after, err := os.ReadFile(configFile) + require.NoError(t, err) + require.Equal(t, string(before), string(after)) + require.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 := newConfigureDockerTestCommand(ctx, "configure-docker", profileName, "--region", "us-west-2") + err := cmd.Execute() + require.ErrorContains(t, err, "requires a profile created by databricks auth login") + require.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: authTypeDatabricksCLI, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", filepath.Join(dir, "docker")) + t.Setenv("HOME", filepath.Join(dir, "home")) + + cases := [][]string{ + {"configure-docker", "DEFAULT", "--region", "us-west-2", "--host", "https://other.cloud.databricks.test"}, + {"configure-docker", "DEFAULT", "--region", "us-west-2", "--account-id", "abc"}, + {"configure-docker", "DEFAULT", "--region", "us-west-2", "--workspace-id", "987654321"}, + } + + for _, args := range cases { + t.Run(args[len(args)-2], func(t *testing.T) { + cmd := newConfigureDockerTestCommand(ctx, args...) + err := cmd.Execute() + require.ErrorContains(t, err, "is not supported for configure-docker") + }) + } +} + +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: 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, 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 := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "one", "--region", "us-west-2") + err := cmd.Execute() + require.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") + require.ErrorContains(t, err, "Remove duplicate workspace_id entries") + require.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandAllowsSameWorkspaceIDInDifferentEnvironment(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: authTypeDatabricksCLI, + }) + writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ + Profile: "dev", + Host: "https://workspace.dev.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: 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, string) (dockercredentials.ShimInstallResult, error) { + return dockercredentials.ShimInstallResult{}, nil + } + var configuredHost string + deps.setCredentialHelper = func(_, registryHost string) error { + configuredHost = registryHost + return nil + } + + cmd := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "prod", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) + require.Equal(t, "123456789.container.us-west-2.cloud.databricks.test", configuredHost) +} + +func TestConfigureDockerAllowsUnsupportedDuplicateProfile(t *testing.T) { + p := profile.Profile{ + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + } + profiler := profile.InMemoryProfiler{Profiles: profile.Profiles{ + p, + { + Name: "m2m", + Host: p.Host, + WorkspaceID: p.WorkspaceID, + HasClientCredentials: true, + }, + }} + registryHost := func(workspaceID, region, _ string) (string, error) { + return workspaceID + ".container." + region + ".cloud.databricks.test", nil + } + + err := ensureConfigureDockerUniqueProfile(t.Context(), profiler, p, p.WorkspaceID, "us-west-2", "123456789.container.us-west-2.cloud.databricks.test", registryHost) + 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: 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, 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 := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "DEFAULT", "--region", "us-west-2") + err := cmd.Execute() + require.ErrorContains(t, err, "install failed") + require.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: 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, 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 := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "DEFAULT", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) + require.Contains(t, stderr.String(), "PATH") + require.Contains(t, stderr.String(), ".EXE is in PATHEXT on Windows") +} diff --git a/cmd/auth/docker_profile.go b/cmd/auth/docker_profile.go new file mode 100644 index 00000000000..8c3a7a0798b --- /dev/null +++ b/cmd/auth/docker_profile.go @@ -0,0 +1,35 @@ +package auth + +import ( + "fmt" + + authlib "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/databricks-sdk-go/config" +) + +// validateDockerCredentialProfile requires metadata eligible for workspace-scoped U2M authentication. +func validateDockerCredentialProfile(p profile.Profile) error { + if p.HasClientCredentials { + return fmt.Errorf("profile %q uses client credentials. Docker credential helper requires a profile created by databricks auth login", p.Name) + } + if p.AuthType != authTypeDatabricksCLI { + return fmt.Errorf("profile %q uses auth_type %q. Docker credential helper requires a profile created by databricks auth login", p.Name, p.AuthType) + } + if isDockerCredentialAccountOnlyProfile(p) { + return fmt.Errorf("profile %q does not target a workspace. Run databricks auth login --host and retry with that profile", p.Name) + } + return nil +} + +// isDockerCredentialAccountOnlyProfile treats classic account hosts and unrouted account profiles as unsafe for workspace requests. +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/token.go b/cmd/auth/token.go index d5e88e64d72..782c9329e31 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -16,6 +16,7 @@ import ( "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/env" "github.com/databricks/cli/libs/flags" "github.com/databricks/cli/libs/log" @@ -31,7 +32,22 @@ func helpfulError(ctx context.Context, profile string, persistentAuth u2m.OAuthA return fmt.Sprintf("Try logging in again with `%s` before retrying. If this fails, please report this issue to the Databricks CLI maintainers at https://github.com/databricks/cli/issues/new", loginMsg) } +type ( + tokenLoader func(context.Context, loadTokenArgs) (*oauth2.Token, error) + registryHostResolver func(string, string, string) (string, error) +) + func newTokenCommand(authArguments *auth.AuthArguments) *cobra.Command { + return newTokenCommandWithLoader(authArguments, loadToken) +} + +// newTokenCommandWithLoader isolates cache-backed token acquisition from command parsing in tests. +func newTokenCommandWithLoader(authArguments *auth.AuthArguments, load tokenLoader) *cobra.Command { + return newTokenCommandWithRegistryHost(authArguments, load, dockercredentials.RegistryHost) +} + +// newTokenCommandWithRegistryHost isolates registry reconstruction so tests can model cloud and environment matches. +func newTokenCommandWithRegistryHost(authArguments *auth.AuthArguments, load tokenLoader, registryHost registryHostResolver) *cobra.Command { cmd := &cobra.Command{ Use: "token [PROFILE]", Short: "Get authentication token", @@ -50,18 +66,31 @@ and secret is not supported.`, cmd.Flags().BoolVar(&forceRefresh, "force-refresh", false, "Force a token refresh even if the cached token is still valid.") - cmd.PreRunE = profileHostConflictCheck + // Docker format is an internal credential-helper contract, not a user-facing output mode. + var format string + cmd.Flags().StringVar(&format, "format", "", "Hidden output format") + _ = cmd.Flags().MarkHidden("format") + + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + if format == "docker" { + return validateDockerTokenRequest(cmd, args) + } + return profileHostConflictCheck(cmd, args) + } cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() + if format != "" && format != "docker" { + return fmt.Errorf("unsupported token format %q", format) + } tokenStore, mode, err := storage.ResolveStore(ctx, "") if err != nil { return err } - t, err := loadToken(ctx, loadTokenArgs{ + loadArgs := loadTokenArgs{ authArguments: authArguments, profileName: profileName, args: args, @@ -71,7 +100,13 @@ and secret is not supported.`, tokenStore: tokenStore, mode: mode, persistentAuthOpts: nil, - }) + } + + if format == "docker" { + return writeDockerTokenOutput(ctx, cmd, loadArgs, load, registryHost) + } + + t, err := load(ctx, loadArgs) if err != nil { return err } @@ -85,6 +120,108 @@ and secret is not supported.`, return cmd } +type dockerGetResponse struct { + Username string `json:"Username"` + Secret string `json:"Secret"` +} + +// writeDockerTokenOutput makes Docker's registry address the sole profile selector for its get response. +// See https://docs.docker.com/reference/cli/docker/login/#credential-helper-protocol. +func writeDockerTokenOutput(ctx context.Context, cmd *cobra.Command, args loadTokenArgs, load tokenLoader, registryHost registryHostResolver) 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 + } + + profileName, err := dockerTokenProfileName(ctx, registry, args.profiler, registryHost) + if err != nil { + return err + } + + args.authArguments = &auth.AuthArguments{} + args.profileName = profileName + args.args = nil + + t, err := load(ctx, args) + if err != nil { + return err + } + + return json.NewEncoder(cmd.OutOrStdout()).Encode(dockerGetResponse{ + Username: dockercredentials.OAuthTokenUsername, + Secret: t.AccessToken, + }) +} + +// validateDockerTokenRequest makes the registry address the sole profile selector for Docker-format requests. +func validateDockerTokenRequest(cmd *cobra.Command, args []string) error { + if len(args) > 0 { + return errors.New("--format=docker 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("--format=docker does not support --%s", name) + } + } + + return nil +} + +// dockerTokenProfileName finds one compatible profile, using the registry DNS zone to disambiguate repeated workspace IDs. +func dockerTokenProfileName(ctx context.Context, registry dockercredentials.Registry, profiler profile.Profiler, registryHost registryHostResolver) (string, error) { + workspaceProfiles, err := profiler.LoadProfiles(ctx, func(p profile.Profile) bool { + return p.WorkspaceID == registry.WorkspaceID + }) + if err != nil { + return "", err + } + if len(workspaceProfiles) == 0 { + return "", 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) + } + if len(workspaceProfiles) == 1 { + return validateDockerTokenProfile(registry, workspaceProfiles[0], registryHost) + } + + var matchingProfiles profile.Profiles + for _, p := range workspaceProfiles { + if validateDockerCredentialProfile(p) != nil { + continue + } + expectedHost, err := registryHost(registry.WorkspaceID, registry.Region, p.Host) + if err == nil && expectedHost == registry.Host { + matchingProfiles = append(matchingProfiles, p) + } + } + if len(matchingProfiles) == 0 { + return "", fmt.Errorf("registry host %s does not match any profile for workspace ID %s. Verify the profile workspace host and workspace_id", registry.Host, registry.WorkspaceID) + } + if len(matchingProfiles) > 1 { + return "", 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 ")) + } + return validateDockerTokenProfile(registry, matchingProfiles[0], registryHost) +} + +// validateDockerTokenProfile verifies U2M eligibility and reconstructs the registry host to prevent cross-environment matches. +func validateDockerTokenProfile(registry dockercredentials.Registry, p profile.Profile, registryHost registryHostResolver) (string, error) { + if err := validateDockerCredentialProfile(p); err != nil { + return "", err + } + // Workspace IDs can repeat across environments, so the registry must also match the profile's DNS zone. + expectedHost, err := registryHost(registry.WorkspaceID, registry.Region, p.Host) + if err != nil { + return "", fmt.Errorf("validate registry host against profile %q: %w", p.Name, err) + } + if expectedHost != registry.Host { + return "", fmt.Errorf("registry host %s does not match profile %q workspace host", registry.Host, p.Name) + } + return p.Name, nil +} + func writeTokenOutput(w io.Writer, t *oauth2.Token, textMode bool) error { if textMode { _, err := fmt.Fprintln(w, t.AccessToken) diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index adda6888a40..1aa775d7530 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -6,17 +6,25 @@ import ( "encoding/json" "errors" "net/http" + "os" + "path/filepath" + "strings" "testing" "time" "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/cli/libs/env" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) @@ -885,6 +893,396 @@ func (e errProfiler) GetPath(context.Context) (string, error) { return "", nil } +func TestTokenDockerFormatEmitsGetResponse(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: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + })) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + + var gotProfile string + loadToken := func(_ context.Context, args loadTokenArgs) (*oauth2.Token, error) { + gotProfile = args.profileName + return &oauth2.Token{AccessToken: "access-token"}, nil + } + + registryHost := "123456789.container.us-west-2.cloud.databricks.com" + var stdout bytes.Buffer + cmd := newTokenCommandWithRegistryHost(&auth.AuthArguments{}, loadToken, func(workspaceID, region, workspaceHost string) (string, error) { + require.Equal(t, "123456789", workspaceID) + require.Equal(t, "us-west-2", region) + require.Equal(t, "https://workspace.cloud.databricks.test", workspaceHost) + return registryHost, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader(registryHost + "\n")) + cmd.SetOut(&stdout) + cmd.SetArgs([]string{"--format=docker"}) + + require.NoError(t, cmd.Execute()) + require.Equal(t, "workspace", gotProfile) + + var got map[string]string + require.NoError(t, json.Unmarshal(stdout.Bytes(), &got)) + require.Equal(t, map[string]string{ + "Username": "oauthtoken", + "Secret": "access-token", + }, got) +} + +func TestWriteDockerTokenOutputUsesConfiguredProfiler(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "workspace", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + }, + }, + } + + var gotProfile string + loadToken := func(_ context.Context, args loadTokenArgs) (*oauth2.Token, error) { + gotProfile = args.profileName + return &oauth2.Token{AccessToken: "access-token"}, nil + } + + cmd := &cobra.Command{Use: "token"} + var stdout bytes.Buffer + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetOut(&stdout) + + err := writeDockerTokenOutput(ctx, cmd, loadTokenArgs{ + authArguments: &auth.AuthArguments{}, + profiler: profiler, + }, loadToken, func(string, string, string) (string, error) { + return "123456789.container.us-west-2.cloud.databricks.com", nil + }) + require.NoError(t, err) + require.Equal(t, "workspace", gotProfile) + + var got dockerGetResponse + require.NoError(t, json.Unmarshal(stdout.Bytes(), &got)) +} + +func TestDockerTokenProfileNameRejectsDifferentEnvironment(t *testing.T) { + registry := dockercredentials.Registry{ + WorkspaceID: "123456789", + Region: "us-west-2", + Host: "123456789.container.us-west-2.cloud.databricks.com", + } + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{{ + Name: "workspace", + Host: "https://workspace.dev.cloud.databricks.test", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }}, + } + registryHost := func(workspaceID, region, workspaceHost string) (string, error) { + require.Equal(t, registry.WorkspaceID, workspaceID) + require.Equal(t, registry.Region, region) + require.Equal(t, "https://workspace.dev.cloud.databricks.test", workspaceHost) + return "123456789.container.us-west-2.dev.cloud.databricks.com", nil + } + + _, err := dockerTokenProfileName(t.Context(), registry, profiler, registryHost) + require.ErrorContains(t, err, "does not match profile") + require.ErrorContains(t, err, "workspace host") +} + +func TestDockerTokenProfileNameAllowsSameWorkspaceIDInDifferentEnvironment(t *testing.T) { + registry := dockercredentials.Registry{ + WorkspaceID: "123456789", + Region: "us-west-2", + Host: "123456789.container.us-west-2.cloud.databricks.com", + } + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "prod", + Host: "https://workspace.cloud.databricks.com", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }, + { + Name: "dev", + Host: "https://workspace.dev.cloud.databricks.com", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }, + }, + } + registryHost := func(workspaceID, region, workspaceHost string) (string, error) { + zone := ".cloud.databricks.com" + if workspaceHost == "https://workspace.dev.cloud.databricks.com" { + zone = ".dev.cloud.databricks.com" + } + return workspaceID + ".container." + region + zone, nil + } + + profileName, err := dockerTokenProfileName(t.Context(), registry, profiler, registryHost) + require.NoError(t, err) + require.Equal(t, "prod", profileName) +} + +func TestDockerTokenProfileNameIgnoresUnsupportedDuplicateProfile(t *testing.T) { + registry := dockercredentials.Registry{ + WorkspaceID: "123456789", + Region: "us-west-2", + Host: "123456789.container.us-west-2.cloud.databricks.com", + } + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "workspace", + Host: "https://workspace.cloud.databricks.com", + WorkspaceID: registry.WorkspaceID, + AuthType: authTypeDatabricksCLI, + }, + { + Name: "m2m", + Host: "https://workspace.cloud.databricks.com", + WorkspaceID: registry.WorkspaceID, + HasClientCredentials: true, + }, + }, + } + registryHost := func(workspaceID, region, _ string) (string, error) { + return workspaceID + ".container." + region + ".cloud.databricks.com", nil + } + + profileName, err := dockerTokenProfileName(t.Context(), registry, profiler, registryHost) + require.NoError(t, err) + require.Equal(t, "workspace", profileName) +} + +func TestTokenDockerFormatRejectsPositionalArgs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", t.TempDir()) + + cmd := newTokenCommandWithLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetArgs([]string{"--format=docker", "DEFAULT"}) + + err := cmd.Execute() + require.ErrorContains(t, err, "--format=docker does not accept positional arguments") +} + +func TestTokenDockerFormatValidatesBeforeResolvingTokenStore(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + t.Setenv(storage.EnvVar, "invalid") + + cmd := newTokenCommandWithLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"--format=docker", "DEFAULT"}) + + err := cmd.Execute() + require.ErrorContains(t, err, "--format=docker does not accept positional arguments") +} + +func TestTokenDockerFormatRejectsAuthSelectionFlags(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + require.NoError(t, databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: "DEFAULT", + Host: "https://profile.cloud.databricks.test", + AuthType: authTypeDatabricksCLI, + })) + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + + cases := [][]string{ + {"--format=docker", "--profile", "DEFAULT"}, + {"--format=docker", "--host", "https://workspace.cloud.databricks.test"}, + {"--format=docker", "--profile", "DEFAULT", "--host", "https://workspace.cloud.databricks.test"}, + {"--format=docker", "--account-id", "abc"}, + {"--format=docker", "--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") + cmd.AddCommand(newTokenCommandWithLoader(&authArgs, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + })) + 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{"token"}, args...)) + + err := cmd.Execute() + require.ErrorContains(t, err, "--format=docker does not support") + }) + } +} + +func TestTokenDockerFormatRejectsNonDARHost(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", t.TempDir()) + + cmd := newTokenCommandWithLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("registry.example.com\n")) + cmd.SetArgs([]string{"--format=docker"}) + + err := cmd.Execute() + require.ErrorContains(t, err, "is not a Databricks Artifact Registry host") +} + +func TestTokenDockerFormatErrorsWithoutMatchingProfile(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) + + cmd := newTokenCommandWithLoader(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetArgs([]string{"--format=docker"}) + + err := cmd.Execute() + require.ErrorContains(t, err, "no Databricks profile found for workspace ID 123456789") + require.ErrorContains(t, err, "databricks auth login --host ") + require.ErrorContains(t, err, "workspace_id") +} + +func TestTokenDockerFormatErrorsWithMultipleMatchingProfiles(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + for _, name := range []string{"one", "two"} { + require.NoError(t, databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: name, + Host: "https://" + name + ".cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authTypeDatabricksCLI, + })) + } + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", dir) + + cmd := newTokenCommandWithRegistryHost(&auth.AuthArguments{}, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }, func(workspaceID, region, _ string) (string, error) { + return workspaceID + ".container." + region + ".cloud.databricks.com", 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{"--format=docker"}) + + err := cmd.Execute() + require.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") + require.ErrorContains(t, err, "one and two") + require.ErrorContains(t, err, "Remove duplicate workspace_id entries") +} + +func TestTokenDockerFormatRejectsUnsupportedProfile(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "pat", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: "pat", + }, + { + Name: "m2m", + Host: "https://m2m.cloud.databricks.test", + WorkspaceID: "987654321", + HasClientCredentials: true, + }, + { + Name: "blank-auth", + Host: "https://blank-auth.cloud.databricks.test", + WorkspaceID: "111222333", + }, + { + Name: "account", + Host: "https://accounts.cloud.databricks.test", + AccountID: "account-id", + WorkspaceID: "444555666", + AuthType: authTypeDatabricksCLI, + }, + }, + } + + for _, tc := range []struct { + registryHost string + wantError string + }{ + {"123456789.container.us-west-2.cloud.databricks.com", "requires a profile created by databricks auth login"}, + {"987654321.container.us-west-2.cloud.databricks.com", "requires a profile created by databricks auth login"}, + {"111222333.container.us-west-2.cloud.databricks.com", "requires a profile created by databricks auth login"}, + {"444555666.container.us-west-2.cloud.databricks.com", "does not target a workspace"}, + } { + t.Run(tc.registryHost, func(t *testing.T) { + cmd := &cobra.Command{Use: "token"} + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader(tc.registryHost + "\n")) + + err := writeDockerTokenOutput(ctx, cmd, loadTokenArgs{ + authArguments: &auth.AuthArguments{}, + profiler: profiler, + }, func(context.Context, loadTokenArgs) (*oauth2.Token, error) { + t.Fatal("loadToken should not be called") + return nil, nil + }, dockercredentials.RegistryHost) + require.ErrorContains(t, err, tc.wantError) + }) + } +} + func TestWriteTokenOutput(t *testing.T) { token := &oauth2.Token{ AccessToken: "my-access-token", diff --git a/libs/dockercredentials/docker_config.go b/libs/dockercredentials/docker_config.go new file mode 100644 index 00000000000..706085ba05b --- /dev/null +++ b/libs/dockercredentials/docker_config.go @@ -0,0 +1,127 @@ +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 + + if err := writeDockerConfig(path, config); err != nil { + return err + } + return nil +} + +// 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 +} + +// readDockerConfig preserves unrelated top-level values as raw JSON while treating a missing file as empty configuration. +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 +} + +// writeDockerConfig writes through an owner-only sibling temporary file before replacing the target. +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 := createOwnerOnlyTempFile(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_permissions_other.go b/libs/dockercredentials/docker_config_permissions_other.go new file mode 100644 index 00000000000..b5f6d0f2abb --- /dev/null +++ b/libs/dockercredentials/docker_config_permissions_other.go @@ -0,0 +1,10 @@ +//go:build !windows + +package dockercredentials + +import "os" + +// createOwnerOnlyTempFile relies on os.CreateTemp's 0600 mode outside Windows. +func createOwnerOnlyTempFile(dir, pattern string) (*os.File, error) { + return os.CreateTemp(dir, pattern) +} diff --git a/libs/dockercredentials/docker_config_permissions_windows.go b/libs/dockercredentials/docker_config_permissions_windows.go new file mode 100644 index 00000000000..8c5aa7e2763 --- /dev/null +++ b/libs/dockercredentials/docker_config_permissions_windows.go @@ -0,0 +1,102 @@ +//go:build windows + +package dockercredentials + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + "unsafe" + + "golang.org/x/sys/windows" +) + +const maxTempFileAttempts = 100 + +// windowsCreateFilePath adds an extended-length prefix at Go's conservative legacy Windows path threshold. +// See https://learn.microsoft.com/en-us/windows/win32/fileio/naming-a-file#maximum-path-length-limitation. +func windowsCreateFilePath(path string) (string, error) { + if strings.HasPrefix(path, `\\?\`) || strings.HasPrefix(path, `\??\`) || strings.HasPrefix(path, `\\.\`) { + return path, nil + } + absolutePath, err := filepath.Abs(path) + if err != nil { + return "", err + } + if strings.HasPrefix(absolutePath, `\\?\`) || strings.HasPrefix(absolutePath, `\??\`) || strings.HasPrefix(absolutePath, `\\.\`) { + return absolutePath, nil + } + if len(absolutePath) < 248 { + return path, nil + } + if strings.HasPrefix(absolutePath, `\\`) { + return `\\?\UNC\` + strings.TrimPrefix(absolutePath, `\\`), nil + } + return `\\?\` + absolutePath, nil +} + +// createOwnerOnlyTempFile creates the Windows temp file with a protected current-user DACL at handle creation. +// See https://learn.microsoft.com/en-us/windows/win32/api/fileapi/nf-fileapi-createfilew. +func createOwnerOnlyTempFile(dir, pattern string) (*os.File, error) { + user, err := windows.GetCurrentProcessToken().GetTokenUser() + if err != nil { + return nil, fmt.Errorf("get current Windows user: %w", err) + } + descriptor, err := windows.SecurityDescriptorFromString(fmt.Sprintf("D:P(A;;GA;;;%s)", user.User.Sid.String())) + if err != nil { + return nil, fmt.Errorf("create owner-only Windows security descriptor: %w", err) + } + + var pinner runtime.Pinner + pinner.Pin(descriptor) + defer pinner.Unpin() + attributes := windows.SecurityAttributes{ + Length: uint32(unsafe.Sizeof(windows.SecurityAttributes{})), + SecurityDescriptor: descriptor, + } + + for range maxTempFileAttempts { + // The empty placeholder borrows os.CreateTemp's naming; CREATE_NEW prevents adopting a raced replacement. + placeholder, err := os.CreateTemp(dir, pattern) + if err != nil { + return nil, err + } + path := placeholder.Name() + if err := placeholder.Close(); err != nil { + _ = os.Remove(path) + return nil, fmt.Errorf("close temporary filename reservation %s: %w", path, err) + } + if err := os.Remove(path); err != nil { + return nil, fmt.Errorf("remove temporary filename reservation %s: %w", path, err) + } + + createPath, err := windowsCreateFilePath(path) + if err != nil { + return nil, fmt.Errorf("resolve temporary Docker config path %s: %w", path, err) + } + pathPtr, err := windows.UTF16PtrFromString(createPath) + if err != nil { + return nil, fmt.Errorf("encode temporary Docker config path %s: %w", path, err) + } + handle, err := windows.CreateFile( + pathPtr, + windows.GENERIC_READ|windows.GENERIC_WRITE, + windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE, + &attributes, + windows.CREATE_NEW, + windows.FILE_ATTRIBUTE_NORMAL, + 0, + ) + if errors.Is(err, windows.ERROR_FILE_EXISTS) || errors.Is(err, windows.ERROR_ALREADY_EXISTS) { + continue + } + if err != nil { + return nil, fmt.Errorf("create owner-only temporary Docker config %s: %w", path, err) + } + return os.NewFile(uintptr(handle), path), nil + } + return nil, fmt.Errorf("create owner-only temporary Docker config in %s: too many name collisions", dir) +} diff --git a/libs/dockercredentials/docker_config_permissions_windows_test.go b/libs/dockercredentials/docker_config_permissions_windows_test.go new file mode 100644 index 00000000000..64dd9f42635 --- /dev/null +++ b/libs/dockercredentials/docker_config_permissions_windows_test.go @@ -0,0 +1,92 @@ +//go:build windows + +package dockercredentials + +import ( + "os" + "path/filepath" + "strings" + "testing" + "unsafe" + + "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" +) + +func TestCreateOwnerOnlyTempFileRestrictsWindowsDACL(t *testing.T) { + file, err := createOwnerOnlyTempFile(t.TempDir(), "config.json.*") + require.NoError(t, err) + t.Cleanup(func() { + _ = file.Close() + }) + + descriptor, err := windows.GetSecurityInfo( + windows.Handle(file.Fd()), + windows.SE_FILE_OBJECT, + windows.DACL_SECURITY_INFORMATION, + ) + require.NoError(t, err) + control, _, err := descriptor.Control() + require.NoError(t, err) + require.NotZero(t, control&windows.SE_DACL_PROTECTED) + + dacl, _, err := descriptor.DACL() + require.NoError(t, err) + require.Equal(t, uint16(1), dacl.AceCount) + var ace *windows.ACCESS_ALLOWED_ACE + require.NoError(t, windows.GetAce(dacl, 0, &ace)) + const fileAllAccess = windows.STANDARD_RIGHTS_REQUIRED | windows.SYNCHRONIZE | 0x1ff + require.True(t, ace.Mask&windows.GENERIC_ALL != 0 || ace.Mask&fileAllAccess == fileAllAccess) + + user, err := windows.GetCurrentProcessToken().GetTokenUser() + require.NoError(t, err) + aceSID := (*windows.SID)(unsafe.Pointer(&ace.SidStart)) + require.True(t, user.User.Sid.Equals(aceSID)) +} + +func TestWindowsCreateFilePathSupportsExtendedLengthPaths(t *testing.T) { + longDrivePath := `C:\` + strings.Repeat("a", 260) + longUNCPath := `\\server\share\` + strings.Repeat("a", 260) + tests := []struct { + name string + path string + want string + }{ + {name: "short", path: `C:\config.json`, want: `C:\config.json`}, + {name: "drive", path: longDrivePath, want: `\\?\` + longDrivePath}, + {name: "UNC", path: longUNCPath, want: `\\?\UNC\server\share\` + strings.Repeat("a", 260)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := windowsCreateFilePath(tt.path) + require.NoError(t, err) + require.Equal(t, tt.want, got) + }) + } +} + +func TestWindowsCreateFilePathResolvesRelativePathBeforeLengthCheck(t *testing.T) { + baseDir := t.TempDir() + componentLength := 247 - len(baseDir) - 1 + if componentLength <= 0 { + t.Skip("temporary directory already exceeds the legacy Windows path threshold") + } + longDir := filepath.Join(baseDir, strings.Repeat("a", componentLength)) + require.Equal(t, 247, len(longDir)) + require.GreaterOrEqual(t, len(filepath.Join(longDir, "config.json")), 248) + require.NoError(t, os.MkdirAll(longDir, 0o755)) + extendedDir := longDir + if !strings.HasPrefix(longDir, `\\?\`) { + if strings.HasPrefix(longDir, `\\`) { + extendedDir = `\\?\UNC\` + strings.TrimPrefix(longDir, `\\`) + } else { + extendedDir = `\\?\` + longDir + } + } + t.Chdir(extendedDir) + + got, err := windowsCreateFilePath("config.json") + require.NoError(t, err) + require.Equal(t, filepath.Join(extendedDir, "config.json"), got) +} diff --git a/libs/dockercredentials/docker_config_test.go b/libs/dockercredentials/docker_config_test.go new file mode 100644 index 00000000000..3e7a0160c2d --- /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/require" +) + +const testRegistryHost = "123.container.us-west-2.cloud.databricks.test" + +// readDockerConfigForTest decodes generically so tests can verify preservation of unrelated fields. +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) + require.Equal(t, map[string]any{ + testRegistryHost: HelperName, + }, got["credHelpers"]) + + info, err := os.Stat(path) + require.NoError(t, err) + if runtime.GOOS != "windows" { + require.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.com": {"auth": "abc"} + }, + "credsStore": "desktop", + "credHelpers": { + "registry.example.com": "desktop" + }, + "experimental": "enabled" +}`), 0o600)) + + require.NoError(t, SetCredentialHelper(path, testRegistryHost)) + + got := readDockerConfigForTest(t, path) + require.Equal(t, "desktop", got["credsStore"]) + require.Equal(t, "enabled", got["experimental"]) + require.Equal(t, map[string]any{ + "registry.example.com": "desktop", + testRegistryHost: HelperName, + }, got["credHelpers"]) + require.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) + require.NotZero(t, info.Mode()&os.ModeSymlink) + got := readDockerConfigForTest(t, target) + require.Equal(t, "desktop", got["credsStore"]) + require.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) + require.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) + require.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) + require.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) + require.ErrorContains(t, err, "read Docker config") +} diff --git a/libs/dockercredentials/registry.go b/libs/dockercredentials/registry.go new file mode 100644 index 00000000000..fee76b0ca64 --- /dev/null +++ b/libs/dockercredentials/registry.go @@ -0,0 +1,192 @@ +package dockercredentials + +import ( + "errors" + "fmt" + "net" + "net/url" + "strconv" + "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, region, and canonical host of a Databricks Artifact Registry endpoint. +type Registry struct { + WorkspaceID string + Region string + Host string +} + +// RegistryHost derives .container.. while preserving the workspace's cloud and environment 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%s%s%s", workspaceID, registryHostInfix, region, dnsZone), nil +} + +// normalizeServerAddress canonicalizes Docker's URL-or-host input and permits only HTTPS on the default registry port. +func normalizeServerAddress(raw string) (string, error) { + value := strings.TrimSpace(raw) + if value == "" { + return "", errors.New("server address is required") + } + + if strings.Contains(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) + } + value = u.Host + } else if i := strings.IndexByte(value, '/'); i >= 0 { + value = value[:i] + } + + if host, port, ok, err := splitOptionalPort(value); err != nil { + return "", err + } else if ok { + value = host + if err := validatePort(port); err != nil { + return "", err + } + if port != "443" { + return "", fmt.Errorf("unsupported registry port %q", port) + } + } + + value = strings.TrimSuffix(strings.ToLower(value), ".") + if value == "" { + return "", errors.New("server address is required") + } + return value, nil +} + +// ParseRegistryHost normalizes a Databricks Artifact Registry address and extracts its workspace and region. +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, + Region: region, + Host: host, + }, nil +} + +// registryDNSZoneForWorkspaceHost derives the registry suffix from the workspace's SDK-known cloud and environment zone. +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 +} + +// matchingDatabricksDNSZone selects the most specific suffix from all SDK-known Databricks environments. +func matchingDatabricksDNSZone(host string) (string, bool) { + return matchingDatabricksDNSZoneInEnvironments(host, environment.AllEnvironments()) +} + +// matchingDatabricksDNSZoneInEnvironments prefers the longest suffix so environment-specific zones beat generic ones. +func matchingDatabricksDNSZoneInEnvironments(host string, envs []environment.DatabricksEnvironment) (string, bool) { + var match string + for _, e := range envs { + dnsZone := strings.ToLower(e.DnsZone) + if dnsZone == "" { + continue + } + if strings.HasSuffix(host, dnsZone) && len(dnsZone) > len(match) { + match = dnsZone + } + } + return match, match != "" +} + +// splitOptionalPort extracts bracketed or plain host ports while leaving non-port colon forms for host validation. +func splitOptionalPort(value string) (host, port string, ok bool, err error) { + host, port, err = net.SplitHostPort(value) + if err == nil { + return host, port, true, nil + } + + if strings.Count(value, ":") == 1 { + host, port, found := strings.Cut(value, ":") + if found && port != "" { + return host, port, true, nil + } + } + + return "", "", false, nil +} + +// validatePort enforces the decimal TCP port range that URL parsing alone does not validate. +func validatePort(port string) error { + n, err := strconv.Atoi(port) + if err != nil || n < 1 || n > 65535 { + return fmt.Errorf("invalid registry port %q", port) + } + return nil +} + +// isDNSLabel restricts interpolated registry components to unescaped lowercase ASCII DNS labels. +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..39f80c1124b --- /dev/null +++ b/libs/dockercredentials/registry_test.go @@ -0,0 +1,187 @@ +package dockercredentials + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/common/environment" + "github.com/stretchr/testify/require" +) + +func TestRegistryHost(t *testing.T) { + cases := []struct { + name string + workspaceHost string + region string + want string + }{ + { + name: "aws prod", + 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 prod", + workspaceHost: "https://adb-123.456.azuredatabricks.net", + region: "eastus", + want: "123456789.container.eastus.azuredatabricks.net", + }, + { + name: "azure dev", + workspaceHost: "https://workspace.dev.azuredatabricks.net", + region: "eastus", + want: "123456789.container.eastus.dev.azuredatabricks.net", + }, + { + name: "gcp prod", + workspaceHost: "https://workspace.gcp.databricks.com", + region: "us-central1", + want: "123456789.container.us-central1.gcp.databricks.com", + }, + { + name: "gcp dev", + workspaceHost: "https://workspace.dev.gcp.databricks.com", + region: "us-central1", + want: "123456789.container.us-central1.dev.gcp.databricks.com", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := RegistryHost("123456789", tc.region, tc.workspaceHost) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } +} + +func TestRegistryHostRejectsEmptyParts(t *testing.T) { + _, err := RegistryHost("", "us-west-2", "https://workspace.cloud.databricks.test") + require.ErrorContains(t, err, "workspace ID is required") + + _, err = RegistryHost("123456789", "", "https://workspace.cloud.databricks.test") + require.ErrorContains(t, err, "region is required") +} + +func TestRegistryHostRejectsUnsupportedWorkspaceHost(t *testing.T) { + _, err := RegistryHost("123456789", "us-west-2", "https://workspace.example.test") + require.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", + "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) + require.Equal(t, Registry{ + WorkspaceID: "123456789", + Region: "us-west-2", + Host: "123456789.container.us-west-2.cloud.databricks.com", + }, got) + }) + } +} + +func TestRegistryHostAndParseRegistryHostSupportAllDatabricksEnvironmentZones(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 + got, err := RegistryHost("123456789", "test-region", "https://workspace"+dnsZone) + require.NoError(t, err) + require.Equal(t, wantHost, got) + + registry, err := ParseRegistryHost("https://" + wantHost + "/v2/") + require.NoError(t, err) + require.Equal(t, Registry{ + WorkspaceID: "123456789", + Region: "test-region", + Host: wantHost, + }, registry) + }) + } +} + +func TestParseRegistryHostUsesLongestDNSZoneSuffix(t *testing.T) { + got, err := ParseRegistryHost("123456789.container.us-west-2.staging.cloud.databricks.com") + require.NoError(t, err) + require.Equal(t, Registry{ + WorkspaceID: "123456789", + Region: "us-west-2", + Host: "123456789.container.us-west-2.staging.cloud.databricks.com", + }, got) +} + +func TestMatchingDatabricksDNSZoneIgnoresEmptyDNSZones(t *testing.T) { + got, ok := matchingDatabricksDNSZoneInEnvironments("workspace.example.test", []environment.DatabricksEnvironment{ + {DnsZone: ""}, + {DnsZone: ".example.test"}, + }) + require.True(t, ok) + require.Equal(t, ".example.test", got) + + _, ok = matchingDatabricksDNSZoneInEnvironments("workspace.invalid", []environment.DatabricksEnvironment{ + {DnsZone: ""}, + }) + require.False(t, ok) +} + +func TestParseRegistryHostRejectsNonDARHost(t *testing.T) { + _, err := ParseRegistryHost("registry.example.com") + require.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") + require.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") + require.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") + require.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) + require.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") + require.ErrorContains(t, err, "unsupported registry URL scheme") +} + +func TestNormalizeServerAddressRejectsInvalidPort(t *testing.T) { + _, err := normalizeServerAddress("https://123.container.us-west-2.cloud.databricks.com:99999") + require.ErrorContains(t, err, "invalid registry port") +} + +func TestNormalizeServerAddressRejectsNonHTTPSPort(t *testing.T) { + _, err := normalizeServerAddress("https://123.container.us-west-2.cloud.databricks.com:5000") + require.ErrorContains(t, err, "unsupported registry port") +} + +func TestNormalizeServerAddressAllowsHTTPSPort(t *testing.T) { + got, err := normalizeServerAddress("https://123.container.us-west-2.cloud.databricks.com:443") + require.NoError(t, err) + require.Equal(t, "123.container.us-west-2.cloud.databricks.com", got) +} diff --git a/libs/dockercredentials/shim.go b/libs/dockercredentials/shim.go new file mode 100644 index 00000000000..642de5f1fd6 --- /dev/null +++ b/libs/dockercredentials/shim.go @@ -0,0 +1,170 @@ +package dockercredentials + +import ( + "errors" + "fmt" + "io" + "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 is the filesystem path of the installed helper. + Path string + // OnPath reports whether Path is the first matching helper in the current PATH, using PATHEXT on Windows. + OnPath bool +} + +// InstallShim installs a Unix wrapper or a Windows CLI copy that dispatches by executable name. +func InstallShim(databricksPath, installDir string) (ShimInstallResult, error) { + return installShimForGOOS(databricksPath, installDir, runtime.GOOS) +} + +// installShimForGOOS writes a Unix wrapper or a Windows CLI copy and reports whether Docker resolves that exact helper first. +func installShimForGOOS(databricksPath, installDir, goos string) (ShimInstallResult, error) { + if strings.TrimSpace(databricksPath) == "" { + return ShimInstallResult{}, errors.New("databricks executable path is required") + } + if strings.TrimSpace(installDir) == "" { + return ShimInstallResult{}, errors.New("install directory is required") + } + + 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) + var err error + if goos == "windows" { + mode = 0o644 + err = copyShimFile(path, databricksPath, mode) + } else { + err = writeShimFile(path, []byte(shimScript(databricksPath)), mode) + } + if 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 +} + +// shimFilename includes .exe on Windows so Docker can discover the helper through PATHEXT. +// See https://go.dev/src/os/exec/lp_windows.go. +func shimFilename(goos string) string { + if goos == "windows" { + return "docker-credential-" + HelperName + ".exe" + } + return "docker-credential-" + HelperName +} + +// shimScript accepts only get and sends CLI logs to stderr so stdout remains valid credential-helper JSON. +// See https://docs.docker.com/reference/cli/docker/login/#credential-helper-protocol. +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 token --format=docker +`, posixShellQuote(databricksPath)) +} + +// writeShimFile stages the Unix wrapper beside the destination before replacing an existing helper. +func writeShimFile(path string, script []byte, mode os.FileMode) error { + return writeShim(path, mode, func(tmp *os.File) error { + _, err := tmp.Write(script) + return err + }) +} + +// copyShimFile stages the Windows CLI copy so a failed copy preserves an existing helper. +func copyShimFile(path, source string, mode os.FileMode) error { + src, err := os.Open(source) + if err != nil { + return err + } + defer src.Close() + + return writeShim(path, mode, func(tmp *os.File) error { + _, err := io.Copy(tmp, src) + return err + }) +} + +// writeShim writes through a sibling temporary file so a failed update cannot truncate an existing helper. +func writeShim(path string, mode os.FileMode, write func(*os.File) error) 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 := write(tmp); 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) +} + +// posixShellQuote returns one shell word even when the executable path contains quotes or metacharacters. +func posixShellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +// helperOnPathForGOOS reports whether Docker's lookup name resolves to the helper just installed. +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) +} + +// samePath compares file identity when possible and falls back to platform-normalized absolute paths. +func samePath(a, b string) bool { + aInfo, aErr := os.Stat(a) + bInfo, bErr := os.Stat(b) + if aErr == nil && bErr == nil { + return os.SameFile(aInfo, bInfo) + } + + absA, err := filepath.Abs(a) + if err == nil { + a = absA + } + absB, err := filepath.Abs(b) + if err == nil { + b = absB + } + a = filepath.Clean(a) + b = filepath.Clean(b) + if runtime.GOOS == "windows" { + return strings.EqualFold(a, b) + } + return a == b +} diff --git a/libs/dockercredentials/shim_test.go b/libs/dockercredentials/shim_test.go new file mode 100644 index 00000000000..f8d44368750 --- /dev/null +++ b/libs/dockercredentials/shim_test.go @@ -0,0 +1,216 @@ +package dockercredentials + +import ( + "bytes" + "os" + "os/exec" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/require" +) + +// writeTestDatabricksExecutable uses the platform suffix so tests exercise Windows helper paths. +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) { + require.Equal(t, "docker-credential-databricks", shimFilename("linux")) + require.Equal(t, "docker-credential-databricks.exe", shimFilename("windows")) +} + +func TestUnixShimScript(t *testing.T) { + got := shimScript("/opt/databricks/bin/databricks") + + require.Contains(t, got, `if [ "$#" -ne 1 ] || [ "$1" != "get" ]; then`) + require.Contains(t, got, `docker-credential-databricks only supports get`) + require.Contains(t, got, `export DATABRICKS_LOG_FILE=stderr`) + require.Contains(t, got, `exec '/opt/databricks/bin/databricks' auth token --format=docker`) +} + +func TestInstallWindowsShimCopiesDatabricksExecutable(t *testing.T) { + dir := t.TempDir() + databricksPath := filepath.Join(dir, "databricks.exe") + require.NoError(t, os.WriteFile(databricksPath, []byte("databricks executable"), 0o755)) + installDir := filepath.Join(dir, "bin") + t.Setenv("PATH", installDir) + + got, err := installShimForGOOS(databricksPath, installDir, "windows") + require.NoError(t, err) + require.Equal(t, filepath.Join(installDir, "docker-credential-databricks.exe"), got.Path) + + raw, err := os.ReadFile(got.Path) + require.NoError(t, err) + require.Equal(t, "databricks executable", string(raw)) +} + +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) + require.JSONEq(t, `{"Username":"oauthtoken","Secret":"secret"}`, string(out)) + + rawArgs, err := os.ReadFile(argsPath) + require.NoError(t, err) + require.Equal(t, "auth token --format=docker", string(rawArgs)) + + rawEnv, err := os.ReadFile(envPath) + require.NoError(t, err) + require.Equal(t, "stderr", string(rawEnv)) + + rawStdin, err := os.ReadFile(stdinPath) + require.NoError(t, err) + require.Equal(t, "registry-host", string(rawStdin)) + + err = exec.Command(shim, "store").Run() + require.Error(t, err) + err = exec.Command(shim, "get", "store").Run() + require.Error(t, err) +} + +func TestInstallShimReportsPathStatus(t *testing.T) { + dir := t.TempDir() + databricksPath := writeTestDatabricksExecutable(t, t.TempDir()) + t.Setenv("PATH", dir) + + got, err := InstallShim(databricksPath, dir) + require.NoError(t, err) + require.Equal(t, filepath.Join(dir, shimFilename(runtime.GOOS)), got.Path) + require.True(t, got.OnPath) + + info, err := os.Stat(got.Path) + require.NoError(t, err) + if runtime.GOOS != "windows" { + require.Equal(t, os.FileMode(0o755), info.Mode().Perm()) + } +} + +func TestInstallShimReportsNotOnPath(t *testing.T) { + dir := t.TempDir() + databricksPath := writeTestDatabricksExecutable(t, t.TempDir()) + t.Setenv("PATH", t.TempDir()) + + got, err := InstallShim(databricksPath, dir) + require.NoError(t, err) + require.Equal(t, filepath.Join(dir, shimFilename(runtime.GOOS)), got.Path) + require.False(t, got.OnPath) +} + +func TestInstallShimReportsNotOnPathWhenHelperIsShadowed(t *testing.T) { + installDir := t.TempDir() + databricksPath := writeTestDatabricksExecutable(t, t.TempDir()) + 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, installDir) + require.NoError(t, err) + require.Equal(t, filepath.Join(installDir, shimFilename(runtime.GOOS)), got.Path) + require.False(t, got.OnPath) +} + +func TestHelperOnPathUsesDockerLookupName(t *testing.T) { + dir := t.TempDir() + helperPath := filepath.Join(dir, "docker-credential-databricks.exe") + 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 + }) + + require.True(t, found) + require.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", "") + require.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)) + + require.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") + databricksPath := writeTestDatabricksExecutable(t, t.TempDir()) + shimPath := filepath.Join(installDir, shimFilename(runtime.GOOS)) + require.NoError(t, os.MkdirAll(installDir, 0o755)) + 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, installDir) + require.Error(t, err) + + raw, readErr := os.ReadFile(shimPath) + require.NoError(t, readErr) + require.Equal(t, "existing helper", string(raw)) +} diff --git a/main.go b/main.go index 6c4dddd40b6..e059f7ee883 100644 --- a/main.go +++ b/main.go @@ -2,8 +2,11 @@ package main import ( "context" + "errors" + "fmt" "os" "path/filepath" + "strings" "github.com/databricks/cli/cmd" "github.com/databricks/cli/cmd/root" @@ -13,19 +16,48 @@ import ( _ "github.com/databricks/cli/libs/hostmetadata" ) +// commandArgs switches copied Windows helpers into Docker get mode and rejects every other helper operation. +// See https://docs.docker.com/reference/cli/docker/login/#credential-helper-protocol. +func commandArgs(executable string, args []string) ([]string, bool, error) { + // Windows installs a copy of this binary as the helper, so argv[0] selects credential-helper mode. + base := executable + if i := strings.LastIndexAny(base, `/\`); i >= 0 { + base = base[i+1:] + } + helperName := "docker-credential-databricks" + if !strings.EqualFold(base, helperName) && !strings.EqualFold(base, helperName+".exe") { + return args, false, nil + } + if len(args) != 1 || args[0] != "get" { + return nil, true, errors.New("docker-credential-databricks only supports get") + } + return []string{"auth", "token", "--format=docker"}, true, nil +} + func main() { + args, dockerHelper, err := commandArgs(os.Args[0], os.Args[1:]) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + if dockerHelper { + _ = os.Setenv("DATABRICKS_LOG_FILE", "stderr") + } + // Configure DATABRICKS_CLI_PATH only if our caller intends to use this specific version of this binary. // Otherwise, if it is equal to its basename, processes can find it in $PATH. // This runs in main rather than in a package init so that importing CLI // packages (e.g. from test binaries or generators) does not mutate the // process environment. arg0 := os.Args[0] - if arg0 != filepath.Base(arg0) { + if !dockerHelper && arg0 != filepath.Base(arg0) { os.Setenv("DATABRICKS_CLI_PATH", arg0) } ctx := context.Background() - err := root.Execute(ctx, cmd.New(ctx)) + cli := cmd.New(ctx) + cli.SetArgs(args) + err = root.Execute(ctx, cli) if err != nil { os.Exit(1) } diff --git a/main_test.go b/main_test.go index 0f93f0236bd..84b21a2c800 100644 --- a/main_test.go +++ b/main_test.go @@ -36,6 +36,37 @@ func TestImportDoesNotSetCliPathEnv(t *testing.T) { assert.NotEqual(t, os.Args[0], os.Getenv("DATABRICKS_CLI_PATH")) } +func TestCommandArgsForDockerCredentialHelper(t *testing.T) { + tests := []struct { + name string + executable string + args []string + want []string + wantHelper bool + wantError string + }{ + {name: "databricks", executable: "databricks", args: []string{"auth", "token"}, want: []string{"auth", "token"}}, + {name: "helper", executable: "/usr/local/bin/docker-credential-databricks", args: []string{"get"}, want: []string{"auth", "token", "--format=docker"}, wantHelper: true}, + {name: "Windows helper", executable: `C:\Program Files\Databricks\docker-credential-databricks.exe`, args: []string{"get"}, want: []string{"auth", "token", "--format=docker"}, wantHelper: true}, + {name: "unsupported operation", executable: "docker-credential-databricks", args: []string{"store"}, wantHelper: true, wantError: "only supports get"}, + {name: "uppercase operation", executable: "docker-credential-databricks.exe", args: []string{"GET"}, wantHelper: true, wantError: "only supports get"}, + {name: "missing operation", executable: "docker-credential-databricks", wantHelper: true, wantError: "only supports get"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, helper, err := commandArgs(tt.executable, tt.args) + if tt.wantError != "" { + require.ErrorContains(t, err, tt.wantError) + return + } + require.NoError(t, err) + require.Equal(t, tt.want, got) + require.Equal(t, tt.wantHelper, helper) + }) + } +} + func TestFilePath(t *testing.T) { // To import this repository as a library, all files must match the // file path constraints made by Go. This test ensures that all files