astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: sync multilayer td integration

Browse source
This commit is contained in:
A 2026-07-15 13:32:06 +08:00
parent 20a310f6ca
commit 766c5db992
491 changed files with 26235 additions and 35340 deletions
Download patch file Download diff file Expand all files Collapse all files

222
internal/store/postgres/auth_identity_lock.go Normal file
View file

@ -0,0 +1,222 @@
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
}

429
internal/store/postgres/auth_identity_lock_integration_test.go Normal file
View file

@ -0,0 +1,429 @@
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)
}
barrier := newAuthStoreQueryBarrier(advanceConn, "auth_identity_hint", "")
msgID := authKeySessionLayerTestMsgID(time.Now().UTC(), 1)
type advanceResult struct {
value store.AuthKeySessionLayer
applied bool
err error
}
result := make(chan advanceResult, 1)
go func() {
value, applied, err := NewAuthKeyStore(barrier).AdvanceSessionLayer(ctx, temp, 8703, 227, msgID)
result <- advanceResult{value: value, applied: applied, err: err}
}()
<-barrier.observed
// The selector already read "unbound". Stage a committed binding behind
// its statement snapshot while retaining P/raw row locks in the outer tx.
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)
}
close(barrier.release)
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)
}
assertTempIdentityAuthKeyMissing(t, ctx, keys, temp)
assertTempIdentityAuthKeyMissing(t, ctx, keys, perm)
assertRevokeTestNoAuthorization(t, ctx, auths, perm)
if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || found {
t.Fatalf("binding after %s found=%v err=%v", op, 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
}

264
internal/store/postgres/authkey.go
View file

@ -8,7 +8,6 @@ import (
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/jackc/pgx/v5/pgtype"
"telesrv/internal/store"
@ -56,24 +55,26 @@ WHERE auth_keys.body = EXCLUDED.body
// 注册进 SessionManager”的窗口被后台清理。
func (s *AuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
var (
body []byte
serverSalt int64
expiresAt int
createdAt pgtype.Timestamptz
layer int
deviceModel string
platform string
systemVersion string
apiID int
appVersion string
body []byte
serverSalt int64
expiresAt int
createdAt pgtype.Timestamptz
layer int
layerObservationID int64
deviceModel string
platform string
systemVersion string
apiID int
appVersion string
)
err := s.db.QueryRow(ctx, `
UPDATE auth_keys
SET last_used_at = now()
WHERE auth_key_id = $1
RETURNING auth_key_id, body, server_salt, created_at,
expires_at, layer, device_model, platform, system_version, api_id, app_version
`, authKeyIDToInt64(id)).Scan(new(int64), &body, &serverSalt, &createdAt, &expiresAt, &layer, &deviceModel, &platform, &systemVersion, &apiID, &appVersion)
expires_at, layer, layer_observation_id,
device_model, platform, system_version, api_id, app_version
`, authKeyIDToInt64(id)).Scan(new(int64), &body, &serverSalt, &createdAt, &expiresAt, &layer, &layerObservationID, &deviceModel, &platform, &systemVersion, &apiID, &appVersion)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return store.AuthKeyData{}, false, nil
@ -84,15 +85,16 @@ RETURNING auth_key_id, body, server_salt, created_at,
return store.AuthKeyData{}, false, fmt.Errorf("auth key body length = %d, want 256", len(body))
}
data := store.AuthKeyData{
ID: id,
ServerSalt: serverSalt,
ExpiresAt: expiresAt,
Layer: layer,
DeviceModel: deviceModel,
Platform: platform,
SystemVersion: systemVersion,
APIID: apiID,
AppVersion: appVersion,
ID: id,
ServerSalt: serverSalt,
ExpiresAt: expiresAt,
Layer: layer,
LayerObservationID: layerObservationID,
DeviceModel: deviceModel,
Platform: platform,
SystemVersion: systemVersion,
APIID: apiID,
AppVersion: appVersion,
}
copy(data.Value[:], body)
if createdAt.Valid {
@ -148,7 +150,8 @@ WHERE auth_key_id = ANY($1::bigint[])`, batch)
}
func (s *AuthKeyStore) UpdateClientInfo(ctx context.Context, id [8]byte, info store.AuthKeyClientInfo) error {
if _, err := s.db.Exec(ctx, `
var updated int
err := s.db.QueryRow(ctx, `
UPDATE auth_keys
SET layer = CASE WHEN $2::integer > 0 THEN $2 ELSE layer END,
device_model = CASE WHEN $3::text <> '' THEN $3 ELSE device_model END,
@ -157,7 +160,31 @@ SET layer = CASE WHEN $2::integer > 0 THEN $2 ELSE layer END,
api_id = CASE WHEN $6::integer <> 0 THEN $6 ELSE api_id END,
app_version = CASE WHEN $7::text <> '' THEN $7 ELSE app_version END
WHERE auth_key_id = $1
`, authKeyIDToInt64(id), info.Layer, info.DeviceModel, info.Platform, info.SystemVersion, info.APIID, info.AppVersion); err != nil {
AND ($2::integer <= 0 OR layer_observation_id = 0 OR layer = $2::integer)
RETURNING 1
`, authKeyIDToInt64(id), info.Layer, info.DeviceModel, info.Platform, info.SystemVersion, info.APIID, info.AppVersion).Scan(&updated)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
var (
currentLayer int
observation int64
)
lookupErr := s.db.QueryRow(ctx, `
SELECT layer, layer_observation_id
FROM auth_keys
WHERE auth_key_id = $1
`, authKeyIDToInt64(id)).Scan(&currentLayer, &observation)
switch {
case errors.Is(lookupErr, pgx.ErrNoRows):
return store.ErrAuthKeyNotFound
case lookupErr != nil:
return fmt.Errorf("classify auth key client info update: %w", lookupErr)
case info.Layer > 0 && observation > 0 && currentLayer != info.Layer:
return store.ErrAuthKeySessionLayerConflict
default:
return fmt.Errorf("update auth key client info: guarded update affected no row")
}
}
return fmt.Errorf("update auth key client info: %w", err)
}
return nil
@ -171,25 +198,20 @@ WHERE auth_key_id = $1
// RESTRICT FK 防止悬空，因此被踢/销毁 perm key 时必须先把关联 temp key 一并删掉。否则 Web/上传连接用
// raw temp key 重连时仍能进入 RPC 层，只得到 AUTH_KEY_UNREGISTERED，而不是连接层 404。
func (s *AuthKeyStore) Delete(ctx context.Context, id [8]byte) error {
for attempt := 0; attempt < 3; attempt++ {
err := s.deleteAuthKeyOnce(ctx, id)
if err == nil {
return nil
}
if !isPermAuthKeyDeleteRace(err) {
return err
}
if _, inTx := s.db.(pgx.Tx); inTx {
return err
}
}
return fmt.Errorf("delete auth key: permanent-key binding changed during all retries")
return withAuthIdentityTx(ctx, s.db, "delete auth key", func(tx pgx.Tx) error {
return deleteAuthKeyTx(ctx, tx, authKeyIDToInt64(id))
})
}
func (s *AuthKeyStore) deleteAuthKeyOnce(ctx context.Context, id [8]byte) error {
keyID := authKeyIDToInt64(id)
func deleteAuthKeyTx(ctx context.Context, tx pgx.Tx, keyID int64) error {
if _, _, _, err := lockRawAuthKeyInIdentityOrder(ctx, tx, keyID); err != nil {
if errors.Is(err, store.ErrAuthKeyNotFound) {
return nil
}
return err
}
var touched int
if err := s.db.QueryRow(ctx, `
if err := tx.QueryRow(ctx, `
WITH doomed_temp AS MATERIALIZED (
SELECT temp_auth_key_id
FROM temp_auth_key_bindings
@ -225,13 +247,6 @@ SELECT
const tempAuthKeyPermFKConstraint = "temp_auth_key_bindings_perm_auth_key_id_fkey"
func isPermAuthKeyDeleteRace(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) &&
pgErr.Code == "23503" &&
pgErr.ConstraintName == tempAuthKeyPermFKConstraint
}
// DeleteOrphaned 回收握手已落库、但从未形成 authorization/temp binding 且当前没有
// 活跃物理连接的旧 auth key。last_used_at 与 Get 的 UPDATE ... RETURNING 行锁配对，封住
// active-key 快照之后新连接开始使用旧 key 的竞态；所有引用条件仍在最终 DELETE 中复核。
@ -248,51 +263,140 @@ func (s *AuthKeyStore) DeleteOrphaned(ctx context.Context, olderThan time.Durati
protectedIDs = append(protectedIDs, authKeyIDToInt64(id))
}
var deleted int
err := s.db.QueryRow(ctx, `
WITH candidates AS MATERIALIZED (
err := withAuthIdentityTx(ctx, s.db, "delete orphaned auth keys", func(tx pgx.Tx) error {
var err error
deleted, err = deleteOrphanedAuthKeysTx(ctx, tx, olderThan, limit, protectedIDs)
return err
})
if err != nil {
return 0, fmt.Errorf("delete orphaned auth keys: %w", err)
}
return deleted, nil
}
func deleteOrphanedAuthKeysTx(
ctx context.Context,
tx pgx.Tx,
olderThan time.Duration,
limit int,
protectedIDs []int64,
) (int, error) {
// Phase 1 is only a bounded hint. It must not lock rows before the complete
// permanent-identity advisory set has been derived and acquired.
rows, err := tx.Query(ctx, `
/* orphan_identity_candidates */
SELECT k.auth_key_id, k.expires_at
FROM auth_keys AS k
WHERE k.last_used_at < now() - make_interval(secs => $1::double precision)
AND NOT (k.auth_key_id = ANY($2::bigint[]))
AND NOT EXISTS (
SELECT 1 FROM authorizations AS a WHERE a.auth_key_id = k.auth_key_id
)
AND NOT EXISTS (
SELECT 1
FROM temp_auth_key_bindings AS b
WHERE b.temp_auth_key_id = k.auth_key_id OR b.perm_auth_key_id = k.auth_key_id
)
ORDER BY k.last_used_at, k.auth_key_id
LIMIT $3`, olderThan.Seconds(), protectedIDs, limit)
if err != nil {
return 0, fmt.Errorf("select orphan auth key candidates: %w", err)
}
candidates := make([]int64, 0, limit)
permanentCandidates := make([]int64, 0, limit)
for rows.Next() {
var (
id int64
expiresAt int
)
if err := rows.Scan(&id, &expiresAt); err != nil {
rows.Close()
return 0, fmt.Errorf("scan orphan auth key candidate: %w", err)
}
candidates = append(candidates, id)
if expiresAt == 0 {
permanentCandidates = append(permanentCandidates, id)
}
}
if err := rows.Err(); err != nil {
rows.Close()
return 0, fmt.Errorf("iterate orphan auth key candidates: %w", err)
}
rows.Close()
if len(candidates) == 0 {
return 0, nil
}
if err := lockPermanentAuthIdentities(ctx, tx, permanentCandidates); err != nil {
return 0, err
}
// Phase 2 locks only the hinted raw rows, in real-ID order, after every P
// gate is held. A temp bind already holding its raw row is skipped; if it
// committed immediately before this lock, phase 3's new READ COMMITTED
// statement sees the binding and excludes it.
lockRows, err := tx.Query(ctx, `
SELECT auth_key_id
FROM auth_keys
WHERE auth_key_id = ANY($1::bigint[])
ORDER BY auth_key_id
FOR UPDATE SKIP LOCKED`, candidates)
if err != nil {
return 0, fmt.Errorf("lock orphan auth key candidates: %w", err)
}
lockedIDs := make([]int64, 0, len(candidates))
for lockRows.Next() {
var id int64
if err := lockRows.Scan(&id); err != nil {
lockRows.Close()
return 0, fmt.Errorf("scan locked orphan auth key: %w", err)
}
lockedIDs = append(lockedIDs, id)
}
if err := lockRows.Err(); err != nil {
lockRows.Close()
return 0, fmt.Errorf("iterate locked orphan auth keys: %w", err)
}
lockRows.Close()
if len(lockedIDs) == 0 {
return 0, nil
}
// Phase 3 is a separate statement snapshot and repeats every ownership,
// activity and protection predicate. Never delete from the phase-1 hint.
var deleted int
err = tx.QueryRow(ctx, `
WITH still_orphaned AS MATERIALIZED (
SELECT k.auth_key_id
FROM auth_keys k
WHERE k.last_used_at < now() - make_interval(secs => $1::double precision)
AND NOT (k.auth_key_id = ANY($2::bigint[]))
AND NOT EXISTS (
SELECT 1 FROM authorizations a WHERE a.auth_key_id = k.auth_key_id
)
AND NOT EXISTS (
SELECT 1
FROM temp_auth_key_bindings b
WHERE b.temp_auth_key_id = k.auth_key_id OR b.perm_auth_key_id = k.auth_key_id
)
ORDER BY k.last_used_at ASC, k.auth_key_id ASC
LIMIT $3
FOR UPDATE OF k SKIP LOCKED
), deleted_update_states AS (
-- Historical authorization-only deletion could leave a cursor without an
-- auth_keys FK. GC owns that stale row once the raw key is proven orphaned.
DELETE FROM update_states s
USING candidates c
WHERE s.auth_key_id = c.auth_key_id
RETURNING s.auth_key_id
), deleted_keys AS (
DELETE FROM auth_keys k
USING candidates c
WHERE k.auth_key_id = c.auth_key_id
FROM auth_keys AS k
WHERE k.auth_key_id = ANY($3::bigint[])
AND k.last_used_at < now() - make_interval(secs => $1::double precision)
AND NOT (k.auth_key_id = ANY($2::bigint[]))
AND NOT EXISTS (
SELECT 1 FROM authorizations a WHERE a.auth_key_id = k.auth_key_id
SELECT 1 FROM authorizations AS a WHERE a.auth_key_id = k.auth_key_id
)
AND NOT EXISTS (
SELECT 1
FROM temp_auth_key_bindings b
FROM temp_auth_key_bindings AS b
WHERE b.temp_auth_key_id = k.auth_key_id OR b.perm_auth_key_id = k.auth_key_id
)
RETURNING k.auth_key_id
), deleted_update_states AS (
DELETE FROM update_states AS state
USING still_orphaned AS orphan
WHERE state.auth_key_id = orphan.auth_key_id
RETURNING state.auth_key_id
), deleted_keys AS (
DELETE FROM auth_keys AS key
USING still_orphaned AS orphan
WHERE key.auth_key_id = orphan.auth_key_id
RETURNING key.auth_key_id
)
SELECT count(*)::int
SELECT count(*)::integer
FROM deleted_keys
CROSS JOIN LATERAL (SELECT count(*) FROM deleted_update_states) AS touched`, olderThan.Seconds(), protectedIDs, limit).Scan(&deleted)
CROSS JOIN LATERAL (SELECT count(*) FROM deleted_update_states) AS touched`,
olderThan.Seconds(), protectedIDs, lockedIDs,
).Scan(&deleted)
if err != nil {
return 0, fmt.Errorf("delete orphaned auth keys: %w", err)
return 0, fmt.Errorf("delete revalidated orphan auth keys: %w", err)
}
return deleted, nil
}

66
internal/store/postgres/authkey_integration_test.go
View file

@ -144,3 +144,69 @@ func TestAuthKeyStoreClientInfoRoundTrip(t *testing.T) {
t.Fatalf("partial client info merge mismatch: %+v", got)
}
}
func TestAuthKeyStoreUpdateClientInfoProtectsObservedLayerPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
keys := NewAuthKeyStore(pool)
var id [8]byte
var value [256]byte
if _, err := rand.Read(id[:]); err != nil {
t.Fatal(err)
}
if _, err := rand.Read(value[:]); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = keys.Delete(ctx, id) })
if err := keys.Save(ctx, store.AuthKeyData{ID: id, Value: value}); err != nil {
t.Fatalf("save auth key: %v", err)
}
if _, err := pool.Exec(ctx, `
UPDATE auth_keys
SET layer = 227, layer_observation_id = 91,
device_model = 'before', platform = 'tdesktop'
WHERE auth_key_id = $1`, authKeyIDToInt64(id)); err != nil {
t.Fatalf("seed ordered layer: %v", err)
}
err := keys.UpdateClientInfo(ctx, id, store.AuthKeyClientInfo{
Layer: 220, DeviceModel: "must-not-merge", AppVersion: "must-not-merge",
})
if !errors.Is(err, store.ErrAuthKeySessionLayerConflict) {
t.Fatalf("conflicting layer update error = %v, want %v", err, store.ErrAuthKeySessionLayerConflict)
}
got, found, err := keys.Get(ctx, id)
if err != nil || !found {
t.Fatalf("get after conflict: found=%v err=%v", found, err)
}
if got.Layer != 227 || got.LayerObservationID != 91 || got.DeviceModel != "before" ||
got.Platform != "tdesktop" || got.AppVersion != "" {
t.Fatalf("conflicting update changed row: %+v", got)
}
if err := keys.UpdateClientInfo(ctx, id, store.AuthKeyClientInfo{
Layer: 227, DeviceModel: "same-layer", AppVersion: "1.0",
}); err != nil {
t.Fatalf("same observed layer metadata merge: %v", err)
}
if err := keys.UpdateClientInfo(ctx, id, store.AuthKeyClientInfo{
Platform: "windows", SystemVersion: "11",
}); err != nil {
t.Fatalf("layerless metadata merge: %v", err)
}
got, found, err = keys.Get(ctx, id)
if err != nil || !found {
t.Fatalf("get guarded metadata merge: found=%v err=%v", found, err)
}
if got.Layer != 227 || got.LayerObservationID != 91 || got.DeviceModel != "same-layer" ||
got.Platform != "windows" || got.SystemVersion != "11" || got.AppVersion != "1.0" {
t.Fatalf("guarded metadata merge = %+v", got)
}
missing := id
missing[0] ^= 0xff
if err := keys.UpdateClientInfo(ctx, missing, store.AuthKeyClientInfo{Layer: 227}); !errors.Is(err, store.ErrAuthKeyNotFound) {
t.Fatalf("missing primary update error = %v, want %v", err, store.ErrAuthKeyNotFound)
}
}

