github.com/hashicorp/vault/sdk@v0.13.0/database/dbplugin/client.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package dbplugin
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"sync"
    10  
    11  	log "github.com/hashicorp/go-hclog"
    12  	plugin "github.com/hashicorp/go-plugin"
    13  	"github.com/hashicorp/vault/sdk/helper/pluginutil"
    14  )
    15  
    16  // DatabasePluginClient embeds a databasePluginRPCClient and wraps it's Close
    17  // method to also call Kill() on the plugin.Client.
    18  type DatabasePluginClient struct {
    19  	client *plugin.Client
    20  	sync.Mutex
    21  
    22  	Database
    23  }
    24  
    25  // This wraps the Close call and ensures we both close the database connection
    26  // and kill the plugin.
    27  func (dc *DatabasePluginClient) Close() error {
    28  	err := dc.Database.Close()
    29  	dc.client.Kill()
    30  
    31  	return err
    32  }
    33  
    34  // NewPluginClient returns a databaseRPCClient with a connection to a running
    35  // plugin. The client is wrapped in a DatabasePluginClient object to ensure the
    36  // plugin is killed on call of Close().
    37  func NewPluginClient(ctx context.Context, sys pluginutil.RunnerUtil, pluginRunner *pluginutil.PluginRunner, logger log.Logger, isMetadataMode bool) (Database, error) {
    38  	// pluginSets is the map of plugins we can dispense.
    39  	pluginSets := map[int]plugin.PluginSet{
    40  		// Version 3 used to supports both protocols. We want to keep it around
    41  		// since it's possible old plugins built against this version will still
    42  		// work with gRPC. There is currently no difference between version 3
    43  		// and version 4.
    44  		3: {
    45  			"database": new(GRPCDatabasePlugin),
    46  		},
    47  		// Version 4 only supports gRPC
    48  		4: {
    49  			"database": new(GRPCDatabasePlugin),
    50  		},
    51  	}
    52  
    53  	var client *plugin.Client
    54  	var err error
    55  	if isMetadataMode {
    56  		client, err = pluginRunner.RunMetadataMode(ctx, sys, pluginSets, handshakeConfig, []string{}, logger)
    57  	} else {
    58  		client, err = pluginRunner.Run(ctx, sys, pluginSets, handshakeConfig, []string{}, logger)
    59  	}
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  
    64  	// Connect via RPC
    65  	rpcClient, err := client.Client()
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  
    70  	// Request the plugin
    71  	raw, err := rpcClient.Dispense("database")
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  
    76  	// We should have a database type now. This feels like a normal interface
    77  	// implementation but is in fact over an RPC connection.
    78  	var db Database
    79  	switch raw.(type) {
    80  	case *gRPCClient:
    81  		db = raw.(*gRPCClient)
    82  	default:
    83  		return nil, errors.New("unsupported client type")
    84  	}
    85  
    86  	// Wrap RPC implementation in DatabasePluginClient
    87  	return &DatabasePluginClient{
    88  		client:   client,
    89  		Database: db,
    90  	}, nil
    91  }