internal/store/postgres/auth_identity_lock.go

Rama main · réplica de lectura del repositorio original.

package postgres

import (
	"context"
	"errors"
	"fmt"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"

	"telesrv/internal/store"
	"telesrv/internal/store/postgres/sqlcgen"
)

const (
	// authIdentityAdvisoryNamespace is deliberately a two-int advisory-lock
	// namespace. PostgreSQL keeps it disjoint from the one-bigint advisory locks
	// used elsewhere in the store. AUTH in ASCII is stable and recognizable in
	// pg_locks diagnostics.
	authIdentityAdvisoryNamespace int32 = 0x41555448
	authIdentityTxMaxAttempts           = 3
)

var errAuthIdentityChanged = errors.New("auth key permanent identity changed while acquiring locks")

type authKeyIdentityHint struct {
	found       bool
	expiresAt   int
	bound       bool
	permID      int64
	identityID  int64
	hasIdentity bool
}

// withAuthIdentityTx gives identity-sensitive stores an explicit transaction
// boundary. Application-level identity visibility changes are retried after a
// savepoint/transaction rollback. PostgreSQL deadlock/serialization retries are
// only safe when this store owns the top-level transaction; an injected pgx.Tx
// is never silently replayed after 40P01/40001.
func withAuthIdentityTx(
	ctx context.Context,
	db sqlcgen.DBTX,
	op string,
	fn func(pgx.Tx) error,
) error {
	_, embedded := db.(pgx.Tx)
	var lastErr error
	for attempt := 0; attempt < authIdentityTxMaxAttempts; attempt++ {
		err := withTx(ctx, db, op, fn)
		if err == nil {
			return nil
		}
		lastErr = err
		switch {
		case errors.Is(err, errAuthIdentityChanged):
			// The attempt owns a nested savepoint even for an injected pgx.Tx,
			// so all row/advisory locks from the stale hint have been released
			// before the next READ COMMITTED statement snapshot is taken.
			continue
		case !embedded && isAuthIdentityRetryableDatabaseError(err):
			// Defensive retry only. The identity gate is the deadlock fix; this
			// does not substitute for the global lock order.
			continue
		default:
			return err
		}
	}
	return fmt.Errorf("%s did not stabilize after %d attempts: %w", op, authIdentityTxMaxAttempts, lastErr)
}

func isAuthIdentityRetryableDatabaseError(err error) bool {
	var pgErr *pgconn.PgError
	return errors.As(err, &pgErr) && (pgErr.Code == "40P01" || pgErr.Code == "40001")
}

