Verify Forgejo PostgreSQL operations and preserve tracing hook context

Assistant: codex
Assistant-Model: gpt-6.1-sol
Assistant-Session: 01a14bfc-9643-7872-aa12-98326dcf25e7

build/forgejo/forgejo-security-extension.patch

@@ -922,6 +922,388 @@
 +	}
 +}
 --- /dev/null
++++ b/models/auth/source_operation_postgres_export_test.go
+@@ -0,0 +1,7 @@
++//go:build source_operation_postgres
++
++package auth
++
++// Test-only seam: the production transaction body still obtains its real locks
++// and private query context. No runtime API or authority callback is added.
++var SourceOperationPostgresTransactionForTest = sourceOperationTx
+--- /dev/null
++++ b/models/auth/source_operation_postgres_test.go
+@@ -0,0 +1,369 @@
++//go:build source_operation_postgres
++
++package auth_test
++
++import (
++	"bytes"
++	"context"
++	"errors"
++	"fmt"
++	"os"
++	"runtime/trace"
++	"strings"
++	"sync"
++	"testing"
++	"time"
++
++	auth "forgejo.org/models/auth"
++	"forgejo.org/models/db"
++	user_model "forgejo.org/models/user"
++	"forgejo.org/modules/log"
++	"forgejo.org/modules/setting"
++	oauth2 "forgejo.org/services/auth/source/oauth2"
++
++	"github.com/lib/pq"
++	"github.com/stretchr/testify/require"
++	"xorm.io/xorm"
++)
++
++type postgresOperationLog struct {
++	log.Logger
++	mu      sync.Mutex
++	entries []string
++}
++
++func (l *postgresOperationLog) Log(_ int, _ log.Level, format string, values ...any) {
++	l.mu.Lock()
++	defer l.mu.Unlock()
++	l.entries = append(l.entries, fmt.Sprintf(format, values...))
++}
++func (l *postgresOperationLog) take() string {
++	l.mu.Lock()
++	defer l.mu.Unlock()
++	result := strings.Join(l.entries, "\n")
++	l.entries = nil
++	return result
++}
++
++// The real Forgejo engine/driver, registered model schemas and public operation
++// functions run. Only the database, application configuration and actors are
++// synthetic. No deployed owner database or authenticated operation is claimed.
++func postgresOperationDB(t *testing.T, schema string) (context.Context, *xorm.Engine, *xorm.Engine) {
++	t.Helper()
++	require.Equal(t, "/run/source-operation-postgres", os.Getenv("CUSTODIAN_PG_SOCKET"), "explicit isolated socket required")
++	oldSettings, oldContext := setting.Database, db.DefaultContext
++	oldEngine, ok := db.GetEngine(oldContext).(*xorm.Engine)
++	require.True(t, ok)
++	expectedSchema := schema
++	if expectedSchema == "" {
++		expectedSchema = "public"
++	}
++	raw, err := xorm.NewEngine("postgres", "host=/run/source-operation-postgres user=custodian_test dbname=custodian_source sslmode=disable search_path="+expectedSchema)
++	require.NoError(t, err)
++	_, err = raw.Exec(`CREATE SCHEMA IF NOT EXISTS "` + expectedSchema + `"`)
++	require.NoError(t, err)
++	setting.Database.Type = "postgres"
++	setting.Database.Host = "/run/source-operation-postgres"
++	setting.Database.Name = "custodian_source"
++	setting.Database.User = "custodian_test"
++	setting.Database.Passwd = ""
++	setting.Database.Schema = schema
++	setting.Database.SSLMode = "disable"
++	setting.Database.LogSQL = false
++	setting.Database.MaxOpenConns, setting.Database.MaxIdleConns = 8, 2
++	setting.Database.SlowQueryThreshold = 0
++	require.NoError(t, db.InitEngine(context.Background()))
++	engine, ok := db.GetEngine(db.DefaultContext).(*xorm.Engine)
++	require.True(t, ok)
++	t.Cleanup(func() {
++		_ = engine.Close()
++		_ = raw.Close()
++		setting.Database = oldSettings
++		db.SetDefaultEngine(oldContext, oldEngine)
++	})
++	version, err := engine.Query("SELECT current_setting('server_version_num')::integer / 10000 AS major, current_schema() AS schema")
++	require.NoError(t, err)
++	require.Equal(t, "16", string(version[0]["major"]))
++	require.Equal(t, expectedSchema, string(version[0]["schema"]))
++	require.NoError(t, db.SyncAllTables())
++	auth.RegisterTypeConfig(auth.OAuth2, new(oauth2.Source))
++	return db.DefaultContext, engine, raw
++}
++
++func resetPostgresOperations(t *testing.T, engine *xorm.Engine) {
++	t.Helper()
++	_, err := engine.Exec(`DROP TRIGGER IF EXISTS custodian_source_rewrite ON login_source`)
++	require.NoError(t, err)
++	_, err = engine.Exec(`TRUNCATE login_source_operation, login_source, "user", external_login_user RESTART IDENTITY CASCADE`)
++	require.NoError(t, err)
++}
++
++func TestSourceOperationPostgres(t *testing.T) {
++	for _, profile := range []struct{ name, schema string }{{"default", ""}, {"public", "public"}, {"named", "custodian_named"}} {
++		t.Run(profile.name, func(t *testing.T) {
++			ctx, engine, raw := postgresOperationDB(t, profile.schema)
++			run := func(name string, f func(*testing.T)) {
++				t.Run(name, func(t *testing.T) { resetPostgresOperations(t, engine); f(t) })
++			}
++			run("create-reconcile-inverse", func(t *testing.T) {
++				before := operationRows(t, ctx)
++				_, err := auth.PrepareSourceOperation(ctx, "postgres-create-operation", "postgres-create-case", 0, operationSource())
++				require.NoError(t, err)
++				require.Equal(t, before, operationRows(t, ctx))
++				r, err := auth.ApplySourceOperation(ctx, "postgres-create-operation", "postgres-create-case")
++				require.NoError(t, err)
++				require.Equal(t, "applied", r.Phase)
++				require.Positive(t, r.SourceID)
++				post := operationRows(t, ctx)
++				again, err := auth.ApplySourceOperation(ctx, "postgres-create-operation", "postgres-create-case")
++				require.NoError(t, err)
++				require.Equal(t, r, again)
++				require.Equal(t, post, operationRows(t, ctx))
++				_, err = auth.InverseSourceOperation(ctx, "postgres-create-operation", "postgres-wrong-case")
++				require.ErrorIs(t, err, auth.ErrSourceOperationConflict)
++				r, err = auth.ReconcileSourceOperation(ctx, "postgres-create-operation", "postgres-create-case")
++				require.NoError(t, err)
++				require.Equal(t, "applied", r.Phase)
++				r, err = auth.InverseSourceOperation(ctx, "postgres-create-operation", "postgres-create-case")
++				require.NoError(t, err)
++				require.Equal(t, "reverted", r.Phase)
++				require.Equal(t, before, operationRows(t, ctx))
++				_, err = auth.ApplySourceOperation(ctx, "postgres-create-operation", "postgres-create-case")
++				require.ErrorIs(t, err, auth.ErrSourceOperationConflict)
++			})
++			run("exact-raw-preexisting-restoration", func(t *testing.T) {
++				s := operationSource()
++				require.NoError(t, auth.CreateSource(ctx, s))
++				_, err := engine.Exec(`UPDATE login_source SET cfg = ? WHERE id = ?`, `{"Provider":"openidConnect","ClientID":"synthetic-client","ClientSecret":"synthetic-private-original","FutureGuard":{"x":true}}`, s.ID)
++				require.NoError(t, err)
++				before := operationRows(t, ctx)
++				s.Cfg.(*oauth2.Source).ClientSecret = "synthetic-private-new"
++				_, err = auth.PrepareSourceOperation(ctx, "postgres-restore-operation", "postgres-restore-case", s.ID, s)
++				require.NoError(t, err)
++				_, err = auth.ApplySourceOperation(ctx, "postgres-restore-operation", "postgres-restore-case")
++				require.NoError(t, err)
++				_, err = auth.InverseSourceOperation(ctx, "postgres-restore-operation", "postgres-restore-case")
++				require.NoError(t, err)
++				require.Equal(t, before, operationRows(t, ctx))
++			})
++			run("unchanged-preexisting-survives", func(t *testing.T) {
++				s := operationSource()
++				require.NoError(t, auth.CreateSource(ctx, s))
++				before := operationRows(t, ctx)
++				_, err := auth.PrepareSourceOperation(ctx, "postgres-unchanged-operation", "postgres-unchanged-case", s.ID, s)
++				require.NoError(t, err)
++				r, err := auth.ApplySourceOperation(ctx, "postgres-unchanged-operation", "postgres-unchanged-case")
++				require.NoError(t, err)
++				require.Equal(t, "unchanged", r.Phase)
++				_, err = auth.InverseSourceOperation(ctx, "postgres-unchanged-operation", "postgres-unchanged-case")
++				require.NoError(t, err)
++				require.Equal(t, before, operationRows(t, ctx))
++			})
++			for _, stage := range []string{"before-apply", "before-inverse"} {
++				run("private-drift-"+stage, func(t *testing.T) {
++					s := operationSource()
++					require.NoError(t, auth.CreateSource(ctx, s))
++					_, err := auth.PrepareSourceOperation(ctx, "postgres-drift-operation", "postgres-drift-case", s.ID, s)
++					require.NoError(t, err)
++					if stage == "before-inverse" {
++						_, err = auth.ApplySourceOperation(ctx, "postgres-drift-operation", "postgres-drift-case")
++						require.NoError(t, err)
++					}
++					_, err = engine.Exec(`UPDATE login_source SET cfg = cfg || ' ' WHERE id = ?`, s.ID)
++					require.NoError(t, err)
++					before := operationRows(t, ctx)
++					if stage == "before-apply" {
++						_, err = auth.ApplySourceOperation(ctx, "postgres-drift-operation", "postgres-drift-case")
++					} else {
++						_, err = auth.InverseSourceOperation(ctx, "postgres-drift-operation", "postgres-drift-case")
++					}
++					require.ErrorIs(t, err, auth.ErrSourceOperationConflict)
++					require.Equal(t, before, operationRows(t, ctx))
++				})
++			}
++			for _, kind := range []string{"user", "external"} {
++				run("linked-"+kind+"-refuses-deletion", func(t *testing.T) {
++					_, err := auth.PrepareSourceOperation(ctx, "postgres-linked-operation", "postgres-linked-case", 0, operationSource())
++					require.NoError(t, err)
++					r, err := auth.ApplySourceOperation(ctx, "postgres-linked-operation", "postgres-linked-case")
++					require.NoError(t, err)
++					if kind == "user" {
++						_, err = engine.Insert(&user_model.User{ID: 1, LowerName: "synthetic", Name: "synthetic", Email: "synthetic@invalid.test", LoginSource: r.SourceID})
++					} else {
++						_, err = engine.Exec(`INSERT INTO external_login_user (external_id, user_id, login_source_id) VALUES ('synthetic-external', 1, ?)`, r.SourceID)
++					}
++					require.NoError(t, err)
++					before := operationRows(t, ctx)
++					_, err = auth.InverseSourceOperation(ctx, "postgres-linked-operation", "postgres-linked-case")
++					require.True(t, auth.IsErrSourceInUse(err))
++					require.Equal(t, before, operationRows(t, ctx))
++					var count int64
++					if kind == "user" {
++						count, err = engine.Table("user").Where("login_source = ?", r.SourceID).Count()
++					} else {
++						count, err = engine.Table("external_login_user").Where("login_source_id = ?", r.SourceID).Count()
++					}
++					require.NoError(t, err)
++					require.EqualValues(t, 1, count)
++				})
++			}
++			run("unexpected-trigger-post-state-rolls-back", func(t *testing.T) {
++				before := operationRows(t, ctx)
++				_, err := auth.PrepareSourceOperation(ctx, "postgres-trigger-operation", "postgres-trigger-case", 0, operationSource())
++				require.NoError(t, err)
++				_, err = engine.Exec(`CREATE OR REPLACE FUNCTION custodian_rewrite_source() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN NEW.name := 'unexpected-rewrite'; RETURN NEW; END $$`)
++				require.NoError(t, err)
++				_, err = engine.Exec(`CREATE TRIGGER custodian_source_rewrite BEFORE INSERT ON login_source FOR EACH ROW EXECUTE FUNCTION custodian_rewrite_source()`)
++				require.NoError(t, err)
++				_, err = auth.ApplySourceOperation(ctx, "postgres-trigger-operation", "postgres-trigger-case")
++				require.ErrorIs(t, err, auth.ErrSourceOperationConflict)
++				require.Equal(t, before, operationRows(t, ctx))
++				r, err := auth.ReconcileSourceOperation(ctx, "postgres-trigger-operation", "postgres-trigger-case")
++				require.NoError(t, err)
++				require.Equal(t, "prepared", r.Phase)
++			})
++			run("driver-error-rollback-and-private-diagnostics", func(t *testing.T) {
++				captured := new(postgresOperationLog)
++				engine.AddHook(&db.SlowQueryHook{Logger: captured})
++				engine.AddHook(&db.ErrorQueryHook{Logger: captured})
++				const control = "synthetic-postgres-unmarked-control"
++				_, err := db.GetEngine(ctx).Exec("SELECT missing_synthetic_function(?)", control)
++				require.Error(t, err)
++				require.Contains(t, captured.take(), control, "actual unmarked diagnostic control must be observable")
++				_, err = engine.Exec(`CREATE OR REPLACE FUNCTION custodian_fail_source() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION '%', NEW.cfg; END $$`)
++				require.NoError(t, err)
++				_, err = engine.Exec(`CREATE TRIGGER custodian_source_rewrite BEFORE INSERT ON login_source FOR EACH ROW EXECUTE FUNCTION custodian_fail_source()`)
++				require.NoError(t, err)
++				before := operationRows(t, ctx)
++				_, err = auth.PrepareSourceOperation(ctx, "postgres-error-operation", "postgres-error-case", 0, operationSource())
++				require.NoError(t, err)
++				captured.take()
++				var tracing bytes.Buffer
++				require.NoError(t, trace.Start(&tracing))
++				_, err = auth.ApplySourceOperation(ctx, "postgres-error-operation", "postgres-error-case")
++				trace.Stop()
++				require.ErrorIs(t, err, auth.ErrSourceOperationStorage)
++				messages := captured.take()
++				require.Contains(t, messages, "[Error SQL Query] [owner-confined query]")
++				require.Contains(t, messages, "[Slow SQL Query] [owner-confined query]")
++				require.NotContains(t, messages, "synthetic-private-original")
++				require.False(t, bytes.Contains(tracing.Bytes(), []byte("synthetic-private-original")), "actual PostgreSQL trace disclosed private cfg")
++				require.Equal(t, before, operationRows(t, ctx))
++				r, err := auth.ReconcileSourceOperation(ctx, "postgres-error-operation", "postgres-error-case")
++				require.NoError(t, err)
++				require.Equal(t, "prepared", r.Phase)
++			})
++			run("uncooperative-writers-excluded", func(t *testing.T) {
++				locked, release, done := make(chan struct{}), make(chan struct{}), make(chan error, 1)
++				var once sync.Once
++				releaseLock := func() { once.Do(func() { close(release) }) }
++				defer releaseLock()
++				go func() {
++					done <- auth.SourceOperationPostgresTransactionForTest(ctx, func(context.Context) error { close(locked); <-release; return nil })
++				}()
++				select {
++				case <-locked:
++				case err := <-done:
++					t.Fatalf("lock acquisition refused: %v", err)
++				case <-time.After(5 * time.Second):
++					t.Fatal("lock acquisition timed out")
++				}
++				statements := []string{`UPDATE login_source_operation SET phase=phase WHERE operation_id='absent'`, `UPDATE login_source SET name=name WHERE id=-1`, `UPDATE "user" SET login_source=login_source WHERE id=-1`, `UPDATE external_login_user SET login_source_id=login_source_id WHERE user_id=-1`}
++				for _, statement := range statements {
++					session := raw.NewSession()
++					require.NoError(t, session.Begin())
++					_, err := session.Exec(`SET LOCAL lock_timeout='150ms'`)
++					require.NoError(t, err)
++					_, err = session.Exec(statement)
++					var pgerr *pq.Error
++					blocked := errors.As(err, &pgerr) && pgerr.Code == "55P03"
++					_ = session.Rollback()
++					_ = session.Close()
++					require.True(t, blocked, "an uncooperative writer escaped the actual table lock")
++				}
++				releaseLock()
++				select {
++				case err := <-done:
++					require.NoError(t, err)
++				case <-time.After(5 * time.Second):
++					t.Fatal("transaction did not finish")
++				}
++				for _, statement := range statements {
++					_, err := raw.Exec(statement)
++					require.NoError(t, err)
++				}
++			})
++			run("concurrent-preparation-has-one-owner", func(t *testing.T) {
++				start, done := make(chan struct{}), make(chan error, 2)
++				for _, id := range []string{"postgres-concurrent-first", "postgres-concurrent-second"} {
++					go func(id string) {
++						<-start
++						_, err := auth.PrepareSourceOperation(ctx, id, "postgres-concurrent-case", 0, operationSource())
++						done <- err
++					}(id)
++				}
++				close(start)
++				ok, conflict := 0, 0
++				for range 2 {
++					select {
++					case err := <-done:
++						if err == nil {
++							ok++
++						} else if errors.Is(err, auth.ErrSourceOperationConflict) {
++							conflict++
++						} else {
++							t.Fatal("unexpected concurrent preparation error")
++						}
++					case <-time.After(5 * time.Second):
++						t.Fatal("concurrent preparation timed out")
++					}
++				}
++				require.Equal(t, 1, ok)
++				require.Equal(t, 1, conflict)
++			})
++		})
++	}
++}
++
++func TestSourceOperationPostgresProcessRecovery(t *testing.T) {
++	step := os.Getenv("CUSTODIAN_PG_STEP")
++	require.NotEmpty(t, step, "explicit synthetic recovery phase required")
++	ctx, engine, _ := postgresOperationDB(t, "custodian_recovery")
++	switch step {
++	case "prepare":
++		resetPostgresOperations(t, engine)
++		_, err := auth.PrepareSourceOperation(ctx, "postgres-process-operation", "postgres-process-case", 0, operationSource())
++		require.NoError(t, err)
++	case "crash":
++		engine.AddHook(operationCrashAfterSourceInsert{})
++		_, _ = auth.ApplySourceOperation(ctx, "postgres-process-operation", "postgres-process-case")
++		t.Fatal("actual source insert did not reach the process-death hook")
++	case "recover-prepared":
++		r, err := auth.ReconcileSourceOperation(ctx, "postgres-process-operation", "postgres-process-case")
++		require.NoError(t, err)
++		require.Equal(t, "prepared", r.Phase)
++		n, err := engine.Table("login_source").Count()
++		require.NoError(t, err)
++		require.Zero(t, n)
++		_, err = auth.ApplySourceOperation(ctx, "postgres-process-operation", "postgres-process-case")
++		require.NoError(t, err)
++	case "recover-applied":
++		r, err := auth.ReconcileSourceOperation(ctx, "postgres-process-operation", "postgres-process-case")
++		require.NoError(t, err)
++		require.Equal(t, "applied", r.Phase)
++		n, err := engine.Table("login_source").Count()
++		require.NoError(t, err)
++		require.EqualValues(t, 1, n)
++		_, err = auth.InverseSourceOperation(ctx, "postgres-process-operation", "postgres-process-case")
++		require.NoError(t, err)
++	case "recover-inverse":
++		r, err := auth.ReconcileSourceOperation(ctx, "postgres-process-operation", "postgres-process-case")
++		require.NoError(t, err)
++		require.Equal(t, "reverted", r.Phase)
++		n, err := engine.Table("login_source").Count()
++		require.NoError(t, err)
++		require.Zero(t, n)
++	default:
++		t.Fatal("unknown synthetic recovery phase")
++	}
++}
+--- /dev/null
 +++ b/models/auth/source_operation_process_test.go
 @@ -0,0 +1,97 @@
 +// Copyright 2026 The Forgejo Authors. All rights reserved.
