internal/store/postgres/auth_key_expiry_migration_integration_test.go
Rama main · réplica de lectura del repositorio original.
package postgres
import (
"context"
"errors"
"fmt"
"strings"
"testing"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"telesrv/deploy"
)
const (
authKeyExpiryMigrationUp = "migrations/0086_auth_key_protocol_expiry.up.sql"
authKeyExpiryMigrationDown = "migrations/0086_auth_key_protocol_expiry.down.sql"
)
func TestAuthKeyProtocolExpiryMigrationBackfillAndRollbackPostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
upSQL := readAuthKeyExpiryMigration(t, authKeyExpiryMigrationUp)
downSQL := readAuthKeyExpiryMigration(t, authKeyExpiryMigrationDown)
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin auth-key expiry migration test: %v", err)
}
defer func() { _ = tx.Rollback(context.Background()) }()
if _, err := tx.Exec(ctx, downSQL); err != nil {
t.Fatalf("return schema to 0085: %v", err)
}
base := authKeyExpiryMigrationBaseID()
tempKeyID := base
permKeyID := base - 1
authorizedPermKeyID := base - 2
unknownKeyID := base - 3
userID := base - 4
const tempExpiresAt = 1_800_086_000
for _, authKeyID := range []int64{tempKeyID, permKeyID, authorizedPermKeyID, unknownKeyID} {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, authKeyID)
}
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, tempKeyID, permKeyID, tempExpiresAt)
insertAuthKeyExpiryMigrationUser(t, ctx, tx, userID)
if _, err := tx.Exec(ctx, `
INSERT INTO public.authorizations (auth_key_id, user_id, hash)
VALUES ($1, $2, $3)`, authorizedPermKeyID, userID, base-5); err != nil {
t.Fatalf("insert authorized permanent key fixture: %v", err)
}
if _, err := tx.Exec(ctx, upSQL); err != nil {
t.Fatalf("apply auth-key expiry migration: %v", err)
}
for _, test := range []struct {
name string
authKeyID int64
want int
}{
{name: "bound temporary", authKeyID: tempKeyID, want: tempExpiresAt},
{name: "binding permanent", authKeyID: permKeyID, want: 0},
{name: "authorized permanent", authKeyID: authorizedPermKeyID, want: 0},
{name: "unclassified legacy", authKeyID: unknownKeyID, want: -1},
} {
t.Run(test.name, func(t *testing.T) {
var got int
if err := tx.QueryRow(ctx, `SELECT expires_at FROM public.auth_keys WHERE auth_key_id = $1`, test.authKeyID).Scan(&got); err != nil {
t.Fatalf("read expires_at: %v", err)
}
if got != test.want {
t.Fatalf("expires_at = %d, want %d", got, test.want)
}
})
}
var (
indexPredicate string
indexValid bool
)
if err := tx.QueryRow(ctx, `
SELECT pg_get_expr(i.indpred, i.indrelid), i.indisvalid
FROM pg_catalog.pg_index AS i
JOIN pg_catalog.pg_class AS c ON c.oid = i.indexrelid
JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = 'public'
AND c.relname = 'auth_keys_temporary_expiry_seek_idx'`).Scan(&indexPredicate, &indexValid); err != nil {
t.Fatalf("inspect temporary expiry partial index: %v", err)
}
if !indexValid || !strings.Contains(indexPredicate, "expires_at > 0") {
t.Fatalf("temporary expiry index valid=%v predicate=%q, want valid partial expires_at > 0", indexValid, indexPredicate)
}
var (
deleteAction string
fkValidated bool
)
if err := tx.QueryRow(ctx, `
SELECT c.confdeltype::text, c.convalidated
FROM pg_catalog.pg_constraint AS c
WHERE c.conrelid = 'public.temp_auth_key_bindings'::regclass
AND c.conname = 'temp_auth_key_bindings_perm_auth_key_id_fkey'`).Scan(&deleteAction, &fkValidated); err != nil {
t.Fatalf("inspect permanent auth-key FK: %v", err)
}
if deleteAction != "r" || !fkValidated {
t.Fatalf("permanent auth-key FK delete action=%q validated=%v, want RESTRICT/true", deleteAction, fkValidated)
}
assertAuthKeyExpiryMigrationForeignKeyViolation(t, ctx, tx, func(nested pgx.Tx) error {
_, err := nested.Exec(ctx, `DELETE FROM public.auth_keys WHERE auth_key_id = $1`, permKeyID)
return err
})
assertAuthKeyExpiryMigrationForeignKeyViolation(t, ctx, tx, func(nested pgx.Tx) error {
_, err := nested.Exec(ctx, `
INSERT INTO public.temp_auth_key_bindings (
temp_auth_key_id, perm_auth_key_id, nonce, expires_at, encrypted_message, temp_session_id
) VALUES ($1, $2, 86, $3, decode('86', 'hex'), 86)`, unknownKeyID, base-86, tempExpiresAt)
return err
})
if _, err := tx.Exec(ctx, downSQL); err != nil {
t.Fatalf("roll back auth-key expiry migration: %v", err)
}
var (
expiresColumnExists bool
expiryIndexExists bool
permFKExists bool
fixtureKeyCount int
)
if err := tx.QueryRow(ctx, `
SELECT
EXISTS (
SELECT 1 FROM information_schema.columns
WHERE table_schema = 'public' AND table_name = 'auth_keys' AND column_name = 'expires_at'
),
to_regclass('public.auth_keys_temporary_expiry_seek_idx') IS NOT NULL,
EXISTS (
SELECT 1 FROM pg_catalog.pg_constraint
WHERE conrelid = 'public.temp_auth_key_bindings'::regclass
AND conname = 'temp_auth_key_bindings_perm_auth_key_id_fkey'
),
(SELECT count(*) FROM public.auth_keys WHERE auth_key_id = ANY($1::bigint[]))
`, []int64{tempKeyID, permKeyID, authorizedPermKeyID, unknownKeyID}).Scan(
&expiresColumnExists,
&expiryIndexExists,
&permFKExists,
&fixtureKeyCount,
); err != nil {
t.Fatalf("inspect 0086 down result: %v", err)
}
if expiresColumnExists || expiryIndexExists || permFKExists || fixtureKeyCount != 4 {
t.Fatalf("0086 down result column=%v index=%v fk=%v keys=%d, want false/false/false/4", expiresColumnExists, expiryIndexExists, permFKExists, fixtureKeyCount)
}
}
func TestAuthKeyProtocolExpiryMigrationRejectsInvalidIdentityStatePostgres(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
upSQL := readAuthKeyExpiryMigration(t, authKeyExpiryMigrationUp)
downSQL := readAuthKeyExpiryMigration(t, authKeyExpiryMigrationDown)
tests := []struct {
name string
wantMessage string
setup func(*testing.T, context.Context, pgx.Tx, int64)
}{
{
name: "nonpositive binding expiry",
wantMessage: "invalid non-positive temporary auth key expiry",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base)
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base-1)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base, base-1, 0)
},
},
{
name: "dangling permanent key",
wantMessage: "temporary auth key binding references missing permanent key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base, base-1, 1_800_086_001)
},
},
{
name: "self binding",
wantMessage: "temporary auth key self-binding",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base, base, 1_800_086_002)
},
},
{
name: "temporary and permanent role overlap",
wantMessage: "auth key appears in both temporary and permanent roles",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
for _, authKeyID := range []int64{base, base - 1, base - 2} {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, authKeyID)
}
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base, base-1, 1_800_086_003)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base-1, base-2, 1_800_086_004)
},
},
{
name: "authorization on bound temporary key",
wantMessage: "invalid authorization on temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base)
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base-1)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base, base-1, 1_800_086_005)
insertAuthKeyExpiryMigrationUser(t, ctx, tx, base-2)
if _, err := tx.Exec(ctx, `
INSERT INTO public.authorizations (auth_key_id, user_id, hash)
VALUES ($1, $2, $3)`, base, base-2, base-3); err != nil {
t.Fatalf("insert temporary-key authorization fixture: %v", err)
}
},
},
{
name: "update state on bound temporary key",
wantMessage: "update state references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base)
insertAuthKeyExpiryMigrationKey(t, ctx, tx, base-1)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, base, base-1, 1_800_086_006)
insertAuthKeyExpiryMigrationUser(t, ctx, tx, base-2)
if _, err := tx.Exec(ctx, `
INSERT INTO public.update_states (auth_key_id, user_id, pts, date)
VALUES ($1, $2, 1, 1)`, base, base-2); err != nil {
t.Fatalf("insert temporary-key update state fixture: %v", err)
}
},
},
{
name: "bootstrap update job on bound temporary key",
wantMessage: "bootstrap update job references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_007)
userID := base - 2
messageID := base - 3
insertAuthKeyExpiryMigrationUser(t, ctx, tx, userID)
if _, err := tx.Exec(ctx, `
INSERT INTO public.private_messages (
id, sender_user_id, recipient_user_id, message_date, body
) VALUES ($1, $2, $2, 1, 'migration-0086')`, messageID, userID); err != nil {
t.Fatalf("insert bootstrap private message fixture: %v", err)
}
if _, err := tx.Exec(ctx, `
INSERT INTO public.message_boxes (
owner_user_id, box_id, private_message_id, message_sender_id,
peer_type, peer_id, from_user_id, message_date, body
) VALUES ($1, 860086, $2, $1, 'user', $1, $1, 1, 'migration-0086')`, userID, messageID); err != nil {
t.Fatalf("insert bootstrap message box fixture: %v", err)
}
if _, err := tx.Exec(ctx, `
INSERT INTO public.bootstrap_update_jobs (
kind, user_id, auth_key_id, session_id, message_box_id
) VALUES ('login_message', $1, $2, 86, 860086)`, userID, base); err != nil {
t.Fatalf("insert temporary-key bootstrap job fixture: %v", err)
}
},
},
{
name: "secret qts watermark on bound temporary key",
wantMessage: "secret qts watermark references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_008)
if _, err := tx.Exec(ctx, `
INSERT INTO public.secret_qts_watermarks (auth_key_id, reserved_qts, confirmed_qts)
VALUES ($1, 1, 1)`, base); err != nil {
t.Fatalf("insert temporary-key secret qts watermark fixture: %v", err)
}
},
},
{
name: "encrypted message queue on bound temporary key",
wantMessage: "encrypted message queue references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_009)
userID := base - 2
insertAuthKeyExpiryMigrationUser(t, ctx, tx, userID)
if _, err := tx.Exec(ctx, `
INSERT INTO public.encrypted_message_queue (
receiver_auth_key_id, qts, receiver_user_id, chat_id, random_id, date, bytes
) VALUES ($1, 1, $2, 860086, $3, 1, decode('86', 'hex'))`, base, userID, base-3); err != nil {
t.Fatalf("insert temporary-key encrypted message queue fixture: %v", err)
}
},
},
{
name: "encrypted state delivery on bound temporary key",
wantMessage: "encrypted state delivery references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_010)
userID := base - 2
eventID := base - 3
insertAuthKeyExpiryMigrationUser(t, ctx, tx, userID)
if _, err := tx.Exec(ctx, `
INSERT INTO public.encrypted_state_events (
id, target_user_id, target_auth_key_id, chat_id, event_type, date
) VALUES ($1, $2, $3, 860086, 1, 1)`, eventID, userID, base-1); err != nil {
t.Fatalf("insert permanent-key encrypted state event fixture: %v", err)
}
if _, err := tx.Exec(ctx, `
INSERT INTO public.encrypted_state_event_delivery (event_id, auth_key_id)
VALUES ($1, $2)`, eventID, base); err != nil {
t.Fatalf("insert temporary-key encrypted state delivery fixture: %v", err)
}
},
},
{
name: "encrypted state event on bound temporary key",
wantMessage: "encrypted state event targets temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_011)
userID := base - 2
insertAuthKeyExpiryMigrationUser(t, ctx, tx, userID)
if _, err := tx.Exec(ctx, `
INSERT INTO public.encrypted_state_events (
id, target_user_id, target_auth_key_id, chat_id, event_type, date
) VALUES ($1, $2, $3, 860086, 1, 1)`, base-3, userID, base); err != nil {
t.Fatalf("insert temporary-key encrypted state event fixture: %v", err)
}
},
},
{
name: "secret chat admin on bound temporary key",
wantMessage: "secret chat references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_012)
adminUserID := base - 2
participantUserID := base - 3
insertAuthKeyExpiryMigrationUser(t, ctx, tx, adminUserID)
insertAuthKeyExpiryMigrationUser(t, ctx, tx, participantUserID)
insertAuthKeyExpiryMigrationSecretChat(t, ctx, tx, base, base-1, adminUserID, participantUserID)
},
},
{
name: "secret chat participant on bound temporary key",
wantMessage: "secret chat references temporary auth key",
setup: func(t *testing.T, ctx context.Context, tx pgx.Tx, base int64) {
insertAuthKeyExpiryMigrationBoundPair(t, ctx, tx, base, base-1, 1_800_086_013)
adminUserID := base - 2
participantUserID := base - 3
insertAuthKeyExpiryMigrationUser(t, ctx, tx, adminUserID)
insertAuthKeyExpiryMigrationUser(t, ctx, tx, participantUserID)
insertAuthKeyExpiryMigrationSecretChat(t, ctx, tx, base-1, base, adminUserID, participantUserID)
},
},
}
for i, test := range tests {
t.Run(test.name, func(t *testing.T) {
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin invalid-state migration test: %v", err)
}
defer func() { _ = tx.Rollback(context.Background()) }()
if _, err := tx.Exec(ctx, downSQL); err != nil {
t.Fatalf("return schema to 0085: %v", err)
}
test.setup(t, ctx, tx, authKeyExpiryMigrationBaseID()-int64(i*100))
_, err = tx.Exec(ctx, upSQL)
if err == nil {
t.Fatal("0086 migration accepted invalid identity state")
}
var pgErr *pgconn.PgError
if !errors.As(err, &pgErr) || pgErr.Code != "P0001" || !strings.Contains(pgErr.Message, test.wantMessage) {
t.Fatalf("0086 migration error = %v, want SQLSTATE P0001 containing %q", err, test.wantMessage)
}
})
}
}
func readAuthKeyExpiryMigration(t *testing.T, name string) string {
t.Helper()
sql, err := deploy.Migrations.ReadFile(name)
if err != nil {
t.Fatalf("read %s: %v", name, err)
}
return string(sql)
}
func authKeyExpiryMigrationBaseID() int64 {
return -(time.Now().UnixNano() & 0x3fffffffffffffff)
}
func insertAuthKeyExpiryMigrationKey(t *testing.T, ctx context.Context, tx pgx.Tx, authKeyID int64) {
t.Helper()
if _, err := tx.Exec(ctx, `
INSERT INTO public.auth_keys (auth_key_id, body, server_salt)
VALUES ($1, decode('86', 'hex'), 86)`, authKeyID); err != nil {
t.Fatalf("insert auth key %d: %v", authKeyID, err)
}
}
func insertAuthKeyExpiryMigrationBinding(t *testing.T, ctx context.Context, tx pgx.Tx, tempKeyID, permKeyID int64, expiresAt int) {
t.Helper()
if _, err := tx.Exec(ctx, `
INSERT INTO public.temp_auth_key_bindings (
temp_auth_key_id, perm_auth_key_id, nonce, expires_at, encrypted_message, temp_session_id
) VALUES ($1, $2, 86, $3, decode('86', 'hex'), 86)`, tempKeyID, permKeyID, expiresAt); err != nil {
t.Fatalf("insert temp auth-key binding %d -> %d: %v", tempKeyID, permKeyID, err)
}
}
func insertAuthKeyExpiryMigrationBoundPair(t *testing.T, ctx context.Context, tx pgx.Tx, tempKeyID, permKeyID int64, expiresAt int) {
t.Helper()
insertAuthKeyExpiryMigrationKey(t, ctx, tx, tempKeyID)
insertAuthKeyExpiryMigrationKey(t, ctx, tx, permKeyID)
insertAuthKeyExpiryMigrationBinding(t, ctx, tx, tempKeyID, permKeyID, expiresAt)
}
func insertAuthKeyExpiryMigrationUser(t *testing.T, ctx context.Context, tx pgx.Tx, userID int64) {
t.Helper()
phone := fmt.Sprintf("+860086%d", -userID)
if _, err := tx.Exec(ctx, `
INSERT INTO public.users (id, access_hash, phone, first_name)
VALUES ($1, $2, $3, 'migration-0086')`, userID, userID-1, phone); err != nil {
t.Fatalf("insert migration user %d: %v", userID, err)
}
}
func insertAuthKeyExpiryMigrationSecretChat(
t *testing.T,
ctx context.Context,
tx pgx.Tx,
adminAuthKeyID int64,
participantAuthKeyID int64,
adminUserID int64,
participantUserID int64,
) {
t.Helper()
if _, err := tx.Exec(ctx, `
INSERT INTO public.secret_chats (
chat_id, admin_access_hash, participant_access_hash,
admin_user_id, admin_auth_key_id, participant_user_id, participant_auth_key_id,
state, random_id, date
) VALUES (
860086, 86, 87,
$1, $2, $3, $4,
'waiting', 860086, 1
)`, adminUserID, adminAuthKeyID, participantUserID, participantAuthKeyID); err != nil {
t.Fatalf("insert temporary-key secret chat fixture: %v", err)
}
}
func assertAuthKeyExpiryMigrationForeignKeyViolation(t *testing.T, ctx context.Context, tx pgx.Tx, action func(pgx.Tx) error) {
t.Helper()
nested, err := tx.Begin(ctx)
if err != nil {
t.Fatalf("begin FK assertion savepoint: %v", err)
}
defer func() { _ = nested.Rollback(context.Background()) }()
err = action(nested)
if err == nil {
t.Fatal("operation bypassed permanent auth-key RESTRICT FK")
}
var pgErr *pgconn.PgError
if !errors.As(err, &pgErr) || pgErr.Code != "23503" || pgErr.ConstraintName != "temp_auth_key_bindings_perm_auth_key_id_fkey" {
t.Fatalf("FK error = %v, want SQLSTATE 23503 from temp_auth_key_bindings_perm_auth_key_id_fkey", err)
}
}