coder · f0ssel · Mar 20, 2024 · Mar 12, 2024 · Mar 12, 2024 · Mar 14, 2024
diff --git a/cli/server.go b/cli/server.go
@@ -64,6 +64,7 @@ import (
 	"github.com/coder/coder/v2/coderd/autobuild"
 	"github.com/coder/coder/v2/coderd/batchstats"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/awsiamrds"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
 	"github.com/coder/coder/v2/coderd/database/dbmetrics"
 	"github.com/coder/coder/v2/coderd/database/dbpurge"
@@ -667,12 +668,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
 				options.Database = dbmem.New()
 				options.Pubsub = pubsub.NewInMemory()
 			} else {
-				dbURL, err := escapePostgresURLUserInfo(vals.PostgresURL.String())
-				if err != nil {
-					return xerrors.Errorf("escaping postgres URL: %w", err)
-				}
-
-				sqlDB, err := ConnectToPostgres(ctx, logger, sqlDriver, dbURL)
+				sqlDB, dbURL, err := getPostgresDB(ctx, logger, vals.PostgresURL.String(), codersdk.PostgresAuth(vals.PostgresAuth), sqlDriver)
 				if err != nil {
 					return xerrors.Errorf("connect to postgres: %w", err)
 				}
@@ -2548,3 +2544,24 @@ func signalNotifyContext(ctx context.Context, inv *serpent.Invocation, sig ...os
 	}
 	return inv.SignalNotifyContext(ctx, sig...)
 }
