Fix migrate lock scope and stop RealIP from bypassing proxy trust.

Hold a session advisory lock for the full migration apply path, and remove Chi RealIP so clientIP can validate the TCP peer before walking X-Forwarded-For.

internal/store/migrate.go

@@ -1,57 +1,70 @@
 package store
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"log"
 )
 
+const migrateLockKey int64 = 0x706c756d5f6d6967 // "plum_mig"
+
 // migrateUserProfileColumns adds avatar_url and state when missing (existing DBs).
-func migrateUserProfileColumns(db *sql.DB) error {
+func migrateUserProfileColumns(ctx context.Context, exec execContext) error {
 	cols := []string{"avatar_url", "state"}
 	for _, col := range cols {
 		stmt := fmt.Sprintf(`ALTER TABLE users ADD COLUMN IF NOT EXISTS %s TEXT NOT NULL DEFAULT ''`, col)
-		if _, err := db.Exec(stmt); err != nil {
+		if _, err := exec.ExecContext(ctx, stmt); err != nil {
 			return fmt.Errorf("add column %s: %w", col, err)
 		}
 	}
 	return nil
 }
 
-const migrateLockKey int64 = 0x706c756d5f6d6967 // "plum_mig"
+type execContext interface {
+	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
+	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
+}
 
-// applyMigrations runs versioned migrations under an advisory lock.
+// applyMigrations runs versioned migrations under a session-level advisory lock
-// Fresh databases apply schemaSQL as version 001; later versions are incremental.
+// held for the entire process (check versions → apply → record).
 func applyMigrations(db *sql.DB, schemaSQL string) error {
-	tx, err := db.Begin()
+	ctx := context.Background()
+	conn, err := db.Conn(ctx)
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback()
+	defer conn.Close()
-	if _, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, migrateLockKey); err != nil {
+
+	if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1)`, migrateLockKey); err != nil {
 		return fmt.Errorf("migrate lock: %w", err)
 	}
-	if _, err := tx.Exec(`
+	defer func() {
+		if _, unlockErr := conn.ExecContext(ctx, `SELECT pg_advisory_unlock($1)`, migrateLockKey); unlockErr != nil {
+			log.Printf("migrate unlock: %v", unlockErr)
+		}
+	}()
+
+	if _, err := conn.ExecContext(ctx, `
 CREATE TABLE IF NOT EXISTS schema_migrations (
   version TEXT PRIMARY KEY,
   applied_at TIMESTAMPTZ NOT NULL DEFAULT now()
 )`); err != nil {
 		return fmt.Errorf("schema_migrations: %w", err)
 	}
-	if err := tx.Commit(); err != nil {
-		return err
-	}
 
-	applied, err := appliedVersions(db)
+	applied, err := appliedVersions(ctx, conn)
 	if err != nil {
 		return err
 	}
 
 	migrations := []struct {
 		version string
-		run     func(*sql.DB) error
+		run     func(context.Context, execContext) error
 	}{
-		{"001_schema", func(db *sql.DB) error { return applySchema(db, schemaSQL) }},
+		{"001_schema", func(ctx context.Context, exec execContext) error {
+			return applySchema(ctx, exec, schemaSQL)
+		}},
 		{"002_user_profile_columns", migrateUserProfileColumns},
 	}
 	for _, m := range migrations {
@@ -59,18 +72,18 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
 			continue
 		}
 		log.Printf("migrate: applying %s", m.version)
-		if err := m.run(db); err != nil {
+		if err := m.run(ctx, conn); err != nil {
 			return fmt.Errorf("migrate %s: %w", m.version, err)
 		}
-		if _, err := db.Exec(`INSERT INTO schema_migrations (version) VALUES ($1)`, m.version); err != nil {
+		if _, err := conn.ExecContext(ctx, `INSERT INTO schema_migrations (version) VALUES ($1)`, m.version); err != nil {
 			return fmt.Errorf("record %s: %w", m.version, err)
 		}
 	}
 	return nil
 }
 
-func appliedVersions(db *sql.DB) (map[string]bool, error) {
+func appliedVersions(ctx context.Context, exec execContext) (map[string]bool, error) {
-	rows, err := db.Query(`SELECT version FROM schema_migrations`)
+	rows, err := exec.QueryContext(ctx, `SELECT version FROM schema_migrations`)
 	if err != nil {
 		return nil, err
 	}

internal/store/postgres.go

@@ -1,6 +1,7 @@
 package store
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"net/url"
@@ -11,13 +12,13 @@ import (
 )
 
 // applySchema runs semicolon-separated DDL statements.
-func applySchema(db *sql.DB, schema string) error {
+func applySchema(ctx context.Context, exec execContext, schema string) error {
 	for _, stmt := range strings.Split(schema, ";") {
 		stmt = strings.TrimSpace(stmt)
 		if stmt == "" {
 			continue
 		}
-		if _, err := db.Exec(stmt); err != nil {
+		if _, err := exec.ExecContext(ctx, stmt); err != nil {
 			return fmt.Errorf("%w: %s", err, stmt)
 		}
 	}

internal/web/server.go

@@ -148,9 +148,9 @@ func New(st store.Store, sessionStore scs.Store, templateFS fs.FS, staticFS fs.F
 func (s *Server) Handler() http.Handler {
 	r := chi.NewRouter()
 	r.Use(middleware.RequestID)
-	if len(s.cfg.TrustedProxies) > 0 {
+	// Do not use middleware.RealIP: it rewrites RemoteAddr from client-controlled
-		r.Use(middleware.RealIP)
+	// forwarding headers before clientIP can validate the TCP peer against
-	}
+	// TrustedProxies. clientIP walks X-Forwarded-For itself when the peer is trusted.
 	r.Use(middleware.Logger)
 	r.Use(middleware.Recoverer)
 	r.Use(func(next http.Handler) http.Handler {

internal/web/server_test.go

@@ -6,10 +6,12 @@ import (
 	"image"
 	"image/png"
 	"mime/multipart"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/alexedwards/scs/v2/memstore"
 	"github.com/google/uuid"
@@ -663,3 +665,52 @@ func csrfFrom(html string) string {
 	}
 	return html[:j]
 }
+
+// TestRegisterThrottleUsesTCPPeerThroughRouter ensures forged X-Forwarded-For
+// cannot bypass rate limits when the direct peer is outside TrustedProxies.
+// This must go through Handler() so middleware ordering bugs are caught.
+func TestRegisterThrottleUsesTCPPeerThroughRouter(t *testing.T) {
+	_, proxyNet, err := net.ParseCIDR("10.0.0.0/8")
+	if err != nil {
+		t.Fatal(err)
+	}
+	srv, _ := newTestServer(t, Config{TrustedProxies: []*net.IPNet{proxyNet}})
+	// Tight window so the test stays fast.
+	srv.registerIP = newThrottle(3, time.Minute, 100)
+	h := srv.Handler()
+
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/register", nil))
+	cookies := rec.Result().Cookies()
+	csrf := csrfFrom(rec.Body.String())
+	if csrf == "" {
+		t.Fatal("missing csrf")
+	}
+
+	post := func(xff string) int {
+		form := strings.NewReader("_csrf=" + csrf + "&username=ab&password=hunter22")
+		req := httptest.NewRequest(http.MethodPost, "/register", form)
+		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+		req.RemoteAddr = "203.0.113.50:9"
+		req.Header.Set("X-Forwarded-For", xff)
+		for _, c := range cookies {
+			req.AddCookie(c)
+		}
+		w := httptest.NewRecorder()
+		h.ServeHTTP(w, req)
+		return w.Code
+	}
+
+	if code := post("198.51.100.1"); code != http.StatusOK {
+		t.Fatalf("attempt 1: got %d want 200 (validation error page)", code)
+	}
+	if code := post("198.51.100.2"); code != http.StatusOK {
+		t.Fatalf("attempt 2: got %d want 200", code)
+	}
+	if code := post("198.51.100.3"); code != http.StatusOK {
+		t.Fatalf("attempt 3: got %d want 200", code)
+	}
+	if code := post("198.51.100.4"); code != http.StatusTooManyRequests {
+		t.Fatalf("forged XFF must not bypass peer throttle, got %d want 429", code)
+	}
+}