Implement reconnections with screen (#25)

* Refactor reconnect test to support sub-tests

Going to add an alternate screen test next.

* Revert "Add reconnecting ptys (#23)"

This partially reverts commit 9120171.

The new method using screen will not share processes which is a
fundamental shift so I think it will be easier to start from scratch.

Even though we could keep the UUID check I removed it because it seems
cool that you could create your own sessions in the terminal then
connect to them in the browser (or vice-versa).

* Add test for alternate screen

The output test waits for EOF; modify that behavior so we can check that
certain strings are not displayed *without* waiting for the timeout.
This means to be accurate we should always check for output that should
exist after the output that should not exist would have shown up.

* Add timeout flag to dev client

This makes it easier to test reconnects manually.

* Add size to initial connection

This way you do not need a subsequent resize and we can have the right
size from the get-go.

* Prevent prompt from rendering twice in tests

* Add Nix flake

* Propagate process close error

* Implement reconnecting TTY with screen

* Encapsulate session logic

* Localize session map

* Consolidate test scaffolding into helpers

I think this helps make the tests a bit more concise.

* Test many connections at once

* Fix errors not propagating through web socket close

Since the server closed the socket the caller has no chance to close
with the right code and reason.

Also abnormal closure is not a valid close code.

* Fix test flake in reading output

Without waiting for the copy you can sometimes get "file already
closed".  I guess process.Wait must have some side effect.

Author: code-asher

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")
 }
 ```
 

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);
 };
 

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

ci/fmt.sh

@@ -1,4 +1,5 @@
-#!/bin/bash
+#!/usr/bin/env bash
+
 echo "Formatting..."
 
 go mod tidy

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

ci/lint.sh

@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
 echo "Linting..."
 

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
 }

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)