feature: multiplexing support for database plugins

builtin/logical/database/backend.go

@@ -110,6 +110,7 @@ func Backend(conf *logical.BackendConfig) *databaseBackend {
 }
 
 type databaseBackend struct {
+	// connections holds configured database connections by config name
 	connections map[string]*dbPluginInstance
 	logger      log.Logger
 

builtin/logical/database/path_config_connection.go

@@ -329,6 +329,8 @@ func (b *databaseBackend) connectionWriteHandler() framework.OperationFunc {
 		}
 		config.ConnectionDetails = initResp.Config
 
+		b.Logger().Debug("created database object", "name", name, "plugin_name", config.PluginName)
+
 		b.Lock()
 		defer b.Unlock()
 
@@ -365,6 +367,9 @@ func (b *databaseBackend) connectionWriteHandler() framework.OperationFunc {
 				"Vault (or the sdk if using a custom plugin) to gain password policy support", config.PluginName))
 		}
 
+		if len(resp.Warnings) == 0 {
+			return nil, nil
+		}
 		return resp, nil
 	}
 }

changelog/14033.txt

@@ -0,0 +1,3 @@
+```release-note:feature
+**Database plugin multiplexing**: manage multiple database connections with a single plugin process
+```

helper/builtinplugins/registry.go

@@ -150,8 +150,8 @@ type registry struct {
 	logicalBackends    map[string]logical.Factory
 }
 
