From a21fe8c4840865bcd5210c8304e5f5d6dd461526 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 2 Jun 2013 23:21:55 -0700
Subject: [PATCH] packer/rpc: Exited fields now work over RPC

---
 packer/rpc/communicator.go      | 53 +++++++++++++++++++++++++++++++++
 packer/rpc/communicator_test.go | 15 +++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/packer/rpc/communicator.go b/packer/rpc/communicator.go
index 5d43406d3..0d0fe8bc1 100644
--- a/packer/rpc/communicator.go
+++ b/packer/rpc/communicator.go
@@ -1,12 +1,14 @@
 package rpc
 
 import (
+	"encoding/gob"
 	"errors"
 	"github.com/mitchellh/packer/packer"
 	"io"
 	"log"
 	"net"
 	"net/rpc"
+	"time"
 )
 
 // An implementation of packer.Communicator where the communicator is actually
@@ -21,11 +23,16 @@ type CommunicatorServer struct {
 	c packer.Communicator
 }
 
+type CommandFinished struct {
+	ExitStatus int
+}
+
 type CommunicatorStartArgs struct {
 	Command       string
 	StdinAddress string
 	StdoutAddress string
 	StderrAddress string
+	ResponseAddress string
 }
 
 type CommunicatorDownloadArgs struct {
@@ -64,6 +71,30 @@ func (c *communicator) Start(cmd *packer.RemoteCmd) (err error) {
 		go serveSingleCopy("stderr", stderrL, cmd.Stderr, nil)
 	}
 
+	responseL := netListenerInRange(portRangeMin, portRangeMax)
+	args.ResponseAddress = responseL.Addr().String()
+
+	go func() {
+		defer responseL.Close()
+
+		conn, err := responseL.Accept()
+		if err != nil {
+			log.Panic(err)
+		}
+
+		defer conn.Close()
+
+		decoder := gob.NewDecoder(conn)
+
+		var finished CommandFinished
+		if err := decoder.Decode(&finished); err != nil {
+			log.Panic(err)
+		}
+
+		cmd.ExitStatus = finished.ExitStatus
+		cmd.Exited = true
+	}()
+
 	err = c.client.Call("Communicator.Start", &args, new(interface{}))
 	return
 }
@@ -149,8 +180,30 @@ func (c *CommunicatorServer) Start(args *CommunicatorStartArgs, reply *interface
 		cmd.Stderr = stderrC
 	}
 
+	// Connect to the response address so we can write our result to it
+	// when ready.
+	responseC, err := net.Dial("tcp", args.ResponseAddress)
+	if err != nil {
+		return err
+	}
+
+	responseWriter := gob.NewEncoder(responseC)
+
 	// Start the actual command
 	err = c.c.Start(&cmd)
+
+	// Start a goroutine to spin and wait for the process to actual
+	// exit. When it does, report it back to caller...
+	go func() {
+		defer responseC.Close()
+
+		for !cmd.Exited {
+			time.Sleep(50 * time.Millisecond)
+		}
+
+		responseWriter.Encode(&CommandFinished{cmd.ExitStatus})
+	}()
+
 	return
 }
 
diff --git a/packer/rpc/communicator_test.go b/packer/rpc/communicator_test.go
index 7b7867d86..003c78473 100644
--- a/packer/rpc/communicator_test.go
+++ b/packer/rpc/communicator_test.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/rpc"
 	"testing"
+	"time"
 )
 
 type testCommunicator struct {
@@ -95,7 +96,19 @@ func TestCommunicatorRPC(t *testing.T) {
 	assert.Equal(data, "infoo\n", "should be correct stdin")
 
 	// Test that we can get the exit status properly
-	// TODO
+	c.startCmd.ExitStatus = 42
+	c.startCmd.Exited = true
+
+	for i := 0; i < 5; i++ {
+		if cmd.Exited {
+			assert.Equal(cmd.ExitStatus, 42, "should have proper exit status")
+			break
+		}
+
+		time.Sleep(50 * time.Millisecond)
+	}
+
+	assert.True(cmd.Exited, "should have exited")
 
 	// Test that we can upload things
 	uploadR, uploadW := io.Pipe()