From 35c8c9f39175c2b867fe8c133935284c4191bcb8 Mon Sep 17 00:00:00 2001 From: codegirl-007 Date: Sat, 22 Aug 2026 12:25:16 -0700 Subject: [PATCH] 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 | 51 +++++++++++++++++++++++-------------- internal/store/postgres.go | 5 ++-- internal/web/server.go | 6 ++--- internal/web/server_test.go | 51 +++++++++++++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 24 deletions(-) diff --git a/internal/store/migrate.go b/internal/store/migrate.go index 49c0f61..abbe1b4 100644 --- a/internal/store/migrate.go +++ b/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. -// Fresh databases apply schemaSQL as version 001; later versions are incremental. +// applyMigrations runs versioned migrations under a session-level advisory lock +// 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() - if _, err := tx.Exec(`SELECT pg_advisory_xact_lock($1)`, migrateLockKey); err != nil { + defer conn.Close() + + 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) { - rows, err := db.Query(`SELECT version FROM schema_migrations`) +func appliedVersions(ctx context.Context, exec execContext) (map[string]bool, error) { + rows, err := exec.QueryContext(ctx, `SELECT version FROM schema_migrations`) if err != nil { return nil, err } diff --git a/internal/store/postgres.go b/internal/store/postgres.go index fefde24..d104baf 100644 --- a/internal/store/postgres.go +++ b/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) } } diff --git a/internal/web/server.go b/internal/web/server.go index 6d484bc..7950bc6 100644 --- a/internal/web/server.go +++ b/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 { - r.Use(middleware.RealIP) - } + // Do not use middleware.RealIP: it rewrites RemoteAddr from client-controlled + // 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 { diff --git a/internal/web/server_test.go b/internal/web/server_test.go index 34bb8a3..7419318 100644 --- a/internal/web/server_test.go +++ b/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) + } +}