Skip to content

Commit 6494ad9

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

File tree

2 files changed

+97
-16
lines changed

2 files changed

+97
-16
lines changed

‎cli/configssh.go

+75-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
}
@@ -185,7 +200,6 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
185200
usePreviousOpts bool
186201
dryRun bool
187202
skipProxyCommand bool
188-
userHostPrefix string
189203
)
190204
client := new(codersdk.Client)
191205
cmd := &clibase.Cmd{
@@ -207,6 +221,13 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
207221
r.InitClient(client),
208222
),
209223
Handler: func(inv *clibase.Invocation) error {
224+
if sshConfigOpts.wait && sshConfigOpts.noWait {
225+
return xerrors.Errorf("cannot specify both --wait and --no-wait")
226+
}
227+
if skipProxyCommand && (sshConfigOpts.wait || sshConfigOpts.noWait) {
228+
return xerrors.Errorf("cannot specify --skip-proxy-command with --wait or --no-wait")
229+
}
230+
210231
recvWorkspaceConfigs := sshPrepareWorkspaceConfigs(inv.Context(), client)
211232

212233
out := inv.Stdout
@@ -295,7 +316,7 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
295316
// Selecting "no" will use the last config.
296317
sshConfigOpts = *lastConfig
297318
} else {
298-
changes = append(changes, "Use new SSH options")
319+
changes = append(changes, "Use new options")
299320
}
300321
// Only print when prompts are shown.
301322
if yes, _ := inv.ParsedFlags().GetBool("yes"); !yes {
@@ -336,9 +357,9 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
336357
coderdConfig.HostnamePrefix = "coder."
337358
}
338359

339-
if userHostPrefix != "" {
360+
if sshConfigOpts.userHostPrefix != "" {
340361
// Override with user flag.
341-
coderdConfig.HostnamePrefix = userHostPrefix
362+
coderdConfig.HostnamePrefix = sshConfigOpts.userHostPrefix
342363
}
343364

344365
// Ensure stable sorting of output.
@@ -363,13 +384,22 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
363384
}
364385

365386
if !skipProxyCommand {
387+
flags := ""
388+
if sshConfigOpts.wait {
389+
flags += " --wait"
390+
} else if sshConfigOpts.noWait {
391+
flags += " --no-wait"
392+
}
366393
defaultOptions = append(defaultOptions, fmt.Sprintf(
367-
"ProxyCommand %s --global-config %s ssh --stdio %s",
368-
escapedCoderBinary, escapedGlobalConfig, workspaceHostname,
394+
"ProxyCommand %s --global-config %s ssh --stdio%s %s",
395+
escapedCoderBinary, escapedGlobalConfig, flags, workspaceHostname,
369396
))
370397
}
371398

372-
var configOptions sshConfigOptions
399+
// Create a copy of the options so we can modify them.
400+
configOptions := sshConfigOpts
401+
configOptions.sshOptions = nil
402+
373403
// Add standard options.
374404
err := configOptions.addOptions(defaultOptions...)
375405
if err != nil {
@@ -507,7 +537,19 @@ func (r *RootCmd) configSSH() *clibase.Cmd {
507537
Flag: "ssh-host-prefix",
508538
Env: "",
509539
Description: "Override the default host prefix.",
510-
Value: clibase.StringOf(&userHostPrefix),
540+
Value: clibase.StringOf(&sshConfigOpts.userHostPrefix),
541+
},
542+
{
543+
Flag: "wait",
544+
Env: "CODER_CONFIGSSH_WAIT", // Not to be mixed with CODER_SSH_WAIT.
545+
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.",
546+
Value: clibase.BoolOf(&sshConfigOpts.wait),
547+
},
548+
{
549+
Flag: "no-wait",
550+
Env: "CODER_CONFIGSSH_NO_WAIT", // Not to be mixed with CODER_SSH_NO_WAIT.
551+
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.",
552+
Value: clibase.BoolOf(&sshConfigOpts.noWait),
511553
},
512554
cliui.SkipPromptOption(),
513555
}
@@ -524,12 +566,25 @@ func sshConfigWriteSectionHeader(w io.Writer, addNewline bool, o sshConfigOption
524566
_, _ = fmt.Fprint(w, nl+sshStartToken+"\n")
525567
_, _ = fmt.Fprint(w, sshConfigSectionHeader)
526568
_, _ = fmt.Fprint(w, sshConfigDocsHeader)
527-
if len(o.sshOptions) > 0 {
569+
570+
var ow strings.Builder
571+
if o.wait {
572+
_, _ = fmt.Fprintf(&ow, "# :%s\n", "wait")
573+
}
574+
if o.noWait {
575+
_, _ = fmt.Fprintf(&ow, "# :%s\n", "no-wait")
576+
}
577+
if o.userHostPrefix != "" {
578+
_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-host-prefix", o.userHostPrefix)
579+
}
580+
for _, opt := range o.sshOptions {
581+
_, _ = fmt.Fprintf(&ow, "# :%s=%s\n", "ssh-option", opt)
582+
}
583+
if ow.Len() > 0 {
528584
_, _ = fmt.Fprint(w, sshConfigOptionsHeader)
529-
for _, opt := range o.sshOptions {
530-
_, _ = fmt.Fprintf(w, "# :%s=%s\n", "ssh-option", opt)
531-
}
585+
_, _ = fmt.Fprint(w, ow.String())
532586
}
587+
533588
_, _ = fmt.Fprint(w, "#\n")
534589
}
535590

@@ -545,6 +600,12 @@ func sshConfigParseLastOptions(r io.Reader) (o sshConfigOptions) {
545600
line = strings.TrimPrefix(line, "# :")
546601
parts := strings.SplitN(line, "=", 2)
547602
switch parts[0] {
603+
case "wait":
604+
o.wait = true
605+
case "no-wait":
606+
o.noWait = true
607+
case "user-host-prefix":
608+
o.userHostPrefix = parts[1]
548609
case "ssh-option":
549610
o.sshOptions = append(o.sshOptions, parts[1])
550611
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)