245
internal/store/postgres/authkey_session_layer.go Normal file
View file

@ -0,0 +1,245 @@
package postgres
import (
"context"
"errors"
"fmt"
"time"
"github.com/jackc/pgx/v5"
"telesrv/internal/store"
)
const maxAuthKeySessionLayerDeleteBatch = 100000
func (s *AuthKeyStore) GetSessionLayer(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
) (store.AuthKeySessionLayer, bool, error) {
var value store.AuthKeySessionLayer
err := s.db.QueryRow(ctx, `
SELECT evidence.layer,
evidence.msg_id,
evidence.observation_id,
evidence.expires_at,
defaults.layer = evidence.layer
AND defaults.layer_observation_id = evidence.observation_id
FROM auth_key_session_layers AS evidence
LEFT JOIN temp_auth_key_bindings AS binding
ON binding.temp_auth_key_id = evidence.raw_auth_key_id
JOIN auth_keys AS defaults
ON defaults.auth_key_id = COALESCE(binding.perm_auth_key_id, evidence.raw_auth_key_id)
WHERE evidence.raw_auth_key_id = $1
AND evidence.session_id = $2
AND evidence.expires_at > now()
`, authKeyIDToInt64(rawAuthKeyID), sessionID).Scan(
&value.Layer,
&value.MessageID,
&value.ObservationID,
&value.ExpiresAt,
&value.SharedDefault,
)
if errors.Is(err, pgx.ErrNoRows) {
return store.AuthKeySessionLayer{}, false, nil
}
if err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("get auth key session layer: %w", err)
}
return value, true, nil
}
// AdvanceSessionLayer enters the permanent identity advisory gate before any
// row lock when rawAuthKeyID is permanent or already-bound temporary. An
// initially-unbound temp key that becomes bound while the raw row is acquired
// rolls the attempt back and retries in the new identity. The session watermark
// and every currently bound shared default then commit in one transaction.
func (s *AuthKeyStore) AdvanceSessionLayer(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
layer int,
msgID int64,
) (store.AuthKeySessionLayer, bool, error) {
expiresAt, validMessageID := store.AuthKeySessionLayerExpiry(msgID)
if layer <= 0 || !validMessageID {
return store.AuthKeySessionLayer{}, false, store.ErrAuthKeySessionLayerInvalid
}
var (
current store.AuthKeySessionLayer
applied bool
)
err := withAuthIdentityTx(ctx, s.db, "advance auth key session layer", func(tx pgx.Tx) error {
var err error
current, applied, err = advanceSessionLayerTx(
ctx, tx, authKeyIDToInt64(rawAuthKeyID), sessionID, layer, msgID, expiresAt,
)
return err
})
if err != nil {
return current, false, err
}
return current, applied, nil
}
func advanceSessionLayerTx(
ctx context.Context,
tx pgx.Tx,
rawID int64,
sessionID int64,
layer int,
msgID int64,
expiresAt time.Time,
) (store.AuthKeySessionLayer, bool, error) {
_, permID, _, err := lockRawAuthKeyInIdentityOrder(ctx, tx, rawID)
if err != nil {
return store.AuthKeySessionLayer{}, false, err
}
var (
current store.AuthKeySessionLayer
now time.Time
)
err = tx.QueryRow(ctx, `
SELECT layer, msg_id, observation_id, expires_at, now()
FROM auth_key_session_layers
WHERE raw_auth_key_id = $1 AND session_id = $2
FOR UPDATE
`, rawID, sessionID).Scan(
&current.Layer,
&current.MessageID,
&current.ObservationID,
&current.ExpiresAt,
&now,
)
if errors.Is(err, pgx.ErrNoRows) {
if err := tx.QueryRow(ctx, `SELECT now()`).Scan(&now); err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("read session layer database time: %w", err)
}
current = store.AuthKeySessionLayer{}
} else if err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("lock auth key session layer: %w", err)
}
if _, fresh := store.AuthKeySessionLayerEvidenceFresh(now, msgID); !fresh {
return store.AuthKeySessionLayer{}, false, store.ErrAuthKeySessionLayerInvalid
}
if current.MessageID != 0 && now.Before(current.ExpiresAt) {
switch {
case msgID < current.MessageID:
if err := tx.QueryRow(ctx, `
SELECT layer = $2 AND layer_observation_id = $3
FROM auth_keys WHERE auth_key_id = $1
`, permID, current.Layer, current.ObservationID).Scan(&current.SharedDefault); err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("compare older session layer with shared default: %w", err)
}
return current, false, nil
case msgID == current.MessageID:
if layer != current.Layer {
return current, false, store.ErrAuthKeySessionLayerConflict
}
if err := tx.QueryRow(ctx, `
SELECT layer = $2 AND layer_observation_id = $3
FROM auth_keys WHERE auth_key_id = $1
`, permID, current.Layer, current.ObservationID).Scan(&current.SharedDefault); err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("compare duplicate session layer with shared default: %w", err)
}
return current, false, nil
}
}
var observationID int64
if err := tx.QueryRow(ctx, `SELECT nextval('auth_key_layer_observation_seq')`).Scan(&observationID); err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("allocate auth key layer observation: %w", err)
}
err = tx.QueryRow(ctx, `
INSERT INTO auth_key_session_layers (
raw_auth_key_id, session_id, layer, msg_id, observation_id, expires_at
) VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (raw_auth_key_id, session_id) DO UPDATE SET
layer = EXCLUDED.layer,
msg_id = EXCLUDED.msg_id,
observation_id = EXCLUDED.observation_id,
expires_at = EXCLUDED.expires_at
RETURNING layer, msg_id, observation_id, expires_at
`, rawID, sessionID, layer, msgID, observationID, expiresAt).Scan(
&current.Layer,
&current.MessageID,
&current.ObservationID,
&current.ExpiresAt,
)
if err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("upsert auth key session layer: %w", err)
}
keyIDs := []int64{rawID}
if permID != rawID {
keyIDs = append(keyIDs, permID)
}
tag, err := tx.Exec(ctx, `
UPDATE auth_keys
SET layer = $2, layer_observation_id = $3
WHERE auth_key_id = ANY($1::bigint[])
AND layer_observation_id < $3
`, keyIDs, layer, observationID)
if err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("publish auth key session layer defaults: %w", err)
}
if tag.RowsAffected() != int64(len(keyIDs)) {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("publish auth key session layer defaults: updated %d of %d locked keys", tag.RowsAffected(), len(keyIDs))
}
if _, err := tx.Exec(ctx, `
UPDATE authorizations
SET layer = $2
WHERE auth_key_id = ANY($1::bigint[])
`, keyIDs, layer); err != nil {
return store.AuthKeySessionLayer{}, false, fmt.Errorf("mirror auth key session layer defaults: %w", err)
}
current.SharedDefault = true
return current, true, nil
}
func (s *AuthKeyStore) DeleteSessionLayer(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
) (bool, error) {
tag, err := s.db.Exec(ctx, `
DELETE FROM auth_key_session_layers
WHERE raw_auth_key_id = $1 AND session_id = $2
`, authKeyIDToInt64(rawAuthKeyID), sessionID)
if err != nil {
return false, fmt.Errorf("delete auth key session layer: %w", err)
}
return tag.RowsAffected() > 0, nil
}
func (s *AuthKeyStore) DeleteExpiredSessionLayers(ctx context.Context, limit int) (int, error) {
if limit <= 0 {
return 0, nil
}
if limit > maxAuthKeySessionLayerDeleteBatch {
limit = maxAuthKeySessionLayerDeleteBatch
}
var deleted int
err := s.db.QueryRow(ctx, `
WITH candidates AS MATERIALIZED (
SELECT raw_auth_key_id, session_id
FROM auth_key_session_layers
WHERE expires_at <= now()
ORDER BY expires_at, raw_auth_key_id, session_id
LIMIT $1
FOR UPDATE SKIP LOCKED
), removed AS (
DELETE FROM auth_key_session_layers AS evidence
USING candidates
WHERE evidence.raw_auth_key_id = candidates.raw_auth_key_id
AND evidence.session_id = candidates.session_id
AND evidence.expires_at <= now()
RETURNING 1
)
SELECT count(*)::integer FROM removed
`, limit).Scan(&deleted)
if err != nil {
return 0, fmt.Errorf("delete expired auth key session layers: %w", err)
}
return deleted, nil
}

