diff --git a/deploy/migrations/0086_auth_key_protocol_expiry.down.sql b/deploy/migrations/0086_auth_key_protocol_expiry.down.sql new file mode 100644 index 00000000..684c2434 --- /dev/null +++ b/deploy/migrations/0086_auth_key_protocol_expiry.down.sql @@ -0,0 +1,10 @@ +ALTER TABLE public.temp_auth_key_bindings + DROP CONSTRAINT IF EXISTS temp_auth_key_bindings_perm_auth_key_id_fkey; + +DROP INDEX IF EXISTS public.auth_keys_temporary_expiry_seek_idx; + +ALTER TABLE public.auth_keys + DROP CONSTRAINT IF EXISTS auth_keys_expires_at_valid; + +ALTER TABLE public.auth_keys + DROP COLUMN IF EXISTS expires_at; diff --git a/deploy/migrations/0086_auth_key_protocol_expiry.up.sql b/deploy/migrations/0086_auth_key_protocol_expiry.up.sql new file mode 100644 index 00000000..f3092f70 --- /dev/null +++ b/deploy/migrations/0086_auth_key_protocol_expiry.up.sql @@ -0,0 +1,161 @@ +-- Preserve the protocol key kind/lifetime established by p_q_inner_data(_temp). +-- A temporary key must expire at the MTProto edge; it must never fall through +-- to the RPC router and be mistaken for an independent permanent identity. +ALTER TABLE public.auth_keys + ADD COLUMN expires_at integer NOT NULL DEFAULT -1; + +ALTER TABLE public.auth_keys + ADD CONSTRAINT auth_keys_expires_at_valid CHECK (expires_at >= -1); + +-- auth_keys.expires_at is the only protocol-lifetime fact. Retention seeks this +-- partial index so unbound temporary handshakes and bound keys follow the same +-- bounded cleanup path; the binding-side expiry index remains only for legacy +-- rollback compatibility. +CREATE INDEX auth_keys_temporary_expiry_seek_idx + ON public.auth_keys (expires_at, auth_key_id) + WHERE expires_at > 0; + +-- Existing bound temporary keys are unambiguous and can be backfilled from the +-- durable bind record. Unclassified historical keys remain -1 and are rejected +-- once with protocol -404 after restart, forcing a clean handshake instead of +-- guessing that they are permanent. New handshakes always write 0 (permanent) +-- or their positive absolute expiry before dh_gen_ok. +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 + FROM public.temp_auth_key_bindings + WHERE expires_at <= 0 + ) THEN + RAISE EXCEPTION 'invalid non-positive temporary auth key expiry; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 + FROM public.temp_auth_key_bindings AS b + LEFT JOIN public.auth_keys AS k ON k.auth_key_id = b.perm_auth_key_id + WHERE k.auth_key_id IS NULL + ) THEN + RAISE EXCEPTION 'temporary auth key binding references missing permanent key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 + FROM public.temp_auth_key_bindings + WHERE temp_auth_key_id = perm_auth_key_id + ) THEN + RAISE EXCEPTION 'temporary auth key self-binding; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 + FROM public.temp_auth_key_bindings AS temp_role + JOIN public.temp_auth_key_bindings AS perm_role + ON perm_role.perm_auth_key_id = temp_role.temp_auth_key_id + ) THEN + RAISE EXCEPTION 'auth key appears in both temporary and permanent roles; repair before migration 0086'; + END IF; +END +$$; + +UPDATE public.auth_keys AS k +SET expires_at = b.expires_at +FROM public.temp_auth_key_bindings AS b +WHERE k.auth_key_id = b.temp_auth_key_id; + +-- Do not normalize an early telesrv bug during reads. If a deployment contains +-- an authorization written against a bound temp key, stop the migration and +-- require an explicit data repair after inspecting the corresponding perm key. +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 + FROM public.authorizations AS a + JOIN public.temp_auth_key_bindings AS b + ON b.temp_auth_key_id = a.auth_key_id + ) THEN + RAISE EXCEPTION 'invalid authorization on temporary auth key; repair before migration 0086'; + END IF; +END +$$; + +-- Every durable table keyed by business/device auth identity must also be free +-- of bound temporary IDs. These tables intentionally do not FK to auth_keys +-- because several retain historical delivery facts; silently deleting the key +-- would therefore strand an identity split instead of repairing it. The +-- physical-session exclusion tuple in dispatch_outbox is deliberately omitted: +-- it stores raw auth_key_id + session_id and a temporary raw key is valid there. +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM public.update_states AS s + JOIN public.temp_auth_key_bindings AS b ON b.temp_auth_key_id = s.auth_key_id + ) THEN + RAISE EXCEPTION 'update state references temporary auth key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 FROM public.bootstrap_update_jobs AS j + JOIN public.temp_auth_key_bindings AS b ON b.temp_auth_key_id = j.auth_key_id + ) THEN + RAISE EXCEPTION 'bootstrap update job references temporary auth key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 FROM public.secret_qts_watermarks AS q + JOIN public.temp_auth_key_bindings AS b ON b.temp_auth_key_id = q.auth_key_id + ) THEN + RAISE EXCEPTION 'secret qts watermark references temporary auth key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 FROM public.encrypted_message_queue AS q + JOIN public.temp_auth_key_bindings AS b ON b.temp_auth_key_id = q.receiver_auth_key_id + ) THEN + RAISE EXCEPTION 'encrypted message queue references temporary auth key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 FROM public.encrypted_state_event_delivery AS d + JOIN public.temp_auth_key_bindings AS b ON b.temp_auth_key_id = d.auth_key_id + ) THEN + RAISE EXCEPTION 'encrypted state delivery references temporary auth key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 FROM public.encrypted_state_events AS e + JOIN public.temp_auth_key_bindings AS b ON b.temp_auth_key_id = e.target_auth_key_id + ) THEN + RAISE EXCEPTION 'encrypted state event targets temporary auth key; repair before migration 0086'; + END IF; + IF EXISTS ( + SELECT 1 FROM public.secret_chats AS c + JOIN public.temp_auth_key_bindings AS b + ON b.temp_auth_key_id = c.admin_auth_key_id + OR b.temp_auth_key_id = c.participant_auth_key_id + ) THEN + RAISE EXCEPTION 'secret chat references temporary auth key; repair before migration 0086'; + END IF; +END +$$; + +-- An authorization or the permanent side of a temp binding proves that the key +-- is permanent. Logged-out, unreferenced legacy keys cannot be proven either +-- way and intentionally keep the -1 sentinel described above. +UPDATE public.auth_keys AS k +SET expires_at = 0 +WHERE k.expires_at = -1 + AND ( + EXISTS ( + SELECT 1 + FROM public.authorizations AS a + WHERE a.auth_key_id = k.auth_key_id + ) + OR EXISTS ( + SELECT 1 + FROM public.temp_auth_key_bindings AS b + WHERE b.perm_auth_key_id = k.auth_key_id + ) + ); + +-- This FK is the durable serialization boundary between auth.bindTempAuthKey +-- and permanent-key revoke/destroy. RESTRICT makes a concurrent delete fail +-- closed; deletion paths remove referenced temp keys first and retry on the +-- narrow FK race, so no committed binding can ever point at a missing perm key. +ALTER TABLE public.temp_auth_key_bindings + ADD CONSTRAINT temp_auth_key_bindings_perm_auth_key_id_fkey + FOREIGN KEY (perm_auth_key_id) + REFERENCES public.auth_keys(auth_key_id) + ON DELETE RESTRICT; diff --git a/internal/app/auth/service.go b/internal/app/auth/service.go index 5c4fd2ad..893b18fa 100644 --- a/internal/app/auth/service.go +++ b/internal/app/auth/service.go @@ -27,6 +27,10 @@ var ( ErrCodeExpired = errors.New("phone code expired or not found") ErrCodeInvalid = errors.New("phone code invalid") ErrEncryptedMessageInvalid = errors.New("encrypted message invalid") + ErrExpiresAtInvalid = errors.New("temporary auth key request expiry invalid") + ErrTempAuthKeyEmpty = errors.New("temporary auth key missing or expired") + ErrTempAuthKeyAlreadyBound = errors.New("temporary auth key already bound") + ErrAuthKeyPermEmpty = errors.New("permanent auth key required") // ErrLoginCodeDeliveryUnavailable 表示已有账号的 app-code 没有可用的 // durable message/event/outbox 投递边界。这是服务端配置错误，不能降级成 // “继续返回 sentCode，等 signIn 后补发”。 @@ -193,26 +197,44 @@ func NewService(users store.UserStore, auths store.AuthorizationStore, codes sto // BindTempAuthKey 校验并记录 TDesktop PFS temp→perm auth key 绑定。 func (s *Service) BindTempAuthKey(ctx context.Context, sessionID int64, binding domain.TempAuthKeyBinding) error { if s.authKeys != nil { - inner, err := s.validateBindTempAuthKey(ctx, sessionID, binding) + inner, protocolExpiresAt, err := s.validateBindTempAuthKey(ctx, sessionID, binding) if err != nil { return err } binding.TempSessionID = inner.TempSessionID + // The bind request's expires_at is a signed client assertion. TDesktop + // intentionally adds a small grace interval, while Android derives its + // value at handshake completion. Retention and edge admission must use the + // server's p_q_inner_data_temp lifetime, never the client value. + binding.ExpiresAt = protocolExpiresAt + } + if binding.ExpiresAt <= int(time.Now().Unix()) { + // The edge may admit the frame immediately before the temporary key's + // absolute boundary and the encrypted proof may cross it. This is a temp-key + // rotation condition, never a destructive permanent-key proof failure. + return ErrTempAuthKeyEmpty } if s.tempKeys == nil { return nil } - return s.tempKeys.Save(ctx, binding) + if err := s.tempKeys.Save(ctx, binding); err != nil { + if errors.Is(err, store.ErrTempAuthKeyAlreadyBound) { + return ErrTempAuthKeyAlreadyBound + } + if errors.Is(err, store.ErrAuthKeyBindingInvalid) { + return s.classifyBindingStoreInvalid(ctx, binding) + } + return err + } + return nil } // ResolveAuthKey 将已绑定的 temp auth_key 解析为对应 perm auth_key。 // -// 过期处理是有意的连续性权衡（见 TestResolveAuthKeyAllowsExpiredTempBindingForAuthorizedPermKey）： -// temp 绑定 expires_at 已过时，仅当 perm key 也未授权才拒绝；perm 仍授权则继续解析， -// 避免已登录会话因 temp key 过期而被强制踢下线。严格 PFS 要求过期 temp key 一律失效 -// （不以 perm 授权豁免），但收紧前需先核实目标客户端（TDesktop/DrKLO）会在过期前主动 -// 轮换 temp key 并优雅处理拒绝，否则会造成在线会话掉线。RetentionWorker 的 DeleteExpired -// 已把残留窗口限制在 expires_at + 宽限（约 24h）内。收紧为显式硬化任务，需客户端验证。 +// temp→perm 是握手/绑定形成的协议身份关系，与 perm 当前是否登录完全无关。即使 +// auth.logOut 已删除 authorization，只要绑定仍存在，后续登录 RPC 也必须继续落到同一 +// perm key，绝不能把 raw temp key 当成新的业务身份。协议过期由 mtprotoedge 在解密/RPC +// 之前返回 -404 并关闭连接；这里不再用 authorization 状态猜测 key 类型。 func (s *Service) ResolveAuthKey(ctx context.Context, authKeyID [8]byte) ([8]byte, bool, error) { if s == nil || s.tempKeys == nil { return [8]byte{}, false, nil @@ -221,19 +243,7 @@ func (s *Service) ResolveAuthKey(ctx context.Context, authKeyID [8]byte) ([8]byt if err != nil || !found { return [8]byte{}, found, err } - permID := authKeyIDFromInt64(binding.PermAuthKeyID) - if binding.ExpiresAt <= int(time.Now().Unix()) && !s.permAuthKeyAuthorized(ctx, permID) { - return [8]byte{}, false, nil - } - return permID, true, nil -} - -func (s *Service) permAuthKeyAuthorized(ctx context.Context, authKeyID [8]byte) bool { - if s == nil || s.auths == nil { - return false - } - _, found, err := s.auths.ByAuthKey(ctx, authKeyID) - return err == nil && found + return authKeyIDFromInt64(binding.PermAuthKeyID), true, nil } // UserID 返回 auth_key 当前绑定的用户。未登录、或两步验证未完成时 found=false。 @@ -1228,11 +1238,29 @@ func (s *Service) authorizationsByUserExcept(ctx context.Context, userID int64, } func (s *Service) bind(ctx context.Context, auth domain.Authorization, userID int64) error { + if s.authKeys != nil { + key, found, err := s.authKeys.Get(ctx, auth.AuthKeyID) + if err != nil { + return err + } + // Defense in depth: Router normally converts a bound temp key to its perm + // identity and edge rejects expired temp keys. Never let an unbound/sticky + // temp key create authorization even if either outer boundary regresses. + if !found || key.ExpiresAt != 0 { + return ErrAuthKeyPermEmpty + } + } auth.UserID = userID // Bind 是授权切换的持久化状态边界：生产 store 会先清同 auth key 的旧用户 // update state，再原子建立新用户 baseline。RPC 层不得在 Bind 成功后清整个 key， // 否则会把刚建立的 retained-floor checkpoint 一并删除。 - return s.auths.Bind(ctx, auth) + if err := s.auths.Bind(ctx, auth); err != nil { + if errors.Is(err, store.ErrAuthKeyNotPermanent) { + return ErrAuthKeyPermEmpty + } + return err + } + return nil } func (s *Service) passwordNeeded(ctx context.Context, userID int64) (bool, error) { @@ -1282,32 +1310,60 @@ func (s *Service) recordLoginMessage(ctx context.Context, userID int64, code str return msg, nil } -func (s *Service) validateBindTempAuthKey(ctx context.Context, sessionID int64, binding domain.TempAuthKeyBinding) (mtcrypto.BindAuthKeyInner, error) { +func (s *Service) validateBindTempAuthKey(ctx context.Context, sessionID int64, binding domain.TempAuthKeyBinding) (mtcrypto.BindAuthKeyInner, int, error) { if binding.ExpiresAt <= int(time.Now().Unix()) { - return mtcrypto.BindAuthKeyInner{}, ErrEncryptedMessageInvalid + return mtcrypto.BindAuthKeyInner{}, 0, ErrExpiresAtInvalid + } + temp, found, err := s.authKeys.Get(ctx, binding.TempAuthKeyID) + if err != nil { + return mtcrypto.BindAuthKeyInner{}, 0, err + } + // expires_at in auth.bindTempAuthKey is client-supplied and must only attest + // to a still-live binding. It may never create or reclassify a protocol key; + // the caller normalizes durable retention to this handshake-authoritative + // temp.ExpiresAt instead of trusting the client value. + if !found || temp.ExpiresAt <= int(time.Now().Unix()) { + return mtcrypto.BindAuthKeyInner{}, 0, ErrTempAuthKeyEmpty } permID := authKeyIDFromInt64(binding.PermAuthKeyID) perm, found, err := s.authKeys.Get(ctx, permID) if err != nil { - return mtcrypto.BindAuthKeyInner{}, err + return mtcrypto.BindAuthKeyInner{}, 0, err } - if !found { - return mtcrypto.BindAuthKeyInner{}, ErrEncryptedMessageInvalid + if !found || perm.ExpiresAt != 0 { + return mtcrypto.BindAuthKeyInner{}, 0, ErrEncryptedMessageInvalid } inner, err := decryptBindAuthKeyInner(perm, binding.EncryptedMessage) if err != nil { - return mtcrypto.BindAuthKeyInner{}, ErrEncryptedMessageInvalid + return mtcrypto.BindAuthKeyInner{}, 0, ErrEncryptedMessageInvalid } if inner.Nonce != binding.Nonce || inner.TempAuthKeyID != authKeyIDInt64(binding.TempAuthKeyID) || inner.PermAuthKeyID != binding.PermAuthKeyID || inner.TempSessionID != sessionID || inner.ExpiresAt != binding.ExpiresAt { - return mtcrypto.BindAuthKeyInner{}, ErrEncryptedMessageInvalid + return mtcrypto.BindAuthKeyInner{}, 0, ErrEncryptedMessageInvalid } - return inner, nil + if temp.ExpiresAt <= int(time.Now().Unix()) { + return mtcrypto.BindAuthKeyInner{}, 0, ErrTempAuthKeyEmpty + } + return inner, temp.ExpiresAt, nil +} + +func (s *Service) classifyBindingStoreInvalid(ctx context.Context, binding domain.TempAuthKeyBinding) error { + if s == nil || s.authKeys == nil { + return ErrEncryptedMessageInvalid + } + temp, found, err := s.authKeys.Get(ctx, binding.TempAuthKeyID) + if err != nil { + return err + } + if !found || temp.ExpiresAt <= int(time.Now().Unix()) { + return ErrTempAuthKeyEmpty + } + return ErrEncryptedMessageInvalid } func decryptBindAuthKeyInner(perm store.AuthKeyData, encrypted []byte) (mtcrypto.BindAuthKeyInner, error) { diff --git a/internal/app/auth/service_test.go b/internal/app/auth/service_test.go index f8016554..c0c802d2 100644 --- a/internal/app/auth/service_test.go +++ b/internal/app/auth/service_test.go @@ -18,11 +18,12 @@ import ( func TestBindTempAuthKeyValidatesEncryptedMessage(t *testing.T) { ctx := context.Background() keys := memory.NewAuthKeyStore() - tempBindings := memory.NewTempAuthKeyBindingStore() + tempBindings := memory.NewTempAuthKeyBindingStore(keys) permKey := testAuthKey(0x11) tempKey := testAuthKey(0x55) + expiresAt := int(time.Now().Add(time.Hour).Unix()) saveAuthKey(t, keys, permKey) - saveAuthKey(t, keys, tempKey) + saveAuthKeyWithExpiry(t, keys, tempKey, expiresAt) svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), keys, tempBindings, "12345") @@ -31,7 +32,6 @@ func TestBindTempAuthKeyValidatesEncryptedMessage(t *testing.T) { sessionID = int64(0x1020304050) msgID = int64(0x0102030405060708) ) - expiresAt := int(time.Now().Add(time.Hour).Unix()) encrypted, err := mtcrypto.EncryptBindMessage( bytes.NewReader(bytes.Repeat([]byte{0xCD}, 128)), permKey, @@ -69,6 +69,70 @@ func TestBindTempAuthKeyValidatesEncryptedMessage(t *testing.T) { if !errors.Is(err, ErrEncryptedMessageInvalid) { t.Fatalf("BindTempAuthKey wrong session err = %v, want ErrEncryptedMessageInvalid", err) } + + // TDesktop intentionally adds a 30-second bind grace to the expiry it + // derived from p_q_inner_data_temp. The request is valid, but the durable + // binding must be normalized back to the server handshake expiry. + extendedExpiry := expiresAt + 30 + extendedEncrypted, err := mtcrypto.EncryptBindMessage( + bytes.NewReader(bytes.Repeat([]byte{0xCE}, 128)), + permKey, + msgID+4, + &mtcrypto.BindAuthKeyInner{ + Nonce: nonce, + TempAuthKeyID: tempKey.IntID(), + PermAuthKeyID: permKey.IntID(), + TempSessionID: sessionID, + ExpiresAt: extendedExpiry, + }, + ) + if err != nil { + t.Fatalf("encrypt extended bind message: %v", err) + } + err = svc.BindTempAuthKey(ctx, sessionID, domain.TempAuthKeyBinding{ + TempAuthKeyID: tempKey.ID, + PermAuthKeyID: permKey.IntID(), + Nonce: nonce, + ExpiresAt: extendedExpiry, + EncryptedMessage: extendedEncrypted, + }) + if err != nil { + t.Fatalf("BindTempAuthKey TDesktop grace expiry: %v", err) + } + stored, found, getErr := tempBindings.GetByTemp(ctx, tempKey.ID) + if getErr != nil || !found || stored.ExpiresAt != expiresAt { + t.Fatalf("stored binding after extension attempt = %+v found=%v err=%v", stored, found, getErr) + } +} + +func TestBindTempAuthKeyClassifiesExpiryWithoutDestroyingPermanentKey(t *testing.T) { + ctx := context.Background() + keys := memory.NewAuthKeyStore() + tempBindings := memory.NewTempAuthKeyBindingStore(keys) + permKey := testAuthKey(0x31) + tempKey := testAuthKey(0x32) + saveAuthKey(t, keys, permKey) + saveAuthKeyWithExpiry(t, keys, tempKey, int(time.Now().Add(-time.Second).Unix())) + svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), keys, tempBindings, "12345") + + request := domain.TempAuthKeyBinding{ + TempAuthKeyID: tempKey.ID, + PermAuthKeyID: permKey.IntID(), + ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + } + if err := svc.BindTempAuthKey(ctx, 1001, request); !errors.Is(err, ErrTempAuthKeyEmpty) { + t.Fatalf("expired protocol temp key err = %v, want ErrTempAuthKeyEmpty", err) + } + + request.TempAuthKeyID = testAuthKey(0x33).ID + if err := svc.BindTempAuthKey(ctx, 1001, request); !errors.Is(err, ErrTempAuthKeyEmpty) { + t.Fatalf("missing protocol temp key err = %v, want ErrTempAuthKeyEmpty", err) + } + + request.ExpiresAt = int(time.Now().Add(-time.Second).Unix()) + if err := svc.BindTempAuthKey(ctx, 1001, request); !errors.Is(err, ErrExpiresAtInvalid) { + t.Fatalf("expired request proof err = %v, want ErrExpiresAtInvalid", err) + } } func TestUpdateAuthKeyClientInfoConvergesAuthorizationMetadata(t *testing.T) { @@ -115,15 +179,19 @@ func TestUpdateAuthKeyClientInfoConvergesAuthorizationMetadata(t *testing.T) { func TestResolveAuthKeyUsesValidTempBinding(t *testing.T) { ctx := context.Background() - tempBindings := memory.NewTempAuthKeyBindingStore() + keys := memory.NewAuthKeyStore() + tempBindings := memory.NewTempAuthKeyBindingStore(keys) permKey := testAuthKey(0x11) tempKey := testAuthKey(0x55) - svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, tempBindings, "12345") + expiresAt := int(time.Now().Add(time.Hour).Unix()) + saveAuthKey(t, keys, permKey) + saveAuthKeyWithExpiry(t, keys, tempKey, expiresAt) + svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), keys, tempBindings, "12345") if err := tempBindings.Save(ctx, domain.TempAuthKeyBinding{ TempAuthKeyID: tempKey.ID, PermAuthKeyID: permKey.IntID(), - ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + ExpiresAt: expiresAt, }); err != nil { t.Fatalf("save temp binding: %v", err) } @@ -139,11 +207,15 @@ func TestResolveAuthKeyUsesValidTempBinding(t *testing.T) { func TestResolveAuthKeyAllowsExpiredTempBindingForAuthorizedPermKey(t *testing.T) { ctx := context.Background() - tempBindings := memory.NewTempAuthKeyBindingStore() + keys := memory.NewAuthKeyStore() + tempBindings := memory.NewTempAuthKeyBindingStore(keys) authz := memory.NewAuthorizationStore() permKey := testAuthKey(0x21) tempKey := testAuthKey(0x65) - svc := NewService(memory.NewUserStore(), authz, memory.NewCodeStore(), nil, tempBindings, "12345") + expiresAt := int(time.Now().Add(-time.Minute).Unix()) + saveAuthKey(t, keys, permKey) + saveAuthKeyWithExpiry(t, keys, tempKey, expiresAt) + svc := NewService(memory.NewUserStore(), authz, memory.NewCodeStore(), keys, tempBindings, "12345") if err := authz.Bind(ctx, domain.Authorization{AuthKeyID: permKey.ID, UserID: 1000000001}); err != nil { t.Fatalf("bind authorization: %v", err) @@ -151,7 +223,7 @@ func TestResolveAuthKeyAllowsExpiredTempBindingForAuthorizedPermKey(t *testing.T if err := tempBindings.Save(ctx, domain.TempAuthKeyBinding{ TempAuthKeyID: tempKey.ID, PermAuthKeyID: permKey.IntID(), - ExpiresAt: int(time.Now().Add(-time.Minute).Unix()), + ExpiresAt: expiresAt, }); err != nil { t.Fatalf("save temp binding: %v", err) } @@ -165,17 +237,21 @@ func TestResolveAuthKeyAllowsExpiredTempBindingForAuthorizedPermKey(t *testing.T } } -func TestResolveAuthKeyRejectsExpiredTempBindingWithoutAuthorizedPermKey(t *testing.T) { +func TestResolveAuthKeyKeepsExpiredBindingCanonicalWithoutAuthorization(t *testing.T) { ctx := context.Background() - tempBindings := memory.NewTempAuthKeyBindingStore() + keys := memory.NewAuthKeyStore() + tempBindings := memory.NewTempAuthKeyBindingStore(keys) permKey := testAuthKey(0x31) tempKey := testAuthKey(0x75) - svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, tempBindings, "12345") + expiresAt := int(time.Now().Add(-time.Minute).Unix()) + saveAuthKey(t, keys, permKey) + saveAuthKeyWithExpiry(t, keys, tempKey, expiresAt) + svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), keys, tempBindings, "12345") if err := tempBindings.Save(ctx, domain.TempAuthKeyBinding{ TempAuthKeyID: tempKey.ID, PermAuthKeyID: permKey.IntID(), - ExpiresAt: int(time.Now().Add(-time.Minute).Unix()), + ExpiresAt: expiresAt, }); err != nil { t.Fatalf("save temp binding: %v", err) } @@ -184,8 +260,87 @@ func TestResolveAuthKeyRejectsExpiredTempBindingWithoutAuthorizedPermKey(t *test if err != nil { t.Fatalf("ResolveAuthKey: %v", err) } - if ok || got != ([8]byte{}) { - t.Fatalf("resolved = %x ok=%v, want expired unresolved", got, ok) + if !ok || got != permKey.ID { + t.Fatalf("resolved = %x ok=%v, want canonical perm %x even while logged out", got, ok, permKey.ID) + } +} + +func TestExpiredTempLogoutReloginNeverAuthorizesRawTempKey(t *testing.T) { + ctx := context.Background() + users := memory.NewUserStore() + authz := memory.NewAuthorizationStore() + keys := memory.NewAuthKeyStore() + tempBindings := memory.NewTempAuthKeyBindingStore(keys) + permKey := testAuthKey(0x41) + tempKey := testAuthKey(0x81) + expiresAt := int(time.Now().Add(-time.Minute).Unix()) + if err := keys.Save(ctx, store.AuthKeyData{ID: permKey.ID}); err != nil { + t.Fatalf("save perm key: %v", err) + } + if err := keys.Save(ctx, store.AuthKeyData{ + ID: tempKey.ID, ExpiresAt: expiresAt, + }); err != nil { + t.Fatalf("save temp key: %v", err) + } + if err := tempBindings.Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: tempKey.ID, + PermAuthKeyID: permKey.IntID(), + ExpiresAt: expiresAt, + }); err != nil { + t.Fatalf("save temp binding: %v", err) + } + bob, err := users.Create(ctx, domain.User{Phone: "15550008101", FirstName: "Bob"}) + if err != nil { + t.Fatalf("create Bob: %v", err) + } + alice, err := users.Create(ctx, domain.User{Phone: "15550008102", FirstName: "Alice"}) + if err != nil { + t.Fatalf("create Alice: %v", err) + } + if err := authz.Bind(ctx, domain.Authorization{AuthKeyID: permKey.ID, UserID: bob.ID}); err != nil { + t.Fatalf("authorize Bob: %v", err) + } + + svc := NewService(users, authz, memory.NewCodeStore(), keys, tempBindings, "12345") + if err := svc.LogOut(ctx, permKey.ID); err != nil { + t.Fatalf("logout Bob: %v", err) + } + resolved, ok, err := svc.ResolveAuthKey(ctx, tempKey.ID) + if err != nil || !ok || resolved != permKey.ID { + t.Fatalf("resolve after logout = %x/%v/%v, want perm", resolved, ok, err) + } + if _, err := svc.BindVerifiedLogin(ctx, domain.Authorization{AuthKeyID: resolved}, alice.ID); err != nil { + t.Fatalf("relogin Alice on canonical perm: %v", err) + } + if a, found, err := authz.ByAuthKey(ctx, permKey.ID); err != nil || !found || a.UserID != alice.ID { + t.Fatalf("perm authorization = %+v found=%v err=%v, want Alice", a, found, err) + } + if a, found, err := authz.ByAuthKey(ctx, tempKey.ID); err != nil || found { + t.Fatalf("temp authorization = %+v found=%v err=%v, want absent", a, found, err) + } +} + +func TestAuthorizationBindRejectsTemporaryProtocolKey(t *testing.T) { + ctx := context.Background() + users := memory.NewUserStore() + authz := memory.NewAuthorizationStore() + keys := memory.NewAuthKeyStore() + tempKey := testAuthKey(0x82) + if err := keys.Save(ctx, store.AuthKeyData{ + ID: tempKey.ID, ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + }); err != nil { + t.Fatalf("save temp key: %v", err) + } + u, err := users.Create(ctx, domain.User{Phone: "15550008201", FirstName: "Alice"}) + if err != nil { + t.Fatalf("create user: %v", err) + } + svc := NewService(users, authz, memory.NewCodeStore(), keys, memory.NewTempAuthKeyBindingStore(keys), "12345") + if _, err := svc.BindVerifiedLogin(ctx, domain.Authorization{AuthKeyID: tempKey.ID}, u.ID); !errors.Is(err, ErrAuthKeyPermEmpty) { + t.Fatalf("bind temp authorization err = %v, want ErrAuthKeyPermEmpty", err) + } + if _, found, err := authz.ByAuthKey(ctx, tempKey.ID); err != nil || found { + t.Fatalf("temp authorization found=%v err=%v, want absent", found, err) } } @@ -666,10 +821,14 @@ func testAuthKey(seed byte) mtcrypto.AuthKey { } func saveAuthKey(t *testing.T, keys store.AuthKeyStore, key mtcrypto.AuthKey) { + saveAuthKeyWithExpiry(t, keys, key, 0) +} + +func saveAuthKeyWithExpiry(t *testing.T, keys store.AuthKeyStore, key mtcrypto.AuthKey, expiresAt int) { t.Helper() var value [256]byte copy(value[:], key.Value[:]) - if err := keys.Save(context.Background(), store.AuthKeyData{ID: key.ID, Value: value}); err != nil { + if err := keys.Save(context.Background(), store.AuthKeyData{ID: key.ID, Value: value, ExpiresAt: expiresAt}); err != nil { t.Fatalf("save auth key: %v", err) } } diff --git a/internal/app/maintenance/retention.go b/internal/app/maintenance/retention.go index 78736c10..a45a1f99 100644 --- a/internal/app/maintenance/retention.go +++ b/internal/app/maintenance/retention.go @@ -12,7 +12,7 @@ type DispatchOutboxRetentionStore interface { DeleteFailed(ctx context.Context, olderThan time.Duration, limit int) (int, error) } -// TempAuthKeyRetentionStore 回收过期的 PFS temp auth key 绑定。 +// TempAuthKeyRetentionStore 回收过期的 PFS temp auth key（含未绑定 key）。 type TempAuthKeyRetentionStore interface { DeleteExpired(ctx context.Context, expiredBefore int64, limit int) (int, error) } @@ -64,9 +64,9 @@ type LoginCodeDeliveryRetentionStore interface { // getUpdates 读取（fromID 恒 > confirmed），宽限仅防御 offset 回拨调试；回收目标是清堆积。 const botAPIConfirmedGrace = 15 * time.Minute -// tempAuthKeyExpiryGrace 是 temp key 过期后的回收宽限：ResolveAuthKey 对 -// 「已过期但 perm 已授权」的绑定是容忍的，立即删除会突然断掉这批宽限中的 -// 连接；回收目标是清堆积，晚一天无妨。 +// tempAuthKeyExpiryGrace 只是一段数据库物理回收宽限。MTProto edge 在 expires_at +// 到点即停止入站 RPC、主动推送和重发，并断开连接；ResolveAuthKey 不容忍过期 key。 +// 晚一天删除用于吸收客户端轮换/诊断窗口，不会延长协议有效期。 const tempAuthKeyExpiryGrace = 24 * time.Hour const ( diff --git a/internal/mtprotoedge/auth_key_expiry_test.go b/internal/mtprotoedge/auth_key_expiry_test.go new file mode 100644 index 00000000..0973b685 --- /dev/null +++ b/internal/mtprotoedge/auth_key_expiry_test.go @@ -0,0 +1,447 @@ +package mtprotoedge + +import ( + "context" + "crypto/rand" + "encoding/binary" + "errors" + "sync" + "testing" + "time" + + "go.uber.org/zap/zaptest" + + "github.com/gotd/log/logzap" + "github.com/gotd/td/bin" + "github.com/gotd/td/clock" + "github.com/gotd/td/crypto" + "github.com/gotd/td/exchange" + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" + "github.com/gotd/td/proto/codec" + "github.com/gotd/td/tg" + "github.com/gotd/td/transport" +) + +func TestAuthKeyProtocolUnavailable(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + tests := []struct { + name string + expiresAt int + want bool + }{ + {name: "legacy unknown", expiresAt: -1, want: true}, + {name: "permanent", expiresAt: 0, want: false}, + {name: "expired temporary", expiresAt: int(now.Unix()), want: true}, + {name: "live temporary", expiresAt: int(now.Add(time.Second).Unix()), want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := authKeyProtocolUnavailable(tt.expiresAt, now); got != tt.want { + t.Fatalf("authKeyProtocolUnavailable(%d) = %v, want %v", tt.expiresAt, got, tt.want) + } + }) + } +} + +// expiryTestClock keeps server protocol time deterministic while retaining real +// timers for transport/RPC deadlines. Expiry admission reads Now before any +// envelope validation, so advancing it exercises the cached active-connection +// boundary without making the test sleep until a wall-clock second rolls over. +type expiryTestClock struct { + mu sync.RWMutex + now time.Time +} + +func newExpiryTestClock(now time.Time) *expiryTestClock { + return &expiryTestClock{now: now} +} + +func (c *expiryTestClock) Now() time.Time { + c.mu.RLock() + defer c.mu.RUnlock() + return c.now +} + +func (c *expiryTestClock) Advance(d time.Duration) { + c.mu.Lock() + c.now = c.now.Add(d) + c.mu.Unlock() +} + +func (*expiryTestClock) Timer(d time.Duration) clock.Timer { return clock.System.Timer(d) } +func (*expiryTestClock) Ticker(d time.Duration) clock.Ticker { return clock.System.Ticker(d) } + +type signalingGuardedLeaseWriter struct { + lease *physicalTransportLease + entered chan struct{} + once sync.Once +} + +func (w *signalingGuardedLeaseWriter) Send(ctx context.Context, b *bin.Buffer) error { + return w.lease.Send(ctx, b) +} + +func (w *signalingGuardedLeaseWriter) SendDeadlineWithScratchGuarded(deadline time.Time, b *bin.Buffer, scratch *[]byte, guard func() error) error { + w.once.Do(func() { close(w.entered) }) + return w.lease.SendDeadlineWithScratchGuarded(deadline, b, scratch, guard) +} + +func dialTemporaryHandshakeForExpiryTest( + t *testing.T, + addr string, + dc, expiresIn int, + pub exchange.PublicKey, +) (transport.Conn, exchange.ClientExchangeResult, crypto.Cipher) { + t.Helper() + conn := dialTransportOnly(t, addr) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + auth, err := exchange.NewExchanger(conn, dc). + WithTempMode(expiresIn). + WithRand(rand.Reader). + WithLogger(logzap.New(zaptest.NewLogger(t).Named("temp-client"))). + Client([]exchange.PublicKey{pub}). + Run(ctx) + if err != nil { + t.Fatalf("temporary client exchange: %v", err) + } + return conn, auth, crypto.NewClientCipher(rand.Reader) +} + +func TestActiveTemporaryAuthKeyExpiresBeforeNextRPCDispatch(t *testing.T) { + const ( + dc = 2 + expiresIn = 60 * 60 + ) + now := time.Now() + testClock := newExpiryTestClock(now) + handler := &admissionCountingRPC{} + addr, pub, srv := startTestServer(t, Options{ + DC: dc, + Clock: testClock, + RPC: handler, + }) + conn, auth, cipher := dialTemporaryHandshakeForExpiryTest(t, addr, dc, expiresIn, pub) + + stored, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID) + if err != nil || !found { + t.Fatalf("temporary auth key after exchange: found=%v err=%v", found, err) + } + wantExpiresAt := int(now.Unix()) + expiresIn + if stored.ExpiresAt != wantExpiresAt { + t.Fatalf("temporary auth key expires_at = %d, want %d", stored.ExpiresAt, wantExpiresAt) + } + + ids := proto.NewMessageIDGen(time.Now) + firstID := ids.New(proto.MessageFromClient) + sendEncrypted(t, conn, cipher, auth, firstID, &tg.HelpGetConfigRequest{}) + collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + proto.ResultTypeID: 1, + mt.MsgsAckTypeID: 1, + }) + waitForAtomicCalls(t, &handler.calls, 1) + + key := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID} + srv.conns.mu.RLock() + active := srv.conns.bySession[key] + srv.conns.mu.RUnlock() + if active == nil || !active.isActive() { + t.Fatalf("temporary session was not active before expiry: %p", active) + } + + // Cross the exact protocol boundary: expires_at <= now is invalid. The next + // frame must be rejected before decrypt/preflight/Dispatch, even though this + // connection already cached the key and completed session activation. + testClock.Advance(time.Duration(expiresIn+1) * time.Second) + sendEncrypted(t, conn, cipher, auth, ids.New(proto.MessageFromClient), &tg.HelpGetConfigRequest{}) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + var response bin.Buffer + err = conn.Recv(ctx, &response) + var protocolErr *codec.ProtocolErr + if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound { + t.Fatalf("expired active temp key recv = %T %v, want protocol -404", err, err) + } + + waitForManagedSessionAbsent(t, srv.conns, key) + if got := handler.calls.Load(); got != 1 { + t.Fatalf("expired active temp key executed %d RPCs, want only the pre-expiry request", got) + } +} + +func TestExpiredTemporaryAuthKeyRejectsServerPushWithoutWireWrite(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + clock := newExpiryTestClock(now) + tr := &failAfterTransport{} + c := newOutboundTestConn(t, tr, nil) + c.now = clock.Now + c.authKeyExpiresAt = int(now.Unix()) + + err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, + &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}, 0) + if !errors.Is(err, ErrConnClosed) { + t.Fatalf("push on expired temp key = %v, want ErrConnClosed", err) + } + if got := tr.sends.Load(); got != 0 { + t.Fatalf("wire sends after expiry = %d, want zero", got) + } + if !c.isRetired() || tr.closes.Load() != 1 { + t.Fatalf("expired connection retired=%v transport_closes=%d, want true/1", c.isRetired(), tr.closes.Load()) + } +} + +func TestQueuedPushCannotCrossTemporaryAuthKeyExpiry(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + clock := newExpiryTestClock(now) + tr := newGatedRecordingTransport() + c := newOutboundTestConn(t, tr, nil) + c.now = clock.Now + c.authKeyExpiresAt = int(now.Add(time.Minute).Unix()) + encoded := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}} + + if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, encoded, 0); err != nil { + t.Fatalf("enqueue first push: %v", err) + } + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("first push did not enter blocked writer") + } + if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, encoded, 0); err != nil { + t.Fatalf("enqueue second push: %v", err) + } + clock.Advance(time.Minute) + tr.once.Do(func() { close(tr.release) }) + + select { + case <-c.outboundDone: + case <-time.After(time.Second): + t.Fatal("expired outbound actor did not stop") + } + if got := len(tr.snapshot()); got != 1 { + t.Fatalf("wire frames across expiry = %d, want only already-writing frame", got) + } +} + +func TestTemporaryAuthKeyExpiryWhileWaitingForPhysicalWriterSkipsRawSend(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + testClock := newExpiryTestClock(now) + raw := newGatedRecordingTransport() + _, lease := newPhysicalTransportOwner(raw) + c := newOutboundTestConn(t, lease, nil) + c.transportLease = lease + c.now = testClock.Now + c.authKeyExpiresAt = int(now.Add(time.Minute).Unix()) + signaling := &signalingGuardedLeaseWriter{lease: lease, entered: make(chan struct{})} + c.writer = signaling + + // Simulate a quick ACK/protocol write that already owns the physical writer. + quickDone := make(chan error, 1) + go func() { + quickDone <- lease.Send(context.Background(), &bin.Buffer{Buf: []byte{1, 2, 3, 4}}) + }() + select { + case <-raw.started: + case <-time.After(time.Second): + t.Fatal("direct protocol write did not acquire physical writer") + } + + encoded := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}} + actorDone := make(chan error, 1) + go func() { + actorDone <- c.SendEncoded(context.Background(), proto.MessageFromServer, encoded) + }() + select { + case <-signaling.entered: + // writeFrame passed its outer expiry check and entered the guarded lease; + // the direct write still owns writeMu, so raw.Send cannot have started. + case <-time.After(time.Second): + t.Fatal("outbound actor did not wait for physical writer ownership") + } + + testClock.Advance(time.Minute) + raw.once.Do(func() { close(raw.release) }) + if err := <-quickDone; err != nil { + t.Fatalf("direct protocol write: %v", err) + } + if err := <-actorDone; !errors.Is(err, ErrConnClosed) { + t.Fatalf("actor write after expiry = %v, want ErrConnClosed", err) + } + if frames := raw.snapshot(); len(frames) != 1 { + t.Fatalf("raw wire frames = %d, want only the pre-expiry direct frame", len(frames)) + } + if !c.isRetired() { + t.Fatal("connection was not fenced after guarded expiry rejection") + } +} + +func TestRetiredActorWaitingForPhysicalWriterDoesNotDefeatLeaseTransfer(t *testing.T) { + raw := newGatedRecordingTransport() + _, lease := newPhysicalTransportOwner(raw) + c := newOutboundTestConn(t, lease, nil) + c.transportLease = lease + signaling := &signalingGuardedLeaseWriter{lease: lease, entered: make(chan struct{})} + c.writer = signaling + + directDone := make(chan error, 1) + go func() { + directDone <- lease.Send(context.Background(), &bin.Buffer{Buf: []byte{5, 6, 7, 8}}) + }() + select { + case <-raw.started: + case <-time.After(time.Second): + t.Fatal("direct protocol write did not acquire physical writer") + } + + actorDone := make(chan error, 1) + go func() { + actorDone <- c.SendEncoded(context.Background(), proto.MessageFromServer, + &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}) + }() + select { + case <-signaling.entered: + case <-time.After(time.Second): + t.Fatal("outbound actor did not reach guarded physical writer") + } + + c.beginTerminalShutdown() + raw.once.Do(func() { close(raw.release) }) + if err := <-directDone; err != nil { + t.Fatalf("direct protocol write: %v", err) + } + select { + case <-c.outboundDone: + case <-time.After(time.Second): + t.Fatal("retired outbound actor did not drain") + } + select { + case err := <-actorDone: + if !errors.Is(err, ErrConnClosed) { + t.Fatalf("retired actor write = %v, want ErrConnClosed", err) + } + default: + } + if frames := raw.snapshot(); len(frames) != 1 { + t.Fatalf("retired actor reached raw writer: frames=%d, want one direct frame", len(frames)) + } + if !lease.IsCurrentOpen() { + t.Fatal("retired actor closed physical lease") + } + if next, ok := lease.Transfer(); !ok || next == nil { + t.Fatal("retired actor defeated physical lease transfer") + } +} + +func TestTerminalAuthKeyNotFoundSurvivesActorWaitingForPhysicalWriter(t *testing.T) { + raw := newGatedRecordingTransport() + _, lease := newPhysicalTransportOwner(raw) + c := newOutboundTestConn(t, lease, nil) + c.transportLease = lease + signaling := &signalingGuardedLeaseWriter{lease: lease, entered: make(chan struct{})} + c.writer = signaling + + directDone := make(chan error, 1) + go func() { + directDone <- lease.Send(context.Background(), &bin.Buffer{Buf: []byte{9, 10, 11, 12}}) + }() + select { + case <-raw.started: + case <-time.After(time.Second): + t.Fatal("direct protocol write did not acquire physical writer") + } + go func() { + _ = c.SendEncoded(context.Background(), proto.MessageFromServer, + &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}) + }() + select { + case <-signaling.entered: + case <-time.After(time.Second): + t.Fatal("outbound actor did not reach guarded physical writer") + } + + srv := New(Options{WriteTimeout: time.Second}) + terminalDone := make(chan error, 1) + go func() { + terminalDone <- srv.sendTerminalProtoError(context.Background(), c, codec.CodeAuthKeyNotFound) + }() + select { + case err := <-terminalDone: + t.Fatalf("terminal error bypassed waiting actor: %v", err) + case <-time.After(50 * time.Millisecond): + } + raw.once.Do(func() { close(raw.release) }) + if err := <-directDone; err != nil { + t.Fatalf("direct protocol write: %v", err) + } + select { + case err := <-terminalDone: + if err != nil { + t.Fatalf("send terminal -404: %v", err) + } + case <-time.After(time.Second): + t.Fatal("terminal -404 did not follow waiting actor drain") + } + + frames := raw.snapshot() + if len(frames) != 2 { + t.Fatalf("wire frames = %d, want direct frame then -404", len(frames)) + } + last := frames[len(frames)-1] + if len(last) != 4 || int32(binary.LittleEndian.Uint32(last)) != -codec.CodeAuthKeyNotFound { + t.Fatalf("last wire frame = %x, want bare -404", last) + } +} + +func TestTerminalAuthKeyNotFoundWaitsForOutboundAndIsLastFrame(t *testing.T) { + now := time.Unix(1_800_000_000, 0) + clock := newExpiryTestClock(now) + tr := newGatedRecordingTransport() + _, lease := newPhysicalTransportOwner(tr) + c := newOutboundTestConn(t, lease, nil) + c.transportLease = lease + c.now = clock.Now + c.authKeyExpiresAt = int(now.Add(time.Minute).Unix()) + encoded := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}} + if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, encoded, 0); err != nil { + t.Fatalf("enqueue blocked push: %v", err) + } + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("push did not enter blocked writer") + } + + clock.Advance(time.Minute) + srv := New(Options{WriteTimeout: time.Second}) + terminalDone := make(chan error, 1) + go func() { + terminalDone <- srv.sendTerminalProtoError(context.Background(), c, codec.CodeAuthKeyNotFound) + }() + select { + case err := <-terminalDone: + t.Fatalf("terminal error bypassed active outbound writer: %v", err) + case <-time.After(50 * time.Millisecond): + } + if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, encoded, 0); !errors.Is(err, ErrConnClosed) { + t.Fatalf("push admitted behind terminal fence: %v", err) + } + tr.once.Do(func() { close(tr.release) }) + select { + case err := <-terminalDone: + if err != nil { + t.Fatalf("send terminal -404: %v", err) + } + case <-time.After(time.Second): + t.Fatal("terminal -404 did not follow drained writer") + } + + frames := tr.snapshot() + if len(frames) != 2 { + t.Fatalf("wire frames = %d, want encrypted frame then -404", len(frames)) + } + last := frames[len(frames)-1] + if len(last) != 4 || int32(binary.LittleEndian.Uint32(last)) != -codec.CodeAuthKeyNotFound { + t.Fatalf("last wire frame = %x, want bare -404", last) + } +} diff --git a/internal/mtprotoedge/bot_callback_e2e_test.go b/internal/mtprotoedge/bot_callback_e2e_test.go index b586161d..12f07d53 100644 --- a/internal/mtprotoedge/bot_callback_e2e_test.go +++ b/internal/mtprotoedge/bot_callback_e2e_test.go @@ -69,7 +69,7 @@ func newBotCallbackEnv(t *testing.T, ctx context.Context) *botCallbackEnv { activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions")) deps := rpc.Deps{ Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, - memory.NewTempAuthKeyBindingStore(), "12345", auth.WithBotLogin(botStore)), + memory.NewTempAuthKeyBindingStore(authKeyStore), "12345", auth.WithBotLogin(botStore)), Account: account.NewService(memory.NewPasswordStore()), Help: help.NewService(helpStore, helpStore), Users: users.NewService(userStore), diff --git a/internal/mtprotoedge/bot_e2e_test.go b/internal/mtprotoedge/bot_e2e_test.go index 8ebe5c44..5fbc0021 100644 --- a/internal/mtprotoedge/bot_e2e_test.go +++ b/internal/mtprotoedge/bot_e2e_test.go @@ -72,7 +72,7 @@ func TestBotManagementRPCFlow(t *testing.T) { activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions")) deps := rpc.Deps{ Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, - memory.NewTempAuthKeyBindingStore(), code, auth.WithBotLogin(botStore)), + memory.NewTempAuthKeyBindingStore(authKeyStore), code, auth.WithBotLogin(botStore)), Account: account.NewService(memory.NewPasswordStore()), Help: help.NewService(helpStore, helpStore), Users: users.NewService(userStore), @@ -308,7 +308,7 @@ func TestBotFatherCreateAndBotLoginFlow(t *testing.T) { activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions")) deps := rpc.Deps{ Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, - memory.NewTempAuthKeyBindingStore(), code, auth.WithBotLogin(botStore)), + memory.NewTempAuthKeyBindingStore(authKeyStore), code, auth.WithBotLogin(botStore)), Account: account.NewService(memory.NewPasswordStore()), Help: help.NewService(helpStore, helpStore), Users: users.NewService(userStore), diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index cc9eddca..946425b7 100644 --- a/internal/mtprotoedge/conn.go +++ b/internal/mtprotoedge/conn.go @@ -50,14 +50,21 @@ type Conn struct { msgID *proto.MessageIDGen writeTimeout time.Duration metrics Metrics + // now shares the Server protocol clock with inbound expiry admission. Tests may + // advance it without sleeping; construction-only Conns fall back to time.Now. + now func() time.Time authKeyID [8]byte // authKeyHex 是 authKeyID 的 hex 缓存：每条 RPC 的结构化日志都会带它， // 建连时算一次，避免热路径反复 hex 编码分配。 authKeyHex string - sessionID int64 - salt int64 - key crypto.AuthKey + // authKeyExpiresAt=0 表示 permanent key；正值是 temporary/media-temporary + // key 在握手时确定的绝对协议失效时间；-1 是仅供迁移的 legacy-unknown + // sentinel（edge 会在创建 Conn 前以 -404 拒绝）。Conn 创建后不可变。 + authKeyExpiresAt int + sessionID int64 + salt int64 + key crypto.AuthKey outbound chan outboundOp outboundControl chan outboundOp @@ -243,6 +250,20 @@ func (c *Conn) SetClientLayer(layer int) { c.clientLayer.Store(int32(layer)) } // AuthKeyID 返回连接的 auth_key_id。 func (c *Conn) AuthKeyID() [8]byte { return c.authKeyID } +// AuthKeyExpiresAt 返回 raw 协议 key 的失效时间；0 表示 permanent key。 +func (c *Conn) AuthKeyExpiresAt() int { return c.authKeyExpiresAt } + +func (c *Conn) authKeyProtocolUnavailableNow() bool { + if c == nil { + return true + } + now := time.Now() + if c.now != nil { + now = c.now() + } + return authKeyProtocolUnavailable(c.authKeyExpiresAt, now) +} + // BusinessAuthKeyID 返回业务视角的 auth_key_id。 // // temp auth_key 绑定后解析为 perm auth_key；第二个返回值表示本连接是否已完成解析， diff --git a/internal/mtprotoedge/e2e_test.go b/internal/mtprotoedge/e2e_test.go index 56062e98..89630939 100644 --- a/internal/mtprotoedge/e2e_test.go +++ b/internal/mtprotoedge/e2e_test.go @@ -46,7 +46,7 @@ func TestTelegramClientEndToEnd(t *testing.T) { authzStore := memory.NewAuthorizationStore() authKeyStore := memory.NewAuthKeyStore() deps := rpc.Deps{ - Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), "12345"), + Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), "12345"), Users: users.NewService(userStore), Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()), } diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index c3e05fe5..8e252af9 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -117,13 +117,16 @@ var errActivationAuthKeyRejected = errors.New("activation auth key no longer exi func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b, plain *bin.Buffer) (*Conn, error) { var key crypto.AuthKey var serverSalt int64 + var authKeyExpiresAt int if fetchedKey != nil { key = crypto.AuthKey{Value: crypto.Key(fetchedKey.Value), ID: fetchedKey.ID} serverSalt = fetchedKey.ServerSalt + authKeyExpiresAt = fetchedKey.ExpiresAt } else { // 快路径：复用已建立连接缓存的密钥与盐（同一 auth key 的后续帧，含同连接换 session）。 key = current.key serverSalt = current.salt + authKeyExpiresAt = current.authKeyExpiresAt } frame, err := decryptClientFrame(key, b, plain) @@ -152,6 +155,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con } else { current = s.newConn(tc, key, frame.sessionID, serverSalt) } + current.authKeyExpiresAt = authKeyExpiresAt // 注册即播种协商 layer：新 Conn 的 clientLayer 为 0（=canonical 227），若等到 // 首条 RPC 的 Dispatch 返回后才刷新，重连老客户端在首条 RPC handler 执行期间 // 收到的 pending flush / 并发 push 会漏降级。进程内重连时 rpc 层留有 @@ -221,10 +225,10 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con if getErr != nil { return current, fmt.Errorf("revalidate activation auth key: %w", getErr) } - if !found || fresh.ID != current.authKeyID || fresh.Value != [256]byte(current.key.Value) { + if !found || fresh.ID != current.authKeyID || fresh.Value != [256]byte(current.key.Value) || authKeyProtocolUnavailable(fresh.ExpiresAt, s.clock.Now()) { // Send the terminal protocol error while the claim still owns a live writer; // the deferred abort then fences and removes it before serveConn returns. - if sendErr := s.sendProtoError(ctx, current.transport, codec.CodeAuthKeyNotFound); sendErr != nil { + if sendErr := s.sendTerminalProtoError(ctx, current, codec.CodeAuthKeyNotFound); sendErr != nil { return current, sendErr } return current, errActivationAuthKeyRejected diff --git a/internal/mtprotoedge/exchange.go b/internal/mtprotoedge/exchange.go index 04221f87..5acba16c 100644 --- a/internal/mtprotoedge/exchange.go +++ b/internal/mtprotoedge/exchange.go @@ -98,12 +98,13 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first } // authKeyData 把握手结果转换为 store 记录。 -func authKeyData(key crypto.AuthKey, salt, createdAt int64) store.AuthKeyData { +func authKeyData(key crypto.AuthKey, salt, createdAt int64, expiresAt int) store.AuthKeyData { return store.AuthKeyData{ ID: key.ID, Value: [256]byte(key.Value), ServerSalt: salt, CreatedAt: createdAt, + ExpiresAt: expiresAt, } } @@ -120,6 +121,24 @@ func (s *Server) sendProtoError(ctx context.Context, conn transport.Conn, code i return nil } +// sendTerminalProtoError serializes a bare transport error after the authenticated +// outbound actor has stopped. A direct write while the actor is still draining can +// otherwise interleave an encrypted update/result after -404 on the same socket. +func (s *Server) sendTerminalProtoError(ctx context.Context, c *Conn, code int32) error { + if c == nil { + return errors.New("send terminal protocol error without logical connection") + } + c.beginTerminalShutdown() + if !c.waitOutboundShutdownUntil(forceCloseBatchTimeout) { + c.closeTransport() + return errors.New("outbound writer did not stop before terminal protocol error") + } + if c.transport == nil { + return ErrConnClosed + } + return s.sendProtoError(ctx, c.transport, code) +} + // maxHandshakeReqPQ 是一次密钥交换内允许的 req_pq(_multi) 帧数上界。正常握手只发 1 个 // req_pq（含个别客户端的「fake+真」也就 2 个）；客户端因 nonce 失步陷入「收到 ResPQ→立刻 // 重启握手换 nonce 重发 req_pq」死循环时，会在同一连接上无限发 req_pq，而委托给 gotd 的 diff --git a/internal/mtprotoedge/exchange_compat.go b/internal/mtprotoedge/exchange_compat.go index cb0c26fe..1e38ead6 100644 --- a/internal/mtprotoedge/exchange_compat.go +++ b/internal/mtprotoedge/exchange_compat.go @@ -50,9 +50,9 @@ func (s *Server) runServerExchange(ctx context.Context, conn transport.Conn) (ex // client is allowed to immediately use the new key, possibly on another TCP // connection. Persisting after the response creates a split-brain window when // storage fails or the process exits between those two operations. -func (s *Server) commitExchangeAuthKey(ctx context.Context, result exchange.ServerExchangeResult) error { +func (s *Server) commitExchangeAuthKey(ctx context.Context, result exchange.ServerExchangeResult, expiresAt int) error { createdAt := s.clock.Now().Unix() - if err := s.authKeys.Save(ctx, authKeyData(result.Key, result.ServerSalt, createdAt)); err != nil { + if err := s.authKeys.Save(ctx, authKeyData(result.Key, result.ServerSalt, createdAt, expiresAt)); err != nil { return fmt.Errorf("persist auth key before DhGenOk: %w", err) } return nil @@ -67,7 +67,7 @@ type serverExchangeCompat struct { dc int log *zap.Logger rng compatServerRNG - commitKey func(context.Context, exchange.ServerExchangeResult) error + commitKey func(context.Context, exchange.ServerExchangeResult, int) error } const pqInnerDataTempTypeID uint32 = 0x3c6a84d4 @@ -129,7 +129,10 @@ SendResPQ: s.log.Debug("Received client ReqDHParamsRequest") } - var innerData mt.PQInnerData + var ( + innerData mt.PQInnerData + authKeyExpiresAt int + ) { if dhParams.DH.Nonce != req.Nonce { return exchange.ServerExchangeResult{}, gofaster.New("req_DH_params nonce does not match req_pq") @@ -161,6 +164,15 @@ SendResPQ: } innerData = d.Data + if d.Temp { + expiresAt := s.clock.Now().Unix() + int64(d.ExpiresIn) + // TL timestamps are signed int32 on the wire. Reject an impossible + // lifetime instead of wrapping a temporary key into a permanent one. + if expiresAt <= 0 || expiresAt > int64(^uint32(0)>>1) { + return exchange.ServerExchangeResult{}, gofaster.New("temporary auth key expiry is out of int32 range") + } + authKeyExpiresAt = int(expiresAt) + } } dhPrime, err := s.rng.DhPrime() @@ -251,7 +263,7 @@ SendResPQ: if s.commitKey == nil { return exchange.ServerExchangeResult{}, gofaster.New("auth key commit hook is required before DhGenOk") } - if err := s.commitKey(ctx, serverResult); err != nil { + if err := s.commitKey(ctx, serverResult, authKeyExpiresAt); err != nil { return exchange.ServerExchangeResult{}, err } diff --git a/internal/mtprotoedge/exchange_test.go b/internal/mtprotoedge/exchange_test.go index a6fc7d4d..57829904 100644 --- a/internal/mtprotoedge/exchange_test.go +++ b/internal/mtprotoedge/exchange_test.go @@ -318,15 +318,19 @@ func TestKeyExchangeAuthKeySaveUsesHandshakeDeadline(t *testing.T) { } func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) { - const dc = 2 + const ( + dc = 2 + expiresIn = 24 * 60 * 60 + ) addr, pub, srv := startTestServer(t, Options{DC: dc}) conn := dialTransportOnly(t, addr) t.Cleanup(func() { _ = conn.Close() }) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() + startedAt := time.Now() res, err := exchange.NewExchanger(conn, -dc). - WithTempMode(24 * 60 * 60). + WithTempMode(expiresIn). WithRand(rand.Reader). WithLogger(logzap.New(zaptest.NewLogger(t).Named("client"))). Client([]exchange.PublicKey{pub}). @@ -334,6 +338,7 @@ func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) { if err != nil { t.Fatalf("client exchange: %v", err) } + completedAt := time.Now() var saved store.AuthKeyData found := false @@ -354,6 +359,11 @@ func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) { if saved.ServerSalt != res.ServerSalt { t.Fatalf("server salt mismatch: server=%d client=%d", saved.ServerSalt, res.ServerSalt) } + minExpiresAt := int(startedAt.Unix()) + expiresIn + maxExpiresAt := int(completedAt.Unix()) + expiresIn + if saved.ExpiresAt < minExpiresAt || saved.ExpiresAt > maxExpiresAt { + t.Fatalf("server temp auth key expires_at = %d, want absolute unix time in [%d, %d]", saved.ExpiresAt, minExpiresAt, maxExpiresAt) + } } func TestKeyExchangeRejectsWrongNegativeTempDC(t *testing.T) { diff --git a/internal/mtprotoedge/login_e2e_test.go b/internal/mtprotoedge/login_e2e_test.go index 27d54486..2af53dcb 100644 --- a/internal/mtprotoedge/login_e2e_test.go +++ b/internal/mtprotoedge/login_e2e_test.go @@ -83,7 +83,7 @@ func TestLoginRegisterFlow(t *testing.T) { t.Fatalf("seed langpack: %v", err) } deps := rpc.Deps{ - Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code), + Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), code), Account: account.NewService(memory.NewPasswordStore()), Help: help.NewService(helpStore, helpStore), Users: users.NewService(userStore), @@ -302,7 +302,7 @@ func TestPrivateMessageRoundTripFlow(t *testing.T) { messageStore := memory.NewMessageStore(dialogStore) activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions")) deps := rpc.Deps{ - Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code), + Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), code), Account: account.NewService(memory.NewPasswordStore()), Help: help.NewService(helpStore, helpStore), Users: users.NewService(userStore), diff --git a/internal/mtprotoedge/login_email_e2e_test.go b/internal/mtprotoedge/login_email_e2e_test.go index e59003cc..924044e7 100644 --- a/internal/mtprotoedge/login_email_e2e_test.go +++ b/internal/mtprotoedge/login_email_e2e_test.go @@ -78,7 +78,7 @@ func TestLoginEmailEndToEnd(t *testing.T) { accountService := account.NewService(passwordStore, account.WithUsers(userStore), account.WithLoginEmailVerification(codeStore, emailSender, 5*time.Minute, 5, 6)) - authService := auth.NewService(userStore, authzStore, codeStore, authKeyStore, memory.NewTempAuthKeyBindingStore(), code, + authService := auth.NewService(userStore, authzStore, codeStore, authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), code, auth.WithLoginMessages(messageStore, dialogStore), auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore)), auth.WithPasswords(passwordStore), diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index 7f397742..d237ce5f 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -567,6 +567,18 @@ func (c *Conn) failTransport() { c.closeTransport() } +// fenceUnavailableAuthKey turns protocol expiry into a connection-level terminal +// boundary. Outbound producers may discover expiry before the read loop sees the +// client's next frame; in that case the socket is closed so the client reconnects +// and receives the ordinary -404 admission response for the stale raw key. +func (c *Conn) fenceUnavailableAuthKey() { + if c == nil || !c.authKeyProtocolUnavailableNow() { + return + } + c.beginTerminalShutdown() + c.closeTransport() +} + // fenceUndeliveredRPCResult is the no-reentry terminal path used from a task's // release callback. That callback may itself run while rpcClose.Do is draining // queued tasks, so calling beginCloseInboundRPCScheduler again would deadlock on @@ -1022,6 +1034,13 @@ func (c *Conn) outboundQueue(op outboundOp) chan outboundOp { } func (c *Conn) beginOutboundEnqueue() bool { + // A temporary key is unusable for both inbound RPCs and server-originated + // updates at the same absolute boundary. Reject before encoding admission; + // the actor and write path repeat this check to close the two race windows. + if c.authKeyProtocolUnavailableNow() { + c.fenceUnavailableAuthKey() + return false + } c.outboundEnqueueMu.Lock() defer c.outboundEnqueueMu.Unlock() if c.outboundClosing || c.isRetired() { @@ -1146,6 +1165,14 @@ func (c *Conn) drainOutbound() { } func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) { + // An operation can sit in a bounded queue across the protocol expiry instant. + // It must be failed and released without touching the wire. + if c.authKeyProtocolUnavailableNow() { + op.releaseReservation(state.budget) + op.finish(outboundResult{err: ErrConnClosed}) + c.fenceUnavailableAuthKey() + return + } if op.kind != outboundSend { defer op.releaseReservation(state.budget) } @@ -1570,7 +1597,20 @@ type deadlineOutboundScratchWriter interface { SendDeadlineWithScratch(deadline time.Time, b *bin.Buffer, scratch *[]byte) error } +type deadlineOutboundGuardedScratchWriter interface { + SendDeadlineWithScratchGuarded(deadline time.Time, b *bin.Buffer, scratch *[]byte, guard func() error) error +} + +var ( + errAuthKeyUnavailableAtPhysicalWrite = errors.New("auth key unavailable at physical write admission") + errConnRetiredAtPhysicalWrite = errors.New("connection retired at physical write admission") +) + func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { + if c.authKeyProtocolUnavailableNow() { + c.fenceUnavailableAuthKey() + return ErrConnClosed + } if ctx == nil { ctx = context.Background() } @@ -1595,12 +1635,28 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { if err := prewriteDeadlineError(ctx, deadline); err != nil { return fmt.Errorf("outbound deadline before write: %w", err) } + // Scratch admission and encryption may straddle expires_at. Recheck at the + // final pre-write barrier so neither fresh sends nor resends use a stale key. + if c.authKeyProtocolUnavailableNow() { + c.fenceUnavailableAuthKey() + return ErrConnClosed + } writer := c.writer if writer == nil { writer = c.transport } - if sw, ok := writer.(deadlineOutboundScratchWriter); ok { + if guarded, ok := writer.(deadlineOutboundGuardedScratchWriter); ok { + err = guarded.SendDeadlineWithScratchGuarded(deadline, out, &scratch.codec, func() error { + if c.isRetired() { + return errConnRetiredAtPhysicalWrite + } + if c.authKeyProtocolUnavailableNow() { + return errAuthKeyUnavailableAtPhysicalWrite + } + return nil + }) + } else if sw, ok := writer.(deadlineOutboundScratchWriter); ok { err = sw.SendDeadlineWithScratch(deadline, out, &scratch.codec) } else if dw, ok := writer.(deadlineOutboundWriter); ok { err = dw.SendDeadline(deadline, out) @@ -1614,6 +1670,18 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { err = writer.Send(sendCtx, out) cancel() } + if errors.Is(err, errAuthKeyUnavailableAtPhysicalWrite) { + // The guarded lease has already released physical write ownership and no + // raw bytes were emitted. Fence outside writeMu so Close cannot deadlock. + c.fenceUnavailableAuthKey() + return ErrConnClosed + } + if errors.Is(err, errConnRetiredAtPhysicalWrite) { + // Terminal shutdown already owns lifecycle/transport close. Do not call + // failTransport here: sendTerminalProtoError is waiting for this actor to + // drain and must retain the lease long enough to write the final bare -404. + return ErrConnClosed + } if err != nil { // 任一 partial write / timeout 都可能破坏 MTProto 帧边界；该 socket // 不可继续复用。这里只发 terminal 信号，不在 actor 内等待自身退出。 diff --git a/internal/mtprotoedge/passkey_e2e_test.go b/internal/mtprotoedge/passkey_e2e_test.go index 006d7e67..6c721d3b 100644 --- a/internal/mtprotoedge/passkey_e2e_test.go +++ b/internal/mtprotoedge/passkey_e2e_test.go @@ -102,7 +102,7 @@ func TestPasskeyEndToEnd(t *testing.T) { passkeyService := passkeyapp.NewService(memory.NewPasskeyStore(), memory.NewPasskeyChallengeStore(), rpID, dc) deps := rpc.Deps{ - Auth: auth.NewService(userStore, memory.NewAuthorizationStore(), memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code, + Auth: auth.NewService(userStore, memory.NewAuthorizationStore(), memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), code, auth.WithLoginMessages(messageStore, dialogStore), auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore))), Account: account.NewService(memory.NewPasswordStore(), account.WithUsers(userStore)), diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index 7c057382..35196c74 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -369,6 +369,7 @@ func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key msgID: proto.NewMessageIDGen(s.clock.Now), writeTimeout: s.writeTimeout, metrics: s.metrics, + now: s.clock.Now, authKeyID: key.ID, authKeyHex: hex.EncodeToString(key.ID[:]), sessionID: sessionID, @@ -770,6 +771,19 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn, remote, loca // 撤销由 SessionManager 主动 Close 连接保证失效，不依赖被动的“下一帧 -404”。 // 尚未进入 SessionManager 的 bad-salt provisional 会在 handleEncrypted 建立 activation claim // 后精确复查一次，既把撤销与激活线性化，也不把 salt storm 放大成 PG 写风暴。 + // temporary key 的绝对 expiry 缓存在 Conn 上，逐帧只做内存比较；到期必须在 + // RPC 前返回 -404，让官方客户端仅轮换 temp key。绝不能落到 Router 后退化为 + // raw business identity，再以会触发整账号退出的 401 结束。 + if current != nil && current.authKeyID == authKeyID && authKeyProtocolUnavailable(current.authKeyExpiresAt, s.clock.Now()) { + s.log.Info("Rejecting unavailable temporary or legacy auth key", + zap.String("auth_key_id", current.authKeyHex), + zap.Int("expires_at", current.authKeyExpiresAt), + ) + if err := s.sendTerminalProtoError(ctx, current, codec.CodeAuthKeyNotFound); err != nil { + return err + } + return nil + } var fetchedKey *store.AuthKeyData if current == nil || current.authKeyID != authKeyID { d, found, err := s.authKeys.Get(ctx, authKeyID) @@ -777,17 +791,35 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn, remote, loca return fmt.Errorf("lookup auth key: %w", err) } if !found { - writer := transport.Conn(conn) + var sendErr error if current != nil { - writer = current.transport + sendErr = s.sendTerminalProtoError(ctx, current, codec.CodeAuthKeyNotFound) + } else { + sendErr = s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound) } - if err := s.sendProtoError(ctx, writer, codec.CodeAuthKeyNotFound); err != nil { - return err + if sendErr != nil { + return sendErr } // -404 对 TDesktop 是 terminal key failure；继续保留 socket 只会允许 // 同一客户端反复触发 AuthKeyStore 查询。回包一次后立即断开。 return nil } + if authKeyProtocolUnavailable(d.ExpiresAt, s.clock.Now()) { + s.log.Info("Rejecting unavailable temporary or legacy auth key", + zap.String("auth_key_id", hex.EncodeToString(d.ID[:])), + zap.Int("expires_at", d.ExpiresAt), + ) + var sendErr error + if current != nil { + sendErr = s.sendTerminalProtoError(ctx, current, codec.CodeAuthKeyNotFound) + } else { + sendErr = s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound) + } + if sendErr != nil { + return sendErr + } + return nil + } fetchedKey = &d } @@ -806,6 +838,12 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn, remote, loca } } +func authKeyProtocolUnavailable(expiresAt int, now time.Time) bool { + // -1 is migration 0086's explicit legacy-unknown sentinel. Reject it once + // instead of guessing permanent and allowing account authorization on a temp key. + return expiresAt < 0 || (expiresAt > 0 && int64(expiresAt) <= now.Unix()) +} + // maxRetainedConnBuffer keeps normal upload/download frames allocation-free while preventing one // exceptional near-16MiB transport frame from pinning that capacity for the lifetime of a long // connection. RPC bodies that outlive dispatch already own a budgeted Copy. diff --git a/internal/mtprotoedge/session_manager.go b/internal/mtprotoedge/session_manager.go index 6185271f..fe7feebf 100644 --- a/internal/mtprotoedge/session_manager.go +++ b/internal/mtprotoedge/session_manager.go @@ -443,6 +443,27 @@ func (m *SessionManager) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID i m.bindAuthKeyLocked(c, key, authKeyID) } +// BindAuthKeyForRawAuthKey 把同一 raw temporary key 的全部活跃 session 绑定到 +// canonical permanent key。Android/TDesktop 会在一个 temp key 上并发创建多个 +// session；bind 只发生在其中一个 session，其他 session 不能继续把 raw temp 当业务 key。 +func (m *SessionManager) BindAuthKeyForRawAuthKey(rawAuthKeyID [8]byte, authKeyID [8]byte) int { + m.mu.Lock() + defer m.mu.Unlock() + bound := 0 + for sessionID, c := range m.byAuthKey[rawAuthKeyID] { + if c == nil { + continue + } + key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID} + if m.bySession[key] != c { + continue + } + m.bindAuthKeyLocked(c, key, authKeyID) + bound++ + } + return bound +} + func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8]byte) { oldAuthKeyID, resolved := c.BusinessAuthKeyID() changed := !resolved || oldAuthKeyID != authKeyID @@ -477,6 +498,17 @@ func (m *SessionManager) AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int return c.BusinessAuthKeyID() } +// AuthKeyExpiresAtForSession 返回 raw key 的握手协议失效时间；0 表示 permanent。 +func (m *SessionManager) AuthKeyExpiresAtForSession(rawAuthKeyID [8]byte, sessionID int64) (int, bool) { + m.mu.RLock() + c, ok := m.bySession[sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}] + m.mu.RUnlock() + if !ok { + return 0, false + } + return c.AuthKeyExpiresAt(), true +} + // CloseSessionsForBusinessAuthKey 强制断开指定业务 auth_key 的全部活跃连接， // 供授权撤销（被踢设备）使用：出站推送用连接持有的密钥加密、不回查密钥库， // 不断开的话被撤销的设备会继续收到推送直至自然断线；perm-key 连接的授权 diff --git a/internal/mtprotoedge/session_manager_test.go b/internal/mtprotoedge/session_manager_test.go index 0f3829ba..18f5c583 100644 --- a/internal/mtprotoedge/session_manager_test.go +++ b/internal/mtprotoedge/session_manager_test.go @@ -109,6 +109,40 @@ func TestSessionManagerRegistry(t *testing.T) { } } +func TestBindAuthKeyForRawAuthKeyUpdatesEveryTemporarySession(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + raw := [8]byte{0x71} + perm := [8]byte{0x31} + expiresAt := int(time.Now().Add(time.Hour).Unix()) + c1 := &Conn{sessionID: 101, authKeyID: raw, authKeyExpiresAt: expiresAt} + c2 := &Conn{sessionID: 102, authKeyID: raw, authKeyExpiresAt: expiresAt} + if err := sm.Register(c1); err != nil { + t.Fatalf("register c1: %v", err) + } + if err := sm.Register(c2); err != nil { + t.Fatalf("register c2: %v", err) + } + sm.BindAuthKeyForSession(raw, c1.sessionID, raw) + sm.BindAuthKeyForSession(raw, c2.sessionID, raw) + sm.BindUserForAuthKey(raw, c1.sessionID, 1001) + sm.BindUserForAuthKey(raw, c2.sessionID, 1001) + + if got := sm.BindAuthKeyForRawAuthKey(raw, perm); got != 2 { + t.Fatalf("bound sessions = %d, want 2", got) + } + for _, sessionID := range []int64{c1.sessionID, c2.sessionID} { + if got, ok := sm.AuthKeyIDForSession(raw, sessionID); !ok || got != perm { + t.Fatalf("session %d business key = %x/%v, want perm %x", sessionID, got, ok, perm) + } + if userID, resolved := sm.UserIDResolvedForAuthKey(raw, sessionID); resolved || userID != 0 { + t.Fatalf("session %d user after identity switch = %d/%v, want unresolved", sessionID, userID, resolved) + } + if got, found := sm.AuthKeyExpiresAtForSession(raw, sessionID); !found || got != expiresAt { + t.Fatalf("session %d raw expiry = %d/%v, want %d", sessionID, got, found, expiresAt) + } + } +} + func TestSessionManagerReplacementClosesOldPhysicalTransport(t *testing.T) { sm := NewSessionManager(zaptest.NewLogger(t)) raw := [8]byte{1, 2, 3} diff --git a/internal/mtprotoedge/transport_ownership.go b/internal/mtprotoedge/transport_ownership.go index a99ebabc..90f317b8 100644 --- a/internal/mtprotoedge/transport_ownership.go +++ b/internal/mtprotoedge/transport_ownership.go @@ -164,7 +164,35 @@ func (l *physicalTransportLease) SendDeadlineWithScratch(deadline time.Time, b * }) } +// SendDeadlineWithScratchGuarded evaluates guard while holding the physical +// write-ownership lock, immediately before entering the raw writer. Conn-level +// checks performed before this call are insufficient: a quick ACK or protocol +// write may hold writeMu across a temporary-key expiry boundary. The guard must +// not close/fence the connection itself because that would re-enter transport +// shutdown while writeMu is held; callers handle its error after the lock drops. +func (l *physicalTransportLease) SendDeadlineWithScratchGuarded(deadline time.Time, b *bin.Buffer, scratch *[]byte, guard func() error) error { + return l.withCurrentWriterGuarded(guard, func(raw transport.Conn) error { + if writer, ok := raw.(deadlineOutboundScratchWriter); ok { + return writer.SendDeadlineWithScratch(deadline, b, scratch) + } + if writer, ok := raw.(deadlineOutboundWriter); ok { + return writer.SendDeadline(deadline, b) + } + ctx := context.Background() + cancel := func() {} + if !deadline.IsZero() { + ctx, cancel = context.WithDeadline(ctx, deadline) + } + defer cancel() + return raw.Send(ctx, b) + }) +} + func (l *physicalTransportLease) withCurrentWriter(send func(transport.Conn) error) error { + return l.withCurrentWriterGuarded(nil, send) +} + +func (l *physicalTransportLease) withCurrentWriterGuarded(guard func() error, send func(transport.Conn) error) error { if l == nil || l.owner == nil || l.owner.raw == nil { return ErrConnClosed } @@ -174,6 +202,11 @@ func (l *physicalTransportLease) withCurrentWriter(send func(transport.Conn) err if owner.state.Load() != l.generation { return ErrConnClosed } + if guard != nil { + if err := guard(); err != nil { + return err + } + } return send(owner.raw) } diff --git a/internal/rpc/auth.go b/internal/rpc/auth.go index efd3b971..7977788b 100644 --- a/internal/rpc/auth.go +++ b/internal/rpc/auth.go @@ -77,7 +77,12 @@ func (r *Router) onAuthBindTempAuthKey(ctx context.Context, req *tg.AuthBindTemp r.tempKeyResolveCache.Delete(id) } if r.deps.Sessions != nil { - r.deps.Sessions.BindAuthKeyForSession(id, sessionID, authKeyIDFromInt64(req.PermAuthKeyID)) + permID := authKeyIDFromInt64(req.PermAuthKeyID) + if all, ok := r.deps.Sessions.(RawAuthKeySessionBinder); ok { + all.BindAuthKeyForRawAuthKey(id, permID) + } else { + r.deps.Sessions.BindAuthKeyForSession(id, sessionID, permID) + } } r.invalidateAuthUserCache(id) return true, nil diff --git a/internal/rpc/deps.go b/internal/rpc/deps.go index d97d6ba8..6f4352e8 100644 --- a/internal/rpc/deps.go +++ b/internal/rpc/deps.go @@ -65,6 +65,20 @@ type SessionBinder interface { PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) } +// RawAuthKeySessionBinder 在 auth.bindTempAuthKey 成功后，把同一 raw temporary key +// 已建立的所有 session 一次性切到 canonical permanent identity。只更新当前 session +// 会让并发启动的其它连接永久粘在 raw identity。 +type RawAuthKeySessionBinder interface { + BindAuthKeyForRawAuthKey(rawAuthKeyID [8]byte, authKeyID [8]byte) int +} + +// RawAuthKeyMetadataProvider 暴露握手时确定的 raw key protocol expiry。0 表示 +// permanent key；正值表示 temporary key。Router 只用它判断 cached==raw 是否可作为 +// permanent 快路径，协议过期的实际拒绝仍由 mtprotoedge 完成。 +type RawAuthKeyMetadataProvider interface { + AuthKeyExpiresAtForSession(rawAuthKeyID [8]byte, sessionID int64) (expiresAt int, found bool) +} + // ImmediateSessionPusher 是可选的登录前信号直推能力。 // 它绕过登录后 updates-ready 队列，只能用于会解锁登录流程本身的握手消息， // 例如 updateLoginToken。 diff --git a/internal/rpc/errors.go b/internal/rpc/errors.go index 5ca60f2e..a6c87ba4 100644 --- a/internal/rpc/errors.go +++ b/internal/rpc/errors.go @@ -399,6 +399,8 @@ func signInErr(err error) error { return tgerr.New(400, "PHONE_CODE_INVALID") case errors.Is(err, auth.ErrCodeExpired): return tgerr.New(400, "PHONE_CODE_EXPIRED") + case errors.Is(err, auth.ErrAuthKeyPermEmpty): + return tgerr.New(401, "AUTH_KEY_PERM_EMPTY") case errors.Is(err, domain.ErrFirstNameInvalid): return firstNameInvalidErr() case errors.Is(err, domain.ErrSessionPasswordNeeded): @@ -455,8 +457,14 @@ func passwordErr(err error) error { // bindTempAuthKeyErr 映射 PFS temp auth key 绑定错误。 func bindTempAuthKeyErr(err error) error { switch { + case errors.Is(err, auth.ErrExpiresAtInvalid): + return tgerr.New(400, "EXPIRES_AT_INVALID") + case errors.Is(err, auth.ErrTempAuthKeyEmpty): + return tgerr.New(400, "TEMP_AUTH_KEY_EMPTY") case errors.Is(err, auth.ErrEncryptedMessageInvalid): return tgerr.New(400, "ENCRYPTED_MESSAGE_INVALID") + case errors.Is(err, auth.ErrTempAuthKeyAlreadyBound): + return tgerr.New(400, "TEMP_AUTH_KEY_ALREADY_BOUND") default: return internalErr() } diff --git a/internal/rpc/errors_test.go b/internal/rpc/errors_test.go index 77124484..657d483f 100644 --- a/internal/rpc/errors_test.go +++ b/internal/rpc/errors_test.go @@ -5,6 +5,7 @@ import ( "github.com/gotd/td/tgerr" + "telesrv/internal/app/auth" "telesrv/internal/domain" ) @@ -13,3 +14,22 @@ func TestPasswordErrMapsOccupiedLoginEmailToNotAllowed(t *testing.T) { t.Fatalf("passwordErr(ErrEmailOccupied) = %v, want EMAIL_NOT_ALLOWED", err) } } + +func TestBindTempAuthKeyErrPreservesRecoverableRotationErrors(t *testing.T) { + tests := []struct { + err error + want string + }{ + {err: auth.ErrExpiresAtInvalid, want: "EXPIRES_AT_INVALID"}, + {err: auth.ErrTempAuthKeyEmpty, want: "TEMP_AUTH_KEY_EMPTY"}, + {err: auth.ErrEncryptedMessageInvalid, want: "ENCRYPTED_MESSAGE_INVALID"}, + {err: auth.ErrTempAuthKeyAlreadyBound, want: "TEMP_AUTH_KEY_ALREADY_BOUND"}, + } + for _, test := range tests { + t.Run(test.want, func(t *testing.T) { + if err := bindTempAuthKeyErr(test.err); !tgerr.Is(err, test.want) { + t.Fatalf("bindTempAuthKeyErr(%v) = %v, want %s", test.err, err, test.want) + } + }) + } +} diff --git a/internal/rpc/router.go b/internal/rpc/router.go index 76c8debb..f44fb8e3 100644 --- a/internal/rpc/router.go +++ b/internal/rpc/router.go @@ -324,9 +324,25 @@ func (r *Router) effectiveAuthKeyID(ctx context.Context, rawAuthKeyID [8]byte, s } } if hasCached { - if cached == rawAuthKeyID || r.deps.Auth == nil { + if r.deps.Auth == nil { return cached, nil } + if cached == rawAuthKeyID { + // raw==business is a permanent-key fast path only when the edge confirms + // the raw key has no protocol expiry. A temporary session may have cached + // raw before another concurrent session completes auth.bindTempAuthKey; + // it must keep resolving until the durable binding becomes visible. + metadata, ok := r.deps.Sessions.(RawAuthKeyMetadataProvider) + if ok { + expiresAt, found := metadata.AuthKeyExpiresAtForSession(rawAuthKeyID, sessionID) + if found && expiresAt == 0 { + return cached, nil + } + } + // Missing metadata and a session lookup miss both fail closed to the + // durable resolver. Treating either as proof of permanence recreates the + // raw-temp identity split when an alternate SessionBinder is installed. + } // temp→perm 解析缓存：PFS 连接每帧都要解析一次 temp key（ResolveAuthKey 打 PG）。TTL 内复用 // 上次解析、跳过 DB。仅当缓存的 perm 仍等于 session binder 当前 perm 才用（rebind 会改 binder // 且 onAuthBindTempAuthKey / 授权撤销都会显式 Delete 缓存，双保险防跨账号串号和被踢滞后）。 diff --git a/internal/rpc/router_dispatch_test.go b/internal/rpc/router_dispatch_test.go index ff3b1cef..5613cc71 100644 --- a/internal/rpc/router_dispatch_test.go +++ b/internal/rpc/router_dispatch_test.go @@ -15,6 +15,7 @@ import ( appauth "telesrv/internal/app/auth" appdialogs "telesrv/internal/app/dialogs" "telesrv/internal/domain" + "telesrv/internal/store" "telesrv/internal/store/memory" "testing" "time" @@ -913,17 +914,25 @@ func TestDispatchUnknownReturnsError(t *testing.T) { func TestDispatchResolvesBoundTempAuthKey(t *testing.T) { var tempAuthKeyID = [8]byte{0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55} var permAuthKeyID = [8]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11} - tempBindings := memory.NewTempAuthKeyBindingStore() + expiresAt := int(time.Now().Add(time.Hour).Unix()) + authKeys := memory.NewAuthKeyStore() + if err := authKeys.Save(context.Background(), store.AuthKeyData{ID: tempAuthKeyID, ExpiresAt: expiresAt}); err != nil { + t.Fatalf("save temporary auth key: %v", err) + } + if err := authKeys.Save(context.Background(), store.AuthKeyData{ID: permAuthKeyID}); err != nil { + t.Fatalf("save permanent auth key: %v", err) + } + tempBindings := memory.NewTempAuthKeyBindingStore(authKeys) if err := tempBindings.Save(context.Background(), domain.TempAuthKeyBinding{ TempAuthKeyID: tempAuthKeyID, PermAuthKeyID: int64(binary.LittleEndian.Uint64(permAuthKeyID[:])), - ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + ExpiresAt: expiresAt, }); err != nil { t.Fatalf("save temp binding: %v", err) } sessions := &captureSessions{} r := New(Config{}, Deps{ - Auth: appauth.NewService(nil, nil, nil, nil, tempBindings, "12345"), + Auth: appauth.NewService(nil, nil, nil, authKeys, tempBindings, "12345"), Sessions: sessions, }, zaptest.NewLogger(t), clock.System) req := &tg.HelpGetConfigRequest{} diff --git a/internal/rpc/router_tempkey_cache_test.go b/internal/rpc/router_tempkey_cache_test.go index 062e8485..0f70f8ff 100644 --- a/internal/rpc/router_tempkey_cache_test.go +++ b/internal/rpc/router_tempkey_cache_test.go @@ -17,6 +17,19 @@ type revokeCaptureSessions struct { closedRawAuthKeyIDs [][8]byte } +type expiringCaptureSessions struct { + *captureSessions + expiresAt int +} + +type metadataBlindSessions struct { + SessionBinder +} + +func (s *expiringCaptureSessions) AuthKeyExpiresAtForSession([8]byte, int64) (int, bool) { + return s.expiresAt, true +} + func (s *revokeCaptureSessions) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int { s.mu.Lock() defer s.mu.Unlock() @@ -31,6 +44,76 @@ func (s *revokeCaptureSessions) CloseSessionsForRawAuthKeyExcept(authKeyID [8]by return 1 } +func TestCachedRawTemporarySessionReResolvesDurableBinding(t *testing.T) { + tempAuthKeyID := [8]byte{0x76, 0x76, 0x76, 0x76, 0x76, 0x76, 0x76, 0x76} + permAuthKeyID := [8]byte{0x36, 0x36, 0x36, 0x36, 0x36, 0x36, 0x36, 0x36} + base := &captureSessions{} + base.BindAuthKeyForSession(tempAuthKeyID, 554, tempAuthKeyID) + sessions := &expiringCaptureSessions{ + captureSessions: base, + expiresAt: int(time.Now().Add(time.Hour).Unix()), + } + auth := &captureAuthService{ + resolvedAuthKeyID: permAuthKeyID, + hasResolved: true, + userID: 1000000001, + } + r := New(Config{TempKeyResolveCacheTTL: time.Minute}, Deps{ + Auth: auth, + Files: &fakeFiles{}, + Sessions: sessions, + }, zaptest.NewLogger(t), clock.System) + + var in bin.Buffer + if err := (&tg.UploadSaveFilePartRequest{FileID: 19, FilePart: 0, Bytes: []byte{1}}).Encode(&in); err != nil { + t.Fatalf("encode: %v", err) + } + if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 554, &in); err != nil { + t.Fatalf("dispatch: %v", err) + } + if auth.resolveCount != 1 { + t.Fatalf("ResolveAuthKey calls = %d, want 1 for cached raw temporary session", auth.resolveCount) + } + got := sessions.snapshot() + if got.authKeyID != permAuthKeyID || got.userID != 1000000001 { + t.Fatalf("session = auth %x user %d, want perm/user", got.authKeyID, got.userID) + } +} + +func TestCachedRawSessionWithoutMetadataFailsClosedToDurableResolver(t *testing.T) { + tempAuthKeyID := [8]byte{0x75, 0x75, 0x75, 0x75, 0x75, 0x75, 0x75, 0x75} + permAuthKeyID := [8]byte{0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35, 0x35} + // captureSessions intentionally has no RawAuthKeyMetadataProvider capability. + // Missing metadata is not evidence that raw is permanent. + base := &captureSessions{} + base.BindAuthKeyForSession(tempAuthKeyID, 553, tempAuthKeyID) + sessions := &metadataBlindSessions{SessionBinder: base} + auth := &captureAuthService{ + resolvedAuthKeyID: permAuthKeyID, + hasResolved: true, + userID: 1000000001, + } + r := New(Config{TempKeyResolveCacheTTL: time.Minute}, Deps{ + Auth: auth, + Files: &fakeFiles{}, + Sessions: sessions, + }, zaptest.NewLogger(t), clock.System) + + var in bin.Buffer + if err := (&tg.UploadSaveFilePartRequest{FileID: 18, FilePart: 0, Bytes: []byte{1}}).Encode(&in); err != nil { + t.Fatalf("encode: %v", err) + } + if _, err := r.Dispatch(context.Background(), tempAuthKeyID, 553, &in); err != nil { + t.Fatalf("dispatch: %v", err) + } + if auth.resolveCount != 1 { + t.Fatalf("ResolveAuthKey calls = %d, want 1 without metadata proof", auth.resolveCount) + } + if got := base.snapshot(); got.authKeyID != permAuthKeyID || got.userID != 1000000001 { + t.Fatalf("session = auth %x user %d, want canonical perm/user", got.authKeyID, got.userID) + } +} + // TestTempKeyResolveCacheHitsWithinTTL 验证：TempKeyResolveCacheTTL>0 时，同一 temp key 的连续 // 请求在 TTL 内只解析一次（首帧走 !hasCached 解析 1 次、次帧 hasCached 解析并填缓存 1 次，之后命中 // 缓存不再打 ResolveAuthKey）。固化「缓存生效」语义，与现有「TTL=0 每帧重校验」的安全测试互补。 diff --git a/internal/rpc/rpc_testkit_sessions_test.go b/internal/rpc/rpc_testkit_sessions_test.go index f88baf74..134d9fee 100644 --- a/internal/rpc/rpc_testkit_sessions_test.go +++ b/internal/rpc/rpc_testkit_sessions_test.go @@ -101,6 +101,12 @@ func (s *captureSessions) AuthKeyIDForSession([8]byte, int64) ([8]byte, bool) { return s.authKeyID, s.authKeyResolved } +// captureSessions models an ordinary permanent-key connection unless a test +// wraps/overrides it with temporary-key metadata. +func (s *captureSessions) AuthKeyExpiresAtForSession([8]byte, int64) (int, bool) { + return 0, true +} + func (s *captureSessions) BindUserForAuthKey(rawAuthKeyID [8]byte, sessionID, userID int64) { s.mu.Lock() defer s.mu.Unlock() diff --git a/internal/store/authkey.go b/internal/store/authkey.go index 493259ba..b2bcb8b6 100644 --- a/internal/store/authkey.go +++ b/internal/store/authkey.go @@ -1,15 +1,44 @@ package store -import "context" +import ( + "context" + "errors" + "math" +) + +var ( + // ErrInvalidAuthKeyProtocolExpiry 表示新握手试图写入 migration-only + // unknown sentinel 或超出 TL int32 时间戳范围的协议寿命。 + ErrInvalidAuthKeyProtocolExpiry = errors.New("invalid auth key protocol expiry") + // ErrAuthKeyProtocolMetadataConflict 表示同一 cryptographic auth_key_id + // 被尝试改写为另一 key body 或另一 permanent/temp 类型/寿命。 + ErrAuthKeyProtocolMetadataConflict = errors.New("auth key protocol metadata conflict") + // ErrAuthKeyNotPermanent 防止 authorization 落到 temporary/legacy-unknown key。 + ErrAuthKeyNotPermanent = errors.New("auth key is not permanent") + // ErrAuthKeyBindingInvalid 表示 temp/perm 引用缺失、类型错误，或 binding + // expiry 未归一到握手权威值。 + ErrAuthKeyBindingInvalid = errors.New("invalid temporary auth key binding") +) + +// ValidNewAuthKeyProtocolExpiry 只允许握手写 permanent(0) 或 TL int32 +// 范围内的 positive temporary expiry。-1 仅能由 migration 0086 写入。 +func ValidNewAuthKeyProtocolExpiry(expiresAt int) bool { + return expiresAt >= 0 && int64(expiresAt) <= math.MaxInt32 +} // AuthKeyData 是一条持久化的 MTProto auth key 记录。 // // 不依赖 td 协议类型：连接层在边界做 crypto.AuthKey ↔ AuthKeyData 转换。 type AuthKeyData struct { - ID [8]byte // auth_key_id（key 的 SHA1 低 64 位） - Value [256]byte // 2048-bit auth key - ServerSalt int64 // 密钥交换产出的初始 server salt - CreatedAt int64 // unix 秒 + ID [8]byte // auth_key_id（key 的 SHA1 低 64 位） + Value [256]byte // 2048-bit auth key + ServerSalt int64 // 密钥交换产出的初始 server salt + CreatedAt int64 // unix 秒 + // ExpiresAt 是 temporary/media-temporary auth key 的协议失效时间（unix 秒）。 + // 0 只允许表示 permanent key；-1 仅表示 migration 0086 无法证明类型的历史 key， + // edge 必须用 -404 拒绝并迫使客户端重握手。key 类型是握手事实，不能由 + // authorization 是否存在推断。 + ExpiresAt int Layer int DeviceModel string Platform string @@ -30,7 +59,7 @@ type AuthKeyClientInfo struct { // AuthKeyStore 持久化 auth key。实现见 store/memory（测试替身）、store/postgres。 type AuthKeyStore interface { - // Save 保存或覆盖一条 auth key 记录。 + // Save 保存一条 auth key 记录；同 ID 重试只能保持 key body 与协议类型/寿命不变。 Save(ctx context.Context, k AuthKeyData) error // Get 按 auth_key_id 查询；不存在时 found=false。 Get(ctx context.Context, id [8]byte) (data AuthKeyData, found bool, err error) diff --git a/internal/store/memory/auth.go b/internal/store/memory/auth.go index 94b5b4e0..edfa2e40 100644 --- a/internal/store/memory/auth.go +++ b/internal/store/memory/auth.go @@ -4,44 +4,60 @@ import ( "context" "encoding/binary" "sync" + "time" + "telesrv/internal/domain" "telesrv/internal/store" - "time" ) +type authKeyState struct { + mu sync.RWMutex + keys map[[8]byte]store.AuthKeyData + bindings map[[8]byte]domain.TempAuthKeyBinding +} + // AuthKeyStore 是 store.AuthKeyStore 的内存实现。 type AuthKeyStore struct { - mu sync.RWMutex - keys map[[8]byte]store.AuthKeyData + state *authKeyState } // NewAuthKeyStore 创建内存 AuthKeyStore。 func NewAuthKeyStore() *AuthKeyStore { - return &AuthKeyStore{keys: make(map[[8]byte]store.AuthKeyData)} + return &AuthKeyStore{state: &authKeyState{ + keys: make(map[[8]byte]store.AuthKeyData), + bindings: make(map[[8]byte]domain.TempAuthKeyBinding), + }} } func (s *AuthKeyStore) Save(_ context.Context, k store.AuthKeyData) error { - s.mu.Lock() - s.keys[k.ID] = k - s.mu.Unlock() + if !store.ValidNewAuthKeyProtocolExpiry(k.ExpiresAt) { + return store.ErrInvalidAuthKeyProtocolExpiry + } + s.state.mu.Lock() + if current, ok := s.state.keys[k.ID]; ok && (current.Value != k.Value || current.ExpiresAt != k.ExpiresAt) { + s.state.mu.Unlock() + return store.ErrAuthKeyProtocolMetadataConflict + } + s.state.keys[k.ID] = k + s.state.mu.Unlock() return nil } func (s *AuthKeyStore) Get(_ context.Context, id [8]byte) (store.AuthKeyData, bool, error) { - s.mu.RLock() - k, ok := s.keys[id] - s.mu.RUnlock() + s.state.mu.RLock() + k, ok := s.state.keys[id] + s.state.mu.RUnlock() return k, ok, nil } func (s *AuthKeyStore) UpdateClientInfo(_ context.Context, id [8]byte, info store.AuthKeyClientInfo) error { - s.mu.Lock() - k, ok := s.keys[id] + s.state.mu.Lock() + k, ok := s.state.keys[id] if ok { mergeAuthKeyClientInfo(&k, info) - s.keys[id] = k + s.state.keys[id] = k } - s.mu.Unlock() + s.state.mu.Unlock() return nil } @@ -67,35 +83,64 @@ func mergeAuthKeyClientInfo(k *store.AuthKeyData, info store.AuthKeyClientInfo) } func (s *AuthKeyStore) Delete(_ context.Context, id [8]byte) error { - s.mu.Lock() - delete(s.keys, id) - s.mu.Unlock() + s.state.mu.Lock() + deleting, exists := s.state.keys[id] + if !exists { + s.state.mu.Unlock() + return nil + } + if deleting.ExpiresAt > 0 { + delete(s.state.bindings, id) + } else { + permID := int64(binary.LittleEndian.Uint64(id[:])) + for tempID, binding := range s.state.bindings { + if binding.PermAuthKeyID != permID { + continue + } + delete(s.state.bindings, tempID) + delete(s.state.keys, tempID) + } + } + delete(s.state.keys, id) + s.state.mu.Unlock() return nil } // TempAuthKeyBindingStore 是 store.TempAuthKeyBindingStore 的内存实现。 type TempAuthKeyBindingStore struct { - mu sync.RWMutex - m map[[8]byte]domain.TempAuthKeyBinding + state *authKeyState } // NewTempAuthKeyBindingStore 创建内存 TempAuthKeyBindingStore。 -func NewTempAuthKeyBindingStore() *TempAuthKeyBindingStore { - return &TempAuthKeyBindingStore{m: make(map[[8]byte]domain.TempAuthKeyBinding)} +func NewTempAuthKeyBindingStore(authKeys *AuthKeyStore) *TempAuthKeyBindingStore { + if authKeys == nil { + panic("memory.NewTempAuthKeyBindingStore requires a non-nil AuthKeyStore") + } + return &TempAuthKeyBindingStore{state: authKeys.state} } func (s *TempAuthKeyBindingStore) Save(_ context.Context, b domain.TempAuthKeyBinding) error { b.EncryptedMessage = append([]byte(nil), b.EncryptedMessage...) - s.mu.Lock() - s.m[b.TempAuthKeyID] = b - s.mu.Unlock() + s.state.mu.Lock() + defer s.state.mu.Unlock() + if current, ok := s.state.bindings[b.TempAuthKeyID]; ok && current.PermAuthKeyID != b.PermAuthKeyID { + return store.ErrTempAuthKeyAlreadyBound + } + temp, tempFound := s.state.keys[b.TempAuthKeyID] + var permID [8]byte + binary.LittleEndian.PutUint64(permID[:], uint64(b.PermAuthKeyID)) + perm, permFound := s.state.keys[permID] + if !tempFound || !permFound || temp.ExpiresAt <= 0 || perm.ExpiresAt != 0 || b.ExpiresAt != temp.ExpiresAt { + return store.ErrAuthKeyBindingInvalid + } + s.state.bindings[b.TempAuthKeyID] = b return nil } func (s *TempAuthKeyBindingStore) GetByTemp(_ context.Context, tempAuthKeyID [8]byte) (domain.TempAuthKeyBinding, bool, error) { - s.mu.RLock() - b, ok := s.m[tempAuthKeyID] - s.mu.RUnlock() + s.state.mu.RLock() + b, ok := s.state.bindings[tempAuthKeyID] + s.state.mu.RUnlock() if !ok { return domain.TempAuthKeyBinding{}, false, nil } @@ -107,17 +152,19 @@ func (s *TempAuthKeyBindingStore) DeleteExpired(_ context.Context, expiredBefore if limit <= 0 { return 0, nil } - s.mu.Lock() - defer s.mu.Unlock() + s.state.mu.Lock() + defer s.state.mu.Unlock() deleted := 0 - for id, b := range s.m { + for id, key := range s.state.keys { if deleted >= limit { break } - if int64(b.ExpiresAt) < expiredBefore { - delete(s.m, id) - deleted++ + if key.ExpiresAt <= 0 || int64(key.ExpiresAt) >= expiredBefore { + continue } + delete(s.state.bindings, id) + delete(s.state.keys, id) + deleted++ } return deleted, nil } diff --git a/internal/store/memory/auth_test.go b/internal/store/memory/auth_test.go new file mode 100644 index 00000000..3a300e57 --- /dev/null +++ b/internal/store/memory/auth_test.go @@ -0,0 +1,289 @@ +package memory + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "testing" + + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestAuthKeyStorePreservesProtocolExpiry(t *testing.T) { + ctx := context.Background() + keys := NewAuthKeyStore() + want := store.AuthKeyData{ + ID: [8]byte{1, 2, 3, 4, 5, 6, 7, 8}, + ServerSalt: 42, + ExpiresAt: 1_799_999_999, + } + want.Value[0] = 0xaa + want.Value[len(want.Value)-1] = 0x55 + + if err := keys.Save(ctx, want); err != nil { + t.Fatalf("save: %v", err) + } + got, found, err := keys.Get(ctx, want.ID) + if err != nil || !found { + t.Fatalf("get: found=%v err=%v", found, err) + } + if got != want { + t.Fatalf("round trip mismatch: got %+v, want %+v", got, want) + } + + conflicting := want + conflicting.ExpiresAt++ + if err := keys.Save(ctx, conflicting); !errors.Is(err, store.ErrAuthKeyProtocolMetadataConflict) { + t.Fatalf("reclassify auth key error = %v, want %v", err, store.ErrAuthKeyProtocolMetadataConflict) + } + got, found, err = keys.Get(ctx, want.ID) + if err != nil || !found || got != want { + t.Fatalf("auth key changed after rejected reclassification: got=%+v found=%v err=%v", got, found, err) + } +} + +func TestTempAuthKeyBindingStoreIsIdempotentAndRejectsCrossPermanentRebind(t *testing.T) { + ctx := context.Background() + keys := NewAuthKeyStore() + bindings := NewTempAuthKeyBindingStore(keys) + handshakeExpiry := 400 + permID := memoryAuthKeyID(101) + otherPermID := memoryAuthKeyID(102) + first := domain.TempAuthKeyBinding{ + TempAuthKeyID: [8]byte{8, 7, 6, 5, 4, 3, 2, 1}, + PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])), + Nonce: 201, + TempSessionID: 301, + ExpiresAt: handshakeExpiry, + EncryptedMessage: []byte("first"), + } + if err := keys.Save(ctx, store.AuthKeyData{ID: permID}); err != nil { + t.Fatalf("save permanent auth key: %v", err) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: otherPermID}); err != nil { + t.Fatalf("save second permanent auth key: %v", err) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: first.TempAuthKeyID, ExpiresAt: handshakeExpiry}); err != nil { + t.Fatalf("save temporary auth key: %v", err) + } + if err := bindings.Save(ctx, first); err != nil { + t.Fatalf("save first: %v", err) + } + assertMemoryAuthKeyExpiry(t, ctx, keys, first.TempAuthKeyID, handshakeExpiry) + + replayed := first + replayed.Nonce = 202 + replayed.TempSessionID = 302 + replayed.ExpiresAt = 402 + replayed.EncryptedMessage = []byte("replayed") + if err := bindings.Save(ctx, replayed); !errors.Is(err, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("replay with changed expiry error = %v, want %v", err, store.ErrAuthKeyBindingInvalid) + } + assertMemoryAuthKeyExpiry(t, ctx, keys, first.TempAuthKeyID, handshakeExpiry) + got, found, err := bindings.GetByTemp(ctx, first.TempAuthKeyID) + if err != nil || !found || got.ExpiresAt != first.ExpiresAt || got.Nonce != first.Nonce { + t.Fatalf("binding after invalid expiry replay = %+v found=%v err=%v, want first binding", got, found, err) + } + + replayed.ExpiresAt = handshakeExpiry + if err := bindings.Save(ctx, replayed); err != nil { + t.Fatalf("replay same normalized binding: %v", err) + } + + forbidden := replayed + forbidden.PermAuthKeyID = int64(binary.LittleEndian.Uint64(otherPermID[:])) + forbidden.ExpiresAt = 999 + forbidden.EncryptedMessage = []byte("must not persist") + if err := bindings.Save(ctx, forbidden); !errors.Is(err, store.ErrTempAuthKeyAlreadyBound) { + t.Fatalf("cross-permanent rebind error = %v, want %v", err, store.ErrTempAuthKeyAlreadyBound) + } + assertMemoryAuthKeyExpiry(t, ctx, keys, first.TempAuthKeyID, handshakeExpiry) + + got, found, err = bindings.GetByTemp(ctx, first.TempAuthKeyID) + if err != nil || !found { + t.Fatalf("get: found=%v err=%v", found, err) + } + if got.TempAuthKeyID != replayed.TempAuthKeyID || got.PermAuthKeyID != replayed.PermAuthKeyID || + got.Nonce != replayed.Nonce || got.TempSessionID != replayed.TempSessionID || got.ExpiresAt != replayed.ExpiresAt || + !bytes.Equal(got.EncryptedMessage, replayed.EncryptedMessage) { + t.Fatalf("binding changed after forbidden rebind: got %+v, want %+v", got, replayed) + } +} + +func TestTempAuthKeyBindingStoreRejectsMissingTypeAndExpiryViolations(t *testing.T) { + ctx := context.Background() + const handshakeExpiry = 500 + tempID := memoryAuthKeyID(201) + permID := memoryAuthKeyID(202) + + tests := []struct { + name string + temp *store.AuthKeyData + perm *store.AuthKeyData + bindingExpiry int + }{ + { + name: "missing temporary key", + perm: &store.AuthKeyData{ID: permID}, + bindingExpiry: handshakeExpiry, + }, + { + name: "missing permanent key", + temp: &store.AuthKeyData{ID: tempID, ExpiresAt: handshakeExpiry}, + bindingExpiry: handshakeExpiry, + }, + { + name: "temporary role uses permanent key", + temp: &store.AuthKeyData{ID: tempID}, + perm: &store.AuthKeyData{ID: permID}, + bindingExpiry: handshakeExpiry, + }, + { + name: "permanent role uses temporary key", + temp: &store.AuthKeyData{ID: tempID, ExpiresAt: handshakeExpiry}, + perm: &store.AuthKeyData{ID: permID, ExpiresAt: handshakeExpiry + 1}, + bindingExpiry: handshakeExpiry, + }, + { + name: "binding expiry differs from handshake", + temp: &store.AuthKeyData{ID: tempID, ExpiresAt: handshakeExpiry}, + perm: &store.AuthKeyData{ID: permID}, + bindingExpiry: handshakeExpiry + 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + keys := NewAuthKeyStore() + bindings := NewTempAuthKeyBindingStore(keys) + if tt.temp != nil { + if err := keys.Save(ctx, *tt.temp); err != nil { + t.Fatalf("save temporary role key: %v", err) + } + } + if tt.perm != nil { + if err := keys.Save(ctx, *tt.perm); err != nil { + t.Fatalf("save permanent role key: %v", err) + } + } + err := bindings.Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: tempID, + PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])), + ExpiresAt: tt.bindingExpiry, + }) + if !errors.Is(err, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("Save error = %v, want %v", err, store.ErrAuthKeyBindingInvalid) + } + if _, found, getErr := bindings.GetByTemp(ctx, tempID); getErr != nil || found { + t.Fatalf("invalid binding found=%v err=%v, want absent", found, getErr) + } + }) + } +} + +func TestAuthKeyStoreDeletePermanentRemovesBoundTemporaryIdentity(t *testing.T) { + ctx := context.Background() + keys := NewAuthKeyStore() + bindings := NewTempAuthKeyBindingStore(keys) + tempID := memoryAuthKeyID(301) + permID := memoryAuthKeyID(302) + const expiresAt = 600 + if err := keys.Save(ctx, store.AuthKeyData{ID: tempID, ExpiresAt: expiresAt}); err != nil { + t.Fatalf("save temp: %v", err) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: permID}); err != nil { + t.Fatalf("save perm: %v", err) + } + if err := bindings.Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: tempID, + PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])), + ExpiresAt: expiresAt, + }); err != nil { + t.Fatalf("save binding: %v", err) + } + + if err := keys.Delete(ctx, permID); err != nil { + t.Fatalf("delete permanent key: %v", err) + } + if _, found, err := keys.Get(ctx, permID); err != nil || found { + t.Fatalf("permanent key found=%v err=%v, want absent", found, err) + } + if _, found, err := keys.Get(ctx, tempID); err != nil || found { + t.Fatalf("bound temporary key found=%v err=%v, want absent", found, err) + } + if _, found, err := bindings.GetByTemp(ctx, tempID); err != nil || found { + t.Fatalf("binding found=%v err=%v, want absent", found, err) + } +} + +func TestTempAuthKeyBindingStoreDeleteExpiredUsesAuthKeyExpiry(t *testing.T) { + ctx := context.Background() + keys := NewAuthKeyStore() + bindings := NewTempAuthKeyBindingStore(keys) + permID := memoryAuthKeyID(401) + boundExpiredID := memoryAuthKeyID(402) + unboundExpiredID := memoryAuthKeyID(403) + liveID := memoryAuthKeyID(404) + if err := keys.Save(ctx, store.AuthKeyData{ID: permID}); err != nil { + t.Fatalf("save perm: %v", err) + } + for id, expiry := range map[[8]byte]int{ + boundExpiredID: 700, + unboundExpiredID: 701, + liveID: 900, + } { + if err := keys.Save(ctx, store.AuthKeyData{ID: id, ExpiresAt: expiry}); err != nil { + t.Fatalf("save temp %x: %v", id, err) + } + } + if err := bindings.Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: boundExpiredID, + PermAuthKeyID: int64(binary.LittleEndian.Uint64(permID[:])), + ExpiresAt: 700, + }); err != nil { + t.Fatalf("save binding: %v", err) + } + + deleted, err := bindings.DeleteExpired(ctx, 800, 10) + if err != nil || deleted != 2 { + t.Fatalf("DeleteExpired = %d, %v; want 2, nil", deleted, err) + } + for _, id := range [][8]byte{boundExpiredID, unboundExpiredID} { + if _, found, getErr := keys.Get(ctx, id); getErr != nil || found { + t.Fatalf("expired key %x found=%v err=%v, want absent", id, found, getErr) + } + } + if _, found, err := bindings.GetByTemp(ctx, boundExpiredID); err != nil || found { + t.Fatalf("expired binding found=%v err=%v, want absent", found, err) + } + for _, id := range [][8]byte{permID, liveID} { + if _, found, getErr := keys.Get(ctx, id); getErr != nil || !found { + t.Fatalf("retained key %x found=%v err=%v, want present", id, found, getErr) + } + } +} + +func memoryAuthKeyID(id int64) [8]byte { + var out [8]byte + binary.LittleEndian.PutUint64(out[:], uint64(id)) + return out +} + +func assertMemoryAuthKeyExpiry( + t *testing.T, + ctx context.Context, + keys store.AuthKeyStore, + id [8]byte, + want int, +) { + t.Helper() + got, found, err := keys.Get(ctx, id) + if err != nil || !found { + t.Fatalf("get auth key: found=%v err=%v", found, err) + } + if got.ExpiresAt != want { + t.Fatalf("auth key expires_at = %d, want handshake expiry %d", got.ExpiresAt, want) + } +} diff --git a/internal/store/postgres/auth_key_expiry_migration_integration_test.go b/internal/store/postgres/auth_key_expiry_migration_integration_test.go new file mode 100644 index 00000000..a29f7772 --- /dev/null +++ b/internal/store/postgres/auth_key_expiry_migration_integration_test.go @@ -0,0 +1,471 @@ +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', 86, 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) + } +} diff --git a/internal/store/postgres/authkey.go b/internal/store/postgres/authkey.go index e0af8c25..4c871d31 100644 --- a/internal/store/postgres/authkey.go +++ b/internal/store/postgres/authkey.go @@ -8,6 +8,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgtype" "telesrv/internal/store" @@ -28,14 +29,24 @@ func NewAuthKeyStore(db sqlcgen.DBTX) *AuthKeyStore { // Save 实现 store.AuthKeyStore。auth_key_id 以小端解释为 int64 存入 BIGINT； // created_at/last_used_at 交由 DB 默认值（now()），故传入的 CreatedAt 不落库。 func (s *AuthKeyStore) Save(ctx context.Context, k store.AuthKeyData) error { - if _, err := s.db.Exec(ctx, ` -INSERT INTO auth_keys (auth_key_id, body, server_salt) -VALUES ($1, $2, $3) + if !store.ValidNewAuthKeyProtocolExpiry(k.ExpiresAt) { + return store.ErrInvalidAuthKeyProtocolExpiry + } + tag, err := s.db.Exec(ctx, ` +INSERT INTO auth_keys (auth_key_id, body, server_salt, expires_at) +VALUES ($1, $2, $3, $4) ON CONFLICT (auth_key_id) DO UPDATE -SET body = EXCLUDED.body, server_salt = EXCLUDED.server_salt, last_used_at = now() -`, authKeyIDToInt64(k.ID), k.Value[:], k.ServerSalt); err != nil { +SET server_salt = EXCLUDED.server_salt, + last_used_at = now() +WHERE auth_keys.body = EXCLUDED.body + AND auth_keys.expires_at = EXCLUDED.expires_at +`, authKeyIDToInt64(k.ID), k.Value[:], k.ServerSalt, k.ExpiresAt) + if err != nil { return fmt.Errorf("upsert auth key: %w", err) } + if tag.RowsAffected() != 1 { + return store.ErrAuthKeyProtocolMetadataConflict + } return nil } @@ -47,6 +58,7 @@ func (s *AuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, var ( body []byte serverSalt int64 + expiresAt int createdAt pgtype.Timestamptz layer int deviceModel string @@ -60,8 +72,8 @@ UPDATE auth_keys SET last_used_at = now() WHERE auth_key_id = $1 RETURNING auth_key_id, body, server_salt, created_at, - layer, device_model, platform, system_version, api_id, app_version -`, authKeyIDToInt64(id)).Scan(new(int64), &body, &serverSalt, &createdAt, &layer, &deviceModel, &platform, &systemVersion, &apiID, &appVersion) + 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) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return store.AuthKeyData{}, false, nil @@ -74,6 +86,7 @@ RETURNING auth_key_id, body, server_salt, created_at, data := store.AuthKeyData{ ID: id, ServerSalt: serverSalt, + ExpiresAt: expiresAt, Layer: layer, DeviceModel: deviceModel, Platform: platform, @@ -154,10 +167,26 @@ WHERE auth_key_id = $1 // 手写 SQL 而非 sqlc 生成：避免触碰 sqlcgen 再生成链路。 // // 同时清理把本 key 当作 perm key 的 temp auth key 行：temp_auth_key_bindings.temp_auth_key_id -// 侧有外键 ON DELETE CASCADE，删除 temp key 会自动清绑定；perm_auth_key_id 列无外键， -// 因此被踢/登出删除 perm key 时必须先把关联 temp key 一并删掉。否则 Web/上传连接用 +// 侧有外键 ON DELETE CASCADE，删除 temp key 会自动清绑定；perm_auth_key_id 侧由 +// 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") +} + +func (s *AuthKeyStore) deleteAuthKeyOnce(ctx context.Context, id [8]byte) error { keyID := authKeyIDToInt64(id) var touched int if err := s.db.QueryRow(ctx, ` @@ -177,17 +206,32 @@ WITH doomed_temp AS MATERIALIZED ( RETURNING auth_key_id ), deleted_temp AS ( DELETE FROM auth_keys - WHERE auth_key_id IN (SELECT auth_key_id FROM doomed_keys) + WHERE auth_key_id IN (SELECT temp_auth_key_id FROM doomed_temp) + RETURNING auth_key_id +), deleted_key AS ( + DELETE FROM auth_keys + WHERE auth_key_id = $1 + AND (SELECT count(*) FROM deleted_temp) >= 0 RETURNING auth_key_id ) SELECT (SELECT count(*) FROM deleted_update_states)::int + - (SELECT count(*) FROM deleted_temp)::int`, keyID).Scan(&touched); err != nil { + (SELECT count(*) FROM deleted_temp)::int + + (SELECT count(*) FROM deleted_key)::int`, keyID).Scan(&touched); err != nil { return fmt.Errorf("delete auth key and temp bindings: %w", err) } return nil } +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 中复核。 diff --git a/internal/store/postgres/authkey_integration_test.go b/internal/store/postgres/authkey_integration_test.go index 041d9967..15599ee3 100644 --- a/internal/store/postgres/authkey_integration_test.go +++ b/internal/store/postgres/authkey_integration_test.go @@ -3,6 +3,7 @@ package postgres import ( "context" "crypto/rand" + "errors" "os" "testing" @@ -47,7 +48,12 @@ func TestAuthKeyStoreRoundTrip(t *testing.T) { _, _ = pool.Exec(ctx, "DELETE FROM auth_keys WHERE auth_key_id = $1", authKeyIDToInt64(id)) }) - want := store.AuthKeyData{ID: id, Value: val, ServerSalt: 0x0badf00d} + want := store.AuthKeyData{ + ID: id, + Value: val, + ServerSalt: 0x0badf00d, + ExpiresAt: 1_799_999_999, + } if err := NewAuthKeyStore(pool).Save(ctx, want); err != nil { t.Fatalf("save: %v", err) } @@ -59,9 +65,18 @@ func TestAuthKeyStoreRoundTrip(t *testing.T) { if !found { t.Fatal("auth key not found after save (重启后丢失)") } - if got.ID != want.ID || got.Value != want.Value || got.ServerSalt != want.ServerSalt { - t.Fatalf("round trip mismatch: got salt=%#x value[:4]=%x, want salt=%#x value[:4]=%x", - got.ServerSalt, got.Value[:4], want.ServerSalt, want.Value[:4]) + if got.ID != want.ID || got.Value != want.Value || got.ServerSalt != want.ServerSalt || got.ExpiresAt != want.ExpiresAt { + t.Fatalf("round trip mismatch: got salt=%#x expires_at=%d value[:4]=%x, want salt=%#x expires_at=%d value[:4]=%x", + got.ServerSalt, got.ExpiresAt, got.Value[:4], want.ServerSalt, want.ExpiresAt, want.Value[:4]) + } + conflicting := want + conflicting.ExpiresAt++ + if err := NewAuthKeyStore(pool).Save(ctx, conflicting); !errors.Is(err, store.ErrAuthKeyProtocolMetadataConflict) { + t.Fatalf("reclassify auth key error = %v, want %v", err, store.ErrAuthKeyProtocolMetadataConflict) + } + got, found, err = NewAuthKeyStore(pool).Get(ctx, id) + if err != nil || !found || got.ExpiresAt != want.ExpiresAt { + t.Fatalf("auth key expiry changed after rejected reclassification: got=%d found=%v err=%v", got.ExpiresAt, found, err) } var missing [8]byte diff --git a/internal/store/postgres/authkey_retention_integration_test.go b/internal/store/postgres/authkey_retention_integration_test.go index ef3dcfaf..1938041e 100644 --- a/internal/store/postgres/authkey_retention_integration_test.go +++ b/internal/store/postgres/authkey_retention_integration_test.go @@ -17,22 +17,23 @@ func TestAuthKeyStoreDeleteOrphanedIsBoundedAndProtectsReferencesPostgres(t *tes auths := NewAuthorizationStore(pool) userID := createRevokeTestUser(t, ctx, pool, "orphan-auth-key") - newKey := func() [8]byte { + newKey := func(expiresAt int) [8]byte { var id [8]byte if _, err := rand.Read(id[:]); err != nil { t.Fatalf("random auth key id: %v", err) } - if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + if err := keys.Save(ctx, store.AuthKeyData{ID: id, ExpiresAt: expiresAt}); err != nil { t.Fatalf("save auth key %x: %v", id, err) } t.Cleanup(func() { _ = keys.Delete(ctx, id) }) return id } - orphanOne, orphanTwo := newKey(), newKey() - recent := newKey() - authorized := newKey() - temp, perm := newKey(), newKey() - active := newKey() + tempExpiry := int(time.Now().Add(time.Hour).Unix()) + orphanOne, orphanTwo := newKey(0), newKey(0) + recent := newKey(0) + authorized := newKey(0) + temp, perm := newKey(tempExpiry), newKey(0) + active := newKey(0) if _, err := pool.Exec(ctx, ` INSERT INTO update_states (auth_key_id, user_id, pts, observed_pts) VALUES ($1, $3, 0, 0), ($2, $3, 0, 0)`, @@ -45,7 +46,7 @@ VALUES ($1, $3, 0, 0), ($2, $3, 0, 0)`, } if err := NewTempAuthKeyBindingStore(pool).Save(ctx, domain.TempAuthKeyBinding{ TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(perm), Nonce: 1, - TempSessionID: 2, ExpiresAt: int(time.Now().Add(time.Hour).Unix()), EncryptedMessage: []byte{1}, + TempSessionID: 2, ExpiresAt: tempExpiry, EncryptedMessage: []byte{1}, }); err != nil { t.Fatalf("save temp binding: %v", err) } @@ -107,8 +108,9 @@ func TestAuthKeyStoreDeleteCleansPermanentAndTempUpdateStatesPostgres(t *testing userID := createRevokeTestUser(t, ctx, pool, "auth-key-delete-state") perm := randomUpdateRetentionAuthKey(t) temp := randomUpdateRetentionAuthKey(t) - for _, id := range [][8]byte{perm, temp} { - if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + tempExpiry := int(time.Now().Add(time.Hour).Unix()) + for id, expiresAt := range map[[8]byte]int{perm: 0, temp: tempExpiry} { + if err := keys.Save(ctx, store.AuthKeyData{ID: id, ExpiresAt: expiresAt}); err != nil { t.Fatalf("save auth key %x: %v", id, err) } id := id @@ -119,7 +121,7 @@ func TestAuthKeyStoreDeleteCleansPermanentAndTempUpdateStatesPostgres(t *testing PermAuthKeyID: authKeyIDToInt64(perm), Nonce: 31, TempSessionID: 32, - ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + ExpiresAt: tempExpiry, EncryptedMessage: []byte{1}, }); err != nil { t.Fatalf("save temp binding: %v", err) @@ -147,6 +149,61 @@ SELECT } } +func TestTempAuthKeyRetentionUsesAuthKeyExpiryForBoundAndUnboundKeysPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + cutoff := int64(time.Now().Add(-time.Hour).Unix()) + unbound := randomUpdateRetentionAuthKey(t) + bound := randomUpdateRetentionAuthKey(t) + live := randomUpdateRetentionAuthKey(t) + perm := randomUpdateRetentionAuthKey(t) + expiries := map[[8]byte]int{ + unbound: int(cutoff - 2), + bound: int(cutoff - 1), + live: int(cutoff + 1), + perm: 0, + } + for id, expiresAt := range expiries { + if err := keys.Save(ctx, store.AuthKeyData{ID: id, ExpiresAt: expiresAt}); err != nil { + t.Fatalf("save key %x: %v", id, err) + } + id := id + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + } + if err := bindings.Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: bound, PermAuthKeyID: authKeyIDToInt64(perm), Nonce: 41, + TempSessionID: 42, ExpiresAt: expiries[bound], EncryptedMessage: []byte{1}, + }); err != nil { + t.Fatalf("save expired bound key: %v", err) + } + + deleted, err := bindings.DeleteExpired(ctx, cutoff, 1) + if err != nil || deleted != 1 { + t.Fatalf("first bounded expiry delete = %d/%v, want 1/nil", deleted, err) + } + if _, found, err := keys.Get(ctx, unbound); err != nil || found { + t.Fatalf("earliest unbound temp found=%v err=%v, want deleted", found, err) + } + if _, found, err := keys.Get(ctx, bound); err != nil || !found { + t.Fatalf("second expired bound temp found=%v err=%v, want retained after limit=1", found, err) + } + + deleted, err = bindings.DeleteExpired(ctx, cutoff, 10) + if err != nil || deleted != 1 { + t.Fatalf("second expiry delete = %d/%v, want 1/nil", deleted, err) + } + if _, found, err := bindings.GetByTemp(ctx, bound); err != nil || found { + t.Fatalf("binding after temp key cascade found=%v err=%v, want absent", found, err) + } + for name, id := range map[string][8]byte{"live temp": live, "permanent": perm} { + if _, found, err := keys.Get(ctx, id); err != nil || !found { + t.Fatalf("%s found=%v err=%v, want retained", name, found, err) + } + } +} + func TestAuthKeyGetTouchPreventsOrphanCollectionPostgres(t *testing.T) { pool := testPool(t) ctx := context.Background() diff --git a/internal/store/postgres/authorization.go b/internal/store/postgres/authorization.go index cf4ed2ff..f3adee28 100644 --- a/internal/store/postgres/authorization.go +++ b/internal/store/postgres/authorization.go @@ -10,6 +10,7 @@ import ( "github.com/jackc/pgx/v5" "telesrv/internal/domain" + "telesrv/internal/store" "telesrv/internal/store/postgres/sqlcgen" ) @@ -54,14 +55,20 @@ 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) - var lockedKeyID int64 + var ( + lockedKeyID int64 + expiresAt int + ) if err := db.QueryRow(ctx, ` -SELECT auth_key_id +SELECT auth_key_id, expires_at FROM auth_keys WHERE auth_key_id = $1 -FOR UPDATE`, keyID).Scan(&lockedKeyID); err != nil { +FOR UPDATE`, keyID).Scan(&lockedKeyID, &expiresAt); err != nil { return fmt.Errorf("lock auth key for authorization: %w", err) } + if expiresAt != 0 { + return store.ErrAuthKeyNotPermanent + } if _, err := db.Exec(ctx, ` INSERT INTO user_update_watermarks (user_id, contiguous_pts) @@ -254,43 +261,81 @@ 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) { - row := s.db.QueryRow(ctx, ` -WITH target AS MATERIALIZED ( - SELECT auth_key_id, user_id, hash, layer, device_model, platform, system_version, api_id, app_version, ip, password_pending, created_at, active_at - FROM authorizations - WHERE user_id = $1 AND hash = $2 -), deleted_temp AS ( - DELETE FROM auth_keys - WHERE auth_key_id IN ( - SELECT temp_auth_key_id - FROM temp_auth_key_bindings - WHERE perm_auth_key_id IN (SELECT auth_key_id FROM target) - ) - RETURNING auth_key_id -), deleted_update_states AS ( - DELETE FROM update_states - WHERE auth_key_id IN (SELECT auth_key_id FROM target) - RETURNING auth_key_id -), deleted_keys AS ( - DELETE FROM auth_keys - WHERE auth_key_id IN (SELECT auth_key_id FROM target) - RETURNING auth_key_id -), touched AS ( - SELECT - (SELECT count(*) FROM deleted_temp) + - (SELECT count(*) FROM deleted_update_states) AS count -) -SELECT target.auth_key_id, target.user_id, target.hash, target.layer, target.device_model, target.platform, - target.system_version, target.api_id, target.app_version, target.ip, target.password_pending, - target.created_at, target.active_at -FROM target -JOIN deleted_keys USING (auth_key_id) -CROSS JOIN touched`, userID, hash) - a, found, err := scanRevokedAuthorization(row) - if err != nil { - return domain.Authorization{}, false, fmt.Errorf("revoke authorization by hash: %w", err) + 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 + } } - return a, found, nil + 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) +} + +// revokeByHashTx deliberately uses separate READ COMMITTED statements. The first +// lookup is only a candidate. Bind locks auth_keys before changing authorization +// ownership, so revocation must lock the same parent row and then re-read the +// owner/hash from a fresh statement snapshot. Otherwise an A->B re-login that +// commits while revoke waits can be deleted using A's stale target snapshot. +func revokeByHashTx(ctx context.Context, tx pgx.Tx, userID, hash int64) (domain.Authorization, bool, error) { + var candidate int64 + if err := tx.QueryRow(ctx, ` +SELECT auth_key_id +FROM authorizations +WHERE user_id = $1 AND hash = $2`, userID, hash).Scan(&candidate); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return domain.Authorization{}, false, nil + } + return domain.Authorization{}, false, fmt.Errorf("select revoke candidate by hash: %w", err) + } + + var locked int64 + if err := tx.QueryRow(ctx, ` +SELECT auth_key_id +FROM auth_keys +WHERE auth_key_id = $1 +FOR UPDATE`, candidate).Scan(&locked); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return domain.Authorization{}, false, nil + } + return domain.Authorization{}, false, fmt.Errorf("lock revoke auth key by hash: %w", err) + } + + a, found, err := scanRevokedAuthorization(tx.QueryRow(ctx, ` +SELECT auth_key_id, user_id, hash, layer, device_model, platform, system_version, + api_id, app_version, ip, password_pending, created_at, active_at +FROM authorizations +WHERE auth_key_id = $1 AND user_id = $2 AND hash = $3 +FOR UPDATE`, candidate, userID, hash)) + if err != nil { + return domain.Authorization{}, false, fmt.Errorf("revalidate revoke authorization by hash: %w", err) + } + if !found { + return domain.Authorization{}, false, nil + } + if err := deleteRevocationTargetsTx(ctx, tx, []int64{candidate}); err != nil { + return domain.Authorization{}, false, err + } + return a, true, nil } func (s *AuthorizationStore) DeleteByUserExcept(ctx context.Context, userID int64, keepAuthKeyID [8]byte) ([]domain.Authorization, error) { @@ -323,57 +368,148 @@ 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) { - rows, err := s.db.Query(ctx, ` -WITH target AS MATERIALIZED ( - SELECT auth_key_id, user_id, hash, layer, device_model, platform, system_version, api_id, app_version, ip, password_pending, created_at, active_at - FROM authorizations - WHERE user_id = $1 AND auth_key_id <> $2 -), deleted_temp AS ( - DELETE FROM auth_keys - WHERE auth_key_id IN ( - SELECT temp_auth_key_id - FROM temp_auth_key_bindings - WHERE perm_auth_key_id IN (SELECT auth_key_id FROM target) - ) - RETURNING auth_key_id -), deleted_update_states AS ( - DELETE FROM update_states - WHERE auth_key_id IN (SELECT auth_key_id FROM target) - RETURNING auth_key_id -), deleted_keys AS ( - DELETE FROM auth_keys - WHERE auth_key_id IN (SELECT auth_key_id FROM target) - RETURNING auth_key_id -), touched AS ( - SELECT - (SELECT count(*) FROM deleted_temp) + - (SELECT count(*) FROM deleted_update_states) AS count -) -SELECT target.auth_key_id, target.user_id, target.hash, target.layer, target.device_model, target.platform, - target.system_version, target.api_id, target.app_version, target.ip, target.password_pending, - target.created_at, target.active_at -FROM target -JOIN deleted_keys USING (auth_key_id) -CROSS JOIN touched -ORDER BY target.created_at, target.auth_key_id`, userID, authKeyIDToInt64(keepAuthKeyID)) - if err != nil { - return nil, fmt.Errorf("revoke authorizations by user: %w", err) + 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") +} + +func revokeByUserExceptTx(ctx context.Context, tx pgx.Tx, userID, keepAuthKeyID int64) ([]domain.Authorization, error) { + candidateRows, err := tx.Query(ctx, ` +SELECT auth_key_id +FROM authorizations +WHERE user_id = $1 AND auth_key_id <> $2 +ORDER BY auth_key_id`, userID, keepAuthKeyID) + if err != nil { + return nil, fmt.Errorf("select revoke candidates by user: %w", err) + } + candidates := make([]int64, 0) + for candidateRows.Next() { + var id int64 + if err := candidateRows.Scan(&id); err != nil { + candidateRows.Close() + return nil, fmt.Errorf("scan revoke candidate by user: %w", err) + } + candidates = append(candidates, id) + } + if err := candidateRows.Err(); err != nil { + candidateRows.Close() + return nil, fmt.Errorf("iterate revoke candidates by user: %w", err) + } + candidateRows.Close() + if len(candidates) == 0 { + return []domain.Authorization{}, nil + } + + // Stable parent-row lock order matches every concurrent batch revocation and + // serializes each candidate with Bind's auth_keys-first ownership change. + 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`, candidates) + if err != nil { + return nil, fmt.Errorf("lock revoke auth keys by user: %w", err) + } + for lockRows.Next() { + var ignored int64 + if err := lockRows.Scan(&ignored); err != nil { + lockRows.Close() + return nil, fmt.Errorf("scan locked revoke auth key: %w", err) + } + } + if err := lockRows.Err(); err != nil { + lockRows.Close() + return nil, fmt.Errorf("iterate locked revoke auth keys: %w", err) + } + lockRows.Close() + + // This is intentionally a new statement snapshot after all parent locks. + // Keys that changed owner while waiting are omitted and must remain intact. + rows, err := tx.Query(ctx, ` +SELECT auth_key_id, user_id, hash, layer, device_model, platform, system_version, + api_id, app_version, ip, password_pending, created_at, active_at +FROM authorizations +WHERE user_id = $1 + AND auth_key_id <> $2 + AND auth_key_id = ANY($3::bigint[]) +ORDER BY created_at, auth_key_id +FOR UPDATE`, userID, keepAuthKeyID, candidates) + if err != nil { + return nil, fmt.Errorf("revalidate revoke authorizations by user: %w", err) } - defer rows.Close() out := make([]domain.Authorization, 0) for rows.Next() { a, err := scanRevokedAuthorizationRow(rows) if err != nil { + rows.Close() return nil, err } out = append(out, a) } if err := rows.Err(); err != nil { + rows.Close() return nil, fmt.Errorf("iterate revoked authorizations: %w", err) } + rows.Close() + if len(out) == 0 { + return out, nil + } + targets := make([]int64, len(out)) + for i := range out { + targets[i] = authKeyIDToInt64(out[i].AuthKeyID) + } + if err := deleteRevocationTargetsTx(ctx, tx, targets); err != nil { + return nil, err + } return out, nil } +func deleteRevocationTargetsTx(ctx context.Context, tx pgx.Tx, authKeyIDs []int64) error { + if len(authKeyIDs) == 0 { + return nil + } + if _, err := tx.Exec(ctx, ` +DELETE FROM auth_keys +WHERE auth_key_id IN ( + SELECT temp_auth_key_id + FROM temp_auth_key_bindings + WHERE perm_auth_key_id = ANY($1::bigint[]) +)`, authKeyIDs); err != nil { + return fmt.Errorf("delete revoked temporary auth keys: %w", err) + } + if _, err := tx.Exec(ctx, ` +DELETE FROM update_states +WHERE auth_key_id = ANY($1::bigint[])`, authKeyIDs); err != nil { + return fmt.Errorf("delete revoked update states: %w", err) + } + tag, err := tx.Exec(ctx, ` +DELETE FROM auth_keys +WHERE auth_key_id = ANY($1::bigint[])`, authKeyIDs) + if err != nil { + return fmt.Errorf("delete revoked permanent auth keys: %w", err) + } + if tag.RowsAffected() != int64(len(authKeyIDs)) { + return fmt.Errorf("delete revoked permanent auth keys: deleted %d of %d locked targets", tag.RowsAffected(), len(authKeyIDs)) + } + return nil +} + func authorizationFromRow(row sqlcgen.Authorization) domain.Authorization { return domain.Authorization{ AuthKeyID: authKeyIDFromInt64(row.AuthKeyID), diff --git a/internal/store/postgres/authorization_revoke_integration_test.go b/internal/store/postgres/authorization_revoke_integration_test.go index d2d2708a..eadb6970 100644 --- a/internal/store/postgres/authorization_revoke_integration_test.go +++ b/internal/store/postgres/authorization_revoke_integration_test.go @@ -2,6 +2,7 @@ package postgres import ( "context" + "errors" "fmt" "testing" "time" @@ -20,9 +21,10 @@ func TestAuthorizationStoreRevokeByHashDeletesProtocolKeyCascadePostgres(t *test temp := revokeTestAuthKeyID(0x92) keys := NewAuthKeyStore(pool) auths := NewAuthorizationStore(pool) + tempExpiry := int(time.Now().Add(time.Hour).Unix()) - saveRevokeTestAuthKey(t, ctx, keys, perm) - saveRevokeTestAuthKey(t, ctx, keys, temp) + saveRevokeTestAuthKey(t, ctx, keys, perm, 0) + saveRevokeTestAuthKey(t, ctx, keys, temp, tempExpiry) if err := auths.Bind(ctx, domain.Authorization{ AuthKeyID: perm, UserID: userID, @@ -46,7 +48,7 @@ func TestAuthorizationStoreRevokeByHashDeletesProtocolKeyCascadePostgres(t *test PermAuthKeyID: authKeyIDToInt64(perm), Nonce: 1, TempSessionID: 2, - ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + ExpiresAt: tempExpiry, EncryptedMessage: []byte{1, 2, 3, 4}, }); err != nil { t.Fatalf("save temp binding: %v", err) @@ -77,20 +79,20 @@ func TestAuthorizationStoreRevokeByUserExceptDeletesOnlyRevokedKeysPostgres(t *t revokedTwo := revokeTestAuthKeyID(0xa3) tempForTwo := revokeTestAuthKeyID(0xa4) - for i, key := range [][8]byte{keep, revokedOne, revokedTwo, tempForTwo} { - saveRevokeTestAuthKey(t, ctx, keys, key) - if i < 3 { - if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: key, UserID: userID, Hash: int64(9100 + i)}); err != nil { - t.Fatalf("bind auth %x: %v", key, err) - } + for i, key := range [][8]byte{keep, revokedOne, revokedTwo} { + saveRevokeTestAuthKey(t, ctx, keys, key, 0) + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: key, UserID: userID, Hash: int64(9100 + i)}); err != nil { + t.Fatalf("bind auth %x: %v", key, err) } } + tempExpiry := int(time.Now().Add(time.Hour).Unix()) + saveRevokeTestAuthKey(t, ctx, keys, tempForTwo, tempExpiry) if err := NewTempAuthKeyBindingStore(pool).Save(ctx, domain.TempAuthKeyBinding{ TempAuthKeyID: tempForTwo, PermAuthKeyID: authKeyIDToInt64(revokedTwo), Nonce: 3, TempSessionID: 4, - ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + ExpiresAt: tempExpiry, EncryptedMessage: []byte{5, 6, 7, 8}, }); err != nil { t.Fatalf("save temp binding: %v", err) @@ -117,7 +119,7 @@ func TestAuthorizationStoreUpdateClientInfoMergesPostgres(t *testing.T) { id := revokeTestAuthKeyID(0xb1) keys := NewAuthKeyStore(pool) auths := NewAuthorizationStore(pool) - saveRevokeTestAuthKey(t, ctx, keys, id) + saveRevokeTestAuthKey(t, ctx, keys, id, 0) if err := auths.Bind(ctx, domain.Authorization{ AuthKeyID: id, UserID: userID, @@ -153,6 +155,382 @@ func TestAuthorizationStoreUpdateClientInfoMergesPostgres(t *testing.T) { } } +func TestAuthorizationStoreRevokeByHashConcurrentTempBindLeavesNoDanglingStatePostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + userID := createRevokeTestUser(t, ctx, pool, "bind-revoke-race") + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + + for attempt := 0; attempt < 24; attempt++ { + tempExpiry := int(time.Now().Add(time.Hour).Unix()) + attempt + perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, tempExpiry) + hash := int64(9300 + attempt) + if err := auths.Bind(ctx, domain.Authorization{ + AuthKeyID: perm, + UserID: userID, + Hash: hash, + }); err != nil { + t.Fatalf("attempt %d bind authorization: %v", attempt, err) + } + candidate := domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(perm), + Nonce: int64(700 + attempt), + TempSessionID: int64(800 + attempt), + ExpiresAt: tempExpiry, + EncryptedMessage: []byte("bind-revoke race"), + } + + start := make(chan struct{}) + bindResult := make(chan error, 1) + type revokeResult struct { + found bool + err error + } + revokeResults := make(chan revokeResult, 1) + go func() { + <-start + bindResult <- bindings.Save(ctx, candidate) + }() + go func() { + <-start + _, found, err := auths.RevokeByHash(ctx, userID, hash) + revokeResults <- revokeResult{found: found, err: err} + }() + close(start) + + bindErr := <-bindResult + if bindErr != nil && !errors.Is(bindErr, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("attempt %d bind/revoke race bind error = %v", attempt, bindErr) + } + revoked := <-revokeResults + if revoked.err != nil || !revoked.found { + t.Fatalf("attempt %d bind/revoke race found=%v err=%v", attempt, revoked.found, revoked.err) + } + if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || found { + t.Fatalf("attempt %d dangling binding found=%v err=%v", attempt, found, err) + } + assertRevokeTestMissingAuthKey(t, ctx, keys, perm) + assertRevokeTestNoAuthorization(t, ctx, auths, perm) + if bindErr == nil { + assertRevokeTestMissingAuthKey(t, ctx, keys, temp) + } else { + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, tempExpiry) + assertRevokeTestNoAuthorization(t, ctx, auths, temp) + if err := keys.Delete(ctx, temp); err != nil { + t.Fatalf("attempt %d clean unbound loser temp: %v", attempt, err) + } + assertRevokeTestMissingAuthKey(t, ctx, keys, temp) + } + } +} + +func TestAuthorizationStoreRevokeByHashSkipsKeyTransferredAfterCandidateReadPostgres(t *testing.T) { + pool := testPool(t) + testCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + states := NewUpdateStateStore(pool) + userA := createRevokeTestUser(t, testCtx, pool, "hash-owner-a") + userB := createRevokeTestUser(t, testCtx, pool, "hash-owner-b") + + tempExpiry := int(time.Now().Add(time.Hour).Unix()) + perm := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, 0) + temp := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, tempExpiry) + const ( + hashA = int64(9501) + hashB = int64(9502) + ) + if err := auths.Bind(testCtx, domain.Authorization{ + AuthKeyID: perm, + UserID: userA, + Hash: hashA, + }); err != nil { + t.Fatalf("bind original A authorization: %v", err) + } + binding := domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(perm), + Nonce: 951, + TempSessionID: 952, + ExpiresAt: tempExpiry, + EncryptedMessage: []byte("owner-transfer binding"), + } + if err := bindings.Save(testCtx, binding); err != nil { + t.Fatalf("save temp binding before owner transfer: %v", err) + } + + // Bind B performs the auth_keys-first ownership change inside an open + // transaction. Its uncommitted row is invisible to A's candidate lookup, but + // the parent FOR UPDATE lock is the deterministic barrier for revocation. + bindB, err := pool.Begin(testCtx) + if err != nil { + t.Fatalf("begin B bind transaction: %v", err) + } + defer func() { _ = bindB.Rollback(context.Background()) }() + wantB := domain.Authorization{ + AuthKeyID: perm, + UserID: userB, + Hash: hashB, + Layer: 227, + DeviceModel: "owner-b-device", + Platform: "android", + SystemVersion: "test", + APIID: 100, + AppVersion: "owner-transfer", + IP: "127.0.0.2", + PasswordPending: true, + } + if err := bindAuthorization(testCtx, bindB, wantB); err != nil { + t.Fatalf("stage B ownership transfer: %v", err) + } + wantBStored, found, err := NewAuthorizationStore(bindB).ByAuthKey(testCtx, perm) + if err != nil || !found { + t.Fatalf("read staged B authorization found=%v err=%v", found, err) + } + wantBState, found, err := NewUpdateStateStore(bindB).Get(testCtx, perm, userB) + if err != nil || !found { + t.Fatalf("read staged B update state found=%v err=%v", found, err) + } + + revokeConn, err := pool.Acquire(testCtx) + if err != nil { + t.Fatalf("acquire dedicated revoke connection: %v", err) + } + t.Cleanup(revokeConn.Release) + var revokePID int + if err := revokeConn.QueryRow(testCtx, "SELECT pg_backend_pid()").Scan(&revokePID); err != nil { + t.Fatalf("get revoke backend pid: %v", err) + } + + type revokeResult struct { + a domain.Authorization + found bool + err error + } + revokeResults := make(chan revokeResult, 1) + go func() { + a, found, err := NewAuthorizationStore(revokeConn).RevokeByHash(testCtx, userA, hashA) + revokeResults <- revokeResult{a: a, found: found, err: err} + }() + waitForPostgresBackendLockWait(t, testCtx, pool, revokePID) + + // Lock wait proves A's revoke already read its candidate and is now serialized + // behind Bind B. After B commits, the fresh owner/hash revalidation must omit + // the key instead of deleting B through A's stale candidate. + if err := bindB.Commit(testCtx); err != nil { + t.Fatalf("commit B ownership transfer: %v", err) + } + + var revoked revokeResult + select { + case revoked = <-revokeResults: + case <-testCtx.Done(): + t.Fatalf("revoke did not finish after releasing FK barrier: %v", testCtx.Err()) + } + if revoked.err != nil || revoked.found { + t.Fatalf("stale A revoke after B transfer found=%v err=%v, want not found", revoked.found, revoked.err) + } + if revoked.a != (domain.Authorization{}) { + t.Fatalf("stale A revoke returned authorization %+v, want zero", revoked.a) + } + + assertRevokeTestPresentAuthKey(t, testCtx, keys, perm) + assertRevokeTestPresentAuthKey(t, testCtx, keys, temp) + assertTempIdentityBinding(t, testCtx, bindings, binding) + gotB, found, err := auths.ByAuthKey(testCtx, perm) + if err != nil || !found || gotB != wantBStored { + t.Fatalf("B authorization after stale A revoke = %+v found=%v err=%v, want %+v", gotB, found, err, wantBStored) + } + gotBState, found, err := states.Get(testCtx, perm, userB) + if err != nil || !found || gotBState != wantBState { + t.Fatalf("B update state after stale A revoke = %+v found=%v err=%v, want %+v", gotBState, found, err, wantBState) + } + if _, found, err := states.Get(testCtx, perm, userA); err != nil || found { + t.Fatalf("stale A update state found=%v err=%v, want absent after B bind", found, err) + } +} + +func TestAuthorizationStoreRevokeByUserExceptPartiallySkipsTransferredCandidatePostgres(t *testing.T) { + pool := testPool(t) + testCtx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + t.Cleanup(cancel) + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + states := NewUpdateStateStore(pool) + userA := createRevokeTestUser(t, testCtx, pool, "bulk-owner-a") + userB := createRevokeTestUser(t, testCtx, pool, "bulk-owner-b") + + keep := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, 0) + transferred := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, 0) + revoked := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, 0) + transferredExpiry := int(time.Now().Add(time.Hour).Unix()) + revokedExpiry := transferredExpiry + 1 + transferredTemp := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, transferredExpiry) + revokedTemp := saveTempIdentityTestAuthKey(t, testCtx, pool, keys, revokedExpiry) + + authorizationsA := []domain.Authorization{ + {AuthKeyID: keep, UserID: userA, Hash: 9601, DeviceModel: "keep-a"}, + {AuthKeyID: transferred, UserID: userA, Hash: 9602, DeviceModel: "transfer-from-a"}, + {AuthKeyID: revoked, UserID: userA, Hash: 9603, DeviceModel: "revoke-a", PasswordPending: true}, + } + for _, authorization := range authorizationsA { + if err := auths.Bind(testCtx, authorization); err != nil { + t.Fatalf("bind A authorization %x: %v", authorization.AuthKeyID, err) + } + } + transferredBinding := domain.TempAuthKeyBinding{ + TempAuthKeyID: transferredTemp, + PermAuthKeyID: authKeyIDToInt64(transferred), + Nonce: 961, + TempSessionID: 962, + ExpiresAt: transferredExpiry, + EncryptedMessage: []byte("transferred candidate binding"), + } + revokedBinding := domain.TempAuthKeyBinding{ + TempAuthKeyID: revokedTemp, + PermAuthKeyID: authKeyIDToInt64(revoked), + Nonce: 963, + TempSessionID: 964, + ExpiresAt: revokedExpiry, + EncryptedMessage: []byte("revoked candidate binding"), + } + for _, binding := range []domain.TempAuthKeyBinding{transferredBinding, revokedBinding} { + if err := bindings.Save(testCtx, binding); err != nil { + t.Fatalf("save candidate binding for perm %d: %v", binding.PermAuthKeyID, err) + } + } + + wantKeepKey, found, err := keys.Get(testCtx, keep) + if err != nil || !found { + t.Fatalf("read keep key before revoke found=%v err=%v", found, err) + } + wantKeepAuth, found, err := auths.ByAuthKey(testCtx, keep) + if err != nil || !found { + t.Fatalf("read keep authorization before revoke found=%v err=%v", found, err) + } + wantKeepState, found, err := states.Get(testCtx, keep, userA) + if err != nil || !found { + t.Fatalf("read keep state before revoke found=%v err=%v", found, err) + } + wantRevokedAuth, found, err := auths.ByAuthKey(testCtx, revoked) + if err != nil || !found { + t.Fatalf("read revocable authorization before revoke found=%v err=%v", found, err) + } + + bindB, err := pool.Begin(testCtx) + if err != nil { + t.Fatalf("begin partial owner-transfer transaction: %v", err) + } + defer func() { _ = bindB.Rollback(context.Background()) }() + wantB := domain.Authorization{ + AuthKeyID: transferred, + UserID: userB, + Hash: 9604, + Layer: 227, + DeviceModel: "bulk-owner-b", + Platform: "android", + SystemVersion: "test", + APIID: 100, + AppVersion: "partial-owner-transfer", + IP: "127.0.0.3", + } + if err := bindAuthorization(testCtx, bindB, wantB); err != nil { + t.Fatalf("stage partial B ownership transfer: %v", err) + } + wantBStored, found, err := NewAuthorizationStore(bindB).ByAuthKey(testCtx, transferred) + if err != nil || !found { + t.Fatalf("read staged partial B authorization found=%v err=%v", found, err) + } + wantBState, found, err := NewUpdateStateStore(bindB).Get(testCtx, transferred, userB) + if err != nil || !found { + t.Fatalf("read staged partial B state found=%v err=%v", found, err) + } + + revokeConn, err := pool.Acquire(testCtx) + if err != nil { + t.Fatalf("acquire dedicated bulk revoke connection: %v", err) + } + t.Cleanup(revokeConn.Release) + var revokePID int + if err := revokeConn.QueryRow(testCtx, "SELECT pg_backend_pid()").Scan(&revokePID); err != nil { + t.Fatalf("get bulk revoke backend pid: %v", err) + } + type bulkRevokeResult struct { + deleted []domain.Authorization + err error + } + revokeResults := make(chan bulkRevokeResult, 1) + go func() { + deleted, err := NewAuthorizationStore(revokeConn).RevokeByUserExcept(testCtx, userA, keep) + revokeResults <- bulkRevokeResult{deleted: deleted, err: err} + }() + waitForPostgresBackendLockWait(t, testCtx, pool, revokePID) + + // The bulk candidate list now contains both old A keys. Releasing B's parent + // lock forces a fresh owner revalidation: transferred must be omitted while + // the unrelated candidate that still belongs to A remains revocable. + if err := bindB.Commit(testCtx); err != nil { + t.Fatalf("commit partial B ownership transfer: %v", err) + } + var result bulkRevokeResult + select { + case result = <-revokeResults: + case <-testCtx.Done(): + t.Fatalf("bulk revoke did not finish after owner transfer: %v", testCtx.Err()) + } + if result.err != nil { + t.Fatalf("bulk revoke after partial owner transfer: %v", result.err) + } + if len(result.deleted) != 1 || result.deleted[0] != wantRevokedAuth { + t.Fatalf("bulk revoked authorizations = %+v, want only %+v", result.deleted, wantRevokedAuth) + } + + gotKeepKey, found, err := keys.Get(testCtx, keep) + if err != nil || !found || gotKeepKey != wantKeepKey { + t.Fatalf("keep key after bulk revoke = %+v found=%v err=%v, want unchanged", gotKeepKey, found, err) + } + gotKeepAuth, found, err := auths.ByAuthKey(testCtx, keep) + if err != nil || !found || gotKeepAuth != wantKeepAuth { + t.Fatalf("keep authorization after bulk revoke = %+v found=%v err=%v, want %+v", gotKeepAuth, found, err, wantKeepAuth) + } + gotKeepState, found, err := states.Get(testCtx, keep, userA) + if err != nil || !found || gotKeepState != wantKeepState { + t.Fatalf("keep state after bulk revoke = %+v found=%v err=%v, want %+v", gotKeepState, found, err, wantKeepState) + } + + assertRevokeTestPresentAuthKey(t, testCtx, keys, transferred) + assertRevokeTestPresentAuthKey(t, testCtx, keys, transferredTemp) + assertTempIdentityBinding(t, testCtx, bindings, transferredBinding) + gotB, found, err := auths.ByAuthKey(testCtx, transferred) + if err != nil || !found || gotB != wantBStored { + t.Fatalf("transferred B authorization = %+v found=%v err=%v, want %+v", gotB, found, err, wantBStored) + } + gotBState, found, err := states.Get(testCtx, transferred, userB) + if err != nil || !found || gotBState != wantBState { + t.Fatalf("transferred B state = %+v found=%v err=%v, want %+v", gotBState, found, err, wantBState) + } + if _, found, err := states.Get(testCtx, transferred, userA); err != nil || found { + t.Fatalf("old A state for transferred key found=%v err=%v, want absent", found, err) + } + + assertRevokeTestMissingAuthKey(t, testCtx, keys, revoked) + assertRevokeTestMissingAuthKey(t, testCtx, keys, revokedTemp) + assertRevokeTestNoAuthorization(t, testCtx, auths, revoked) + if _, found, err := bindings.GetByTemp(testCtx, revokedTemp); err != nil || found { + t.Fatalf("revoked temp binding found=%v err=%v, want absent", found, err) + } + if _, found, err := states.Get(testCtx, revoked, userA); err != nil || found { + t.Fatalf("revoked A state found=%v err=%v, want absent", found, err) + } +} + func createRevokeTestUser(t *testing.T, ctx context.Context, db *pgxpool.Pool, suffix string) int64 { t.Helper() phone := fmt.Sprintf("+1555%09d", time.Now().UnixNano()%1_000_000_000) @@ -173,9 +551,9 @@ func revokeTestAuthKeyID(seed byte) [8]byte { return [8]byte{seed, seed, seed, seed, seed, seed, seed, seed} } -func saveRevokeTestAuthKey(t *testing.T, ctx context.Context, keys store.AuthKeyStore, id [8]byte) { +func saveRevokeTestAuthKey(t *testing.T, ctx context.Context, keys store.AuthKeyStore, id [8]byte, expiresAt int) { t.Helper() - if err := keys.Save(ctx, store.AuthKeyData{ID: id, ServerSalt: int64(id[0])}); err != nil { + if err := keys.Save(ctx, store.AuthKeyData{ID: id, ServerSalt: int64(id[0]), ExpiresAt: expiresAt}); err != nil { t.Fatalf("save auth key %x: %v", id, err) } t.Cleanup(func() { diff --git a/internal/store/postgres/bot_integration_test.go b/internal/store/postgres/bot_integration_test.go index ba01343c..cd3c5c4a 100644 --- a/internal/store/postgres/bot_integration_test.go +++ b/internal/store/postgres/bot_integration_test.go @@ -205,7 +205,7 @@ func TestBotStoreRoundTripPostgres(t *testing.T) { appauth.WithBotLogin(bots)) var authKeyID [8]byte copy(authKeyID[:], fmt.Sprintf("%08d", suffix%100000000)) - if _, err := pool.Exec(ctx, "INSERT INTO auth_keys (auth_key_id, body, server_salt) VALUES ($1, $2, 0) ON CONFLICT DO NOTHING", + if _, err := pool.Exec(ctx, "INSERT INTO auth_keys (auth_key_id, body, server_salt, expires_at) VALUES ($1, $2, 0, 0) ON CONFLICT DO NOTHING", authKeyIDToInt64(authKeyID), make([]byte, 256)); err != nil { t.Fatalf("seed auth key: %v", err) } diff --git a/internal/store/postgres/business_integration_test.go b/internal/store/postgres/business_integration_test.go index ef156ee9..489f0317 100644 --- a/internal/store/postgres/business_integration_test.go +++ b/internal/store/postgres/business_integration_test.go @@ -8,6 +8,7 @@ import ( "fmt" "strings" "testing" + "time" "telesrv/internal/domain" "telesrv/internal/store" @@ -26,15 +27,32 @@ func TestBusinessStoresRoundTrip(t *testing.T) { if _, err := rand.Read(authBody[:]); err != nil { t.Fatal(err) } - if err := NewAuthKeyStore(pool).Save(ctx, store.AuthKeyData{ + authExpiry := int(time.Now().Add(time.Hour).Unix()) + keys := NewAuthKeyStore(pool) + if err := keys.Save(ctx, store.AuthKeyData{ ID: authID, Value: authBody, ServerSalt: 42, + ExpiresAt: authExpiry, }); err != nil { t.Fatalf("save auth key: %v", err) } t.Cleanup(func() { - _, _ = pool.Exec(ctx, "DELETE FROM auth_keys WHERE auth_key_id = $1", authKeyIDToInt64(authID)) + _ = keys.Delete(ctx, authID) + }) + var permAuthID [8]byte + var permAuthBody [256]byte + if _, err := rand.Read(permAuthID[:]); err != nil { + t.Fatal(err) + } + if _, err := rand.Read(permAuthBody[:]); err != nil { + t.Fatal(err) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: permAuthID, Value: permAuthBody}); err != nil { + t.Fatalf("save permanent auth key: %v", err) + } + t.Cleanup(func() { + _ = keys.Delete(ctx, permAuthID) }) users := NewUserStore(pool) @@ -164,10 +182,10 @@ func TestBusinessStoresRoundTrip(t *testing.T) { if err := NewTempAuthKeyBindingStore(pool).Save(ctx, domain.TempAuthKeyBinding{ TempAuthKeyID: authID, - PermAuthKeyID: 12345, + PermAuthKeyID: authKeyIDToInt64(permAuthID), Nonce: 67890, TempSessionID: 24680, - ExpiresAt: 111, + ExpiresAt: authExpiry, EncryptedMessage: []byte("binding"), }); err != nil { t.Fatalf("save temp auth key binding: %v", err) diff --git a/internal/store/postgres/queries/authkey.sql b/internal/store/postgres/queries/authkey.sql deleted file mode 100644 index 2d8dbe4a..00000000 --- a/internal/store/postgres/queries/authkey.sql +++ /dev/null @@ -1,10 +0,0 @@ --- name: GetAuthKey :one -SELECT auth_key_id, body, server_salt, created_at -FROM auth_keys -WHERE auth_key_id = $1; - --- name: UpsertAuthKey :exec -INSERT INTO auth_keys (auth_key_id, body, server_salt) -VALUES ($1, $2, $3) -ON CONFLICT (auth_key_id) DO UPDATE -SET body = EXCLUDED.body, server_salt = EXCLUDED.server_salt; diff --git a/internal/store/postgres/queries/temp_auth_key.sql b/internal/store/postgres/queries/temp_auth_key.sql index 8820c19a..4278da84 100644 --- a/internal/store/postgres/queries/temp_auth_key.sql +++ b/internal/store/postgres/queries/temp_auth_key.sql @@ -1,15 +1,21 @@ --- name: UpsertTempAuthKeyBinding :exec +-- name: UpsertTempAuthKeyBinding :execrows INSERT INTO temp_auth_key_bindings ( temp_auth_key_id, perm_auth_key_id, nonce, temp_session_id, expires_at, encrypted_message ) -VALUES ($1, $2, $3, $4, $5, $6) +SELECT $1, $2, $3, $4, $5, $6 +FROM auth_keys AS temp_key +JOIN auth_keys AS perm_key ON perm_key.auth_key_id = $2 +WHERE temp_key.auth_key_id = $1 + AND temp_key.expires_at = $5 + AND temp_key.expires_at > 0 + AND perm_key.expires_at = 0 ON CONFLICT (temp_auth_key_id) DO UPDATE SET - perm_auth_key_id = EXCLUDED.perm_auth_key_id, nonce = EXCLUDED.nonce, temp_session_id = EXCLUDED.temp_session_id, expires_at = EXCLUDED.expires_at, encrypted_message = EXCLUDED.encrypted_message, - created_at = now(); + created_at = now() +WHERE temp_auth_key_bindings.perm_auth_key_id = EXCLUDED.perm_auth_key_id; -- name: GetTempAuthKeyBinding :one SELECT @@ -23,10 +29,15 @@ FROM temp_auth_key_bindings WHERE temp_auth_key_id = $1; -- name: DeleteExpiredTempAuthKeys :execrows -DELETE FROM auth_keys -WHERE auth_key_id IN ( - SELECT temp_auth_key_id - FROM temp_auth_key_bindings - WHERE expires_at < $1 +WITH candidates AS ( + SELECT candidate_key.auth_key_id + FROM auth_keys AS candidate_key + WHERE candidate_key.expires_at > 0 + AND candidate_key.expires_at < $1 + ORDER BY candidate_key.expires_at, candidate_key.auth_key_id LIMIT $2 -); + FOR UPDATE SKIP LOCKED +) +DELETE FROM auth_keys AS k +USING candidates AS c +WHERE k.auth_key_id = c.auth_key_id; diff --git a/internal/store/postgres/sqlcgen/authkey.sql.go b/internal/store/postgres/sqlcgen/authkey.sql.go deleted file mode 100644 index 53186025..00000000 --- a/internal/store/postgres/sqlcgen/authkey.sql.go +++ /dev/null @@ -1,55 +0,0 @@ -// Code generated by sqlc. DO NOT EDIT. -// versions: -// sqlc v1.31.1 -// source: authkey.sql - -package sqlcgen - -import ( - "context" - - "github.com/jackc/pgx/v5/pgtype" -) - -const getAuthKey = `-- name: GetAuthKey :one -SELECT auth_key_id, body, server_salt, created_at -FROM auth_keys -WHERE auth_key_id = $1 -` - -type GetAuthKeyRow struct { - AuthKeyID int64 - Body []byte - ServerSalt int64 - CreatedAt pgtype.Timestamptz -} - -func (q *Queries) GetAuthKey(ctx context.Context, authKeyID int64) (GetAuthKeyRow, error) { - row := q.db.QueryRow(ctx, getAuthKey, authKeyID) - var i GetAuthKeyRow - err := row.Scan( - &i.AuthKeyID, - &i.Body, - &i.ServerSalt, - &i.CreatedAt, - ) - return i, err -} - -const upsertAuthKey = `-- name: UpsertAuthKey :exec -INSERT INTO auth_keys (auth_key_id, body, server_salt) -VALUES ($1, $2, $3) -ON CONFLICT (auth_key_id) DO UPDATE -SET body = EXCLUDED.body, server_salt = EXCLUDED.server_salt -` - -type UpsertAuthKeyParams struct { - AuthKeyID int64 - Body []byte - ServerSalt int64 -} - -func (q *Queries) UpsertAuthKey(ctx context.Context, arg UpsertAuthKeyParams) error { - _, err := q.db.Exec(ctx, upsertAuthKey, arg.AuthKeyID, arg.Body, arg.ServerSalt) - return err -} diff --git a/internal/store/postgres/sqlcgen/models.go b/internal/store/postgres/sqlcgen/models.go index 71e73d9d..8df00bd6 100644 --- a/internal/store/postgres/sqlcgen/models.go +++ b/internal/store/postgres/sqlcgen/models.go @@ -175,6 +175,7 @@ type AuthKey struct { ApiID int32 AppVersion string LastUsedAt pgtype.Timestamptz + ExpiresAt int32 } type Authorization struct { diff --git a/internal/store/postgres/sqlcgen/temp_auth_key.sql.go b/internal/store/postgres/sqlcgen/temp_auth_key.sql.go index 21c2d33c..940f2ad7 100644 --- a/internal/store/postgres/sqlcgen/temp_auth_key.sql.go +++ b/internal/store/postgres/sqlcgen/temp_auth_key.sql.go @@ -10,13 +10,18 @@ import ( ) const deleteExpiredTempAuthKeys = `-- name: DeleteExpiredTempAuthKeys :execrows -DELETE FROM auth_keys -WHERE auth_key_id IN ( - SELECT temp_auth_key_id - FROM temp_auth_key_bindings - WHERE expires_at < $1 +WITH candidates AS ( + SELECT candidate_key.auth_key_id + FROM auth_keys AS candidate_key + WHERE candidate_key.expires_at > 0 + AND candidate_key.expires_at < $1 + ORDER BY candidate_key.expires_at, candidate_key.auth_key_id LIMIT $2 + FOR UPDATE SKIP LOCKED ) +DELETE FROM auth_keys AS k +USING candidates AS c +WHERE k.auth_key_id = c.auth_key_id ` type DeleteExpiredTempAuthKeysParams struct { @@ -67,18 +72,24 @@ func (q *Queries) GetTempAuthKeyBinding(ctx context.Context, tempAuthKeyID int64 return i, err } -const upsertTempAuthKeyBinding = `-- name: UpsertTempAuthKeyBinding :exec +const upsertTempAuthKeyBinding = `-- name: UpsertTempAuthKeyBinding :execrows INSERT INTO temp_auth_key_bindings ( temp_auth_key_id, perm_auth_key_id, nonce, temp_session_id, expires_at, encrypted_message ) -VALUES ($1, $2, $3, $4, $5, $6) +SELECT $1, $2, $3, $4, $5, $6 +FROM auth_keys AS temp_key +JOIN auth_keys AS perm_key ON perm_key.auth_key_id = $2 +WHERE temp_key.auth_key_id = $1 + AND temp_key.expires_at = $5 + AND temp_key.expires_at > 0 + AND perm_key.expires_at = 0 ON CONFLICT (temp_auth_key_id) DO UPDATE SET - perm_auth_key_id = EXCLUDED.perm_auth_key_id, nonce = EXCLUDED.nonce, temp_session_id = EXCLUDED.temp_session_id, expires_at = EXCLUDED.expires_at, encrypted_message = EXCLUDED.encrypted_message, created_at = now() +WHERE temp_auth_key_bindings.perm_auth_key_id = EXCLUDED.perm_auth_key_id ` type UpsertTempAuthKeyBindingParams struct { @@ -90,8 +101,8 @@ type UpsertTempAuthKeyBindingParams struct { EncryptedMessage []byte } -func (q *Queries) UpsertTempAuthKeyBinding(ctx context.Context, arg UpsertTempAuthKeyBindingParams) error { - _, err := q.db.Exec(ctx, upsertTempAuthKeyBinding, +func (q *Queries) UpsertTempAuthKeyBinding(ctx context.Context, arg UpsertTempAuthKeyBindingParams) (int64, error) { + result, err := q.db.Exec(ctx, upsertTempAuthKeyBinding, arg.TempAuthKeyID, arg.PermAuthKeyID, arg.Nonce, @@ -99,5 +110,8 @@ func (q *Queries) UpsertTempAuthKeyBinding(ctx context.Context, arg UpsertTempAu arg.ExpiresAt, arg.EncryptedMessage, ) - return err + if err != nil { + return 0, err + } + return result.RowsAffected(), nil } diff --git a/internal/store/postgres/temp_auth_key.go b/internal/store/postgres/temp_auth_key.go index ccd8c636..75ace452 100644 --- a/internal/store/postgres/temp_auth_key.go +++ b/internal/store/postgres/temp_auth_key.go @@ -4,10 +4,13 @@ import ( "context" "errors" "fmt" + "math" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "telesrv/internal/domain" + "telesrv/internal/store" "telesrv/internal/store/postgres/sqlcgen" ) @@ -22,25 +25,45 @@ func NewTempAuthKeyBindingStore(db sqlcgen.DBTX) *TempAuthKeyBindingStore { } func (s *TempAuthKeyBindingStore) Save(ctx context.Context, b domain.TempAuthKeyBinding) error { - if err := s.q.UpsertTempAuthKeyBinding(ctx, sqlcgen.UpsertTempAuthKeyBindingParams{ + if b.ExpiresAt <= 0 || int64(b.ExpiresAt) > math.MaxInt32 { + return store.ErrAuthKeyBindingInvalid + } + n, err := s.q.UpsertTempAuthKeyBinding(ctx, sqlcgen.UpsertTempAuthKeyBindingParams{ TempAuthKeyID: authKeyIDToInt64(b.TempAuthKeyID), PermAuthKeyID: b.PermAuthKeyID, Nonce: b.Nonce, TempSessionID: b.TempSessionID, ExpiresAt: int32(b.ExpiresAt), EncryptedMessage: b.EncryptedMessage, - }); err != nil { + }) + if err != nil { + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "23503" { + return store.ErrAuthKeyBindingInvalid + } 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 + } return nil } -// DeleteExpired 实现 store.TempAuthKeyBindingStore：删除 auth_keys 中过期的 temp key， -// temp_auth_key_bindings 经 ON DELETE CASCADE 一并清除，过期 key 的入站帧随之失效。 +// DeleteExpired 实现 store.TempAuthKeyBindingStore：按 auth_keys.expires_at 的部分索引 +// 有界删除所有过期 temp key（含从未绑定的握手 key），binding 经 CASCADE 一并清除。 +// Edge 已在准确协议时刻停止使用 key；这里的 24h 宽限只控制数据库物理回收。 func (s *TempAuthKeyBindingStore) DeleteExpired(ctx context.Context, expiredBefore int64, limit int) (int, error) { if limit <= 0 { return 0, nil } + if expiredBefore <= 0 || expiredBefore > math.MaxInt32 { + return 0, fmt.Errorf("delete expired temp auth keys: invalid expiry cutoff %d", expiredBefore) + } n, err := s.q.DeleteExpiredTempAuthKeys(ctx, sqlcgen.DeleteExpiredTempAuthKeysParams{ ExpiresAt: int32(expiredBefore), Limit: int32(limit), diff --git a/internal/store/postgres/temp_auth_key_identity_integration_test.go b/internal/store/postgres/temp_auth_key_identity_integration_test.go new file mode 100644 index 00000000..79605a87 --- /dev/null +++ b/internal/store/postgres/temp_auth_key_identity_integration_test.go @@ -0,0 +1,554 @@ +package postgres + +import ( + "bytes" + "context" + "crypto/rand" + "errors" + "fmt" + "math" + "strconv" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgxpool" + + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestTempAuthKeyBindingStoreRejectsIntegerWraparoundPostgres(t *testing.T) { + if strconv.IntSize < 64 { + t.Skip("64-bit int required to construct an out-of-int32 expiry") + } + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + handshakeExpiry := int(time.Now().Add(time.Hour).Unix()) + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, handshakeExpiry) + perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + + overflow := domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(perm), + Nonce: 901, + TempSessionID: 902, + ExpiresAt: int(int64(handshakeExpiry) + (int64(1) << 32)), + EncryptedMessage: []byte("wraparound"), + } + if err := bindings.Save(ctx, overflow); !errors.Is(err, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("overflow binding expiry error = %v, want %v", err, store.ErrAuthKeyBindingInvalid) + } + if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || found { + t.Fatalf("binding after overflow found=%v err=%v, want absent", found, err) + } + if _, err := bindings.DeleteExpired(ctx, int64(math.MaxInt32)+1, 1); err == nil { + t.Fatal("overflow retention cutoff succeeded, want explicit rejection") + } +} + +func TestAuthorizationStoreRejectsTemporaryProtocolKeyPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, int(time.Now().Add(time.Hour).Unix())) + phone := fmt.Sprintf("15558%015d", time.Now().UnixNano()) + user, err := NewUserStore(pool).Create(ctx, domain.User{Phone: phone, FirstName: "TempKeyGuard"}) + if err != nil { + t.Fatalf("create user: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", user.ID) + }) + + err = NewAuthorizationStore(pool).Bind(ctx, domain.Authorization{AuthKeyID: temp, UserID: user.ID}) + if !errors.Is(err, store.ErrAuthKeyNotPermanent) { + t.Fatalf("bind authorization to temp key error = %v, want %v", err, store.ErrAuthKeyNotPermanent) + } + if _, found, getErr := NewAuthorizationStore(pool).ByAuthKey(ctx, temp); getErr != nil || found { + t.Fatalf("temporary authorization found=%v err=%v, want absent", found, getErr) + } +} + +func TestTempAuthKeyBindingStorePreservesHandshakeExpiryAndRejectsRebindPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + + handshakeExpiry := int(time.Now().Add(time.Hour).Unix()) + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, handshakeExpiry) + permA := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + permB := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + + first := domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(permA), + Nonce: 101, + TempSessionID: 201, + ExpiresAt: handshakeExpiry, + EncryptedMessage: []byte("first binding"), + } + if err := bindings.Save(ctx, first); err != nil { + t.Fatalf("save first binding: %v", err) + } + assertTempIdentityBinding(t, ctx, bindings, first) + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry) + + // The app service accepts client-specific proof expiry (TDesktop adds 30s) + // but must normalize what it passes to the store. A direct caller cannot + // persist proof metadata whose lifetime differs from the handshake key. + mismatched := first + mismatched.Nonce = 102 + mismatched.TempSessionID = 202 + mismatched.ExpiresAt = handshakeExpiry + 60 + mismatched.EncryptedMessage = []byte("mismatched replay") + if err := bindings.Save(ctx, mismatched); !errors.Is(err, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("mismatched replay error = %v, want %v", err, store.ErrAuthKeyBindingInvalid) + } + assertTempIdentityBinding(t, ctx, bindings, first) + + replayed := first + replayed.Nonce = 103 + replayed.TempSessionID = 203 + replayed.EncryptedMessage = []byte("normalized replay") + if err := bindings.Save(ctx, replayed); err != nil { + t.Fatalf("replay normalized binding: %v", err) + } + assertTempIdentityBinding(t, ctx, bindings, replayed) + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry) + + forbidden := replayed + forbidden.PermAuthKeyID = authKeyIDToInt64(permB) + forbidden.Nonce = 999 + forbidden.ExpiresAt = handshakeExpiry + forbidden.EncryptedMessage = []byte("must not persist") + if err := bindings.Save(ctx, forbidden); !errors.Is(err, store.ErrTempAuthKeyAlreadyBound) { + t.Fatalf("cross-permanent rebind error = %v, want %v", err, store.ErrTempAuthKeyAlreadyBound) + } + assertTempIdentityBinding(t, ctx, bindings, replayed) + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry) +} + +func TestTempAuthKeyBindingStoreConcurrentFirstBindKeepsHandshakeExpiryPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + + // A temp key has exactly one permanent identity even when two valid bind + // proofs race. Repeat with fresh rows so the test covers the contended first + // bind path instead of only the already-bound fast path. + for attempt := 0; attempt < 16; attempt++ { + handshakeExpiry := int(time.Now().Add(30*time.Minute).Unix()) + attempt + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, handshakeExpiry) + permA := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + permB := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + candidates := []domain.TempAuthKeyBinding{ + { + TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(permA), Nonce: 301, + ExpiresAt: handshakeExpiry, EncryptedMessage: []byte("candidate-a"), + }, + { + TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(permB), Nonce: 302, + ExpiresAt: handshakeExpiry, EncryptedMessage: []byte("candidate-b"), + }, + } + + start := make(chan struct{}) + results := make(chan error, len(candidates)) + for _, candidate := range candidates { + candidate := candidate + go func() { + <-start + results <- bindings.Save(ctx, candidate) + }() + } + close(start) + + var success, rejected int + for range candidates { + err := <-results + switch { + case err == nil: + success++ + case errors.Is(err, store.ErrTempAuthKeyAlreadyBound): + rejected++ + default: + t.Fatalf("attempt %d concurrent bind: unexpected error %v", attempt, err) + } + } + if success != 1 || rejected != 1 { + t.Fatalf("attempt %d concurrent bind outcomes: success=%d rejected=%d, want 1/1", attempt, success, rejected) + } + + got, found, err := bindings.GetByTemp(ctx, temp) + if err != nil || !found { + t.Fatalf("attempt %d get winner: found=%v err=%v", attempt, found, err) + } + var winner domain.TempAuthKeyBinding + switch got.PermAuthKeyID { + case candidates[0].PermAuthKeyID: + winner = candidates[0] + case candidates[1].PermAuthKeyID: + winner = candidates[1] + default: + t.Fatalf("attempt %d winner perm auth key = %d, want one of the candidates", attempt, got.PermAuthKeyID) + } + assertTempIdentityBinding(t, ctx, bindings, winner) + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry) + } +} + +func TestTempAuthKeyBindingStoreRejectsMissingPermanentKeyPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + + handshakeExpiry := int(time.Now().Add(time.Hour).Unix()) + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, handshakeExpiry) + missingPerm := randomTempIdentityAuthKeyID(t) + candidate := domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(missingPerm), + Nonce: 401, + TempSessionID: 402, + ExpiresAt: handshakeExpiry, + EncryptedMessage: []byte("missing permanent key"), + } + if err := bindings.Save(ctx, candidate); !errors.Is(err, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("missing permanent key error = %v, want %v", err, store.ErrAuthKeyBindingInvalid) + } + _, rawInsertErr := pool.Exec(ctx, ` +INSERT INTO temp_auth_key_bindings ( + temp_auth_key_id, perm_auth_key_id, nonce, temp_session_id, expires_at, encrypted_message +) VALUES ($1, $2, $3, $4, $5, $6)`, + authKeyIDToInt64(temp), candidate.PermAuthKeyID, candidate.Nonce, + candidate.TempSessionID, candidate.ExpiresAt, candidate.EncryptedMessage, + ) + var pgErr *pgconn.PgError + if !errors.As(rawInsertErr, &pgErr) || pgErr.Code != "23503" || pgErr.ConstraintName != tempAuthKeyPermFKConstraint { + t.Fatalf("raw missing-perm FK error = %v, want 23503/%s", rawInsertErr, tempAuthKeyPermFKConstraint) + } + if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || found { + t.Fatalf("binding with missing permanent key found=%v err=%v, want absent", found, err) + } + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry) +} + +func TestTempAuthKeyBindingConcurrentWithPermanentDeleteLeavesNoDanglingStatePostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + bindings := NewTempAuthKeyBindingStore(pool) + + for attempt := 0; attempt < 32; attempt++ { + handshakeExpiry := int(time.Now().Add(time.Hour).Unix()) + attempt + temp := saveTempIdentityTestAuthKey(t, ctx, pool, keys, handshakeExpiry) + perm := saveTempIdentityTestAuthKey(t, ctx, pool, keys, 0) + candidate := domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(perm), + Nonce: int64(500 + attempt), + TempSessionID: int64(600 + attempt), + ExpiresAt: handshakeExpiry, + EncryptedMessage: []byte("bind-delete race"), + } + + start := make(chan struct{}) + bindResult := make(chan error, 1) + deleteResult := make(chan error, 1) + go func() { + <-start + bindResult <- bindings.Save(ctx, candidate) + }() + go func() { + <-start + deleteResult <- keys.Delete(ctx, perm) + }() + close(start) + + bindErr := <-bindResult + if bindErr != nil && !errors.Is(bindErr, store.ErrAuthKeyBindingInvalid) { + t.Fatalf("attempt %d bind/delete race bind error = %v", attempt, bindErr) + } + if err := <-deleteResult; err != nil { + t.Fatalf("attempt %d bind/delete race delete: %v", attempt, err) + } + if _, found, err := bindings.GetByTemp(ctx, temp); err != nil || found { + t.Fatalf("attempt %d dangling binding found=%v err=%v", attempt, found, err) + } + assertTempIdentityAuthKeyMissing(t, ctx, keys, perm) + if bindErr == nil { + // The binding committed first, so permanent-key deletion must have + // observed it (or retried after the FK race) and deleted the temp key. + assertTempIdentityAuthKeyMissing(t, ctx, keys, temp) + } else { + // Deletion won before the binding existed. The loser remains a valid, + // unbound protocol temp key until its own expiry collector runs; it is + // not allowed to acquire a binding or authorization to the deleted perm. + assertTempIdentityAuthKeyExpiry(t, ctx, keys, temp, handshakeExpiry) + if _, found, err := NewAuthorizationStore(pool).ByAuthKey(ctx, temp); err != nil || found { + t.Fatalf("attempt %d loser temp authorization found=%v err=%v", attempt, found, err) + } + if err := keys.Delete(ctx, temp); err != nil { + t.Fatalf("attempt %d clean unbound loser temp: %v", attempt, err) + } + assertTempIdentityAuthKeyMissing(t, ctx, keys, temp) + } + } +} + +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, + pool *pgxpool.Pool, + backendPID int, +) { + t.Helper() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + var waiting bool + if err := pool.QueryRow(ctx, ` +SELECT EXISTS ( + SELECT 1 + FROM pg_stat_activity AS activity + WHERE activity.pid = $1 + AND activity.state = 'active' + AND activity.wait_event_type = 'Lock' + AND EXISTS ( + SELECT 1 + FROM pg_locks AS waiting_lock + WHERE waiting_lock.pid = activity.pid + AND NOT waiting_lock.granted + ) +)`, backendPID).Scan(&waiting); err != nil { + t.Fatalf("observe delete backend lock wait: %v", err) + } + if waiting { + return + } + select { + case <-ctx.Done(): + t.Fatalf("backend %d never entered a PostgreSQL lock wait: %v", backendPID, ctx.Err()) + case <-ticker.C: + } + } +} + +func saveTempIdentityTestAuthKey( + t *testing.T, + ctx context.Context, + pool *pgxpool.Pool, + keys store.AuthKeyStore, + expiresAt int, +) [8]byte { + t.Helper() + 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) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: id, Value: value, ExpiresAt: expiresAt}); err != nil { + t.Fatalf("save auth key: %v", err) + } + t.Cleanup(func() { + _ = NewAuthKeyStore(pool).Delete(ctx, id) + }) + return id +} + +func randomTempIdentityAuthKeyID(t *testing.T) [8]byte { + t.Helper() + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + t.Fatal(err) + } + return id +} + +func assertTempIdentityBinding( + t *testing.T, + ctx context.Context, + bindings store.TempAuthKeyBindingStore, + want domain.TempAuthKeyBinding, +) { + t.Helper() + got, found, err := bindings.GetByTemp(ctx, want.TempAuthKeyID) + if err != nil { + t.Fatalf("get binding: %v", err) + } + if !found { + t.Fatal("binding not found") + } + if got.TempAuthKeyID != want.TempAuthKeyID || got.PermAuthKeyID != want.PermAuthKeyID || + got.Nonce != want.Nonce || got.TempSessionID != want.TempSessionID || got.ExpiresAt != want.ExpiresAt || + !bytes.Equal(got.EncryptedMessage, want.EncryptedMessage) { + t.Fatalf("binding mismatch: got %+v, want %+v", got, want) + } +} + +func assertTempIdentityAuthKeyExpiry( + t *testing.T, + ctx context.Context, + keys store.AuthKeyStore, + id [8]byte, + want int, +) { + t.Helper() + got, found, err := keys.Get(ctx, id) + if err != nil { + t.Fatalf("get auth key: %v", err) + } + if !found { + t.Fatal("auth key not found") + } + if got.ExpiresAt != want { + t.Fatalf("auth key expires_at = %d, want %d", got.ExpiresAt, want) + } +} + +func assertTempIdentityAuthKeyMissing( + t *testing.T, + ctx context.Context, + keys store.AuthKeyStore, + id [8]byte, +) { + t.Helper() + if _, found, err := keys.Get(ctx, id); err != nil || found { + t.Fatalf("auth key %x found=%v err=%v, want absent", id, found, err) + } +} diff --git a/internal/store/postgres/update_event_retention_integration_test.go b/internal/store/postgres/update_event_retention_integration_test.go index 67d79828..8c7b198d 100644 --- a/internal/store/postgres/update_event_retention_integration_test.go +++ b/internal/store/postgres/update_event_retention_integration_test.go @@ -505,8 +505,8 @@ RETURNING id _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", userIDs) }) if _, err := pool.Exec(ctx, ` -INSERT INTO auth_keys (auth_key_id, body, server_salt) -SELECT id, decode(repeat('00', 256), 'hex'), 0 +INSERT INTO auth_keys (auth_key_id, body, server_salt, expires_at) +SELECT id, decode(repeat('00', 256), 'hex'), 0, 0 FROM unnest($1::bigint[]) AS id`, authKeyIDs); err != nil { t.Fatalf("bulk insert old-tail auth keys: %v", err) } diff --git a/internal/store/temp_auth_key.go b/internal/store/temp_auth_key.go index 14fd14b0..27924812 100644 --- a/internal/store/temp_auth_key.go +++ b/internal/store/temp_auth_key.go @@ -2,17 +2,21 @@ package store import ( "context" + "errors" "telesrv/internal/domain" ) +// ErrTempAuthKeyAlreadyBound 表示同一 temporary key 已绑定到另一个 permanent key。 +// temp key 的 canonical identity 在首次 bind 后不可漂移，重放同一绑定才允许幂等成功。 +var ErrTempAuthKeyAlreadyBound = errors.New("temporary auth key already bound") + // TempAuthKeyBindingStore 持久化 auth.bindTempAuthKey 的 temp→perm 绑定。 type TempAuthKeyBindingStore interface { Save(ctx context.Context, binding domain.TempAuthKeyBinding) error GetByTemp(ctx context.Context, tempAuthKeyID [8]byte) (domain.TempAuthKeyBinding, bool, error) - // DeleteExpired 回收过期早于 expiredBefore（unix 秒）的 temp 绑定，单次最多 limit 条， - // 返回回收数。PFS temp key 定期轮换，无回收时绑定表无界堆积。 - // postgres 实现删除 auth_keys 中的 temp key 行（绑定经 ON DELETE CASCADE 一并清除）， - // 让过期 temp key 的入站帧立即失效；memory 替身仅删绑定。 + // DeleteExpired 以 auth_keys 的握手协议 expiry 为唯一事实源，回收早于 + // expiredBefore（unix 秒）的 temporary key，单次最多 limit 条并返回 key 数。 + // 未绑定与已绑定的 PFS key 必须走同一路径；绑定经 ON DELETE CASCADE 清除。 DeleteExpired(ctx context.Context, expiredBefore int64, limit int) (int, error) }