Refactor pty package to support Windows spawn

coder · kylecarbs · Feb 17, 2022 · Feb 15, 2022 · Feb 16, 2022 · Feb 16, 2022
commit 722be6c98288db4102e7f617c020c13f2edb0ce8
diff --git a/agent/server.go b/agent/server.go
@@ -7,13 +7,12 @@ import (
 	"errors"
 	"io"
 	"net"
-	"os"
+	"os/exec"
 	"sync"
-	"syscall"
 	"time"
 
 	"cdr.dev/slog"
-	"github.com/ActiveState/termtest/conpty"
+	"github.com/coder/coder/console/pty"
 	"github.com/coder/coder/peer"
 	"github.com/coder/coder/peerbroker"
 	"github.com/coder/retry"
@@ -71,31 +70,24 @@ func (s *server) init(ctx context.Context) {
 			sshLogger.Info(ctx, "ssh connection ended", slog.Error(err))
 		},
 		Handler: func(session ssh.Session) {
-			sshPty, windowSize, isPty := session.Pty()
+			_, windowSize, isPty := session.Pty()
 			if isPty {
-				cpty, err := conpty.New(int16(sshPty.Window.Width), int16(sshPty.Window.Height))
-				if err != nil {
-					panic(err)
-				}
-				_, _, err = cpty.Spawn("C:\\WINDOWS\\System32\\WindowsPowerShell\\v1.0\\powershell.exe", []string{}, &syscall.ProcAttr{
-					Env: os.Environ(),
-				})
+				pty, err := pty.Start(exec.Command("powershell.exe"))
 				if err != nil {
 					panic(err)
 				}
 				go func() {
 					for win := range windowSize {
-						err := cpty.Resize(uint16(win.Width), uint16(win.Height))
+						err := pty.Resize(uint16(win.Width), uint16(win.Height))
 						if err != nil {
 							panic(err)
 						}
 					}
 				}()
-
 				go func() {
-					io.Copy(session, cpty)
+					io.Copy(session, pty.Output())
 				}()
-				io.Copy(cpty, session)
+				io.Copy(pty.Input(), session)
 			}
 		},
 		HostSigners: []ssh.Signer{randomSigner},

diff --git a/cli/login_test.go b/cli/login_test.go
@@ -5,7 +5,7 @@ import (
 
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
-	"github.com/coder/coder/console"
+	"github.com/coder/coder/pty/ptytest"
 	"github.com/stretchr/testify/require"
 )
 
@@ -26,7 +26,9 @@ func TestLogin(t *testing.T) {
 		// accurately detect Windows ptys when they are not attached to a process:
 		// https://github.com/mattn/go-isatty/issues/59
 		root, _ := clitest.New(t, "login", client.URL.String(), "--force-tty")
-		cons := console.New(t, root)
+		pty := ptytest.New(t)
+		root.SetIn(pty.Input())
+		root.SetOut(pty.Output())
 		go func() {
 			err := root.Execute()
 			require.NoError(t, err)
@@ -42,12 +44,9 @@ func TestLogin(t *testing.T) {
 		for i := 0; i < len(matches); i += 2 {
 			match := matches[i]
 			value := matches[i+1]
-			_, err := cons.ExpectString(match)
-			require.NoError(t, err)
-			_, err = cons.SendLine(value)
-			require.NoError(t, err)
+			pty.ExpectMatch(match)
+			pty.WriteLine(value)
 		}
-		_, err := cons.ExpectString("Welcome to Coder")
-		require.NoError(t, err)
+		pty.ExpectMatch("Welcome to Coder")
 	})
 }
diff --git a/cli/projectcreate_test.go b/cli/projectcreate_test.go
@@ -7,10 +7,10 @@ import (
 
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
-	"github.com/coder/coder/console"
 	"github.com/coder/coder/database"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
+	"github.com/coder/coder/pty/ptytest"
 )
 
 func TestProjectCreate(t *testing.T) {
@@ -26,7 +26,9 @@ func TestProjectCreate(t *testing.T) {
 		cmd, root := clitest.New(t, "projects", "create", "--directory", source, "--provisioner", string(database.ProvisionerTypeEcho))
 		clitest.SetupConfig(t, client, root)
 		_ = coderdtest.NewProvisionerDaemon(t, client)
-		console := console.New(t, cmd)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
 		closeChan := make(chan struct{})
 		go func() {
 			err := cmd.Execute()
@@ -43,10 +45,8 @@ func TestProjectCreate(t *testing.T) {
 		for i := 0; i < len(matches); i += 2 {
 			match := matches[i]
 			value := matches[i+1]
-			_, err := console.ExpectString(match)
-			require.NoError(t, err)
-			_, err = console.SendLine(value)
-			require.NoError(t, err)
+			pty.ExpectMatch(match)
+			pty.WriteLine(value)
 		}
 		<-closeChan
 	})
@@ -73,7 +73,9 @@ func TestProjectCreate(t *testing.T) {
 		cmd, root := clitest.New(t, "projects", "create", "--directory", source, "--provisioner", string(database.ProvisionerTypeEcho))
 		clitest.SetupConfig(t, client, root)
 		coderdtest.NewProvisionerDaemon(t, client)
-		cons := console.New(t, cmd)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
 		closeChan := make(chan struct{})
 		go func() {
 			err := cmd.Execute()
@@ -91,10 +93,8 @@ func TestProjectCreate(t *testing.T) {
 		for i := 0; i < len(matches); i += 2 {
 			match := matches[i]
 			value := matches[i+1]
-			_, err := cons.ExpectString(match)
-			require.NoError(t, err)
-			_, err = cons.SendLine(value)
-			require.NoError(t, err)
+			pty.ExpectMatch(match)
+			pty.WriteLine(value)
 		}
 		<-closeChan
 	})

diff --git a/cli/root.go b/cli/root.go
@@ -139,14 +139,21 @@ func isTTY(cmd *cobra.Command) bool {
 
 func prompt(cmd *cobra.Command, prompt *promptui.Prompt) (string, error) {
 	var ok bool
-	prompt.Stdin, ok = cmd.InOrStdin().(io.ReadCloser)
+	reader, ok := cmd.InOrStdin().(io.Reader)
 	if !ok {
 		return "", xerrors.New("stdin must be a readcloser")
 	}
-	prompt.Stdout, ok = cmd.OutOrStdout().(io.WriteCloser)
+	prompt.Stdin = readWriteCloser{
+		Reader: reader,
+	}
+
+	writer, ok := cmd.OutOrStdout().(io.Writer)
 	if !ok {
 		return "", xerrors.New("stdout must be a readcloser")
 	}
+	prompt.Stdout = readWriteCloser{
+		Writer: writer,
+	}
 
 	// The prompt library displays defaults in a jarring way for the user
 	// by attempting to autocomplete it. This sets no default enabling us
@@ -199,3 +206,10 @@ func prompt(cmd *cobra.Command, prompt *promptui.Prompt) (string, error) {
 
 	return value, err
 }
+
+// readWriteCloser fakes reads, writes, and closing!
+type readWriteCloser struct {
+	io.Reader
+	io.Writer
+	io.Closer
+}
diff --git a/cli/workspacecreate_test.go b/cli/workspacecreate_test.go
@@ -5,9 +5,9 @@ import (
 
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
-	"github.com/coder/coder/console"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
+	"github.com/coder/coder/pty/ptytest"
 	"github.com/stretchr/testify/require"
 )
 
@@ -36,7 +36,9 @@ func TestWorkspaceCreate(t *testing.T) {
 		cmd, root := clitest.New(t, "workspaces", "create", project.Name)
 		clitest.SetupConfig(t, client, root)
 
-		cons := console.New(t, cmd)
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
 		closeChan := make(chan struct{})
 		go func() {
 			err := cmd.Execute()
@@ -51,13 +53,10 @@ func TestWorkspaceCreate(t *testing.T) {
 		for i := 0; i < len(matches); i += 2 {
 			match := matches[i]
 			value := matches[i+1]
-			_, err := cons.ExpectString(match)
-			require.NoError(t, err)
-			_, err = cons.SendLine(value)
-			require.NoError(t, err)
+			pty.ExpectMatch(match)
+			pty.WriteLine(value)
 		}
-		_, err := cons.ExpectString("Create")
-		require.NoError(t, err)
+		pty.ExpectMatch("Create")
 		<-closeChan
 	})
 }
diff --git a/console/conpty/conpty.go b/console/conpty/conpty.go