Skip to content

Update newpool to add IAM auth beforeConnect #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mamundsen-specter wants to merge 20 commits into from
Closed

Update newpool to add IAM auth beforeConnect #32

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
401e9a7
update newpool
mamundsen-specter Jan 13, 2026
02903c3
update logging and variable names for coderabbit
mamundsen-specter Feb 10, 2026
e3bffd3
move DatabaseConfiguration from Bloodhound repo to dawgs
mamundsen-specter Feb 10, 2026
ef65e44
Merge branch 'main' into bi-1348-main
mamundsen-specter Feb 10, 2026
b66b204
update go.mod go.sum
mamundsen-specter Feb 10, 2026
3e01779
urlencode password, update region call
mamundsen-specter Feb 10, 2026
7d3b4dd
chore: clean up composition in config.go - make BeforeConnect in pool…
zinic Feb 11, 2026
7c43cf9
Merge branch 'main' of github.com:SpecterOps/DAWGS into bi-1348-main
mamundsen-specter Feb 24, 2026
96678e6
Merge branch 'main' of github.com:SpecterOps/DAWGS into bi-1348-main
mamundsen-specter Mar 20, 2026
b473e89
merge with main
mamundsen-specter Mar 27, 2026
b34cd9c
Merge branch 'main' of github.com:SpecterOps/DAWGS into bi-1348-main
mamundsen-specter Apr 2, 2026
837bc43
fixup test NewPool calls
mamundsen-specter Apr 2, 2026
ad81ebb
move CNAME lookup to pool creation
mamundsen-specter Apr 8, 2026
5d44793
Merge branch 'main' of github.com:SpecterOps/DAWGS into bi-1348-main
mamundsen-specter Apr 8, 2026
e3fcf9e
move CNAME lookup to before poolCfg creation
mamundsen-specter Apr 8, 2026
d8fedab
move back cname handling, quiet logs
mamundsen-specter Apr 9, 2026
c00ee81
update comments
mamundsen-specter Apr 17, 2026
eb97c09
Merge branch 'main' of github.com:SpecterOps/DAWGS into bi-1348-main
mamundsen-specter Apr 17, 2026
7265ebc
add log for postgres connection string fallback
mamundsen-specter Apr 22, 2026
7129bea
extract endpoint lookup and run before NewPool
mamundsen-specter Apr 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 14 additions & 5 deletions drivers/pg/pg.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (

"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/specterops/bloodhound/cmd/api/src/config"
"github.com/specterops/dawgs"
"github.com/specterops/dawgs/cypher/models/pgsql"
"github.com/specterops/dawgs/graph"
Expand Down Expand Up @@ -50,15 +51,12 @@ func afterPooledConnectionRelease(conn *pgx.Conn) bool {
return true
}

func NewPool(connectionString string) (*pgxpool.Pool, error) {
if connectionString == "" {
return nil, fmt.Errorf("graph connection requires a connection url to be set")
}
func NewPool(cfg config.DatabaseConfiguration) (*pgxpool.Pool, error) {

poolCtx, done := context.WithTimeout(context.Background(), poolInitConnectionTimeout)
defer done()

poolCfg, err := pgxpool.ParseConfig(connectionString)
poolCfg, err := pgxpool.ParseConfig(cfg.PostgreSQLConnectionString())
if err != nil {
return nil, err
}
Expand All @@ -73,6 +71,17 @@ func NewPool(connectionString string) (*pgxpool.Pool, error) {
poolCfg.AfterConnect = afterPooledConnectionEstablished
poolCfg.AfterRelease = afterPooledConnectionRelease

poolCfg.BeforeConnect = func(ctx context.Context, connCfg *pgx.ConnConfig) error {
slog.Info("RDS credential beforeConnect(), creating new IAM credentials")
refreshConnectionString := cfg.PostgreSQLConnectionString()
newPoolCfg, err := pgxpool.ParseConfig(refreshConnectionString)
if err != nil {
return err
}
Comment on lines +73 to +84
Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Apr 22, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

# Find the relevant files
find . -name "pg.go" -path "*/drivers/*" | head -20

Repository: SpecterOps/DAWGS

Length of output: 78

🏁 Script executed:

# Find config.go in drivers
find . -name "config.go" -path "*/drivers/*" | head -20

Repository: SpecterOps/DAWGS

Length of output: 79

🏁 Script executed:

# Check overall structure
ls -la drivers/ 2>/dev/null || echo "drivers not at root"

Repository: SpecterOps/DAWGS

Length of output: 402

🏁 Script executed:

cat -n drivers/pg/pg.go

Repository: SpecterOps/DAWGS

Length of output: 3664

🏁 Script executed:

cat -n drivers/config.go

Repository: SpecterOps/DAWGS

Length of output: 3021

Cache static connection parameters during pool setup; refresh only auth tokens in BeforeConnect.

PostgreSQLConnectionString() loads AWS config (line 35 in drivers/config.go) and performs DNS lookup via net.LookupCNAME() (line 40) on every invocation. When called inside BeforeConnect for each new connection, this amplifies latency and failure surface under pool churn. Extract host, port, and database once during NewPool(), then have BeforeConnect use only auth.BuildAuthToken() to generate the fresh time-limited credential.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@drivers/pg/pg.go` around lines 73 - 83, The BeforeConnect currently calls
cfg.PostgreSQLConnectionString() (which does AWS config and DNS lookups) for
every new connection; instead, in NewPool() parse the connection string once and
cache static parameters (host, port, user, database) into the pool configuration
(e.g., set poolCfg.ConnConfig.Host/Port/Database/User from the parsed result),
and then change poolCfg.BeforeConnect to only refresh the RDS IAM password by
calling auth.BuildAuthToken() (or otherwise generate the short-lived credential)
and assign that to connCfg.Password; keep using cfg.EnableRDSIAMAuth to guard
this logic and reference the cached parsed result rather than calling
PostgreSQLConnectionString() inside BeforeConnect.

Copy link
Copy Markdown
Contributor Author

@mamundsen-specter mamundsen-specter Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added LookupEndpoint() to run before newpool creation and if endpoint field is empty. This should reduce the number of lookups, while still allowing IAM PostgreSQLConnectionString() calls outside of pool context.

Copy link
Copy Markdown

@coderabbitai coderabbitai Bot Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
cat -n drivers/config.go

Length of output: 3340

@mamundsen-specter, the fix looks good — the DNS concern from the original comment is now fully addressed. Here's the breakdown:

CNAME lookup — ✅ resolved
cfg.Endpoint = cfg.LookupEndpoint() is called once before the closure is created. Since cfg is a value type and Endpoint is pre-populated, RDSIAMAuthConnectionString() will take the if s.Endpoint != "" branch (line 41–42 of config.go) on every subsequent BeforeConnect call, skipping the net.LookupCNAME entirely.

awsConfig.LoadDefaultConfig — still called per connection
RDSIAMAuthConnectionString() still invokes awsConfig.LoadDefaultConfig(context.TODO()) on every BeforeConnect. This is the remaining overhead, but it's likely intentional and acceptable — reloading the AWS config on each call allows credential rotation (IMDS, IAM role assumption, env-var changes) to be picked up without restarting the pool. Caching it would risk using stale credentials for token generation.

The fix strikes a reasonable balance: the expensive/durable operation (CNAME lookup) is cached, while the security-sensitive operation (credential + token refresh) stays live per-connection.

✏️ Learnings added 
Learnt from: mamundsen-specter
Repo: SpecterOps/DAWGS PR: 32
File: drivers/pg/pg.go:73-84
Timestamp: 2026-04-22T17:09:34.255Z
Learning: In `drivers/config.go` (SpecterOps/DAWGS), `DatabaseConfiguration.Endpoint` is pre-populated via `cfg.Endpoint = cfg.LookupEndpoint()` in `NewPool()` before the `BeforeConnect` closure is set. This caches the CNAME lookup. `RDSIAMAuthConnectionString()` checks `if s.Endpoint != ""` to skip re-doing the DNS lookup. `awsConfig.LoadDefaultConfig` is intentionally still called on every `BeforeConnect` invocation to support credential rotation.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

🧠 Learnings used 
Learnt from: mamundsen-specter
Repo: SpecterOps/DAWGS PR: 32
File: drivers/config.go:61-63
Timestamp: 2026-04-22T16:52:55.097Z
Learning: In `drivers/config.go` (SpecterOps/DAWGS), `DatabaseConfiguration.PostgreSQLConnectionString()` intentionally ignores the `Connection` field when `EnableRDSIAMAuth` is true. IAM auth always takes precedence; the `Connection` string is not parsed or used as a fallback input for IAM token generation in that code path.

Learnt from: ykaiboussiSO
Repo: SpecterOps/DAWGS PR: 53
File: .github/workflows/go-test.yml:75-76
Timestamp: 2026-03-31T18:22:51.105Z
Learning: In `ops/ops_test.go` (SpecterOps/DAWGS), the `dbSetup` helper function reads `Driver_CONNECTION_STRING` and supports both `bolt://` and `neo4j://` schemes for Neo4j — it rewrites `bolt://` to `neo4j://` before calling the driver. Do not flag `bolt://` connection strings in test matrix configs as incorrect; the scheme normalization is intentional and handled in `dbSetup`.

Learnt from: bsheth711
Repo: SpecterOps/DAWGS PR: 38
File: drivers/pg/batch.go:114-125
Timestamp: 2026-03-17T23:49:14.270Z
Learning: In the DAWGS codebase, when using pgx.CopyFrom in batch-like scenarios, prefer prebuilding the full payload (as in LargeNodeUpdateRows) rather than streaming rows lazily if it reduces per-row encoding overhead during COPY. This trade-off is acceptable since memory usage is O(N) and does not change the overall memory footprint. Apply this guideline to similar code paths in the drivers/pg directory where a full batch can be preconstructed efficiently; avoid streaming-only patterns if prebuilding provides measurable performance gains, and back the choice with simple benchmarks or profiling.

connCfg.Password = newPoolCfg.ConnConfig.Password
return nil
}

pool, err := pgxpool.NewWithConfig(poolCtx, poolCfg)
if err != nil {
return nil, err
Expand Down
1 change: 1 addition & 0 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ module github.com/specterops/dawgs
go 1.24.0

require (
github.com/specterops/BloodHound v0.0.0-580f0c7776b5ab0f354c45337cc286bce7c4dffd
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
cuelang.org/go v0.15.1
github.com/RoaringBitmap/roaring/v2 v2.14.4
github.com/antlr4-go/antlr/v4 v4.13.1
Expand Down