database: Avoid race condition in connection creation

builtin/logical/database/backend.go

@@ -161,8 +161,9 @@ func (b *databaseBackend) collectPluginInstanceGaugeValues(context.Context) ([]m
 
 type databaseBackend struct {
 	// connections holds configured database connections by config name
-	connections *syncmap.SyncMap[string, *dbPluginInstance]
-	logger      log.Logger
+	createConnectionLock sync.Mutex
+	connections          *syncmap.SyncMap[string, *dbPluginInstance]
+	logger               log.Logger
 
 	*framework.Backend
 	// credRotationQueue is an in-memory priority queue used to track Static Roles
@@ -291,11 +292,23 @@ func (b *databaseBackend) GetConnection(ctx context.Context, s logical.Storage,
 }
 
 func (b *databaseBackend) GetConnectionWithConfig(ctx context.Context, name string, config *DatabaseConfig) (*dbPluginInstance, error) {
+	// fast path, reuse the existing connection
 	dbi := b.connections.Get(name)
 	if dbi != nil {
 		return dbi, nil
 	}
 
+	// slow path, create a new connection
+	// if we don't lock the rest of the operation, there is a race condition for multiple callers of this function
+	b.createConnectionLock.Lock()
+	defer b.createConnectionLock.Unlock()
+
+	// check again in case we lost the race
+	dbi = b.connections.Get(name)
+	if dbi != nil {
+		return dbi, nil
+	}
+
 	id, err := uuid.GenerateUUID()
 	if err != nil {
 		return nil, err
@@ -332,14 +345,17 @@ func (b *databaseBackend) GetConnectionWithConfig(ctx context.Context, name stri
 		name:                 name,
 		runningPluginVersion: pluginVersion,
 	}
-	oldConn := b.connections.Put(name, dbi)
-	if oldConn != nil {
-		err := oldConn.Close()
+	conn, ok := b.connections.PutIfEmpty(name, dbi)
+	if !ok {
+		// this is a bug
+		b.Logger().Error("Error: there was a race condition adding to the database connection map")
+		// There was already an existing connection, so we will use that and close our new one to avoid a race condition.
+		err := dbi.Close()
 		if err != nil {
-			b.Logger().Warn("Error closing database connection", "error", err)
+			b.Logger().Warn("Error closing new database connection", "error", err)
 		}
 	}
-	return dbi, nil
+	return conn, nil
 }
 
 // ClearConnection closes the database connection and

builtin/logical/database/backend_get_test.go

@@ -0,0 +1,106 @@
+package database
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/hashicorp/vault/sdk/helper/consts"
+	"github.com/hashicorp/vault/sdk/helper/pluginutil"
+	"github.com/hashicorp/vault/sdk/logical"
+	"github.com/hashicorp/vault/sdk/queue"
+)
+
+func newSystemViewWrapper(view logical.SystemView) logical.SystemView {
+	return &systemViewWrapper{
+		view,
+	}
+}
+
+type systemViewWrapper struct {
+	logical.SystemView
+}
+
+var _ logical.ExtendedSystemView = (*systemViewWrapper)(nil)
+
+func (s *systemViewWrapper) RequestWellKnownRedirect(ctx context.Context, src, dest string) error {
+	panic("nope")
+}
+
+func (s *systemViewWrapper) DeregisterWellKnownRedirect(ctx context.Context, src string) bool {
+	panic("nope")
+}
+
+func (s *systemViewWrapper) Auditor() logical.Auditor {
+	panic("nope")
+}
+
+func (s *systemViewWrapper) ForwardGenericRequest(ctx context.Context, request *logical.Request) (*logical.Response, error) {
+	panic("nope")
+}
+
+func (s *systemViewWrapper) APILockShouldBlockRequest() (bool, error) {
+	panic("nope")
+}
+
+func (s *systemViewWrapper) GetPinnedPluginVersion(ctx context.Context, pluginType consts.PluginType, pluginName string) (*pluginutil.PinnedVersion, error) {
+	return nil, pluginutil.ErrPinnedVersionNotFound
+}
+
+func (s *systemViewWrapper) LookupPluginVersion(ctx context.Context, pluginName string, pluginType consts.PluginType, version string) (*pluginutil.PluginRunner, error) {
+	return &pluginutil.PluginRunner{
+		Name:           "mockv5",
+		Type:           consts.PluginTypeDatabase,
+		Builtin:        true,
+		BuiltinFactory: New,
+	}, nil
+}
+
+func getDbBackend(t *testing.T) (*databaseBackend, logical.Storage) {
+	t.Helper()
+	config := logical.TestBackendConfig()
+	config.System = newSystemViewWrapper(config.System)
+	config.StorageView = &logical.InmemStorage{}
+	// Create and init the backend ourselves instead of using a Factory because
+	// the factory function kicks off threads that cause racy tests.
+	b := Backend(config)
+	if err := b.Setup(context.Background(), config); err != nil {
+		t.Fatal(err)
+	}
+	b.schedule = &TestSchedule{}
+	b.credRotationQueue = queue.New()
+	b.populateQueue(context.Background(), config.StorageView)
+
+	return b, config.StorageView
+}
+
+// TestGetConnectionRaceCondition checks that GetConnection always returns the same instance, even when asked
+// by multiple goroutines in parallel.
+func TestGetConnectionRaceCondition(t *testing.T) {
+	ctx := context.Background()
+	b, s := getDbBackend(t)
+	defer b.Cleanup(ctx)
+	configureDBMount(t, s)
+
+	goroutines := 16
+
+	wg := sync.WaitGroup{}
+	wg.Add(goroutines)
+	dbis := make([]*dbPluginInstance, goroutines)
+	errs := make([]error, goroutines)
+	for i := 0; i < goroutines; i++ {
+		go func(i int) {
+			dbis[i], errs[i] = b.GetConnection(ctx, s, "mockv5")
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+	for i := 0; i < goroutines; i++ {
+		if errs[i] != nil {
+			t.Fatal(errs[i])
+		}
+		if dbis[0] != dbis[i] {
+			t.Fatal("Error: database instances did not match")
+		}
+	}
+}

builtin/logical/database/mockv5.go

@@ -51,6 +51,9 @@ func (m MockDatabaseV5) Initialize(ctx context.Context, req v5.InitializeRequest
 		"req", req)
 
 	config := req.Config
+	if config == nil {
+		config = map[string]interface{}{}
+	}
 	config["from-plugin"] = "this value is from the plugin itself"
 
 	resp := v5.InitializeResponse{

helper/syncmap/syncmap.go

@@ -62,6 +62,20 @@ func (m *SyncMap[K, V]) Put(k K, v V) V {
 	return oldV
 }
 
+// PutIfEmpty adds the given key-value pair to the map only if there is no value already in it,
+// and returns the new value and true if so.
+// If there is already a value, it returns the existing value and false.
+func (m *SyncMap[K, V]) PutIfEmpty(k K, v V) (V, bool) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	oldV, ok := m.data[k]
+	if ok {
+		return oldV, false
+	}
+	m.data[k] = v
+	return v, true
+}
+
 // Clear deletes all entries from the map, and returns the previous map.
 func (m *SyncMap[K, V]) Clear() map[K]V {
 	m.lock.Lock()