diff options
Diffstat (limited to 'vendor/github.com/hashicorp/go-plugin/rpc_client.go')
-rw-r--r-- | vendor/github.com/hashicorp/go-plugin/rpc_client.go | 123 |
1 files changed, 123 insertions, 0 deletions
diff --git a/vendor/github.com/hashicorp/go-plugin/rpc_client.go b/vendor/github.com/hashicorp/go-plugin/rpc_client.go new file mode 100644 index 0000000..29f9bf0 --- /dev/null +++ b/vendor/github.com/hashicorp/go-plugin/rpc_client.go | |||
@@ -0,0 +1,123 @@ | |||
1 | package plugin | ||
2 | |||
3 | import ( | ||
4 | "fmt" | ||
5 | "io" | ||
6 | "net" | ||
7 | "net/rpc" | ||
8 | |||
9 | "github.com/hashicorp/yamux" | ||
10 | ) | ||
11 | |||
12 | // RPCClient connects to an RPCServer over net/rpc to dispense plugin types. | ||
13 | type RPCClient struct { | ||
14 | broker *MuxBroker | ||
15 | control *rpc.Client | ||
16 | plugins map[string]Plugin | ||
17 | |||
18 | // These are the streams used for the various stdout/err overrides | ||
19 | stdout, stderr net.Conn | ||
20 | } | ||
21 | |||
22 | // NewRPCClient creates a client from an already-open connection-like value. | ||
23 | // Dial is typically used instead. | ||
24 | func NewRPCClient(conn io.ReadWriteCloser, plugins map[string]Plugin) (*RPCClient, error) { | ||
25 | // Create the yamux client so we can multiplex | ||
26 | mux, err := yamux.Client(conn, nil) | ||
27 | if err != nil { | ||
28 | conn.Close() | ||
29 | return nil, err | ||
30 | } | ||
31 | |||
32 | // Connect to the control stream. | ||
33 | control, err := mux.Open() | ||
34 | if err != nil { | ||
35 | mux.Close() | ||
36 | return nil, err | ||
37 | } | ||
38 | |||
39 | // Connect stdout, stderr streams | ||
40 | stdstream := make([]net.Conn, 2) | ||
41 | for i, _ := range stdstream { | ||
42 | stdstream[i], err = mux.Open() | ||
43 | if err != nil { | ||
44 | mux.Close() | ||
45 | return nil, err | ||
46 | } | ||
47 | } | ||
48 | |||
49 | // Create the broker and start it up | ||
50 | broker := newMuxBroker(mux) | ||
51 | go broker.Run() | ||
52 | |||
53 | // Build the client using our broker and control channel. | ||
54 | return &RPCClient{ | ||
55 | broker: broker, | ||
56 | control: rpc.NewClient(control), | ||
57 | plugins: plugins, | ||
58 | stdout: stdstream[0], | ||
59 | stderr: stdstream[1], | ||
60 | }, nil | ||
61 | } | ||
62 | |||
63 | // SyncStreams should be called to enable syncing of stdout, | ||
64 | // stderr with the plugin. | ||
65 | // | ||
66 | // This will return immediately and the syncing will continue to happen | ||
67 | // in the background. You do not need to launch this in a goroutine itself. | ||
68 | // | ||
69 | // This should never be called multiple times. | ||
70 | func (c *RPCClient) SyncStreams(stdout io.Writer, stderr io.Writer) error { | ||
71 | go copyStream("stdout", stdout, c.stdout) | ||
72 | go copyStream("stderr", stderr, c.stderr) | ||
73 | return nil | ||
74 | } | ||
75 | |||
76 | // Close closes the connection. The client is no longer usable after this | ||
77 | // is called. | ||
78 | func (c *RPCClient) Close() error { | ||
79 | // Call the control channel and ask it to gracefully exit. If this | ||
80 | // errors, then we save it so that we always return an error but we | ||
81 | // want to try to close the other channels anyways. | ||
82 | var empty struct{} | ||
83 | returnErr := c.control.Call("Control.Quit", true, &empty) | ||
84 | |||
85 | // Close the other streams we have | ||
86 | if err := c.control.Close(); err != nil { | ||
87 | return err | ||
88 | } | ||
89 | if err := c.stdout.Close(); err != nil { | ||
90 | return err | ||
91 | } | ||
92 | if err := c.stderr.Close(); err != nil { | ||
93 | return err | ||
94 | } | ||
95 | if err := c.broker.Close(); err != nil { | ||
96 | return err | ||
97 | } | ||
98 | |||
99 | // Return back the error we got from Control.Quit. This is very important | ||
100 | // since we MUST return non-nil error if this fails so that Client.Kill | ||
101 | // will properly try a process.Kill. | ||
102 | return returnErr | ||
103 | } | ||
104 | |||
105 | func (c *RPCClient) Dispense(name string) (interface{}, error) { | ||
106 | p, ok := c.plugins[name] | ||
107 | if !ok { | ||
108 | return nil, fmt.Errorf("unknown plugin type: %s", name) | ||
109 | } | ||
110 | |||
111 | var id uint32 | ||
112 | if err := c.control.Call( | ||
113 | "Dispenser.Dispense", name, &id); err != nil { | ||
114 | return nil, err | ||
115 | } | ||
116 | |||
117 | conn, err := c.broker.Dial(id) | ||
118 | if err != nil { | ||
119 | return nil, err | ||
120 | } | ||
121 | |||
122 | return p.Client(c.broker, rpc.NewClient(conn)) | ||
123 | } | ||