feat: update db pooling to handle iam rds auth, move DatabaseConfiguration to DAWGS - BI-1306

cmd/api/src/api/tools/pg.go

@@ -448,7 +448,7 @@ func (s *PGMigrator) MigrationStatus(response http.ResponseWriter, request *http
 }
 
 func (s *PGMigrator) OpenPostgresGraphConnection() (graph.Database, error) {
-	if pool, err := pg.NewPool(s.Cfg.Database.PostgreSQLConnectionString()); err != nil {
+	if pool, err := pg.NewPool(s.Cfg.Database); err != nil {
 		return nil, err
 	} else {
 		return dawgs.Open(s.ServerCtx, pg.DriverName, dawgs.Config{

cmd/api/src/bootstrap/util.go

@@ -98,7 +98,7 @@ func ConnectGraph(ctx context.Context, cfg config.Configuration) (*graph.Databas
 		slog.InfoContext(ctx, "Connecting to graph using PostgreSQL")
 		connectionString = cfg.Database.PostgreSQLConnectionString()
 
-		pool, err = pg.NewPool(connectionString)
+		pool, err = pg.NewPool(cfg.Database)
 		if err != nil {
 			return nil, err
 		}

cmd/api/src/cmd/dawgs-harness/main.go

@@ -30,9 +30,11 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/specterops/bloodhound/cmd/api/src/cmd/dawgs-harness/tests"
+	"github.com/specterops/bloodhound/cmd/api/src/config"
 	"github.com/specterops/bloodhound/packages/go/bhlog"
 	schema "github.com/specterops/bloodhound/packages/go/graphschema"
 	"github.com/specterops/dawgs"
+	"github.com/specterops/dawgs/drivers"
 	"github.com/specterops/dawgs/drivers/neo4j"
 	"github.com/specterops/dawgs/drivers/pg"
 	"github.com/specterops/dawgs/graph"
@@ -44,14 +46,14 @@ func fatalf(format string, args ...any) {
 	os.Exit(1)
 }
 
-func RunTestSuite(ctx context.Context, connectionStr, driverName string) tests.TestSuite {
+func RunTestSuite(ctx context.Context, connectionStr, driverName string, cfg drivers.DatabaseConfiguration) tests.TestSuite {
 	var (
 		pool *pgxpool.Pool
 		err  error
 	)
 
 	if driverName == pg.DriverName {
-		pool, err = pg.NewPool(connectionStr)
+		pool, err = pg.NewPool(cfg)
 		if err != nil {
 			fatalf("Failed creating a new pgxpool: %s", err)
 		}
@@ -141,10 +143,18 @@ func main() {
 
 	bhlog.ConfigureDefaultText(os.Stdout)
 
+	cfg, err := config.NewDefaultConfiguration()
+	if err != nil {
+		fmt.Errorf("failed to create default configuration: %w", err)
+	}
+	cfg.Neo4J.Connection = neo4jConnectionStr
+	cfg.Database.Connection = pgConnectionStr
+
 	switch testType {
 	case "both":
+
 		n4jTestSuite := execSuite(neo4j.DriverName, func() tests.TestSuite {
-			return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName)
+			return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Neo4J)
 		})
 
 		fmt.Println()
@@ -153,20 +163,20 @@ func main() {
 		time.Sleep(time.Second * 3)
 
 		pgTestSuite := execSuite(pg.DriverName, func() tests.TestSuite {
-			return RunTestSuite(ctx, pgConnectionStr, pg.DriverName)
+			return RunTestSuite(ctx, pgConnectionStr, pg.DriverName, cfg.Database)
 		})
 		fmt.Println()
 
 		OutputTestSuiteDeltas(pgTestSuite, n4jTestSuite)
 
 	case "postgres":
 		execSuite(pg.DriverName, func() tests.TestSuite {
-			return RunTestSuite(ctx, pgConnectionStr, pg.DriverName)
+			return RunTestSuite(ctx, pgConnectionStr, pg.DriverName, cfg.Database)
 		})
 
 	case "neo4j":
 		execSuite(neo4j.DriverName, func() tests.TestSuite {
-			return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName)
+			return RunTestSuite(ctx, neo4jConnectionStr, neo4j.DriverName, cfg.Neo4J)
 		})
 	}
 }

cmd/api/src/config/config.go

@@ -31,6 +31,7 @@ import (
 	"github.com/specterops/bloodhound/cmd/api/src/serde"
 	"github.com/specterops/bloodhound/packages/go/bhlog/attr"
 	"github.com/specterops/bloodhound/packages/go/crypto"
+	Dawgs "github.com/specterops/dawgs/drivers"
 )
 
 const (
@@ -50,15 +51,6 @@ func (s TLSConfiguration) Enabled() bool {
 	return s.CertFile != "" && s.KeyFile != ""
 }
 
-type DatabaseConfiguration struct {
-	Connection            string `json:"connection"`
-	Address               string `json:"addr"`
-	Database              string `json:"database"`
-	Username              string `json:"username"`
-	Secret                string `json:"secret"`
-	MaxConcurrentSessions int    `json:"max_concurrent_sessions"`
-}
-
 type CollectorManifest struct {
 	Latest   string             `json:"latest"`
 	Versions []CollectorVersion `json:"versions"`
@@ -72,22 +64,6 @@ type CollectorVersion struct {
 
 type CollectorManifests map[string]CollectorManifest
 
-func (s DatabaseConfiguration) PostgreSQLConnectionString() string {
-	if s.Connection == "" {
-		return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database)
-	}
-
-	return s.Connection
-}
-
-func (s DatabaseConfiguration) Neo4jConnectionString() string {
-	if s.Connection == "" {
-		return fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database)
-	}
-
-	return s.Connection
-}
-
 type CryptoConfiguration struct {
 	JWT    JWTConfiguration    `json:"jwt"`
 	Argon2 Argon2Configuration `json:"argon2"`
@@ -135,39 +111,40 @@ type DefaultAdminConfiguration struct {
 }
 
 type Configuration struct {
-	Version                         int                       `json:"version"`
-	BindAddress                     string                    `json:"bind_addr"`
-	SlowQueryThreshold              int64                     `json:"slow_query_threshold"`
-	MaxGraphQueryCacheSize          int                       `json:"max_graphdb_cache_size"`
-	MaxAPICacheSize                 int                       `json:"max_api_cache_size"`
-	MetricsPort                     string                    `json:"metrics_port"`
-	RootURL                         serde.URL                 `json:"root_url"`
-	WorkDir                         string                    `json:"work_dir"`
-	LogLevel                        string                    `json:"log_level"`
-	LogPath                         string                    `json:"log_path"`
-	TLS                             TLSConfiguration          `json:"tls"`
-	GraphDriver                     string                    `json:"graph_driver"`
-	Database                        DatabaseConfiguration     `json:"database"`
-	Neo4J                           DatabaseConfiguration     `json:"neo4j"`
-	Crypto                          CryptoConfiguration       `json:"crypto"`
-	SAML                            SAMLConfiguration         `json:"saml"`
-	DefaultAdmin                    DefaultAdminConfiguration `json:"default_admin"`
-	CollectorsBucketURL             serde.URL                 `json:"collectors_bucket_url"`
-	CollectorsBasePath              string                    `json:"collectors_base_path"`
-	DatapipeInterval                int                       `json:"datapipe_interval"`
-	EnableStartupWaitPeriod         bool                      `json:"enable_startup_wait_period"`
-	EnableAPILogging                bool                      `json:"enable_api_logging"`
-	EnableCypherMutations           bool                      `json:"enable_cypher_mutations"`
-	DisableAnalysis                 bool                      `json:"disable_analysis"`
-	DisableCypherComplexityLimit    bool                      `json:"disable_cypher_complexity_limit"`
-	DisableIngest                   bool                      `json:"disable_ingest"`
-	DisableMigrations               bool                      `json:"disable_migrations"`
-	GraphQueryMemoryLimit           uint16                    `json:"graph_query_memory_limit"`
-	EnableTextLogger                bool                      `json:"enable_text_logger"`
-	RecreateDefaultAdmin            bool                      `json:"recreate_default_admin"`
-	EnableUserAnalytics             bool                      `json:"enable_user_analytics"`
-	ForceDownloadEmbeddedCollectors bool                      `json:"force_download_embedded_collectors"`
-	EnableAuditLogStdout            bool                      `json:"enable_audit_log_stdout"`
+	Version                         int                         `json:"version"`
+	BindAddress                     string                      `json:"bind_addr"`
+	SlowQueryThreshold              int64                       `json:"slow_query_threshold"`
+	MaxGraphQueryCacheSize          int                         `json:"max_graphdb_cache_size"`
+	MaxAPICacheSize                 int                         `json:"max_api_cache_size"`
+	MetricsPort                     string                      `json:"metrics_port"`
+	RootURL                         serde.URL                   `json:"root_url"`
+	WorkDir                         string                      `json:"work_dir"`
+	LogLevel                        string                      `json:"log_level"`
+	LogPath                         string                      `json:"log_path"`
+	TLS                             TLSConfiguration            `json:"tls"`
+	GraphDriver                     string                      `json:"graph_driver"`
+	Database                        Dawgs.DatabaseConfiguration `json:"database"`
+	Neo4J                           Dawgs.DatabaseConfiguration `json:"neo4j"`
+	Crypto                          CryptoConfiguration         `json:"crypto"`
+	SAML                            SAMLConfiguration           `json:"saml"`
+	DefaultAdmin                    DefaultAdminConfiguration   `json:"default_admin"`
+	CollectorsBucketURL             serde.URL                   `json:"collectors_bucket_url"`
+	CollectorsBasePath              string                      `json:"collectors_base_path"`
+	DatapipeInterval                int                         `json:"datapipe_interval"`
+	EnableStartupWaitPeriod         bool                        `json:"enable_startup_wait_period"`
+	EnableAPILogging                bool                        `json:"enable_api_logging"`
+	EnableCypherMutations           bool                        `json:"enable_cypher_mutations"`
+	DisableAnalysis                 bool                        `json:"disable_analysis"`
+	DisableCypherComplexityLimit    bool                        `json:"disable_cypher_complexity_limit"`
+	DisableIngest                   bool                        `json:"disable_ingest"`
+	DisableMigrations               bool                        `json:"disable_migrations"`
+	DisableTimeoutLimit             bool                        `json:"disable_timeout_limit"`
+	GraphQueryMemoryLimit           uint16                      `json:"graph_query_memory_limit"`
+	EnableTextLogger                bool                        `json:"enable_text_logger"`
+	RecreateDefaultAdmin            bool                        `json:"recreate_default_admin"`
+	EnableUserAnalytics             bool                        `json:"enable_user_analytics"`
+	ForceDownloadEmbeddedCollectors bool                        `json:"force_download_embedded_collectors"`
+	EnableAuditLogStdout            bool                        `json:"enable_audit_log_stdout"`
 }
 
 func (s Configuration) TempDirectory() string {

cmd/api/src/config/default.go

@@ -19,6 +19,7 @@ package config
 import (
 	"fmt"
 
+	Dawgs "github.com/specterops/dawgs/drivers"
 	"github.com/specterops/dawgs/drivers/neo4j"
 
 	"github.com/specterops/bloodhound/cmd/api/src/serde"
@@ -73,10 +74,10 @@ func NewDefaultConfiguration() (Configuration, error) {
 			TLS:                             TLSConfiguration{},
 			SAML:                            SAMLConfiguration{},
 			GraphDriver:                     neo4j.DriverName, // Default to PG as the graph driver
-			Database: DatabaseConfiguration{
+			Database: Dawgs.DatabaseConfiguration{
 				MaxConcurrentSessions: 10,
 			},
-			Neo4J: DatabaseConfiguration{
+			Neo4J: Dawgs.DatabaseConfiguration{
 				MaxConcurrentSessions: 10,
 			},
 			Crypto: CryptoConfiguration{

cmd/api/src/daemons/changelog/changelog_e2e/main.go

@@ -79,7 +79,14 @@ func newHarness() *Harness {
 		os.Exit(1)
 	}
 
-	pool, err := pg.NewPool(connStr)
+	cfg, err := config.NewDefaultConfiguration()
+	if err != nil {
+		slog.Error("Error creating new default configuration")
+		os.Exit(1)
+	}
+	cfg.Database.Connection = connStr
+
+	pool, err := pg.NewPool(cfg.Database)
 	if err != nil {
 		slog.Error("Failed to connect", attr.Error(err))
 		os.Exit(1)
@@ -91,7 +98,7 @@ func newHarness() *Harness {
 		os.Exit(1)
 	}
 
-	gormDB, err := database.OpenDatabase(connStr)
+	gormDB, err := database.OpenDatabase(cfg.Database)
 	if err != nil {
 		slog.Error("Failed to open", attr.Error(err))
 		os.Exit(1)

cmd/api/src/daemons/changelog/ingestion_integration_test.go

@@ -97,8 +97,14 @@ func setupIntegrationTest(t *testing.T) IntegrationTestSuite {
 		connConf = pgtestdb.Custom(t, getPostgresConfig(t), pgtestdb.NoopMigrator{})
 	)
 
+	cfg, err := config.NewDefaultConfiguration()
+	if err != nil {
+		t.Logf("Error creating new default configuration: %v", err)
+	}
+	cfg.Database.Connection = connConf.URL()
+
 	// Create connection pool
-	pool, err := pg.NewPool(connConf.URL())
+	pool, err := pg.NewPool(cfg.Database)
 	require.NoError(t, err)
 
 	// Open graph database
@@ -113,7 +119,7 @@ func setupIntegrationTest(t *testing.T) IntegrationTestSuite {
 	err = graphDB.AssertSchema(ctx, schema())
 	require.NoError(t, err)
 
-	gormDB, err := database.OpenDatabase(connConf.URL())
+	gormDB, err := database.OpenDatabase(cfg.Database)
 	require.NoError(t, err)
 
 	db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver(), cfg)

cmd/api/src/daemons/datapipe/datapipe_integration_test.go

@@ -68,11 +68,17 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
 		workDir  = t.TempDir()
 	)
 
+	cfg, err := config.NewDefaultConfiguration()
+	if err != nil {
+		t.Logf("Error creating new default configuration: %v", err)
+	}
+	cfg.Database.Connection = connConf.URL()
+
 	//#region Setup for dbs
-	pool, err := pg.NewPool(connConf.URL())
+	pool, err := pg.NewPool(cfg.Database)
 	require.NoError(t, err)
 
-	gormDB, err := database.OpenDatabase(connConf.URL())
+	gormDB, err := database.OpenDatabase(cfg.Database)
 	require.NoError(t, err)
 
 	db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver(), config.Configuration{})
@@ -107,9 +113,7 @@ func setupIntegrationTestSuite(t *testing.T, fixturesPath string) IntegrationTes
 	err = os.Mkdir(path.Join(workDir, "tmp"), 0755)
 	require.NoError(t, err)
 
-	cfg := config.Configuration{
-		WorkDir: workDir,
-	}
+	cfg.WorkDir = workDir
 
 	cl := changelog.NewChangelog(graphDB, db, changelog.DefaultOptions())
 

cmd/api/src/database/database_integration_test.go

@@ -56,9 +56,15 @@ func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite {
 		cfg      = config.Configuration{}
 	)
 
+	cfg, err := config.NewDefaultConfiguration()
+	if err != nil {
+		t.Logf("Error creating new default configuration: %v", err)
+	}
+	cfg.Database.Connection = connConf.URL()
+
 	// #region Setup for dbs
 
-	gormDB, err := database.OpenDatabase(connConf.URL())
+	gormDB, err := database.OpenDatabase(cfg.Database)
 	require.NoError(t, err)
 
 	db := database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver(), cfg)

cmd/api/src/database/db.go

@@ -27,13 +27,16 @@ import (
 	"time"
 
 	"github.com/gofrs/uuid"
+	"github.com/jackc/pgx/v5/stdlib"
 	"github.com/specterops/bloodhound/cmd/api/src/auth"
 	"github.com/specterops/bloodhound/cmd/api/src/config"
 	"github.com/specterops/bloodhound/cmd/api/src/database/migration"
 	"github.com/specterops/bloodhound/cmd/api/src/model"
 	"github.com/specterops/bloodhound/cmd/api/src/model/appcfg"
 	"github.com/specterops/bloodhound/cmd/api/src/services/upload"
 	"github.com/specterops/bloodhound/packages/go/bhlog/attr"
+	"github.com/specterops/dawgs/drivers"
+	"github.com/specterops/dawgs/drivers/pg"
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 )
@@ -227,15 +230,22 @@ func (s *BloodhoundDB) Transaction(ctx context.Context, fn func(tx *BloodhoundDB
 	}, opts...)
 }
 
-func OpenDatabase(connection string) (*gorm.DB, error) {
+func OpenDatabase(cfg drivers.DatabaseConfiguration) (*gorm.DB, error) {
 	gormConfig := &gorm.Config{
 		Logger: &GormLogAdapter{
 			SlowQueryErrorThreshold: time.Second * 30,
 			SlowQueryWarnThreshold:  time.Second * 10,
 		},
 	}
 
-	if db, err := gorm.Open(postgres.Open(connection), gormConfig); err != nil {
+	pool, err := pg.NewPool(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	dbPool := stdlib.OpenDBFromPool(pool)
+
+	if db, err := gorm.Open(postgres.New(postgres.Config{Conn: dbPool}), gormConfig); err != nil {
 		return nil, err
 	} else {
 		return db, nil

cmd/api/src/database/migration/extensions_integration_test.go

@@ -55,7 +55,13 @@ func setupIntegrationTestSuite(t *testing.T) IntegrationTestSuite {
 
 	// #region Setup for dbs
 
-	gormDB, err = database.OpenDatabase(connConf.URL())
+	cfg, err := config.NewDefaultConfiguration()
+	if err != nil {
+		t.Logf("Error creating new default configuration: %v", err)
+	}
+	cfg.Database.Connection = connConf.URL()
+
+	gormDB, err = database.OpenDatabase(cfg.Database)
 	require.NoError(t, err)
 
 	db = database.NewBloodhoundDB(gormDB, auth.NewIdentityResolver(), config.Configuration{})