coder · f0ssel · Aug 26, 2024 · Aug 13, 2024 · Aug 13, 2024 · Aug 13, 2024
diff --git a/coderd/database/awsiamrds/awsiamrds.go b/coderd/database/awsiamrds/awsiamrds.go
@@ -10,15 +10,21 @@ import (
 	"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"
+	"github.com/lib/pq"
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/coderd/database"
 )
 
 type awsIamRdsDriver struct {
 	parent driver.Driver
 	cfg    aws.Config
 }
 
-var _ driver.Driver = &awsIamRdsDriver{}
+var (
+	_ driver.Driver             = &awsIamRdsDriver{}
+	_ database.ConnectorCreator = &awsIamRdsDriver{}
+)
 
 // Register initializes and registers our aws iam rds wrapped database driver.
 func Register(ctx context.Context, parentName string) (string, error) {
@@ -65,6 +71,16 @@ func (d *awsIamRdsDriver) Open(name string) (driver.Conn, error) {
 	return conn, nil
 }
 
+// Connector returns a driver.Connector that fetches a new authentication token for each connection.
+func (d *awsIamRdsDriver) Connector(name string) (driver.Connector, error) {
+	connector := &connector{
+		url: name,
+		cfg: d.cfg,
+	}
+
+	return connector, nil
+}
+
 func getAuthenticatedURL(cfg aws.Config, dbURL string) (string, error) {
 	nURL, err := url.Parse(dbURL)
 	if err != nil {
@@ -82,3 +98,37 @@ func getAuthenticatedURL(cfg aws.Config, dbURL string) (string, error) {
 
 	return nURL.String(), nil
 }
+
+type connector struct {
+	url    string
+	cfg    aws.Config
+	dialer pq.Dialer
+}
+
+var _ database.DialerConnector = &connector{}
+
+func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
+	nURL, err := getAuthenticatedURL(c.cfg, c.url)
+	if err != nil {
+		return nil, xerrors.Errorf("assigning authentication token to url: %w", err)
+	}
+
+	nc, err := pq.NewConnector(nURL)
+	if err != nil {
+		return nil, xerrors.Errorf("creating new connector: %w", err)
+	}
+
+	if c.dialer != nil {
+		nc.Dialer(c.dialer)
+	}
+
+	return nc.Connect(ctx)
+}
+
+func (*connector) Driver() driver.Driver {
+	return &pq.Driver{}
+}
+
+func (c *connector) Dialer(dialer pq.Dialer) {
+	c.dialer = dialer
+}
diff --git a/coderd/database/awsiamrds/awsiamrds_test.go b/coderd/database/awsiamrds/awsiamrds_test.go
@@ -7,10 +7,11 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"cdr.dev/slog"
 	"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/coderd/database/awsiamrds"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/testutil"
 )
 
@@ -22,13 +23,15 @@ func TestDriver(t *testing.T) {
 	// export DBAWSIAMRDS_TEST_URL="postgres://user@host:5432/dbname";
 	url := os.Getenv("DBAWSIAMRDS_TEST_URL")
 	if url == "" {
+		t.Log("skipping test; no DBAWSIAMRDS_TEST_URL set")
 		t.Skip()
 	}
 
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
 	defer cancel()
 
-	sqlDriver, err := awsrdsiam.Register(ctx, "postgres")
+	sqlDriver, err := awsiamrds.Register(ctx, "postgres")
 	require.NoError(t, err)
 
 	db, err := cli.ConnectToPostgres(ctx, slogtest.Make(t, nil), sqlDriver, url)
@@ -47,4 +50,23 @@ func TestDriver(t *testing.T) {
 	var one int
 	require.NoError(t, i.Scan(&one))
 	require.Equal(t, 1, one)
+
+	ps, err := pubsub.New(ctx, logger, db, url)
+	require.NoError(t, err)
+
+	gotChan := make(chan struct{})
+	subCancel, err := ps.Subscribe("test", func(_ context.Context, _ []byte) {
+		close(gotChan)
+	})
+	defer subCancel()
+	require.NoError(t, err)
+
+	err = ps.Publish("test", []byte("hello"))
+	require.NoError(t, err)
+
+	select {
+	case <-gotChan:
+	case <-ctx.Done():
+		require.Fail(t, "timed out waiting for message")
+	}
 }
diff --git a/coderd/database/connector.go b/coderd/database/connector.go
@@ -0,0 +1,19 @@
+package database
+
+import (
+	"context"
+	"database/sql/driver"
+
+	"github.com/lib/pq"
+)
+
+// ConnectorCreator can create a driver.Connector.
+type ConnectorCreator interface {
+	Connector(name string) (driver.Connector, error)
+}
+
+// DialerConnector can create a driver.Connector and set a pq.Dialer.
+type DialerConnector interface {
+	Connect(context.Context) (driver.Conn, error)
+	Dialer(dialer pq.Dialer)
+}
diff --git a/coderd/database/pubsub/pubsub.go b/coderd/database/pubsub/pubsub.go
@@ -3,6 +3,7 @@ package pubsub
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"errors"
 	"io"
 	"net"
@@ -15,6 +16,8 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/v2/coderd/database"
+
 	"cdr.dev/slog"
 )
 
@@ -432,9 +435,32 @@ func (p *PGPubsub) startListener(ctx context.Context, connectURL string) error {
 			// pq.defaultDialer uses a zero net.Dialer as well.
 			d: net.Dialer{},
 		}
+		connector driver.Connector
+		err       error
 	)
