1 files changed, 41 insertions, 23 deletions
diff --git a/pkg/varlinkapi/virtwriter/virtwriter.go b/pkg/varlinkapi/virtwriter/virtwriter.go
index 0da2a91fc..27ecd1f52 100644
--- a/pkg/varlinkapi/virtwriter/virtwriter.go
+++ b/pkg/varlinkapi/virtwriter/virtwriter.go
@@ -89,10 +89,14 @@ func (v VirtWriteCloser) Write(input []byte) (int, error) {
 }
 
 // Reader decodes the content that comes over the wire and directs it to the proper destination.
-func Reader(r *bufio.Reader, output io.Writer, errput io.Writer, input io.Writer, resize chan remotecommand.TerminalSize) error {
+func Reader(r *bufio.Reader, output, errput, input io.Writer, resize chan remotecommand.TerminalSize, execEcChan chan int) error {
 	var messageSize int64
 	headerBytes := make([]byte, 8)
 
+	if r == nil {
+		return errors.Errorf("Reader must not be nil")
+	}
+
 	for {
 		n, err := io.ReadFull(r, headerBytes)
 		if err != nil {
@@ -106,35 +110,43 @@ func Reader(r *bufio.Reader, output io.Writer, errput io.Writer, input io.Writer
 
 		switch IntToSocketDest(int(headerBytes[0])) {
 		case ToStdout:
-			_, err := io.CopyN(output, r, messageSize)
-			if err != nil {
-				return err
+			if output != nil {
+				_, err := io.CopyN(output, r, messageSize)
+				if err != nil {
+					return err
+				}
 			}
 		case ToStderr:
-			_, err := io.CopyN(errput, r, messageSize)
-			if err != nil {
-				return err
+			if errput != nil {
+				_, err := io.CopyN(errput, r, messageSize)
+				if err != nil {
+					return err
+				}
 			}
 		case ToStdin:
-			_, err := io.CopyN(input, r, messageSize)
-			if err != nil {
-				return err
-			}
-		case TerminalResize:
-			out := make([]byte, messageSize)
-			if messageSize > 0 {
-				_, err = io.ReadFull(r, out)
-
+			if input != nil {
+				_, err := io.CopyN(input, r, messageSize)
 				if err != nil {
 					return err
 				}
 			}
-			// Resize events come over in bytes, need to be reserialized
-			resizeEvent := remotecommand.TerminalSize{}
-			if err := json.Unmarshal(out, &resizeEvent); err != nil {
-				return err
+		case TerminalResize:
+			if resize != nil {
+				out := make([]byte, messageSize)
+				if messageSize > 0 {
+					_, err = io.ReadFull(r, out)
+
+					if err != nil {
+						return err
+					}
+				}
+				// Resize events come over in bytes, need to be reserialized
+				resizeEvent := remotecommand.TerminalSize{}
+				if err := json.Unmarshal(out, &resizeEvent); err != nil {
+					return err
+				}
+				resize <- resizeEvent
 			}
-			resize <- resizeEvent
 		case Quit:
 			out := make([]byte, messageSize)
 			if messageSize > 0 {
@@ -144,6 +156,10 @@ func Reader(r *bufio.Reader, output io.Writer, errput io.Writer, input io.Writer
 					return err
 				}
 			}
+			if execEcChan != nil {
+				ecInt := binary.BigEndian.Uint32(out)
+				execEcChan <- int(ecInt)
+			}
 			return nil
 
 		default:
@@ -154,9 +170,11 @@ func Reader(r *bufio.Reader, output io.Writer, errput io.Writer, input io.Writer
 }
 
 // HangUp sends message to peer to close connection
-func HangUp(writer *bufio.Writer) (err error) {
+func HangUp(writer *bufio.Writer, ec uint32) (err error) {
 	n := 0
-	msg := []byte("HANG-UP")
+	msg := make([]byte, 4)
+
+	binary.BigEndian.PutUint32(msg, ec)
 
 	writeQuit := NewVirtWriteCloser(writer, Quit)
 	if n, err = writeQuit.Write(msg); err != nil {