fix: use authenticated urls for pubsub

f0ssel · f0ssel · commit d2cd8d7cb88f · 2024-08-13T17:02:43.000Z
diff --git a/coderd/database/awsiamrds/awsiamrds.go b/coderd/database/awsiamrds/awsiamrds.go
@@ -10,7 +10,10 @@ 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 {
@@ -19,6 +22,7 @@ type awsIamRdsDriver struct {
 }
 
 var _ driver.Driver = &awsIamRdsDriver{}
+var _ 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 +69,15 @@ func (d *awsIamRdsDriver) Open(name string) (driver.Conn, error) {
 	return conn, nil
 }
 
+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 +95,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 (c *connector) Driver() driver.Driver {
+	return &pq.Driver{}
+}
+
+func (c *connector) Dialer(dialer pq.Dialer) {
+	c.dialer = dialer
+}
diff --git a/coderd/database/connector.go b/coderd/database/connector.go
@@ -0,0 +1,17 @@
+package database
+
+import (
+	"context"
+	"database/sql/driver"
+
+	"github.com/lib/pq"
+)
+
+type ConnectorCreator interface {
+	Connector(name string) (driver.Connector, error)
+}
+
+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,31 @@ 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 = pq.NewConnector(connectURL)
+		if err != nil {
+			return xerrors.Errorf("create pq connector: %w", err)
+		}
+	} else {
+		connector, err = connectorCreator.Connector(connectURL)
+		if err != nil {
+			return xerrors.Errorf("create custom 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")
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.20240813145306-1ce661cfa68d
+
 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.20240813145306-1ce661cfa68d h1:pv+JacyCHoHAr2kh6HltHdFlWqVeWHCvaQDqra5Aff4=
+github.com/coder/pq v1.10.5-0.20240813145306-1ce661cfa68d/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=