// lockPermanentAuthIdentities acquires the complete batch before any auth-key,
// binding, authorization, or update-state row lock. Ordering is by the final
// int32 hashint8 key, not by the source bigint identity: hash collisions are
// intentionally one lock and cannot create an opposite acquisition order.
func lockPermanentAuthIdentities(ctx context.Context, tx pgx.Tx, permIDs []int64) error {
	if len(permIDs) == 0 {
		return nil
	}
	rows, err := tx.Query(ctx, `
SELECT DISTINCT hashint8(identity_id)::integer AS lock_key
FROM unnest($1::bigint[]) AS identities(identity_id)
ORDER BY lock_key`, permIDs)
	if err != nil {
		return fmt.Errorf("derive permanent auth identity lock keys: %w", err)
	}
	lockKeys := make([]int32, 0, len(permIDs))
	for rows.Next() {
		var key int32
		if err := rows.Scan(&key); err != nil {
			rows.Close()
			return fmt.Errorf("scan permanent auth identity lock key: %w", err)
		}
		lockKeys = append(lockKeys, key)
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return fmt.Errorf("iterate permanent auth identity lock keys: %w", err)
	}
	rows.Close()
	for _, key := range lockKeys {
		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1::integer, $2::integer)`, authIdentityAdvisoryNamespace, key); err != nil {
			return fmt.Errorf("lock permanent auth identity %d: %w", key, err)
		}
	}
	return nil
}

// lookupAuthKeyIdentityHint is intentionally lock-free. A positive-expiry raw
// key has no permanent identity until a binding is committed; a permanent raw
// key is its own identity. Callers must re-read after the raw row is locked.
func lookupAuthKeyIdentityHint(ctx context.Context, tx pgx.Tx, rawID int64) (authKeyIdentityHint, error) {
	var hint authKeyIdentityHint
	err := tx.QueryRow(ctx, `
/* auth_identity_hint */
SELECT key.expires_at,
       binding.temp_auth_key_id IS NOT NULL,
       COALESCE(binding.perm_auth_key_id, 0)
FROM auth_keys AS key
LEFT JOIN temp_auth_key_bindings AS binding
  ON binding.temp_auth_key_id = key.auth_key_id
WHERE key.auth_key_id = $1`, rawID).Scan(&hint.expiresAt, &hint.bound, &hint.permID)
	if errors.Is(err, pgx.ErrNoRows) {
		return authKeyIdentityHint{}, nil
	}
	if err != nil {
		return authKeyIdentityHint{}, fmt.Errorf("resolve auth key identity hint: %w", err)
	}
	hint.found = true
	switch {
	case hint.bound:
		hint.identityID = hint.permID
		hint.hasIdentity = true
	case hint.expiresAt == 0:
		hint.identityID = rawID
		hint.hasIdentity = true
	}
	return hint, nil
}

// lockRawAuthKeyInIdentityOrder establishes the only cross-identity row-lock
// order used by bind, selector advance and direct key deletion:
//
//	permanent identity advisory gate -> raw auth-key row -> permanent row
//
// If an initially-unbound temp key becomes bound before the raw lock is
// acquired, taking its newly discovered identity advisory lock at that point
// would recreate raw->identity inversion. The caller must roll back and retry.
func lockRawAuthKeyInIdentityOrder(
	ctx context.Context,
	tx pgx.Tx,
	rawID int64,
) (rawExpiry int, permID int64, bound bool, err error) {
	hint, err := lookupAuthKeyIdentityHint(ctx, tx, rawID)
	if err != nil || !hint.found {
		if err != nil {
			return 0, 0, false, err
		}
		return 0, 0, false, store.ErrAuthKeyNotFound
	}
	if hint.hasIdentity {
		if err := lockPermanentAuthIdentities(ctx, tx, []int64{hint.identityID}); err != nil {
			return 0, 0, false, err
		}
	}
	if err := tx.QueryRow(ctx, `
SELECT expires_at
FROM auth_keys
WHERE auth_key_id = $1
FOR UPDATE`, rawID).Scan(&rawExpiry); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return 0, 0, false, store.ErrAuthKeyNotFound
		}
		return 0, 0, false, fmt.Errorf("lock raw auth key: %w", err)
	}

	var actualPermID int64
	err = tx.QueryRow(ctx, `
SELECT perm_auth_key_id
FROM temp_auth_key_bindings
WHERE temp_auth_key_id = $1`, rawID).Scan(&actualPermID)
	switch {
	case err == nil:
		bound = true
		permID = actualPermID
	case errors.Is(err, pgx.ErrNoRows):
		permID = rawID
	default:
		return 0, 0, false, fmt.Errorf("revalidate auth key permanent identity: %w", err)
	}

	actualHasIdentity := bound || rawExpiry == 0
	actualIdentityID := permID
	if actualHasIdentity != hint.hasIdentity ||
		(actualHasIdentity && actualIdentityID != hint.identityID) ||
		bound != hint.bound || rawExpiry != hint.expiresAt {
		return 0, 0, false, errAuthIdentityChanged
	}
	if !bound {
		return rawExpiry, permID, false, nil
	}

	var permExpiry int
	if err := tx.QueryRow(ctx, `
SELECT expires_at
FROM auth_keys
WHERE auth_key_id = $1
FOR UPDATE`, permID).Scan(&permExpiry); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return 0, 0, false, store.ErrAuthKeyBindingInvalid
		}
		return 0, 0, false, fmt.Errorf("lock permanent auth key: %w", err)
	}
	if rawExpiry <= 0 || permExpiry != 0 {
		return 0, 0, false, store.ErrAuthKeyBindingInvalid
	}
	return rawExpiry, permID, true, nil
}