From 37372bac934fc8f06cb1803a2648356c1d38d7da Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Thu, 23 May 2013 21:37:16 -0700
Subject: [PATCH] packer/plugin: Support provisioners

---
 packer/plugin/plugin.go           | 12 +++++
 packer/plugin/plugin_test.go      |  2 +
 packer/plugin/provisioner.go      | 75 +++++++++++++++++++++++++++++++
 packer/plugin/provisioner_test.go | 28 ++++++++++++
 4 files changed, 117 insertions(+)
 create mode 100644 packer/plugin/provisioner.go
 create mode 100644 packer/plugin/provisioner_test.go

diff --git a/packer/plugin/plugin.go b/packer/plugin/plugin.go
index d0c94d1ee..44113bf5e 100644
--- a/packer/plugin/plugin.go
+++ b/packer/plugin/plugin.go
@@ -110,3 +110,15 @@ func ServeHook(hook packer.Hook) {
 		log.Panic(err)
 	}
 }
+
+// Serves a provisioner from a plugin.
+func ServeProvisioner(p packer.Provisioner) {
+	log.Println("Preparing to serve a provisioner plugin...")
+
+	server := rpc.NewServer()
+	packrpc.RegisterProvisioner(server, p)
+
+	if err := serve(server); err != nil {
+		log.Panic(err)
+	}
+}
diff --git a/packer/plugin/plugin_test.go b/packer/plugin/plugin_test.go
index b1fb4bf78..939e2e8fd 100644
--- a/packer/plugin/plugin_test.go
+++ b/packer/plugin/plugin_test.go
@@ -56,6 +56,8 @@ func TestHelperProcess(*testing.T) {
 		ServeHook(new(helperHook))
 	case "invalid-rpc-address":
 		fmt.Println("lolinvalid")
+	case "provisioner":
+		ServeProvisioner(new(helperProvisioner))
 	case "start-timeout":
 		time.Sleep(1 * time.Minute)
 		os.Exit(1)
diff --git a/packer/plugin/provisioner.go b/packer/plugin/provisioner.go
new file mode 100644
index 000000000..7391543b4
--- /dev/null
+++ b/packer/plugin/provisioner.go
@@ -0,0 +1,75 @@
+package plugin
+
+import (
+	"github.com/mitchellh/packer/packer"
+	packrpc "github.com/mitchellh/packer/packer/rpc"
+	"log"
+	"net/rpc"
+	"os/exec"
+)
+
+type cmdProvisioner struct {
+	p   packer.Provisioner
+	client *client
+}
+
+func (c *cmdProvisioner) Prepare(config interface{}, ui packer.Ui) {
+	defer func() {
+		r := recover()
+		c.checkExit(r, nil)
+	}()
+
+	c.p.Prepare(config, ui)
+}
+
+func (c *cmdProvisioner) Provision(ui packer.Ui, comm packer.Communicator) {
+	defer func() {
+		r := recover()
+		c.checkExit(r, nil)
+	}()
+
+	c.p.Provision(ui, comm)
+}
+
+func (c *cmdProvisioner) checkExit(p interface{}, cb func()) {
+	if c.client.Exited() {
+		cb()
+	} else if p != nil {
+		log.Panic(p)
+	}
+}
+
+// Returns a valid packer.Provisioner where the hook is executed via RPC
+// to a plugin that is within a subprocess.
+//
+// This method will start the given exec.Cmd, which should point to
+// the plugin binary to execute. Some configuration will be done to
+// the command, such as overriding Stdout and some environmental variables.
+//
+// This function guarantees the subprocess will end in a timely manner.
+func Provisioner(cmd *exec.Cmd) (result packer.Provisioner, err error) {
+	cmdClient := NewManagedClient(cmd)
+	address, err := cmdClient.Start()
+	if err != nil {
+		return
+	}
+
+	defer func() {
+		// Make sure the command is properly killed in the case of an error
+		if err != nil {
+			cmdClient.Kill()
+		}
+	}()
+
+	client, err := rpc.Dial("tcp", address)
+	if err != nil {
+		return
+	}
+
+	result = &cmdProvisioner{
+		packrpc.Provisioner(client),
+		cmdClient,
+	}
+
+	return
+}
diff --git a/packer/plugin/provisioner_test.go b/packer/plugin/provisioner_test.go
new file mode 100644
index 000000000..39d489360
--- /dev/null
+++ b/packer/plugin/provisioner_test.go
@@ -0,0 +1,28 @@
+package plugin
+
+import (
+	"cgl.tideland.biz/asserts"
+	"github.com/mitchellh/packer/packer"
+	"os/exec"
+	"testing"
+)
+
+type helperProvisioner byte
+
+func (helperProvisioner) Prepare(interface{}, packer.Ui) {}
+
+func (helperProvisioner) Provision(packer.Ui, packer.Communicator) {}
+
+func TestProvisioner_NoExist(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	_, err := Provisioner(exec.Command("i-should-never-ever-ever-exist"))
+	assert.NotNil(err, "should have an error")
+}
+
+func TestProvisioner_Good(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	_, err := Provisioner(helperProcess("provisioner"))
+	assert.Nil(err, "should start provisioner properly")
+}