coder · code-asher · Nov 22, 2022 · Aug 23, 2022 · Sep 13, 2022 · Aug 23, 2022
diff --git a/README.md b/README.md
@@ -12,7 +12,7 @@ Error handling is omitted for brevity.
 
 ```golang
 conn, _, _ := websocket.Dial(ctx, "ws://remote.exec.addr", nil)
-defer conn.Close(websocket.StatusAbnormalClosure, "terminate process")
+defer conn.Close(websocket.StatusNormalClosure, "normal closure")
 
 execer := wsep.RemoteExecer(conn)
 process, _ := execer.Start(ctx, wsep.Command{
@@ -25,18 +25,16 @@ go io.Copy(os.Stderr, process.Stderr())
 go io.Copy(os.Stdout, process.Stdout())
 
 process.Wait()
-conn.Close(websocket.StatusNormalClosure, "normal closure")
 ```
 
 ### Server
 
 ```golang
 func (s server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
   conn, _ := websocket.Accept(w, r, nil)
+  defer conn.Close(websocket.StatusNormalClosure, "normal closure")
 
   wsep.Serve(r.Context(), conn, wsep.LocalExecer{})
-
-  ws.Close(websocket.StatusNormalClosure, "normal closure")
 }
 ```
 

diff --git a/browser/client.ts b/browser/client.ts
@@ -14,7 +14,7 @@ export interface Command {
 }
 
 export type ClientHeader =
-  | { type: 'start'; command: Command }
+  | { type: 'start'; id: string; command: Command; cols: number; rows: number; }
   | { type: 'stdin' }
   | { type: 'close_stdin' }
   | { type: 'resize'; cols: number; rows: number };
@@ -42,8 +42,14 @@ export const closeStdin = (ws: WebSocket) => {
   ws.send(msg.buffer);
 };
 
