Skip to content

Commit 5b67067

Browse files
mamundsen-spectermamundsen-specter
committed
move DatabaseConfiguration from Bloodhound repo to dawgs, update newpool
1 parent 2380370 commit 5b67067

6 files changed

Lines changed: 151 additions & 9 deletions

File tree

‎cmd/benchmark/main.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"time"
2727

2828
"github.com/specterops/dawgs"
29+
"github.com/specterops/dawgs/drivers"
2930
"github.com/specterops/dawgs/drivers/pg"
3031
"github.com/specterops/dawgs/graph"
3132
"github.com/specterops/dawgs/opengraph"
@@ -39,10 +40,11 @@ func main() {
3940
driver = flag.String("driver", "pg", "database driver (pg, neo4j)")
4041
connStr = flag.String("connection", "", "database connection string (or PG_CONNECTION_STRING)")
4142
iterations = flag.Int("iterations", 10, "timed iterations per scenario")
42-
output = flag.String("output", "", "markdown output file (default: stdout)")
43-
datasetDir = flag.String("dataset-dir", "integration/testdata", "path to testdata directory")
43+
output = flag.String("output", "", "markdown output file (default: stdout)")
44+
datasetDir = flag.String("dataset-dir", "integration/testdata", "path to testdata directory")
4445
localDataset = flag.String("local-dataset", "", "additional local dataset (e.g. local/phantom)")
4546
onlyDataset = flag.String("dataset", "", "run only this dataset (e.g. diamond, local/phantom)")
47+
dbcfg = drivers.DatabaseConfiguration{}
4648
)
4749

4850
flag.Parse()
@@ -62,8 +64,10 @@ func main() {
6264
ConnectionString: conn,
6365
}
6466

67+
dbcfg.Connection = conn
68+
6569
if *driver == pg.DriverName {
66-
pool, err := pg.NewPool(conn)
70+
pool, err := pg.NewPool(dbcfg)
6771
if err != nil {
6872
fatal("failed to create pool: %v", err)
6973
}

‎cmd/export/main.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"os"
77

8+
"github.com/specterops/dawgs/drivers"
89
"github.com/specterops/dawgs/drivers/pg"
910
"github.com/specterops/dawgs/opengraph"
1011
"github.com/specterops/dawgs/util/size"
@@ -16,7 +17,10 @@ func main() {
1617
connStr = "postgresql://bloodhound:bloodhoundcommunityedition@localhost:5432/bloodhound"
1718
}
1819

19-
pool, err := pg.NewPool(connStr)
20+
dbcfg := drivers.DatabaseConfiguration{}
21+
dbcfg.Connection = connStr
22+
23+
pool, err := pg.NewPool(dbcfg)
2024
if err != nil {
2125
fmt.Fprintf(os.Stderr, "failed to connect: %v\n", err)
2226
os.Exit(1)

‎drivers/config.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package drivers
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"log/slog"
7+
"net"
8+
"net/url"
9+
"strings"
10+
11+
awsConfig "github.com/aws/aws-sdk-go-v2/config"
12+
"github.com/aws/aws-sdk-go-v2/feature/rds/auth"
13+
)
14+
15+
type DatabaseConfiguration struct {
16+
Connection string `json:"connection"`
17+
Address string `json:"addr"`
18+
Database string `json:"database"`
19+
Username string `json:"username"`
20+
Secret string `json:"secret"`
21+
MaxConcurrentSessions int `json:"max_concurrent_sessions"`
22+
EnableRDSIAMAuth bool `json:"enable_rds_iam_auth"`
23+
}
24+
25+
func (s DatabaseConfiguration) defaultPostgreSQLConnectionString() string {
26+
if s.Connection != "" {
27+
return s.Connection
28+
}
29+
30+
return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(s.Secret), s.Address, s.Database)
31+
}
32+
33+
func (s DatabaseConfiguration) RDSIAMAuthConnectionString() string {
34+
slog.Info("Loading RDS Configuration With IAM Auth")
35+
36+
if cfg, err := awsConfig.LoadDefaultConfig(context.TODO()); err != nil {
37+
slog.Error("AWS Config Loading Error", slog.String("err", err.Error()))
38+
} else {
39+
host := s.Address
40+
41+
if hostCName, err := net.LookupCNAME(s.Address); err != nil {
42+
slog.Warn("Error looking up CNAME for DB host. Using original address.", slog.String("err", err.Error()))
43+
} else {
44+
host = hostCName
45+
}
46+
47+
endpoint := strings.TrimSuffix(host, ".") + ":5432"
48+
49+
slog.Info("Requesting RDS IAM Auth Token")
50+
51+
if authenticationToken, err := auth.BuildAuthToken(context.TODO(), endpoint, cfg.Region, s.Username, cfg.Credentials); err != nil {
52+
slog.Error("RDS IAM Auth Token Request Error", slog.String("err", err.Error()))
53+
} else {
54+
slog.Info("RDS IAM Auth Token Created")
55+
return fmt.Sprintf("postgresql://%s:%s@%s/%s", s.Username, url.QueryEscape(authenticationToken), endpoint, s.Database)
56+
}
57+
}
58+
59+
return s.defaultPostgreSQLConnectionString()
60+
}
61+
62+
func (s DatabaseConfiguration) PostgreSQLConnectionString() string {
63+
if s.EnableRDSIAMAuth {
64+
return s.RDSIAMAuthConnectionString()
65+
}
66+
67+
return s.defaultPostgreSQLConnectionString()
68+
}
69+
70+
func (s DatabaseConfiguration) Neo4jConnectionString() string {
71+
if s.Connection == "" {
72+
return fmt.Sprintf("neo4j://%s:%s@%s/%s", s.Username, s.Secret, s.Address, s.Database)
73+
}
74+
75+
return s.Connection
76+
}