+
+	// Create a custom connector if the database driver supports it.
+	connectorCreator, ok := p.db.Driver().(database.ConnectorCreator)
+	if ok {
+		connector, err = connectorCreator.Connector(connectURL)
+		if err != nil {
+			return xerrors.Errorf("create custom connector: %w", err)
+		}
+	} else {
+		// use the default pq connector otherwise
+		connector, err = pq.NewConnector(connectURL)
+		if err != nil {
+			return xerrors.Errorf("create pq connector: %w", err)
+		}
+	}
+
+	// Set the dialer if the connector supports it.
+	if dc, ok := connector.(database.DialerConnector); ok {
+		dc.Dialer(dialer)
+	}
+
 	p.pgListener = pqListenerShim{
-		Listener: pq.NewDialListener(dialer, connectURL, time.Second, time.Minute, func(t pq.ListenerEventType, err error) {
+		Listener: pq.NewConnectorListener(connector, connectURL, time.Second, time.Minute, func(t pq.ListenerEventType, err error) {
 			switch t {
 			case pq.ListenerEventConnected:
 				p.logger.Info(ctx, "pubsub connected to postgres")
@@ -583,8 +609,8 @@ func (p *PGPubsub) Collect(metrics chan<- prometheus.Metric) {
 }
 
 // New creates a new Pubsub implementation using a PostgreSQL connection.
-func New(startCtx context.Context, logger slog.Logger, database *sql.DB, connectURL string) (*PGPubsub, error) {
-	p := newWithoutListener(logger, database)
+func New(startCtx context.Context, logger slog.Logger, db *sql.DB, connectURL string) (*PGPubsub, error) {
+	p := newWithoutListener(logger, db)
 	if err := p.startListener(startCtx, connectURL); err != nil {
 		return nil, err
 	}
@@ -594,11 +620,11 @@ func New(startCtx context.Context, logger slog.Logger, database *sql.DB, connect
 }
 
 // newWithoutListener creates a new PGPubsub without creating the pqListener.
-func newWithoutListener(logger slog.Logger, database *sql.DB) *PGPubsub {
+func newWithoutListener(logger slog.Logger, db *sql.DB) *PGPubsub {
 	return &PGPubsub{
 		logger:          logger,
 		listenDone:      make(chan struct{}),
-		db:              database,
+		db:              db,
 		queues:          make(map[string]map[uuid.UUID]*msgQueue),
 		latencyMeasurer: NewLatencyMeasurer(logger.Named("latency-measurer")),
 

diff --git a/flake.nix b/flake.nix
@@ -117,7 +117,7 @@
             name = "coder-${osArch}";
             # Updated with ./scripts/update-flake.sh`.
             # This should be updated whenever go.mod changes!
-            vendorHash = "sha256-AZ0qzh7H+UwnZNyg2iaNMSUWlGgomI/mo70T+FdF7ws=";
+            vendorHash = "sha256-0tsJY4/afWbsaFgq4j0F4kUSVXsWCQOwCk2VeU32vfM=";
             proxyVendor = true;
             src = ./.;
             nativeBuildInputs = with pkgs; [ getopt openssl zstd ];

diff --git a/go.mod b/go.mod
@@ -62,6 +62,11 @@ replace github.com/imulab/go-scim/pkg/v2 => github.com/coder/go-scim/pkg/v2 v2.0
 // Fixes https://github.com/coder/coder/issues/6685
 replace github.com/pkg/sftp => github.com/mafredri/sftp v1.13.6-0.20231212144145-8218e927edb0
 
+// Adds support for a new Listener from a driver.Connector
+// This lets us use rotating authentication tokens for passwords in connection strings
+// which we use in the awsiamrds package.
+replace github.com/lib/pq => github.com/coder/pq v1.10.5-0.20240813183442-0c420cb5a048
+
 require (
 	cdr.dev/slog v1.6.2-0.20240126064726-20367d4aede6
 	cloud.google.com/go/compute/metadata v0.5.0

diff --git a/go.sum b/go.sum
@@ -215,6 +215,8 @@ github.com/coder/go-httpstat v0.0.0-20230801153223-321c88088322 h1:m0lPZjlQ7vdVp
 github.com/coder/go-httpstat v0.0.0-20230801153223-321c88088322/go.mod h1:rOLFDDVKVFiDqZFXoteXc97YXx7kFi9kYqR+2ETPkLQ=
 github.com/coder/go-scim/pkg/v2 v2.0.0-20230221055123-1d63c1222136 h1:0RgB61LcNs24WOxc3PBvygSNTQurm0PYPujJjLLOzs0=
 github.com/coder/go-scim/pkg/v2 v2.0.0-20230221055123-1d63c1222136/go.mod h1:VkD1P761nykiq75dz+4iFqIQIZka189tx1BQLOp0Skc=
+github.com/coder/pq v1.10.5-0.20240813183442-0c420cb5a048 h1:3jzYUlGH7ZELIH4XggXhnTnP05FCYiAFeQpoN+gNR5I=
+github.com/coder/pq v1.10.5-0.20240813183442-0c420cb5a048/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/coder/pretty v0.0.0-20230908205945-e89ba86370e0 h1:3A0ES21Ke+FxEM8CXx9n47SZOKOpgSE1bbJzlE4qPVs=
 github.com/coder/pretty v0.0.0-20230908205945-e89ba86370e0/go.mod h1:5UuS2Ts+nTToAMeOjNlnHFkPahrtDkmpydBen/3wgZc=
 github.com/coder/quartz v0.1.0 h1:cLL+0g5l7xTf6ordRnUMMiZtRE8Sq5LxpghS63vEXrQ=
@@ -670,8 +672,6 @@ github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80/go.mod h1:imJHygn/1
 github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII=
 github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
 github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
-github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
-github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
 github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/mafredri/sftp v1.13.6-0.20231212144145-8218e927edb0 h1:lG2o/EWMEOlV/RfQrf3zYfQStjnUj0Mg2gmbcBcoxFI=