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/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..a3ecc3ef00e --- /dev/null +++ b/cmd/auth/configure_docker.go @@ -0,0 +1,324 @@ +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" +) + +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(context.Context, string, string) (dockercredentials.ShimInstallResult, error) + setCredentialHelper func(string, string) error +} + +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, + } +} + +func (d configureDockerDeps) withDefaults() configureDockerDeps { + defaults := defaultConfigureDockerDeps() + if d.profiler == nil { + d.profiler = defaults.profiler + } + if d.newWorkspaceClient == nil { + d.newWorkspaceClient = defaults.newWorkspaceClient + } + if d.resolveWorkspaceID == nil { + d.resolveWorkspaceID = defaults.resolveWorkspaceID + } + if d.executable == nil { + d.executable = defaults.executable + } + if d.registryHost == nil { + d.registryHost = defaults.registryHost + } + if d.installShim == nil { + d.installShim = defaults.installShim + } + if d.setCredentialHelper == nil { + d.setCredentialHelper = defaults.setCredentialHelper + } + return d +} + +func newConfigureDockerCommand() *cobra.Command { + return newConfigureDockerCommandWithDeps(defaultConfigureDockerDeps()) +} + +func newConfigureDockerCommandWithDeps(deps configureDockerDeps) *cobra.Command { + deps = deps.withDefaults() + var region string + + cmd := &cobra.Command{ + Use: "configure-docker [PROFILE]", + 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.`, + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + if err := rejectConfigureDockerAuthFlags(cmd); err != nil { + return err + } + if region == "" { + return errors.New("--region is required because workspace region cannot be inferred from this profile") + } + + profileName, err := configureDockerProfileName(ctx, cmd, args, deps.profiler) + if err != nil { + return err + } + + p, err := loadConfigureDockerProfile(ctx, profileName, deps.profiler) + if err != nil { + return err + } + if err := validateConfigureDockerProfile(p); err != nil { + return err + } + if isConfigureDockerAccountOnlyProfile(p) { + return fmt.Errorf("profile %q does not target a workspace. Run databricks auth login --host and retry with that profile", p.Name) + } + + workspaceID, err := resolveConfigureDockerWorkspaceID(ctx, p, deps) + if err != nil { + return err + } + if err := ensureConfigureDockerUniqueProfile(ctx, deps.profiler, p, workspaceID); err != nil { + return err + } + registryHost, err := deps.registryHost(workspaceID, region, p.Host) + if 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) + } + } + + executable, err := deps.executable() + if err != nil { + return fmt.Errorf("locate databricks executable: %w", err) + } + installDir, err := configureDockerShimInstallDir(ctx) + if err != nil { + return err + } + shim, err := deps.installShim(ctx, 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, fmt.Sprintf("Configured Docker credential helper for %s", registryHost)) + cmdio.LogString(ctx, fmt.Sprintf("Updated Docker config: %s", dockerConfigPath)) + cmdio.LogString(ctx, fmt.Sprintf("Installed Docker credential helper: %s", shim.Path)) + if !shim.OnPath { + cmdio.LogString(ctx, fmt.Sprintf("Warning: ensure %s is on PATH before any other docker-credential-databricks helper so Docker can find it", installDir)) + } + return nil + }, + } + cmd.Flags().StringVar(®ion, "region", "", "Cloud region for the Databricks Artifact Registry host") + + return cmd +} + +func rejectConfigureDockerAuthFlags(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 +} + +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 +} + +func loadConfigureDockerProfile(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) + } + return profiles[0], nil +} + +func validateConfigureDockerProfile(p profile.Profile) error { + if p.HasClientCredentials { + return fmt.Errorf("profile %q uses client credentials. databricks auth configure-docker requires a profile created by databricks auth login", p.Name) + } + if p.AuthType != authTypeDatabricksCLI { + return fmt.Errorf("profile %q uses auth_type %q. databricks auth configure-docker requires a profile created by databricks auth login", p.Name, p.AuthType) + } + return nil +} + +func isConfigureDockerAccountOnlyProfile(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) +} + +func resolveConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, 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, + WorkspaceID: p.WorkspaceID, + ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), + } + 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) + } + 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 +} + +func ensureConfigureDockerUniqueProfile(ctx context.Context, profiler profile.Profiler, p profile.Profile, workspaceID string) error { + matches, err := profiler.LoadProfiles(ctx, func(candidate profile.Profile) bool { + return candidate.WorkspaceID == workspaceID + }) + if err != nil { + return err + } + + names := matches.Names() + 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. Make the profile selection unambiguous before using Docker credential helper", workspaceID, strings.Join(names, " and ")) +} + +func persistConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, workspaceID string) error { + return databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), + Profile: p.Name, + Host: p.Host, + AccountID: p.AccountID, + WorkspaceID: workspaceID, + AuthType: p.AuthType, + ClusterID: p.ClusterID, + Scopes: splitScopes(p.Scopes), + AzureTenantID: "", + }) +} + +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 +} + +func configureDockerShimInstallDir(ctx context.Context) (string, error) { + home, err := env.UserHomeDir(ctx) + if err != nil { + return "", err + } + return filepath.Join(home, ".databricks", "bin"), nil +} diff --git a/cmd/auth/configure_docker_test.go b/cmd/auth/configure_docker_test.go new file mode 100644 index 00000000000..64f70b52367 --- /dev/null +++ b/cmd/auth/configure_docker_test.go @@ -0,0 +1,369 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/cli/libs/dockercredentials" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func 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 +} + +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 +} + +func writeConfigureDockerProfile(t *testing.T, ctx context.Context, configFile string, cfg *config.Config) { + t.Helper() + cfg.ConfigFile = configFile + require.NoError(t, databrickscfg.SaveToProfile(ctx, cfg)) +} + +func readCredentialHelpers(t *testing.T, path string) map[string]string { + t.Helper() + raw, err := os.ReadFile(path) + require.NoError(t, err) + + var cfg struct { + CredHelpers map[string]string `json:"credHelpers"` + } + require.NoError(t, json.Unmarshal(raw, &cfg)) + return cfg.CredHelpers +} + +func configureDockerRegistryHostStub(t *testing.T, wantWorkspaceID, wantRegion, wantWorkspaceHost, registryHost string) func(string, string, string) (string, error) { + t.Helper() + return func(workspaceID, region, workspaceHost string) (string, error) { + require.Equal(t, wantWorkspaceID, workspaceID) + require.Equal(t, wantRegion, region) + require.Equal(t, wantWorkspaceHost, workspaceHost) + return registryHost, nil + } +} + +func TestConfigureDockerCommandWritesDockerConfigAndShim(t *testing.T) { + ctx, stderr := cmdio.NewTestContextWithStderr(t.Context()) + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := filepath.Join(dir, "docker") + homeDir := filepath.Join(dir, "home") + binDir := filepath.Join(homeDir, ".databricks", "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("HOME", homeDir) + t.Setenv("PATH", binDir) + + registryHost := "123456789.container.us-west-2.staging.cloud.databricks.test" + deps := defaultConfigureDockerDeps() + deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, registryHost) + + cmd := newConfigureDockerTestCommandWithDeps(ctx, deps, "configure-docker", "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]) + + _, err := os.Stat(filepath.Join(binDir, "docker-credential-databricks")) + require.NoError(t, err) + assert.Contains(t, stderr.String(), registryHost) + assert.Contains(t, stderr.String(), filepath.Join(dockerDir, "config.json")) +} + +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") +} + +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, + AuthType: authTypeDatabricksCLI, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", homeDir) + + 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 "999999", nil + } + 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(context.Context, 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.installShim = func(context.Context, 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.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +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(context.Context, 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")) +} diff --git a/cmd/auth/token.go b/cmd/auth/token.go index d5e88e64d72..b7c93d5ac17 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,13 @@ 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) + func newTokenCommand(authArguments *auth.AuthArguments) *cobra.Command { + return newTokenCommandWithLoader(authArguments, loadToken) +} + +func newTokenCommandWithLoader(authArguments *auth.AuthArguments, load tokenLoader) *cobra.Command { cmd := &cobra.Command{ Use: "token [PROFILE]", Short: "Get authentication token", @@ -50,6 +57,10 @@ and secret is not supported.`, cmd.Flags().BoolVar(&forceRefresh, "force-refresh", false, "Force a token refresh even if the cached token is still valid.") + var format string + cmd.Flags().StringVar(&format, "format", "", "Hidden output format") + _ = cmd.Flags().MarkHidden("format") + cmd.PreRunE = profileHostConflictCheck cmd.RunE = func(cmd *cobra.Command, args []string) error { @@ -61,7 +72,7 @@ and secret is not supported.`, return err } - t, err := loadToken(ctx, loadTokenArgs{ + loadArgs := loadTokenArgs{ authArguments: authArguments, profileName: profileName, args: args, @@ -71,7 +82,16 @@ and secret is not supported.`, tokenStore: tokenStore, mode: mode, persistentAuthOpts: nil, - }) + } + + if format == "docker" { + return writeDockerTokenOutput(ctx, cmd, loadArgs, load) + } + if format != "" { + return fmt.Errorf("unsupported token format %q", format) + } + + t, err := load(ctx, loadArgs) if err != nil { return err } @@ -85,6 +105,70 @@ and secret is not supported.`, return cmd } +type dockerGetResponse struct { + Username string `json:"Username"` + Secret string `json:"Secret"` +} + +func writeDockerTokenOutput(ctx context.Context, cmd *cobra.Command, args loadTokenArgs, load tokenLoader) error { + if len(args.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) + } + } + + 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) + 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, + }) +} + +func dockerTokenProfileName(ctx context.Context, registry dockercredentials.Registry, profiler profile.Profiler) (string, error) { + matchingProfiles, err := profiler.LoadProfiles(ctx, func(p profile.Profile) bool { + return p.WorkspaceID == registry.WorkspaceID + }) + if err != nil { + return "", err + } + if len(matchingProfiles) == 0 { + return "", fmt.Errorf("no Databricks profile found for workspace ID %s from registry host %s. Run databricks auth login --host or databricks auth configure-docker", registry.WorkspaceID, registry.Host) + } + if len(matchingProfiles) > 1 { + return "", fmt.Errorf("multiple Databricks profiles match workspace ID %s: %s. Make the profile selection unambiguous before using Docker credential helper", registry.WorkspaceID, strings.Join(matchingProfiles.Names(), " and ")) + } + if err := validateConfigureDockerProfile(matchingProfiles[0]); err != nil { + return "", err + } + return matchingProfiles[0].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..9412f8e3f7f 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -6,17 +6,24 @@ 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/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 +892,255 @@ 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 + } + + var stdout bytes.Buffer + cmd := newTokenCommandWithLoader(&auth.AuthArguments{}, loadToken) + cmd.Flags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader("123456789.container.us-west-2.cloud.databricks.com\n")) + cmd.SetOut(&stdout) + 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) + require.NoError(t, err) + require.Equal(t, "workspace", gotProfile) + + var got dockerGetResponse + require.NoError(t, json.Unmarshal(stdout.Bytes(), &got)) +} + +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 TestTokenDockerFormatRejectsAuthSelectionFlags(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + t.Setenv(storage.EnvVar, string(storage.StorageModePlaintext)) + t.Setenv("HOME", t.TempDir()) + + cases := [][]string{ + {"--format=docker", "--profile", "DEFAULT"}, + {"--format=docker", "--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 configure-docker") +} + +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 := 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, "multiple Databricks profiles match workspace ID 123456789") + require.ErrorContains(t, err, "one and two") +} + +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", + }, + }, + } + + for _, registryHost := range []string{ + "123456789.container.us-west-2.cloud.databricks.com", + "987654321.container.us-west-2.cloud.databricks.com", + "111222333.container.us-west-2.cloud.databricks.com", + } { + t.Run(registryHost, func(t *testing.T) { + cmd := &cobra.Command{Use: "token"} + cmd.SetContext(ctx) + cmd.SetIn(strings.NewReader(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 + }) + require.ErrorContains(t, err, "requires a profile created by databricks auth login") + }) + } +} + 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..2dea312f25f --- /dev/null +++ b/libs/dockercredentials/docker_config.go @@ -0,0 +1,100 @@ +package dockercredentials + +import ( + "encoding/json" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" +) + +func SetCredentialHelper(path, registryHost string) error { + 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 +} + +func readDockerConfig(path string) (map[string]json.RawMessage, error) { + raw, err := os.ReadFile(path) + if errors.Is(err, fs.ErrNotExist) { + return map[string]json.RawMessage{}, nil + } + if err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + + var config map[string]json.RawMessage + if err := json.Unmarshal(raw, &config); err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + if config == nil { + config = map[string]json.RawMessage{} + } + return config, nil +} + +func writeDockerConfig(path string, config map[string]json.RawMessage) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("create Docker config directory %s: %w", dir, err) + } + + raw, err := json.MarshalIndent(config, "", " ") + if err != nil { + return err + } + raw = append(raw, '\n') + + tmp, err := os.CreateTemp(dir, ".config.json.*") + if err != nil { + return fmt.Errorf("create temporary Docker config in %s: %w", dir, err) + } + tmpPath := tmp.Name() + defer func() { + _ = os.Remove(tmpPath) + }() + + if _, err := tmp.Write(raw); err != nil { + _ = tmp.Close() + return fmt.Errorf("write temporary Docker config %s: %w", tmpPath, err) + } + if err := tmp.Chmod(0o600); err != nil { + _ = tmp.Close() + return fmt.Errorf("set permissions on temporary Docker config %s: %w", tmpPath, err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("close temporary Docker config %s: %w", tmpPath, err) + } + if err := os.Rename(tmpPath, path); err != nil { + return fmt.Errorf("write Docker config %s: %w", path, err) + } + return nil +} diff --git a/libs/dockercredentials/docker_config_test.go b/libs/dockercredentials/docker_config_test.go new file mode 100644 index 00000000000..b1aacdd16e9 --- /dev/null +++ b/libs/dockercredentials/docker_config_test.go @@ -0,0 +1,125 @@ +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" + +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 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/executable_unix.go b/libs/dockercredentials/executable_unix.go new file mode 100644 index 00000000000..2df8037258e --- /dev/null +++ b/libs/dockercredentials/executable_unix.go @@ -0,0 +1,25 @@ +//go:build !windows + +package dockercredentials + +import ( + "errors" + "os" + + "golang.org/x/sys/unix" +) + +func isExecutableForPath(path string, mode os.FileMode, goos string) bool { + if goos == "windows" { + return true + } + + err := unix.Faccessat(unix.AT_FDCWD, path, unix.X_OK, unix.AT_EACCESS) + if err == nil { + return true + } + if errors.Is(err, unix.ENOSYS) { + return mode&0o111 != 0 + } + return false +} diff --git a/libs/dockercredentials/executable_windows.go b/libs/dockercredentials/executable_windows.go new file mode 100644 index 00000000000..4ed3da0edc2 --- /dev/null +++ b/libs/dockercredentials/executable_windows.go @@ -0,0 +1,12 @@ +//go:build windows + +package dockercredentials + +import "os" + +func isExecutableForPath(_ string, mode os.FileMode, goos string) bool { + if goos == "windows" { + return true + } + return mode&0o111 != 0 +} diff --git a/libs/dockercredentials/registry.go b/libs/dockercredentials/registry.go new file mode 100644 index 00000000000..87ab7316c81 --- /dev/null +++ b/libs/dockercredentials/registry.go @@ -0,0 +1,176 @@ +package dockercredentials + +import ( + "errors" + "fmt" + "net" + "net/url" + "strconv" + "strings" + "unicode" + + "github.com/databricks/databricks-sdk-go/common/environment" +) + +const ( + HelperName = "databricks" + OAuthTokenUsername = "oauthtoken" + registryHostInfix = ".container." +) + +type Registry struct { + WorkspaceID string + Region string + Host string +} + +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 +} + +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) + } + 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 + } + } + + value = strings.TrimSuffix(strings.ToLower(value), ".") + if value == "" { + return "", errors.New("server address is required") + } + return value, nil +} + +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 +} + +func registryDNSZoneForWorkspaceHost(raw string) (string, error) { + host, err := normalizeServerAddress(raw) + if err != nil { + return "", fmt.Errorf("parse workspace host: %w", err) + } + dnsZone, ok := matchingDatabricksDNSZone(host) + if !ok { + return "", fmt.Errorf("%q is not a supported Databricks workspace host", host) + } + return dnsZone, nil +} + +func matchingDatabricksDNSZone(host string) (string, bool) { + return matchingDatabricksDNSZoneInEnvironments(host, environment.AllEnvironments()) +} + +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 != "" +} + +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 +} + +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 +} + +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..5e6040cad30 --- /dev/null +++ b/libs/dockercredentials/registry_test.go @@ -0,0 +1,171 @@ +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 TestNormalizeServerAddressRejectsInvalidPort(t *testing.T) { + _, err := normalizeServerAddress("https://123.container.us-west-2.cloud.databricks.com:99999") + require.ErrorContains(t, err, "invalid registry port") +} diff --git a/libs/dockercredentials/shim.go b/libs/dockercredentials/shim.go new file mode 100644 index 00000000000..107b55f7452 --- /dev/null +++ b/libs/dockercredentials/shim.go @@ -0,0 +1,207 @@ +package dockercredentials + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + + "github.com/databricks/cli/libs/env" +) + +type ShimInstallResult struct { + Path string + OnPath bool +} + +func InstallShim(ctx context.Context, databricksPath, installDir 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(runtime.GOOS)) + mode := os.FileMode(0o755) + if runtime.GOOS == "windows" { + mode = 0o644 + } + if err := writeShimFile(path, []byte(shimScript(databricksPath, runtime.GOOS)), mode); err != nil { + return ShimInstallResult{}, fmt.Errorf("write Docker credential helper %s: %w", path, err) + } + + return ShimInstallResult{ + Path: path, + OnPath: helperOnPath(ctx, path), + }, nil +} + +func shimFilename(goos string) string { + if goos == "windows" { + return "docker-credential-" + HelperName + ".cmd" + } + return "docker-credential-" + HelperName +} + +func shimScript(databricksPath, goos string) string { + if goos == "windows" { + return fmt.Sprintf(`@echo off +setlocal DisableDelayedExpansion +if /I not "%%~1"=="get" ( + echo docker-credential-databricks only supports get 1>&2 + exit /b 1 +) +set "DATABRICKS_LOG_FILE=stderr" +shift /1 +%s auth token --format=docker +`, windowsBatchQuote(databricksPath)) + } + + return fmt.Sprintf(`#!/bin/sh +if [ "${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)) +} + +func writeShimFile(path string, script []byte, mode os.FileMode) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*") + if err != nil { + return err + } + tmpPath := tmp.Name() + defer func() { + _ = os.Remove(tmpPath) + }() + + if _, err := tmp.Write(script); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Chmod(mode); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + return os.Rename(tmpPath, path) +} + +func posixShellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +func windowsBatchQuote(value string) string { + var b strings.Builder + b.WriteByte('"') + for _, r := range value { + switch r { + case '%': + b.WriteString("%%") + case '^', '&', '|', '<', '>': + b.WriteByte('^') + b.WriteRune(r) + default: + b.WriteRune(r) + } + } + b.WriteByte('"') + return b.String() +} + +func helperOnPath(ctx context.Context, helperPath string) bool { + return helperOnPathForGOOS(helperPath, runtime.GOOS, func(key string) string { + return env.Get(ctx, key) + }) +} + +func helperOnPathForGOOS(helperPath, goos string, getenv func(string) string) bool { + helperName := helperLookupName(helperPath, goos) + pathEntries := filepath.SplitList(getenv("PATH")) + if len(pathEntries) == 0 && goos != "windows" { + pathEntries = []string{""} + } + for _, entry := range pathEntries { + if entry == "" { + if goos == "windows" { + continue + } + entry = "." + } + for _, candidateName := range helperLookupNames(helperName, goos, getenv("PATHEXT")) { + candidate := filepath.Join(entry, candidateName) + info, err := os.Stat(candidate) + if err != nil || info.IsDir() || !isExecutableForPath(candidate, info.Mode(), goos) { + continue + } + return samePath(candidate, helperPath) + } + } + return false +} + +func helperLookupName(helperPath, goos string) string { + name := filepath.Base(helperPath) + if goos == "windows" { + return strings.TrimSuffix(name, filepath.Ext(name)) + } + return name +} + +func helperLookupNames(name, goos, pathext string) []string { + if goos != "windows" || filepath.Ext(name) != "" { + return []string{name} + } + if pathext == "" { + pathext = ".COM;.EXE;.BAT;.CMD" + } + + names := []string{} + for _, ext := range strings.Split(pathext, ";") { + if ext == "" { + continue + } + if !strings.HasPrefix(ext, ".") { + ext = "." + ext + } + names = append(names, name+strings.ToLower(ext)) + } + return names +} + +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..f5cca086078 --- /dev/null +++ b/libs/dockercredentials/shim_test.go @@ -0,0 +1,283 @@ +package dockercredentials + +import ( + "bytes" + "os" + "os/exec" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestShimFilename(t *testing.T) { + require.Equal(t, "docker-credential-databricks", shimFilename("linux")) + require.Equal(t, "docker-credential-databricks.cmd", shimFilename("windows")) +} + +func TestUnixShimScript(t *testing.T) { + got := shimScript("/opt/databricks/bin/databricks", "linux") + + require.Contains(t, got, `if [ "${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 TestWindowsShimScript(t *testing.T) { + got := shimScript(`C:\Program Files (x86)\Data%bricks & CLI\!DATABRICKS_LOG_FILE!\databricks.exe`, "windows") + + require.Contains(t, got, `setlocal DisableDelayedExpansion`) + require.Contains(t, got, `if /I not "%~1"=="get"`) + require.Contains(t, got, `docker-credential-databricks only supports get`) + require.Contains(t, got, `shift /1`) + require.Contains(t, got, `set "DATABRICKS_LOG_FILE=stderr"`) + require.Contains(t, got, `"C:\Program Files (x86)\Data%%bricks ^& CLI\!DATABRICKS_LOG_FILE!\databricks.exe" auth token --format=docker`) +} + +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, "linux")), 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) +} + +func TestInstallShimReportsPathStatus(t *testing.T) { + dir := t.TempDir() + t.Setenv("PATH", dir) + + got, err := InstallShim(t.Context(), "/usr/local/bin/databricks", 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() + t.Setenv("PATH", t.TempDir()) + + got, err := InstallShim(t.Context(), "/usr/local/bin/databricks", 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() + 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(t.Context(), "/usr/local/bin/databricks", installDir) + require.NoError(t, err) + require.Equal(t, filepath.Join(installDir, shimFilename(runtime.GOOS)), got.Path) + require.False(t, got.OnPath) +} + +func TestHelperOnPathUsesWindowsPathExtOrder(t *testing.T) { + installDir := t.TempDir() + shadowDir := t.TempDir() + helperPath := filepath.Join(installDir, "docker-credential-databricks.cmd") + shadowPath := filepath.Join(shadowDir, "docker-credential-databricks.exe") + require.NoError(t, os.WriteFile(shadowPath, []byte("shadow"), 0o755)) + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o644)) + + env := map[string]string{ + "PATH": shadowDir + string(os.PathListSeparator) + installDir, + "PATHEXT": ".EXE;.CMD", + } + require.False(t, helperOnPathForGOOS(helperPath, "windows", func(key string) string { + return env[key] + })) +} + +func TestHelperOnPathNormalizesWindowsPathExt(t *testing.T) { + installDir := t.TempDir() + helperPath := filepath.Join(installDir, "docker-credential-databricks.cmd") + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o644)) + + env := map[string]string{ + "PATH": installDir, + "PATHEXT": "EXE;CMD", + } + require.True(t, helperOnPathForGOOS(helperPath, "windows", func(key string) string { + return env[key] + })) +} + +func TestHelperOnPathIgnoresWindowsExtensionlessShadow(t *testing.T) { + installDir := t.TempDir() + shadowDir := t.TempDir() + helperPath := filepath.Join(installDir, "docker-credential-databricks.cmd") + shadowPath := filepath.Join(shadowDir, "docker-credential-databricks") + require.NoError(t, os.WriteFile(shadowPath, []byte("shadow"), 0o755)) + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o644)) + + env := map[string]string{ + "PATH": shadowDir + string(os.PathListSeparator) + installDir, + "PATHEXT": ".COM;.EXE;.BAT;.CMD", + } + require.True(t, helperOnPathForGOOS(helperPath, "windows", func(key string) string { + return env[key] + })) +} + +func TestHelperOnPathRequiresUnixExecutablePermission(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix mode-bit test") + } + + installDir := t.TempDir() + shadowDir := t.TempDir() + helperPath := filepath.Join(installDir, "docker-credential-databricks") + shadowPath := filepath.Join(shadowDir, "docker-credential-databricks") + require.NoError(t, os.WriteFile(shadowPath, []byte("shadow"), 0o644)) + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o755)) + + env := map[string]string{ + "PATH": shadowDir + string(os.PathListSeparator) + installDir, + } + require.True(t, helperOnPathForGOOS(helperPath, "linux", func(key string) string { + return env[key] + })) +} + +func TestHelperOnPathUsesUnixEffectiveUserExecutePermission(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix execute permission test") + } + if os.Geteuid() == 0 { + t.Skip("root can execute files that normal users cannot") + } + + installDir := t.TempDir() + shadowDir := t.TempDir() + helperPath := filepath.Join(installDir, "docker-credential-databricks") + shadowPath := filepath.Join(shadowDir, "docker-credential-databricks") + require.NoError(t, os.WriteFile(shadowPath, []byte("shadow"), 0o001)) + require.NoError(t, os.Chmod(shadowPath, 0o001)) + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o755)) + + env := map[string]string{ + "PATH": shadowDir + string(os.PathListSeparator) + installDir, + } + require.True(t, helperOnPathForGOOS(helperPath, "linux", func(key string) string { + return env[key] + })) +} + +func TestHelperOnPathTreatsEmptyUnixPathEntryAsWorkingDirectory(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix PATH semantics test") + } + + dir := t.TempDir() + helperPath := filepath.Join(dir, "docker-credential-databricks") + require.NoError(t, os.WriteFile(helperPath, []byte("helper"), 0o755)) + + oldwd, err := os.Getwd() + require.NoError(t, err) + require.NoError(t, os.Chdir(dir)) + t.Cleanup(func() { + _ = os.Chdir(oldwd) + }) + + env := map[string]string{"PATH": ""} + require.True(t, helperOnPathForGOOS(helperPath, "linux", func(key string) string { + return env[key] + })) +} + +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") + } + + installDir := filepath.Join(t.TempDir(), "missing") + 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(t.Context(), "/usr/local/bin/databricks", installDir) + if os.Geteuid() == 0 { + require.NoError(t, err) + return + } + require.Error(t, err) + + raw, readErr := os.ReadFile(shimPath) + require.NoError(t, readErr) + require.Equal(t, "existing helper", string(raw)) +}