Update newpool to add IAM auth beforeConnect#32
Update newpool to add IAM auth beforeConnect#32mamundsen-specter wants to merge 20 commits intomainfrom
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughNew DatabaseConfiguration is introduced and passed into pg.NewPool; the pool construction now derives the DSN via cfg.PostgreSQLConnectionString(), and when RDS IAM auth is enabled the pool sets a BeforeConnect hook to refresh IAM-derived credentials before each connection. go.mod adds AWS SDK v2 rds/auth deps. Changes
Sequence DiagramsequenceDiagram
participant App as Application
participant Driver as PostgreSQL Driver
participant Config as DatabaseConfiguration
participant IAM as AWS RDS IAM
participant DB as PostgreSQL DB
App->>Driver: NewPool(cfg)
Driver->>Config: PostgreSQLConnectionString()
Driver->>Driver: Parse pool config (pgxpool.ParseConfig)
App->>Driver: Acquire connection
Driver->>Driver: BeforeConnect hook triggered
Driver->>Config: Re-parse connection info
alt EnableRDSIAMAuth
Driver->>IAM: Request RDS auth token
IAM-->>Driver: Return token/password
end
Driver->>DB: Establish connection using refreshed credentials
DB-->>Driver: Connection established
Driver-->>App: Return connection
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@drivers/pg/pg.go`:
- Around line 74-83: Fix the BeforeConnect hook: correct the log text typo
("credentional"→"credential"), remove the unnecessary fmt.Sprint and lower the
verbosity by using slog.Debug instead of slog.Info, avoid parsing the whole pool
config by using pgx.ParseConfig to extract the password (replace
pgxpool.ParseConfig), and rename the closure parameter from poolCfg to connCfg
to avoid shadowing the outer poolCfg variable; ensure you still call
cfg.PostgreSQLConnectionString() each time and assign connCfg.Password from the
parsed pgx config.
🧹 Nitpick comments (1)
drivers/pg/pg.go (1)
64-66: TODO: Make pool size configurable.The hardcoded
MinConns = 5andMaxConns = 50may not suit all environments (dev vs. production, small vs. large instances). Since the function already accepts a config object, these values could be sourced from it.Would you like me to open an issue to track making these configurable via
DatabaseConfiguration?
| poolCfg.BeforeConnect = func(ctx context.Context, poolCfg *pgx.ConnConfig) error { | ||
| slog.Info(fmt.Sprint("RDS credentional beforeConnect(), creating new IAM credentials")) | ||
| refreshConnectionString := cfg.PostgreSQLConnectionString() | ||
| newPoolCfg, err := pgxpool.ParseConfig(refreshConnectionString) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| poolCfg.Password = newPoolCfg.ConnConfig.Password | ||
| return nil | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
ast-grep --pattern $'func ($_ DatabaseConfiguration) PostgreSQLConnectionString() string {
$$$
}'Repository: SpecterOps/DAWGS
Length of output: 42
🏁 Script executed:
rg "PostgreSQLConnectionString" --type go -A 20Repository: SpecterOps/DAWGS
Length of output: 1994
🏁 Script executed:
rg "func NewPool" --type go -A 5Repository: SpecterOps/DAWGS
Length of output: 407
🏁 Script executed:
rg "func.*PostgreSQLConnectionString" --type go -A 15Repository: SpecterOps/DAWGS
Length of output: 42
🏁 Script executed:
rg "PostgreSQLConnectionString" --type go -B 2 -A 15 | head -100Repository: SpecterOps/DAWGS
Length of output: 1780
🏁 Script executed:
fd "config" --type d | grep -E "^[^/]*config" | head -10Repository: SpecterOps/DAWGS
Length of output: 42
🏁 Script executed:
head -30 drivers/pg/pg.goRepository: SpecterOps/DAWGS
Length of output: 908
🏁 Script executed:
rg "RDS|IAM|credential|token" --type go -i -B 2 -A 2 | head -50Repository: SpecterOps/DAWGS
Length of output: 2751
🏁 Script executed:
rg "credentional|credential" --type go -B 5 -A 10 drivers/pg/pg.goRepository: SpecterOps/DAWGS
Length of output: 662
Several issues in the BeforeConnect hook.
-
Typo:
"credentional"→"credential"(Line 75). -
Unnecessary
fmt.Sprint:fmt.Sprint("...")on a literal string is redundant — pass the string directly toslog.Info. -
Log level too noisy:
slog.Infofires on every new connection. Under load or pool churn this will be very chatty. Considerslog.Debug. -
Heavier parse than needed:
pgxpool.ParseConfigparses pool-level settings you don't use.pgx.ParseConfigis sufficient to extract the password. -
Shadowed variable name: The closure parameter
poolCfg *pgx.ConnConfigshadows the outerpoolCfg *pgxpool.Config. Renaming the parameter (e.g.,connCfg) would improve clarity. -
IAM refresh pattern:
cfg.PostgreSQLConnectionString()is correctly called on everyBeforeConnectto rotate expiring IAM tokens. This is the proper pattern for AWS RDS IAM authentication with time-limited credentials.
Proposed fix (items 1–5)
- poolCfg.BeforeConnect = func(ctx context.Context, poolCfg *pgx.ConnConfig) error {
- slog.Info(fmt.Sprint("RDS credentional beforeConnect(), creating new IAM credentials"))
- refreshConnectionString := cfg.PostgreSQLConnectionString()
- newPoolCfg, err := pgxpool.ParseConfig(refreshConnectionString)
- if err != nil {
- return err
- }
- poolCfg.Password = newPoolCfg.ConnConfig.Password
- return nil
- }
+ poolCfg.BeforeConnect = func(ctx context.Context, connCfg *pgx.ConnConfig) error {
+ slog.Debug("BeforeConnect: refreshing IAM credentials")
+ refreshConnStr := cfg.PostgreSQLConnectionString()
+ refreshCfg, err := pgx.ParseConfig(refreshConnStr)
+ if err != nil {
+ return fmt.Errorf("failed to parse refreshed connection string: %w", err)
+ }
+ connCfg.Password = refreshCfg.Password
+ return nil
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| poolCfg.BeforeConnect = func(ctx context.Context, poolCfg *pgx.ConnConfig) error { | |
| slog.Info(fmt.Sprint("RDS credentional beforeConnect(), creating new IAM credentials")) | |
| refreshConnectionString := cfg.PostgreSQLConnectionString() | |
| newPoolCfg, err := pgxpool.ParseConfig(refreshConnectionString) | |
| if err != nil { | |
| return err | |
| } | |
| poolCfg.Password = newPoolCfg.ConnConfig.Password | |
| return nil | |
| } | |
| poolCfg.BeforeConnect = func(ctx context.Context, connCfg *pgx.ConnConfig) error { | |
| slog.Debug("BeforeConnect: refreshing IAM credentials") | |
| refreshConnStr := cfg.PostgreSQLConnectionString() | |
| refreshCfg, err := pgx.ParseConfig(refreshConnStr) | |
| if err != nil { | |
| return fmt.Errorf("failed to parse refreshed connection string: %w", err) | |
| } | |
| connCfg.Password = refreshCfg.Password | |
| return nil | |
| } |
🤖 Prompt for AI Agents
In `@drivers/pg/pg.go` around lines 74 - 83, Fix the BeforeConnect hook: correct
the log text typo ("credentional"→"credential"), remove the unnecessary
fmt.Sprint and lower the verbosity by using slog.Debug instead of slog.Info,
avoid parsing the whole pool config by using pgx.ParseConfig to extract the
password (replace pgxpool.ParseConfig), and rename the closure parameter from
poolCfg to connCfg to avoid shadowing the outer poolCfg variable; ensure you
still call cfg.PostgreSQLConnectionString() each time and assign
connCfg.Password from the parsed pgx config.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@go.mod`:
- Line 6: The dawgs module currently imports BloodHound's config type causing a
circular dependency; change the PG driver API to avoid that by modifying NewPool
in drivers/pg/pg.go to accept a simple factory (e.g., connectionStringFn func()
string) or a minimal interface defined in dawgs (e.g., type DBConnStringProvider
interface { PostgreSQLConnectionString() string }) instead of the concrete
BloodHound type (DatabaseConfiguration/cfg), update the BeforeConnect closure to
call connectionStringFn() (or use the local interface) to obtain the DSN, and
remove the import of github.com/specterops/bloodhound/cmd/api/src/config so the
go.mod entry is no longer required.
|
coderabbit has showed me an error in my ways. relying on config from bloodhound has created a circular dependency. looking at moving DatabaseConfig to DAWGs to cut the loop |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@drivers/config.go`:
- Line 39: The code that builds the DB endpoint uses a hardcoded ":5432" when
creating dbinput from cname (dbinput := strings.TrimSuffix(cname, ".") +
":5432"), which breaks if Address already contains a port; change this to parse
the existing address/port and only add a default port when none is present.
Locate the dbinput construction (variable dbinput, cname) in drivers/config.go
and use net.SplitHostPort or net.ParseIP logic to detect an existing port on the
Address/cname and, if absent, append the default (5432) via net.JoinHostPort or
equivalent; alternatively expose a configurable port field on the config and
prefer that value when present. Ensure the final dbinput is a valid host:port
string in all cases.
- Line 41: Replace the hardcoded region string in the auth.BuildAuthToken call
with the resolved AWS region from the loaded AWS config or the
DatabaseConfiguration (e.g., use cfg.Region or the awsConfig.Region value loaded
earlier) so deployments in other regions work; update the call to
auth.BuildAuthToken(context.TODO(), dbinput, <resolvedRegion>, s.Username,
cfg.Credentials) and ensure resolvedRegion falls back to the standard AWS SDK
resolution if cfg.Region is empty.
- Line 52: The non-IAM Postgres connection string builder is interpolating
s.Username and s.Secret raw, which can break the URI when they contain reserved
characters; update the Postgres connection string code (the line returning
fmt.Sprintf("postgresql://%s:%s@%s/%s", ...)) to apply URL-encoding (e.g.,
url.QueryEscape) to s.Username and s.Secret before formatting, and apply the
same encoding fix to Neo4jConnectionString so credentials are escaped
consistently.
- Around line 25-53: The PostgreSQLConnectionString method on
DatabaseConfiguration currently panics on transient AWS failures
(config.LoadDefaultConfig and auth.BuildAuthToken); change the signature of
PostgreSQLConnectionString to return (string, error), remove the panic calls and
instead return descriptive errors when config.LoadDefaultConfig or
auth.BuildAuthToken fail (include context like which step failed and the wrapped
error), and propagate the error paths where you currently fall back to returning
a string; update all call sites (notably in pg.go) to accept the (string, error)
result and handle retries/errors accordingly.
🧹 Nitpick comments (3)
drivers/config.go (1)
35-35: Inconsistent logging:fmt.Printfhere vsslogeverywhere else.Use
slog.Warn(orslog.Info) for consistency with the rest of the file.Proposed fix
- fmt.Printf("Error looking up CNAME for %s: %v. Using original address", s.Address, err) + slog.Warn("CNAME lookup failed, using original address", "address", s.Address, "error", err)drivers/pg/pg.go (2)
64-66:MaxConcurrentSessionsfrom config is ignored; pool limits are hardcoded.
DatabaseConfigurationhas aMaxConcurrentSessionsfield but Lines 65-66 hardcodeMinConns = 5andMaxConns = 50. The TODO on Line 64 acknowledges this. Consider wiring it up now since the config struct is already available:Proposed fix
- // TODO: Min and Max connections for the pool should be configurable - poolCfg.MinConns = 5 - poolCfg.MaxConns = 50 + if cfg.MaxConcurrentSessions > 0 { + poolCfg.MaxConns = int32(cfg.MaxConcurrentSessions) + } else { + poolCfg.MaxConns = 50 + } + poolCfg.MinConns = 5
54-62: No validation ofcfgbefore use.If
cfghas emptyAddress/Database/Usernameand noConnectionstring,PostgreSQLConnectionString()will silently produce a malformed DSN likepostgresql://://:5432/that will fail with an unhelpful error frompgxpool.ParseConfig. An early validation or a meaningful error message would improve debuggability.
| func (s DatabaseConfiguration) PostgreSQLConnectionString() string { | ||
| if s.IamAuth { | ||
| slog.Info("loading default config for rds auth") | ||
| cfg, err := config.LoadDefaultConfig(context.TODO()) | ||
| if err != nil { | ||
| panic("configuration error: " + err.Error()) | ||
| } | ||
|
|
||
| cname, err := net.LookupCNAME(s.Address) | ||
| if err != nil { | ||
| fmt.Printf("Error looking up CNAME for %s: %v. Using original address", s.Address, err) | ||
| cname = s.Address | ||
| } | ||
|
|
||
| dbinput := strings.TrimSuffix(cname, ".") + ":5432" | ||
| slog.Info("requesting auth token") | ||
| authenticationToken, err := auth.BuildAuthToken(context.TODO(), dbinput, "us-east-1", s.Username, cfg.Credentials) | ||
| if err != nil { | ||
| panic("failed to create authentication token: " + err.Error()) | ||
| } | ||
| slog.Info("auth token successfully created") | ||
| encodedToken := url.QueryEscape(authenticationToken) | ||
|
|
||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, encodedToken, dbinput, s.Database) | ||
| } else if s.Connection != "" { | ||
| return s.Connection | ||
| } else { | ||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database) | ||
| } |
There was a problem hiding this comment.
panic in a library function will crash callers on transient AWS errors.
PostgreSQLConnectionString() panics on AWS config-load failure (Line 30) and token-creation failure (Line 43). These can be transient (network blip, expired credentials, metadata-service timeout). A library should return (string, error) so callers can handle failures gracefully (retry, degrade, alert) instead of bringing down the process.
Proposed signature change
-func (s DatabaseConfiguration) PostgreSQLConnectionString() string {
+func (s DatabaseConfiguration) PostgreSQLConnectionString() (string, error) {
if s.IamAuth {
slog.Info("loading default config for rds auth")
cfg, err := config.LoadDefaultConfig(context.TODO())
if err != nil {
- panic("configuration error: " + err.Error())
+ return "", fmt.Errorf("loading AWS config: %w", err)
}
...
authenticationToken, err := auth.BuildAuthToken(...)
if err != nil {
- panic("failed to create authentication token: " + err.Error())
+ return "", fmt.Errorf("building RDS auth token: %w", err)
}
...
- return fmt.Sprintf(...)
+ return fmt.Sprintf(...), nil
} else if s.Connection != "" {
- return s.Connection
+ return s.Connection, nil
} else {
- return fmt.Sprintf(...)
+ return fmt.Sprintf(...), nil
}
}This will require updating callers in pg.go accordingly.
🤖 Prompt for AI Agents
In `@drivers/config.go` around lines 25 - 53, The PostgreSQLConnectionString
method on DatabaseConfiguration currently panics on transient AWS failures
(config.LoadDefaultConfig and auth.BuildAuthToken); change the signature of
PostgreSQLConnectionString to return (string, error), remove the panic calls and
instead return descriptive errors when config.LoadDefaultConfig or
auth.BuildAuthToken fail (include context like which step failed and the wrapped
error), and propagate the error paths where you currently fall back to returning
a string; update all call sites (notably in pg.go) to accept the (string, error)
result and handle retries/errors accordingly.
| } else if s.Connection != "" { | ||
| return s.Connection | ||
| } else { | ||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database) |
There was a problem hiding this comment.
Secret is not URL-encoded in the non-IAM connection string.
On Line 46-48, the IAM token is url.QueryEscape'd, but on Line 52 the Secret (and Username) are interpolated raw. If the secret contains characters like @, /, #, or ?, the resulting URI will be malformed.
Proposed fix
- return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database)
+ return fmt.Sprintf("postgresql://%s:%s@%s/%s", url.QueryEscape(s.Username), url.QueryEscape(s.Secret), s.Address, s.Database)Same applies to Neo4jConnectionString on Line 58.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database) | |
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", url.QueryEscape(s.Username), url.QueryEscape(s.Secret), s.Address, s.Database) |
🤖 Prompt for AI Agents
In `@drivers/config.go` at line 52, The non-IAM Postgres connection string builder
is interpolating s.Username and s.Secret raw, which can break the URI when they
contain reserved characters; update the Postgres connection string code (the
line returning fmt.Sprintf("postgresql://%s:%s@%s/%s", ...)) to apply
URL-encoding (e.g., url.QueryEscape) to s.Username and s.Secret before
formatting, and apply the same encoding fix to Neo4jConnectionString so
credentials are escaped consistently.
There was a problem hiding this comment.
This is by design. Tokens can return special characters, so the token must be escaped to be used in the connection string.
There was a problem hiding this comment.
@mamundsen-specter, I completely agree that the IAM token escaping on Line 46 is correct and necessary!
My concern is specifically about Line 52 in the non-IAM path:
return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database)Here, s.Secret (and s.Username) are interpolated raw without URL encoding. Secrets can also contain special characters like @, :, /, or # that would break the connection string, just like tokens can. For example:
- A secret like
myP@sswordwould break at the@ - A secret like
pass:word/123would be misinterpreted
Since you're already URL-encoding the IAM token for this exact reason, it would be safer and more consistent to also encode the credentials in the non-IAM path. However, if there's a specific reason why secrets are guaranteed not to contain special characters in your use case, that's fine too!
mamundsen-specter
force-pushed
the
bi-1348-main
branch
from
6bda635 to
09369ad
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@drivers/config.go`:
- Line 43: The fmt.Errorf call is using the %w verb with a string (err.Error()),
which fails to compile; update the fmt.Errorf invocation in the authentication
token creation path to pass the original error value (err) to %w (or switch to
%v/%s if you intend to embed the string), i.e., replace the use of err.Error()
with err in the fmt.Errorf call so the error can be wrapped properly (the
affected call is the fmt.Errorf("failed to create authentication token: %w",
...) in drivers/config.go).
🧹 Nitpick comments (1)
drivers/config.go (1)
33-36: Inconsistent logging:fmt.Printfhere vsslog.Infoeverywhere else.Use
slog.Warn(orslog.Info) for the CNAME fallback message to maintain consistent structured logging. Also, the format string on Line 35 is missing a trailing newline (if you keepfmt.Printf).Proposed fix
cname, err := net.LookupCNAME(s.Address) if err != nil { - fmt.Printf("Error looking up CNAME for %s: %v. Using original address", s.Address, err) + slog.Warn("CNAME lookup failed, using original address", "address", s.Address, "error", err) cname = s.Address }
mamundsen-specter
force-pushed
the
bi-1348-main
branch
from
09369ad to
8a4e8dd
Compare
mamundsen-specter
force-pushed
the
bi-1348-main
branch
from
8a4e8dd to
3e01779
Compare
… creation fully optional based on the config binary setting for RDS IAM auth
|
cleanup looks good and makes sense. shouldn't have to enable beforeConnect if we aren't using IAM auth and tidies up config.go 👍 working on bloodhound changes and testing before merging in |
mvlipka
left a comment
mvlipka
left a comment
There was a problem hiding this comment.
Looks good to me!
Only thing I'd say is that we should comment exported methods/functions, but these seem pretty self explanatory
|
Updated comments where it felt necessary and pulled in main one last time. Thanks for the review @mvlipka |
|
@CodeRabbit review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
drivers/config.go (1)
46-46:⚠️ Potential issue | 🟡 MinorPreserve configured ports when building the IAM endpoint.
Line 46 still appends
:5432unconditionally. Ifs.Addressalready includes a port and CNAME lookup falls back, this produces values likedb.example:5433:5432; it also prevents non-default RDS ports.Proposed fix
- endpoint := strings.TrimSuffix(host, ".") + ":5432" + host = strings.TrimSuffix(host, ".") + endpointHost, endpointPort, err := net.SplitHostPort(host) + if err != nil { + endpointHost = host + endpointPort = "5432" + } + endpoint := net.JoinHostPort(endpointHost, endpointPort)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@drivers/config.go` at line 46, The code builds endpoint := strings.TrimSuffix(host, ".") + ":5432" but unconditionally appends :5432, which breaks when s.Address/host already includes a port; change the logic to trim the trailing dot from host, then detect whether host already contains a port (use net.SplitHostPort or strings.Contains(host, ":") with care for IPv6) and only append ":5432" when no port is present; update the assignment to the endpoint variable accordingly and keep using host/s.Address and strings.TrimSuffix to locate the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@drivers/config.go`:
- Around line 60-62: The PostgreSQLConnectionString branch unconditionally calls
RDSIAMAuthConnectionString when EnableRDSIAMAuth is true, ignoring a provided
Connection DSN; update PostgreSQLConnectionString to handle the Connection field
explicitly: if EnableRDSIAMAuth is true and Connection is set, either
(preferred) parse the DSN in Connection to extract Address, Username and
Database and feed those values into RDSIAMAuthConnectionString (or a new helper
that accepts parsed inputs), or (alternatively) return an explicit validation
error rejecting the combination of Connection + EnableRDSIAMAuth. Modify
PostgreSQLConnectionString (and/or add a small DSN parsing helper used by
RDSIAMAuthConnectionString) to ensure IAM token generation uses the parsed
Connection values instead of empty Address/Username/Database.
- Around line 35-45: The IAM auth branch in
PostgreSQLConnectionString/defaultPostgreSQLConnectionString currently logs
failures from awsConfig.LoadDefaultConfig, net.LookupCNAME and token generation
and then returns the password fallback silently; change these sites (the
awsConfig.LoadDefaultConfig call, the net.LookupCNAME lookup using s.Address,
and the IAM token build) to return descriptive errors up the stack instead of
just slog.Error/slog.Warn so callers can detect IAM setup failures—update
PostgreSQLConnectionString() to propagate those errors (or wrap them with
context) rather than ignoring them and only falling back to
defaultPostgreSQLConnectionString().
In `@drivers/pg/pg.go`:
- Around line 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.
---
Duplicate comments:
In `@drivers/config.go`:
- Line 46: The code builds endpoint := strings.TrimSuffix(host, ".") + ":5432"
but unconditionally appends :5432, which breaks when s.Address/host already
includes a port; change the logic to trim the trailing dot from host, then
detect whether host already contains a port (use net.SplitHostPort or
strings.Contains(host, ":") with care for IPv6) and only append ":5432" when no
port is present; update the assignment to the endpoint variable accordingly and
keep using host/s.Address and strings.TrimSuffix to locate the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8c768b7e-f36f-4811-8a1d-2e942d028bc7
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (6)
cmd/benchmark/main.gocmd/export/main.godrivers/config.godrivers/pg/pg.gogo.modintegration/harness.go
🚧 Files skipped from review as they are similar to previous changes (1)
- go.mod
| if cfg.EnableRDSIAMAuth { | ||
| // Only enable the BeforeConnect handler if RDS IAM Auth is enabled | ||
| poolCfg.BeforeConnect = func(ctx context.Context, connCfg *pgx.ConnConfig) error { | ||
| if newPoolCfg, err := pgxpool.ParseConfig(cfg.PostgreSQLConnectionString()); err != nil { | ||
| return err | ||
| } else { | ||
| connCfg.Password = newPoolCfg.ConnConfig.Password | ||
| } | ||
|
|
||
| return nil | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# Find the relevant files
find . -name "pg.go" -path "*/drivers/*" | head -20Repository: SpecterOps/DAWGS
Length of output: 78
🏁 Script executed:
# Find config.go in drivers
find . -name "config.go" -path "*/drivers/*" | head -20Repository: 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.goRepository: SpecterOps/DAWGS
Length of output: 3664
🏁 Script executed:
cat -n drivers/config.goRepository: 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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
🧩 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.
65370c1 to
e92eb64
Compare
Signed-off-by: mamundsen-specter <mamundsen@specterops.io>
Signed-off-by: mamundsen-specter <mamundsen@specterops.io>
mamundsen-specter
force-pushed
the
bi-1348-main
branch
from
e92eb64 to
7129bea
Compare
3 participants
Summary by CodeRabbit
New Features
Improvements