Fix pty tests on Windows

kylecarbs · kylecarbs · commit 11e5274b7ba5 · 2022-02-17T23:31:28.000-06:00
diff --git a/agent/agent.go b/agent/agent.go
@@ -9,6 +9,8 @@ import (
 	"io"
 	"net"
 	"os/exec"
+	"os/user"
+	"sync"
 	"time"
 
 	"cdr.dev/slog"
@@ -39,7 +41,6 @@ func DialSSHClient(conn *peer.Conn) (*gossh.Client, error) {
 		return nil, err
 	}
 	sshConn, channels, requests, err := gossh.NewClientConn(netConn, "localhost:22", &gossh.ClientConfig{
-		User: "kyle",
 		Config: gossh.Config{
 			Ciphers: []string{"arcfour"},
 		},
@@ -66,6 +67,7 @@ func New(dialer Dialer, options *Options) io.Closer {
 		clientDialer: dialer,
 		options:      options,
 		closeCancel:  cancelFunc,
+		closed:       make(chan struct{}),
 	}
 	server.init(ctx)
 	return server
@@ -76,6 +78,7 @@ type server struct {
 	options      *Options
 
 	closeCancel context.CancelFunc
+	closeMutex  sync.Mutex
 	closed      chan struct{}
 
 	sshServer *ssh.Server
@@ -153,10 +156,19 @@ func (*server) handleSSHSession(session ssh.Session) error {
 		err     error
 	)
 
+	username := session.User()
+	if username == "" {
+		currentUser, err := user.Current()
+		if err != nil {
+			return xerrors.Errorf("get current user: %w", err)
+		}
+		username = currentUser.Username
+	}
+
 	// gliderlabs/ssh returns a command slice of zero
 	// when a shell is requested.
 	if len(session.Command()) == 0 {
-		command, err = usershell.Get(session.User())
+		command, err = usershell.Get(username)
 		if err != nil {
 			return xerrors.Errorf("get user shell: %w", err)
 		}
@@ -208,6 +220,7 @@ func (*server) handleSSHSession(session ssh.Session) error {
 			_, _ = io.Copy(session, ptty.Output())
 		}()
 		_, _ = process.Wait()
+		_ = ptty.Close()
 		return nil
 	}
 
@@ -254,7 +267,11 @@ func (s *server) run(ctx context.Context) {
 	for {
 		conn, err := peerListener.Accept()
 		if err != nil {
-			// This is closed!
+			if s.isClosed() {
+				return
+			}
+			s.options.Logger.Debug(ctx, "peer listener accept exited; restarting connection", slog.Error(err))
+			s.run(ctx)
 			return
 		}
 		go s.handlePeerConn(ctx, conn)
@@ -265,15 +282,16 @@ func (s *server) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 	for {
 		channel, err := conn.Accept(ctx)
 		if err != nil {
-			// TODO: Log here!
+			if s.isClosed() {
+				return
+			}
+			s.options.Logger.Debug(ctx, "accept channel from peer connection", slog.Error(err))
 			return
 		}
 
 		switch channel.Protocol() {
 		case "ssh":
 			s.sshServer.HandleConn(channel.NetConn())
-		case "proxy":
-			// Proxy the port provided.
 		}
 	}
 }
@@ -289,6 +307,13 @@ func (s *server) isClosed() bool {
 }
 
 func (s *server) Close() error {
-	s.sshServer.Close()
+	s.closeMutex.Lock()
+	defer s.closeMutex.Unlock()
+	if s.isClosed() {
+		return nil
+	}
+	close(s.closed)
+	s.closeCancel()
+	_ = s.sshServer.Close()
 	return nil
 }
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -11,6 +11,7 @@ import (
 	"go.uber.org/goleak"
 	"golang.org/x/crypto/ssh"
 
+	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/peer"
@@ -35,7 +36,9 @@ func TestAgent(t *testing.T) {
 			Logger: slogtest.Make(t, nil),
 		})
 		require.NoError(t, err)
-		defer conn.Close()
+		t.Cleanup(func() {
+			_ = conn.Close()
+		})
 		sshClient, err := agent.DialSSHClient(conn)
 		require.NoError(t, err)
 		session, err := sshClient.NewSession()
@@ -58,7 +61,9 @@ func TestAgent(t *testing.T) {
 			Logger: slogtest.Make(t, nil),
 		})
 		require.NoError(t, err)
-		defer conn.Close()
+		t.Cleanup(func() {
+			_ = conn.Close()
+		})
 		sshClient, err := agent.DialSSHClient(conn)
 		require.NoError(t, err)
 		session, err := sshClient.NewSession()
@@ -94,7 +99,7 @@ func setup(t *testing.T) proto.DRPCPeerBrokerClient {
 			Logger: slogtest.Make(t, nil),
 		})
 	}, &agent.Options{
-		Logger: slogtest.Make(t, nil),
+		Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
 	})
 	t.Cleanup(func() {
 		_ = client.Close()
diff --git a/agent/usershell/usershell_darwin.go b/agent/usershell/usershell_darwin.go
@@ -0,0 +1,8 @@
+package usershell
+
+// Get returns the $SHELL environment variable.
+// TODO: This should use "dscl" to fetch the proper value. See:
+// https://stackoverflow.com/questions/16375519/how-to-get-the-default-shell
+func Get(username string) (string, error) {
+	return os.Getenv("SHELL"), nil
+}
diff --git a/agent/usershell/usershell_other.go b/agent/usershell/usershell_other.go
@@ -1,5 +1,5 @@
-//go:build !windows
-// +build !windows
+//go:build !windows && !darwin
+// +build !windows,!darwin
 
 package usershell
 
diff --git a/agent/usershell/usershell_other_test.go b/agent/usershell/usershell_other_test.go
@@ -1,5 +1,5 @@
-//go:build !windows
-// +build !windows
+//go:build !windows && !darwin
+// +build !windows,!darwin
 
 package usershell_test
 
diff --git a/go.mod b/go.mod
@@ -41,6 +41,7 @@ require (
 	github.com/pion/logging v0.2.2
 	github.com/pion/transport v0.13.0
 	github.com/pion/webrtc/v3 v3.1.23
+	github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8
 	github.com/psanford/memfs v0.0.0-20210214183328-a001468d78ef
 	github.com/quasilyte/go-ruleguard/dsl v0.3.17
 	github.com/spf13/cobra v1.3.0
@@ -115,7 +116,6 @@ require (
 	github.com/pion/stun v0.3.5 // indirect
 	github.com/pion/turn/v2 v2.0.6 // indirect
 	github.com/pion/udp v0.1.1 // indirect
-	github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/sirupsen/logrus v1.8.1 // indirect
diff --git a/pty/start_other.go b/pty/start_other.go
@@ -32,9 +32,5 @@ func startPty(cmd *exec.Cmd) (PTY, *os.Process, error) {
 		pty: ptty,
 		tty: tty,
 	}
-	go func() {
-		_ = cmd.Wait()
-		_ = oPty.Close()
-	}()
 	return oPty, cmd.Process, nil
 }
diff --git a/pty/start_windows.go b/pty/start_windows.go
@@ -83,10 +83,6 @@ func startPty(cmd *exec.Cmd) (PTY, *os.Process, error) {
 	if err != nil {
 		return nil, nil, xerrors.Errorf("find process %d: %w", processInfo.ProcessId, err)
 	}
-	go func() {
-		_, _ = process.Wait()
-		_ = pty.Close()
-	}()
 	return pty, process, nil
 }
 

Original file line number	Diff line number	Diff line change
`@@ -32,9 +32,5 @@ func startPty(cmd exec.Cmd) (PTY, os.Process, error) {`
`32`	`32`	`pty: ptty,`
`33`	`33`	`tty: tty,`
`34`	`34`	`}`
`35`		`- go func() {`
`36`		`- _ = cmd.Wait()`
`37`		`- _ = oPty.Close()`
`38`		`- }()`
`39`	`35`	`return oPty, cmd.Process, nil`
`40`	`36`	`}`
Original file line number	Diff line number	Diff line change
`@@ -83,10 +83,6 @@ func startPty(cmd exec.Cmd) (PTY, os.Process, error) {`
`83`	`83`	`if err != nil {`
`84`	`84`	`return nil, nil, xerrors.Errorf("find process %d: %w", processInfo.ProcessId, err)`
`85`	`85`	`}`
`86`		`- go func() {`
`87`		`- _, _ = process.Wait()`
`88`		`- _ = pty.Close()`
`89`		`- }()`
`90`	`86`	`return pty, process, nil`
`91`	`87`	`}`
`92`	`88`