140
internal/store/postgres/authkey_session_layer_integration_test.go Normal file
View file

@ -0,0 +1,140 @@
package postgres
import (
"context"
"crypto/rand"
"encoding/binary"
"errors"
"sync"
"testing"
"time"
"telesrv/internal/domain"
"telesrv/internal/store"
)
func TestAuthKeySessionLayerTransactionAndRestartPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
temp := randomLayerTestAuthKeyID(t)
perm := randomLayerTestAuthKeyID(t)
for perm == temp {
perm = randomLayerTestAuthKeyID(t)
}
const sessionID = int64(87001)
t.Cleanup(func() {
_ = NewAuthKeyStore(pool).Delete(ctx, perm)
_ = NewAuthKeyStore(pool).Delete(ctx, temp)
})
keys := NewAuthKeyStore(pool)
expiresAt := int(time.Now().Add(time.Hour).Unix())
if err := keys.Save(ctx, store.AuthKeyData{ID: temp, ExpiresAt: expiresAt}); err != nil {
t.Fatal(err)
}
if err := keys.Save(ctx, store.AuthKeyData{ID: perm, ExpiresAt: 0}); err != nil {
t.Fatal(err)
}
now := time.Now().UTC()
firstMsgID := authKeySessionLayerTestMsgID(now, 1)
newerMsgID := authKeySessionLayerTestMsgID(now, 2)
concurrentLowMsgID := authKeySessionLayerTestMsgID(now, 3)
concurrentHighMsgID := authKeySessionLayerTestMsgID(now, 4)
for _, invalidMsgID := range []int64{
authKeySessionLayerTestMsgID(now.Add(-302*time.Second), 1),
authKeySessionLayerTestMsgID(now.Add(31*time.Second), 1),
firstMsgID + 1,
} {
if _, _, err := keys.AdvanceSessionLayer(ctx, temp, sessionID, 220, invalidMsgID); !errors.Is(err, store.ErrAuthKeySessionLayerInvalid) {
t.Fatalf("invalid msg_id %d advance err = %v", invalidMsgID, err)
}
}
if _, found, err := keys.GetSessionLayer(ctx, temp, sessionID); err != nil || found {
t.Fatalf("rejected evidence created session row: found=%v err=%v", found, err)
}
first, applied, err := keys.AdvanceSessionLayer(ctx, temp, sessionID, 220, firstMsgID)
if err != nil || !applied || !first.SharedDefault || first.ObservationID <= 0 {
t.Fatalf("first advance = (%+v,%v,%v)", first, applied, err)
}
if err := NewTempAuthKeyBindingStore(pool).Save(ctx, domain.TempAuthKeyBinding{
TempAuthKeyID: temp,
PermAuthKeyID: int64(binary.LittleEndian.Uint64(perm[:])),
Nonce: 87,
TempSessionID: sessionID,
ExpiresAt: expiresAt,
EncryptedMessage: []byte{8, 7},
}); err != nil {
t.Fatal(err)
}
for _, id := range [][8]byte{temp, perm} {
got, found, err := NewAuthKeyStore(pool).Get(ctx, id)
if err != nil || !found || got.Layer != 220 || got.LayerObservationID != first.ObservationID {
t.Fatalf("bound default %x = (%+v,%v,%v)", id, got, found, err)
}
}
newer, applied, err := NewAuthKeyStore(pool).AdvanceSessionLayer(ctx, temp, sessionID, 227, newerMsgID)
if err != nil || !applied || !newer.SharedDefault || newer.ObservationID <= first.ObservationID {
t.Fatalf("newer advance = (%+v,%v,%v)", newer, applied, err)
}
old, applied, err := NewAuthKeyStore(pool).AdvanceSessionLayer(ctx, temp, sessionID, 220, firstMsgID)
if err != nil || applied || old.Layer != 227 || old.MessageID != newerMsgID || !old.SharedDefault {
t.Fatalf("old replay = (%+v,%v,%v)", old, applied, err)
}
if _, _, err := NewAuthKeyStore(pool).AdvanceSessionLayer(ctx, temp, sessionID, 220, newerMsgID); !errors.Is(err, store.ErrAuthKeySessionLayerConflict) {
t.Fatalf("same-msg conflict = %v", err)
}
// Two independent store instances model two server processes. The raw-key
// row lock and session CAS must converge on the greater selector msg_id.
type candidate struct {
layer int
msgID int64
}
candidates := []candidate{{layer: 225, msgID: concurrentLowMsgID}, {layer: 227, msgID: concurrentHighMsgID}}
start := make(chan struct{})
errs := make(chan error, len(candidates))
var wg sync.WaitGroup
for _, item := range candidates {
item := item
wg.Add(1)
go func() {
defer wg.Done()
<-start
_, _, err := NewAuthKeyStore(pool).AdvanceSessionLayer(ctx, temp, sessionID, item.layer, item.msgID)
errs <- err
}()
}
close(start)
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatal(err)
}
}
restarted := NewAuthKeyStore(pool)
current, found, err := restarted.GetSessionLayer(ctx, temp, sessionID)
if err != nil || !found || current.Layer != 227 || current.MessageID != concurrentHighMsgID || !current.SharedDefault {
t.Fatalf("restart authoritative row = (%+v,%v,%v)", current, found, err)
}
for _, id := range [][8]byte{temp, perm} {
got, found, err := restarted.Get(ctx, id)
if err != nil || !found || got.Layer != 227 || got.LayerObservationID != current.ObservationID {
t.Fatalf("transactional shared default %x = (%+v,%v,%v)", id, got, found, err)
}
}
}
func authKeySessionLayerTestMsgID(at time.Time, order uint32) int64 {
return int64((uint64(at.Unix()) << 32) | uint64(order)<<2)
}
func randomLayerTestAuthKeyID(t *testing.T) [8]byte {
t.Helper()
var id [8]byte
if _, err := rand.Read(id[:]); err != nil {
t.Fatal(err)
}
return id
}

