diff --git a/.nextchanges/cli/configure-docker.md b/.nextchanges/cli/configure-docker.md new file mode 100644 index 0000000000..80b374df4d --- /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 7ef3a9f72a..6d803b321d 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 0000000000..1819b5e7c5 --- /dev/null +++ b/cmd/auth/configure_docker.go @@ -0,0 +1,272 @@ +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(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 newConfigureDockerCommand() *cobra.Command { + return newConfigureDockerCommandWithDeps(defaultConfigureDockerDeps()) +} + +func newConfigureDockerCommandWithDeps(deps configureDockerDeps) *cobra.Command { + var region string + + 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.`, + Args: cobra.MaximumNArgs(1), + } + 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 + } + + workspaceID, err := resolveConfigureDockerWorkspaceID(ctx, p, 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) + } + } + + executable, err := deps.executable() + if err != nil { + return fmt.Errorf("locate databricks executable: %w", 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 so Docker can find it", installDir)) + } + return nil + } + + return cmd +} + +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 +} + +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 loadAndValidateConfigureDockerProfile(ctx context.Context, profileName string, profiler profile.Profiler) (profile.Profile, error) { + profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName)) + if err != nil { + return profile.Profile{}, err + } + if len(profiles) == 0 { + return profile.Profile{}, fmt.Errorf("profile %q not found", profileName) + } + if err := validateDockerCredentialProfile(profiles[0]); err != nil { + return profile.Profile{}, err + } + return profiles[0], nil +} + +func resolveConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, 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: authlib.WorkspaceIDNone, + AuthType: p.AuthType, + 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, 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 + // Workspace IDs can repeat across environments, so only profiles producing this registry host are duplicates. + 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 ")) +} + +func persistConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, workspaceID string) error { + return databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: env.Get(ctx, "DATABRICKS_CONFIG_FILE"), + Profile: p.Name, + WorkspaceID: workspaceID, + }) +} + +func configureDockerConfigPath(ctx context.Context) (string, error) { + if dockerConfig := env.Get(ctx, "DOCKER_CONFIG"); dockerConfig != "" { + return filepath.Join(dockerConfig, "config.json"), nil + } + home, err := env.UserHomeDir(ctx) + if err != nil { + return "", err + } + return filepath.Join(home, ".docker", "config.json"), nil +} diff --git a/cmd/auth/configure_docker_test.go b/cmd/auth/configure_docker_test.go new file mode 100644 index 0000000000..789b50ba68 --- /dev/null +++ b/cmd/auth/configure_docker_test.go @@ -0,0 +1,477 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "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/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 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", "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, + AuthType: authTypeDatabricksCLI, + }) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DATABRICKS_WORKSPACE_ID", "ambient-workspace") + t.Setenv("DOCKER_CONFIG", dockerDir) + t.Setenv("HOME", homeDir) + + 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, authlib.WorkspaceIDNone, cfg.WorkspaceID) + require.Equal(t, authTypeDatabricksCLI, cfg.AuthType) + 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(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")) +} diff --git a/libs/dockercredentials/docker_config.go b/libs/dockercredentials/docker_config.go new file mode 100644 index 0000000000..c44b05c29d --- /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. +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 +} + +func resolveDockerConfigPath(path string) (string, error) { + info, err := os.Lstat(path) + if errors.Is(err, fs.ErrNotExist) { + return path, nil + } + if err != nil { + return "", fmt.Errorf("inspect Docker config %s: %w", path, err) + } + if info.Mode()&os.ModeSymlink == 0 { + return path, nil + } + + resolved, err := filepath.EvalSymlinks(path) + if err != nil { + return "", fmt.Errorf("resolve Docker config symlink %s: %w", path, err) + } + return resolved, nil +} + +func readDockerConfig(path string) (map[string]json.RawMessage, error) { + raw, err := os.ReadFile(path) + if errors.Is(err, fs.ErrNotExist) { + return map[string]json.RawMessage{}, nil + } + if err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + + var config map[string]json.RawMessage + if err := json.Unmarshal(raw, &config); err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + if config == nil { + config = map[string]json.RawMessage{} + } + return config, nil +} + +func writeDockerConfig(path string, config map[string]json.RawMessage) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("create Docker config directory %s: %w", dir, err) + } + + raw, err := json.MarshalIndent(config, "", " ") + if err != nil { + return err + } + raw = append(raw, '\n') + + tmp, err := os.CreateTemp(dir, ".config.json.*") + if err != nil { + return fmt.Errorf("create temporary Docker config in %s: %w", dir, err) + } + tmpPath := tmp.Name() + defer func() { + _ = os.Remove(tmpPath) + }() + + if _, err := tmp.Write(raw); err != nil { + _ = tmp.Close() + return fmt.Errorf("write temporary Docker config %s: %w", tmpPath, err) + } + if err := tmp.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 0000000000..c41cafe7de --- /dev/null +++ b/libs/dockercredentials/docker_config_test.go @@ -0,0 +1,144 @@ +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 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/shim.go b/libs/dockercredentials/shim.go new file mode 100644 index 0000000000..8c560f895a --- /dev/null +++ b/libs/dockercredentials/shim.go @@ -0,0 +1,159 @@ +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 helper in installDir that delegates credential requests to databricksPath. +func InstallShim(databricksPath, installDir string) (ShimInstallResult, error) { + return installShimForGOOS(databricksPath, installDir, runtime.GOOS) +} + +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 +} + +func shimFilename(goos string) string { + if goos == "windows" { + return "docker-credential-" + HelperName + ".exe" + } + return "docker-credential-" + HelperName +} + +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)) +} + +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 + }) +} + +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 + }) +} + +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) +} + +func posixShellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +func helperOnPathForGOOS(helperPath, goos string, lookPath func(string) (string, error)) bool { + name := filepath.Base(helperPath) + if goos == "windows" { + name = strings.TrimSuffix(name, filepath.Ext(name)) + } + candidate, err := lookPath(name) + if err != nil { + return false + } + return samePath(candidate, helperPath) +} + +func samePath(a, b string) bool { + aInfo, aErr := os.Stat(a) + 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 0000000000..4c542a5d8f --- /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" +) + +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") + } + + 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) + 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)) +} diff --git a/main.go b/main.go index 6c4dddd40b..9eead3d23a 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,46 @@ import ( _ "github.com/databricks/cli/libs/hostmetadata" ) +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 0f93f0236b..84b21a2c80 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