Fix more race conditions

Author: cmoog

README.md

@@ -1,6 +1,7 @@
 # wsep
 
-`wsep` is a high performance, ***W***eb***S***ocket command ***e***xecution ***p***rotocol. It can be thought of as SSH without encryption.
+`wsep` is a high performance,
+<strong style="font-size: 1.5em; text-decoration: underline;">w</strong>eb <strong style="font-size: 1.5em;text-decoration: underline;">s</strong>ocket command <strong style="font-size: 1.5em;text-decoration: underline;">e</strong>xecution <strong style="font-size: 1.5em;text-decoration: underline;">p</strong>rotocol. It can be thought of as SSH without encryption.
 
 It's useful in cases where you want to provide a command exec interface into a remote environment. It's implemented
 with WebSocket so it may be used directly by a browser frontend.
@@ -25,9 +26,7 @@ if err != nil {
   // handle error
 }
 
-go io.Copy(os.Stdout, process.Stdout())
-go io.Copy(os.Stderr, process.Stderr())
-go io.Copy(process.Stdin(), os.Stdin)
+io.Copy(os.Stderr, process.Stderr())
 
 err = process.Wait()
 if err != nil {

client_test.go

@@ -76,7 +76,6 @@ func TestRemoteExec(t *testing.T) {
 	defer cancel()
 
 	ws, server := mockConn(ctx, t)
-	defer ws.Close(websocket.StatusAbnormalClosure, "abnormal closure")
 	defer server.Close()
 
 	execer := RemoteExecer(ws)

exec.go

@@ -19,9 +19,14 @@ func (e ExitError) Error() string {
 
 // Process represents a started command.
 type Process interface {
+	// Pid is populated immediately during a successfull start with the process ID.
 	Pid() int
+	// Stdout returns an io.WriteCloser that will pipe writes to the remote command.
+	// Closure of stdin sends the correspoding close messsage.
 	Stdin() io.WriteCloser
+	// Stdout returns an io.Reader that is connected to the command's standard output.
 	Stdout() io.Reader
+	// Stderr returns an io.Reader that is connected to the command's standard error.
 	Stderr() io.Reader
 	// Resize resizes the TTY if a TTY is enabled.
 	Resize(ctx context.Context, rows, cols uint16) error

internal/proto/protocol.go

@@ -46,5 +46,5 @@ func (h headerWriter) Write(b []byte) (int, error) {
 	if err != nil {
 		return 0, err
 	}
-	return len(b), nil // TODO: potential buggy
+	return len(b), nil
 }

localexec_test.go

@@ -24,30 +24,33 @@ func testExecer(ctx context.Context, t *testing.T, execer Execer) {
 		Command: "pwd",
 	})
 	assert.Success(t, "start local cmd", err)
+	var (
+		stderr = process.Stderr()
+		stdout = process.Stdout()
+		wg     sync.WaitGroup
+	)
 
-	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
 
-		stdout, err := ioutil.ReadAll(process.Stdout())
+		stdoutByt, err := ioutil.ReadAll(stdout)
 		assert.Success(t, "read stdout", err)
 		wd, err := os.Getwd()
 		assert.Success(t, "get real working dir", err)
 
-		assert.Equal(t, "stdout", wd, strings.TrimSuffix(string(stdout), "\n"))
+		assert.Equal(t, "stdout", wd, strings.TrimSuffix(string(stdoutByt), "\n"))
 	}()
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
 
-		stderr, err := ioutil.ReadAll(process.Stderr())
+		stderrByt, err := ioutil.ReadAll(stderr)
 		assert.Success(t, "read stderr", err)
-		assert.True(t, "len stderr", len(stderr) == 0)
+		assert.True(t, "len stderr", len(stderrByt) == 0)
 	}()
 
+	wg.Wait()
 	err = process.Wait()
 	assert.Success(t, "wait for process to complete", err)
-
-	wg.Wait()
 }

server.go

@@ -7,9 +7,9 @@ import (
 	"errors"
 	"io"
 	"net"
-	"sync"
 
 	"go.coder.com/flog"
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
 
@@ -68,13 +68,23 @@ func Serve(ctx context.Context, c *websocket.Conn, execer Execer) error {
 			}
 
 			sendPID(ctx, process.Pid(), wsNetConn)
-			var wg sync.WaitGroup
-			go pipeProcessOutput(ctx, process, wsNetConn, &wg)
+
+			var outputgroup errgroup.Group
+			outputgroup.Go(func() error {
+				return copyWithHeader(process.Stdout(), wsNetConn, proto.Header{Type: proto.TypeStdout})
+			})
+			outputgroup.Go(func() error {
+				return copyWithHeader(process.Stderr(), wsNetConn, proto.Header{Type: proto.TypeStdout})
+			})
 
 			go func() {
 				defer wsNetConn.Close()
+				err := outputgroup.Wait()
+				if err != nil {
+					// connection should close without an exit code if copy fails
+					return
+				}
 				err = process.Wait()
-				wg.Wait()
 				if exitErr, ok := err.(*ExitError); ok {
 					sendExitCode(ctx, exitErr.Code, wsNetConn)
 					return
@@ -131,22 +141,15 @@ func sendPID(ctx context.Context, pid int, conn net.Conn) {
 	proto.WithHeader(conn, header).Write(nil)
 }
 
-func pipeProcessOutput(ctx context.Context, process Process, conn net.Conn, wg *sync.WaitGroup) {
-	var (
-		stdout = process.Stdout()
-		stderr = process.Stderr()
-	)
-	wg.Add(2)
-	go copyWithHeader(stdout, conn, proto.Header{Type: proto.TypeStdout}, wg)
-	go copyWithHeader(stderr, conn, proto.Header{Type: proto.TypeStderr}, wg)
-}
-
-func copyWithHeader(r io.Reader, w io.Writer, header proto.Header, wg *sync.WaitGroup) {
-	defer wg.Done()
+func copyWithHeader(r io.Reader, w io.Writer, header proto.Header) error {
 	headerByt, err := json.Marshal(header)
 	if err != nil {
-		return
+		return err
 	}
 	wr := proto.WithHeader(w, headerByt)
-	io.Copy(wr, r)
+	_, err = io.Copy(wr, r)
+	if err != nil {
+		return err
+	}
+	return nil
 }