-// Get returns the BuiltinFactory func for a particular backend plugin
-// from the plugins map.
+// Get returns the Factory func for a particular backend plugin from the
+// plugins map.
 func (r *registry) Get(name string, pluginType consts.PluginType) (func() (interface{}, error), bool) {
 	switch pluginType {
 	case consts.PluginTypeCredential:

plugins/database/postgresql/postgresql.go

@@ -48,7 +48,6 @@ var (
 	singleQuotedPhrases = regexp.MustCompile(`('.*?')`)
 )
 
-// New implements builtinplugins.BuiltinFactory
 func New() (interface{}, error) {
 	db := new()
 	// Wrap the plugin with middleware to sanitize errors

sdk/database/dbplugin/v5/grpc_database_plugin.go

@@ -12,14 +12,17 @@ import (
 // a plugin and host. If the handshake fails, a user friendly error is shown.
 // This prevents users from executing bad plugins or executing a plugin
 // directory. It is a UX feature, not a security feature.
-var handshakeConfig = plugin.HandshakeConfig{
-	ProtocolVersion:  5,
+var HandshakeConfig = plugin.HandshakeConfig{
 	MagicCookieKey:   "VAULT_DATABASE_PLUGIN",
 	MagicCookieValue: "926a0820-aea2-be28-51d6-83cdf00e8edb",
 }
 
+// Factory is the factory function to create a dbplugin Database.
+type Factory func() (interface{}, error)
+
 type GRPCDatabasePlugin struct {
-	Impl Database
+	FactoryFunc Factory
+	Impl        Database
 
 	// Embeding this will disable the netRPC protocol
 	plugin.NetRPCUnsupportedPlugin
@@ -31,7 +34,19 @@ var (
 )
 
 func (d GRPCDatabasePlugin) GRPCServer(_ *plugin.GRPCBroker, s *grpc.Server) error {
-	proto.RegisterDatabaseServer(s, gRPCServer{impl: d.Impl})
+	var server gRPCServer
+
+	if d.Impl != nil {
+		server = gRPCServer{singleImpl: d.Impl}
+	} else {
+		// multiplexing is supported
+		server = gRPCServer{
+			factoryFunc: d.FactoryFunc,
+			instances:   make(map[string]Database),
+		}
+	}
+
+	proto.RegisterDatabaseServer(s, &server)
 	return nil
 }
 

sdk/database/dbplugin/v5/grpc_server.go

@@ -3,32 +3,121 @@ package dbplugin
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/golang/protobuf/ptypes"
 	"github.com/hashicorp/vault/sdk/database/dbplugin/v5/proto"
+	"github.com/hashicorp/vault/sdk/helper/pluginutil"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
-var _ proto.DatabaseServer = gRPCServer{}
+var _ proto.DatabaseServer = &gRPCServer{}
 
 type gRPCServer struct {
 	proto.UnimplementedDatabaseServer
 
-	impl Database
+	// holds the non-multiplexed Database
+	// when this is set the plugin does not support multiplexing
+	singleImpl Database
+
+	// instances holds the multiplexed Databases
+	instances   map[string]Database
+	factoryFunc func() (interface{}, error)
+
+	sync.RWMutex
+}
+
+func getMultiplexIDFromContext(ctx context.Context) (string, error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return "", fmt.Errorf("missing plugin multiplexing metadata")
+	}
+
+	multiplexIDs := md[pluginutil.MultiplexingCtxKey]
+	if len(multiplexIDs) != 1 {
+		return "", fmt.Errorf("unexpected number of IDs in metadata: (%d)", len(multiplexIDs))
+	}
+
+	multiplexID := multiplexIDs[0]
+	if multiplexID == "" {
+		return "", fmt.Errorf("empty multiplex ID in metadata")
+	}
+
+	return multiplexID, nil
+}
+
+func (g *gRPCServer) getOrCreateDatabase(ctx context.Context) (Database, error) {
+	g.Lock()
+	defer g.Unlock()
+
+	if g.singleImpl != nil {
+		return g.singleImpl, nil
+	}
+
+	id, err := getMultiplexIDFromContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if db, ok := g.instances[id]; ok {
+		return db, nil
+	}
+
+	db, err := g.factoryFunc()
+	if err != nil {
+		return nil, err
+	}
+
+	database := db.(Database)
+	g.instances[id] = database
+
+	return database, nil
+}
+
+// getDatabaseInternal returns the database but does not hold a lock
+func (g *gRPCServer) getDatabaseInternal(ctx context.Context) (Database, error) {
+	if g.singleImpl != nil {
+		return g.singleImpl, nil
+	}
+
+	id, err := getMultiplexIDFromContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if db, ok := g.instances[id]; ok {
+		return db, nil
+	}
+
+	return nil, fmt.Errorf("no database instance found")
+}
+
+// getDatabase holds a read lock and returns the database
+func (g *gRPCServer) getDatabase(ctx context.Context) (Database, error) {
+	g.RLock()
+	impl, err := g.getDatabaseInternal(ctx)
+	g.RUnlock()
+	return impl, err
 }
 
 // Initialize the database plugin
-func (g gRPCServer) Initialize(ctx context.Context, request *proto.InitializeRequest) (*proto.InitializeResponse, error) {
+func (g *gRPCServer) Initialize(ctx context.Context, request *proto.InitializeRequest) (*proto.InitializeResponse, error) {
+	impl, err := g.getOrCreateDatabase(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	rawConfig := structToMap(request.ConfigData)
 
 	dbReq := InitializeRequest{
 		Config:           rawConfig,
 		VerifyConnection: request.VerifyConnection,
 	}
 
-	dbResp, err := g.impl.Initialize(ctx, dbReq)
+	dbResp, err := impl.Initialize(ctx, dbReq)
 	if err != nil {
 		return &proto.InitializeResponse{}, status.Errorf(codes.Internal, "failed to initialize: %s", err)
 	}
@@ -45,7 +134,7 @@ func (g gRPCServer) Initialize(ctx context.Context, request *proto.InitializeReq
 	return resp, nil
 }
 
-func (g gRPCServer) NewUser(ctx context.Context, req *proto.NewUserRequest) (*proto.NewUserResponse, error) {
+func (g *gRPCServer) NewUser(ctx context.Context, req *proto.NewUserRequest) (*proto.NewUserResponse, error) {
 	if req.GetUsernameConfig() == nil {
 		return &proto.NewUserResponse{}, status.Errorf(codes.InvalidArgument, "missing username config")
 	}
@@ -60,6 +149,11 @@ func (g gRPCServer) NewUser(ctx context.Context, req *proto.NewUserRequest) (*pr
 		expiration = exp
 	}
 
+	impl, err := g.getDatabase(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	dbReq := NewUserRequest{
 		UsernameConfig: UsernameMetadata{
 			DisplayName: req.GetUsernameConfig().GetDisplayName(),
@@ -71,7 +165,7 @@ func (g gRPCServer) NewUser(ctx context.Context, req *proto.NewUserRequest) (*pr
 		RollbackStatements: getStatementsFromProto(req.GetRollbackStatements()),
 	}
 
-	dbResp, err := g.impl.NewUser(ctx, dbReq)
+	dbResp, err := impl.NewUser(ctx, dbReq)
 	if err != nil {
 		return &proto.NewUserResponse{}, status.Errorf(codes.Internal, "unable to create new user: %s", err)
 	}
@@ -82,7 +176,7 @@ func (g gRPCServer) NewUser(ctx context.Context, req *proto.NewUserRequest) (*pr
 	return resp, nil
 }
 
-func (g gRPCServer) UpdateUser(ctx context.Context, req *proto.UpdateUserRequest) (*proto.UpdateUserResponse, error) {
+func (g *gRPCServer) UpdateUser(ctx context.Context, req *proto.UpdateUserRequest) (*proto.UpdateUserResponse, error) {
 	if req.GetUsername() == "" {
 		return &proto.UpdateUserResponse{}, status.Errorf(codes.InvalidArgument, "no username provided")
 	}
@@ -92,7 +186,12 @@ func (g gRPCServer) UpdateUser(ctx context.Context, req *proto.UpdateUserRequest
 		return &proto.UpdateUserResponse{}, status.Errorf(codes.InvalidArgument, err.Error())
 	}
 
-	_, err = g.impl.UpdateUser(ctx, dbReq)
+	impl, err := g.getDatabase(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = impl.UpdateUser(ctx, dbReq)
 	if err != nil {
 		return &proto.UpdateUserResponse{}, status.Errorf(codes.Internal, "unable to update user: %s", err)
 	}
@@ -144,7 +243,7 @@ func hasChange(dbReq UpdateUserRequest) bool {
 	return false
 }
 
-func (g gRPCServer) DeleteUser(ctx context.Context, req *proto.DeleteUserRequest) (*proto.DeleteUserResponse, error) {
+func (g *gRPCServer) DeleteUser(ctx context.Context, req *proto.DeleteUserRequest) (*proto.DeleteUserResponse, error) {
 	if req.GetUsername() == "" {
 		return &proto.DeleteUserResponse{}, status.Errorf(codes.InvalidArgument, "no username provided")
 	}
@@ -153,15 +252,25 @@ func (g gRPCServer) DeleteUser(ctx context.Context, req *proto.DeleteUserRequest
 		Statements: getStatementsFromProto(req.GetStatements()),
 	}
 
-	_, err := g.impl.DeleteUser(ctx, dbReq)
+	impl, err := g.getDatabase(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = impl.DeleteUser(ctx, dbReq)
 	if err != nil {
 		return &proto.DeleteUserResponse{}, status.Errorf(codes.Internal, "unable to delete user: %s", err)
 	}
 	return &proto.DeleteUserResponse{}, nil
 }
 
-func (g gRPCServer) Type(ctx context.Context, _ *proto.Empty) (*proto.TypeResponse, error) {
-	t, err := g.impl.Type()
+func (g *gRPCServer) Type(ctx context.Context, _ *proto.Empty) (*proto.TypeResponse, error) {
+	impl, err := g.getOrCreateDatabase(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	t, err := impl.Type()
 	if err != nil {
 		return &proto.TypeResponse{}, status.Errorf(codes.Internal, "unable to retrieve type: %s", err)
 	}
@@ -172,11 +281,29 @@ func (g gRPCServer) Type(ctx context.Context, _ *proto.Empty) (*proto.TypeRespon
 	return resp, nil
 }
 
-func (g gRPCServer) Close(ctx context.Context, _ *proto.Empty) (*proto.Empty, error) {
-	err := g.impl.Close()
+func (g *gRPCServer) Close(ctx context.Context, _ *proto.Empty) (*proto.Empty, error) {
+	g.Lock()
+	defer g.Unlock()
+
+	impl, err := g.getDatabaseInternal(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	err = impl.Close()
 	if err != nil {
 		return &proto.Empty{}, status.Errorf(codes.Internal, "unable to close database plugin: %s", err)
 	}
+
+	if g.singleImpl == nil {
+		// only cleanup instances map when multiplexing is supported
+		id, err := getMultiplexIDFromContext(ctx)
+		if err != nil {
+			return nil, err
+		}
+		delete(g.instances, id)
+	}
+
 	return &proto.Empty{}, nil
 }