Skip to content

Commit a47b44c

Browse files
mafredrimafredri
committed
feat(cli): Add wait and no-wait support to config-ssh
Refs #7768
1 parent a77b48a commit a47b44c

File tree

2 files changed

+98
-16
lines changed

2 files changed

+98
-16
lines changed

‎cli/configssh.go

+76-14
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,10 @@ const (
4545
// sshConfigOptions represents options that can be stored and read
4646
// from the coder config in ~/.ssh/coder.
4747
type sshConfigOptions struct {
48-
sshOptions []string
48+
wait bool
49+
noWait bool
50+
userHostPrefix string
51+
sshOptions []string
4952
}
5053

5154
// addOptions expects options in the form of "option=value" or "option value".
@@ -100,10 +103,22 @@ func (o sshConfigOptions) equal(other sshConfigOptions) bool {
100103
sort.Strings(opt1)
101104
opt2 := slices.Clone(other.sshOptions)
102105
sort.Strings(opt2)
103-
return slices.Equal(opt1, opt2)
106+
if !slices.Equal(opt1, opt2) {
107+
return false
108+
}
109+
return o.wait == other.wait && o.noWait == other.noWait && o.userHostPrefix == other.userHostPrefix
104110
}
105111

106112
func (o sshConfigOptions) asList() (list []string) {
113+
if o.wait {
114+
list = append(list, "wait")
115+
}
116+
if o.noWait {
117+
list = append(list, "no-wait")
118+
}
119+
if o.userHostPrefix != "" {
120+
list = append(list, fmt.Sprintf("ssh-host-prefix: %s", o.userHostPrefix))
121+
}
107122
for _, opt := range o.sshOptions {
108123
list = append(list, fmt.Sprintf("ssh-option: %s", opt))
109124
}
@@ -178,14 +193,14 @@ func sshPrepareWorkspaceConfigs(ctx context.Context, client *codersdk.Client) (r
178193
}
179194
}
180195

196+
//nolint:gocyclo
181197
func (r *RootCmd) configSSH() *clibase.Cmd {
182198
var (
183199
sshConfigFile string
184200
sshConfigOpts sshConfigOptions
185201
usePreviousOpts bool
186202
dryRun bool
187203
skipProxyCommand bool
188-
userHostPrefix string
189204
)
190205
client := new(codersdk.Client)
191206
cmd := &clibase.Cmd{
@@ -207,6 +222,13 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
207222
r.InitClient(client),
208223
),
209224
Handler: func(inv *clibase.Invocation) error {
225+
if sshConfigOpts.wait && sshConfigOpts.noWait {
226+
return xerrors.Errorf("cannot specify both --wait and --no-wait")
227+
}
228+
if skipProxyCommand && (sshConfigOpts.wait || sshConfigOpts.noWait) {
229+
return xerrors.Errorf("cannot specify --skip-proxy-command with --wait or --no-wait")
230+
}
231+
210232
recvWorkspaceConfigs := sshPrepareWorkspaceConfigs(inv.Context(), client)
211233

212234
out := inv.Stdout
@@ -295,7 +317,7 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
295317
// Selecting "no" will use the last config.
296318
sshConfigOpts = *lastConfig
297319
} else {
298-
changes = append(changes, "Use new SSH options")
320+
changes = append(changes, "Use new options")
299321
}
300322
// Only print when prompts are shown.
301323
if yes, _ := inv.ParsedFlags().GetBool("yes"); !yes {
@@ -336,9 +358,9 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
336358
coderdConfig.HostnamePrefix = "coder."
337359
}
338360

339-
if userHostPrefix != "" {
361+
if sshConfigOpts.userHostPrefix != "" {
340362
// Override with user flag.
341-
coderdConfig.HostnamePrefix = userHostPrefix
363+
coderdConfig.HostnamePrefix = sshConfigOpts.userHostPrefix
342364
}
343365

344366
// Ensure stable sorting of output.
@@ -363,13 +385,22 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
363385
}
364386

365387
if !skipProxyCommand {
388+
flags := ""
389+
if sshConfigOpts.wait {
390+
flags += " --wait"
391+
} else if sshConfigOpts.noWait {
392+
flags += " --no-wait"
393+
}
366394
defaultOptions = append(defaultOptions, fmt.Sprintf(
367-
"ProxyCommand %s --global-config %s ssh --stdio %s",
368-
escapedCoderBinary, escapedGlobalConfig, workspaceHostname,
395+
"ProxyCommand %s --global-config %s ssh --stdio%s %s",
396+
escapedCoderBinary, escapedGlobalConfig, flags, workspaceHostname,
369397
))
370398
}
371399

372-
var configOptions sshConfigOptions
400+
// Create a copy of the options so we can modify them.
401+
configOptions := sshConfigOpts
402+
configOptions.sshOptions = nil
403+
373404
// Add standard options.
374405
err := configOptions.addOptions(defaultOptions...)
375406
if err != nil {
@@ -507,7 +538,19 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
507538
Flag: "ssh-host-prefix",
508539
Env: "",
509540
Description: "Override the default host prefix.",
510-
Value: clibase.StringOf(&userHostPrefix),
541+
Value: clibase.StringOf(&sshConfigOpts.userHostPrefix),
542+
},
543+
{
544+
Flag: "wait",
545+
Env: "CODER_CONFIGSSH_WAIT", // Not to be mixed with CODER_SSH_WAIT.
546+
Description: "Wait for the the startup script to finish executing. This is the default if the template has configured the agent startup script behavior as blocking. Can not be used together with --no-wait.",
547+
Value: clibase.BoolOf(&sshConfigOpts.wait),
548+
},
549+
{
550+
Flag: "no-wait",
551+
Env: "CODER_CONFIGSSH_NO_WAIT", // Not to be mixed with CODER_SSH_NO_WAIT.
552+
Description: "Enter workspace immediately after the agent has connected. This is the default if the template has configured the agent startup script behavior as non-blocking. Can not be used together with --wait.",
553+
Value: clibase.BoolOf(&sshConfigOpts.noWait),
511554
},
512555
cliui.SkipPromptOption(),
513556
}
@@ -524,12 +567,25 @@ func sshConfigWriteSectionHeader(w io.Writer, addNewline bool, o sshConfigOption
524567
_, _ = fmt.Fprint(w, nl+sshStartToken+"\n")
525568
_, _ = fmt.Fprint(w, sshConfigSectionHeader)
526569
_, _ = fmt.Fprint(w, sshConfigDocsHeader)
527-
if len(o.sshOptions) > 0 {
570+
571+
var ow strings.Builder
572+
if o.wait {
573+
_, _ = fmt.Fprintf(&ow, "# :%s\n", "wait")
574+
}
575+
if o.noWait {
576+
_, _ = fmt.Fprintf(&ow, "# :%s\n", "no-wait")
577+
}
578+
if o.userHostPrefix != "" {
579+
_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-host-prefix", o.userHostPrefix)
580+
}
581+
for _, opt := range o.sshOptions {
582+
_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-option", opt)
583+
}
584+
if ow.Len() > 0 {
528585
_, _ = fmt.Fprint(w, sshConfigOptionsHeader)
529-
for _, opt := range o.sshOptions {
530-
_, _ = fmt.Fprintf(w, "# :%s=%s\n", "ssh-option", opt)
531-
}
586+
_, _ = fmt.Fprint(w, ow.String())
532587
}
588+
533589
_, _ = fmt.Fprint(w, "#\n")
534590
}
535591

@@ -545,6 +601,12 @@ func sshConfigParseLastOptions(r io.Reader) (o sshConfigOptions) {
545601
line = strings.TrimPrefix(line, "# :")
546602
parts := strings.SplitN(line, "=", 2)
547603
switch parts[0] {
604+
case "wait":
605+
o.wait = true
606+
case "no-wait":
607+
o.noWait = true
608+
case "user-host-prefix":
609+
o.userHostPrefix = parts[1]
548610
case "ssh-option":
549611
o.sshOptions = append(o.sshOptions, parts[1])
550612
default:

‎cli/configssh_test.go

+22-2
Original file line numberDiff line numberDiff line change
@@ -481,12 +481,32 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
481481
},
482482
args: []string{"--yes"},
483483
},
484+
{
485+
name: "Serialize supported flags",
486+
wantConfig: wantConfig{
487+
ssh: strings.Join([]string{
488+
headerStart,
489+
"# Last config-ssh options:",
490+
"# :wait",
491+
"# :ssh-host-prefix=coder-test.",
492+
"#",
493+
headerEnd,
494+
"",
495+
}, "\n"),
496+
},
497+
args: []string{
498+
"--yes",
499+
"--wait",
500+
"--ssh-host-prefix", "coder-test.",
501+
},
502+
},
484503
{
485504
name: "Do not prompt for new options when prev opts flag is set",
486505
writeConfig: writeConfig{
487506
ssh: strings.Join([]string{
488507
headerStart,
489508
"# Last config-ssh options:",
509+
"# :no-wait",
490510
"# :ssh-option=ForwardAgent=yes",
491511
"#",
492512
headerEnd,
@@ -497,6 +517,7 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
497517
ssh: strings.Join([]string{
498518
headerStart,
499519
"# Last config-ssh options:",
520+
"# :no-wait",
500521
"# :ssh-option=ForwardAgent=yes",
501522
"#",
502523
headerEnd,
@@ -589,8 +610,7 @@ func TestConfigSSH_FileWriteAndOptionsFlow(t *testing.T) {
589610
clitest.SetupConfig(t, client, root)
590611

591612
pty := ptytest.New(t)
592-
inv.Stdin = pty.Input()
593-
inv.Stdout = pty.Output()
613+
pty.Attach(inv)
594614
done := tGo(t, func() {
595615
err := inv.Run()
596616
if !tt.wantErr {

0 commit comments

Comments
 (0)