Cleanup exit error handling

Author: cmoog

client.go

@@ -156,11 +156,7 @@ func (r remoteProcess) listen(ctx context.Context) {
 		defer r.stderr.w.Close()
 
 		buf := make([]byte, maxMessageSize) // max size of one websocket message
-		for {
-			if err := ctx.Err(); err != nil {
-				r.done <- xerrors.Errorf("process canceled: %w", err)
-				break
-			}
+		for ctx.Err() == nil {
 			_, payload, err := r.conn.Read(ctx)
 			if err != nil {
 				continue
@@ -195,7 +191,7 @@ func (r remoteProcess) listen(ctx context.Context) {
 				return nil
 			}
 		}
-		return nil
+		return ctx.Err()
 	})
 
 	err := eg.Wait()

client_test.go

@@ -81,3 +81,31 @@ func TestRemoteExec(t *testing.T) {
 	execer := RemoteExecer(ws)
 	testExecer(ctx, t, execer)
 }
+
+func TestRemoteExecFail(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+
+	ws, server := mockConn(ctx, t)
+	defer server.Close()
+
+	execer := RemoteExecer(ws)
+	testExecerFail(ctx, t, execer)
+}
+
+func testExecerFail(ctx context.Context, t *testing.T, execer Execer) {
+	process, err := execer.Start(ctx, Command{
+		Command: "ls",
+		Args:    []string{"/doesnotexist"},
+	})
+	assert.Success(t, "start local cmd", err)
+
+	go io.Copy(ioutil.Discard, process.Stderr())
+	go io.Copy(ioutil.Discard, process.Stdout())
+
+	err = process.Wait()
+	code, ok := err.(ExitError)
+	assert.True(t, "is exit error", ok)
+	assert.True(t, "exit code is nonzero", code.Code != 0)
+	assert.Error(t, "wait for process to error", err)
+}

localexec.go

@@ -42,7 +42,7 @@ func (l *localProcess) Stderr() io.Reader {
 func (l *localProcess) Wait() error {
 	err := l.cmd.Wait()
 	if exitErr, ok := err.(*exec.ExitError); ok {
-		return &ExitError{
+		return ExitError{
 			Code: exitErr.ExitCode(),
 		}
 	}

localexec_test.go

@@ -67,8 +67,8 @@ func TestExitCode(t *testing.T) {
 	assert.Success(t, "start local cmd", err)
 
 	err = process.Wait()
-	exitErr, ok := err.(*ExitError)
-	assert.True(t, "error is *ExitError", ok)
+	exitErr, ok := err.(ExitError)
+	assert.True(t, "error is ExitError", ok)
 	assert.Equal(t, "exit error", exitErr.Code, 127)
 }
 

server.go

@@ -71,7 +71,7 @@ func Serve(ctx context.Context, c *websocket.Conn, execer Execer) error {
 				return err
 			}
 
-			sendPID(ctx, process.Pid(), wsNetConn)
+			_ = sendPID(ctx, process.Pid(), wsNetConn)
 
 			var outputgroup errgroup.Group
 			outputgroup.Go(func() error {
@@ -85,11 +85,11 @@ func Serve(ctx context.Context, c *websocket.Conn, execer Execer) error {
 				defer wsNetConn.Close()
 				_ = outputgroup.Wait()
 				err = process.Wait()
-				if exitErr, ok := err.(*ExitError); ok {
-					sendExitCode(ctx, exitErr.Code, wsNetConn)
+				if exitErr, ok := err.(ExitError); ok {
+					_ = sendExitCode(ctx, exitErr.Code, wsNetConn)
 					return
 				}
-				sendExitCode(ctx, 0, wsNetConn)
+				_ = sendExitCode(ctx, 0, wsNetConn)
 			}()
 		case proto.TypeResize:
 			if process == nil {
@@ -122,23 +122,25 @@ func Serve(ctx context.Context, c *websocket.Conn, execer Execer) error {
 	}
 }
 
-func sendExitCode(ctx context.Context, exitCode int, conn net.Conn) {
+func sendExitCode(ctx context.Context, exitCode int, conn net.Conn) error {
 	header, err := json.Marshal(proto.ServerExitCodeHeader{
 		Type:     proto.TypeExitCode,
 		ExitCode: exitCode,
 	})
 	if err != nil {
-		return
+		return err
 	}
-	proto.WithHeader(conn, header).Write(nil)
+	_, err = proto.WithHeader(conn, header).Write(nil)
+	return err
 }
 
-func sendPID(ctx context.Context, pid int, conn net.Conn) {
+func sendPID(ctx context.Context, pid int, conn net.Conn) error {
 	header, err := json.Marshal(proto.ServerPidHeader{Type: proto.TypePid, Pid: pid})
 	if err != nil {
-		return
+		return err
 	}
-	proto.WithHeader(conn, header).Write(nil)
+	_, err = proto.WithHeader(conn, header).Write(nil)
+	return err
 }
 
 func copyWithHeader(r io.Reader, w io.Writer, header proto.Header) error {