+
+func getPostgresDB(ctx context.Context, logger slog.Logger, postgresURL string, auth codersdk.PostgresAuth, sqlDriver string) (*sql.DB, string, error) {
+	dbURL, err := escapePostgresURLUserInfo(postgresURL)
+	if err != nil {
+		return nil, "", xerrors.Errorf("escaping postgres URL: %w", err)
+	}
+
+	if auth == codersdk.PostgresAuthAWSIAMRDS {
+		sqlDriver, err = awsiamrds.Register(ctx, sqlDriver)
+		if err != nil {
+			return nil, "", xerrors.Errorf("register aws rds iam auth: %w", err)
+		}
+	}
+
+	sqlDB, err := ConnectToPostgres(ctx, logger, sqlDriver, dbURL)
+	if err != nil {
+		return nil, "", xerrors.Errorf("connect to postgres: %w", err)
+	}
+
+	return sqlDB, dbURL, nil
+}
diff --git a/cli/server_createadminuser.go b/cli/server_createadminuser.go
@@ -13,6 +13,7 @@ import (
 	"cdr.dev/slog/sloggers/sloghuman"
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/awsiamrds"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/gitsshkey"
 	"github.com/coder/coder/v2/coderd/httpapi"
@@ -25,6 +26,7 @@ import (
 func (r *RootCmd) newCreateAdminUserCommand() *serpent.Command {
 	var (
 		newUserDBURL              string
+		newUserPgAuth             string
 		newUserSSHKeygenAlgorithm string
 		newUserUsername           string
 		newUserEmail              string
@@ -62,7 +64,15 @@ func (r *RootCmd) newCreateAdminUserCommand() *serpent.Command {
 				newUserDBURL = url
 			}
 
-			sqlDB, err := ConnectToPostgres(ctx, logger, "postgres", newUserDBURL)
+			sqlDriver := "postgres"
+			if codersdk.PostgresAuth(newUserPgAuth) == codersdk.PostgresAuthAWSIAMRDS {
+				sqlDriver, err = awsiamrds.Register(inv.Context(), sqlDriver)
+				if err != nil {
+					return xerrors.Errorf("register aws rds iam auth: %w", err)
+				}
+			}
+
+			sqlDB, err := ConnectToPostgres(ctx, logger, sqlDriver, newUserDBURL)
 			if err != nil {
 				return xerrors.Errorf("connect to postgres: %w", err)
 			}
@@ -243,6 +253,14 @@ func (r *RootCmd) newCreateAdminUserCommand() *serpent.Command {
 			Description: "URL of a PostgreSQL database. If empty, the built-in PostgreSQL deployment will be used (Coder must not be already running in this case).",
 			Value:       serpent.StringOf(&newUserDBURL),
 		},
+		serpent.Option{
+			Name:        "Postgres Connection Auth",
+			Description: "Type of auth to use when connecting to postgres.",
+			Flag:        "postgres-connection-auth",
+			Env:         "CODER_PG_CONNECTION_AUTH",
+			Default:     "password",
+			Value:       serpent.EnumOf(&newUserPgAuth, codersdk.PostgresAuthDrivers...),
+		},
 		serpent.Option{
 			Env:         "CODER_SSH_KEYGEN_ALGORITHM",
 			Flag:        "ssh-keygen-algorithm",

diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden
@@ -44,6 +44,9 @@ OPTIONS:
           Separate multiple experiments with commas, or enter '*' to opt-in to
           all available experiments.
 
+      --postgres-auth password|awsiamrds, $CODER_PG_AUTH (default: password)
+          Type of auth to use when connecting to postgres.
+
       --postgres-url string, $CODER_PG_CONNECTION_URL
           URL of a PostgreSQL database. If empty, PostgreSQL binaries will be
           downloaded from Maven (https://repo1.maven.org/maven2) and store all

diff --git a/cli/testdata/coder_server_create-admin-user_--help.golden b/cli/testdata/coder_server_create-admin-user_--help.golden
@@ -7,6 +7,9 @@ USAGE:
   it to every organization.
 
 OPTIONS:
+      --postgres-connection-auth password|awsiamrds, $CODER_PG_CONNECTION_AUTH (default: password)
+          Type of auth to use when connecting to postgres.
+
       --email string, $CODER_EMAIL
           The email of the new user. If not specified, you will be prompted via
           stdin.

diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden
@@ -411,6 +411,9 @@ cacheDir: [cache dir]
 # Controls whether data will be stored in an in-memory database.
 # (default: <unset>, type: bool)
 inMemoryDatabase: false
+# Type of auth to use when connecting to postgres.
+# (default: password, type: enum[password\|awsiamrds])
+pgAuth: password
 # The algorithm to use for generating ssh keys. Accepted values are "ed25519",
 # "ecdsa", or "rsa4096".
 # (default: ed25519, type: string)

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/awsiamrds/awsiamrds.go b/coderd/database/awsiamrds/awsiamrds.go
@@ -0,0 +1,84 @@
+package awsiamrds
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"fmt"
+	"net/url"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/rds/auth"
+	"golang.org/x/xerrors"
+)
+
+type awsIamRdsDriver struct {
+	parent driver.Driver
+	cfg    aws.Config
+}
+
+var _ driver.Driver = &awsIamRdsDriver{}
+
+// Register initializes and registers our aws iam rds wrapped database driver.
+func Register(ctx context.Context, parentName string) (string, error) {
+	cfg, err := config.LoadDefaultConfig(ctx)
+	if err != nil {
+		return "", err
+	}
+
+	db, err := sql.Open(parentName, "")
+	if err != nil {
+		return "", err
+	}
+
+	// create a new aws iam rds driver
+	d := newDriver(db.Driver(), cfg)
+	name := fmt.Sprintf("%s-awsiamrds", parentName)
+	sql.Register(fmt.Sprintf("%s-awsiamrds", parentName), d)
+
+	return name, nil
+}
+
+// newDriver will create a new *AwsIamRdsDriver using the environment aws session.
+func newDriver(parentDriver driver.Driver, cfg aws.Config) *awsIamRdsDriver {
+	return &awsIamRdsDriver{
+		parent: parentDriver,
+		cfg:    cfg,
+	}
+}
+
+// Open creates a new connection to the database using the provided name.
+func (d *awsIamRdsDriver) Open(name string) (driver.Conn, error) {
+	// set password with signed aws authentication token for the rds instance
+	nURL, err := getAuthenticatedURL(d.cfg, name)
+	if err != nil {
+		return nil, xerrors.Errorf("assigning authentication token to url: %w", err)
+	}
+
+	// make connection
+	conn, err := d.parent.Open(nURL)
+	if err != nil {
+		return nil, xerrors.Errorf("opening connection with %s: %w", nURL, err)
+	}
+
+	return conn, nil
+}
+
+func getAuthenticatedURL(cfg aws.Config, dbURL string) (string, error) {
+	nURL, err := url.Parse(dbURL)
+	if err != nil {
+		return "", xerrors.Errorf("parsing dbURL: %w", err)
+	}
+
+	// generate a new rds session auth tokenized URL
+	rdsEndpoint := fmt.Sprintf("%s:%s", nURL.Hostname(), nURL.Port())
+	token, err := auth.BuildAuthToken(context.Background(), rdsEndpoint, cfg.Region, nURL.User.Username(), cfg.Credentials)
+	if err != nil {
+		return "", xerrors.Errorf("building rds auth token: %w", err)
+	}
+	// set token as user password
+	nURL.User = url.UserPassword(nURL.User.Username(), token)
+
+	return nURL.String(), nil
+}
diff --git a/coderd/database/awsiamrds/awsiamrds_test.go b/coderd/database/awsiamrds/awsiamrds_test.go
@@ -0,0 +1,50 @@
+package awsiamrds_test
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog/sloggers/slogtest"
+
+	"github.com/coder/coder/v2/cli"
+	awsrdsiam "github.com/coder/coder/v2/coderd/database/awsiamrds"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestDriver(t *testing.T) {
+	t.Parallel()
+	// Be sure to set AWS_DEFAULT_REGION to the database region as well.
+	// Example:
+	// export AWS_DEFAULT_REGION=us-east-2;
+	// export DBAWSIAMRDS_TEST_URL="postgres://user@host:5432/dbname";
+	url := os.Getenv("DBAWSIAMRDS_TEST_URL")
+	if url == "" {
+		t.Skip()
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	sqlDriver, err := awsrdsiam.Register(ctx, "postgres")
+	require.NoError(t, err)
+
+	db, err := cli.ConnectToPostgres(ctx, slogtest.Make(t, nil), sqlDriver, url)
+	require.NoError(t, err)
+	defer func() {
+		_ = db.Close()
+	}()
+
+	i, err := db.QueryContext(ctx, "select 1;")
+	require.NoError(t, err)
+	defer func() {
+		_ = i.Close()
+	}()
+
+	require.True(t, i.Next())
+	var one int
+	require.NoError(t, i.Scan(&one))
+	require.Equal(t, 1, one)
+}
diff --git a/codersdk/deployment.go b/codersdk/deployment.go
@@ -135,6 +135,18 @@ func (c *Client) Entitlements(ctx context.Context) (Entitlements, error) {
 	return ent, json.NewDecoder(res.Body).Decode(&ent)
 }
 
+type PostgresAuth string
+
+const (
+	PostgresAuthPassword  PostgresAuth = "password"
+	PostgresAuthAWSIAMRDS PostgresAuth = "awsiamrds"
+)
+
+var PostgresAuthDrivers = []string{
+	string(PostgresAuthPassword),
+	string(PostgresAuthAWSIAMRDS),
+}
+
 // DeploymentValues is the central configuration values the coder server.
 type DeploymentValues struct {
 	Verbose             serpent.Bool   `json:"verbose,omitempty"`
@@ -154,6 +166,7 @@ type DeploymentValues struct {
 	CacheDir                        serpent.String                       `json:"cache_directory,omitempty" typescript:",notnull"`
 	InMemoryDatabase                serpent.Bool                         `json:"in_memory_database,omitempty" typescript:",notnull"`
 	PostgresURL                     serpent.String                       `json:"pg_connection_url,omitempty" typescript:",notnull"`
+	PostgresAuth                    string                               `json:"pg_auth,omitempty" typescript:",notnull"`
 	OAuth2                          OAuth2Config                         `json:"oauth2,omitempty" typescript:",notnull"`
 	OIDC                            OIDCConfig                           `json:"oidc,omitempty" typescript:",notnull"`
 	Telemetry                       TelemetryConfig                      `json:"telemetry,omitempty" typescript:",notnull"`
@@ -1630,6 +1643,15 @@ when required by your organization's security policy.`,
 			Annotations: serpent.Annotations{}.Mark(annotationSecretKey, "true"),
 			Value:       &c.PostgresURL,
 		},
+		{
+			Name:        "Postgres Auth",
+			Description: "Type of auth to use when connecting to postgres.",
+			Flag:        "postgres-auth",
+			Env:         "CODER_PG_AUTH",
+			Default:     "password",
+			Value:       serpent.EnumOf(&c.PostgresAuth, PostgresAuthDrivers...),
+			YAML:        "pgAuth",
+		},
 		{
 			Name:        "Secure Auth Cookie",
 			Description: "Controls if the 'Secure' property is set on browser session cookies.",

diff --git a/docs/api/general.md b/docs/api/general.md
diff --git a/docs/api/schemas.md b/docs/api/schemas.md
diff --git a/docs/cli/server.md b/docs/cli/server.md
diff --git a/docs/cli/server_create-admin-user.md b/docs/cli/server_create-admin-user.md