‎drivers/pg/pg.go‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"github.com/jackc/pgx/v5/pgxpool"
1111
"github.com/specterops/dawgs"
1212
"github.com/specterops/dawgs/cypher/models/pgsql"
13+
"github.com/specterops/dawgs/drivers"
1314
"github.com/specterops/dawgs/graph"
1415
)
1516

@@ -50,15 +51,12 @@ func afterPooledConnectionRelease(conn *pgx.Conn) bool {
5051
return true
5152
}
5253

53-
func NewPool(connectionString string) (*pgxpool.Pool, error) {
54-
if connectionString == "" {
55-
return nil, fmt.Errorf("graph connection requires a connection url to be set")
56-
}
54+
func NewPool(cfg drivers.DatabaseConfiguration) (*pgxpool.Pool, error) {
5755

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

61-
poolCfg, err := pgxpool.ParseConfig(connectionString)
59+
poolCfg, err := pgxpool.ParseConfig(cfg.PostgreSQLConnectionString())
6260
if err != nil {
6361
return nil, err
6462
}
@@ -73,6 +71,21 @@ func NewPool(connectionString string) (*pgxpool.Pool, error) {
7371
poolCfg.AfterConnect = afterPooledConnectionEstablished
7472
poolCfg.AfterRelease = afterPooledConnectionRelease
7573

74+
if cfg.EnableRDSIAMAuth {
75+
// Only enable the BeforeConnect handler if RDS IAM Auth is enabled
76+
poolCfg.BeforeConnect = func(ctx context.Context, connCfg *pgx.ConnConfig) error {
77+
slog.Debug("New Connection RDS IAM Auth")
78+
79+
if newPoolCfg, err := pgxpool.ParseConfig(cfg.PostgreSQLConnectionString()); err != nil {
80+
return err
81+
} else {
82+
connCfg.Password = newPoolCfg.ConnConfig.Password
83+
}
84+
85+
return nil
86+
}
87+
}
88+
7689
pool, err := pgxpool.NewWithConfig(poolCtx, poolCfg)
7790
if err != nil {
7891
return nil, err

‎go.mod‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,21 @@ require (
1717
)
1818

1919
require (
20+
github.com/aws/aws-sdk-go-v2 v1.41.5 // indirect
21+
github.com/aws/aws-sdk-go-v2/config v1.32.13 // indirect
22+
github.com/aws/aws-sdk-go-v2/credentials v1.19.13 // indirect
23+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 // indirect
24+
github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.21 // indirect
25+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect
26+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect
27+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 // indirect
28+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect
29+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect
30+
github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 // indirect
31+
github.com/aws/aws-sdk-go-v2/service/sso v1.30.14 // indirect
32+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.18 // indirect
33+
github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 // indirect
34+
github.com/aws/smithy-go v1.24.2 // indirect
2035
github.com/cockroachdb/apd/v3 v3.2.2 // indirect
2136
github.com/davecgh/go-spew v1.1.1 // indirect
2237
github.com/dgryski/go-metro v0.0.0-20250106013310-edb8663e5e33 // indirect

‎go.sum‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,36 @@ github.com/RoaringBitmap/roaring/v2 v2.16.0 h1:Kys1UNf49d5W8Tq3bpuAhIr/Z8/yPB+59
88
github.com/RoaringBitmap/roaring/v2 v2.16.0/go.mod h1:eq4wdNXxtJIS/oikeCzdX1rBzek7ANzbth041hrU8Q4=
99
github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ=
1010
github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw=
11+
github.com/aws/aws-sdk-go-v2 v1.41.5 h1:dj5kopbwUsVUVFgO4Fi5BIT3t4WyqIDjGKCangnV/yY=
12+
github.com/aws/aws-sdk-go-v2 v1.41.5/go.mod h1:mwsPRE8ceUUpiTgF7QmQIJ7lgsKUPQOUl3o72QBrE1o=
13+
github.com/aws/aws-sdk-go-v2/config v1.32.13 h1:5KgbxMaS2coSWRrx9TX/QtWbqzgQkOdEa3sZPhBhCSg=
14+
github.com/aws/aws-sdk-go-v2/config v1.32.13/go.mod h1:8zz7wedqtCbw5e9Mi2doEwDyEgHcEE9YOJp6a8jdSMY=
15+
github.com/aws/aws-sdk-go-v2/credentials v1.19.13 h1:mA59E3fokBvyEGHKFdnpNNrvaR351cqiHgRg+JzOSRI=
16+
github.com/aws/aws-sdk-go-v2/credentials v1.19.13/go.mod h1:yoTXOQKea18nrM69wGF9jBdG4WocSZA1h38A+t/MAsk=
17+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 h1:NUS3K4BTDArQqNu2ih7yeDLaS3bmHD0YndtA6UP884g=
18+
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21/go.mod h1:YWNWJQNjKigKY1RHVJCuupeWDrrHjRqHm0N9rdrWzYI=
19+
github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.21 h1:HFn8sVT87KWnGs2Q2gO/brPZc2bR0RXD++cYKRmABzk=
20+
github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.21/go.mod h1:BGZ/K6gLGJt8K36j6gcsD7WVxmWt0MGBYtr57iLweio=
21+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 h1:Rgg6wvjjtX8bNHcvi9OnXWwcE0a2vGpbwmtICOsvcf4=
22+
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21/go.mod h1:A/kJFst/nm//cyqonihbdpQZwiUhhzpqTsdbhDdRF9c=
23+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 h1:PEgGVtPoB6NTpPrBgqSE5hE/o47Ij9qk/SEZFbUOe9A=
24+
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21/go.mod h1:p+hz+PRAYlY3zcpJhPwXlLC4C+kqn70WIHwnzAfs6ps=
25+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 h1:qYQ4pzQ2Oz6WpQ8T3HvGHnZydA72MnLuFK9tJwmrbHw=
26+
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6/go.mod h1:O3h0IK87yXci+kg6flUKzJnWeziQUKciKrLjcatSNcY=
27+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 h1:5EniKhLZe4xzL7a+fU3C2tfUN4nWIqlLesfrjkuPFTY=
28+
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7/go.mod h1:x0nZssQ3qZSnIcePWLvcoFisRXJzcTVvYpAAdYX8+GI=
29+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 h1:c31//R3xgIJMSC8S6hEVq+38DcvUlgFY0FM6mSI5oto=
30+
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21/go.mod h1:r6+pf23ouCB718FUxaqzZdbpYFyDtehyZcmP5KL9FkA=
31+
github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 h1:QKZH0S178gCmFEgst8hN0mCX1KxLgHBKKY/CLqwP8lg=
32+
github.com/aws/aws-sdk-go-v2/service/signin v1.0.9/go.mod h1:7yuQJoT+OoH8aqIxw9vwF+8KpvLZ8AWmvmUWHsGQZvI=
33+
github.com/aws/aws-sdk-go-v2/service/sso v1.30.14 h1:GcLE9ba5ehAQma6wlopUesYg/hbcOhFNWTjELkiWkh4=
34+
github.com/aws/aws-sdk-go-v2/service/sso v1.30.14/go.mod h1:WSvS1NLr7JaPunCXqpJnWk1Bjo7IxzZXrZi1QQCkuqM=
35+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.18 h1:mP49nTpfKtpXLt5SLn8Uv8z6W+03jYVoOSAl/c02nog=
36+
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.18/go.mod h1:YO8TrYtFdl5w/4vmjL8zaBSsiNp3w0L1FfKVKenZT7w=
37+
github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 h1:p8ogvvLugcR/zLBXTXrTkj0RYBUdErbMnAFFp12Lm/U=
38+
github.com/aws/aws-sdk-go-v2/service/sts v1.41.10/go.mod h1:60dv0eZJfeVXfbT1tFJinbHrDfSJ2GZl4Q//OSSNAVw=
39+
github.com/aws/smithy-go v1.24.2 h1:FzA3bu/nt/vDvmnkg+R8Xl46gmzEDam6mZ1hzmwXFng=
40+
github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc=
1141
github.com/axiomhq/hyperloglog v0.2.6 h1:sRhvvF3RIXWQgAXaTphLp4yJiX4S0IN3MWTaAgZoRJw=
1242
github.com/axiomhq/hyperloglog v0.2.6/go.mod h1:YjX/dQqCR/7QYX0g8mu8UZAjpIenz1FKM71UEsjFoTo=
1343
github.com/bits-and-blooms/bitset v1.24.4 h1:95H15Og1clikBrKr/DuzMXkQzECs1M6hhoGXLwLQOZE=

0 commit comments

Comments
 (0)