@@ -1276,12 +1658,15 @@
  	return x.Context(ctx)
 --- a/models/db/engine.go
 +++ b/models/db/engine.go
-@@ -330,8 +330,12 @@
+@@ -330,8 +330,15 @@
  func (TracingHook) BeforeProcess(c *contexts.ContextHook) (context.Context, error) {
  	ctx, task := trace.NewTask(c.Ctx, "sql")
  	ctx = context.WithValue(ctx, sqlTask{}, task)
 -	trace.Log(ctx, "query", c.SQL)
 -	trace.Logf(ctx, "args", "%v", c.Args)
++	// XORM passes the same ContextHook to every BeforeProcess hook. Publish
++	// the task there too, so a later hook returning c.Ctx retains it.
++	c.Ctx = ctx
 +	if isSensitiveSQL(c.Ctx) {
 +		trace.Log(ctx, "query", "[owner-confined query]")
 +	} else {
@@ -1291,7 +1676,7 @@
  	return ctx, nil
  }
  
-@@ -353,7 +357,11 @@
+@@ -353,7 +360,11 @@
  
  func (h *SlowQueryHook) AfterProcess(c *contexts.ContextHook) error {
  	if c.ExecuteTime >= h.Treshold {
@@ -1304,7 +1689,7 @@
  	}
  	return nil
  }
-@@ -370,7 +378,11 @@
+@@ -370,7 +381,11 @@
  
  func (h *ErrorQueryHook) AfterProcess(c *contexts.ContextHook) error {
  	if c.Err != nil && !errors.Is(c.Err, context.Canceled) {
@@ -1520,6 +1905,63 @@
 +		}
 +	}
 +}
