Fix race conditions with stderr/stdout streams

Author: cmoog

client.go

@@ -1,13 +1,14 @@
 package wsep
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"io"
 	"net"
 
 	"cdr.dev/wsep/internal/proto"
-	"go.coder.com/flog"
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
 )
@@ -123,46 +124,67 @@ func newPipe() pipe {
 
 func (r remoteProcess) listen(ctx context.Context) {
 	defer r.conn.Close(websocket.StatusNormalClosure, "normal closure")
-	defer r.stdout.w.Close()
-	defer r.stderr.w.Close()
 
-	for {
-		if err := ctx.Err(); err != nil {
-			r.done <- xerrors.Errorf("process canceled: %w", err)
-			break
-		}
-		_, payload, err := r.conn.Read(ctx)
-		if err != nil {
-			continue
-		}
-		headerByt, body := proto.SplitMessage(payload)
+	exitCode := make(chan int, 1)
+	var eg errgroup.Group
 
-		var header proto.Header
-		err = json.Unmarshal(headerByt, &header)
-		if err != nil {
-			continue
-		}
+	eg.Go(func() error {
+		defer r.stdout.w.Close()
+		defer r.stderr.w.Close()
 
-		switch header.Type {
-		case proto.TypeStderr:
-			go r.stderr.w.Write(body)
-		case proto.TypeStdout:
-			go r.stdout.w.Write(body)
-		case proto.TypeExitCode:
-			var exitMsg proto.ServerExitCodeHeader
-			err = json.Unmarshal(headerByt, &exitMsg)
+		buf := make([]byte, 32<<10) // max size of one websocket message
+		for {
+			if err := ctx.Err(); err != nil {
+				r.done <- xerrors.Errorf("process canceled: %w", err)
+				break
+			}
+			_, payload, err := r.conn.Read(ctx)
 			if err != nil {
-				flog.Error("failed to unmarshal exit code message: %v", err)
 				continue
 			}
+			headerByt, body := proto.SplitMessage(payload)
 
-			var err error = ExitError{Code: exitMsg.ExitCode}
-			if exitMsg.ExitCode == 0 {
-				err = nil
+			var header proto.Header
+			err = json.Unmarshal(headerByt, &header)
+			if err != nil {
+				continue
+			}
+
+			switch header.Type {
+			case proto.TypeStderr:
+				_, err = io.CopyBuffer(r.stderr.w, bytes.NewReader(body), buf)
+				if err != nil {
+					return err
+				}
+			case proto.TypeStdout:
+				_, err = io.CopyBuffer(r.stdout.w, bytes.NewReader(body), buf)
+				if err != nil {
+					return err
+				}
+			case proto.TypeExitCode:
+				var exitMsg proto.ServerExitCodeHeader
+				err = json.Unmarshal(headerByt, &exitMsg)
+				if err != nil {
+					continue
+				}
+
+				exitCode <- exitMsg.ExitCode
+				return nil
 			}
-			r.done <- err
+		}
+		return nil
+	})
+
+	err := eg.Wait()
+	select {
+	case exitCode := <-exitCode:
+		if exitCode != 0 {
+			r.done <- ExitError{Code: int(exitCode)}
 			return
 		}
+		r.done <- nil
+	default:
+		r.done <- err
 	}
 }
 

go.mod

@@ -10,6 +10,7 @@ require (
 	go.coder.com/cli v0.4.0
 	go.coder.com/flog v0.0.0-20190906214207-47dd47ea0512
 	golang.org/x/crypto v0.0.0-20191206172530-e9b2fee46413
+	golang.org/x/sync v0.0.0-20190423024810-112230192c58
 	golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543
 	nhooyr.io/websocket v1.8.6
 )

go.sum

@@ -206,6 +206,7 @@ golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58 h1:8gQV6CLnAEikrhgkHFbMAEhagSSnXWGV915qUMm9mrU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20181128092732-4ed8d59d0b35/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=

server.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync"
 
 	"go.coder.com/flog"
 	"golang.org/x/xerrors"
@@ -67,11 +68,13 @@ func Serve(ctx context.Context, c *websocket.Conn, execer Execer) error {
 			}
 
 			sendPID(ctx, process.Pid(), wsNetConn)
-			go pipeProcessOutput(ctx, process, wsNetConn)
+			var wg sync.WaitGroup
+			go pipeProcessOutput(ctx, process, wsNetConn, &wg)
 
 			go func() {
 				defer wsNetConn.Close()
 				err = process.Wait()
+				wg.Wait()
 				if exitErr, ok := err.(*ExitError); ok {
 					sendExitCode(ctx, exitErr.Code, wsNetConn)
 					return
@@ -128,16 +131,18 @@ 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) {
+func pipeProcessOutput(ctx context.Context, process Process, conn net.Conn, wg *sync.WaitGroup) {
 	var (
 		stdout = process.Stdout()
 		stderr = process.Stderr()
 	)
-	go copyWithHeader(stdout, conn, proto.Header{Type: proto.TypeStdout})
-	go copyWithHeader(stderr, conn, proto.Header{Type: proto.TypeStderr})
+	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) {
+func copyWithHeader(r io.Reader, w io.Writer, header proto.Header, wg *sync.WaitGroup) {
+	defer wg.Done()
 	headerByt, err := json.Marshal(header)
 	if err != nil {
 		return