From 8d4ba1fc2bbef53245549fc3c7a524bc7003cf2b Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 10 Dec 2013 13:23:07 -0800
Subject: [PATCH] packer/rpc: complete command

---
 packer/rpc/command.go             | 23 +++++++++++++----------
 packer/rpc/command_test.go        |  5 -----
 packer/rpc/post_processor.go      |  4 ----
 packer/rpc/post_processor_test.go |  2 +-
 4 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/packer/rpc/command.go b/packer/rpc/command.go
index 945b8d58b..6d70fdace 100644
--- a/packer/rpc/command.go
+++ b/packer/rpc/command.go
@@ -20,8 +20,8 @@ type CommandServer struct {
 }
 
 type CommandRunArgs struct {
-	RPCAddress string
 	Args       []string
+	StreamId   uint32
 }
 
 type CommandSynopsisArgs byte
@@ -40,11 +40,15 @@ func (c *command) Help() (result string) {
 }
 
 func (c *command) Run(env packer.Environment, args []string) (result int) {
-	// Create and start the server for the Environment
-	server := rpc.NewServer()
-	RegisterEnvironment(server, env)
+	nextId := c.mux.NextId()
+	server := NewServerWithMux(c.mux, nextId)
+	server.RegisterEnvironment(env)
+	go server.Serve()
 
-	rpcArgs := &CommandRunArgs{serveSingleConn(server), args}
+	rpcArgs := &CommandRunArgs{
+		Args: args,
+		StreamId: nextId,
+	}
 	err := c.client.Call("Command.Run", rpcArgs, &result)
 	if err != nil {
 		panic(err)
@@ -68,14 +72,13 @@ func (c *CommandServer) Help(args *interface{}, reply *string) error {
 }
 
 func (c *CommandServer) Run(args *CommandRunArgs, reply *int) error {
-	client, err := rpcDial(args.RPCAddress)
+	client, err := NewClientWithMux(c.mux, args.StreamId)
 	if err != nil {
-		return err
+		return NewBasicError(err)
 	}
+	defer client.Close()
 
-	env := &Environment{client: client}
-
-	*reply = c.command.Run(env, args.Args)
+	*reply = c.command.Run(client.Environment(), args.Args)
 	return nil
 }
 
diff --git a/packer/rpc/command_test.go b/packer/rpc/command_test.go
index 976b70e6f..a6c3c9592 100644
--- a/packer/rpc/command_test.go
+++ b/packer/rpc/command_test.go
@@ -59,11 +59,6 @@ func TestRPCCommand(t *testing.T) {
 		t.Fatal("runEnv should not be nil")
 	}
 
-	command.runEnv.Ui()
-	if !testEnv.uiCalled {
-		t.Fatal("ui should be called")
-	}
-
 	// Test Synopsis
 	synopsis := commClient.Synopsis()
 	if synopsis != "foo" {
diff --git a/packer/rpc/post_processor.go b/packer/rpc/post_processor.go
index 2a0dfe62a..ad8fd0b9d 100644
--- a/packer/rpc/post_processor.go
+++ b/packer/rpc/post_processor.go
@@ -30,10 +30,6 @@ type PostProcessorProcessResponse struct {
 	StreamId uint32
 }
 
-func PostProcessor(client *rpc.Client) *postProcessor {
-	return &postProcessor{client: client}
-}
-
 func (p *postProcessor) Configure(raw ...interface{}) (err error) {
 	args := &PostProcessorConfigureArgs{Configs: raw}
 	if cerr := p.client.Call("PostProcessor.Configure", args, &err); cerr != nil {
diff --git a/packer/rpc/post_processor_test.go b/packer/rpc/post_processor_test.go
index d469ee4e0..dabd8c03a 100644
--- a/packer/rpc/post_processor_test.go
+++ b/packer/rpc/post_processor_test.go
@@ -83,7 +83,7 @@ func TestPostProcessorRPC(t *testing.T) {
 
 func TestPostProcessor_Implements(t *testing.T) {
 	var raw interface{}
-	raw = PostProcessor(nil)
+	raw = new(postProcessor)
 	if _, ok := raw.(packer.PostProcessor); !ok {
 		t.Fatal("not a postprocessor")
 	}