Update newpool to add IAM auth beforeConnect#32
Update newpool to add IAM auth beforeConnect#32mamundsen-specter wants to merge 7 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:
WalkthroughThe PostgreSQL pool now accepts a 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: Create pool (MinConns/MaxConns) with BeforeConnect
App->>Driver: Acquire connection
Driver->>Driver: BeforeConnect triggered
Driver->>Config: Re-parse connection info
alt IAM enabled
Driver->>IAM: Request RDS auth token
IAM-->>Driver: Return token (password)
end
Driver->>DB: Establish connection using refreshed creds
DB-->>Driver: Connection established
Driver-->>App: Return connection
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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?
drivers/pg/pg.go
Outdated
| 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.
drivers/config.go
Outdated
| 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.
drivers/config.go
Outdated
| cname = s.Address | ||
| } | ||
|
|
||
| dbinput := strings.TrimSuffix(cname, ".") + ":5432" |
There was a problem hiding this comment.
Hardcoded port :5432 — should be parsed from Address or configurable.
If Address already contains a port (e.g., host:5433), appending :5432 will produce an invalid endpoint. Consider parsing the address or adding a port field.
🤖 Prompt for AI Agents
In `@drivers/config.go` at 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.
drivers/config.go
Outdated
| } 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 }
drivers/config.go
Outdated
| slog.Info("requesting auth token") | ||
| authenticationToken, err := auth.BuildAuthToken(context.TODO(), dbinput, cfg.Region, s.Username, cfg.Credentials) | ||
| if err != nil { | ||
| return "", fmt.Errorf("failed to create authentication token: %w", err.Error()) |
There was a problem hiding this comment.
fmt.Errorf with %w requires an error, not a string — this will not compile.
err.Error() returns a string, but the %w verb expects a value implementing error. The static analysis pipeline confirms this failure.
Proposed fix
- return "", fmt.Errorf("failed to create authentication token: %w", err.Error())
+ return "", fmt.Errorf("failed to create authentication token: %w", err)📝 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.Errorf("failed to create authentication token: %w", err.Error()) | |
| return "", fmt.Errorf("failed to create authentication token: %w", err) |
🧰 Tools
🪛 GitHub Check: test
[failure] 43-43:
fmt.Errorf format %w has arg err.Error() of wrong type string
🤖 Prompt for AI Agents
In `@drivers/config.go` at 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).
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 |
2 participants
Summary by CodeRabbit
New Features
Improvements