//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")
	}
}