forked from hashicorp/vault
-
Notifications
You must be signed in to change notification settings - Fork 0
/
client.go
75 lines (63 loc) · 1.98 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
package dbplugin
import (
"context"
"errors"
"sync"
log "github.com/hashicorp/go-hclog"
"github.com/hashicorp/go-plugin"
"github.com/hashicorp/vault/helper/pluginutil"
)
// DatabasePluginClient embeds a databasePluginRPCClient and wraps it's Close
// method to also call Kill() on the plugin.Client.
type DatabasePluginClient struct {
client *plugin.Client
sync.Mutex
Database
}
// This wraps the Close call and ensures we both close the database connection
// and kill the plugin.
func (dc *DatabasePluginClient) Close() error {
err := dc.Database.Close()
dc.client.Kill()
return err
}
// newPluginClient returns a databaseRPCClient with a connection to a running
// plugin. The client is wrapped in a DatabasePluginClient object to ensure the
// plugin is killed on call of Close().
func newPluginClient(ctx context.Context, sys pluginutil.RunnerUtil, pluginRunner *pluginutil.PluginRunner, logger log.Logger) (Database, error) {
// pluginMap is the map of plugins we can dispense.
var pluginMap = map[string]plugin.Plugin{
"database": new(DatabasePlugin),
}
client, err := pluginRunner.Run(ctx, sys, pluginMap, handshakeConfig, []string{}, logger)
if err != nil {
return nil, err
}
// Connect via RPC
rpcClient, err := client.Client()
if err != nil {
return nil, err
}
// Request the plugin
raw, err := rpcClient.Dispense("database")
if err != nil {
return nil, err
}
// We should have a database type now. This feels like a normal interface
// implementation but is in fact over an RPC connection.
var db Database
switch raw.(type) {
case *gRPCClient:
db = raw.(*gRPCClient)
case *databasePluginRPCClient:
logger.Warn("plugin is using deprecated net RPC transport, recompile plugin to upgrade to gRPC", "plugin", pluginRunner.Name)
db = raw.(*databasePluginRPCClient)
default:
return nil, errors.New("unsupported client type")
}
// Wrap RPC implementation in DatabasePluginClient
return &DatabasePluginClient{
client: client,
Database: db,
}, nil
}