@@ -61,38 +61,48 @@ func TestMain(m *testing.M) {
6161goleak .VerifyTestMain (m ,testutil .GoleakOptions ... )
6262}
6363
64+ var sshPorts = []uint16 {workspacesdk .AgentSSHPort ,workspacesdk .AgentStandardSSHPort }
65+
6466// NOTE: These tests only work when your default shell is bash for some reason.
6567
6668func TestAgent_Stats_SSH (t * testing.T ) {
6769t .Parallel ()
68- ctx ,cancel := context .WithTimeout (context .Background (),testutil .WaitLong )
69- defer cancel ()
7070
71- //nolint:dogsled
72- conn ,_ ,stats ,_ ,_ := setupAgent (t , agentsdk.Manifest {},0 )
71+ for _ ,port := range sshPorts {
72+ port := port
73+ t .Run (fmt .Sprintf ("(:%d)" ,port ),func (t * testing.T ) {
74+ t .Parallel ()
7375
74- sshClient ,err := conn .SSHClient (ctx )
75- require .NoError (t ,err )
76- defer sshClient .Close ()
77- session ,err := sshClient .NewSession ()
78- require .NoError (t ,err )
79- defer session .Close ()
80- stdin ,err := session .StdinPipe ()
81- require .NoError (t ,err )
82- err = session .Shell ()
83- require .NoError (t ,err )
76+ ctx ,cancel := context .WithTimeout (context .Background (),testutil .WaitLong )
77+ defer cancel ()
8478
85- var s * proto.Stats
86- require .Eventuallyf (t ,func ()bool {
87- var ok bool
88- s ,ok = <- stats
89- return ok && s .ConnectionCount > 0 && s .RxBytes > 0 && s .TxBytes > 0 && s .SessionCountSsh == 1
90- },testutil .WaitLong ,testutil .IntervalFast ,
91- "never saw stats: %+v" ,s ,
92- )
93- _ = stdin .Close ()
94- err = session .Wait ()
95- require .NoError (t ,err )
79+ //nolint:dogsled
80+ conn ,_ ,stats ,_ ,_ := setupAgent (t , agentsdk.Manifest {},0 )
81+
82+ sshClient ,err := conn .SSHClientOnPort (ctx ,port )
83+ require .NoError (t ,err )
84+ defer sshClient .Close ()
85+ session ,err := sshClient .NewSession ()
86+ require .NoError (t ,err )
87+ defer session .Close ()
88+ stdin ,err := session .StdinPipe ()
89+ require .NoError (t ,err )
90+ err = session .Shell ()
91+ require .NoError (t ,err )
92+
93+ var s * proto.Stats
94+ require .Eventuallyf (t ,func ()bool {
95+ var ok bool
96+ s ,ok = <- stats
97+ return ok && s .ConnectionCount > 0 && s .RxBytes > 0 && s .TxBytes > 0 && s .SessionCountSsh == 1
98+ },testutil .WaitLong ,testutil .IntervalFast ,
99+ "never saw stats: %+v" ,s ,
100+ )
101+ _ = stdin .Close ()
102+ err = session .Wait ()
103+ require .NoError (t ,err )
104+ })
105+ }
96106}
97107
98108func TestAgent_Stats_ReconnectingPTY (t * testing.T ) {
@@ -266,15 +276,23 @@ func TestAgent_Stats_Magic(t *testing.T) {
266276
267277func TestAgent_SessionExec (t * testing.T ) {
268278t .Parallel ()
269- session := setupSSHSession (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {},nil )
270279
271- command := "echo test"
272- if runtime .GOOS == "windows" {
273- command = "cmd.exe /c echo test"
280+ for _ ,port := range sshPorts {
281+ port := port
282+ t .Run (fmt .Sprintf ("(:%d)" ,port ),func (t * testing.T ) {
283+ t .Parallel ()
284+
285+ session := setupSSHSessionOnPort (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {},nil ,port )
286+
287+ command := "echo test"
288+ if runtime .GOOS == "windows" {
289+ command = "cmd.exe /c echo test"
290+ }
291+ output ,err := session .Output (command )
292+ require .NoError (t ,err )
293+ require .Equal (t ,"test" ,strings .TrimSpace (string (output )))
294+ })
274295}
275- output ,err := session .Output (command )
276- require .NoError (t ,err )
277- require .Equal (t ,"test" ,strings .TrimSpace (string (output )))
278296}
279297
280298//nolint:tparallel // Sub tests need to run sequentially.
@@ -384,25 +402,33 @@ func TestAgent_SessionTTYShell(t *testing.T) {
384402// it seems like it could be either.
385403t .Skip ("ConPTY appears to be inconsistent on Windows." )
386404}
387- session := setupSSHSession (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {},nil )
388- command := "sh"
389- if runtime .GOOS == "windows" {
390- command = "cmd.exe"
405+
406+ for _ ,port := range sshPorts {
407+ port := port
408+ t .Run (fmt .Sprintf ("(%d)" ,port ),func (t * testing.T ) {
409+ t .Parallel ()
410+
411+ session := setupSSHSessionOnPort (t , agentsdk.Manifest {}, codersdk.ServiceBannerConfig {},nil ,port )
412+ command := "sh"
413+ if runtime .GOOS == "windows" {
414+ command = "cmd.exe"
415+ }
416+ err := session .RequestPty ("xterm" ,128 ,128 , ssh.TerminalModes {})
417+ require .NoError (t ,err )
418+ ptty := ptytest .New (t )
419+ session .Stdout = ptty .Output ()
420+ session .Stderr = ptty .Output ()
421+ session .Stdin = ptty .Input ()
422+ err = session .Start (command )
423+ require .NoError (t ,err )
424+ _ = ptty .Peek (ctx ,1 )// wait for the prompt
425+ ptty .WriteLine ("echo test" )
426+ ptty .ExpectMatch ("test" )
427+ ptty .WriteLine ("exit" )
428+ err = session .Wait ()
429+ require .NoError (t ,err )
430+ })
391431}
392- err := session .RequestPty ("xterm" ,128 ,128 , ssh.TerminalModes {})
393- require .NoError (t ,err )
394- ptty := ptytest .New (t )
395- session .Stdout = ptty .Output ()
396- session .Stderr = ptty .Output ()
397- session .Stdin = ptty .Input ()
398- err = session .Start (command )
399- require .NoError (t ,err )
400- _ = ptty .Peek (ctx ,1 )// wait for the prompt
401- ptty .WriteLine ("echo test" )
402- ptty .ExpectMatch ("test" )
403- ptty .WriteLine ("exit" )
404- err = session .Wait ()
405- require .NoError (t ,err )
406432}
407433
408434func TestAgent_SessionTTYExitCode (t * testing.T ) {
@@ -596,37 +622,41 @@ func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
596622//nolint:dogsled // Allow the blank identifiers.
597623conn ,client ,_ ,_ ,_ := setupAgent (t , agentsdk.Manifest {},0 ,setSBInterval )
598624
599- sshClient ,err := conn .SSHClient (ctx )
600- require .NoError (t ,err )
601- t .Cleanup (func () {
602- _ = sshClient .Close ()
603- })
604-
605625//nolint:paralleltest // These tests need to swap the banner func.
606- for i ,test := range tests {
607- test := test
608- t .Run (fmt .Sprintf ("%d" ,i ),func (t * testing.T ) {
609- // Set new banner func and wait for the agent to call it to update the
610- // banner.
611- ready := make (chan struct {},2 )
612- client .SetAnnouncementBannersFunc (func () ([]codersdk.BannerConfig ,error ) {
613- select {
614- case ready <- struct {}{}:
615- default :
616- }
617- return []codersdk.BannerConfig {test .banner },nil
618- })
619- <- ready
620- <- ready // Wait for two updates to ensure the value has propagated.
621-
622- session ,err := sshClient .NewSession ()
623- require .NoError (t ,err )
624- t .Cleanup (func () {
625- _ = session .Close ()
626- })
626+ for _ ,port := range sshPorts {
627+ port := port
627628
628- testSessionOutput (t ,session ,test .expected ,test .unexpected ,nil )
629+ sshClient ,err := conn .SSHClientOnPort (ctx ,port )
630+ require .NoError (t ,err )
631+ t .Cleanup (func () {
632+ _ = sshClient .Close ()
629633})
634+
635+ for i ,test := range tests {
636+ test := test
637+ t .Run (fmt .Sprintf ("(:%d)/%d" ,port ,i ),func (t * testing.T ) {
638+ // Set new banner func and wait for the agent to call it to update the
639+ // banner.
640+ ready := make (chan struct {},2 )
641+ client .SetAnnouncementBannersFunc (func () ([]codersdk.BannerConfig ,error ) {
642+ select {
643+ case ready <- struct {}{}:
644+ default :
645+ }
646+ return []codersdk.BannerConfig {test .banner },nil
647+ })
648+ <- ready
649+ <- ready // Wait for two updates to ensure the value has propagated.
650+
651+ session ,err := sshClient .NewSession ()
652+ require .NoError (t ,err )
653+ t .Cleanup (func () {
654+ _ = session .Close ()
655+ })
656+
657+ testSessionOutput (t ,session ,test .expected ,test .unexpected ,nil )
658+ })
659+ }
630660}
631661}
632662
@@ -2313,6 +2343,17 @@ func setupSSHSession(
23132343banner codersdk.BannerConfig ,
23142344prepareFS func (fs afero.Fs ),
23152345opts ... func (* agenttest.Client ,* agent.Options ),
2346+ )* ssh.Session {
2347+ return setupSSHSessionOnPort (t ,manifest ,banner ,prepareFS ,workspacesdk .AgentSSHPort ,opts ... )
2348+ }
2349+
2350+ func setupSSHSessionOnPort (
2351+ t * testing.T ,
2352+ manifest agentsdk.Manifest ,
2353+ banner codersdk.BannerConfig ,
2354+ prepareFS func (fs afero.Fs ),
2355+ port uint16 ,
2356+ opts ... func (* agenttest.Client ,* agent.Options ),
23162357)* ssh.Session {
23172358ctx ,cancel := context .WithTimeout (context .Background (),testutil .WaitLong )
23182359defer cancel ()
@@ -2326,7 +2367,7 @@ func setupSSHSession(
23262367if prepareFS != nil {
23272368prepareFS (fs )
23282369}
2329- sshClient ,err := conn .SSHClient (ctx )
2370+ sshClient ,err := conn .SSHClientOnPort (ctx , port )
23302371require .NoError (t ,err )
23312372t .Cleanup (func () {
23322373_ = sshClient .Close ()