125
internal/store/postgres/authorization.go
View file

@ -29,17 +29,9 @@ func (s *AuthorizationStore) Bind(ctx context.Context, a domain.Authorization) e
if a.Hash == 0 {
a.Hash = authorizationHash(a.AuthKeyID)
}
bind := func(db sqlcgen.DBTX) error {
return bindAuthorization(ctx, db, a)
}
var err error
if tx, ok := s.db.(pgx.Tx); ok {
err = bind(tx)
} else {
err = withTx(ctx, s.db, "bind authorization", func(tx pgx.Tx) error {
return bind(tx)
})
}
err := withAuthIdentityTx(ctx, s.db, "bind authorization", func(tx pgx.Tx) error {
return bindAuthorization(ctx, tx, a)
})
if err != nil {
return fmt.Errorf("upsert authorization: %w", err)
}
@ -55,20 +47,35 @@ func (s *AuthorizationStore) Bind(ctx context.Context, a domain.Authorization) e
// raw auth key 的并发登录/换号。
func bindAuthorization(ctx context.Context, db sqlcgen.DBTX, a domain.Authorization) error {
keyID := authKeyIDToInt64(a.AuthKeyID)
tx, ok := db.(pgx.Tx)
if !ok {
return fmt.Errorf("bind authorization requires a transaction")
}
if err := lockPermanentAuthIdentities(ctx, tx, []int64{keyID}); err != nil {
return err
}
var (
lockedKeyID int64
expiresAt int
lockedKeyID int64
expiresAt int
authLayer int
layerObservationID int64
)
if err := db.QueryRow(ctx, `
SELECT auth_key_id, expires_at
SELECT auth_key_id, expires_at, layer, layer_observation_id
FROM auth_keys
WHERE auth_key_id = $1
FOR UPDATE`, keyID).Scan(&lockedKeyID, &expiresAt); err != nil {
FOR UPDATE`, keyID).Scan(&lockedKeyID, &expiresAt, &authLayer, &layerObservationID); err != nil {
return fmt.Errorf("lock auth key for authorization: %w", err)
}
if expiresAt != 0 {
return store.ErrAuthKeyNotPermanent
}
if authLayer < 0 || layerObservationID < 0 || (layerObservationID > 0 && authLayer == 0) {
return fmt.Errorf(
"authorization auth-key layer invariant violation: auth key %x has layer %d observation %d",
a.AuthKeyID, authLayer, layerObservationID,
)
}
if _, err := db.Exec(ctx, `
INSERT INTO user_update_watermarks (user_id, contiguous_pts)
@ -154,7 +161,7 @@ ON CONFLICT (auth_key_id) DO UPDATE SET
ip = EXCLUDED.ip,
password_pending = EXCLUDED.password_pending,
active_at = now()`,
keyID, a.UserID, a.Hash, int32(a.Layer), a.DeviceModel, a.Platform, a.SystemVersion, int32(a.APIID), a.AppVersion, a.IP, a.PasswordPending,
keyID, a.UserID, a.Hash, int32(authLayer), a.DeviceModel, a.Platform, a.SystemVersion, int32(a.APIID), a.AppVersion, a.IP, a.PasswordPending,
); err != nil {
return fmt.Errorf("write authorization: %w", err)
}
@ -178,18 +185,6 @@ FROM authorizations WHERE auth_key_id = $1`, authKeyIDToInt64(id))
return a, true, nil
}
func (s *AuthorizationStore) UpdateLayer(ctx context.Context, id [8]byte, layer int) error {
if layer <= 0 {
return nil
}
if _, err := s.db.Exec(ctx, `
UPDATE authorizations SET layer = $2, active_at = now() WHERE auth_key_id = $1`,
authKeyIDToInt64(id), int32(layer)); err != nil {
return fmt.Errorf("update authorization layer: %w", err)
}
return nil
}
func (s *AuthorizationStore) UpdateClientInfo(ctx context.Context, id [8]byte, info domain.AuthKeyClientInfo) error {
if _, err := s.db.Exec(ctx, `
UPDATE authorizations SET
@ -261,34 +256,19 @@ RETURNING auth_key_id, user_id, hash, layer, device_model, platform, system_vers
// authorizations 通过 FK cascade 删除；update_states 没有 auth_keys FK，必须显式清理；
// 关联 temp auth key 也显式删除，避免 raw temp key 重连。
func (s *AuthorizationStore) RevokeByHash(ctx context.Context, userID, hash int64) (domain.Authorization, bool, error) {
for attempt := 0; attempt < 3; attempt++ {
var (
a domain.Authorization
found bool
)
err := s.withRevocationTx(ctx, "revoke authorization by hash", func(tx pgx.Tx) error {
var err error
a, found, err = revokeByHashTx(ctx, tx, userID, hash)
return err
})
if err == nil {
return a, found, nil
}
if !isPermAuthKeyDeleteRace(err) {
return domain.Authorization{}, false, err
}
if _, inTx := s.db.(pgx.Tx); inTx {
return domain.Authorization{}, false, err
}
var (
a domain.Authorization
found bool
)
err := withAuthIdentityTx(ctx, s.db, "revoke authorization by hash", func(tx pgx.Tx) error {
var err error
a, found, err = revokeByHashTx(ctx, tx, userID, hash)
return err
})
if err != nil {
return domain.Authorization{}, false, err
}
return domain.Authorization{}, false, fmt.Errorf("revoke authorization by hash: permanent-key binding changed during all retries")
}
func (s *AuthorizationStore) withRevocationTx(ctx context.Context, op string, fn func(pgx.Tx) error) error {
if tx, ok := s.db.(pgx.Tx); ok {
return fn(tx)
}
return withTx(ctx, s.db, op, fn)
return a, found, nil
}
// revokeByHashTx deliberately uses separate READ COMMITTED statements. The first
@ -307,6 +287,9 @@ WHERE user_id = $1 AND hash = $2`, userID, hash).Scan(&candidate); err != nil {
}
return domain.Authorization{}, false, fmt.Errorf("select revoke candidate by hash: %w", err)
}
if err := lockPermanentAuthIdentities(ctx, tx, []int64{candidate}); err != nil {
return domain.Authorization{}, false, err
}
var locked int64
if err := tx.QueryRow(ctx, `
@ -368,24 +351,13 @@ RETURNING auth_key_id, user_id, hash, layer, device_model, platform, system_vers
// RevokeByUserExcept 批量删除协议 auth_key，保留 keepAuthKeyID 对应的当前设备。
func (s *AuthorizationStore) RevokeByUserExcept(ctx context.Context, userID int64, keepAuthKeyID [8]byte) ([]domain.Authorization, error) {
for attempt := 0; attempt < 3; attempt++ {
var out []domain.Authorization
err := s.withRevocationTx(ctx, "revoke authorizations by user", func(tx pgx.Tx) error {
var err error
out, err = revokeByUserExceptTx(ctx, tx, userID, authKeyIDToInt64(keepAuthKeyID))
return err
})
if err == nil {
return out, nil
}
if !isPermAuthKeyDeleteRace(err) {
return nil, err
}
if _, inTx := s.db.(pgx.Tx); inTx {
return nil, err
}
}
return nil, fmt.Errorf("revoke authorizations by user: permanent-key binding changed during all retries")
var out []domain.Authorization
err := withAuthIdentityTx(ctx, s.db, "revoke authorizations by user", func(tx pgx.Tx) error {
var err error
out, err = revokeByUserExceptTx(ctx, tx, userID, authKeyIDToInt64(keepAuthKeyID))
return err
})
return out, err
}
func revokeByUserExceptTx(ctx context.Context, tx pgx.Tx, userID, keepAuthKeyID int64) ([]domain.Authorization, error) {
@ -414,9 +386,12 @@ ORDER BY auth_key_id`, userID, keepAuthKeyID)
if len(candidates) == 0 {
return []domain.Authorization{}, nil
}
if err := lockPermanentAuthIdentities(ctx, tx, candidates); err != nil {
return nil, err
}
// Stable parent-row lock order matches every concurrent batch revocation and
// serializes each candidate with Bind's auth_keys-first ownership change.
// Advisory keys are already all held in final int32-hash order. Parent rows
// are then locked by their real bigint IDs for deterministic batch behavior.
lockRows, err := tx.Query(ctx, `
SELECT auth_key_id
FROM auth_keys

111
internal/store/postgres/temp_auth_key.go
View file

@ -16,19 +16,99 @@ import (
// TempAuthKeyBindingStore 用 PostgreSQL 实现 store.TempAuthKeyBindingStore。
type TempAuthKeyBindingStore struct {
q *sqlcgen.Queries
db sqlcgen.DBTX
q *sqlcgen.Queries
}
// NewTempAuthKeyBindingStore 基于 pgx 连接池（或事务）创建 TempAuthKeyBindingStore。
func NewTempAuthKeyBindingStore(db sqlcgen.DBTX) *TempAuthKeyBindingStore {
return &TempAuthKeyBindingStore{q: sqlcgen.New(db)}
return &TempAuthKeyBindingStore{db: db, q: sqlcgen.New(db)}
}
func (s *TempAuthKeyBindingStore) Save(ctx context.Context, b domain.TempAuthKeyBinding) error {
if b.ExpiresAt <= 0 || int64(b.ExpiresAt) > math.MaxInt32 {
return store.ErrAuthKeyBindingInvalid
}
n, err := s.q.UpsertTempAuthKeyBinding(ctx, sqlcgen.UpsertTempAuthKeyBindingParams{
return withAuthIdentityTx(ctx, s.db, "save temp auth key binding", func(tx pgx.Tx) error {
return s.saveTx(ctx, tx, b)
})
}
func (s *TempAuthKeyBindingStore) saveTx(ctx context.Context, tx pgx.Tx, b domain.TempAuthKeyBinding) error {
rawID := authKeyIDToInt64(b.TempAuthKeyID)
permID := b.PermAuthKeyID
// Every operation that may bridge temp and permanent rows enters the
// permanent identity gate before taking the raw-key row lock. This is the
// same gate/order used by selector advance and permanent revocation.
if err := lockPermanentAuthIdentities(ctx, tx, []int64{permID}); err != nil {
return err
}
var (
tempExpiry int
tempLayer int
tempObservationID int64
)
if err := tx.QueryRow(ctx, `
SELECT expires_at, layer, layer_observation_id
FROM auth_keys
WHERE auth_key_id = $1
FOR UPDATE
`, rawID).Scan(&tempExpiry, &tempLayer, &tempObservationID); err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return store.ErrAuthKeyBindingInvalid
}
return fmt.Errorf("lock temporary auth key for binding: %w", err)
}
if tempExpiry <= 0 || tempExpiry != b.ExpiresAt || rawID == permID {
return store.ErrAuthKeyBindingInvalid
}
// The raw row serializes first bind and rebind attempts. Read the binding
// only after taking that lock, so a concurrent winner is either visible or
// still waiting behind us. A different permanent identity is immutable.
var currentPermID int64
err := tx.QueryRow(ctx, `
SELECT perm_auth_key_id
FROM temp_auth_key_bindings
WHERE temp_auth_key_id = $1
`, rawID).Scan(&currentPermID)
switch {
case err == nil && currentPermID != permID:
return store.ErrTempAuthKeyAlreadyBound
case err != nil && !errors.Is(err, pgx.ErrNoRows):
return fmt.Errorf("read existing temporary auth key binding: %w", err)
}
var (
permExpiry int
permLayer int
permObservationID int64
)
if err := tx.QueryRow(ctx, `
SELECT expires_at, layer, layer_observation_id
FROM auth_keys
WHERE auth_key_id = $1
FOR UPDATE
`, permID).Scan(&permExpiry, &permLayer, &permObservationID); err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return store.ErrAuthKeyBindingInvalid
}
return fmt.Errorf("lock permanent auth key for binding: %w", err)
}
if permExpiry != 0 {
return store.ErrAuthKeyBindingInvalid
}
mergedLayer, mergedObservationID, err := store.MergeAuthKeyLayerObservations(
tempLayer, tempObservationID,
permLayer, permObservationID,
)
if err != nil {
return err
}
q := s.q.WithTx(tx)
n, err := q.UpsertTempAuthKeyBinding(ctx, sqlcgen.UpsertTempAuthKeyBindingParams{
TempAuthKeyID: authKeyIDToInt64(b.TempAuthKeyID),
PermAuthKeyID: b.PermAuthKeyID,
Nonce: b.Nonce,
@ -44,13 +124,28 @@ func (s *TempAuthKeyBindingStore) Save(ctx context.Context, b domain.TempAuthKey
return fmt.Errorf("upsert temp auth key binding: %w", err)
}
if n == 0 {
if current, found, getErr := s.GetByTemp(ctx, b.TempAuthKeyID); getErr != nil {
return getErr
} else if found && current.PermAuthKeyID != b.PermAuthKeyID {
return store.ErrTempAuthKeyAlreadyBound
}
return store.ErrAuthKeyBindingInvalid
}
keyIDs := []int64{rawID, permID}
tag, err := tx.Exec(ctx, `
UPDATE auth_keys
SET layer = $2,
layer_observation_id = $3
WHERE auth_key_id = ANY($1::bigint[])
`, keyIDs, mergedLayer, mergedObservationID)
if err != nil {
return fmt.Errorf("merge bound auth key layer defaults: %w", err)
}
if tag.RowsAffected() != int64(len(keyIDs)) {
return fmt.Errorf("merge bound auth key layer defaults: updated %d of %d locked keys", tag.RowsAffected(), len(keyIDs))
}
if _, err := tx.Exec(ctx, `
UPDATE authorizations
SET layer = $2
WHERE auth_key_id = ANY($1::bigint[])
`, keyIDs, mergedLayer); err != nil {
return fmt.Errorf("mirror bound auth key layer default: %w", err)
}
return nil
}

254
internal/store/postgres/temp_auth_key_identity_integration_test.go
View file

@ -11,7 +11,6 @@ import (
"testing"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/jackc/pgx/v5/pgxpool"
@ -133,6 +132,90 @@ func TestTempAuthKeyBindingStorePreservesHandshakeExpiryAndRejectsRebindPostgres
assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry)
}
func TestTempAuthKeyBindingStoreAtomicallyMergesLayerObservationsPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
tests := []struct {
name string
tempLayer int
tempObs int64
permLayer int
permObs int64
wantLayer int
wantObs int64
wantErr error
}{
{name: "temporary newer", tempLayer: 227, tempObs: 20, permLayer: 220, permObs: 10, wantLayer: 227, wantObs: 20},
{name: "permanent newer", tempLayer: 220, tempObs: 10, permLayer: 227, permObs: 20, wantLayer: 227, wantObs: 20},
{name: "equal ordered same layer", tempLayer: 225, tempObs: 30, permLayer: 225, permObs: 30, wantLayer: 225, wantObs: 30},
{name: "equal ordered conflict", tempLayer: 220, tempObs: 30, permLayer: 227, permObs: 30, wantErr: store.ErrAuthKeySessionLayerConflict},
{name: "legacy permanent wins", tempLayer: 220, permLayer: 227, wantLayer: 227},
}
for i, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
userID := createRevokeTestUser(t, ctx, pool, fmt.Sprintf("layer-merge-%d", i))
keys := NewAuthKeyStore(pool)
bindings := NewTempAuthKeyBindingStore(pool)
auths := NewAuthorizationStore(pool)
handshakeExpiry := int(time.Now().Add(time.Hour).Unix()) + i
tempID := saveTempIdentityTestAuthKey(t, ctx, pool, keys, handshakeExpiry)
permID := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0)
if err := auths.Bind(ctx, domain.Authorization{
AuthKeyID: permID, UserID: userID, Layer: tt.permLayer,
}); err != nil {
t.Fatalf("bind permanent authorization: %v", err)
}
if _, err := pool.Exec(ctx, `
UPDATE auth_keys SET layer = $2, layer_observation_id = $3
WHERE auth_key_id = $1`, authKeyIDToInt64(tempID), tt.tempLayer, tt.tempObs); err != nil {
t.Fatalf("seed temporary layer observation: %v", err)
}
if _, err := pool.Exec(ctx, `
UPDATE auth_keys SET layer = $2, layer_observation_id = $3
WHERE auth_key_id = $1`, authKeyIDToInt64(permID), tt.permLayer, tt.permObs); err != nil {
t.Fatalf("seed permanent layer observation: %v", err)
}
if _, err := pool.Exec(ctx, `
UPDATE authorizations SET layer = $2 WHERE auth_key_id = $1`, authKeyIDToInt64(permID), tt.permLayer); err != nil {
t.Fatalf("seed authorization layer mirror: %v", err)
}
binding := domain.TempAuthKeyBinding{
TempAuthKeyID: tempID,
PermAuthKeyID: authKeyIDToInt64(permID),
Nonce: int64(1_000 + i),
TempSessionID: int64(2_000 + i),
ExpiresAt: handshakeExpiry,
EncryptedMessage: []byte("layer merge proof"),
}
err := bindings.Save(ctx, binding)
if tt.wantErr != nil {
if !errors.Is(err, tt.wantErr) {
t.Fatalf("bind error = %v, want %v", err, tt.wantErr)
}
if _, found, getErr := bindings.GetByTemp(ctx, tempID); getErr != nil || found {
t.Fatalf("conflicting binding found=%v err=%v, want absent", found, getErr)
}
assertTempIdentityLayerTuple(t, ctx, pool, tempID, tt.tempLayer, tt.tempObs)
assertTempIdentityLayerTuple(t, ctx, pool, permID, tt.permLayer, tt.permObs)
assertTempIdentityAuthorizationLayer(t, ctx, pool, permID, tt.permLayer)
return
}
if err != nil {
t.Fatalf("bind: %v", err)
}
// A normalized proof replay is idempotent and repeats the same merge.
binding.Nonce++
if err := bindings.Save(ctx, binding); err != nil {
t.Fatalf("replay merged binding: %v", err)
}
assertTempIdentityLayerTuple(t, ctx, pool, tempID, tt.wantLayer, tt.wantObs)
assertTempIdentityLayerTuple(t, ctx, pool, permID, tt.wantLayer, tt.wantObs)
assertTempIdentityAuthorizationLayer(t, ctx, pool, permID, tt.wantLayer)
})
}
}
func TestTempAuthKeyBindingStoreConcurrentFirstBindKeepsHandshakeExpiryPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
@ -303,131 +386,6 @@ func TestTempAuthKeyBindingConcurrentWithPermanentDeleteLeavesNoDanglingStatePos
}
}
func TestAuthKeyStoreDeleteRetriesDeterministicPermanentBindingFKRacePostgres(t *testing.T) {
pool := testPool(t)
testCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
t.Cleanup(cancel)
keys := NewAuthKeyStore(pool)
bindings := NewTempAuthKeyBindingStore(pool)
auths := NewAuthorizationStore(pool)
handshakeExpiry := int(time.Now().Add(time.Hour).Unix())
temp := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, handshakeExpiry)
perm := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, 0)
userID := createRevokeTestUser(t, testCtx, pool, "deterministic-bind-delete-race")
if err := auths.Bind(testCtx, domain.Authorization{
AuthKeyID: perm,
UserID: userID,
Hash: 9401,
}); err != nil {
t.Fatalf("bind permanent authorization: %v", err)
}
candidate := domain.TempAuthKeyBinding{
TempAuthKeyID: temp,
PermAuthKeyID: authKeyIDToInt64(perm),
Nonce: 901,
TempSessionID: 902,
ExpiresAt: handshakeExpiry,
EncryptedMessage: []byte("deterministic FK retry barrier"),
}
deleteConn, err := pool.Acquire(testCtx)
if err != nil {
t.Fatalf("acquire dedicated delete connection: %v", err)
}
defer deleteConn.Release()
var deletePID int
if err := deleteConn.QueryRow(testCtx, "SELECT pg_backend_pid()").Scan(&deletePID); err != nil {
t.Fatalf("get delete backend pid: %v", err)
}
blocker, err := pool.Begin(testCtx)
if err != nil {
t.Fatalf("begin key-share blocker: %v", err)
}
defer func() { _ = blocker.Rollback(context.Background()) }()
var lockedPermID int64
if err := blocker.QueryRow(testCtx, `
SELECT auth_key_id
FROM auth_keys
WHERE auth_key_id = $1
FOR KEY SHARE`, authKeyIDToInt64(perm)).Scan(&lockedPermID); err != nil {
t.Fatalf("lock permanent key FOR KEY SHARE: %v", err)
}
if lockedPermID != authKeyIDToInt64(perm) {
t.Fatalf("locked permanent key = %d, want %d", lockedPermID, authKeyIDToInt64(perm))
}
observedDeleteDB := &permanentKeyFKRetryObservingDB{Conn: deleteConn}
deleteResult := make(chan error, 1)
go func() {
deleteResult <- NewAuthKeyStore(observedDeleteDB).Delete(testCtx, perm)
}()
waitForPostgresBackendLockWait(t, testCtx, pool, deletePID)
// This transaction already owns the compatible KEY SHARE lock needed by the
// FK check, so it can commit a new binding while the first DELETE statement
// remains blocked with a snapshot that cannot see that binding.
if err := NewTempAuthKeyBindingStore(blocker).Save(testCtx, candidate); err != nil {
t.Fatalf("save binding behind delete snapshot barrier: %v", err)
}
if err := blocker.Commit(testCtx); err != nil {
t.Fatalf("commit binding and release delete blocker: %v", err)
}
select {
case err := <-deleteResult:
if err != nil {
t.Fatalf("delete after deterministic FK retry: %v", err)
}
case <-testCtx.Done():
t.Fatalf("delete did not finish after releasing FK barrier: %v", testCtx.Err())
}
if observedDeleteDB.attempts != 2 || observedDeleteDB.fkViolations != 1 {
t.Fatalf(
"delete attempts/FK violations = %d/%d, want 2/1",
observedDeleteDB.attempts,
observedDeleteDB.fkViolations,
)
}
if _, found, err := bindings.GetByTemp(testCtx, temp); err != nil || found {
t.Fatalf("binding after deterministic retry found=%v err=%v, want absent", found, err)
}
assertTempIdentityAuthKeyMissing(t, testCtx, keys, temp)
assertTempIdentityAuthKeyMissing(t, testCtx, keys, perm)
assertRevokeTestNoAuthorization(t, testCtx, auths, perm)
}
type permanentKeyFKRetryObservingDB struct {
*pgxpool.Conn
attempts int
fkViolations int
}
func (db *permanentKeyFKRetryObservingDB) QueryRow(ctx context.Context, sql string, args ...interface{}) pgx.Row {
db.attempts++
return &permanentKeyFKRetryObservingRow{Row: db.Conn.QueryRow(ctx, sql, args...), db: db}
}
func (db *permanentKeyFKRetryObservingDB) observeFKViolation(err error) {
var pgErr *pgconn.PgError
if errors.As(err, &pgErr) && pgErr.Code == "23503" && pgErr.ConstraintName == tempAuthKeyPermFKConstraint {
db.fkViolations++
}
}
type permanentKeyFKRetryObservingRow struct {
pgx.Row
db *permanentKeyFKRetryObservingDB
}
func (row *permanentKeyFKRetryObservingRow) Scan(dest ...any) error {
err := row.Row.Scan(dest...)
row.db.observeFKViolation(err)
return err
}
func waitForPostgresBackendLockWait(
t *testing.T,
ctx context.Context,
@ -552,3 +510,47 @@ func assertTempIdentityAuthKeyMissing(
t.Fatalf("auth key %x found=%v err=%v, want absent", id, found, err)
}
}
func assertTempIdentityLayerTuple(
t *testing.T,
ctx context.Context,
pool *pgxpool.Pool,
authKeyID [8]byte,
wantLayer int,
wantObservationID int64,
) {
t.Helper()
var (
layer int
observationID int64
)
if err := pool.QueryRow(ctx, `
SELECT layer, layer_observation_id
FROM auth_keys
WHERE auth_key_id = $1`, authKeyIDToInt64(authKeyID)).Scan(&layer, &observationID); err != nil {
t.Fatalf("read auth-key layer tuple: %v", err)
}
if layer != wantLayer || observationID != wantObservationID {
t.Fatalf("auth-key layer tuple = (%d,%d), want (%d,%d)", layer, observationID, wantLayer, wantObservationID)
}
}
func assertTempIdentityAuthorizationLayer(
t *testing.T,
ctx context.Context,
pool *pgxpool.Pool,
authKeyID [8]byte,
wantLayer int,
) {
t.Helper()
var layer int
if err := pool.QueryRow(ctx, `
SELECT layer
FROM authorizations
WHERE auth_key_id = $1`, authKeyIDToInt64(authKeyID)).Scan(&layer); err != nil {
t.Fatalf("read authorization layer mirror: %v", err)
}
if layer != wantLayer {
t.Fatalf("authorization layer mirror = %d, want %d", layer, wantLayer)
}
}

27
internal/store/postgres/updatestate.go
View file

@ -55,16 +55,39 @@ func (s *UpdateStateStore) Save(ctx context.Context, id [8]byte, userID int64, s
return nil
}
func (s *UpdateStateStore) CommitDeliveredState(ctx context.Context, id [8]byte, userID int64, st domain.UpdateState, mode domain.UpdateStateCommitMode) error {
establishObserved := mode == domain.UpdateStateCommitDeliveredAndObservedBaseline
if mode != domain.UpdateStateCommitDeliveredOnly && !establishObserved {
return fmt.Errorf("commit delivered update state: invalid mode %d", mode)
}
if _, err := s.db.Exec(ctx, `
INSERT INTO update_states (auth_key_id, user_id, pts, qts, date, seq, observed_pts)
VALUES ($1, $2, $3, $4, $5, $6, CASE WHEN $7 THEN $3 ELSE 0 END)
ON CONFLICT (auth_key_id, user_id) DO UPDATE SET
pts = GREATEST(update_states.pts, EXCLUDED.pts),
qts = GREATEST(update_states.qts, EXCLUDED.qts),
date = GREATEST(update_states.date, EXCLUDED.date),
seq = GREATEST(update_states.seq, EXCLUDED.seq),
observed_pts = CASE
WHEN $7 THEN GREATEST(update_states.observed_pts, EXCLUDED.observed_pts)
ELSE update_states.observed_pts
END,
updated_at = now()`, authKeyIDToInt64(id), userID, st.Pts, st.Qts, st.Date, st.Seq, establishObserved); err != nil {
return fmt.Errorf("commit delivered update state: %w", err)
}
return nil
}
func (s *UpdateStateStore) ObserveClientState(ctx context.Context, id [8]byte, userID int64, st domain.UpdateState) error {
if st.Pts < 0 {
st.Pts = 0
}
if _, err := s.db.Exec(ctx, `
INSERT INTO update_states (auth_key_id, user_id, pts, qts, date, seq, observed_pts)
VALUES ($1, $2, $3, $4, $5, $6, $3)
VALUES ($1, $2, 0, 0, 0, 0, $3)
ON CONFLICT (auth_key_id, user_id) DO UPDATE SET
observed_pts = GREATEST(update_states.observed_pts, EXCLUDED.observed_pts),
updated_at = now()`, authKeyIDToInt64(id), userID, st.Pts, st.Qts, st.Date, st.Seq); err != nil {
updated_at = now()`, authKeyIDToInt64(id), userID, st.Pts); err != nil {
return fmt.Errorf("observe client update state: %w", err)
}
return nil

