feat: move DatabaseConfiguration from Bloodhound repo to dawgs allowing for IAM auth option #57
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,76 @@ | ||
| package drivers | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "log/slog" | ||
| "net" | ||
| "net/url" | ||
| "strings" | ||
|
|
||
| awsConfig "github.com/aws/aws-sdk-go-v2/config" | ||
| "github.com/aws/aws-sdk-go-v2/feature/rds/auth" | ||
| ) | ||
|
|
||
| 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"` | ||
| EnableRDSIAMAuth bool `json:"enable_rds_iam_auth"` | ||
| } | ||
|
|
||
| func (s DatabaseConfiguration) defaultPostgreSQLConnectionString() string { | ||
| if s.Connection != "" { | ||
| return s.Connection | ||
| } | ||
|
|
||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database) | ||
| } | ||
|
|
||
| func (s DatabaseConfiguration) RDSIAMAuthConnectionString() string { | ||
| slog.Info("Loading RDS Configuration With IAM Auth") | ||
|
|
||
| if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil { | ||
| slog.Error("AWS Config Loading Error", slog.String("err", err.Error())) | ||
| } else { | ||
| host := s.Address | ||
|
|
||
| if hostCName, err := net.LookupCNAME(s.Address); err != nil { | ||
| slog.Warn("Error looking up CNAME for DB host. Using original address.", slog.String("err", err.Error())) | ||
| } else { | ||
| host = hostCName | ||
| } | ||
|
|
||
| endpoint := strings.TrimSuffix(host, ".") + ":5432" | ||
|
There was a problem hiding this comment. Hardcoded port 5432 may not match the configured address. If 🔧 Suggested approach+func (s DatabaseConfiguration) getHostPort() (string, string) {
+ host := s.Address
+ port := "5432"
+ if h, p, err := net.SplitHostPort(s.Address); err == nil {
+ host = h
+ port = p
+ }
+ return host, port
+}Then use it in - endpoint := strings.TrimSuffix(host, ".") + ":5432"
+ _, port := s.getHostPort()
+ endpoint := strings.TrimSuffix(host, ".") + ":" + port🤖 Prompt for AI Agents |
||
|
|
||
| slog.Info("Requesting RDS IAM Auth Token") | ||
|
|
||
| if authenticationToken, err := auth.BuildAuthToken(context.TODO(), endpoint, cfg.Region, s.Username, cfg.Credentials); err != nil { | ||
|
Comment on lines
+36
to
+51
There was a problem hiding this comment. Avoid repeated AWS config/DNS work on every connection attempt. Lines 36-51 perform 💡 Refactor direction- if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil {
+ // load once (e.g., during pool initialization) and reuse:
+ // - aws.Config
+ // - normalized host/endpoint
+ // per-connection: only mint the IAM token🤖 Prompt for AI Agents |
||
| slog.Error("RDS IAM Auth Token Request Error", slog.String("err", err.Error())) | ||
| } else { | ||
| slog.Info("RDS IAM Auth Token Created") | ||
| return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(authenticationToken), endpoint, s.Database) | ||
| } | ||
| } | ||
|
|
||
| return s.defaultPostgreSQLConnectionString() | ||
| } | ||
|
Comment on lines
+33
to
+60
There was a problem hiding this comment. Silent error fallback may mask configuration issues. When AWS config loading or IAM token generation fails, the function logs an error but silently falls back to Consider returning an error or making the fallback behavior explicit/configurable. 🔧 Suggested approach to surface errors-func (s DatabaseConfiguration) RDSIAMAuthConnectionString() string {
+func (s DatabaseConfiguration) RDSIAMAuthConnectionString() (string, error) {
slog.Info("Loading RDS Configuration With IAM Auth")
if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil {
- slog.Error("AWS Config Loading Error", slog.String("err", err.Error()))
+ return "", fmt.Errorf("failed to load AWS config: %w", err)
- } else {
+ }
// ... rest of implementation
- }
-
- return s.defaultPostgreSQLConnectionString()
+ return connString, nil
}Then update 🤖 Prompt for AI Agents |
||
|
|
||
| func (s DatabaseConfiguration) PostgreSQLConnectionString() string { | ||
| if s.EnableRDSIAMAuth { | ||
| return s.RDSIAMAuthConnectionString() | ||
| } | ||
|
|
||
| return s.defaultPostgreSQLConnectionString() | ||
| } | ||
|
|
||
| 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 | ||
| } | ||
|
Comment on lines
+70
to
+76
There was a problem hiding this comment.
Unlike 🔧 Proposed fix 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 fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database)
}
return s.Connection
}🤖 Prompt for AI Agents |
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🧩 Analysis chain
🌐 Web query:
In Go's net/url package, is url.QueryEscape appropriate for encoding URL userinfo credentials, or should url.UserPassword/url.URL be used? Please include how spaces and reserved characters are encoded in each case.💡 Result:
url.UserPassword/url.URL should be used for encoding URL userinfo credentials, not url.QueryEscape. url.UserPassword(username, password) creates a *Userinfo that, when used in a url.URL and stringified via u.String, properly encodes the username and password using the package's internal userinfo encoding rules (encodeUserPassword mode). This escapes only the characters that must be escaped in userinfo per RFC 3986 §3.2.1: '@', '/', '?', and ':' (to handle parsing). Other reserved characters like '&', '=', '+', ';', ',', '$' are left unescaped, and spaces are encoded as '+'. In contrast, url.QueryEscape is for query parameters (encodeQueryComponent mode). It escapes all reserved characters (including '&' as %26, '=' as %3D, etc.) and encodes spaces as '+'. Using QueryEscape for userinfo can produce incorrectly encoded URLs because it over-escapes characters safe in userinfo, potentially causing parsing issues or non-standard URLs. Official docs recommend UserPassword for constructing userinfo, noting it's for legacy sites but handles encoding correctly. Since Go 1.9.4, url.Parse requires userinfo to be percent-encoded if it contains invalid characters. Examples: For username="user name", password="pass&word?": - url.UserPassword: "user+name:pass%26word%3F" (spaces -> '+', '&'->%26, '?'->%3F; '+' and '=' unescaped if present) - url.QueryEscape on both then join: "user+name:pass%26word%3F" (same for these chars, but would escape '=' to %3D, which userinfo allows unescaped) Code snippet: package main import ( "fmt" "net/url" ) func main { u := &url.URL{ Scheme: "https", Host: "example.com", User: url.UserPassword("user name", "pass&word?"), } fmt.Println(u.String) // https://user+name:pass%26word%3F@example.com }
Citations:
🏁 Script executed:
Repository: SpecterOps/DAWGS
Length of output: 42
🏁 Script executed:
Repository: SpecterOps/DAWGS
Length of output: 77
🏁 Script executed:
Repository: SpecterOps/DAWGS
Length of output: 68
🏁 Script executed:
cat -n drivers/config.go | head -50Repository: SpecterOps/DAWGS
Length of output: 1846
Use
url.UserPassword()to properly encode URL credentials instead ofurl.QueryEscape().Line 30 uses
url.QueryEscape()for encoding credentials in the userinfo portion of the PostgreSQL URL, which is incorrect.url.QueryEscapeis designed for query parameters and over-escapes characters that are safe in userinfo (like=,;). Useurl.UserPassword(username, password)withurl.URLinstead, which applies the correct userinfo encoding per RFC 3986.🔧 Suggested fix
func (s DatabaseConfiguration) defaultPostgreSQLConnectionString() string { if s.Connection != "" { return s.Connection } - return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database) + u := &url.URL{ + Scheme: "postgresql", + Host: s.Address, + Path: "/" + s.Database, + } + u.User = url.UserPassword(s.Username, s.Secret) + return u.String() }🤖 Prompt for AI Agents