From cb1e0cbabfde0e8b92bdce88ce104f513e6cc47d Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 11 May 2013 10:46:17 -0700
Subject: [PATCH] packer/plugin: Support hooks

---
 packer/plugin/hook.go        | 66 ++++++++++++++++++++++++++++++++++++
 packer/plugin/hook_test.go   | 26 ++++++++++++++
 packer/plugin/plugin.go      | 12 +++++++
 packer/plugin/plugin_test.go |  2 ++
 4 files changed, 106 insertions(+)
 create mode 100644 packer/plugin/hook.go
 create mode 100644 packer/plugin/hook_test.go

diff --git a/packer/plugin/hook.go b/packer/plugin/hook.go
new file mode 100644
index 000000000..532b3eea8
--- /dev/null
+++ b/packer/plugin/hook.go
@@ -0,0 +1,66 @@
+package plugin
+
+import (
+	"github.com/mitchellh/packer/packer"
+	packrpc "github.com/mitchellh/packer/packer/rpc"
+	"log"
+	"net/rpc"
+	"os/exec"
+)
+
+type cmdHook struct {
+	hook packer.Hook
+	client  *client
+}
+
+func (c *cmdHook) Run(name string, data interface{}, ui packer.Ui) {
+	defer func() {
+		r := recover()
+		c.checkExit(r, nil)
+	}()
+
+	c.hook.Run(name, data, ui)
+}
+
+func (c *cmdHook) checkExit(p interface{}, cb func()) {
+	if c.client.Exited() {
+		cb()
+	} else if p != nil {
+		log.Panic(p)
+	}
+}
+
+// Returns a valid packer.Hook 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 Hook(cmd *exec.Cmd) (result packer.Hook, 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 = &cmdHook{
+		packrpc.Hook(client),
+		cmdClient,
+	}
+
+	return
+}
diff --git a/packer/plugin/hook_test.go b/packer/plugin/hook_test.go
new file mode 100644
index 000000000..1f3ce25c1
--- /dev/null
+++ b/packer/plugin/hook_test.go
@@ -0,0 +1,26 @@
+package plugin
+
+import (
+	"cgl.tideland.biz/asserts"
+	"github.com/mitchellh/packer/packer"
+	"os/exec"
+	"testing"
+)
+
+type helperHook byte
+
+func (helperHook) Run(string, interface{}, packer.Ui) {}
+
+func TestHook_NoExist(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	_, err := Hook(exec.Command("i-should-never-ever-ever-exist"))
+	assert.NotNil(err, "should have an error")
+}
+
+func TestHook_Good(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	_, err := Hook(helperProcess("hook"))
+	assert.Nil(err, "should start hook properly")
+}
diff --git a/packer/plugin/plugin.go b/packer/plugin/plugin.go
index 175f62bed..d0c94d1ee 100644
--- a/packer/plugin/plugin.go
+++ b/packer/plugin/plugin.go
@@ -98,3 +98,15 @@ func ServeCommand(command packer.Command) {
 		log.Panic(err)
 	}
 }
+
+// Serves a hook from a plugin.
+func ServeHook(hook packer.Hook) {
+	log.Println("Preparing to serve a hook plugin...")
+
+	server := rpc.NewServer()
+	packrpc.RegisterHook(server, hook)
+
+	if err := serve(server); err != nil {
+		log.Panic(err)
+	}
+}
diff --git a/packer/plugin/plugin_test.go b/packer/plugin/plugin_test.go
index 30dbeec5c..b1fb4bf78 100644
--- a/packer/plugin/plugin_test.go
+++ b/packer/plugin/plugin_test.go
@@ -52,6 +52,8 @@ func TestHelperProcess(*testing.T) {
 		ServeBuilder(new(helperBuilder))
 	case "command":
 		ServeCommand(new(helperCommand))
+	case "hook":
+		ServeHook(new(helperHook))
 	case "invalid-rpc-address":
 		fmt.Println("lolinvalid")
 	case "start-timeout":