Code review fixes

Signed-off-by: Spike Curtis <spike@coder.com>
coder · spikecurtis · Apr 24, 2023 · Apr 11, 2023 · Apr 12, 2023 · Apr 12, 2023
commit c09083ef56915c900b93707dcd86b267d8cf40f1
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -24,8 +24,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/coder/coder/pty"
-
 	scp "github.com/bramvdbogaerde/go-scp"
 	"github.com/google/uuid"
 	"github.com/pion/udp"
@@ -47,6 +45,7 @@ import (
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
+	"github.com/coder/coder/pty"
 	"github.com/coder/coder/pty/ptytest"
 	"github.com/coder/coder/tailnet"
 	"github.com/coder/coder/tailnet/tailnettest"

diff --git a/go.mod b/go.mod
@@ -108,6 +108,7 @@ require (
 	github.com/hashicorp/terraform-config-inspect v0.0.0-20211115214459-90acf1ca460f
 	github.com/hashicorp/terraform-json v0.14.0
 	github.com/hashicorp/yamux v0.0.0-20220718163420-dd80a7ee44ce
+	github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02
 	github.com/imulab/go-scim/pkg/v2 v2.2.0
 	github.com/jedib0t/go-pretty/v6 v6.4.0
 	github.com/jmoiron/sqlx v1.3.5
@@ -174,8 +175,6 @@ require (
 	tailscale.com v1.32.2
 )
 
-require github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02 // indirect
-
 require (
 	cloud.google.com/go/compute v1.18.0 // indirect
 	cloud.google.com/go/logging v1.6.1 // indirect

diff --git a/go.sum b/go.sum
@@ -1086,7 +1086,6 @@ github.com/hashicorp/yamux v0.0.0-20220718163420-dd80a7ee44ce h1:7FO+LmZwiG/eDsB
 github.com/hashicorp/yamux v0.0.0-20220718163420-dd80a7ee44ce/go.mod h1:CtWFDAQgb7dxtzFs4tWbplKIe2jSi3+5vKbgIO0SLnQ=
 github.com/hdevalence/ed25519consensus v0.0.0-20220222234857-c00d1f31bab3 h1:aSVUgRRRtOrZOC1fYmY9gV0e9z/Iu+xNVSASWjsuyGU=
 github.com/hdevalence/ed25519consensus v0.0.0-20220222234857-c00d1f31bab3/go.mod h1:5PC6ZNPde8bBqU/ewGZig35+UIZtw9Ytxez8/q5ZyFE=
-github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec h1:qv2VnGeEQHchGaZ/u7lxST/RaJw+cv273q79D81Xbog=
 github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec/go.mod h1:Q48J4R4DvxnHolD5P8pOtXigYlRuPLGl6moFx3ulM68=
 github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02 h1:AgcIVYPa6XJnU3phs104wLj8l5GEththEw6+F79YsIY=
 github.com/hinshun/vt10x v0.0.0-20220301184237-5011da428d02/go.mod h1:Q48J4R4DvxnHolD5P8pOtXigYlRuPLGl6moFx3ulM68=

diff --git a/pty/pty_other.go b/pty/pty_other.go
@@ -10,11 +10,10 @@ import (
 	"runtime"
 	"sync"
 
-	"golang.org/x/xerrors"
-
 	"github.com/creack/pty"
 	"github.com/u-root/u-root/pkg/termios"
 	"golang.org/x/sys/unix"
+	"golang.org/x/xerrors"
 )
 
 func newPty(opt ...Option) (retPTY *otherPty, err error) {

diff --git a/pty/pty_windows.go b/pty/pty_windows.go
@@ -160,13 +160,15 @@ func (p *ptyWindows) Close() error {
 	if p.outputWrite != nil {
 		_ = p.outputWrite.Close()
 	}
-	if p.inputRead.Close() != nil {
+	if p.inputRead != nil {
 		_ = p.inputRead.Close()
 	}
 	return nil
 }
 
 func (p *windowsProcess) waitInternal() {
+	// put this on the bottom of the defer stack since the next defer can write to p.cmdErr
+	defer close(p.cmdDone)
 	defer func() {
 		// close the pseudoconsole handle when the process exits, if it hasn't already been closed.
 		// this is important because the PseudoConsole (conhost.exe) holds the write-end
@@ -186,7 +188,7 @@ func (p *windowsProcess) waitInternal() {
 			p.pw.console = windows.InvalidHandle
 		}
 	}()
-	defer close(p.cmdDone)
+
 	state, err := p.proc.Wait()
 	if err != nil {
 		p.cmdErr = err

diff --git a/pty/ptytest/ptytest.go b/pty/ptytest/ptytest.go
@@ -32,7 +32,7 @@ func New(t *testing.T, opts ...pty.Option) *PTY {
 
 	e := newExpecter(t, ptty.Output(), "cmd")
 	r := &PTY{
-		outExpecter: *e,
+		outExpecter: e,
 		PTY:         ptty,
 	}
 	// Ensure pty is cleaned up at the end of test.
@@ -56,7 +56,7 @@ func Start(t *testing.T, cmd *exec.Cmd, opts ...pty.StartOption) (*PTYCmd, pty.P
 	ex := newExpecter(t, ptty.OutputReader(), cmd.Args[0])
 
 	r := &PTYCmd{
-		outExpecter: *ex,
+		outExpecter: ex,
 		PTYCmd:      ptty,
 	}
 	t.Cleanup(func() {
@@ -65,7 +65,7 @@ func Start(t *testing.T, cmd *exec.Cmd, opts ...pty.StartOption) (*PTYCmd, pty.P
 	return r, ps
 }
 
-func newExpecter(t *testing.T, r io.Reader, name string) *outExpecter {
+func newExpecter(t *testing.T, r io.Reader, name string) outExpecter {
 	// Use pipe for logging.
 	logDone := make(chan struct{})
 	logr, logw := io.Pipe()
@@ -75,7 +75,7 @@ func newExpecter(t *testing.T, r io.Reader, name string) *outExpecter {
 	out := newStdbuf()
 	w := io.MultiWriter(logw, out)
 
-	ex := &outExpecter{
+	ex := outExpecter{
 		t:    t,
 		out:  out,
 		name: name,

diff --git a/pty/start_test.go b/pty/start_test.go
@@ -29,7 +29,7 @@ func Test_Start_copy(t *testing.T) {
 	pc, cmd, err := pty.Start(exec.CommandContext(ctx, cmdEcho, argEcho...))
 	require.NoError(t, err)
 	b := &bytes.Buffer{}
-	readDone := make(chan error)
+	readDone := make(chan error, 1)
 	go func() {
 		_, err := io.Copy(b, pc.OutputReader())
 		readDone <- err
@@ -43,7 +43,7 @@ func Test_Start_copy(t *testing.T) {
 	}
 	assert.Contains(t, b.String(), "test")
 
-	cmdDone := make(chan error)
+	cmdDone := make(chan error, 1)
 	go func() {
 		cmdDone <- cmd.Wait()
 	}()
@@ -58,7 +58,7 @@ func Test_Start_copy(t *testing.T) {
 
 // Test_Start_truncation tests that we can read command output without truncation
 // even after the command has exited.
-func Test_Start_trucation(t *testing.T) {
+func Test_Start_truncation(t *testing.T) {
 	t.Parallel()
 
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
@@ -94,7 +94,7 @@ func Test_Start_trucation(t *testing.T) {
 		assert.NoError(t, err)
 	}()
 
-	cmdDone := make(chan error)
+	cmdDone := make(chan error, 1)
 	go func() {
 		cmdDone <- cmd.Wait()
 	}()
@@ -103,14 +103,14 @@ func Test_Start_trucation(t *testing.T) {
 	case err := <-cmdDone:
 		require.NoError(t, err)
 	case <-ctx.Done():
-		t.Error("cmd.Wait() timed out")
+		t.Fatal("cmd.Wait() timed out")
 	}
 
 	select {
 	case <-readDone:
 		// OK!
 	case <-ctx.Done():
-		t.Error("read timed out")
+		t.Fatal("read timed out")
 	}
 }
 

diff --git a/pty/start_windows.go b/pty/start_windows.go
@@ -17,7 +17,7 @@ import (
 
 // Allocates a PTY and starts the specified command attached to it.
 // See: https://docs.microsoft.com/en-us/windows/console/creating-a-pseudoconsole-session#creating-the-hosted-process
-func startPty(cmd *exec.Cmd, opt ...StartOption) (PTY, Process, error) {
+func startPty(cmd *exec.Cmd, opt ...StartOption) (_ PTY, _ Process, retErr error) {
 	var opts startOptions
 	for _, o := range opt {
 		o(&opts)
@@ -45,10 +45,18 @@ func startPty(cmd *exec.Cmd, opt ...StartOption) (PTY, Process, error) {
 	if err != nil {
 		return nil, nil, err
 	}
+
 	pty, err := newPty(opts.ptyOpts...)
 	if err != nil {
 		return nil, nil, err
 	}
+	defer func() {
+		if retErr != nil {
+			// we hit some error finishing setup; close pty, so
+			// we don't leak the kernel resources associated with it
+			_ := pty.Close()
+		}
+	}()
 	winPty := pty.(*ptyWindows)
 	if winPty.opts.sshReq != nil {
 		cmd.Env = append(cmd.Env, fmt.Sprintf("SSH_TTY=%s", winPty.Name()))
@@ -87,6 +95,24 @@ func startPty(cmd *exec.Cmd, opt ...StartOption) (PTY, Process, error) {
 	}
 	defer windows.CloseHandle(processInfo.Thread)
 	defer windows.CloseHandle(processInfo.Process)
+
+	process, err := os.FindProcess(int(processInfo.ProcessId))
+	if err != nil {
+		return nil, nil, xerrors.Errorf("find process %d: %w", processInfo.ProcessId, err)
+	}
+	wp := &windowsProcess{
+		cmdDone: make(chan any),
+		proc:    process,
+		pw:      winPty,
+	}
+	defer func() {
+		if retErr != nil {
+			// if we later error out, kill the process since
+			// the caller will have no way to interact with it
+			_ = process.Kill()
+		}
+	}()
+
 	// Now that we've started the command, and passed the pseudoconsole to it,
 	// close the output write and input read files, so that the other process
 	// has the only handles to them.  Once the process closes the console, there
@@ -103,16 +129,6 @@ func startPty(cmd *exec.Cmd, opt ...StartOption) (PTY, Process, error) {
 	if errI != nil {
 		return nil, nil, errI
 	}
-
-	process, err := os.FindProcess(int(processInfo.ProcessId))
-	if err != nil {
-		return nil, nil, xerrors.Errorf("find process %d: %w", processInfo.ProcessId, err)
-	}
-	wp := &windowsProcess{
-		cmdDone: make(chan any),
-		proc:    process,
-		pw:      winPty,
-	}
 	go wp.waitInternal()
 	return pty, wp, nil
 }