internal/store/postgres/auth_identity_lock_integration_test.go
Rama main · réplica de lectura del repositorio original.
package postgres
import (
"context"
"errors"
"strings"
"sync"
"testing"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"telesrv/internal/domain"
"telesrv/internal/store"
)
func TestAuthIdentitySelectorRetriesUncommittedFirstBindSnapshotPostgres(t *testing.T) {
pool := testPool(t)
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
t.Cleanup(cancel)
keys := NewAuthKeyStore(pool)
bindings := NewTempAuthKeyBindingStore(pool)
expiresAt := int(time.Now().Add(time.Hour).Unix())
temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, expiresAt)
perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
binding := domain.TempAuthKeyBinding{
TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(perm),
Nonce: 8701, TempSessionID: 8702, ExpiresAt: expiresAt,
EncryptedMessage: []byte("first bind snapshot"),
}
advanceConn, err := pool.Acquire(ctx)
if err != nil {
t.Fatal(err)
}
defer advanceConn.Release()
var advancePID int
if err := advanceConn.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&advancePID); err != nil {
t.Fatal(err)
}
msgID := authKeySessionLayerTestMsgID(time.Now().UTC(), 1)
type advanceResult struct {
value store.AuthKeySessionLayer
applied bool
err error
}
// Stage the first binding but do not commit it. Save holds the permanent
// identity gate and raw row, so the selector sees the old unbound hint and
// then waits on the raw row inside the server-side advance function.
bindTx, err := pool.Begin(ctx)
if err != nil {
t.Fatal(err)
}
defer func() { _ = bindTx.Rollback(context.Background()) }()
if err := NewTempAuthKeyBindingStore(bindTx).Save(ctx, binding); err != nil {
t.Fatalf("stage first bind: %v", err)
}
result := make(chan advanceResult, 1)
go func() {
value, applied, err := NewAuthKeyStore(advanceConn).AdvanceSessionLayer(ctx, temp, 8703, 227, msgID)
result <- advanceResult{value: value, applied: applied, err: err}
}()
waitForPostgresBackendLockWait(t, ctx, pool, advancePID)
if err := bindTx.Commit(ctx); err != nil {
t.Fatalf("commit first bind: %v", err)
}
got := <-result
if got.err != nil || !got.applied || got.value.Layer != 227 || !got.value.SharedDefault {
t.Fatalf("advance after identity retry = (%+v,%v,%v)", got.value, got.applied, got.err)
}
assertTempIdentityBinding(t, ctx, bindings, binding)
for _, id := range [][8]byte{temp, perm} {
stored, found, err := keys.Get(ctx, id)
if err != nil || !found || stored.Layer != 227 || stored.LayerObservationID != got.value.ObservationID {
t.Fatalf("shared tuple %x = (%+v,%v,%v)", id, stored, found, err)
}
}
}
func TestAuthIdentitySelectorSerializesWithPermanentRevocationAndDeletePostgres(t *testing.T) {
for _, op := range []string{"revoke", "delete"} {
t.Run(op, func(t *testing.T) {
pool := testPool(t)
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
t.Cleanup(cancel)
keys := NewAuthKeyStore(pool)
auths := NewAuthorizationStore(pool)
bindings := NewTempAuthKeyBindingStore(pool)
expiresAt := int(time.Now().Add(time.Hour).Unix())
temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, expiresAt)
perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
userID := createRevokeTestUser(t, ctx, pool, "selector-"+op)
hash := int64(8800)
if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: perm, UserID: userID, Hash: hash}); err != nil {
t.Fatal(err)
}
if err := bindings.Save(ctx, domain.TempAuthKeyBinding{
TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(perm), ExpiresAt: expiresAt,
EncryptedMessage: []byte("identity serialization"),
}); err != nil {
t.Fatal(err)
}
blocker, err := pool.Begin(ctx)
if err != nil {
t.Fatal(err)
}
defer func() { _ = blocker.Rollback(context.Background()) }()
if err := lockPermanentAuthIdentities(ctx, blocker, []int64{authKeyIDToInt64(perm)}); err != nil {
t.Fatal(err)
}
selectorConn, err := pool.Acquire(ctx)
if err != nil {
t.Fatal(err)
}
defer selectorConn.Release()
opConn, err := pool.Acquire(ctx)
if err != nil {
t.Fatal(err)
}
defer opConn.Release()
var selectorPID, opPID int
if err := selectorConn.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&selectorPID); err != nil {
t.Fatal(err)
}
if err := opConn.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&opPID); err != nil {
t.Fatal(err)
}
selectorResult := make(chan error, 1)
go func() {
_, _, err := NewAuthKeyStore(selectorConn).AdvanceSessionLayer(
ctx, temp, 8801, 227, authKeySessionLayerTestMsgID(time.Now().UTC(), 1),
)
selectorResult <- err
}()
waitForPostgresBackendLockWait(t, ctx, pool, selectorPID)
opResult := make(chan error, 1)
go func() {
if op == "revoke" {
_, found, err := NewAuthorizationStore(opConn).RevokeByHash(ctx, userID, hash)
if err == nil && !found {
err = errors.New("revoke target disappeared")
}
opResult <- err
return
}
opResult <- NewAuthKeyStore(opConn).Delete(ctx, perm)
}()
waitForPostgresBackendLockWait(t, ctx, pool, opPID)
if err := blocker.Commit(ctx); err != nil {
t.Fatal(err)
}
selectorErr := <-selectorResult
if selectorErr != nil &&
!errors.Is(selectorErr, store.ErrAuthKeyNotFound) &&
!errors.Is(selectorErr, store.ErrAuthKeyBindingInvalid) {
t.Fatalf("selector error = %v", selectorErr)
}
if err := <-opResult; err != nil {
t.Fatalf("%s error = %v", op, err)
}
assertRevokeTestNoAuthorization(t, ctx, auths, perm)
if op == "revoke" {
// Remote authorization revocation deliberately preserves protocol
// keys and their binding so reconnect reaches the RPC authorization
// gate and receives AUTH_KEY_UNREGISTERED rather than transport -404.
assertRevokeTestPresentAuthKey(t, ctx, keys, temp)
assertRevokeTestPresentAuthKey(t, ctx, keys, perm)
if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || !found {
t.Fatalf("binding after revoke found=%v err=%v, want present", found, err)
}
return
}
assertTempIdentityAuthKeyMissing(t, ctx, keys, temp)
assertTempIdentityAuthKeyMissing(t, ctx, keys, perm)
if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || found {
t.Fatalf("binding after delete found=%v err=%v, want absent", found, err)
}
})
}
}
func TestAuthIdentityAuthorizationMirrorUsesLockedPrimaryLayerPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
keys := NewAuthKeyStore(pool)
auths := NewAuthorizationStore(pool)
t.Run("advance before stale bind", func(t *testing.T) {
perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
userID := createRevokeTestUser(t, ctx, pool, "layer-advance-before-bind")
if _, _, err := keys.AdvanceSessionLayer(
ctx, perm, 8901, 227, authKeySessionLayerTestMsgID(time.Now().UTC(), 1),
); err != nil {
t.Fatal(err)
}
if err := auths.Bind(ctx, domain.Authorization{
AuthKeyID: perm, UserID: userID, Hash: 8902, Layer: 220,
}); err != nil {
t.Fatal(err)
}
got, found, err := auths.ByAuthKey(ctx, perm)
if err != nil || !found || got.Layer != 227 {
t.Fatalf("stale bind mirror = (%+v,%v,%v)", got, found, err)
}
})
t.Run("bind before advance", func(t *testing.T) {
perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
userID := createRevokeTestUser(t, ctx, pool, "layer-bind-before-advance")
if err := auths.Bind(ctx, domain.Authorization{
AuthKeyID: perm, UserID: userID, Hash: 8903, Layer: 220,
}); err != nil {
t.Fatal(err)
}
if _, _, err := keys.AdvanceSessionLayer(
ctx, perm, 8904, 227, authKeySessionLayerTestMsgID(time.Now().UTC(), 2),
); err != nil {
t.Fatal(err)
}
got, found, err := auths.ByAuthKey(ctx, perm)
if err != nil || !found || got.Layer != 227 {
t.Fatalf("advanced mirror = (%+v,%v,%v)", got, found, err)
}
})
}
func TestDeleteOrphanedRevalidatesUncommittedAuthorizationAndTempBindPostgres(t *testing.T) {
t.Run("authorization bind", func(t *testing.T) {
pool := testPool(t)
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
t.Cleanup(cancel)
keys := NewAuthKeyStore(pool)
auths := NewAuthorizationStore(pool)
perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
if _, err := pool.Exec(ctx, `UPDATE auth_keys SET last_used_at = now() - interval '72 hours' WHERE auth_key_id = $1`, authKeyIDToInt64(perm)); err != nil {
t.Fatal(err)
}
userID := createRevokeTestUser(t, ctx, pool, "orphan-auth-bind")
gcConn, err := pool.Acquire(ctx)
if err != nil {
t.Fatal(err)
}
defer gcConn.Release()
var gcPID int
if err := gcConn.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&gcPID); err != nil {
t.Fatal(err)
}
barrier := newAuthStoreQueryBarrier(gcConn, "", "orphan_identity_candidates")
gcResult := make(chan struct {
deleted int
err error
}, 1)
go func() {
deleted, err := NewAuthKeyStore(barrier).DeleteOrphaned(ctx, 24*time.Hour, 1, nil)
gcResult <- struct {
deleted int
err error
}{deleted: deleted, err: err}
}()
<-barrier.observed
bindTx, err := pool.Begin(ctx)
if err != nil {
t.Fatal(err)
}
defer func() { _ = bindTx.Rollback(context.Background()) }()
if err := NewAuthorizationStore(bindTx).Bind(ctx, domain.Authorization{
AuthKeyID: perm, UserID: userID, Hash: 9001,
}); err != nil {
t.Fatal(err)
}
close(barrier.release)
waitForPostgresBackendLockWait(t, ctx, pool, gcPID)
if err := bindTx.Commit(ctx); err != nil {
t.Fatal(err)
}
got := <-gcResult
if got.err != nil || got.deleted != 0 {
t.Fatalf("orphan GC after authorization bind = (%d,%v)", got.deleted, got.err)
}
assertRevokeTestPresentAuthKey(t, ctx, keys, perm)
assertRevokeTestPresentAuthorization(t, ctx, auths, perm)
})
t.Run("temporary bind", func(t *testing.T) {
pool := testPool(t)
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
t.Cleanup(cancel)
keys := NewAuthKeyStore(pool)
bindings := NewTempAuthKeyBindingStore(pool)
expiresAt := int(time.Now().Add(time.Hour).Unix())
temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, expiresAt)
perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
if _, err := pool.Exec(ctx, `UPDATE auth_keys SET last_used_at = now() - interval '72 hours' WHERE auth_key_id = $1`, authKeyIDToInt64(temp)); err != nil {
t.Fatal(err)
}
binding := domain.TempAuthKeyBinding{
TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(perm),
Nonce: 9002, TempSessionID: 9003, ExpiresAt: expiresAt,
EncryptedMessage: []byte("orphan bind revalidation"),
}
gcConn, err := pool.Acquire(ctx)
if err != nil {
t.Fatal(err)
}
defer gcConn.Release()
barrier := newAuthStoreQueryBarrier(gcConn, "", "orphan_identity_candidates")
gcResult := make(chan struct {
deleted int
err error
}, 1)
go func() {
deleted, err := NewAuthKeyStore(barrier).DeleteOrphaned(ctx, 24*time.Hour, 1, nil)
gcResult <- struct {
deleted int
err error
}{deleted: deleted, err: err}
}()
<-barrier.observed
bindTx, err := pool.Begin(ctx)
if err != nil {
t.Fatal(err)
}
defer func() { _ = bindTx.Rollback(context.Background()) }()
if err := NewTempAuthKeyBindingStore(bindTx).Save(ctx, binding); err != nil {
t.Fatal(err)
}
close(barrier.release)
var gc struct {
deleted int
err error
}
select {
case gc = <-gcResult:
case <-time.After(2 * time.Second):
_ = bindTx.Commit(context.Background())
t.Fatal("orphan GC blocked on an uncommitted temp bind despite SKIP LOCKED")
}
if gc.err != nil || gc.deleted != 0 {
t.Fatalf("orphan GC during temp bind = (%d,%v)", gc.deleted, gc.err)
}
if err := bindTx.Commit(ctx); err != nil {
t.Fatal(err)
}
assertTempIdentityBinding(t, ctx, bindings, binding)
assertRevokeTestPresentAuthKey(t, ctx, keys, temp)
assertRevokeTestPresentAuthKey(t, ctx, keys, perm)
})
}
type authStoreQueryBarrier struct {
*pgxpool.Conn
queryRowMarker string
queryMarker string
observed chan struct{}
release chan struct{}
once sync.Once
}
func newAuthStoreQueryBarrier(conn *pgxpool.Conn, queryRowMarker, queryMarker string) *authStoreQueryBarrier {
return &authStoreQueryBarrier{
Conn: conn, queryRowMarker: queryRowMarker, queryMarker: queryMarker,
observed: make(chan struct{}), release: make(chan struct{}),
}
}
func (db *authStoreQueryBarrier) Begin(ctx context.Context) (pgx.Tx, error) {
tx, err := db.Conn.Begin(ctx)
if err != nil {
return nil, err
}
return &authStoreBarrierTx{Tx: tx, owner: db}, nil
}
type authStoreBarrierTx struct {
pgx.Tx
owner *authStoreQueryBarrier
}
func (tx *authStoreBarrierTx) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row {
row := tx.Tx.QueryRow(ctx, sql, args...)
if tx.owner.queryRowMarker != "" && strings.Contains(sql, tx.owner.queryRowMarker) {
return &authStoreBarrierRow{Row: row, owner: tx.owner}
}
return row
}
func (tx *authStoreBarrierTx) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) {
rows, err := tx.Tx.Query(ctx, sql, args...)
if err != nil {
return nil, err
}
if tx.owner.queryMarker != "" && strings.Contains(sql, tx.owner.queryMarker) {
return &authStoreBarrierRows{Rows: rows, owner: tx.owner}, nil
}
return rows, nil
}
type authStoreBarrierRow struct {
pgx.Row
owner *authStoreQueryBarrier
}
func (row *authStoreBarrierRow) Scan(dest ...any) error {
err := row.Row.Scan(dest...)
if err == nil {
row.owner.once.Do(func() {
close(row.owner.observed)
<-row.owner.release
})
}
return err
}
type authStoreBarrierRows struct {
pgx.Rows
owner *authStoreQueryBarrier
}
func (rows *authStoreBarrierRows) Next() bool {
next := rows.Rows.Next()
if !next {
rows.owner.once.Do(func() {
close(rows.owner.observed)
<-rows.owner.release
})
}
return next
}