+--- /dev/null
++++ b/models/db/tracing_context_test.go
+@@ -0,0 +1,54 @@
++package db
++
++import (
++	"bytes"
++	"context"
++	"runtime/trace"
++	"testing"
++
++	"xorm.io/xorm"
++	"xorm.io/xorm/contexts"
++)
++
++type followingQueryHook struct{}
++
++func (followingQueryHook) BeforeProcess(c *contexts.ContextHook) (context.Context, error) {
++	return c.Ctx, nil
++}
++func (followingQueryHook) AfterProcess(*contexts.ContextHook) error { return nil }
++
++func TestTracingContextSurvivesFollowingHooks(t *testing.T) {
++	for _, sensitive := range []bool{false, true} {
++		for _, fail := range []bool{false, true} {
++			engine, err := xorm.NewEngine("sqlite3", ":memory:")
++			if err != nil {
++				t.Fatal(err)
++			}
++			engine.AddHook(TracingHook{})
++			engine.AddHook(followingQueryHook{})
++			ctx := context.Background()
++			if sensitive {
++				ctx = WithSensitiveSQL(ctx)
++			}
++			var output bytes.Buffer
++			if err := trace.Start(&output); err != nil {
++				t.Fatal(err)
++			}
++			func() {
++				defer trace.Stop()
++				query := "SELECT ?"
++				if fail {
++					query = "SELECT missing_synthetic_function(?)"
++				}
++				_, err = engine.Context(ctx).Exec(query, "synthetic-following-hook-sentinel")
++			}()
++			_ = engine.Close()
++			if (err != nil) != fail {
++				t.Fatal("unexpected real SQLite query outcome")
++			}
++			if bytes.Contains(output.Bytes(), []byte("synthetic-following-hook-sentinel")) == sensitive {
++				t.Fatal("following hook changed private or ordinary trace disclosure")
++			}
++		}
++	}
++}
 --- a/routers/web/admin/auths.go
 +++ b/routers/web/admin/auths.go
 @@ -322,6 +322,13 @@