Mar 24, 2023 · Mar 24, 2023 · Mar 24, 2023 · Mar 24, 2023
diff --git a/agent/agent.go b/agent/agent.go
 _=session.Exit(MagicSessionErrorCode)
 return
 }
 _=session.Exit(0)
 },
 HostSigners: []ssh.Signer{randomSigner},
 LocalPortForwardingCallback:func(ctx ssh.Context,destinationHoststring,destinationPortuint32)bool {
 iferr!=nil {
 returnxerrors.Errorf("start command: %w",err)
 }
 varwg sync.WaitGroup
 deferfunc() {
 deferwg.Wait()
 closeErr:=ptty.Close()
 ifcloseErr!=nil {
 a.logger.Warn(ctx,"failed to close tty",slog.Error(closeErr))
 }
 }
 }()
 // We don't add input copy to wait group because
 // it won't return until the session is closed.
 gofunc() {
 _,_=io.Copy(ptty.Input(),session)
 }()
 wg.Add(1)
 gofunc() {
 // Ensure data is flushed to session on command exit, if we
 // close the session too soon, we might lose data.
 deferwg.Done()
 _,_=io.Copy(session,ptty.Output())
 }()
 err=process.Wait()
diff --git a/agent/agent_test.go b/agent/agent_test.go
 require.NotContains(t,stdout.String(),wantNotMOTD,"should not show motd")
 }

 funcTestAgent_Session_TTY_FastCommandHasOutput(t*testing.T) {
 t.Parallel()
 ifruntime.GOOS=="windows" {
 // This might be our implementation, or ConPTY itself.
 // It's difficult to find extensive tests for it, so
 // it seems like it could be either.
 t.Skip("ConPTY appears to be inconsistent on Windows.")
 }

 // This test is here to prevent regressions where quickly executing
 // commands (with TTY) don't flush their output to the SSH session.
 //
 // See: https://github.com/coder/coder/issues/6656
 ctx,cancel:=context.WithTimeout(context.Background(),testutil.WaitLong)
 defercancel()
 //nolint:dogsled
 conn,_,_,_,_:=setupAgent(t, agentsdk.Metadata{},0)
 sshClient,err:=conn.SSHClient(ctx)
 require.NoError(t,err)
 defersshClient.Close()

 ptty:=ptytest.New(t)

 varstdout bytes.Buffer
 // NOTE(mafredri): Increase iterations to increase chance of failure,
 //                 assuming bug is present.
 // Using 1000 iterations is basically a guaranteed failure (but let's
 // not increase test times needlessly).
 fori:=0;i<5;i++ {
 func() {
 stdout.Reset()

 session,err:=sshClient.NewSession()
 require.NoError(t,err)
 defersession.Close()
 err=session.RequestPty("xterm",128,128, ssh.TerminalModes{})
 require.NoError(t,err)

 session.Stdout=&stdout
 session.Stderr=ptty.Output()
 session.Stdin=ptty.Input()
 err=session.Start("echo wazzup")
 require.NoError(t,err)

 err=session.Wait()
 require.NoError(t,err)
 require.Contains(t,stdout.String(),"wazzup","should output greeting")
 }()
 }
 }

 //nolint:paralleltest // This test reserves a port.
 funcTestAgent_TCPLocalForwarding(t*testing.T) {
 random,err:=net.Listen("tcp","127.0.0.1:0")
Original file line number	Diff line number	Diff line change
Expand Up		@@ -844,6 +844,7 @@ func (a *agent) init(ctx context.Context) {
		_=session.Exit(MagicSessionErrorCode)
		return
		}
		_=session.Exit(0)
		},
		HostSigners: []ssh.Signer{randomSigner},
		LocalPortForwardingCallback:func(ctx ssh.Context,destinationHoststring,destinationPortuint32)bool {
Expand DownExpand Up		@@ -1100,7 +1101,9 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
		iferr!=nil {
		returnxerrors.Errorf("start command: %w",err)
		}
		varwg sync.WaitGroup
		deferfunc() {
		deferwg.Wait()
		closeErr:=ptty.Close()
		ifcloseErr!=nil {
		a.logger.Warn(ctx,"failed to close tty",slog.Error(closeErr))
Expand All		@@ -1117,10 +1120,16 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
		}
		}
		}()
		// We don't add input copy to wait group because
		// it won't return until the session is closed.
		gofunc() {
		_,_=io.Copy(ptty.Input(),session)
		}()
		wg.Add(1)
		gofunc() {
		// Ensure data is flushed to session on command exit, if we
		// close the session too soon, we might lose data.
		deferwg.Done()
		_,_=io.Copy(session,ptty.Output())
		}()
		err=process.Wait()
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -348,6 +348,57 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
		require.NotContains(t,stdout.String(),wantNotMOTD,"should not show motd")
		}

		funcTestAgent_Session_TTY_FastCommandHasOutput(t*testing.T) {
		t.Parallel()
		ifruntime.GOOS=="windows" {
		// This might be our implementation, or ConPTY itself.
		// It's difficult to find extensive tests for it, so
		// it seems like it could be either.
		t.Skip("ConPTY appears to be inconsistent on Windows.")
		}

		// This test is here to prevent regressions where quickly executing
		// commands (with TTY) don't flush their output to the SSH session.
		//
		// See: https://github.com/coder/coder/issues/6656
		ctx,cancel:=context.WithTimeout(context.Background(),testutil.WaitLong)
		defercancel()
		//nolint:dogsled
		conn,_,_,_,_:=setupAgent(t, agentsdk.Metadata{},0)
		sshClient,err:=conn.SSHClient(ctx)
		require.NoError(t,err)
		defersshClient.Close()

		ptty:=ptytest.New(t)

		varstdout bytes.Buffer
		// NOTE(mafredri): Increase iterations to increase chance of failure,
		// assuming bug is present.
		// Using 1000 iterations is basically a guaranteed failure (but let's
		// not increase test times needlessly).
Comment on lines +375 to +378 Copy link Member kylecarbsMar 24, 2023 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. I really enjoy comments like this. I should do more of these! mafredri reacted with heart emoji
		fori:=0;i<5;i++ {
		func() {
		stdout.Reset()

		session,err:=sshClient.NewSession()
		require.NoError(t,err)
		defersession.Close()
		err=session.RequestPty("xterm",128,128, ssh.TerminalModes{})
		require.NoError(t,err)

		session.Stdout=&stdout
		session.Stderr=ptty.Output()
		session.Stdin=ptty.Input()
		err=session.Start("echo wazzup")
		require.NoError(t,err)

		err=session.Wait()
		require.NoError(t,err)
		require.Contains(t,stdout.String(),"wazzup","should output greeting")
		}()
		}
		}

		//nolint:paralleltest // This test reserves a port.
		funcTestAgent_TCPLocalForwarding(t*testing.T) {
		random,err:=net.Listen("tcp","127.0.0.1:0")
Expand Down