fix: use authenticated urls for pubsub (#14261)

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,35 @@ 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.
+	dc, ok := connector.(database.DialerConnector)
+	if !ok {
+		p.logger.Critical(ctx, "connector does not support setting log dialer, database connection debug logs will be missing")
+	} else {
+		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 +612,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 +623,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")),
 

coderd/database/pubsub/pubsub_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
@@ -51,7 +52,7 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	event := "test"
 	data := "testing"
 	messageChannel := make(chan []byte)
-	unsub0, err := uut.Subscribe(event, func(ctx context.Context, message []byte) {
+	unsub0, err := uut.Subscribe(event, func(_ context.Context, message []byte) {
 		messageChannel <- message
 	})
 	require.NoError(t, err)
@@ -86,7 +87,7 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	for i := range colossalData {
 		colossalData[i] = 'q'
 	}
-	unsub1, err := uut.Subscribe(event, func(ctx context.Context, message []byte) {
+	unsub1, err := uut.Subscribe(event, func(_ context.Context, message []byte) {
 		messageChannel <- message
 	})
 	require.NoError(t, err)
@@ -119,3 +120,74 @@ func TestPGPubsub_Metrics(t *testing.T) {
 			!testutil.PromCounterGathered(t, metrics, "coder_pubsub_latency_measure_errs_total")
 	}, testutil.WaitShort, testutil.IntervalFast)
 }
+
+func TestPGPubsubDriver(t *testing.T) {
+	t.Parallel()
+	if !dbtestutil.WillUsePostgres() {
+		t.Skip("test only with postgres")
+	}
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	logger := slogtest.Make(t, &slogtest.Options{
+		IgnoreErrors: true,
+	}).Leveled(slog.LevelDebug)
+
+	connectionURL, closePg, err := dbtestutil.Open()
+	require.NoError(t, err)
+	defer closePg()
+
+	// use a separate subber and pubber so we can keep track of listener connections
+	db, err := sql.Open("postgres", connectionURL)
+	require.NoError(t, err)
+	pubber, err := pubsub.New(ctx, logger, db, connectionURL)
+	require.NoError(t, err)
+	defer pubber.Close()
+
+	// use a connector that sends us the connections for the subber
+	subDriver := dbtestutil.NewDriver()
+	defer subDriver.Close()
+	tconn, err := subDriver.Connector(connectionURL)
+	require.NoError(t, err)
+	tcdb := sql.OpenDB(tconn)
+	subber, err := pubsub.New(ctx, logger, tcdb, connectionURL)
+	require.NoError(t, err)
+	defer subber.Close()
+
+	// test that we can publish and subscribe
+	gotChan := make(chan struct{}, 1)
+	defer close(gotChan)
+	subCancel, err := subber.Subscribe("test", func(_ context.Context, _ []byte) {
+		gotChan <- struct{}{}
+	})
+	require.NoError(t, err)
+	defer subCancel()
+
+	// send a message
+	err = pubber.Publish("test", []byte("hello"))
+	require.NoError(t, err)
+
+	// wait for the message
+	_ = testutil.RequireRecvCtx(ctx, t, gotChan)
+
+	// read out first connection
+	firstConn := testutil.RequireRecvCtx(ctx, t, subDriver.Connections)
+
+	// drop the underlying connection being used by the pubsub
+	// the pq.Listener should reconnect and repopulate it's listeners
+	// so old subscriptions should still work
+	err = firstConn.Close()
+	require.NoError(t, err)
+
+	// wait for the reconnect
+	_ = testutil.RequireRecvCtx(ctx, t, subDriver.Connections)
+	// we need to sleep because the raw connection notification
+	// is sent before the pq.Listener can reestablish it's listeners
+	time.Sleep(1 * time.Second)
+
+	// ensure our old subscription still fires
+	err = pubber.Publish("test", []byte("hello-again"))
+	require.NoError(t, err)
+
+	// wait for the message on the old subscription
+	_ = testutil.RequireRecvCtx(ctx, t, gotChan)
+}