coder · Emyrk · Mar 16, 2023 · Mar 14, 2023 · Mar 15, 2023 · Mar 15, 2023
diff --git a/cli/configssh.go b/cli/configssh.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"net/http"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -48,6 +49,43 @@ type sshConfigOptions struct {
 	sshOptions []string
 }
 
+// addOptions expects options in the form of "option=value" or "option value".
+// It will override any existing option with the same key to prevent duplicates.
+// Invalid options will return an error.
+func (o *sshConfigOptions) addOptions(options ...string) error {
+	for _, option := range options {
+		err := o.addOption(option)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (o *sshConfigOptions) addOption(option string) error {
+	key, _, err := codersdk.ParseSSHConfigOption(option)
+	if err != nil {
+		return err
+	}
+	for i, existing := range o.sshOptions {
+		// Override existing option if they share the same key.
+		// This is case-insensitive. Parsing each time might be a little slow,
+		// but it is ok.
+		existingKey, _, err := codersdk.ParseSSHConfigOption(existing)
+		if err != nil {
+			// Don't mess with original values if there is an error.
+			// This could have come from the user's manual edits.
+			continue
+		}
+		if strings.EqualFold(existingKey, key) {
+			o.sshOptions[i] = option
+			return nil
+		}
+	}
+	o.sshOptions = append(o.sshOptions, option)
+	return nil
+}
+
 func (o sshConfigOptions) equal(other sshConfigOptions) bool {
 	// Compare without side-effects or regard to order.
 	opt1 := slices.Clone(o.sshOptions)
@@ -139,6 +177,7 @@ func configSSH() *cobra.Command {
 		usePreviousOpts  bool
 		dryRun           bool
 		skipProxyCommand bool
+		userHostPrefix   string
 	)
 	cmd := &cobra.Command{
 		Annotations: workspaceCommand,
@@ -156,12 +195,13 @@ func configSSH() *cobra.Command {
 		),
 		Args: cobra.ExactArgs(0),
 		RunE: func(cmd *cobra.Command, _ []string) error {
+			ctx := cmd.Context()
 			client, err := CreateClient(cmd)
 			if err != nil {
 				return err
 			}
 
-			recvWorkspaceConfigs := sshPrepareWorkspaceConfigs(cmd.Context(), client)
+			recvWorkspaceConfigs := sshPrepareWorkspaceConfigs(ctx, client)
 
 			out := cmd.OutOrStdout()
 			if dryRun {
@@ -220,6 +260,13 @@ func configSSH() *cobra.Command {
 			if usePreviousOpts && lastConfig != nil {
 				sshConfigOpts = *lastConfig
 			} else if lastConfig != nil && !sshConfigOpts.equal(*lastConfig) {
+				for _, v := range sshConfigOpts.sshOptions {
+					// If the user passes an invalid option, we should catch
+					// this early.
+					if _, _, err := codersdk.ParseSSHConfigOption(v); err != nil {
+						return xerrors.Errorf("invalid option from flag: %w", err)
+					}
+				}
 				newOpts := sshConfigOpts.asList()
 				newOptsMsg := "\n\n  New options: none"
 				if len(newOpts) > 0 {
@@ -269,42 +316,85 @@ func configSSH() *cobra.Command {
 			if err != nil {
 				return xerrors.Errorf("fetch workspace configs failed: %w", err)
 			}
+
+			coderdConfig, err := client.SSHConfiguration(ctx)
+			if err != nil {
+				// If the error is 404, this deployment does not support
+				// this endpoint yet. Do not error, just assume defaults.
+				// TODO: Remove this in 2 months (May 31, 2023). Just return the error
+				// 	and remove this 404 check.
+				var sdkErr *codersdk.Error
+				if !(xerrors.As(err, &sdkErr) && sdkErr.StatusCode() == http.StatusNotFound) {
+					return xerrors.Errorf("fetch coderd config failed: %w", err)
+				}
+				coderdConfig.HostnamePrefix = "coder."
+			}
+
+			if userHostPrefix != "" {
+				// Override with user flag.
+				coderdConfig.HostnamePrefix = userHostPrefix
+			}
+
 			// Ensure stable sorting of output.
 			slices.SortFunc(workspaceConfigs, func(a, b sshWorkspaceConfig) bool {
 				return a.Name < b.Name
 			})
 			for _, wc := range workspaceConfigs {
 				sort.Strings(wc.Hosts)
 				// Write agent configuration.
-				for _, hostname := range wc.Hosts {
-					configOptions := []string{
-						"Host coder." + hostname,
-					}
-					for _, option := range sshConfigOpts.sshOptions {
-						configOptions = append(configOptions, "\t"+option)
-					}
-					configOptions = append(configOptions,
-						"\tHostName coder."+hostname,
-						"\tConnectTimeout=0",
-						"\tStrictHostKeyChecking=no",
+				for _, workspaceHostname := range wc.Hosts {
+					sshHostname := fmt.Sprintf("%s%s", coderdConfig.HostnamePrefix, workspaceHostname)
+					defaultOptions := []string{
+						"HostName " + sshHostname,
+						"ConnectTimeout=0",
+						"StrictHostKeyChecking=no",
 						// Without this, the "REMOTE HOST IDENTITY CHANGED"
 						// message will appear.
-						"\tUserKnownHostsFile=/dev/null",
+						"UserKnownHostsFile=/dev/null",
 						// This disables the "Warning: Permanently added 'hostname' (RSA) to the list of known hosts."
 						// message from appearing on every SSH. This happens because we ignore the known hosts.
-						"\tLogLevel ERROR",
-					)
+						"LogLevel ERROR",
+					}
+
 					if !skipProxyCommand {
-						configOptions = append(
-							configOptions,
-							fmt.Sprintf(
-								"\tProxyCommand %s --global-config %s ssh --stdio %s",
-								escapedCoderBinary, escapedGlobalConfig, hostname,
-							),
-						)
+						defaultOptions = append(defaultOptions, fmt.Sprintf(
+							"ProxyCommand %s --global-config %s ssh --stdio %s",
+							escapedCoderBinary, escapedGlobalConfig, workspaceHostname,
+						))
+					}
+
+					var configOptions sshConfigOptions
+					// Add standard options.
+					err := configOptions.addOptions(defaultOptions...)
+					if err != nil {
+						return err
+					}
+
+					// Override with deployment options
+					for k, v := range coderdConfig.SSHConfigOptions {
+						opt := fmt.Sprintf("%s %s", k, v)
+						err := configOptions.addOptions(opt)
+						if err != nil {
+							return xerrors.Errorf("add coderd config option %q: %w", opt, err)
+						}
+					}
+					// Override with flag options
+					for _, opt := range sshConfigOpts.sshOptions {
+						err := configOptions.addOptions(opt)
+						if err != nil {
+							return xerrors.Errorf("add flag config option %q: %w", opt, err)
+						}
+					}
+
+					hostBlock := []string{
+						"Host " + sshHostname,
+					}
+					// Prefix with '\t'
+					for _, v := range configOptions.sshOptions {
+						hostBlock = append(hostBlock, "\t"+v)
 					}
 
-					_, _ = buf.WriteString(strings.Join(configOptions, "\n"))
+					_, _ = buf.WriteString(strings.Join(hostBlock, "\n"))
 					_ = buf.WriteByte('\n')
 				}
 			}
@@ -363,7 +453,7 @@ func configSSH() *cobra.Command {
 
 			if len(workspaceConfigs) > 0 {
 				_, _ = fmt.Fprintln(out, "You should now be able to ssh into your workspace.")
-				_, _ = fmt.Fprintf(out, "For example, try running:\n\n\t$ ssh coder.%s\n", workspaceConfigs[0].Name)
+				_, _ = fmt.Fprintf(out, "For example, try running:\n\n\t$ ssh %s%s\n", coderdConfig.HostnamePrefix, workspaceConfigs[0].Name)
 			} else {
 				_, _ = fmt.Fprint(out, "You don't have any workspaces yet, try creating one with:\n\n\t$ coder create <workspace>\n")
 			}
@@ -376,6 +466,7 @@ func configSSH() *cobra.Command {
 	cmd.Flags().BoolVarP(&skipProxyCommand, "skip-proxy-command", "", false, "Specifies whether the ProxyCommand option should be skipped. Useful for testing.")
 	_ = cmd.Flags().MarkHidden("skip-proxy-command")
 	cliflag.BoolVarP(cmd.Flags(), &usePreviousOpts, "use-previous-options", "", "CODER_SSH_USE_PREVIOUS_OPTIONS", false, "Specifies whether or not to keep options from previous run of config-ssh.")
+	cmd.Flags().StringVarP(&userHostPrefix, "ssh-host-prefix", "", "", "Override the default host prefix.")
 	cliui.AllowSkipPrompt(cmd)
 
 	return cmd

diff --git a/cli/configssh_internal_test.go b/cli/configssh_internal_test.go
@@ -5,6 +5,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"sort"
 	"strings"
 	"testing"
 
@@ -179,3 +180,80 @@ func Test_sshConfigExecEscape(t *testing.T) {
 		})
 	}
 }
+
+func Test_sshConfigOptions_addOption(t *testing.T) {
+	t.Parallel()
+	testCases := []struct {
+		Name        string
+		Start       []string
+		Add         []string
+		Expect      []string
+		ExpectError bool
+	}{
+		{
+			Name: "Empty",
+		},
+		{
+			Name: "AddOne",
+			Add:  []string{"foo bar"},
+			Expect: []string{
+				"foo bar",
+			},
+		},
+		{
+			Name: "Replace",
+			Start: []string{
+				"foo bar",
+			},
+			Add: []string{"Foo baz"},
+			Expect: []string{
+				"Foo baz",
+			},
+		},
+		{
+			Name: "AddAndReplace",
+			Start: []string{
+				"a b",
+				"foo bar",
+				"buzz bazz",
+			},
+			Add: []string{
+				"b c",
+				"A hello",
+				"hello world",
+			},
+			Expect: []string{
+				"foo bar",
+				"buzz bazz",
+				"b c",
+				"A hello",
+				"hello world",
+			},
+		},
+		{
+			Name:        "Error",
+			Add:         []string{"novalue"},
+			ExpectError: true,
+		},
+	}
+
+	for _, tt := range testCases {
+		tt := tt
+		t.Run(tt.Name, func(t *testing.T) {
+			t.Parallel()
+
+			o := sshConfigOptions{
+				sshOptions: tt.Start,
+			}
+			err := o.addOptions(tt.Add...)
+			if tt.ExpectError {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			sort.Strings(tt.Expect)
+			sort.Strings(o.sshOptions)
+			require.Equal(t, tt.Expect, o.sshOptions)
+		})
+	}
+}
diff --git a/cli/configssh_test.go b/cli/configssh_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
@@ -63,7 +64,18 @@ func sshConfigFileRead(t *testing.T, name string) string {
 func TestConfigSSH(t *testing.T) {
 	t.Parallel()
 
-	client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+	const hostname = "test-coder."
+	const expectedKey = "ConnectionAttempts"
+	client := coderdtest.New(t, &coderdtest.Options{
+		IncludeProvisionerDaemon: true,
+		ConfigSSH: codersdk.SSHConfigResponse{
+			HostnamePrefix: hostname,
+			SSHConfigOptions: map[string]string{
+				// Something we can test for
+				expectedKey: "3",
+			},
+		},
+	})
 	user := coderdtest.CreateFirstUser(t, client)
 	authToken := uuid.NewString()
 	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
@@ -181,9 +193,13 @@ func TestConfigSSH(t *testing.T) {
 
 	<-doneChan
 
+	fileContents, err := os.ReadFile(sshConfigFile)
+	require.NoError(t, err, "read ssh config file")
+	require.Contains(t, string(fileContents), expectedKey, "ssh config file contains expected key")
+
 	home := filepath.Dir(filepath.Dir(sshConfigFile))
 	// #nosec
-	sshCmd := exec.Command("ssh", "-F", sshConfigFile, "coder."+workspace.Name, "echo", "test")
+	sshCmd := exec.Command("ssh", "-F", sshConfigFile, hostname+workspace.Name, "echo", "test")
 	pty = ptytest.New(t)
 	// Set HOME because coder config is included from ~/.ssh/coder.
 	sshCmd.Env = append(sshCmd.Env, fmt.Sprintf("HOME=%s", home))

diff --git a/cli/server.go b/cli/server.go
@@ -672,6 +672,11 @@ flags, and YAML configuration. The precedence is as follows:
 				return xerrors.Errorf("parse real ip config: %w", err)
 			}
 
+			configSSHOptions, err := cfg.SSHConfig.ParseOptions()
+			if err != nil {
+				return xerrors.Errorf("parse ssh config options %q: %w", cfg.SSHConfig.SSHConfigOptions.String(), err)
+			}
+
 			options := &coderd.Options{
 				AccessURL:                   cfg.AccessURL.Value(),
 				AppHostname:                 appHostname,
@@ -696,6 +701,10 @@ flags, and YAML configuration. The precedence is as follows:
 				LoginRateLimit:              loginRateLimit,
 				FilesRateLimit:              filesRateLimit,
 				HTTPClient:                  httpClient,
+				SSHConfig: codersdk.SSHConfigResponse{
+					HostnamePrefix:   cfg.SSHConfig.DeploymentName.String(),
+					SSHConfigOptions: configSSHOptions,
+				},
 			}
 			if tlsConfig != nil {
 				options.TLSCertificates = tlsConfig.Certificates

diff --git a/cli/testdata/coder_config-ssh_--help.golden b/cli/testdata/coder_config-ssh_--help.golden
@@ -19,6 +19,7 @@ Flags:
   -h, --help                     help for config-ssh
       --ssh-config-file string   Specifies the path to an SSH config.
                                  Consumes $CODER_SSH_CONFIG_FILE (default "~/.ssh/config")
+      --ssh-host-prefix string   Override the default host prefix.
   -o, --ssh-option stringArray   Specifies additional SSH options to embed in each host stanza.
       --use-previous-options     Specifies whether or not to keep options from previous run of
                                  config-ssh.