feat: add awsiamrds db auth driver (#12566)

coderd/apidoc/docs.go

@@ -9597,6 +9597,9 @@ const docTemplate = `{
                 "oidc": {
                     "$ref": "#/definitions/codersdk.OIDCConfig"
                 },
+                "pg_auth": {
+                    "type": "string"
+                },
                 "pg_connection_url": {
                     "type": "string"
                 },

coderd/apidoc/swagger.json

@@ -8582,6 +8582,9 @@
         "oidc": {
           "$ref": "#/definitions/codersdk.OIDCConfig"
         },
+        "pg_auth": {
+          "type": "string"
+        },
         "pg_connection_url": {
           "type": "string"
         },

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
+}

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)
+}