From bc6db2bd1b7f7270adf4e2bf06411122df6d2b26 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 28 Sep 2014 10:55:12 -0700
Subject: [PATCH] rpc: ResourceProvisioner support on client/server

---
 rpc/client.go      | 18 ++++++++++++++++++
 rpc/client_test.go | 34 ++++++++++++++++++++++++++++++++++
 rpc/rpc_test.go    |  6 ++++++
 rpc/server.go      | 33 ++++++++++++++++++++++++++++++---
 4 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/rpc/client.go b/rpc/client.go
index 259be4d76..a42b8f569 100644
--- a/rpc/client.go
+++ b/rpc/client.go
@@ -81,3 +81,21 @@ func (c *Client) ResourceProvider() (terraform.ResourceProvider, error) {
 		Name:   "ResourceProvider",
 	}, nil
 }
+
+func (c *Client) ResourceProvisioner() (terraform.ResourceProvisioner, error) {
+	var id uint32
+	if err := c.control.Call(
+		"Dispenser.ResourceProvisioner", new(interface{}), &id); err != nil {
+		return nil, err
+	}
+
+	conn, err := c.broker.Dial(id)
+	if err != nil {
+		return nil, err
+	}
+
+	return &ResourceProvisioner{
+		Client: rpc.NewClient(conn),
+		Name:   "ResourceProvisioner",
+	}, nil
+}
diff --git a/rpc/client_test.go b/rpc/client_test.go
index 5bb9bc3b2..5e3feb619 100644
--- a/rpc/client_test.go
+++ b/rpc/client_test.go
@@ -40,3 +40,37 @@ func TestClient_ResourceProvider(t *testing.T) {
 		t.Fatalf("bad: %#v", e)
 	}
 }
+
+
+func TestClient_ResourceProvisioner(t *testing.T) {
+	clientConn, serverConn := testConn(t)
+
+	p := new(terraform.MockResourceProvisioner)
+	server := &Server{ProvisionerFunc: testProvisionerFixed(p)}
+	go server.ServeConn(serverConn)
+
+	client, err := NewClient(clientConn)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	defer client.Close()
+
+	provisioner, err := client.ResourceProvisioner()
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Apply
+	state := &terraform.InstanceState{}
+	conf := &terraform.ResourceConfig{}
+	err = provisioner.Apply(state, conf)
+	if !p.ApplyCalled {
+		t.Fatal("apply should be called")
+	}
+	if !reflect.DeepEqual(p.ApplyConfig, conf) {
+		t.Fatalf("bad: %#v", p.ApplyConfig)
+	}
+	if err != nil {
+		t.Fatalf("bad: %#v", err)
+	}
+}
diff --git a/rpc/rpc_test.go b/rpc/rpc_test.go
index 3006dfbec..d8550d84a 100644
--- a/rpc/rpc_test.go
+++ b/rpc/rpc_test.go
@@ -51,3 +51,9 @@ func testProviderFixed(p terraform.ResourceProvider) ProviderFunc {
 		return p
 	}
 }
+
+func testProvisionerFixed(p terraform.ResourceProvisioner) ProvisionerFunc {
+	return func() terraform.ResourceProvisioner {
+		return p
+	}
+}
diff --git a/rpc/server.go b/rpc/server.go
index 705e6e0fa..0ad92366c 100644
--- a/rpc/server.go
+++ b/rpc/server.go
@@ -13,13 +13,18 @@ import (
 // Server listens for network connections and then dispenses interface
 // implementations for Terraform over net/rpc.
 type Server struct {
-	ProviderFunc ProviderFunc
+	ProviderFunc    ProviderFunc
+	ProvisionerFunc ProvisionerFunc
 }
 
 // ProviderFunc creates terraform.ResourceProviders when they're requested
 // from the server.
 type ProviderFunc func() terraform.ResourceProvider
 
+// ProvisionerFunc creates terraform.ResourceProvisioners when they're requested
+// from the server.
+type ProvisionerFunc func() terraform.ResourceProvisioner
+
 // Accept accepts connections on a listener and serves requests for
 // each incoming connection. Accept blocks; the caller typically invokes
 // it in a go statement.
@@ -63,7 +68,8 @@ func (s *Server) ServeConn(conn io.ReadWriteCloser) {
 	// connection.
 	server := rpc.NewServer()
 	server.RegisterName("Dispenser", &dispenseServer{
-		ProviderFunc: s.ProviderFunc,
+		ProviderFunc:    s.ProviderFunc,
+		ProvisionerFunc: s.ProvisionerFunc,
 
 		broker: broker,
 	})
@@ -72,7 +78,8 @@ func (s *Server) ServeConn(conn io.ReadWriteCloser) {
 
 // dispenseServer dispenses variousinterface implementations for Terraform.
 type dispenseServer struct {
-	ProviderFunc ProviderFunc
+	ProviderFunc    ProviderFunc
+	ProvisionerFunc ProvisionerFunc
 
 	broker *muxBroker
 }
@@ -97,6 +104,26 @@ func (d *dispenseServer) ResourceProvider(
 	return nil
 }
 
+func (d *dispenseServer) ResourceProvisioner(
+	args interface{}, response *uint32) error {
+	id := d.broker.NextId()
+	*response = id
+
+	go func() {
+		conn, err := d.broker.Accept(id)
+		if err != nil {
+			log.Printf("[ERR] Plugin dispense: %s", err)
+			return
+		}
+
+		d.serve(conn, "ResourceProvisioner", &ResourceProvisionerServer{
+			Provisioner: d.ProvisionerFunc(),
+		})
+	}()
+
+	return nil
+}
+
 func (d *dispenseServer) serve(conn io.ReadWriteCloser, name string, v interface{}) {
 	server := rpc.NewServer()
 	if err := server.RegisterName(name, v); err != nil {