79
internal/store/postgres/updatestate_delivery_test.go Normal file
View file

@ -0,0 +1,79 @@
package postgres
import (
"context"
"strings"
"testing"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"telesrv/internal/domain"
)
type captureUpdateStateDB struct {
sql string
args []any
}
func (d *captureUpdateStateDB) Exec(_ context.Context, sql string, args ...any) (pgconn.CommandTag, error) {
d.sql = sql
d.args = append([]any(nil), args...)
return pgconn.NewCommandTag("INSERT 0 1"), nil
}
func (*captureUpdateStateDB) Query(context.Context, string, ...any) (pgx.Rows, error) {
panic("unexpected Query")
}
func (*captureUpdateStateDB) QueryRow(context.Context, string, ...any) pgx.Row {
panic("unexpected QueryRow")
}
func TestCommitDeliveredStateUsesOneAtomicBaselineUpsert(t *testing.T) {
db := &captureUpdateStateDB{}
store := NewUpdateStateStore(db)
state := domain.UpdateState{Pts: 9, Qts: 2, Date: 90, Seq: 3}
if err := store.CommitDeliveredState(context.Background(), [8]byte{4}, 1004, state, domain.UpdateStateCommitDeliveredAndObservedBaseline); err != nil {
t.Fatalf("commit baseline: %v", err)
}
if len(db.args) != 7 || db.args[6] != true {
t.Fatalf("baseline args = %#v, want final true mode", db.args)
}
for _, fragment := range []string{
"pts = GREATEST(update_states.pts, EXCLUDED.pts)",
"WHEN $7 THEN GREATEST(update_states.observed_pts, EXCLUDED.observed_pts)",
} {
if !strings.Contains(db.sql, fragment) {
t.Fatalf("atomic commit SQL missing %q:\n%s", fragment, db.sql)
}
}
}
func TestCommitDeliveredOnlyLeavesObservedUntouched(t *testing.T) {
db := &captureUpdateStateDB{}
store := NewUpdateStateStore(db)
if err := store.CommitDeliveredState(context.Background(), [8]byte{5}, 1005, domain.UpdateState{Pts: 7}, domain.UpdateStateCommitDeliveredOnly); err != nil {
t.Fatalf("commit delivered-only: %v", err)
}
if len(db.args) != 7 || db.args[6] != false {
t.Fatalf("delivered-only args = %#v, want final false mode", db.args)
}
if !strings.Contains(db.sql, "ELSE update_states.observed_pts") {
t.Fatalf("delivered-only SQL can overwrite observed:\n%s", db.sql)
}
}
func TestObserveClientStateDoesNotFabricateConfirmedCursor(t *testing.T) {
db := &captureUpdateStateDB{}
store := NewUpdateStateStore(db)
if err := store.ObserveClientState(context.Background(), [8]byte{6}, 1006, domain.UpdateState{Pts: 11, Qts: 4, Date: 110, Seq: 2}); err != nil {
t.Fatalf("observe request: %v", err)
}
if !strings.Contains(db.sql, "VALUES ($1, $2, 0, 0, 0, 0, $3)") {
t.Fatalf("observed-only insert fabricated confirmed values:\n%s", db.sql)
}
if len(db.args) != 3 || db.args[2] != 11 {
t.Fatalf("observed-only args = %#v, want auth/user/pts", db.args)
}
}