Agent-System-Backend/internal/store/account_serialization_pg_test.go

package store

import (
	"context"
	"testing"
	"time"
)

// The trigger holds a transaction after its authorization read. A second
// connection must wait for that owner's lock before replacing credentials.
func holdPGAccountWrite(t *testing.T, s *Store, table, event, suffix string) func() {
	t.Helper()
	lock := "u7_account_" + suffix
	_, err := s.db.Exec(`CREATE FUNCTION u7_pause_account() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN PERFORM pg_advisory_xact_lock(hashtext(current_schema() || ':` + lock + `')); RETURN NEW; END $$`)
	if err != nil {
		t.Fatal(err)
	}
	_, err = s.db.Exec(`CREATE TRIGGER u7_pause_account BEFORE ` + event + ` ON ` + table + ` FOR EACH ROW EXECUTE FUNCTION u7_pause_account()`)
	if err != nil {
		t.Fatal(err)
	}
	conn, err := s.db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if _, err := conn.ExecContext(context.Background(), `SELECT pg_advisory_lock(hashtext(current_schema() || ':`+lock+`'))`); err != nil {
		conn.Close()
		t.Fatal(err)
	}
	released := false
	release := func() {
		if !released {
			released = true
			_, _ = conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock(hashtext(current_schema() || ':`+lock+`'))`)
			_ = conn.Close()
		}
	}
	t.Cleanup(release)
	return release
}

func waitForPGAccountWrite(t *testing.T, s *Store, suffix string) {
	t.Helper()
	lock := "u7_account_" + suffix
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		var n int
		err := s.db.QueryRow(`SELECT count(*) FROM pg_locks WHERE locktype='advisory' AND NOT granted
			AND objid=(hashtext(current_schema() || ':` + lock + `')::bigint & 4294967295)::oid`).Scan(&n)
		if err != nil {
			t.Fatal(err)
		}
		if n > 0 {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("account write did not reach the deterministic PostgreSQL pause")
}

func awaitPGAccountResult[T any](t *testing.T, done <-chan T) T {
	t.Helper()
	select {
	case result := <-done:
		return result
	case <-time.After(5 * time.Second):
		t.Fatal("serialized account transaction did not finish")
		var zero T
		return zero
	}
}

func TestPGVerifiedSessionAndPasswordChangeSerialize(t *testing.T) {
	s := activeAccount(t)
	if s.Dialect() != "postgres" {
		t.Skip("real PostgreSQL required")
	}
	_, version, ok := s.VerifyLoginWithVersion("ada", "old-password")
	if !ok {
		t.Fatal("initial login failed")
	}
	release := holdPGAccountWrite(t, s, "sessions", "INSERT", "session")
	type sessionResult struct {
		token string
		err   error
	}
	sessionDone := make(chan sessionResult, 1)
	go func() {
		token, err := s.CreateSessionVerified("ada", SessionScopeFull, time.Hour, version)
		sessionDone <- sessionResult{token, err}
	}()
	waitForPGAccountWrite(t, s, "session")
	hash, err := HashPassword("replacement-password")
	if err != nil {
		t.Fatal(err)
	}
	changeDone := make(chan error, 1)
	go func() {
		_, err := s.ApplyPasswordChange(AccountWrite{Owner: "ada", PasswordHash: hash, Actor: "admin"})
		changeDone <- err
	}()
	select {
	case err := <-changeDone:
		t.Fatalf("replacement committed ahead of the in-flight verified session: %v", err)
	case <-time.After(100 * time.Millisecond):
	}
	release()
	session := awaitPGAccountResult(t, sessionDone)
	if session.err != nil {
		t.Fatal(session.err)
	}
	if err := awaitPGAccountResult(t, changeDone); err != nil {
		t.Fatal(err)
	}
	if _, _, ok := s.SessionInfo(session.token); ok {
		t.Fatal("old verified session survived later password replacement")
	}
}

func TestPGRestrictedPasswordSetAndKeyReissueSerialize(t *testing.T) {
	s := openTestStore(t)
	if s.Dialect() != "postgres" {
		t.Skip("real PostgreSQL required")
	}
	if err := s.EnsureSystemIdentity("test-pg-grant-serialize", "pg"); err != nil {
		t.Fatal(err)
	}
	out, err := s.BootstrapFirstAdminWrite(AccountWrite{Owner: "ada", Name: "Ada", Actor: "test"})
	if err != nil {
		t.Fatal(err)
	}
	account, err := s.GetAccountByOwner("ada")
	if err != nil {
		t.Fatal(err)
	}
	token, _, err := s.RedeemLoginKey(out.Plain)
	if err != nil {
		t.Fatal(err)
	}
	hash, err := HashPassword("replacement-password")
	if err != nil {
		t.Fatal(err)
	}
	release := holdPGAccountWrite(t, s, "users", "UPDATE", "password")
	setDone := make(chan error, 1)
	go func() {
		_, err := s.CompleteSetPasswordChecked(token, hash)
		setDone <- err
	}()
	waitForPGAccountWrite(t, s, "password")
	issueDone := make(chan error, 1)
	go func() {
		_, err := s.InitializePendingAccount(account.ID, account.StateVersion, account.InitializationGeneration)
		issueDone <- err
	}()
	select {
	case err := <-issueDone:
		t.Fatalf("key replacement committed ahead of the in-flight password set: %v", err)
	case <-time.After(100 * time.Millisecond):
	}
	release()
	if err := awaitPGAccountResult(t, setDone); err != nil {
		t.Fatal(err)
	}
	if err := awaitPGAccountResult(t, issueDone); err != ErrConflict && err != ErrAccountForbidden {
		t.Fatalf("completed initialization was reissued: %v", err)
	}
	if _, _, ok := s.SessionInfo(token); ok {
		t.Fatal("old restricted session survived password set and key reissue")
	}
}