feat(cli/configssh): add support for wait yes/no/auto (#7893)

mafredri · web-flow · commit a1c32954d9ac · 2023-06-08T14:06:50.000Z
Refs #7768
diff --git a/cli/configssh.go b/cli/configssh.go
@@ -45,7 +45,9 @@ const (
 // sshConfigOptions represents options that can be stored and read
 // from the coder config in ~/.ssh/coder.
 type sshConfigOptions struct {
-	sshOptions []string
+	waitEnum       string
+	userHostPrefix string
+	sshOptions     []string
 }
 
 // addOptions expects options in the form of "option=value" or "option value".
@@ -100,10 +102,19 @@ func (o sshConfigOptions) equal(other sshConfigOptions) bool {
 	sort.Strings(opt1)
 	opt2 := slices.Clone(other.sshOptions)
 	sort.Strings(opt2)
-	return slices.Equal(opt1, opt2)
+	if !slices.Equal(opt1, opt2) {
+		return false
+	}
+	return o.waitEnum == other.waitEnum && o.userHostPrefix == other.userHostPrefix
 }
 
 func (o sshConfigOptions) asList() (list []string) {
+	if o.waitEnum != "auto" {
+		list = append(list, fmt.Sprintf("wait: %s", o.waitEnum))
+	}
+	if o.userHostPrefix != "" {
+		list = append(list, fmt.Sprintf("ssh-host-prefix: %s", o.userHostPrefix))
+	}
 	for _, opt := range o.sshOptions {
 		list = append(list, fmt.Sprintf("ssh-option: %s", opt))
 	}
@@ -178,14 +189,14 @@ func sshPrepareWorkspaceConfigs(ctx context.Context, client *codersdk.Client) (r
 	}
 }
 
+//nolint:gocyclo
 func (r *RootCmd) configSSH() *clibase.Cmd {
 	var (
 		sshConfigFile    string
 		sshConfigOpts    sshConfigOptions
 		usePreviousOpts  bool
 		dryRun           bool
 		skipProxyCommand bool
-		userHostPrefix   string
 	)
 	client := new(codersdk.Client)
 	cmd := &clibase.Cmd{
@@ -207,6 +218,10 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
 			r.InitClient(client),
 		),
 		Handler: func(inv *clibase.Invocation) error {
+			if sshConfigOpts.waitEnum != "auto" && skipProxyCommand {
+				return xerrors.Errorf("cannot specify both --skip-proxy-command and --wait")
+			}
+
 			recvWorkspaceConfigs := sshPrepareWorkspaceConfigs(inv.Context(), client)
 
 			out := inv.Stdout
@@ -295,7 +310,7 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
 					// Selecting "no" will use the last config.
 					sshConfigOpts = *lastConfig
 				} else {
-					changes = append(changes, "Use new SSH options")
+					changes = append(changes, "Use new options")
 				}
 				// Only print when prompts are shown.
 				if yes, _ := inv.ParsedFlags().GetBool("yes"); !yes {
@@ -336,9 +351,9 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
 				coderdConfig.HostnamePrefix = "coder."
 			}
 
-			if userHostPrefix != "" {
+			if sshConfigOpts.userHostPrefix != "" {
 				// Override with user flag.
-				coderdConfig.HostnamePrefix = userHostPrefix
+				coderdConfig.HostnamePrefix = sshConfigOpts.userHostPrefix
 			}
 
 			// Ensure stable sorting of output.
@@ -363,13 +378,20 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
 					}
 
 					if !skipProxyCommand {
+						flags := ""
+						if sshConfigOpts.waitEnum != "auto" {
+							flags += " --wait=" + sshConfigOpts.waitEnum
+						}
 						defaultOptions = append(defaultOptions, fmt.Sprintf(
-							"ProxyCommand %s --global-config %s ssh --stdio %s",
-							escapedCoderBinary, escapedGlobalConfig, workspaceHostname,
+							"ProxyCommand %s --global-config %s ssh --stdio%s %s",
+							escapedCoderBinary, escapedGlobalConfig, flags, workspaceHostname,
 						))
 					}
 
-					var configOptions sshConfigOptions
+					// Create a copy of the options so we can modify them.
+					configOptions := sshConfigOpts
+					configOptions.sshOptions = nil
+
 					// Add standard options.
 					err := configOptions.addOptions(defaultOptions...)
 					if err != nil {
@@ -505,9 +527,16 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
 		},
 		{
 			Flag:        "ssh-host-prefix",
-			Env:         "",
+			Env:         "CODER_CONFIGSSH_SSH_HOST_PREFIX",
 			Description: "Override the default host prefix.",
-			Value:       clibase.StringOf(&userHostPrefix),
+			Value:       clibase.StringOf(&sshConfigOpts.userHostPrefix),
+		},
+		{
+			Flag:        "wait",
+			Env:         "CODER_CONFIGSSH_WAIT", // Not to be mixed with CODER_SSH_WAIT.
+			Description: "Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior configured in the workspace template is used.",
+			Default:     "auto",
+			Value:       clibase.EnumOf(&sshConfigOpts.waitEnum, "yes", "no", "auto"),
 		},
 		cliui.SkipPromptOption(),
 	}
@@ -524,12 +553,22 @@ func sshConfigWriteSectionHeader(w io.Writer, addNewline bool, o sshConfigOption
 	_, _ = fmt.Fprint(w, nl+sshStartToken+"\n")
 	_, _ = fmt.Fprint(w, sshConfigSectionHeader)
 	_, _ = fmt.Fprint(w, sshConfigDocsHeader)
-	if len(o.sshOptions) > 0 {
+
+	var ow strings.Builder
+	if o.waitEnum != "auto" {
+		_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "wait", o.waitEnum)
+	}
+	if o.userHostPrefix != "" {
+		_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-host-prefix", o.userHostPrefix)
+	}
+	for _, opt := range o.sshOptions {
+		_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-option", opt)
+	}
+	if ow.Len() > 0 {
 		_, _ = fmt.Fprint(w, sshConfigOptionsHeader)
-		for _, opt := range o.sshOptions {
-			_, _ = fmt.Fprintf(w, "# :%s=%s\n", "ssh-option", opt)
-		}
+		_, _ = fmt.Fprint(w, ow.String())
 	}
+
 	_, _ = fmt.Fprint(w, "#\n")
 }
 
@@ -538,13 +577,20 @@ func sshConfigWriteSectionEnd(w io.Writer) {
 }
 
 func sshConfigParseLastOptions(r io.Reader) (o sshConfigOptions) {
+	// Default values.
+	o.waitEnum = "auto"
+
 	s := bufio.NewScanner(r)
 	for s.Scan() {
 		line := s.Text()
 		if strings.HasPrefix(line, "# :") {
 			line = strings.TrimPrefix(line, "# :")
 			parts := strings.SplitN(line, "=", 2)
 			switch parts[0] {
+			case "wait":
+				o.waitEnum = parts[1]
+			case "ssh-host-prefix":
+				o.userHostPrefix = parts[1]
 			case "ssh-option":
 				o.sshOptions = append(o.sshOptions, parts[1])
 			default:
diff --git a/cli/configssh_test.go b/cli/configssh_test.go
@@ -481,12 +481,32 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
 			},
 			args: []string{"--yes"},
 		},
+		{
+			name: "Serialize supported flags",
+			wantConfig: wantConfig{
+				ssh: strings.Join([]string{
+					headerStart,
+					"# Last config-ssh options:",
+					"# :wait=yes",
+					"# :ssh-host-prefix=coder-test.",
+					"#",
+					headerEnd,
+					"",
+				}, "\n"),
+			},
+			args: []string{
+				"--yes",
+				"--wait=yes",
+				"--ssh-host-prefix", "coder-test.",
+			},
+		},
 		{
 			name: "Do not prompt for new options when prev opts flag is set",
 			writeConfig: writeConfig{
 				ssh: strings.Join([]string{
 					headerStart,
 					"# Last config-ssh options:",
+					"# :wait=no",
 					"# :ssh-option=ForwardAgent=yes",
 					"#",
 					headerEnd,
@@ -497,6 +517,7 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
 				ssh: strings.Join([]string{
 					headerStart,
 					"# Last config-ssh options:",
+					"# :wait=no",
 					"# :ssh-option=ForwardAgent=yes",
 					"#",
 					headerEnd,
@@ -589,8 +610,7 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
 			clitest.SetupConfig(t, client, root)
 
 			pty := ptytest.New(t)
-			inv.Stdin = pty.Input()
-			inv.Stdout = pty.Output()
+			pty.Attach(inv)
 			done := tGo(t, func() {
 				err := inv.Run()
 				if !tt.wantErr {
diff --git a/cli/testdata/coder_config-ssh_--help.golden b/cli/testdata/coder_config-ssh_--help.golden
@@ -18,7 +18,7 @@ Add an SSH Host entry for your workspaces "ssh coder.workspace"
       --ssh-config-file string, $CODER_SSH_CONFIG_FILE (default: ~/.ssh/config)
           Specifies the path to an SSH config.
 
-      --ssh-host-prefix string
+      --ssh-host-prefix string, $CODER_CONFIGSSH_SSH_HOST_PREFIX
           Override the default host prefix.
 
   -o, --ssh-option string-array, $CODER_SSH_CONFIG_OPTS
@@ -28,6 +28,11 @@ Add an SSH Host entry for your workspaces "ssh coder.workspace"
           Specifies whether or not to keep options from previous run of
           config-ssh.
 
+      --wait yes|no|auto, $CODER_CONFIGSSH_WAIT (default: auto)
+          Specifies whether or not to wait for the startup script to finish
+          executing. Auto means that the agent startup script behavior
+          configured in the workspace template is used.
+
   -y, --yes bool
           Bypass prompts.
 
diff --git a/docs/cli/config-ssh.md b/docs/cli/config-ssh.md
@@ -46,9 +46,10 @@ Specifies the path to an SSH config.
 
 ### --ssh-host-prefix
 
-|      |                     |
-| ---- | ------------------- |
-| Type | <code>string</code> |
+|             |                                               |
+| ----------- | --------------------------------------------- |
+| Type        | <code>string</code>                           |
+| Environment | <code>$CODER_CONFIGSSH_SSH_HOST_PREFIX</code> |
 
 Override the default host prefix.
 
@@ -70,6 +71,16 @@ Specifies additional SSH options to embed in each host stanza.
 
 Specifies whether or not to keep options from previous run of config-ssh.
 
+### --wait
+
+|             |                                    |
+| ----------- | ---------------------------------- | --- | ------------ |
+| Type        | <code>enum[yes                     | no  | auto]</code> |
+| Environment | <code>$CODER_CONFIGSSH_WAIT</code> |
+| Default     | <code>auto</code>                  |
+
+Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior configured in the workspace template is used.
+
 ### -y, --yes
 
 |      |                   |