-export const startCommand = (ws: WebSocket, command: Command) => {
-  const msg = joinMessage({ type: 'start', command: command });
+export const startCommand = (
+  ws: WebSocket,
+  command: Command,
+  id: string,
+  rows: number,
+  cols: number
+) => {
+  const msg = joinMessage({ type: 'start', command, id, rows, cols });
   ws.send(msg.buffer);
 };
 

diff --git a/ci/alt.sh b/ci/alt.sh
@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+# Script for testing the alt screen.
+
+# Enter alt screen.
+tput smcup
+
+function display() {
+  # Clear the screen.
+  tput clear
+  # Move cursor to the top left.
+  tput cup 0 0
+  # Display content.
+  echo "ALT SCREEN"
+}
+
+function redraw() {
+  display
+  echo "redrawn"
+}
+
+# Re-display on resize.
+trap 'redraw' WINCH
+
+display
+
+# The trap will not run while waiting for a command so read input in a loop with
+# a timeout.
+while true ; do
+  if read -n 1 -t .1 ; then
+    # Clear the screen.
+    tput clear
+    # Exit alt screen.
+    tput rmcup
+    exit
+  fi
+done
diff --git a/ci/fmt.sh b/ci/fmt.sh
@@ -1,4 +1,5 @@
-#!/bin/bash
+#!/usr/bin/env bash
+
 echo "Formatting..."
 
 go mod tidy

diff --git a/ci/image/Dockerfile b/ci/image/Dockerfile
@@ -3,6 +3,8 @@ FROM golang:1
 ENV GOFLAGS="-mod=readonly"
 ENV CI=true
 
+RUN apt update && apt install -y screen
+
 RUN go install golang.org/x/tools/cmd/goimports@latest
 RUN go install golang.org/x/lint/golint@latest
 RUN go install github.com/mattn/goveralls@latest
diff --git a/ci/lint.sh b/ci/lint.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
 echo "Linting..."
 

diff --git a/client.go b/client.go
@@ -28,10 +28,13 @@ func RemoteExecer(conn *websocket.Conn) Execer {
 // Command represents an external command to be run
 type Command struct {
 	// ID allows reconnecting commands that have a TTY.
-	ID         string
-	Command    string
-	Args       []string
+	ID      string
+	Command string
+	Args    []string
+	// Commands with a TTY also require Rows and Cols.
 	TTY        bool
+	Rows       uint16
+	Cols       uint16
 	Stdin      bool
 	UID        uint32
 	GID        uint32
@@ -103,7 +106,7 @@ type remoteProcess struct {
 	pid          int
 	done         chan struct{}
 	closeErr     error
-	exitCode     *int
+	exitMsg      *proto.ServerExitCodeHeader
 	readErr      error
 	stdin        io.WriteCloser
 	stdout       pipe
@@ -267,8 +270,7 @@ func (r *remoteProcess) listen(ctx context.Context) {
 				r.readErr = err
 				return
 			}
-
-			r.exitCode = &exitMsg.ExitCode
+			r.exitMsg = &exitMsg
 			return
 		}
 	}
@@ -319,11 +321,10 @@ func (r *remoteProcess) Wait() error {
 	if r.readErr != nil {
 		return r.readErr
 	}
-	// when listen() closes r.done, either there must be a read error
-	// or exitCode is set non-nil, so it's safe to dereference the pointer
-	// here
-	if *r.exitCode != 0 {
-		return ExitError{Code: *r.exitCode}
+	// when listen() closes r.done, either there must be a read error or exitMsg
+	// is set non-nil, so it's safe to access members here.
+	if r.exitMsg.ExitCode != 0 {
+		return ExitError{code: r.exitMsg.ExitCode, error: r.exitMsg.Error}
 	}
 	return nil
 }

diff --git a/client_test.go b/client_test.go
@@ -51,17 +51,25 @@ func TestRemoteStdin(t *testing.T) {
 	}
 }
 
-func mockConn(ctx context.Context, t *testing.T, options *Options) (*websocket.Conn, *httptest.Server) {
+func mockConn(ctx context.Context, t *testing.T, wsepServer *Server, options *Options) (*websocket.Conn, *httptest.Server) {
 	mockServerHandler := func(w http.ResponseWriter, r *http.Request) {
 		ws, err := websocket.Accept(w, r, nil)
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
 			return
 		}
-		err = Serve(r.Context(), ws, LocalExecer{}, options)
+		if wsepServer != nil {
+			err = wsepServer.Serve(r.Context(), ws, LocalExecer{}, options)
+		} else {
+			err = Serve(r.Context(), ws, LocalExecer{}, options)
+		}
 		if err != nil {
-			t.Errorf("failed to serve execer: %v", err)
-			ws.Close(websocket.StatusAbnormalClosure, "failed to serve execer")
+			// Max reason string length is 123.
+			errStr := err.Error()
+			if len(errStr) > 123 {
+				errStr = errStr[:123]
+			}
+			ws.Close(websocket.StatusInternalError, errStr)
 			return
 		}
 		ws.Close(websocket.StatusNormalClosure, "normal closure")
@@ -79,7 +87,11 @@ func TestRemoteExec(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
-	ws, server := mockConn(ctx, t, nil)
+	wsepServer := NewServer()
+	defer wsepServer.Close()
+	defer assert.Equal(t, "no leaked sessions", 0, wsepServer.SessionCount())
+
+	ws, server := mockConn(ctx, t, wsepServer, nil)
 	defer server.Close()
 
 	execer := RemoteExecer(ws)
@@ -92,14 +104,20 @@ func TestRemoteClose(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
-	ws, server := mockConn(ctx, t, nil)
+	wsepServer := NewServer()
+	defer wsepServer.Close()
+	defer assert.Equal(t, "no leaked sessions", 0, wsepServer.SessionCount())
+
+	ws, server := mockConn(ctx, t, wsepServer, nil)
 	defer server.Close()
 
 	execer := RemoteExecer(ws)
 	cmd := Command{
-		Command: "/bin/bash",
+		Command: "sh",
 		TTY:     true,
 		Stdin:   true,
+		Cols:    100,
+		Rows:    100,
 		Env:     []string{"TERM=linux"},
 	}
 
@@ -138,14 +156,20 @@ func TestRemoteCloseNoData(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
-	ws, server := mockConn(ctx, t, nil)
+	wsepServer := NewServer()
+	defer wsepServer.Close()
+	defer assert.Equal(t, "no leaked sessions", 0, wsepServer.SessionCount())
+
+	ws, server := mockConn(ctx, t, wsepServer, nil)
 	defer server.Close()
 
 	execer := RemoteExecer(ws)
 	cmd := Command{
-		Command: "/bin/bash",
+		Command: "sh",
 		TTY:     true,
 		Stdin:   true,
+		Cols:    100,
+		Rows:    100,
 		Env:     []string{"TERM=linux"},
 	}
 
@@ -171,14 +195,20 @@ func TestRemoteClosePartialRead(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
-	ws, server := mockConn(ctx, t, nil)
+	wsepServer := NewServer()
+	defer wsepServer.Close()
+	defer assert.Equal(t, "no leaked sessions", 0, wsepServer.SessionCount())
+
+	ws, server := mockConn(ctx, t, wsepServer, nil)
 	defer server.Close()
 
 	execer := RemoteExecer(ws)
 	cmd := Command{
-		Command: "/bin/bash",
+		Command: "sh",
 		TTY:     true,
 		Stdin:   true,
+		Cols:    100,
+		Rows:    100,
 		Env:     []string{"TERM=linux"},
 	}
 
@@ -205,7 +235,11 @@ func TestRemoteExecFail(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
-	ws, server := mockConn(ctx, t, nil)
+	wsepServer := NewServer()
+	defer wsepServer.Close()
+	defer assert.Equal(t, "no leaked sessions", 0, wsepServer.SessionCount())
+
+	ws, server := mockConn(ctx, t, wsepServer, nil)
 	defer server.Close()
 
 	execer := RemoteExecer(ws)
@@ -223,9 +257,10 @@ func testExecerFail(ctx context.Context, t *testing.T, execer Execer) {
 	go io.Copy(ioutil.Discard, process.Stdout())
 
 	err = process.Wait()
-	code, ok := err.(ExitError)
+	exitErr, ok := err.(ExitError)
 	assert.True(t, "is exit error", ok)
-	assert.True(t, "exit code is nonzero", code.Code != 0)
+	assert.True(t, "exit code is nonzero", exitErr.ExitCode() != 0)
+	assert.Equal(t, "exit error", exitErr.Error(), "exit status 2")
 	assert.Error(t, "wait for process to error", err)
 }
 
@@ -239,7 +274,11 @@ func TestStderrVsStdout(t *testing.T) {
 		stderr bytes.Buffer
 	)
 
-	ws, server := mockConn(ctx, t, nil)
+	wsepServer := NewServer()
+	defer wsepServer.Close()
+	defer assert.Equal(t, "no leaked sessions", 0, wsepServer.SessionCount())
+
+	ws, server := mockConn(ctx, t, wsepServer, nil)
 	defer server.Close()
 
 	execer := RemoteExecer(ws)