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

perf: sync protocol and core hardening updates

Browse source
This commit is contained in:
A 2026-07-11 19:48:26 +08:00
parent 152fed3b87
commit 4390ebf5a9
283 changed files with 29231 additions and 2295 deletions
Download patch file Download diff file Expand all files Collapse all files

283
internal/app/account/login_email_test.go
View file

@ -3,6 +3,8 @@ package account
import (
"context"
"errors"
"fmt"
"sync"
"testing"
"time"
@ -32,6 +34,88 @@ type captureMailSender struct {
code string
}
type blockingCodeCAS struct {
store.CodeStore
mu sync.Mutex
blockRevision string
blockUpdate bool
blockDelete bool
entered chan struct{}
release chan struct{}
once sync.Once
}
type switchableEmailOwnerStore struct {
store.UserStore
mu sync.RWMutex
phone string
override bool
owner domain.User
found bool
}
func (s *switchableEmailOwnerStore) ByPhone(ctx context.Context, phone string) (domain.User, bool, error) {
s.mu.RLock()
if s.override && domain.NormalizePhone(phone) == s.phone {
owner, found := s.owner, s.found
s.mu.RUnlock()
return owner, found, nil
}
s.mu.RUnlock()
return s.UserStore.ByPhone(ctx, phone)
}
func (s *switchableEmailOwnerStore) switchOwner(phone string, owner domain.User) {
s.mu.Lock()
s.phone = domain.NormalizePhone(phone)
s.owner = owner
s.found = true
s.override = true
s.mu.Unlock()
}
type afterSavePasswordStore struct {
store.PasswordStore
once sync.Once
afterSave func(userID int64, settings domain.PasswordSettings)
}
func (s *afterSavePasswordStore) Save(ctx context.Context, userID int64, settings domain.PasswordSettings) error {
if err := s.PasswordStore.Save(ctx, userID, settings); err != nil {
return err
}
if s.afterSave != nil {
s.once.Do(func() { s.afterSave(userID, settings) })
}
return nil
}
func (s *blockingCodeCAS) shouldBlock(revision string, update bool) bool {
s.mu.Lock()
defer s.mu.Unlock()
return revision == s.blockRevision && ((update && s.blockUpdate) || (!update && s.blockDelete))
}
func (s *blockingCodeCAS) waitIfBlocked(revision string, update bool) {
if !s.shouldBlock(revision, update) {
return
}
s.once.Do(func() {
close(s.entered)
<-s.release
})
}
func (s *blockingCodeCAS) CompareAndUpdate(ctx context.Context, key, revision string, next store.PhoneCode) (bool, error) {
s.waitIfBlocked(revision, true)
return s.CodeStore.CompareAndUpdate(ctx, key, revision, next)
}
func (s *blockingCodeCAS) CompareAndDelete(ctx context.Context, key, revision string) (bool, error) {
s.waitIfBlocked(revision, false)
return s.CodeStore.CompareAndDelete(ctx, key, revision)
}
func (s *captureMailSender) SendLoginCode(_ context.Context, to, code string, _ time.Duration) error {
s.to = to
s.code = code
@ -66,22 +150,22 @@ func TestSetLoginEmailPersistsAndMasks(t *testing.T) {
}
}
// TestLoginEmailByPhoneAndClear 验证按手机号读取/清除登录邮箱（sendCode 检测 + reset 用）。
// TestLoginEmailByPhoneAndClear 验证 sendCode 可按手机号读取，但 reset 只按已锁定 userID 清除。
func TestLoginEmailByPhoneAndClear(t *testing.T) {
ctx := context.Background()
svc, users := newLoginEmailService(t)
createUser(t, users, "15550010002")
u := createUser(t, users, "15550010002")
if err := svc.SetLoginEmailByPhone(ctx, "+1 555 001 0002", "bob@mail.com"); err != nil {
t.Fatalf("SetLoginEmailByPhone: %v", err)
if err := svc.SetLoginEmail(ctx, u.ID, "bob@mail.com"); err != nil {
t.Fatalf("SetLoginEmail: %v", err)
}
email, found, err := svc.LoginEmailByPhone(ctx, "15550010002")
if err != nil || !found || email != "bob@mail.com" {
t.Fatalf("LoginEmailByPhone = %q found=%v err=%v", email, found, err)
}
if err := svc.ClearLoginEmailByPhone(ctx, "15550010002"); err != nil {
t.Fatalf("ClearLoginEmailByPhone: %v", err)
if err := svc.ClearLoginEmail(ctx, u.ID); err != nil {
t.Fatalf("ClearLoginEmail: %v", err)
}
if _, found, _ := svc.LoginEmailByPhone(ctx, "15550010002"); found {
t.Fatal("login email still present after clear")
@ -182,7 +266,7 @@ func TestLoginEmailSetupRejectsAlreadyOwnedEmailForNewPhone(t *testing.T) {
if err := svc.SetLoginEmail(ctx, owner.ID, "owner@example.test"); err != nil {
t.Fatalf("SetLoginEmail owner: %v", err)
}
if err := codes.Set(ctx, "new-phone-hash", store.PhoneCode{Phone: "15550010108", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
if err := codes.Set(ctx, "new-phone-hash", store.PhoneCode{Version: store.PhoneCodeVersionCurrent, Phone: "15550010108", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
t.Fatalf("seed phone code: %v", err)
}
@ -234,7 +318,7 @@ func TestLoginEmailSetupStoresPendingEmailOnPhoneCodeHash(t *testing.T) {
svc := NewService(memory.NewPasswordStore(),
WithUsers(memory.NewUserStore()),
WithLoginEmailVerification(codes, sender, time.Minute, 2, 6))
if err := codes.Set(ctx, "phone-hash", store.PhoneCode{Phone: "15550010006", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
if err := codes.Set(ctx, "phone-hash", store.PhoneCode{Version: store.PhoneCodeVersionCurrent, Phone: "15550010006", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
t.Fatalf("seed phone code: %v", err)
}
@ -252,7 +336,7 @@ func TestLoginEmailSetupStoresPendingEmailOnPhoneCodeHash(t *testing.T) {
if err != nil || !found {
t.Fatalf("phone code found=%v err=%v", found, err)
}
if rec.Channel != "email_login" || rec.Code != sender.code || rec.Email != "new@example.test" || !rec.VerifiedEmail || rec.PendingEmail != "new@example.test" {
if rec.Channel != "email_login" || rec.Code != sender.code || rec.Email != "new@example.test" || !rec.VerifiedEmail || !rec.SignUpVerified || rec.PendingEmail != "new@example.test" {
t.Fatalf("phone code after verify = %+v", rec)
}
}
@ -291,3 +375,184 @@ func TestVerifyLoginEmailDeletesCodeAfterMaxAttempts(t *testing.T) {
t.Fatal("login email was set after exhausted verification code")
}
}
func TestStaleLoginEmailVerificationCannotMutateResentCode(t *testing.T) {
ctx := context.Background()
for _, tc := range []struct {
name string
blockUpdate bool
blockDelete bool
verificationCode func(old string) string
}{
{
name: "wrong-code-update",
blockUpdate: true,
verificationCode: func(old string) string {
if old != "000000" {
return "000000"
}
return "111111"
},
},
{
name: "correct-code-delete",
blockDelete: true,
verificationCode: func(old string) string { return old },
},
} {
t.Run(tc.name, func(t *testing.T) {
users := memory.NewUserStore()
baseCodes := memory.NewCodeStore()
codes := &blockingCodeCAS{CodeStore: baseCodes}
passwords := memory.NewPasswordStore()
sender := &captureMailSender{}
svc := NewService(passwords,
WithUsers(users),
WithLoginEmailVerification(codes, sender, time.Minute, 3, 6))
u := createUser(t, users, "155500102"+fmt.Sprint(10+len(tc.name)))
if _, _, err := svc.SendLoginEmailCode(ctx, u.ID, "", "", "cas@example.test", false); err != nil {
t.Fatalf("first SendLoginEmailCode: %v", err)
}
oldCode := sender.code
key := loginEmailVerifyChangePrefix + fmt.Sprint(u.ID)
oldSnapshot, found, err := baseCodes.GetSnapshot(ctx, key)
if err != nil || !found {
t.Fatalf("old snapshot found=%v err=%v", found, err)
}
codes.blockRevision = oldSnapshot.Revision
codes.blockUpdate = tc.blockUpdate
codes.blockDelete = tc.blockDelete
codes.entered = make(chan struct{})
codes.release = make(chan struct{})
verifyErr := make(chan error, 1)
go func() {
_, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", tc.verificationCode(oldCode), false)
verifyErr <- err
}()
<-codes.entered
for attempts := 0; attempts < 5; attempts++ {
if _, _, err := svc.SendLoginEmailCode(ctx, u.ID, "", "", "cas@example.test", false); err != nil {
t.Fatalf("resent SendLoginEmailCode: %v", err)
}
if sender.code != oldCode {
break
}
}
newCode := sender.code
if newCode == oldCode {
t.Fatal("random resend repeatedly produced the old code")
}
newSnapshot, found, err := baseCodes.GetSnapshot(ctx, key)
if err != nil || !found || newSnapshot.Revision == oldSnapshot.Revision {
t.Fatalf("new snapshot=%+v found=%v err=%v", newSnapshot, found, err)
}
close(codes.release)
if err := <-verifyErr; !errors.Is(err, domain.ErrEmailCodeInvalid) {
t.Fatalf("stale verification err=%v, want ErrEmailCodeInvalid", err)
}
current, found, err := baseCodes.GetSnapshot(ctx, key)
if err != nil || !found || current.Revision != newSnapshot.Revision || current.Record.Code != newCode || current.Record.Attempts != 0 {
t.Fatalf("current code after stale verifier=%+v found=%v err=%v", current, found, err)
}
if _, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", newCode, false); err != nil {
t.Fatalf("VerifyLoginEmail new code: %v", err)
}
})
}
}
func TestConcurrentWrongLoginEmailCodesNeverAuthorize(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
codes := memory.NewCodeStore()
passwords := memory.NewPasswordStore()
sender := &captureMailSender{}
svc := NewService(passwords,
WithUsers(users),
WithLoginEmailVerification(codes, sender, time.Minute, 2, 6))
u := createUser(t, users, "15550010231")
if _, _, err := svc.SendLoginEmailCode(ctx, u.ID, "", "", "wrong@example.test", false); err != nil {
t.Fatalf("SendLoginEmailCode: %v", err)
}
correct := sender.code
wrong := "000000"
if wrong == correct {
wrong = "111111"
}
const workers = 32
start := make(chan struct{})
errs := make(chan error, workers)
for i := 0; i < workers; i++ {
go func() {
<-start
_, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", wrong, false)
errs <- err
}()
}
close(start)
for i := 0; i < workers; i++ {
if err := <-errs; !errors.Is(err, domain.ErrEmailCodeInvalid) {
t.Fatalf("wrong concurrent verification err=%v", err)
}
}
if _, found, err := svc.LoginEmail(ctx, u.ID); err != nil || found {
t.Fatalf("LoginEmail after wrong codes found=%v err=%v, want absent", found, err)
}
key := loginEmailVerifyChangePrefix + fmt.Sprint(u.ID)
for attempts := 0; attempts < 3; attempts++ {
if _, found, err := codes.GetSnapshot(ctx, key); err != nil {
t.Fatalf("GetSnapshot after concurrent attempts: %v", err)
} else if !found {
break
}
if _, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", wrong, false); !errors.Is(err, domain.ErrEmailCodeInvalid) {
t.Fatalf("final wrong verification err=%v", err)
}
}
if _, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", correct, false); !errors.Is(err, domain.ErrEmailCodeInvalid) {
t.Fatalf("correct code after exhausted attempts err=%v, want invalid", err)
}
}
func TestEmailSetupOwnerTransferDuringSaveNeverWritesFactorToNewOwner(t *testing.T) {
ctx := context.Background()
baseUsers := memory.NewUserStore()
ownerA := createUser(t, baseUsers, "15550010241")
ownerB := createUser(t, baseUsers, "15550010242")
users := &switchableEmailOwnerStore{UserStore: baseUsers}
basePasswords := memory.NewPasswordStore()
passwords := &afterSavePasswordStore{PasswordStore: basePasswords}
codes := memory.NewCodeStore()
sender := &captureMailSender{}
svc := NewService(passwords,
WithUsers(users),
WithLoginEmailVerification(codes, sender, time.Minute, 3, 6))
hash := "owner-save-race"
if err := codes.Set(ctx, hash, store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: ownerA.ID,
Phone: ownerA.Phone,
Channel: codeChannelEmailSetupRequired,
MaxAttempts: 3,
}, time.Minute); err != nil {
t.Fatalf("seed phone code: %v", err)
}
if _, _, err := svc.SendLoginEmailCode(ctx, 0, ownerA.Phone, hash, "owner-a@example.test", true); err != nil {
t.Fatalf("SendLoginEmailCode: %v", err)
}
passwords.afterSave = func(userID int64, settings domain.PasswordSettings) {
if userID == ownerA.ID && settings.LoginEmail == "owner-a@example.test" {
users.switchOwner(ownerA.Phone, ownerB)
}
}
if _, err := svc.VerifyLoginEmail(ctx, 0, ownerA.Phone, hash, sender.code, true); !errors.Is(err, domain.ErrEmailCodeInvalid) {
t.Fatalf("VerifyLoginEmail across save-time owner transfer err=%v, want invalid", err)
}
if settings, found, err := basePasswords.GetByUser(ctx, ownerB.ID); err != nil || (found && settings.LoginEmail != "") {
t.Fatalf("new owner settings=%+v found=%v err=%v, SMTP factor leaked to B", settings, found, err)
}
if _, found, err := codes.GetSnapshot(ctx, hash); err != nil || found {
t.Fatalf("owner-drift phone hash found=%v err=%v, want invalidated", found, err)
}
}

70
internal/app/account/phone_change.go
View file

@ -3,7 +3,6 @@ package account
import (
"context"
"crypto/rand"
"crypto/subtle"
"encoding/hex"
"fmt"
"strings"
@ -51,9 +50,10 @@ func (s *Service) SendChangePhoneCode(ctx context.Context, userID int64, authKey
return "", domain.AuthCodeDelivery{}, err
}
rec := store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
Phone: phone,
Code: s.phoneChangeCode,
Channel: "phone",
Channel: store.PhoneCodeChannelPhone,
Purpose: store.PhoneCodePurposeChangePhone,
UserID: userID,
AuthKeyID: authKeyID,
@ -68,7 +68,7 @@ func (s *Service) SendChangePhoneCode(ctx context.Context, userID int64, authKey
// ChangePhone 验证作用域和验证码后执行原子改号。返回事件用于当前 session 的
// pts 簿记；其它 session 由 transactional outbox 投递 updateUserPhone。
func (s *Service) ChangePhone(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error) {
func (s *Service) ChangePhone(ctx context.Context, userID int64, authKeyID, originRawAuthKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error) {
if strings.TrimSpace(phoneCodeHash) == "" || strings.TrimSpace(code) == "" {
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeEmpty
}
@ -82,51 +82,45 @@ func (s *Service) ChangePhone(ctx context.Context, userID int64, authKeyID [8]by
if s.codes == nil || s.phoneChanges == nil {
return domain.PhoneChangeResult{}, fmt.Errorf("phone change service is not configured")
}
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
scope := store.PhoneCodeScope{
Purpose: store.PhoneCodePurposeChangePhone,
UserID: userID,
AuthKeyID: authKeyID,
Phone: phone,
}
verified, err := s.codes.VerifyScoped(ctx, phoneCodeHash, scope, strings.TrimSpace(code), s.phoneChangeMaxAttempts)
if err != nil {
return domain.PhoneChangeResult{}, err
}
if !found {
switch verified.Status {
case store.LoginCodeVerifyMissing:
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeExpired
}
if rec.Purpose != store.PhoneCodePurposeChangePhone || rec.Phone != phone || rec.UserID != userID || rec.AuthKeyID != authKeyID {
case store.LoginCodeVerifyInvalid:
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeInvalid
case store.LoginCodeVerifyAccepted:
default:
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeInvalid
}
code = strings.TrimSpace(code)
if subtle.ConstantTimeCompare([]byte(rec.Code), []byte(code)) != 1 {
return domain.PhoneChangeResult{}, s.rejectPhoneChangeCode(ctx, phoneCodeHash, rec)
}
if existing, occupied, err := s.users.ByPhone(ctx, phone); err != nil {
return domain.PhoneChangeResult{}, err
} else if occupied && existing.ID != userID {
return domain.PhoneChangeResult{}, domain.ErrPhoneNumberOccupied
}
// 正确 code 必须在进入持久化事务前原子消费。并发重放中只有一个请求能
// 获得记录，其余请求不得再次推进 pts 或追加 user_phone event。
consumed, found, err := s.codes.ConsumeScoped(ctx, phoneCodeHash, store.PhoneCodeScope{
Purpose: store.PhoneCodePurposeChangePhone,
UserID: userID,
AuthKeyID: authKeyID,
Phone: phone,
})
if err != nil {
return domain.PhoneChangeResult{}, err
}
if !found {
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeExpired
}
if consumed.Purpose != rec.Purpose || consumed.UserID != rec.UserID || consumed.AuthKeyID != rec.AuthKeyID || consumed.Phone != rec.Phone ||
subtle.ConstantTimeCompare([]byte(consumed.Code), []byte(code)) != 1 {
consumed := verified.Record
if consumed.Version != store.PhoneCodeVersionCurrent || consumed.Scope() != scope || consumed.Channel != store.PhoneCodeChannelPhone {
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeInvalid
}
if date == 0 {
date = int(time.Now().Unix())
}
result, err := s.phoneChanges.ChangePhone(ctx, domain.PhoneChangeRequest{
UserID: userID,
Phone: phone,
Date: date,
ExcludeAuthKeyID: authKeyID,
UserID: userID,
Phone: phone,
Date: date,
// Authorization/code scope is the stable business (perm) key, while dispatch exclusion
// must use the physical raw key. They differ on PFS/temp connections; conflating them
// echoes updateUserPhone back to the initiating device and suppresses the wrong session.
ExcludeAuthKeyID: originRawAuthKeyID,
ExcludeSessionID: sessionID,
})
if err != nil {
@ -162,20 +156,6 @@ func (s *Service) phoneChangeCaller(ctx context.Context, userID int64, authKeyID
return u, nil
}
func (s *Service) rejectPhoneChangeCode(ctx context.Context, hash string, rec store.PhoneCode) error {
rec.Attempts++
max := rec.MaxAttempts
if max <= 0 {
max = s.phoneChangeMaxAttempts
}
if max > 0 && rec.Attempts >= max {
_ = s.codes.Del(ctx, hash)
return domain.ErrPhoneCodeInvalid
}
_ = s.codes.Update(ctx, hash, rec)
return domain.ErrPhoneCodeInvalid
}
func phoneChangeHash() (string, error) {
var raw [8]byte
if _, err := rand.Read(raw[:]); err != nil {

45
internal/app/account/phone_change_test.go
View file

@ -21,6 +21,26 @@ type phoneChangeFixture struct {
events *memory.UpdateEventStore
user domain.User
authKeyID [8]byte
changes *recordingPhoneChangeStore
}
type recordingPhoneChangeStore struct {
mu sync.Mutex
inner store.PhoneChangeStore
last domain.PhoneChangeRequest
}
func (s *recordingPhoneChangeStore) ChangePhone(ctx context.Context, req domain.PhoneChangeRequest) (domain.PhoneChangeResult, error) {
s.mu.Lock()
s.last = req
s.mu.Unlock()
return s.inner.ChangePhone(ctx, req)
}
func (s *recordingPhoneChangeStore) lastRequest() domain.PhoneChangeRequest {
s.mu.Lock()
defer s.mu.Unlock()
return s.last
}
func newPhoneChangeFixture(t *testing.T) phoneChangeFixture {
@ -38,12 +58,13 @@ func newPhoneChangeFixture(t *testing.T) phoneChangeFixture {
if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: authKeyID, UserID: u.ID, CreatedAt: time.Now().Add(-48 * time.Hour)}); err != nil {
t.Fatalf("bind auth: %v", err)
}
changes := &recordingPhoneChangeStore{inner: memory.NewPhoneChangeStore(users, events)}
service := NewService(
memory.NewPasswordStore(),
WithUsers(users),
WithPhoneChange(memory.NewPhoneChangeStore(users, events), auths, codes, nil, "12345", time.Minute, 3),
WithPhoneChange(changes, auths, codes, nil, "12345", time.Minute, 3),
)
return phoneChangeFixture{ctx: ctx, service: service, users: users, auths: auths, codes: codes, events: events, user: u, authKeyID: authKeyID}
return phoneChangeFixture{ctx: ctx, service: service, users: users, auths: auths, codes: codes, events: events, user: u, authKeyID: authKeyID, changes: changes}
}
func TestPhoneChangeScopesCodeAndPersistsDurableEvent(t *testing.T) {
@ -59,17 +80,21 @@ func TestPhoneChangeScopesCodeAndPersistsDurableEvent(t *testing.T) {
if err != nil || !found {
t.Fatalf("load code found=%v err=%v", found, err)
}
if rec.Purpose != store.PhoneCodePurposeChangePhone || rec.Phone != "15550012002" || rec.UserID != f.user.ID || rec.AuthKeyID != f.authKeyID || rec.SessionID != 77 {
if rec.Version != store.PhoneCodeVersionCurrent || rec.Purpose != store.PhoneCodePurposeChangePhone || rec.Phone != "15550012002" || rec.UserID != f.user.ID || rec.AuthKeyID != f.authKeyID || rec.SessionID != 77 {
t.Fatalf("scoped code = %+v", rec)
}
result, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 88, "+1 555 001 2002", hash, "12345", 1700000000)
rawAuthKeyID := [8]byte{8, 8, 8, 8}
result, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, rawAuthKeyID, 88, "+1 555 001 2002", hash, "12345", 1700000000)
if err != nil {
t.Fatalf("change phone after session reconnect: %v", err)
}
if !result.Changed || result.User.Phone != "15550012002" || result.Event.Type != domain.UpdateEventUserPhone || result.Event.Phone != "15550012002" || result.Event.Pts != 1 {
t.Fatalf("change result = %+v", result)
}
if got := f.changes.lastRequest().ExcludeAuthKeyID; got != rawAuthKeyID {
t.Fatalf("outbox exclusion auth key = %x, want physical raw %x", got, rawAuthKeyID)
}
if _, found, _ := f.users.ByPhone(f.ctx, "15550012001"); found {
t.Fatal("old phone still resolves")
}
@ -103,7 +128,7 @@ func TestPhoneChangeRejectsOccupiedAndCrossAuthCode(t *testing.T) {
if err := f.auths.Bind(f.ctx, domain.Authorization{AuthKeyID: otherKey, UserID: occupied.ID}); err != nil {
t.Fatalf("bind other auth: %v", err)
}
if _, err := f.service.ChangePhone(f.ctx, occupied.ID, otherKey, 99, "15550012004", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeInvalid) {
if _, err := f.service.ChangePhone(f.ctx, occupied.ID, otherKey, otherKey, 99, "15550012004", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeExpired) {
t.Fatalf("cross-auth change err = %v", err)
}
if got, found, _ := f.users.ByID(f.ctx, occupied.ID); !found || got.Phone != "15550012003" {
@ -118,11 +143,11 @@ func TestPhoneChangeWrongCodeExhaustsAttempts(t *testing.T) {
t.Fatalf("send code: %v", err)
}
for i := 0; i < 3; i++ {
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 77, "15550012005", hash, "00000", 0); !errors.Is(err, domain.ErrPhoneCodeInvalid) {
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 77, "15550012005", hash, "00000", 0); !errors.Is(err, domain.ErrPhoneCodeInvalid) {
t.Fatalf("wrong attempt %d err = %v", i+1, err)
}
}
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 77, "15550012005", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeExpired) {
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 77, "15550012005", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeExpired) {
t.Fatalf("exhausted code err = %v", err)
}
if got, _, _ := f.users.ByID(f.ctx, f.user.ID); got.Phone != "15550012001" {
@ -143,10 +168,10 @@ func TestPhoneChangeNewSendInvalidatesPreviousHash(t *testing.T) {
if oldHash == newHash {
t.Fatalf("hash was not rotated: %q", oldHash)
}
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 99, "15550012006", oldHash, "12345", 1700000001); !errors.Is(err, domain.ErrPhoneCodeExpired) {
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 99, "15550012006", oldHash, "12345", 1700000001); !errors.Is(err, domain.ErrPhoneCodeExpired) {
t.Fatalf("old hash replay err = %v", err)
}
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 99, "15550012006", newHash, "12345", 1700000002); err != nil {
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 99, "15550012006", newHash, "12345", 1700000002); err != nil {
t.Fatalf("new hash change: %v", err)
}
events, err := f.events.ListAfter(f.ctx, f.user.ID, 0, 10)
@ -168,7 +193,7 @@ func TestPhoneChangeConcurrentReplayAppendsOneEvent(t *testing.T) {
wg.Add(1)
go func() {
defer wg.Done()
_, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 88, "15550012007", hash, "12345", 1700000003)
_, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 88, "15550012007", hash, "12345", 1700000003)
errs <- err
}()
}

167
internal/app/account/service.go
View file

@ -17,13 +17,14 @@ import (
var defaultSecureRandom = []byte("telesrv-tdesktop-dev-secure-rand")
const (
passwordResetWait = 7 * 24 * time.Hour
passwordResetRetry = 24 * time.Hour
loginEmailVerifyChangePrefix = "login-email-change:"
loginEmailVerifySetupPrefix = "login-email-setup:"
codeChannelEmailSetup = "email_setup"
codeChannelEmailChange = "email_change"
codeChannelEmailLogin = "email_login"
passwordResetWait = 7 * 24 * time.Hour
passwordResetRetry = 24 * time.Hour
loginEmailVerifyChangePrefix = "login-email-change:"
loginEmailVerifySetupPrefix = "login-email-setup:"
codeChannelEmailSetup = "email_setup"
codeChannelEmailChange = "email_change"
codeChannelEmailLogin = "email_login"
codeChannelEmailSetupRequired = "email_setup_required"
)
// Service 提供账号安全配置查询。
@ -568,12 +569,16 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
}
key := loginEmailVerifyChangePrefix + fmt.Sprint(userID)
rec := store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
Code: "",
Channel: codeChannelEmailChange,
PendingEmail: email,
MaxAttempts: s.loginEmailCodeMaxAttempts,
}
if setup {
if s.users == nil {
return "", 0, domain.ErrEmailNotAllowed
}
phone = domain.NormalizePhone(phone)
phoneRec, found, err := s.codes.Get(ctx, phoneCodeHash)
if err != nil {
@ -582,7 +587,8 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
if !found {
return "", 0, domain.ErrEmailCodeInvalid
}
if phoneRec.Phone != phone {
if phoneRec.Version != store.PhoneCodeVersionCurrent || phoneRec.Purpose != "" || phoneRec.Phone != phone ||
phoneRec.Channel != codeChannelEmailSetupRequired || phoneRec.SignUpVerified {
return "", 0, domain.ErrEmailInvalid
}
targetUserID := int64(0)
@ -591,6 +597,9 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
} else if found {
targetUserID = existingUserID
}
if phoneRec.IssuedUserID != targetUserID {
return "", 0, domain.ErrEmailInvalid
}
if err := s.ensureLoginEmailAvailable(ctx, targetUserID, email); err != nil {
return "", 0, err
}
@ -611,7 +620,9 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
return "", 0, err
}
if err := s.loginEmailSender.SendLoginCode(ctx, email, code, s.loginEmailCodeTTL); err != nil {
_ = s.codes.Del(ctx, key)
// Set does not expose its generated revision. A blind Del here could
// remove a newer concurrent resend; leave the unreachable random code
// to expire or be replaced by the retry instead.
return "", 0, err
}
return emailPattern(email), len(code), nil
@ -625,28 +636,43 @@ func (s *Service) VerifyLoginEmail(ctx context.Context, userID int64, phone, pho
if setup {
key = loginEmailVerifySetupPrefix + phoneCodeHash
}
rec, found, err := s.codes.Get(ctx, key)
snapshot, found, err := s.codes.GetSnapshot(ctx, key)
if err != nil {
return "", err
}
if !found {
return "", domain.ErrEmailCodeInvalid
}
rec := snapshot.Record
if strings.TrimSpace(code) == "" || subtle.ConstantTimeCompare([]byte(rec.Code), []byte(strings.TrimSpace(code))) != 1 {
return "", s.rejectEmailCode(ctx, key, rec)
return "", s.rejectEmailCode(ctx, key, snapshot)
}
email := normalizeLoginEmail(rec.PendingEmail)
if !validLoginEmail(email) {
_ = s.codes.Del(ctx, key)
applied, deleteErr := s.codes.CompareAndDelete(ctx, key, snapshot.Revision)
if deleteErr != nil {
return "", deleteErr
}
if !applied {
return "", domain.ErrEmailCodeInvalid
}
return "", domain.ErrEmailInvalid
}
if setup {
if s.users == nil || rec.Channel != codeChannelEmailSetup {
return "", domain.ErrEmailCodeInvalid
}
phone = domain.NormalizePhone(phone)
phoneRec, found, err := s.codes.Get(ctx, phoneCodeHash)
if rec.Phone != phone {
return "", domain.ErrEmailCodeInvalid
}
phoneSnapshot, found, err := s.codes.GetSnapshot(ctx, phoneCodeHash)
if err != nil {
return "", err
}
if !found || phoneRec.Phone != phone {
phoneRec := phoneSnapshot.Record
if !found || phoneRec.Version != store.PhoneCodeVersionCurrent || phoneRec.Purpose != "" ||
phoneRec.Phone != phone || phoneRec.Channel != codeChannelEmailSetupRequired || phoneRec.SignUpVerified {
return "", domain.ErrEmailCodeInvalid
}
targetUserID := int64(0)
@ -655,11 +681,23 @@ func (s *Service) VerifyLoginEmail(ctx context.Context, userID int64, phone, pho
} else if found {
targetUserID = existingUserID
}
if phoneRec.IssuedUserID != targetUserID {
s.invalidateLoginCode(ctx, phoneCodeHash, phone)
return "", domain.ErrEmailCodeInvalid
}
if err := s.ensureLoginEmailAvailable(ctx, targetUserID, email); err != nil {
_ = s.codes.Del(ctx, key)
return "", err
}
_ = s.codes.Del(ctx, key)
// Claim this exact email-code revision before mutating the phone login
// state. A concurrent resend rotates the revision, so an old verifier
// can neither consume the new code nor authorize the phone hash.
claimed, err := s.codes.CompareAndDelete(ctx, key, snapshot.Revision)
if err != nil {
return "", err
}
if !claimed {
return "", domain.ErrEmailCodeInvalid
}
phoneRec.Channel = codeChannelEmailLogin
phoneRec.Code = strings.TrimSpace(code)
phoneRec.Email = email
@ -667,43 +705,94 @@ func (s *Service) VerifyLoginEmail(ctx context.Context, userID int64, phone, pho
phoneRec.VerifiedEmail = true
phoneRec.Attempts = 0
phoneRec.MaxAttempts = s.loginEmailCodeMaxAttempts
if err := s.codes.Update(ctx, phoneCodeHash, phoneRec); err != nil {
updated, err := s.codes.CompareAndUpdate(ctx, phoneCodeHash, phoneSnapshot.Revision, phoneRec)
if err != nil {
return "", err
}
if _, found, err := s.userIDByPhone(ctx, phone); err != nil {
if !updated {
return "", domain.ErrEmailCodeInvalid
}
if targetUserID == 0 {
verified, err := s.codes.VerifyLogin(ctx, phoneCodeHash, phone, phoneRec.Code, true, s.loginEmailCodeMaxAttempts)
if err != nil {
return "", err
}
if verified.Status != store.LoginCodeVerifyAccepted || verified.Record.IssuedUserID != 0 || !verified.Record.SignUpVerified {
return "", domain.ErrEmailCodeInvalid
}
phoneRec = verified.Record
}
afterUserID := int64(0)
if existingUserID, found, err := s.userIDByPhone(ctx, phone); err != nil {
return "", err
} else if found {
if err := s.SetLoginEmailByPhone(ctx, phone, email); err != nil {
afterUserID = existingUserID
}
if afterUserID != targetUserID || phoneRec.IssuedUserID != afterUserID {
s.invalidateLoginCode(ctx, phoneCodeHash, phone)
return "", domain.ErrEmailCodeInvalid
}
if targetUserID != 0 {
// Keep the identity selected before SMTP verification. Re-resolving
// phone at this write boundary would let an A→B transfer attach A's
// verified factor to B.
if err := s.SetLoginEmail(ctx, targetUserID, email); err != nil {
return "", err
}
finalUserID := int64(0)
if existingUserID, found, err := s.userIDByPhone(ctx, phone); err != nil {
return "", err
} else if found {
finalUserID = existingUserID
}
if finalUserID != targetUserID {
s.invalidateLoginCode(ctx, phoneCodeHash, phone)
return "", domain.ErrEmailCodeInvalid
}
}
return email, nil
}
if err := s.ensureLoginEmailAvailable(ctx, userID, email); err != nil {
_ = s.codes.Del(ctx, key)
return "", err
}
_ = s.codes.Del(ctx, key)
claimed, err := s.codes.CompareAndDelete(ctx, key, snapshot.Revision)
if err != nil {
return "", err
}
if !claimed {
return "", domain.ErrEmailCodeInvalid
}
if err := s.SetLoginEmail(ctx, userID, email); err != nil {
return "", err
}
return email, nil
}
func (s *Service) rejectEmailCode(ctx context.Context, key string, rec store.PhoneCode) error {
func (s *Service) rejectEmailCode(ctx context.Context, key string, snapshot store.PhoneCodeSnapshot) error {
rec := snapshot.Record
rec.Attempts++
max := rec.MaxAttempts
if max <= 0 {
max = s.loginEmailCodeMaxAttempts
}
if max > 0 && rec.Attempts >= max {
_ = s.codes.Del(ctx, key)
if _, err := s.codes.CompareAndDelete(ctx, key, snapshot.Revision); err != nil {
return err
}
return domain.ErrEmailCodeInvalid
}
_ = s.codes.Update(ctx, key, rec)
if _, err := s.codes.CompareAndUpdate(ctx, key, snapshot.Revision, rec); err != nil {
return err
}
return domain.ErrEmailCodeInvalid
}
func (s *Service) invalidateLoginCode(ctx context.Context, hash, phone string) {
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Second)
defer cancel()
_, _ = s.codes.InvalidateLoginCode(cleanupCtx, hash, phone)
}
// SetLoginEmail 为已登录用户写入登录邮箱（authed 的 emailVerifyPurposeLoginChange）。
// 账号无 2FA 也可设置：account_passwords 行可在 has_password=false 下仅承载登录邮箱。
func (s *Service) SetLoginEmail(ctx context.Context, userID int64, email string) error {
@ -726,19 +815,6 @@ func (s *Service) SetLoginEmail(ctx context.Context, userID int64, email string)
return s.passwords.Save(ctx, userID, settings)
}
// SetLoginEmailByPhone 为某手机号对应的账号写入登录邮箱（登录流程中的
// emailVerifyPurposeLoginSetup，此时尚未鉴权，只能凭 phone 定位用户）。
func (s *Service) SetLoginEmailByPhone(ctx context.Context, phone, email string) error {
userID, found, err := s.userIDByPhone(ctx, phone)
if err != nil {
return err
}
if !found {
return domain.ErrEmailInvalid
}
return s.SetLoginEmail(ctx, userID, email)
}
// LoginEmail 返回已登录用户的登录邮箱原始地址（用于 verifyEmail 回显 emailVerified.email）。
func (s *Service) LoginEmail(ctx context.Context, userID int64) (string, bool, error) {
if s == nil || s.passwords == nil || userID == 0 {
@ -754,8 +830,7 @@ func (s *Service) LoginEmail(ctx context.Context, userID int64) (string, bool, e
return normalizeLoginEmail(settings.LoginEmail), true, nil
}
// LoginEmailByPhone 按手机号返回登录邮箱原始地址（供 auth.sendCode 检测是否改投邮箱、
// login-setup 回显、reset 回显使用）。
// LoginEmailByPhone 按手机号返回登录邮箱原始地址，供 auth.sendCode 检测是否改投邮箱。
func (s *Service) LoginEmailByPhone(ctx context.Context, phone string) (string, bool, error) {
userID, found, err := s.userIDByPhone(ctx, phone)
if err != nil || !found {
@ -764,14 +839,12 @@ func (s *Service) LoginEmailByPhone(ctx context.Context, phone string) (string,
return s.LoginEmail(ctx, userID)
}
// ClearLoginEmailByPhone 清除某手机号账号的登录邮箱（auth.resetLoginEmail）。
func (s *Service) ClearLoginEmailByPhone(ctx context.Context, phone string) error {
userID, found, err := s.userIDByPhone(ctx, phone)
if err != nil {
return err
}
if !found {
return nil
// ClearLoginEmail clears the factor on the exact account selected by the
// preceding reset-code consume. Authentication factors must never be mutated
// through a second phone→user lookup.
func (s *Service) ClearLoginEmail(ctx context.Context, userID int64) error {
if s == nil || s.passwords == nil || userID == 0 {
return domain.ErrEmailInvalid
}
settings, found, err := s.passwords.GetByUser(ctx, userID)
if err != nil || !found {

32
internal/app/auth/change_phone_resend_test.go
View file

@ -15,6 +15,7 @@ func TestResendCodePreservesChangePhoneScopeAndSMSDelivery(t *testing.T) {
codes := memory.NewCodeStore()
authKeyID := [8]byte{8, 7, 6}
rec := store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
Phone: "15550014001",
Code: "old",
Channel: codeChannelPhone,
@ -69,3 +70,34 @@ func TestResendCodePreservesChangePhoneScopeAndSMSDelivery(t *testing.T) {
t.Fatal("scoped cancel left hash valid")
}
}
func TestResendAndCancelRejectLegacyChangePhoneCode(t *testing.T) {
ctx := context.Background()
codes := memory.NewCodeStore()
authKeyID := [8]byte{8, 8, 8}
legacy := store.PhoneCode{
Version: 0, Phone: "15550014002", Code: "12345", Channel: codeChannelPhone,
Purpose: store.PhoneCodePurposeChangePhone, UserID: 43, AuthKeyID: authKeyID,
}
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithCodeTTL(time.Minute))
if err := codes.Set(ctx, "legacy-resend", legacy, time.Minute); err != nil {
t.Fatal(err)
}
if _, err := svc.ResendCodeForAuthKey(ctx, authKeyID, legacy.Phone, "legacy-resend"); err != ErrCodeExpired {
t.Fatalf("legacy resend err=%v, want ErrCodeExpired", err)
}
if _, found, _ := codes.Get(ctx, "legacy-resend"); found {
t.Fatal("legacy resend left code active")
}
if err := codes.Set(ctx, "legacy-cancel", legacy, time.Minute); err != nil {
t.Fatal(err)
}
if err := svc.CancelCodeForAuthKey(ctx, authKeyID, legacy.Phone, "legacy-cancel"); err != ErrCodeExpired {
t.Fatalf("legacy cancel err=%v, want ErrCodeExpired", err)
}
if _, found, _ := codes.Get(ctx, "legacy-cancel"); found {
t.Fatal("legacy cancel left code active")
}
}

366
internal/app/auth/login_code_delivery_test.go Normal file
View file

@ -0,0 +1,366 @@
package auth
import (
"context"
"errors"
"testing"
"time"
"telesrv/internal/domain"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
type captureLoginCodeDelivery struct {
requests []domain.LoginCodeDeliveryRequest
result domain.LoginCodeDeliveryResult
err error
failAt int
}
func (d *captureLoginCodeDelivery) DeliverLoginCodeMessage(_ context.Context, req domain.LoginCodeDeliveryRequest) (domain.LoginCodeDeliveryResult, error) {
d.requests = append(d.requests, req)
if d.err != nil && (d.failAt == 0 || len(d.requests) == d.failAt) {
return domain.LoginCodeDeliveryResult{}, d.err
}
return d.result, nil
}
type trackingCodeStore struct {
store.CodeStore
lastSetHash string
deleted []string
deleteCtx []error
deleteErr error
}
func (s *trackingCodeStore) Set(ctx context.Context, hash string, code store.PhoneCode, ttl time.Duration) error {
s.lastSetHash = hash
return s.CodeStore.Set(ctx, hash, code, ttl)
}
func (s *trackingCodeStore) Del(ctx context.Context, hash string) error {
s.deleted = append(s.deleted, hash)
s.deleteCtx = append(s.deleteCtx, ctx.Err())
if s.deleteErr != nil {
return s.deleteErr
}
return s.CodeStore.Del(ctx, hash)
}
func TestExistingAccountSendCodeDeliversBeforeSignInAndDoesNotRedeliver(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
authz := memory.NewAuthorizationStore()
codes := memory.NewCodeStore()
u, err := users.Create(ctx, domain.User{Phone: "15550009201", FirstName: "Existing"})
if err != nil {
t.Fatalf("create user: %v", err)
}
delivery := &captureLoginCodeDelivery{result: domain.LoginCodeDeliveryResult{Created: true}}
svc := NewService(users, authz, codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
before := int(time.Now().Unix())
hash, err := svc.SendCode(ctx, "+1 555 000 9201")
if err != nil {
t.Fatalf("SendCode: %v", err)
}
if hash == "" || len(delivery.requests) != 1 {
t.Fatalf("SendCode hash=%q delivery calls=%d, want non-empty/1", hash, len(delivery.requests))
}
req := delivery.requests[0]
if req.UserID != u.ID || req.PhoneCodeHash != hash || req.Code != "12345" || req.Date < before || req.ExpiresAt < int64(before)+int64((5*time.Minute)/time.Second)-1 {
t.Fatalf("delivery request = %+v, want user=%d hash=%q code=12345 date>=%d", req, u.ID, hash, before)
}
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || rec.Code != "12345" {
t.Fatalf("code after synchronous delivery = %+v found=%v err=%v", rec, found, err)
}
var key [8]byte
key[0] = 0x92
got, lateMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "+15550009201", hash, "12345")
if err != nil || needSignUp || got.ID != u.ID {
t.Fatalf("SignIn user=%d needSignUp=%v err=%v, want %d/false", got.ID, needSignUp, err, u.ID)
}
if lateMessage.ID != 0 || len(delivery.requests) != 1 {
t.Fatalf("SignIn lateMessage=%+v delivery calls=%d, want zero/unchanged", lateMessage, len(delivery.requests))
}
}
func TestDeliveredLoginCodeSurvivesWrongSignInAndCancelWithoutDuplicate(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
u, err := users.Create(ctx, domain.User{Phone: "15550009208", FirstName: "Cancel"})
if err != nil {
t.Fatalf("create user: %v", err)
}
codes := memory.NewCodeStore()
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
events := memory.NewUpdateEventStore()
delivery := memory.NewLoginCodeDeliveryStore(messages, events)
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
hash, err := svc.SendCode(ctx, "15550009208")
if err != nil {
t.Fatalf("SendCode: %v", err)
}
assertFacts := func(stage string) {
t.Helper()
history, historyErr := messages.ListByUser(ctx, u.ID, domain.MessageFilter{
HasPeer: true,
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID},
Limit: 10,
})
durable, eventErr := events.ListAfter(ctx, u.ID, 0, 10)
if historyErr != nil || eventErr != nil || len(history.Messages) != 1 || len(durable) != 1 {
t.Fatalf("%s messages=%d events=%d historyErr=%v eventErr=%v, want 1/1", stage, len(history.Messages), len(durable), historyErr, eventErr)
}
}
assertFacts("after SendCode")
if _, late, _, err := svc.SignIn(ctx, domain.Authorization{}, "15550009208", hash, "00000"); !errors.Is(err, ErrCodeInvalid) || late.ID != 0 {
t.Fatalf("wrong SignIn late=%+v err=%v, want ErrCodeInvalid/no message", late, err)
}
assertFacts("after wrong SignIn")
if err := svc.CancelCode(ctx, "15550009208", hash); err != nil {
t.Fatalf("CancelCode: %v", err)
}
assertFacts("after CancelCode")
if _, late, _, err := svc.SignIn(ctx, domain.Authorization{}, "15550009208", hash, "12345"); !errors.Is(err, ErrCodeExpired) || late.ID != 0 {
t.Fatalf("SignIn after cancel late=%+v err=%v, want ErrCodeExpired/no message", late, err)
}
assertFacts("after canceled SignIn")
}
func TestCodeIssuedBeforeConcurrentOwnerCreationIsRejected(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
authz := memory.NewAuthorizationStore()
codes := memory.NewCodeStore()
delivery := &captureLoginCodeDelivery{}
svc := NewService(users, authz, codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
hash, err := svc.SendCode(ctx, "15550009209")
if err != nil {
t.Fatalf("SendCode before signup: %v", err)
}
rec, found, err := codes.Get(ctx, hash)
if err != nil || !found || rec.Version != store.PhoneCodeVersionCurrent || rec.IssuedUserID != 0 || rec.SignUpVerified || len(delivery.requests) != 0 {
t.Fatalf("pre-signup code=%+v found=%v err=%v deliveries=%d", rec, found, err, len(delivery.requests))
}
u, err := users.Create(ctx, domain.User{Phone: "15550009209", FirstName: "Concurrent"})
if err != nil {
t.Fatalf("concurrent create user: %v", err)
}
var key [8]byte
key[0] = 0x93
got, lateMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "15550009209", hash, "12345")
if !errors.Is(err, ErrCodeInvalid) || needSignUp || got.ID != 0 || lateMessage.ID != 0 {
t.Fatalf("SignIn after owner creation got=%+v late=%+v needSignUp=%v err=%v, want invalid", got, lateMessage, needSignUp, err)
}
if len(delivery.requests) != 0 {
t.Fatalf("owner-transfer code was delivered to new owner: %+v", delivery.requests)
}
if bound, ok, err := svc.UserID(ctx, key); err != nil || ok || bound != 0 {
t.Fatalf("bound user=%d ok=%v err=%v, want no authorization (created uid=%d)", bound, ok, err, u.ID)
}
}
func TestExistingAccountRepeatedSendCodeDeliversEachIssuedHashWithoutSignIn(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009202"}); err != nil {
t.Fatalf("create user: %v", err)
}
delivery := &captureLoginCodeDelivery{}
svc := NewService(users, memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithLoginCodeDelivery(delivery))
first, err := svc.SendCode(ctx, "15550009202")
if err != nil {
t.Fatalf("first SendCode: %v", err)
}
second, err := svc.SendCode(ctx, "15550009202")
if err != nil {
t.Fatalf("second SendCode: %v", err)
}
if first == second || len(delivery.requests) != 2 {
t.Fatalf("hashes=%q/%q delivery calls=%d, want distinct/2", first, second, len(delivery.requests))
}
if delivery.requests[0].PhoneCodeHash != first || delivery.requests[1].PhoneCodeHash != second {
t.Fatalf("delivery hashes = %q/%q, want %q/%q", delivery.requests[0].PhoneCodeHash, delivery.requests[1].PhoneCodeHash, first, second)
}
}
func TestExistingAccountSendCodeDeliveryFailureRevokesCode(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009203"}); err != nil {
t.Fatalf("create user: %v", err)
}
baseCodes := memory.NewCodeStore()
codes := &trackingCodeStore{CodeStore: baseCodes}
deliveryCause := errors.New("durable write failed")
delivery := &captureLoginCodeDelivery{err: deliveryCause}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
hash, err := svc.SendCode(ctx, "15550009203")
if hash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) || !errors.Is(err, deliveryCause) {
t.Fatalf("SendCode hash=%q err=%v, want empty ErrLoginCodeDeliveryFailed+cause", hash, err)
}
if codes.lastSetHash == "" || len(codes.deleted) != 1 || codes.deleted[0] != codes.lastSetHash {
t.Fatalf("set hash=%q deleted=%v, want exact rollback", codes.lastSetHash, codes.deleted)
}
if _, found, getErr := baseCodes.Get(ctx, codes.lastSetHash); getErr != nil || found {
t.Fatalf("rolled-back hash found=%v err=%v", found, getErr)
}
}
func TestExistingAccountAmbiguousDeliveryPreservesCodeForIdempotentRetry(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009213"}); err != nil {
t.Fatalf("create user: %v", err)
}
baseCodes := memory.NewCodeStore()
codes := &trackingCodeStore{CodeStore: baseCodes}
delivery := &captureLoginCodeDelivery{err: domain.ErrLoginCodeDeliveryCommitAmbiguous}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
hash, err := svc.SendCode(ctx, "15550009213")
if hash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) || !errors.Is(err, domain.ErrLoginCodeDeliveryCommitAmbiguous) {
t.Fatalf("SendCode hash=%q err=%v, want ambiguous delivery failure", hash, err)
}
if codes.lastSetHash == "" || len(codes.deleted) != 0 {
t.Fatalf("ambiguous delivery set=%q deleted=%v, want code preserved", codes.lastSetHash, codes.deleted)
}
if rec, found, getErr := baseCodes.Get(ctx, codes.lastSetHash); getErr != nil || !found || rec.Code != "12345" {
t.Fatalf("ambiguous delivery code=%+v found=%v err=%v", rec, found, getErr)
}
}
func TestExplicitDeliveryFailureRollsBackWithDetachedContext(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
users := memory.NewUserStore()
if _, err := users.Create(context.Background(), domain.User{Phone: "15550009214"}); err != nil {
t.Fatalf("create user: %v", err)
}
baseCodes := memory.NewCodeStore()
codes := &trackingCodeStore{CodeStore: baseCodes}
delivery := &captureLoginCodeDelivery{err: errors.New("definite rollback")}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
if hash, err := svc.SendCode(ctx, "15550009214"); hash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) {
t.Fatalf("SendCode hash=%q err=%v, want definite failure", hash, err)
}
if len(codes.deleted) != 1 || len(codes.deleteCtx) != 1 || codes.deleteCtx[0] != nil {
t.Fatalf("rollback deleted=%v ctxErr=%v, want one detached delete", codes.deleted, codes.deleteCtx)
}
if _, found, err := baseCodes.Get(context.Background(), codes.lastSetHash); err != nil || found {
t.Fatalf("detached rollback found=%v err=%v", found, err)
}
}
func TestExistingAccountMissingDeliveryFailsClosedAndRevokesCode(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009204"}); err != nil {
t.Fatalf("create user: %v", err)
}
baseCodes := memory.NewCodeStore()
codes := &trackingCodeStore{CodeStore: baseCodes}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
hash, err := svc.SendCode(ctx, "15550009204")
if hash != "" || !errors.Is(err, ErrLoginCodeDeliveryUnavailable) {
t.Fatalf("SendCode hash=%q err=%v, want unavailable", hash, err)
}
if codes.lastSetHash == "" {
t.Fatal("missing delivery was checked before code creation; want rollback path covered")
}
if _, found, getErr := baseCodes.Get(ctx, codes.lastSetHash); getErr != nil || found {
t.Fatalf("unavailable delivery hash found=%v err=%v", found, getErr)
}
}
func TestExistingAccountResendDeliversNewHashAndInvalidatesOld(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009205"}); err != nil {
t.Fatalf("create user: %v", err)
}
codes := memory.NewCodeStore()
delivery := &captureLoginCodeDelivery{}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
oldHash, err := svc.SendCode(ctx, "15550009205")
if err != nil {
t.Fatalf("SendCode: %v", err)
}
newHash, err := svc.ResendCode(ctx, "15550009205", oldHash)
if err != nil {
t.Fatalf("ResendCode: %v", err)
}
if oldHash == newHash || len(delivery.requests) != 2 || delivery.requests[1].PhoneCodeHash != newHash {
t.Fatalf("old/new=%q/%q deliveries=%+v", oldHash, newHash, delivery.requests)
}
if _, found, err := codes.Get(ctx, oldHash); err != nil || found {
t.Fatalf("old code found=%v err=%v", found, err)
}
if _, found, err := codes.Get(ctx, newHash); err != nil || !found {
t.Fatalf("new code found=%v err=%v", found, err)
}
}
func TestExistingAccountResendDeliveryFailureLeavesNoUsableCode(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009206"}); err != nil {
t.Fatalf("create user: %v", err)
}
codes := memory.NewCodeStore()
delivery := &captureLoginCodeDelivery{err: errors.New("second delivery failed"), failAt: 2}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
oldHash, err := svc.SendCode(ctx, "15550009206")
if err != nil {
t.Fatalf("SendCode: %v", err)
}
newHash, err := svc.ResendCode(ctx, "15550009206", oldHash)
if newHash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) || len(delivery.requests) != 2 {
t.Fatalf("ResendCode hash=%q err=%v deliveries=%d", newHash, err, len(delivery.requests))
}
failedHash := delivery.requests[1].PhoneCodeHash
for _, hash := range []string{oldHash, failedHash} {
if _, found, getErr := codes.Get(ctx, hash); getErr != nil || found {
t.Fatalf("failed resend hash %q found=%v err=%v", hash, found, getErr)
}
}
}
func TestConfiguredEmailLoginDoesNotLeakCodeThroughAppDelivery(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
if _, err := users.Create(ctx, domain.User{Phone: "15550009207"}); err != nil {
t.Fatalf("create user: %v", err)
}
emails := &testLoginEmailStore{emails: map[string]string{"15550009207": "secure@example.test"}}
mailSender := &testMailSender{}
delivery := &captureLoginCodeDelivery{}
svc := NewService(users, memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
WithLoginEmail(LoginEmailOptions{Enabled: true, CodeLength: 6, Store: emails, Sender: mailSender}),
WithLoginCodeDelivery(delivery),
)
if _, err := svc.SendCode(ctx, "15550009207"); err != nil {
t.Fatalf("SendCode: %v", err)
}
if mailSender.to != "secure@example.test" || mailSender.code == "" {
t.Fatalf("email delivery = %q/%q", mailSender.to, mailSender.code)
}
if len(delivery.requests) != 0 {
t.Fatalf("email code leaked into app delivery: %+v", delivery.requests)
}
}

3
internal/app/auth/login_email_config_test.go
View file

@ -19,8 +19,7 @@ func (s *testLoginEmailStore) LoginEmailByPhone(_ context.Context, phone string)
return email, ok, nil
}
func (s *testLoginEmailStore) SetLoginEmailByPhone(_ context.Context, phone, email string) error {
s.emails[domain.NormalizePhone(phone)] = email
func (s *testLoginEmailStore) SetLoginEmail(_ context.Context, _ int64, _ string) error {
return nil
}

28
internal/app/auth/login_email_test.go
View file

@ -9,13 +9,17 @@ import (
"telesrv/internal/store/memory"
)
// TestSignInWithEmailCompletesLogin 验证带 email_verification 的登录：注册账号→登出→
// 重新 sendCode→用任意邮箱验证码经 SignInWithEmail 完成登录。
// TestSignInWithEmailCompletesLogin 验证旧客户端把 phone channel 放进
// email_verification 时仍可登录，但验证码必须精确匹配，不能用任意非空值绕过。
func TestSignInWithEmailCompletesLogin(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
authz := memory.NewAuthorizationStore()
svc := NewService(users, authz, memory.NewCodeStore(), nil, nil, "12345")
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
svc := NewService(users, authz, memory.NewCodeStore(), nil, nil, "12345",
WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messages, memory.NewUpdateEventStore())),
)
var key [8]byte
key[0] = 0x42
@ -23,6 +27,7 @@ func TestSignInWithEmailCompletesLogin(t *testing.T) {
if err != nil {
t.Fatalf("SendCode signup: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550009001", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "Email", "Login")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -35,7 +40,10 @@ func TestSignInWithEmailCompletesLogin(t *testing.T) {
if err != nil {
t.Fatalf("SendCode signin: %v", err)
}
got, _, needSignUp, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "anything-goes")
if _, _, _, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "anything-goes"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("SignInWithEmail arbitrary nonempty code err=%v, want ErrCodeInvalid", err)
}
got, _, needSignUp, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "12345")
if err != nil {
t.Fatalf("SignInWithEmail: %v", err)
}
@ -48,7 +56,7 @@ func TestSignInWithEmailCompletesLogin(t *testing.T) {
}
}
// TestSignInWithEmailRejectsEmptyCode 空邮箱验证码必须被拒（即使开发环境码任意，也不能空）。
// TestSignInWithEmailRejectsEmptyCode 空邮箱验证码必须被拒。
func TestSignInWithEmailRejectsEmptyCode(t *testing.T) {
ctx := context.Background()
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345")
@ -66,7 +74,12 @@ func TestSignInWithEmailRejectsEmptyCode(t *testing.T) {
func TestSignInWithEmailStillHonorsTwoFactor(t *testing.T) {
ctx := context.Background()
passwords := memory.NewPasswordStore()
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithPasswords(passwords))
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
WithPasswords(passwords),
WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messages, memory.NewUpdateEventStore())),
)
var key [8]byte
key[0] = 0x43
@ -74,6 +87,7 @@ func TestSignInWithEmailStillHonorsTwoFactor(t *testing.T) {
if err != nil {
t.Fatalf("SendCode signup: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550009003", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550009003", hash, "Two", "Factor")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -89,7 +103,7 @@ func TestSignInWithEmailStillHonorsTwoFactor(t *testing.T) {
if err != nil {
t.Fatalf("SendCode signin: %v", err)
}
got, _, _, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009003", hash, "any-email-code")
got, _, _, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009003", hash, "12345")
if !errors.Is(err, domain.ErrSessionPasswordNeeded) {
t.Fatalf("SignInWithEmail err = %v, want ErrSessionPasswordNeeded", err)
}

2
internal/app/auth/premium_grant_test.go
View file

@ -19,6 +19,7 @@ func TestSignUpPremiumGrant(t *testing.T) {
if err != nil {
t.Fatalf("SendCode: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550004401", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{}, "+15550004401", hash, "Prem", "User")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -41,6 +42,7 @@ func TestSignUpPremiumGrantDisabled(t *testing.T) {
if err != nil {
t.Fatalf("SendCode: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550004402", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{}, "+15550004402", hash, "Free", "User")
if err != nil {
t.Fatalf("SignUp: %v", err)

606
internal/app/auth/service.go
View file

@ -27,6 +27,13 @@ var (
ErrCodeExpired = errors.New("phone code expired or not found")
ErrCodeInvalid = errors.New("phone code invalid")
ErrEncryptedMessageInvalid = errors.New("encrypted message invalid")
// ErrLoginCodeDeliveryUnavailable 表示已有账号的 app-code 没有可用的
// durable message/event/outbox 投递边界。这是服务端配置错误，不能降级成
// “继续返回 sentCode，等 signIn 后补发”。
ErrLoginCodeDeliveryUnavailable = errors.New("login code durable delivery unavailable")
// ErrLoginCodeDeliveryFailed 表示 durable 投递未成功。SendCode/ResendCode
// 必须同时撤销刚写入的 CodeStore hash，防止客户拿到无法送达的码。
ErrLoginCodeDeliveryFailed = errors.New("login code durable delivery failed")
// ErrPhoneNumberInvalid 表示手机号为空或非纯数字/长度越界。
// 0090 把 users.phone 唯一约束改为忽略空串的部分索引（bot 行 phone=''），
// 因此 phone 校验必须前移到 auth 入口，否则 sendCode/signUp 可无限铸造
@ -40,6 +47,7 @@ const (
codeChannelPhone = "phone"
codeChannelEmailLogin = "email_login"
codeChannelEmailSetupRequired = "email_setup_required"
loginCodeRollbackTimeout = 2 * time.Second
)
// validPhone 校验规范化后的手机号：5-32 位纯数字（上限对齐 users.phone 列宽）。
@ -68,6 +76,7 @@ type Service struct {
passwords store.PasswordStore
messages store.MessageStore
dialogs store.DialogStore
loginCodeDelivery store.LoginCodeDeliveryStore
bots store.BotStore
fixedCode string
codeTTL time.Duration
@ -83,7 +92,7 @@ type Service struct {
type loginEmailStore interface {
LoginEmailByPhone(ctx context.Context, phone string) (string, bool, error)
SetLoginEmailByPhone(ctx context.Context, phone, email string) error
SetLoginEmail(ctx context.Context, userID int64, email string) error
}
type LoginEmailOptions struct {
@ -102,7 +111,9 @@ type authorizationRevoker interface {
// Option 调整登录服务的可选依赖。
type Option func(*Service)
// WithLoginMessages 在登录成功后写入官方系统账号的登录消息与会话摘要。
// WithLoginMessages 在新用户注册成功后写入官方系统账号的首条登录消息与会话摘要。
// 已有账号的 app 验证码必须在 auth.sendCode/resendCode 阶段通过
// WithLoginCodeDelivery 持久化，禁止在 signIn 成功后补发。
func WithLoginMessages(messages store.MessageStore, dialogs store.DialogStore) Option {
return func(s *Service) {
s.messages = messages
@ -110,6 +121,15 @@ func WithLoginMessages(messages store.MessageStore, dialogs store.DialogStore) O
}
}
// WithLoginCodeDelivery 注入已有账号 app-code 的 durable 投递边界。
// 实现必须以 user_id + phone_code_hash 幂等，并原子写入 777000
// message/dialog/user update event/dispatch outbox。
func WithLoginCodeDelivery(delivery store.LoginCodeDeliveryStore) Option {
return func(s *Service) {
s.loginCodeDelivery = delivery
}
}
// WithPasswords lets sign-in stop at SESSION_PASSWORD_NEEDED for 2FA accounts.
func WithPasswords(passwords store.PasswordStore) Option {
return func(s *Service) {
@ -271,53 +291,157 @@ func (s *Service) SendCode(ctx context.Context, phone string) (string, error) {
if systemLoginPhoneForbidden(phone) {
return "", ErrSystemUserLoginForbidden
}
existing, found, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return "", fmt.Errorf("lookup login-code recipient: %w", err)
}
if found && systemUserLoginForbidden(existing) {
return "", ErrSystemUserLoginForbidden
}
issuedUserID := int64(0)
if found {
issuedUserID = existing.ID
}
if s.loginEmailEnabled && s.loginEmails != nil {
email, found, err := s.loginEmails.LoginEmailByPhone(ctx, phone)
if err != nil {
return "", err
}
if found && strings.TrimSpace(email) != "" {
return s.createEmailLoginCode(ctx, phone, email)
return s.createEmailLoginCode(ctx, phone, email, issuedUserID)
}
if s.loginEmailRequireSetup {
return s.createSetupRequiredCode(ctx, phone)
return s.createSetupRequiredCode(ctx, phone, issuedUserID)
}
}
return s.createPhoneCode(ctx, phone)
return s.createPhoneCode(ctx, phone, issuedUserID)
}
func (s *Service) createPhoneCode(ctx context.Context, phone string) (string, error) {
func (s *Service) currentPhoneOwner(ctx context.Context, phone string) (domain.User, bool, error) {
if s == nil || s.users == nil {
return domain.User{}, false, fmt.Errorf("user store is not configured")
}
return s.users.ByPhone(ctx, phone)
}
func (s *Service) issuedOwnerMatches(ctx context.Context, phone string, issuedUserID int64) (bool, error) {
current, found, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return false, err
}
currentUserID := int64(0)
if found {
currentUserID = current.ID
}
return currentUserID == issuedUserID, nil
}
func (s *Service) ensureIssuedOwnerAfterSet(ctx context.Context, hash string, rec store.PhoneCode) error {
matches, err := s.issuedOwnerMatches(ctx, rec.Phone, rec.IssuedUserID)
if err == nil && matches {
return nil
}
cause := err
if cause == nil {
cause = ErrCodeInvalid
}
return s.rollbackUndeliveredCode(ctx, hash, cause)
}
func (s *Service) invalidateLoginCodeDetached(ctx context.Context, hash, phone string) {
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginCodeRollbackTimeout)
defer cancel()
_, _ = s.codes.InvalidateLoginCode(cleanupCtx, hash, phone)
}
func (s *Service) createPhoneCode(ctx context.Context, phone string, existingUserID int64) (string, error) {
hash, err := randomHex(8)
if err != nil {
return "", err
}
if err := s.codes.Set(ctx, hash, store.PhoneCode{
Phone: phone,
Code: s.fixedCode,
Channel: codeChannelPhone,
MaxAttempts: s.codeMaxAttempts,
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: existingUserID,
Phone: phone,
Code: s.fixedCode,
Channel: codeChannelPhone,
MaxAttempts: s.codeMaxAttempts,
}, s.codeTTL); err != nil {
return "", fmt.Errorf("store code: %w", err)
}
rec := store.PhoneCode{Phone: phone, IssuedUserID: existingUserID}
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
return "", err
}
// 新手机号还没有 owner/dialog，只能在 SignUp 创建用户后写第一条
// 777000 消息。已有账号则必须在 sendCode RPC 返回前把 app-code
// 作为普通 incoming message + durable update/outbox 提交；登录成功不再补发。
if existingUserID == 0 {
return hash, nil
}
if err := s.deliverLoginCode(ctx, existingUserID, hash, s.fixedCode); err != nil {
return "", s.rollbackUndeliveredCode(ctx, hash, err)
}
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
return "", err
}
return hash, nil
}
func (s *Service) createSetupRequiredCode(ctx context.Context, phone string) (string, error) {
func (s *Service) deliverLoginCode(ctx context.Context, userID int64, phoneCodeHash, code string) error {
if s.loginCodeDelivery == nil {
return ErrLoginCodeDeliveryUnavailable
}
now := time.Now()
if _, err := s.loginCodeDelivery.DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{
UserID: userID,
PhoneCodeHash: phoneCodeHash,
Code: code,
Date: int(now.Unix()),
ExpiresAt: now.Add(s.codeTTL).Unix(),
}); err != nil {
return errors.Join(ErrLoginCodeDeliveryFailed, err)
}
return nil
}
func (s *Service) rollbackUndeliveredCode(ctx context.Context, phoneCodeHash string, cause error) error {
// lib/pq can report an I/O failure after COMMIT reached PostgreSQL. In that
// state deleting the code could turn an already delivered 777000 message
// into an unusable login attempt. Preserve it and let the delivery receipt
// make the retry idempotent.
if errors.Is(cause, domain.ErrLoginCodeDeliveryCommitAmbiguous) {
return cause
}
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginCodeRollbackTimeout)
defer cancel()
if err := s.codes.Del(cleanupCtx, phoneCodeHash); err != nil {
return errors.Join(cause, fmt.Errorf("rollback undelivered login code: %w", err))
}
return cause
}
func (s *Service) createSetupRequiredCode(ctx context.Context, phone string, issuedUserID int64) (string, error) {
hash, err := randomHex(8)
if err != nil {
return "", err
}
if err := s.codes.Set(ctx, hash, store.PhoneCode{
Phone: phone,
Channel: codeChannelEmailSetupRequired,
MaxAttempts: s.codeMaxAttempts,
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: issuedUserID,
Phone: phone,
Channel: codeChannelEmailSetupRequired,
MaxAttempts: s.codeMaxAttempts,
}, s.codeTTL); err != nil {
return "", fmt.Errorf("store code: %w", err)
}
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, store.PhoneCode{Phone: phone, IssuedUserID: issuedUserID}); err != nil {
return "", err
}
return hash, nil
}
func (s *Service) createEmailLoginCode(ctx context.Context, phone, email string) (string, error) {
func (s *Service) createEmailLoginCode(ctx context.Context, phone, email string, issuedUserID int64) (string, error) {
hash, err := randomHex(8)
if err != nil {
return "", err
@ -327,22 +451,28 @@ func (s *Service) createEmailLoginCode(ctx context.Context, phone, email string)
return "", err
}
rec := store.PhoneCode{
Phone: phone,
Code: code,
Channel: codeChannelEmailLogin,
Email: strings.TrimSpace(email),
MaxAttempts: s.codeMaxAttempts,
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: issuedUserID,
Phone: phone,
Code: code,
Channel: codeChannelEmailLogin,
Email: strings.TrimSpace(email),
MaxAttempts: s.codeMaxAttempts,
}
if err := s.codes.Set(ctx, hash, rec, s.codeTTL); err != nil {
return "", fmt.Errorf("store email code: %w", err)
}
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
return "", err
}
if s.loginEmailSender == nil {
_ = s.codes.Del(ctx, hash)
return "", fmt.Errorf("login email sender is not configured")
return "", s.rollbackUndeliveredCode(ctx, hash, fmt.Errorf("login email sender is not configured"))
}
if err := s.loginEmailSender.SendLoginCode(ctx, rec.Email, code, s.codeTTL); err != nil {
_ = s.codes.Del(ctx, hash)
return "", fmt.Errorf("send login email code: %w", err)
return "", s.rollbackUndeliveredCode(ctx, hash, fmt.Errorf("send login email code: %w", err))
}
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
return "", err
}
return hash, nil
}
@ -396,20 +526,53 @@ func (s *Service) resendCode(ctx context.Context, authKeyID [8]byte, phone, phon
if rec.Phone != phone {
return "", ErrCodeInvalid
}
if rec.Purpose == store.PhoneCodePurposeChangePhone && (authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID) {
if rec.Purpose == store.PhoneCodePurposeChangePhone {
if authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID {
return "", ErrCodeInvalid
}
consumed, ok, err := s.codes.ConsumeScoped(ctx, phoneCodeHash, rec.Scope())
if err != nil {
return "", err
}
if !ok {
return "", ErrCodeExpired
}
return s.recreateChangePhoneCode(ctx, consumed)
}
if rec.Version != store.PhoneCodeVersionCurrent {
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
return "", ErrCodeExpired
}
if matches, err := s.issuedOwnerMatches(ctx, phone, rec.IssuedUserID); err != nil {
return "", err
} else if !matches {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return "", ErrCodeInvalid
}
_ = s.codes.Del(ctx, phoneCodeHash)
if rec.Purpose == store.PhoneCodePurposeChangePhone {
return s.recreateChangePhoneCode(ctx, rec)
consumed, ok, err := s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
if err != nil {
return "", err
}
if !ok {
return "", ErrCodeExpired
}
rec = consumed
if matches, err := s.issuedOwnerMatches(ctx, phone, rec.IssuedUserID); err != nil {
return "", err
} else if !matches {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return "", ErrCodeInvalid
}
if rec.Channel == codeChannelEmailLogin && strings.TrimSpace(rec.Email) != "" {
return s.createEmailLoginCode(ctx, phone, rec.Email)
return s.createEmailLoginCode(ctx, phone, rec.Email, rec.IssuedUserID)
}
if rec.Channel == codeChannelEmailSetupRequired {
return s.createSetupRequiredCode(ctx, phone)
return s.createSetupRequiredCode(ctx, phone, rec.IssuedUserID)
}
return s.SendCode(ctx, phone)
if rec.Channel != codeChannelPhone {
return "", ErrCodeInvalid
}
return s.createPhoneCode(ctx, phone, rec.IssuedUserID)
}
func (s *Service) recreateChangePhoneCode(ctx context.Context, rec store.PhoneCode) (string, error) {
@ -439,6 +602,75 @@ func (s *Service) CancelCodeForAuthKey(ctx context.Context, authKeyID [8]byte, p
return s.cancelCode(ctx, authKeyID, phone, phoneCodeHash)
}
// ConsumeLoginEmailReset authorizes auth.resetLoginEmail with the exact
// email-login hash previously issued for this phone owner. Possession of only
// a phone number is never sufficient to remove an authentication factor.
func (s *Service) ConsumeLoginEmailReset(ctx context.Context, phone, phoneCodeHash string) (int64, error) {
phone = normalizePhone(phone)
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
if err != nil {
return 0, err
}
if !found {
return 0, ErrCodeExpired
}
if rec.Version != store.PhoneCodeVersionCurrent {
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
return 0, ErrCodeExpired
}
if rec.Purpose != "" || rec.Phone != phone || rec.Channel != codeChannelEmailLogin || rec.SignUpVerified {
return 0, ErrCodeInvalid
}
before, beforeFound, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return 0, err
}
if !beforeFound || systemUserLoginForbidden(before) || rec.IssuedUserID == 0 || rec.IssuedUserID != before.ID {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return 0, ErrCodeInvalid
}
consumed, consumedOK, err := s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
if err != nil {
return 0, err
}
if !consumedOK {
return 0, ErrCodeExpired
}
if consumed.Channel != codeChannelEmailLogin || consumed.IssuedUserID != before.ID {
return 0, ErrCodeInvalid
}
after, afterFound, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return 0, err
}
if !afterFound || after.ID != before.ID {
return 0, ErrCodeInvalid
}
return before.ID, nil
}
// SendPhoneCodeAfterLoginEmailReset issues the replacement app code only for
// the exact user selected by ConsumeLoginEmailReset. It deliberately bypasses
// SendCode's phone→owner reclassification so an A→B transfer cannot send B a
// code and return that hash to A's reset flow.
func (s *Service) SendPhoneCodeAfterLoginEmailReset(ctx context.Context, phone string, expectedUserID int64) (string, error) {
phone = normalizePhone(phone)
if !validPhone(phone) {
return "", ErrPhoneNumberInvalid
}
if expectedUserID == 0 || systemLoginPhoneForbidden(phone) {
return "", ErrCodeInvalid
}
owner, found, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return "", err
}
if !found || owner.ID != expectedUserID || systemUserLoginForbidden(owner) {
return "", ErrCodeInvalid
}
return s.createPhoneCode(ctx, phone, expectedUserID)
}
func (s *Service) cancelCode(ctx context.Context, authKeyID [8]byte, phone, phoneCodeHash string) error {
phone = normalizePhone(phone)
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
@ -451,10 +683,37 @@ func (s *Service) cancelCode(ctx context.Context, authKeyID [8]byte, phone, phon
if rec.Phone != phone {
return ErrCodeInvalid
}
if rec.Purpose == store.PhoneCodePurposeChangePhone && (authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID) {
if rec.Purpose == store.PhoneCodePurposeChangePhone {
if authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID {
return ErrCodeInvalid
}
_, consumed, err := s.codes.ConsumeScoped(ctx, phoneCodeHash, rec.Scope())
if err != nil {
return err
}
if !consumed {
return ErrCodeExpired
}
return nil
}
if rec.Version != store.PhoneCodeVersionCurrent {
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
return ErrCodeExpired
}
if matches, err := s.issuedOwnerMatches(ctx, phone, rec.IssuedUserID); err != nil {
return err
} else if !matches {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return ErrCodeInvalid
}
return s.codes.Del(ctx, phoneCodeHash)
_, consumed, err := s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
if err != nil {
return err
}
if !consumed {
return ErrCodeExpired
}
return nil
}
// SignIn 校验验证码并尝试登录。
@ -464,102 +723,164 @@ func (s *Service) SignIn(ctx context.Context, auth domain.Authorization, phone,
if systemLoginPhoneForbidden(phone) {
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
}
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
if err != nil {
return domain.User{}, domain.Message{}, false, err
}
if !found {
return domain.User{}, domain.Message{}, false, ErrCodeExpired
}
if rec.Phone != phone || rec.Channel == codeChannelEmailSetupRequired {
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
}
if rec.Channel == codeChannelEmailLogin {
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
}
if rec.Code != code {
return domain.User{}, domain.Message{}, false, s.rejectCode(ctx, phoneCodeHash, rec, ErrCodeInvalid)
}
existing, found, err := s.users.ByPhone(ctx, phone)
if err != nil {
return domain.User{}, domain.Message{}, false, err
}
if !found {
return domain.User{}, domain.Message{}, true, nil // 验证码对、但需注册
}
return s.finishSignIn(ctx, auth, existing, phoneCodeHash, rec.Code)
}
// SignInWithEmail 处理带 email_verification 的 auth.signIn：账号设置了登录邮箱后，新设备
// 的验证码改投递到邮箱，客户端凭邮箱码（而非短信码）登录。开启真实登录邮箱后必须匹配
// 随机邮箱码；未开启该特性时仅保留旧开发路径的任意非空兼容。仍校验 phone_code_hash
// 有效、手机号匹配，并与短信登录共用 2FA 门控——即便走邮箱验证，开启了两步验证的账号
// 同样会停在 SESSION_PASSWORD_NEEDED。
func (s *Service) SignInWithEmail(ctx context.Context, auth domain.Authorization, phone, phoneCodeHash, code string) (domain.User, domain.Message, bool, error) {
phone = normalizePhone(phone)
if systemLoginPhoneForbidden(phone) {
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
}
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
if err != nil {
return domain.User{}, domain.Message{}, false, err
}
if !found {
return domain.User{}, domain.Message{}, false, ErrCodeExpired
}
if rec.Phone != phone {
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
}
if rec.Channel != codeChannelEmailLogin {
if s.loginEmailEnabled {
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
}
if strings.TrimSpace(code) == "" {
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
}
} else if rec.Code != strings.TrimSpace(code) {
return domain.User{}, domain.Message{}, false, s.rejectCode(ctx, phoneCodeHash, rec, ErrCodeInvalid)
}
existing, found, err := s.users.ByPhone(ctx, phone)
_, existing, found, err := s.verifyLoginCode(ctx, phone, phoneCodeHash, code, false)
if err != nil {
return domain.User{}, domain.Message{}, false, err
}
if !found {
return domain.User{}, domain.Message{}, true, nil
}
return s.finishSignIn(ctx, auth, existing, phoneCodeHash, rec.Code)
return s.finishSignIn(ctx, auth, existing)
}
// SignInWithEmail 处理带 email_verification 的 auth.signIn：账号设置了登录邮箱后，新设备
// 的验证码改投递到邮箱，客户端凭邮箱码（而非短信码）登录。开启真实登录邮箱后必须匹配
// 随机邮箱码；未开启该特性时仍允许旧客户端把 phone channel 放进
// email_verification，但必须精确匹配该 phone code，不能再接受任意非空值。
// 两条路径共用 owner 绑定、原子尝试计数与 2FA 门控。
func (s *Service) SignInWithEmail(ctx context.Context, auth domain.Authorization, phone, phoneCodeHash, code string) (domain.User, domain.Message, bool, error) {
phone = normalizePhone(phone)
if systemLoginPhoneForbidden(phone) {
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
}
_, existing, found, err := s.verifyLoginCode(ctx, phone, phoneCodeHash, strings.TrimSpace(code), true)
if err != nil {
return domain.User{}, domain.Message{}, false, err
}
if !found {
return domain.User{}, domain.Message{}, true, nil
}
return s.finishSignIn(ctx, auth, existing)
}
// verifyLoginCode closes the login-code state transition around one atomic
// CodeStore verification. The phone owner is read both before and after that
// linearization point. A hash issued for an unregistered number therefore can
// never authorize whichever account happens to acquire that number later.
func (s *Service) verifyLoginCode(ctx context.Context, phone, phoneCodeHash, code string, emailPath bool) (store.PhoneCode, domain.User, bool, error) {
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
if err != nil {
return store.PhoneCode{}, domain.User{}, false, err
}
if !found {
return store.PhoneCode{}, domain.User{}, false, ErrCodeExpired
}
if rec.Version != store.PhoneCodeVersionCurrent {
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
return store.PhoneCode{}, domain.User{}, false, ErrCodeExpired
}
if rec.Phone != phone || rec.Purpose != "" {
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
channelAllowed := rec.Channel == codeChannelPhone && !emailPath
if emailPath {
channelAllowed = rec.Channel == codeChannelEmailLogin || (!s.loginEmailEnabled && rec.Channel == codeChannelPhone)
}
if !channelAllowed {
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
before, beforeFound, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return store.PhoneCode{}, domain.User{}, false, err
}
beforeUserID := int64(0)
if beforeFound {
if systemUserLoginForbidden(before) {
return store.PhoneCode{}, domain.User{}, false, ErrSystemUserLoginForbidden
}
beforeUserID = before.ID
}
if rec.IssuedUserID != beforeUserID {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
// A verified sign-up marker may precede auth.signIn on the email-setup
// path, and a normal signIn response can be lost and retried. The marker is
// already the durable authorization fact; return signUpRequired
// idempotently without asking CodeStore to verify it a second time.
if rec.SignUpVerified {
if beforeFound || rec.IssuedUserID != 0 || subtle.ConstantTimeCompare([]byte(rec.Code), []byte(code)) != 1 {
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
after, afterFound, err := s.currentPhoneOwner(ctx, phone)
if err != nil {
return store.PhoneCode{}, domain.User{}, false, err
}
if afterFound || after.ID != 0 {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
return rec, domain.User{}, false, nil
}
result, err := s.codes.VerifyLogin(ctx, phoneCodeHash, phone, code, !beforeFound, s.codeMaxAttempts)
if err != nil {
return store.PhoneCode{}, domain.User{}, false, err
}
after, afterFound, ownerErr := s.currentPhoneOwner(ctx, phone)
if ownerErr != nil {
return store.PhoneCode{}, domain.User{}, false, ownerErr
}
afterUserID := int64(0)
if afterFound {
afterUserID = after.ID
}
recordOwnerMismatch := result.Status != store.LoginCodeVerifyMissing && result.Record.IssuedUserID != rec.IssuedUserID
if beforeUserID != afterUserID || recordOwnerMismatch {
// keepForSignUp may have left a verified marker behind. Remove it on
// owner drift so a later transfer-back cannot resurrect authorization.
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
switch result.Status {
case store.LoginCodeVerifyMissing:
return store.PhoneCode{}, domain.User{}, false, ErrCodeExpired
case store.LoginCodeVerifyInvalid:
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
case store.LoginCodeVerifyAccepted:
if result.Record.Version != store.PhoneCodeVersionCurrent || result.Record.Phone != phone || result.Record.IssuedUserID != afterUserID {
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
if afterFound && systemUserLoginForbidden(after) {
return store.PhoneCode{}, domain.User{}, false, ErrSystemUserLoginForbidden
}
return result.Record, after, afterFound, nil
default:
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
}
}
// finishSignIn 是短信/邮箱两条登录路径在「验证码已通过、用户已存在」之后的共用收尾：
// 处理 2FA password_pending 绑定、写登录消息、消费验证码。
func (s *Service) finishSignIn(ctx context.Context, auth domain.Authorization, existing domain.User, phoneCodeHash, loginCode string) (domain.User, domain.Message, bool, error) {
// 验证码已由 VerifyLogin 原子消费；这里只处理 2FA password_pending 绑定。已有账号的 app-code 消息已在
// SendCode/ResendCode 返回前持久化与入 outbox，这里绝不能再创建或补发；
// 否则未完成登录/2FA 的真实验证码反而不会及时到达旧设备。
func (s *Service) finishSignIn(ctx context.Context, auth domain.Authorization, existing domain.User) (domain.User, domain.Message, bool, error) {
if systemUserLoginForbidden(existing) {
_ = s.codes.Del(ctx, phoneCodeHash)
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
}
// 开启两步验证的账号：把授权标记为 password_pending 再写入，业务鉴权据此拒绝该 auth_key，
// 直到 auth.checkPassword 通过。绝不能先以完全授权写入再返回 SESSION_PASSWORD_NEEDED，
// 否则客户端忽略该错误即可直接调用业务 RPC 绕过两步验证。
passwordNeeded := s.passwordNeeded(ctx, existing.ID)
passwordNeeded, err := s.passwordNeeded(ctx, existing.ID)
if err != nil {
// Password state is part of the authentication decision. Treat store
// failures as fail-closed and leave the auth key entirely unbound.
return domain.User{}, domain.Message{}, false, err
}
auth.PasswordPending = passwordNeeded
if err := s.bind(ctx, auth, existing.ID); err != nil {
return domain.User{}, domain.Message{}, false, err
}
if passwordNeeded {
_ = s.codes.Del(ctx, phoneCodeHash)
return existing, domain.Message{}, false, domain.ErrSessionPasswordNeeded
}
loginMessage, err := s.recordLoginMessage(ctx, existing.ID, loginCode)
if err != nil {
return domain.User{}, domain.Message{}, false, err
}
_ = s.codes.Del(ctx, phoneCodeHash)
return existing, loginMessage, false, nil
return existing, domain.Message{}, false, nil
}
// SignUp 在 SignIn 判定需注册后创建用户并绑定授权。
// signUp 的 TL 请求不带验证码，这里校验 phone_code_hash 仍有效且手机号匹配。
// signUp 的 TL 请求不带验证码，因此只消费由正确 SignIn/email setup 原子
// 标记过的 hash。直接 SendCode→SignUp 永远不能创建账号。
func (s *Service) SignUp(ctx context.Context, auth domain.Authorization, phone, phoneCodeHash, firstName, lastName string) (domain.User, domain.Message, error) {
phone = normalizePhone(phone)
if !validPhone(phone) {
@ -580,15 +901,48 @@ func (s *Service) SignUp(ctx context.Context, auth domain.Authorization, phone,
if !found {
return domain.User{}, domain.Message{}, ErrCodeExpired
}
if rec.Phone != phone {
if rec.Version != store.PhoneCodeVersionCurrent {
_, _, _ = s.codes.ConsumeSignUpVerified(ctx, phoneCodeHash, phone)
return domain.User{}, domain.Message{}, ErrCodeExpired
}
if rec.Phone != phone || rec.Purpose != "" {
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
if rec.Channel == codeChannelEmailSetupRequired {
if !rec.SignUpVerified {
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
if rec.IssuedUserID != 0 {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
if rec.Channel != codeChannelPhone && rec.Channel != codeChannelEmailLogin {
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
if s.loginEmailRequireSetup && !rec.VerifiedEmail && strings.TrimSpace(rec.PendingEmail) == "" {
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
if current, currentFound, err := s.currentPhoneOwner(ctx, phone); err != nil {
return domain.User{}, domain.Message{}, err
} else if currentFound || current.ID != 0 {
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
consumed, consumedOK, err := s.codes.ConsumeSignUpVerified(ctx, phoneCodeHash, phone)
if err != nil {
return domain.User{}, domain.Message{}, err
}
if !consumedOK {
return domain.User{}, domain.Message{}, ErrCodeExpired
}
rec = consumed
if rec.IssuedUserID != 0 || !rec.SignUpVerified || (rec.Channel != codeChannelPhone && rec.Channel != codeChannelEmailLogin) {
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
if current, currentFound, err := s.currentPhoneOwner(ctx, phone); err != nil {
return domain.User{}, domain.Message{}, err
} else if currentFound || current.ID != 0 {
return domain.User{}, domain.Message{}, ErrCodeInvalid
}
accessHash, err := randomInt64()
if err != nil {
@ -610,18 +964,22 @@ func (s *Service) SignUp(ctx context.Context, auth domain.Authorization, phone,
return domain.User{}, domain.Message{}, err
}
if rec.VerifiedEmail && strings.TrimSpace(rec.PendingEmail) != "" && s.loginEmails != nil {
if err := s.loginEmails.SetLoginEmailByPhone(ctx, phone, rec.PendingEmail); err != nil {
if err := s.loginEmails.SetLoginEmail(ctx, u.ID, rec.PendingEmail); err != nil {
return domain.User{}, domain.Message{}, err
}
}
if err := s.bind(ctx, auth, u.ID); err != nil {
return domain.User{}, domain.Message{}, err
}
loginMessage, err := s.recordLoginMessage(ctx, u.ID, rec.Code)
if err != nil {
return domain.User{}, domain.Message{}, err
loginMessage := domain.Message{}
// SMTP setup/login codes are secret factors, not 777000 app messages. Only
// the normal phone/app-code registration path creates the bootstrap dialog.
if rec.Channel == codeChannelPhone {
loginMessage, err = s.recordLoginMessage(ctx, u.ID, rec.Code)
if err != nil {
return domain.User{}, domain.Message{}, err
}
}
_ = s.codes.Del(ctx, phoneCodeHash)
return u, loginMessage, nil
}
@ -865,15 +1223,21 @@ func (s *Service) authorizationsByUserExcept(ctx context.Context, userID int64,
func (s *Service) bind(ctx context.Context, auth domain.Authorization, userID int64) error {
auth.UserID = userID
// Bind 是授权切换的持久化状态边界：生产 store 会先清同 auth key 的旧用户
// update state，再原子建立新用户 baseline。RPC 层不得在 Bind 成功后清整个 key，
// 否则会把刚建立的 retained-floor checkpoint 一并删除。
return s.auths.Bind(ctx, auth)
}
func (s *Service) passwordNeeded(ctx context.Context, userID int64) bool {
func (s *Service) passwordNeeded(ctx context.Context, userID int64) (bool, error) {
if s.passwords == nil {
return false
return false, nil
}
settings, found, err := s.passwords.GetByUser(ctx, userID)
return err == nil && found && settings.HasPassword
if err != nil {
return false, err
}
return found && settings.HasPassword, nil
}
const loginMessageTpl = `Login code: %s. Do not give this code to anyone, even if they say they are from Telegram!
@ -1005,20 +1369,6 @@ func authKeyIDInt64(id [8]byte) int64 {
return int64(binary.LittleEndian.Uint64(id[:]))
}
func (s *Service) rejectCode(ctx context.Context, hash string, rec store.PhoneCode, ret error) error {
rec.Attempts++
max := rec.MaxAttempts
if max <= 0 {
max = s.codeMaxAttempts
}
if max > 0 && rec.Attempts >= max {
_ = s.codes.Del(ctx, hash)
return ret
}
_ = s.codes.Update(ctx, hash, rec)
return ret
}
func normalizePhone(phone string) string {
return domain.NormalizePhone(phone)
}

104
internal/app/auth/service_test.go
View file

@ -178,6 +178,14 @@ func TestPhoneCodeAcceptsTDesktopDigitsOnlySignIn(t *testing.T) {
}
}
func verifyCodeForSignUp(t *testing.T, svc *Service, phone, hash, code string) {
t.Helper()
got, msg, needSignUp, err := svc.SignIn(context.Background(), domain.Authorization{}, phone, hash, code)
if err != nil || !needSignUp || got.ID != 0 || msg.ID != 0 {
t.Fatalf("SignIn before SignUp user=%+v message=%+v needSignUp=%v err=%v, want empty/empty/true/nil", got, msg, needSignUp, err)
}
}
func TestSystemUserPhoneCannotLoginOrSignUp(t *testing.T) {
ctx := context.Background()
codes := memory.NewCodeStore()
@ -257,6 +265,7 @@ func TestMultipleAuthKeysKeepSeparateUsers(t *testing.T) {
if err != nil {
t.Fatalf("SendCode user1: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550005001", hash1, "12345")
user1, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key1}, "+15550005001", hash1, "One", "")
if err != nil {
t.Fatalf("SignUp user1: %v", err)
@ -265,6 +274,7 @@ func TestMultipleAuthKeysKeepSeparateUsers(t *testing.T) {
if err != nil {
t.Fatalf("SendCode user2: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550005002", hash2, "12345")
user2, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key2}, "+15550005002", hash2, "Two", "")
if err != nil {
t.Fatalf("SignUp user2: %v", err)
@ -294,6 +304,7 @@ func TestLogOutThenSignInSameAuthKeySwitchesUser(t *testing.T) {
if err != nil {
t.Fatalf("SendCode user1: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550006001", hash1, "12345")
user1, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550006001", hash1, "One", "")
if err != nil {
t.Fatalf("SignUp user1: %v", err)
@ -312,6 +323,7 @@ func TestLogOutThenSignInSameAuthKeySwitchesUser(t *testing.T) {
if err != nil {
t.Fatalf("SendCode user2: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550006002", hash2, "12345")
user2, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550006002", hash2, "Two", "")
if err != nil {
t.Fatalf("SignUp user2: %v", err)
@ -337,6 +349,7 @@ func TestResetAuthorizationDeletesProtocolAuthKey(t *testing.T) {
if err != nil {
t.Fatalf("SendCode: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550007001", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550007001", hash, "One", "")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -375,6 +388,7 @@ func TestResetAuthorizationsDeletesOnlyRevokedProtocolAuthKeys(t *testing.T) {
if err != nil {
t.Fatalf("SendCode: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550007002", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: keep}, "+15550007002", hash, "Two", "")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -399,16 +413,27 @@ func TestSignUpWritesOfficialLoginMessage(t *testing.T) {
ctx := context.Background()
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithLoginMessages(messages, dialogs))
delivery := &captureLoginCodeDelivery{}
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
WithLoginMessages(messages, dialogs),
WithLoginCodeDelivery(delivery),
)
hash, err := svc.SendCode(ctx, "+15550004311")
if err != nil {
t.Fatalf("SendCode: %v", err)
}
if len(delivery.requests) != 0 {
t.Fatalf("unregistered SendCode delivered before user exists: %+v", delivery.requests)
}
verifyCodeForSignUp(t, svc, "+15550004311", hash, "12345")
u, msg, err := svc.SignUp(ctx, domain.Authorization{}, "+15550004311", hash, "Test", "User")
if err != nil {
t.Fatalf("SignUp: %v", err)
}
if len(delivery.requests) != 0 {
t.Fatalf("SignUp unexpectedly used existing-account delivery: %+v", delivery.requests)
}
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
if err != nil {
@ -431,17 +456,23 @@ func TestSignUpWritesOfficialLoginMessage(t *testing.T) {
}
}
func TestSignInLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
func TestSendCodeLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
ctx := context.Background()
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithLoginMessages(messages, dialogs))
events := memory.NewUpdateEventStore()
delivery := memory.NewLoginCodeDeliveryStore(messages, events)
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
WithLoginMessages(messages, dialogs),
WithLoginCodeDelivery(delivery),
)
phone := "+15550004312"
hash, err := svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode signup: %v", err)
}
verifyCodeForSignUp(t, svc, phone, hash, "12345")
u, first, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Test", "User")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -452,15 +483,6 @@ func TestSignInLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
} else if read.MaxID != first.ID || read.StillUnreadCount != 0 {
t.Fatalf("read first login message = %+v, want max_id %d unread 0", read, first.ID)
}
hash, err = svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode signin second: %v", err)
}
_, second, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
if err != nil || needSignUp {
t.Fatalf("SignIn second needSignUp=%v err=%v", needSignUp, err)
}
assertOfficialDialog := func(wantTop, wantRead, wantUnread int) {
t.Helper()
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
@ -475,23 +497,67 @@ func TestSignInLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
t.Fatalf("dialog = %+v, want top=%d read=%d unread=%d", got, wantTop, wantRead, wantUnread)
}
}
latestLoginMessage := func(wantCount int) domain.Message {
t.Helper()
history, err := messages.ListByUser(ctx, u.ID, domain.MessageFilter{
HasPeer: true,
Peer: peer,
Limit: 10,
})
if err != nil || len(history.Messages) != wantCount {
t.Fatalf("official history count=%d err=%v, want %d", len(history.Messages), err, wantCount)
}
latest := history.Messages[0]
for _, msg := range history.Messages[1:] {
if msg.ID > latest.ID {
latest = msg
}
}
return latest
}
hash, err = svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode signin second: %v", err)
}
second := latestLoginMessage(2)
// 核心时序：SendCode 返回时 message/dialog/unread 已提交，尚未 SignIn。
assertOfficialDialog(second.ID, first.ID, 1)
_, signInMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
if err != nil || needSignUp {
t.Fatalf("SignIn second needSignUp=%v err=%v", needSignUp, err)
}
if signInMessage.ID != 0 {
t.Fatalf("SignIn second returned a late login message %+v", signInMessage)
}
assertOfficialDialog(second.ID, first.ID, 1)
hash, err = svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode signin third: %v", err)
}
_, third, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
third := latestLoginMessage(3)
assertOfficialDialog(third.ID, first.ID, 2)
_, signInMessage, needSignUp, err = svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
if err != nil || needSignUp {
t.Fatalf("SignIn third needSignUp=%v err=%v", needSignUp, err)
}
if signInMessage.ID != 0 {
t.Fatalf("SignIn third returned a late login message %+v", signInMessage)
}
assertOfficialDialog(third.ID, first.ID, 2)
}
func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
ctx := context.Background()
passwords := memory.NewPasswordStore()
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithPasswords(passwords))
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
delivery := memory.NewLoginCodeDeliveryStore(messages, memory.NewUpdateEventStore())
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
WithPasswords(passwords),
WithLoginCodeDelivery(delivery),
)
var key [8]byte
key[0] = 7
@ -499,6 +565,7 @@ func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
if err != nil {
t.Fatalf("SendCode signup: %v", err)
}
verifyCodeForSignUp(t, svc, "+15550004312", hash, "12345")
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550004312", hash, "Two", "Factor")
if err != nil {
t.Fatalf("SignUp: %v", err)
@ -514,13 +581,16 @@ func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
if err != nil {
t.Fatalf("SendCode signin: %v", err)
}
got, _, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "+15550004312", hash, "12345")
got, signInMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "+15550004312", hash, "12345")
if !errors.Is(err, domain.ErrSessionPasswordNeeded) {
t.Fatalf("SignIn err = %v, want ErrSessionPasswordNeeded", err)
}
if needSignUp || got.ID != u.ID {
t.Fatalf("SignIn user=%+v needSignUp=%v, want existing 2FA user", got, needSignUp)
}
if signInMessage.ID != 0 {
t.Fatalf("2FA SignIn returned a late login message %+v", signInMessage)
}
// 两步验证未完成：业务鉴权（UserID）必须视为未登录，避免绕过 2FA。
bound, found, err := svc.UserID(ctx, key)
if err != nil || found || bound != 0 {
@ -539,6 +609,10 @@ func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
if err != nil || !found || bound != u.ID {
t.Fatalf("UserID after 2FA passed = %d found=%v err=%v, want %d", bound, found, err, u.ID)
}
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
if err != nil || len(list.Messages) != 1 {
t.Fatalf("2FA login-code messages after password = %+v err=%v, want exactly the SendCode message", list.Messages, err)
}
}
func testAuthKey(seed byte) mtcrypto.AuthKey {

554
internal/app/auth/signup_state_test.go Normal file
View file

@ -0,0 +1,554 @@
package auth
import (
"context"
"errors"
"sync"
"testing"
"time"
accountapp "telesrv/internal/app/account"
"telesrv/internal/domain"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
func TestSignUpRequiresCorrectSignInAndConsumesMarkerOnce(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
codes := memory.NewCodeStore()
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
phone := "15550009301"
hash, err := svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
if _, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Direct", "Bypass"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("direct SignUp err=%v, want ErrCodeInvalid", err)
}
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "00000"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("wrong SignIn err=%v, want ErrCodeInvalid", err)
}
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || rec.SignUpVerified {
t.Fatalf("wrong code marker=%v found=%v err=%v, want live/unverified", rec.SignUpVerified, found, err)
}
if _, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Wrong", "Code"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("SignUp after wrong code err=%v, want ErrCodeInvalid", err)
}
verifyCodeForSignUp(t, svc, phone, hash, "12345")
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || !rec.SignUpVerified || rec.IssuedUserID != 0 {
t.Fatalf("verified record=%+v found=%v err=%v", rec, found, err)
}
if _, msg, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); err != nil || !needSignUp || msg.ID != 0 {
t.Fatalf("idempotent SignIn needSignUp=%v message=%+v err=%v", needSignUp, msg, err)
}
u, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Verified", "User")
if err != nil || u.Phone != phone {
t.Fatalf("verified SignUp user=%+v err=%v", u, err)
}
if _, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Replay", "User"); !errors.Is(err, ErrCodeExpired) {
t.Fatalf("replayed SignUp err=%v, want ErrCodeExpired", err)
}
}
func TestConcurrentSignUpConsumesVerifiedHashExactlyOnce(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
svc := NewService(users, memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345")
phone := "15550009302"
hash, err := svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
verifyCodeForSignUp(t, svc, phone, hash, "12345")
const workers = 16
start := make(chan struct{})
errs := make(chan error, workers)
for i := 0; i < workers; i++ {
go func(i int) {
<-start
var key [8]byte
key[0] = byte(i + 1)
_, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, phone, hash, "Concurrent", "User")
errs <- err
}(i)
}
close(start)
successes := 0
for i := 0; i < workers; i++ {
err := <-errs
switch {
case err == nil:
successes++
case errors.Is(err, ErrCodeExpired), errors.Is(err, ErrCodeInvalid):
default:
t.Fatalf("concurrent SignUp err=%v", err)
}
}
if successes != 1 {
t.Fatalf("successful SignUp calls=%d, want 1", successes)
}
}
type afterVerifyCodeStore struct {
store.CodeStore
once sync.Once
afterVerify func()
}
type failingPasswordStore struct {
store.PasswordStore
err error
}
func (s *failingPasswordStore) GetByUser(context.Context, int64) (domain.PasswordSettings, bool, error) {
return domain.PasswordSettings{}, false, s.err
}
type switchablePhoneOwnerStore struct {
store.UserStore
mu sync.RWMutex
phone string
override bool
owner domain.User
found bool
}
func (s *switchablePhoneOwnerStore) ByPhone(ctx context.Context, phone string) (domain.User, bool, error) {
s.mu.RLock()
if s.override && domain.NormalizePhone(phone) == s.phone {
owner, found := s.owner, s.found
s.mu.RUnlock()
return owner, found, nil
}
s.mu.RUnlock()
return s.UserStore.ByPhone(ctx, phone)
}
func (s *switchablePhoneOwnerStore) setOwnerView(phone string, owner domain.User, found bool) {
s.mu.Lock()
s.phone = domain.NormalizePhone(phone)
s.owner = owner
s.found = found
s.override = true
s.mu.Unlock()
}
func (s *switchablePhoneOwnerStore) resetOwnerView() {
s.mu.Lock()
s.override = false
s.mu.Unlock()
}
func (s *afterVerifyCodeStore) VerifyLogin(ctx context.Context, hash, phone, code string, keep bool, maxAttempts int) (store.LoginCodeVerifyResult, error) {
result, err := s.CodeStore.VerifyLogin(ctx, hash, phone, code, keep, maxAttempts)
if err == nil && result.Status == store.LoginCodeVerifyAccepted && s.afterVerify != nil {
s.once.Do(s.afterVerify)
}
return result, err
}
func TestOwnerTransferAcrossVerifyInvalidatesHashPermanently(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
baseCodes := memory.NewCodeStore()
var createErr error
codes := &afterVerifyCodeStore{CodeStore: baseCodes}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
phone := "15550009303"
hash, err := svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
codes.afterVerify = func() {
_, createErr = users.Create(ctx, domain.User{Phone: phone, FirstName: "NewOwner"})
}
if _, _, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); !errors.Is(err, ErrCodeInvalid) || needSignUp {
t.Fatalf("SignIn across owner transfer needSignUp=%v err=%v, want invalid", needSignUp, err)
}
if createErr != nil {
t.Fatalf("create concurrent owner: %v", createErr)
}
if _, found, err := baseCodes.Get(ctx, hash); err != nil || found {
t.Fatalf("owner-drift hash found=%v err=%v, want invalidated", found, err)
}
}
func TestPasswordLookupFailureNeverCreatesOrChangesAuthorization(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
target, err := users.Create(ctx, domain.User{Phone: "15550009320", FirstName: "Target"})
if err != nil {
t.Fatalf("create target: %v", err)
}
previous, err := users.Create(ctx, domain.User{Phone: "15550009321", FirstName: "Previous"})
if err != nil {
t.Fatalf("create previous: %v", err)
}
authz := memory.NewAuthorizationStore()
lookupErr := errors.New("password store unavailable")
passwords := &failingPasswordStore{PasswordStore: memory.NewPasswordStore(), err: lookupErr}
svc := NewService(users, authz, memory.NewCodeStore(), nil, nil, "12345",
WithPasswords(passwords),
WithLoginCodeDelivery(&captureLoginCodeDelivery{}),
)
t.Run("unbound-key-remains-unbound", func(t *testing.T) {
key := [8]byte{0xC1}
hash, err := svc.SendCode(ctx, target.Phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, target.Phone, hash, "12345"); !errors.Is(err, lookupErr) {
t.Fatalf("SignIn err=%v, want password lookup failure", err)
}
if got, found, err := authz.ByAuthKey(ctx, key); err != nil || found {
t.Fatalf("authorization=%+v found=%v err=%v, want absent", got, found, err)
}
})
t.Run("previous-binding-remains-unchanged", func(t *testing.T) {
key := [8]byte{0xC2}
original := domain.Authorization{AuthKeyID: key, UserID: previous.ID, Hash: 987654321}
if err := authz.Bind(ctx, original); err != nil {
t.Fatalf("bind previous authorization: %v", err)
}
hash, err := svc.SendCode(ctx, target.Phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, target.Phone, hash, "12345"); !errors.Is(err, lookupErr) {
t.Fatalf("SignIn err=%v, want password lookup failure", err)
}
got, found, err := authz.ByAuthKey(ctx, key)
if err != nil || !found || got.UserID != previous.ID || got.Hash != original.Hash || got.PasswordPending != original.PasswordPending {
t.Fatalf("authorization after failure=%+v found=%v err=%v, want unchanged %+v", got, found, err, original)
}
})
}
func TestOwnerTransferAwayAndBackCannotReviveLoginHash(t *testing.T) {
ctx := context.Background()
t.Run("unregistered-signin", func(t *testing.T) {
baseUsers := memory.NewUserStore()
other, err := baseUsers.Create(ctx, domain.User{Phone: "15550009311", FirstName: "Other"})
if err != nil {
t.Fatalf("create other owner: %v", err)
}
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
codes := memory.NewCodeStore()
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
phone := "15550009310"
hash, err := svc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
users.setOwnerView(phone, other, true)
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("SignIn after 0->B owner transfer err=%v, want invalid", err)
}
users.resetOwnerView()
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); !errors.Is(err, ErrCodeExpired) {
t.Fatalf("SignIn after 0->B->0 err=%v, want expired", err)
}
})
t.Run("existing-resend", func(t *testing.T) {
baseUsers := memory.NewUserStore()
ownerA, err := baseUsers.Create(ctx, domain.User{Phone: "15550009312", FirstName: "A"})
if err != nil {
t.Fatalf("create owner A: %v", err)
}
ownerB, err := baseUsers.Create(ctx, domain.User{Phone: "15550009313", FirstName: "B"})
if err != nil {
t.Fatalf("create owner B: %v", err)
}
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
codes := memory.NewCodeStore()
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(&captureLoginCodeDelivery{}))
hash, err := svc.SendCode(ctx, ownerA.Phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
users.setOwnerView(ownerA.Phone, ownerB, true)
if _, err := svc.ResendCode(ctx, ownerA.Phone, hash); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("ResendCode after A->B err=%v, want invalid", err)
}
users.resetOwnerView()
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, ownerA.Phone, hash, "12345"); !errors.Is(err, ErrCodeExpired) {
t.Fatalf("SignIn after A->B->A err=%v, want expired", err)
}
})
t.Run("existing-cancel", func(t *testing.T) {
baseUsers := memory.NewUserStore()
ownerA, err := baseUsers.Create(ctx, domain.User{Phone: "15550009314", FirstName: "A"})
if err != nil {
t.Fatalf("create owner A: %v", err)
}
ownerB, err := baseUsers.Create(ctx, domain.User{Phone: "15550009315", FirstName: "B"})
if err != nil {
t.Fatalf("create owner B: %v", err)
}
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
codes := memory.NewCodeStore()
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(&captureLoginCodeDelivery{}))
hash, err := svc.SendCode(ctx, ownerA.Phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
users.setOwnerView(ownerA.Phone, ownerB, true)
if err := svc.CancelCode(ctx, ownerA.Phone, hash); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("CancelCode after A->B err=%v, want invalid", err)
}
users.resetOwnerView()
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, ownerA.Phone, hash, "12345"); !errors.Is(err, ErrCodeExpired) {
t.Fatalf("SignIn after canceled A->B->A err=%v, want expired", err)
}
})
}
func TestEmailSetupVerificationAuthorizesSignUpWithout777000Message(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
codes := memory.NewCodeStore()
passwords := memory.NewPasswordStore()
sender := &testMailSender{}
accountSvc := accountapp.NewService(passwords,
accountapp.WithUsers(users),
accountapp.WithLoginEmailVerification(codes, sender, time.Minute, 3, 6),
)
dialogs := memory.NewDialogStore()
messages := memory.NewMessageStore(dialogs)
authSvc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345",
WithLoginMessages(messages, dialogs),
WithLoginEmail(LoginEmailOptions{
Enabled: true,
RequireSetup: true,
CodeLength: 6,
Store: accountSvc,
Sender: sender,
}),
)
phone := "15550009304"
hash, err := authSvc.SendCode(ctx, phone)
if err != nil {
t.Fatalf("SendCode: %v", err)
}
if _, _, err := authSvc.SignUp(ctx, domain.Authorization{}, phone, hash, "Direct", "Email"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("SignUp before email setup err=%v, want ErrCodeInvalid", err)
}
if _, _, err := accountSvc.SendLoginEmailCode(ctx, 0, phone, hash, "new@example.test", true); err != nil {
t.Fatalf("SendLoginEmailCode: %v", err)
}
bad := wrongCode(sender.code, '0')
if _, err := accountSvc.VerifyLoginEmail(ctx, 0, phone, hash, bad, true); !errors.Is(err, domain.ErrEmailCodeInvalid) {
t.Fatalf("wrong VerifyLoginEmail err=%v, want ErrEmailCodeInvalid", err)
}
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || rec.SignUpVerified {
t.Fatalf("wrong SMTP code marker=%v found=%v err=%v", rec.SignUpVerified, found, err)
}
if _, err := accountSvc.VerifyLoginEmail(ctx, 0, phone, hash, sender.code, true); err != nil {
t.Fatalf("VerifyLoginEmail: %v", err)
}
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || !rec.SignUpVerified || rec.Channel != codeChannelEmailLogin {
t.Fatalf("email-verified phone code=%+v found=%v err=%v", rec, found, err)
}
if _, msg, needSignUp, err := authSvc.SignInWithEmail(ctx, domain.Authorization{}, phone, hash, sender.code); err != nil || !needSignUp || msg.ID != 0 {
t.Fatalf("SignInWithEmail after setup needSignUp=%v message=%+v err=%v", needSignUp, msg, err)
}
u, msg, err := authSvc.SignUp(ctx, domain.Authorization{}, phone, hash, "Email", "User")
if err != nil {
t.Fatalf("SignUp after email setup: %v", err)
}
if msg.ID != 0 || msg.Body != "" {
t.Fatalf("email SignUp returned SMTP code message: %+v", msg)
}
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
if err != nil {
t.Fatalf("ListByUser: %v", err)
}
if len(list.Dialogs) != 0 || len(list.Messages) != 0 {
t.Fatalf("email SignUp created 777000 bootstrap state: dialogs=%+v messages=%+v", list.Dialogs, list.Messages)
}
if email, found, err := accountSvc.LoginEmailByPhone(ctx, phone); err != nil || !found || email != "new@example.test" {
t.Fatalf("LoginEmailByPhone email=%q found=%v err=%v", email, found, err)
}
}
func TestConsumeLoginEmailResetRequiresExactIssuedHash(t *testing.T) {
ctx := context.Background()
baseUsers := memory.NewUserStore()
owner, err := baseUsers.Create(ctx, domain.User{Phone: "15550009330", FirstName: "Owner"})
if err != nil {
t.Fatalf("create owner: %v", err)
}
other, err := baseUsers.Create(ctx, domain.User{Phone: "15550009331", FirstName: "Other"})
if err != nil {
t.Fatalf("create other: %v", err)
}
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
codes := memory.NewCodeStore()
delivery := &captureLoginCodeDelivery{}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
seed := func(hash, channel string) {
t.Helper()
if err := codes.Set(ctx, hash, store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: owner.ID,
Phone: owner.Phone,
Code: "654321",
Channel: channel,
MaxAttempts: 5,
}, time.Minute); err != nil {
t.Fatalf("seed %s: %v", hash, err)
}
}
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "arbitrary-missing"); !errors.Is(err, ErrCodeExpired) {
t.Fatalf("arbitrary hash err=%v, want expired", err)
}
seed("wrong-phone", codeChannelEmailLogin)
if _, err := svc.ConsumeLoginEmailReset(ctx, other.Phone, "wrong-phone"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("wrong phone err=%v, want invalid", err)
}
if _, found, err := codes.Get(ctx, "wrong-phone"); err != nil || !found {
t.Fatalf("wrong-phone probe destroyed valid hash found=%v err=%v", found, err)
}
seed("wrong-channel", codeChannelPhone)
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "wrong-channel"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("wrong channel err=%v, want invalid", err)
}
seed("owner-drift", codeChannelEmailLogin)
users.setOwnerView(owner.Phone, other, true)
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "owner-drift"); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("A->B reset err=%v, want invalid", err)
}
users.resetOwnerView()
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "owner-drift"); !errors.Is(err, ErrCodeExpired) {
t.Fatalf("A->B->A reset err=%v, want expired", err)
}
seed("successful-reset", codeChannelEmailLogin)
resetUserID, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "successful-reset")
if err != nil || resetUserID != owner.ID {
t.Fatalf("successful reset consume uid=%d err=%v", resetUserID, err)
}
replacementHash, err := svc.SendPhoneCodeAfterLoginEmailReset(ctx, owner.Phone, resetUserID)
if err != nil || replacementHash == "" {
t.Fatalf("replacement hash=%q err=%v", replacementHash, err)
}
if len(delivery.requests) != 1 || delivery.requests[0].UserID != owner.ID || delivery.requests[0].PhoneCodeHash != replacementHash {
t.Fatalf("replacement delivery=%+v", delivery.requests)
}
if rec, found, err := codes.Get(ctx, replacementHash); err != nil || !found || rec.Version != store.PhoneCodeVersionCurrent || rec.IssuedUserID != owner.ID || rec.Channel != codeChannelPhone {
t.Fatalf("replacement code=%+v found=%v err=%v", rec, found, err)
}
}
func TestConcurrentLoginEmailResetHasSingleConsumer(t *testing.T) {
ctx := context.Background()
users := memory.NewUserStore()
owner, err := users.Create(ctx, domain.User{Phone: "15550009332", FirstName: "Owner"})
if err != nil {
t.Fatalf("create owner: %v", err)
}
codes := memory.NewCodeStore()
hash := "concurrent-email-reset"
if err := codes.Set(ctx, hash, store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: owner.ID,
Phone: owner.Phone,
Code: "654321",
Channel: codeChannelEmailLogin,
MaxAttempts: 5,
}, time.Minute); err != nil {
t.Fatalf("seed code: %v", err)
}
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
const workers = 24
start := make(chan struct{})
errs := make(chan error, workers)
for i := 0; i < workers; i++ {
go func() {
<-start
_, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, hash)
errs <- err
}()
}
close(start)
successes := 0
for i := 0; i < workers; i++ {
err := <-errs
if err == nil {
successes++
continue
}
if !errors.Is(err, ErrCodeExpired) {
t.Fatalf("concurrent reset err=%v", err)
}
}
if successes != 1 {
t.Fatalf("successful reset consumers=%d, want 1", successes)
}
}
func TestLoginEmailResetLocksUserAcrossOwnerTransfer(t *testing.T) {
ctx := context.Background()
baseUsers := memory.NewUserStore()
ownerA, err := baseUsers.Create(ctx, domain.User{Phone: "15550009340", FirstName: "A"})
if err != nil {
t.Fatalf("create A: %v", err)
}
ownerB, err := baseUsers.Create(ctx, domain.User{Phone: "15550009341", FirstName: "B"})
if err != nil {
t.Fatalf("create B: %v", err)
}
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
passwords := memory.NewPasswordStore()
accountSvc := accountapp.NewService(passwords, accountapp.WithUsers(users))
if err := accountSvc.SetLoginEmail(ctx, ownerA.ID, "a@example.test"); err != nil {
t.Fatalf("SetLoginEmail A: %v", err)
}
if err := accountSvc.SetLoginEmail(ctx, ownerB.ID, "b@example.test"); err != nil {
t.Fatalf("SetLoginEmail B: %v", err)
}
codes := memory.NewCodeStore()
hash := "locked-reset-user"
if err := codes.Set(ctx, hash, store.PhoneCode{
Version: store.PhoneCodeVersionCurrent,
IssuedUserID: ownerA.ID,
Phone: ownerA.Phone,
Code: "654321",
Channel: codeChannelEmailLogin,
MaxAttempts: 5,
}, time.Minute); err != nil {
t.Fatalf("seed reset code: %v", err)
}
delivery := &captureLoginCodeDelivery{}
authSvc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
resetUserID, err := authSvc.ConsumeLoginEmailReset(ctx, ownerA.Phone, hash)
if err != nil || resetUserID != ownerA.ID {
t.Fatalf("ConsumeLoginEmailReset uid=%d err=%v", resetUserID, err)
}
users.setOwnerView(ownerA.Phone, ownerB, true)
if err := accountSvc.ClearLoginEmail(ctx, resetUserID); err != nil {
t.Fatalf("ClearLoginEmail exact A: %v", err)
}
if _, err := authSvc.SendPhoneCodeAfterLoginEmailReset(ctx, ownerA.Phone, resetUserID); !errors.Is(err, ErrCodeInvalid) {
t.Fatalf("SendPhoneCodeAfterLoginEmailReset across A->B err=%v, want invalid", err)
}
if _, found, err := accountSvc.LoginEmail(ctx, ownerA.ID); err != nil || found {
t.Fatalf("A login email found=%v err=%v, want cleared", found, err)
}
if email, found, err := accountSvc.LoginEmail(ctx, ownerB.ID); err != nil || !found || email != "b@example.test" {
t.Fatalf("B login email=%q found=%v err=%v, want unchanged", email, found, err)
}
if len(delivery.requests) != 0 {
t.Fatalf("owner B received reset replacement code: %+v", delivery.requests)
}
}

78
internal/app/channels/service.go
View file

@ -1242,6 +1242,25 @@ func (s *Service) SendMessage(ctx context.Context, userID int64, req domain.Send
if req.UserID != userID {
return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid
}
if req.RandomID != 0 && !req.IdempotencyPreflighted {
fingerprint, err := store.ChannelSendFingerprint(req)
if err != nil {
return domain.SendChannelMessageResult{}, err
}
req.IdempotencyFingerprint = fingerprint
if replayStore, ok := s.channels.(store.ChannelSendReplayStore); ok {
replay, found, err := replayStore.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{
ChannelID: req.ChannelID,
SenderUserID: req.UserID,
RandomID: req.RandomID,
IdempotencyFingerprint: fingerprint,
})
if err != nil || found {
return replay, err
}
req.IdempotencyPreflighted = true
}
}
if err := s.ensureCanSend(ctx, req.UserID); err != nil {
return domain.SendChannelMessageResult{}, err
}
@ -1253,6 +1272,25 @@ func (s *Service) SendMessage(ctx context.Context, userID int64, req domain.Send
return s.channels.SendChannelMessage(ctx, req)
}
// LookupChannelSendReplay reads a regular-channel or monoforum receipt without current
// membership/send-gate checks. The authenticated caller remains bound to SenderUserID.
func (s *Service) LookupChannelSendReplay(ctx context.Context, userID int64, req domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) {
if s == nil || s.channels == nil || userID == 0 {
return domain.SendChannelMessageResult{}, false, nil
}
if req.SenderUserID == 0 {
req.SenderUserID = userID
}
if req.SenderUserID != userID || req.ChannelID == 0 || req.RandomID == 0 {
return domain.SendChannelMessageResult{}, false, domain.ErrChannelInvalid
}
replayStore, ok := s.channels.(store.ChannelSendReplayStore)
if !ok {
return domain.SendChannelMessageResult{}, false, nil
}
return replayStore.LookupChannelSendReplay(ctx, req)
}
func (s *Service) ensureCanSend(ctx context.Context, userID int64) error {
if s == nil || s.sendGate == nil || userID == 0 {
return nil
@ -1754,6 +1792,26 @@ func (s *Service) SendMonoforumMessage(ctx context.Context, req domain.SendMonof
if s == nil || s.channels == nil || req.MonoforumID == 0 || req.SenderUserID == 0 || req.SavedPeer.ID == 0 {
return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid
}
if req.RandomID != 0 && !req.IdempotencyPreflighted {
fingerprint, err := store.MonoforumSendFingerprint(req)
if err != nil {
return domain.SendChannelMessageResult{}, err
}
req.IdempotencyFingerprint = fingerprint
if replayStore, ok := s.channels.(store.ChannelSendReplayStore); ok {
replay, found, err := replayStore.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{
ChannelID: req.MonoforumID,
SenderUserID: req.SenderUserID,
SavedPeer: req.SavedPeer,
RandomID: req.RandomID,
IdempotencyFingerprint: fingerprint,
})
if err != nil || found {
return replay, err
}
req.IdempotencyPreflighted = true
}
}
if err := s.ensureCanSend(ctx, req.SenderUserID); err != nil {
return domain.SendChannelMessageResult{}, err
}
@ -2021,6 +2079,26 @@ func (s *Service) DirtyActiveChannelsForUser(ctx context.Context, userID int64,
return s.channels.ListDirtyActiveChannelsForUser(ctx, userID, sinceDate, afterChannelID, limit)
}
// MaxChannelPts returns the durable channel watermark used by the fan-out saturation recovery
// sweep. It intentionally performs no viewer access check: target visibility is derived from the
// process-local joined-membership index, while getChannelDifference performs authoritative access
// validation when a client consumes the nudge.
func (s *Service) MaxChannelPts(ctx context.Context, channelID int64) (int, error) {
if s == nil || s.channels == nil || channelID == 0 {
return 0, domain.ErrChannelInvalid
}
return s.channels.MaxChannelPts(ctx, channelID)
}
// MaxChannelPtsBatch reloads a bounded recovery page in one store call. Missing ids are omitted:
// they represent channels deleted after the process-local online-membership snapshot was taken.
func (s *Service) MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error) {
if s == nil || s.channels == nil {
return nil, domain.ErrChannelInvalid
}
return s.channels.MaxChannelPtsBatch(ctx, channelIDs)
}
// ActiveMemberIDs returns a bounded list for transient online fanout such as typing.
func (s *Service) ActiveMemberIDs(ctx context.Context, userID, channelID int64, limit int) ([]int64, error) {
if s == nil || s.channels == nil || userID == 0 || channelID == 0 {

95
internal/app/channels/service_test.go
View file

@ -30,6 +30,46 @@ func TestServiceSendMessageHonorsSendPermissionGate(t *testing.T) {
}
}
func TestServiceChannelReplayPrecedesCurrentSendPermissionGate(t *testing.T) {
ctx := context.Background()
channels := memory.NewChannelStore()
created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{
CreatorUserID: 1001,
Title: "replay gate",
Megagroup: true,
Date: 1_700_000_000,
})
if err != nil {
t.Fatalf("CreateChannel: %v", err)
}
req := domain.SendChannelMessageRequest{
ChannelID: created.Channel.ID,
RandomID: 92,
Message: "committed before restriction",
Date: 1_700_000_001,
}
allowed := NewService(channels)
first, err := allowed.SendMessage(ctx, 1001, req)
if err != nil {
t.Fatalf("first SendMessage: %v", err)
}
denied := NewService(channels, WithSendPermissionChecker(channelDenySendChecker{}))
req.Date++
replay, err := denied.SendMessage(ctx, 1001, req)
if err != nil {
t.Fatalf("replay through denied gate: %v", err)
}
if !replay.Duplicate || replay.Message.ID != first.Message.ID {
t.Fatalf("replay = %+v, want committed duplicate %d", replay, first.Message.ID)
}
req.Message = "different intent"
if _, err := denied.SendMessage(ctx, 1001, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
t.Fatalf("conflicting replay err=%v, want ErrMessageRandomIDDuplicate before send gate", err)
}
}
func TestServiceSendMonoforumMessageHonorsSendPermissionGate(t *testing.T) {
ctx := context.Background()
svc := NewService(memory.NewChannelStore(), WithSendPermissionChecker(channelDenySendChecker{}))
@ -44,6 +84,52 @@ func TestServiceSendMonoforumMessageHonorsSendPermissionGate(t *testing.T) {
}
}
func TestServiceMonoforumReplayPrecedesCurrentSendPermissionGate(t *testing.T) {
ctx := context.Background()
channels := memory.NewChannelStore()
parent, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{
CreatorUserID: 1001,
Title: "direct messages",
Broadcast: true,
Date: 1_700_000_010,
})
if err != nil {
t.Fatalf("CreateChannel: %v", err)
}
enabled, err := channels.SetPaidMessagesPrice(ctx, 1001, parent.Channel.ID, 0, true)
if err != nil {
t.Fatalf("SetPaidMessagesPrice: %v", err)
}
req := domain.SendMonoforumMessageRequest{
MonoforumID: enabled.Channel.LinkedMonoforumID,
SenderUserID: 1002,
SavedPeer: domain.Peer{Type: domain.PeerTypeUser, ID: 1002},
RandomID: 93,
Message: "committed direct message",
Date: 1_700_000_011,
}
allowed := NewService(channels)
first, err := allowed.SendMonoforumMessage(ctx, req)
if err != nil {
t.Fatalf("first SendMonoforumMessage: %v", err)
}
denied := NewService(channels, WithSendPermissionChecker(channelDenySendChecker{}))
req.Date++
replay, err := denied.SendMonoforumMessage(ctx, req)
if err != nil {
t.Fatalf("monoforum replay through denied gate: %v", err)
}
if !replay.Duplicate || replay.Message.ID != first.Message.ID {
t.Fatalf("monoforum replay = %+v, want committed duplicate %d", replay, first.Message.ID)
}
req.Message = "different intent"
if _, err := denied.SendMonoforumMessage(ctx, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
t.Fatalf("conflicting monoforum replay err=%v, want ErrMessageRandomIDDuplicate before send gate", err)
}
}
type channelDenySendChecker struct{}
func (channelDenySendChecker) CanSendMessages(context.Context, int64) error {
@ -813,7 +899,8 @@ func TestCreateChatCreatesMegagroupWithChannelPts(t *testing.T) {
duplicate, err := service.SendMessage(ctx, 1001, domain.SendChannelMessageRequest{
ChannelID: created.Channel.ID,
RandomID: 99,
Message: "hello again",
Message: "hello",
ViaBotID: 1003,
Date: 12,
})
if err != nil {
@ -2032,12 +2119,12 @@ func TestChannelEditDeleteAndLocalClearUseChannelPts(t *testing.T) {
if edited.Event.Type != domain.ChannelUpdateEditMessage || edited.Event.Pts != 4 || edited.Event.PtsCount != 1 {
t.Fatalf("edit event = %+v, want channel edit pts=4 count=1", edited.Event)
}
duplicate, err := service.SendMessage(ctx, 1002, domain.SendChannelMessageRequest{ChannelID: created.Channel.ID, RandomID: 2, Message: "two retry", Date: 13})
duplicate, err := service.SendMessage(ctx, 1002, domain.SendChannelMessageRequest{ChannelID: created.Channel.ID, RandomID: 2, Message: "two", Date: 13})
if err != nil {
t.Fatalf("duplicate SendMessage after edit: %v", err)
}
if !duplicate.Duplicate || duplicate.Event.Type != domain.ChannelUpdateNewMessage || duplicate.Message.Body != "two" || duplicate.Event.Message.Body != "two" {
t.Fatalf("duplicate after edit = %+v, want original new-message snapshot", duplicate)
if !duplicate.Duplicate || duplicate.Event.Type != domain.ChannelUpdateNewMessage || duplicate.Message.Body != "two edited" || duplicate.Event.Message.Body != "two edited" {
t.Fatalf("duplicate after edit = %+v, want current message in new-message replay", duplicate)
}
deleted, err := service.DeleteMessages(ctx, 1001, domain.DeleteChannelMessagesRequest{

67
internal/app/files/photos.go
View file

@ -8,10 +8,10 @@ import (
"fmt"
"image"
"image/color"
"io"
stddraw "image/draw"
_ "image/jpeg" // 注册 jpeg DecodeConfig，用于读取上传头像/图片尺寸
"image/png"
"io"
"math"
"strings"
"time"
@ -48,14 +48,40 @@ func (s *Service) UploadProfilePhotoKind(ctx context.Context, ownerType domain.P
// CreatePhotoFromUpload 把已上传文件组装成 Photo（不绑定 profile_photos），用于频道头像 / 图片消息。
func (s *Service) CreatePhotoFromUpload(ctx context.Context, file domain.UploadedFileRef) (domain.Photo, error) {
data, err := s.assembleUpload(ctx, file.OwnerUserID, file.FileID, file.Parts)
intentHash, err := uploadedMediaIntentHash(domain.UploadedMediaPhoto, file, nil)
if err != nil {
return domain.Photo{}, err
}
if photo, found, err := s.replayUploadedPhoto(ctx, file, intentHash); err != nil || found {
return photo, err
}
data, err := s.readUploadBytes(ctx, file.OwnerUserID, file.FileID, file.Parts)
if err != nil {
return domain.Photo{}, err
}
if len(data) == 0 {
return domain.Photo{}, domain.ErrPhotoInvalid
}
return s.createPhoto(ctx, data, photoSizeSpecsForMessage(data))
photo, err := s.createPhoto(ctx, data, photoSizeSpecsForMessage(data))
if err != nil {
return domain.Photo{}, err
}
receipt, err := s.commitUploadedMediaReceipt(ctx, file, domain.UploadedMediaPhoto, intentHash, photo.ID)
if err != nil {
return domain.Photo{}, err
}
if receipt.MediaID != photo.ID {
winner, found, err := s.media.GetPhoto(ctx, receipt.MediaID)
if err != nil {
return domain.Photo{}, err
}
if !found {
return domain.Photo{}, fmt.Errorf("concurrent upload receipt references missing photo %d", receipt.MediaID)
}
photo = winner
}
s.cleanupMaterializedUpload(ctx, file, "photo materialized")
return photo, nil
}
// CreatePhotoFromBytes stores already-fetched image bytes as a message Photo.
@ -202,6 +228,13 @@ func validateAvatarMarkupSize(size domain.PhotoSize) error {
// CreateDocumentFromUpload 把已上传文件组装成 Document（文件/视频/音频/gif/贴纸消息），落 blob + documents。
func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.UploadedFileRef, spec domain.DocumentSpec) (domain.Document, error) {
intentHash, err := uploadedMediaIntentHash(domain.UploadedMediaDocument, file, &spec)
if err != nil {
return domain.Document{}, err
}
if doc, found, err := s.replayUploadedDocument(ctx, file, intentHash); err != nil || found {
return doc, err
}
body, err := s.assembleUploadBlob(ctx, file.OwnerUserID, file.FileID, file.Parts)
if err != nil {
return domain.Document{}, err
@ -234,11 +267,13 @@ func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.Uplo
DCID: s.dc,
Attributes: spec.Attributes,
}
thumbMaterialized := false
if spec.Thumb != nil {
thumbData, err := s.assembleUpload(ctx, spec.Thumb.OwnerUserID, spec.Thumb.FileID, spec.Thumb.Parts)
thumbData, err := s.readUploadBytes(ctx, spec.Thumb.OwnerUserID, spec.Thumb.FileID, spec.Thumb.Parts)
if err == nil && len(thumbData) > 0 {
if thumb, err := s.putDocumentThumb(ctx, docID, thumbData); err == nil {
doc.Thumbs = []domain.PhotoSize{thumb}
thumbMaterialized = true
}
}
}
@ -250,12 +285,23 @@ func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.Uplo
if err := s.media.PutDocument(ctx, doc); err != nil {
return domain.Document{}, err
}
if err := s.cleanupUploadParts(ctx, file.OwnerUserID, file.FileID); err != nil {
s.log.Warn("cleanup assembled document upload parts failed",
zap.Int64("owner_user_id", file.OwnerUserID),
zap.Int64("file_id", file.FileID),
zap.Int64("document_id", docID),
zap.Error(err))
receipt, err := s.commitUploadedMediaReceipt(ctx, file, domain.UploadedMediaDocument, intentHash, doc.ID)
if err != nil {
return domain.Document{}, err
}
if receipt.MediaID != doc.ID {
winner, found, err := s.media.GetDocument(ctx, receipt.MediaID)
if err != nil {
return domain.Document{}, err
}
if !found {
return domain.Document{}, fmt.Errorf("concurrent upload receipt references missing document %d", receipt.MediaID)
}
doc = winner
}
s.cleanupMaterializedUpload(ctx, file, "document materialized")
if spec.Thumb != nil && thumbMaterialized {
s.cleanupMaterializedUpload(ctx, *spec.Thumb, "document thumbnail materialized")
}
return doc, nil
}
@ -277,6 +323,7 @@ var faststartVideoMimes = map[string]bool{
// 此时只发生几次 16 字节读，不读整段媒体。
// 2. 仅 moov 在末尾时才重写；且优先走流式（仅 ftyp+moov 进内存，mdat 大块分块流式拼接），
// 不把整段视频 2× 驻留内存。moov 非末尾的罕见排布回退到全量重排。
//
// 任何不适用/失败都返回原 body，绝不让上传失败或损坏数据。
func (s *Service) maybeFaststartVideoBlob(ctx context.Context, mimeType string, body assembledUploadBlob) assembledUploadBlob {
if !faststartVideoMimes[strings.ToLower(strings.TrimSpace(mimeType))] {

27
internal/app/files/seed_test.go
View file

@ -26,6 +26,7 @@ type fakeMediaStore struct {
parts map[string][]domain.UploadPart
webPages map[int64]domain.MessageWebPage
seedState map[string]string
receipts map[string]domain.UploadedMediaReceipt
}
func newFakeMediaStore() *fakeMediaStore {
@ -36,9 +37,35 @@ func newFakeMediaStore() *fakeMediaStore {
sets: map[int64]domain.StickerSet{},
parts: map[string][]domain.UploadPart{},
seedState: map[string]string{},
receipts: map[string]domain.UploadedMediaReceipt{},
}
}
func fakeUploadReceiptKey(ownerUserID, fileID int64) string {
return fmt.Sprintf("%d/%d", ownerUserID, fileID)
}
func (f *fakeMediaStore) GetUploadedMediaReceipt(_ context.Context, ownerUserID, fileID int64) (domain.UploadedMediaReceipt, bool, error) {
f.mu.Lock()
defer f.mu.Unlock()
receipt, ok := f.receipts[fakeUploadReceiptKey(ownerUserID, fileID)]
receipt.IntentHash = append([]byte(nil), receipt.IntentHash...)
return receipt, ok, nil
}
func (f *fakeMediaStore) PutUploadedMediaReceipt(_ context.Context, receipt domain.UploadedMediaReceipt) (domain.UploadedMediaReceipt, bool, error) {
f.mu.Lock()
defer f.mu.Unlock()
key := fakeUploadReceiptKey(receipt.OwnerUserID, receipt.FileID)
if stored, ok := f.receipts[key]; ok {
stored.IntentHash = append([]byte(nil), stored.IntentHash...)
return stored, false, nil
}
receipt.IntentHash = append([]byte(nil), receipt.IntentHash...)
f.receipts[key] = receipt
return receipt, true, nil
}
func (f *fakeMediaStore) SaveFilePart(_ context.Context, part domain.UploadPart) error {
f.mu.Lock()
defer f.mu.Unlock()

17
internal/app/files/service.go
View file

@ -514,6 +514,20 @@ func orderDocuments(docs []domain.Document, ids []int64) []domain.Document {
// assembleUpload 把已上传分片按 part 顺序拼成完整字节，并清理分片。
// expectedParts>0 时校验分片连续且齐全。
func (s *Service) assembleUpload(ctx context.Context, ownerUserID, fileID int64, expectedParts int) ([]byte, error) {
buf, err := s.readUploadBytes(ctx, ownerUserID, fileID, expectedParts)
if err != nil {
return nil, err
}
if err := s.cleanupUploadParts(ctx, ownerUserID, fileID); err != nil {
return nil, err
}
return buf, nil
}
// readUploadBytes validates and reads all parts without consuming them. Message-media
// materialization persists an upload receipt before cleanup; callers that do not need replayability
// continue to use assembleUpload.
func (s *Service) readUploadBytes(ctx context.Context, ownerUserID, fileID int64, expectedParts int) ([]byte, error) {
parts, _, err := s.loadAndValidateUploadParts(ctx, ownerUserID, fileID, expectedParts)
if err != nil {
return nil, err
@ -532,9 +546,6 @@ func (s *Service) assembleUpload(ctx context.Context, ownerUserID, fileID int64,
}
buf = append(buf, data...)
}
if err := s.cleanupUploadParts(ctx, ownerUserID, fileID); err != nil {
return nil, err
}
return buf, nil
}

45
internal/app/files/upload_parts_test.go
View file

@ -124,6 +124,51 @@ func TestCreateDocumentFromUploadStreamsBodyAndCleansParts(t *testing.T) {
if string(body) != strings.Join(parts, "") {
t.Fatalf("body blob mismatch")
}
replayed, err := svc.CreateDocumentFromUpload(ctx,
domain.UploadedFileRef{OwnerUserID: 10, FileID: 200, Parts: len(parts), Name: "large.bin", Big: true},
domain.DocumentSpec{MimeType: "application/octet-stream"},
)
if err != nil {
t.Fatalf("replay CreateDocumentFromUpload after part cleanup: %v", err)
}
if replayed.ID != doc.ID || replayed.AccessHash != doc.AccessHash {
t.Fatalf("replayed document = %d/%d, want original %d/%d", replayed.ID, replayed.AccessHash, doc.ID, doc.AccessHash)
}
if _, err := svc.CreateDocumentFromUpload(ctx,
domain.UploadedFileRef{OwnerUserID: 10, FileID: 200, Parts: len(parts), Name: "large.bin", Big: true},
domain.DocumentSpec{MimeType: "text/plain"},
); !errors.Is(err, domain.ErrFilePartsInvalid) {
t.Fatalf("changed materialization intent err = %v, want ErrFilePartsInvalid", err)
}
}
func TestCreatePhotoFromUploadReceiptReplaysAfterPartCleanup(t *testing.T) {
ctx := context.Background()
media := newFakeMediaStore()
svc, _ := newUploadPartTestService(t, media, domain.UploadPartQuota{})
file := domain.UploadedFileRef{OwnerUserID: 10, FileID: 201, Parts: 1, Name: "photo.jpg"}
if _, err := svc.SaveFilePart(ctx, file.OwnerUserID, file.FileID, 0, []byte("image-bytes")); err != nil {
t.Fatalf("SaveFilePart: %v", err)
}
first, err := svc.CreatePhotoFromUpload(ctx, file)
if err != nil {
t.Fatalf("CreatePhotoFromUpload: %v", err)
}
if remaining, err := media.LoadFileParts(ctx, file.OwnerUserID, file.FileID); err != nil || len(remaining) != 0 {
t.Fatalf("upload parts after photo materialization = %+v err=%v", remaining, err)
}
replayed, err := svc.CreatePhotoFromUpload(ctx, file)
if err != nil {
t.Fatalf("replay CreatePhotoFromUpload: %v", err)
}
if replayed.ID != first.ID || replayed.AccessHash != first.AccessHash {
t.Fatalf("replayed photo = %d/%d, want original %d/%d", replayed.ID, replayed.AccessHash, first.ID, first.AccessHash)
}
changed := file
changed.Name = "different.jpg"
if _, err := svc.CreatePhotoFromUpload(ctx, changed); !errors.Is(err, domain.ErrFilePartsInvalid) {
t.Fatalf("changed photo intent err = %v, want ErrFilePartsInvalid", err)
}
}
type countingUploadPartBackend struct {

106
internal/app/files/upload_receipt.go Normal file
View file

@ -0,0 +1,106 @@
package files
import (
"bytes"
"context"
"crypto/sha256"
"encoding/json"
"fmt"
"go.uber.org/zap"
"telesrv/internal/domain"
)
const uploadedMediaIntentVersion = 1
type uploadedMediaIntent struct {
Version int `json:"version"`
Kind domain.UploadedMediaKind `json:"kind"`
File domain.UploadedFileRef `json:"file"`
Spec *domain.DocumentSpec `json:"spec,omitempty"`
}
func uploadedMediaIntentHash(kind domain.UploadedMediaKind, file domain.UploadedFileRef, spec *domain.DocumentSpec) ([]byte, error) {
payload, err := json.Marshal(uploadedMediaIntent{
Version: uploadedMediaIntentVersion,
Kind: kind,
File: file,
Spec: spec,
})
if err != nil {
return nil, fmt.Errorf("marshal uploaded media intent: %w", err)
}
sum := sha256.Sum256(payload)
return sum[:], nil
}
func sameUploadedMediaReceipt(receipt domain.UploadedMediaReceipt, kind domain.UploadedMediaKind, intentHash []byte) bool {
return receipt.Kind == kind && len(intentHash) == sha256.Size && bytes.Equal(receipt.IntentHash, intentHash)
}
func (s *Service) replayUploadedPhoto(ctx context.Context, file domain.UploadedFileRef, intentHash []byte) (domain.Photo, bool, error) {
receipt, found, err := s.media.GetUploadedMediaReceipt(ctx, file.OwnerUserID, file.FileID)
if err != nil || !found {
return domain.Photo{}, false, err
}
if !sameUploadedMediaReceipt(receipt, domain.UploadedMediaPhoto, intentHash) {
return domain.Photo{}, false, domain.ErrFilePartsInvalid
}
photo, found, err := s.media.GetPhoto(ctx, receipt.MediaID)
if err != nil {
return domain.Photo{}, false, err
}
if !found {
return domain.Photo{}, false, fmt.Errorf("uploaded photo receipt %d/%d references missing photo %d", file.OwnerUserID, file.FileID, receipt.MediaID)
}
s.cleanupMaterializedUpload(ctx, file, "photo replay")
return photo, true, nil
}
func (s *Service) replayUploadedDocument(ctx context.Context, file domain.UploadedFileRef, intentHash []byte) (domain.Document, bool, error) {
receipt, found, err := s.media.GetUploadedMediaReceipt(ctx, file.OwnerUserID, file.FileID)
if err != nil || !found {
return domain.Document{}, false, err
}
if !sameUploadedMediaReceipt(receipt, domain.UploadedMediaDocument, intentHash) {
return domain.Document{}, false, domain.ErrFilePartsInvalid
}
doc, found, err := s.media.GetDocument(ctx, receipt.MediaID)
if err != nil {
return domain.Document{}, false, err
}
if !found {
return domain.Document{}, false, fmt.Errorf("uploaded document receipt %d/%d references missing document %d", file.OwnerUserID, file.FileID, receipt.MediaID)
}
s.cleanupMaterializedUpload(ctx, file, "document replay")
return doc, true, nil
}
func (s *Service) commitUploadedMediaReceipt(ctx context.Context, file domain.UploadedFileRef, kind domain.UploadedMediaKind, intentHash []byte, mediaID int64) (domain.UploadedMediaReceipt, error) {
receipt, _, err := s.media.PutUploadedMediaReceipt(ctx, domain.UploadedMediaReceipt{
OwnerUserID: file.OwnerUserID,
FileID: file.FileID,
IntentHash: intentHash,
Kind: kind,
MediaID: mediaID,
})
if err != nil {
return domain.UploadedMediaReceipt{}, err
}
if !sameUploadedMediaReceipt(receipt, kind, intentHash) {
return domain.UploadedMediaReceipt{}, domain.ErrFilePartsInvalid
}
return receipt, nil
}
func (s *Service) cleanupMaterializedUpload(ctx context.Context, file domain.UploadedFileRef, reason string) {
if err := s.cleanupUploadParts(ctx, file.OwnerUserID, file.FileID); err != nil {
s.log.Warn("cleanup materialized upload parts failed",
zap.String("reason", reason),
zap.Int64("owner_user_id", file.OwnerUserID),
zap.Int64("file_id", file.FileID),
zap.Error(err),
)
}
}

277
internal/app/maintenance/retention.go
View file

@ -17,12 +17,49 @@ type TempAuthKeyRetentionStore interface {
DeleteExpired(ctx context.Context, expiredBefore int64, limit int) (int, error)
}
// OrphanAuthKeyRetentionStore 回收从未形成授权/temp binding 的旧握手 key。
// protected 是当前连接注册表实际使用的 raw auth_key_id 快照。
type OrphanAuthKeyRetentionStore interface {
DeleteOrphaned(ctx context.Context, olderThan time.Duration, limit int, protected [][8]byte) (int, error)
}
type ActiveRawAuthKeyProvider interface {
ActiveRawAuthKeyIDs() [][8]byte
}
// ActiveAuthKeyHeartbeatStore 把本实例仍在使用的 raw auth key 活性持久化。多实例下
// orphan GC 不能只看当前进程的 active 快照；其它实例的 heartbeat 会推进数据库
// last_used_at，使它们不会被误判为孤儿。
type ActiveAuthKeyHeartbeatStore interface {
TouchActiveRawAuthKeys(ctx context.Context, ids [][8]byte) error
}
// BotAPIUpdateRetentionStore 回收 Bot API getUpdates 投递队列的死行（性能审计 H1）：
// 已确认且超过宽限期的行 + 按消息 date 超过保留期的行（官方 Bot API updates 最多保留 24h）。
type BotAPIUpdateRetentionStore interface {
DeleteDeliveredOrExpired(ctx context.Context, confirmedGrace, maxAge time.Duration, limit int) (int, error)
}
// UserUpdateEventRetentionStore 只回收所有当前授权设备都明确确认过的账号事件前缀。
// 它不是普通 TTL：任一授权缺 state 时确认水位为 0，不得删除该设备可能仍需的事件。
type UserUpdateEventRetentionStore interface {
DeleteConfirmedPrefix(ctx context.Context, olderThan time.Duration, limit int) (int, error)
}
// ChannelUpdateEventRetentionStore 回收超过保留期的 channel durable update 连续前缀。
// 具体 store 必须在同一事务内删除事件并推进 retained floor，低于 floor 的客户端由
// updates.getChannelDifference 走 channelDifferenceTooLong 快照恢复。
type ChannelUpdateEventRetentionStore interface {
DeleteExpiredChannelUpdateEvents(ctx context.Context, olderThan time.Duration, limit int) (int, error)
}
// LoginCodeDeliveryRetentionStore reclaims only compact idempotency receipts
// after their associated opaque code lifetime. It must not delete the message,
// durable update event, or outbox facts created by the delivery transaction.
type LoginCodeDeliveryRetentionStore interface {
DeleteExpiredLoginCodeDeliveries(ctx context.Context, expiredBefore time.Time, limit int) (int, error)
}
// botAPIConfirmedGrace 是已确认 Bot API update 行的删除宽限：确认水位之下的行不会再被
// getUpdates 读取（fromID 恒 > confirmed），宽限仅防御 offset 回拨调试；回收目标是清堆积。
const botAPIConfirmedGrace = 15 * time.Minute
@ -32,23 +69,39 @@ const botAPIConfirmedGrace = 15 * time.Minute
// 连接；回收目标是清堆积，晚一天无妨。
const tempAuthKeyExpiryGrace = 24 * time.Hour
const (
// terminal failed outbox 只承担短期诊断隔离；它不是 durable update log。
// 删除该任务会由 head trigger 立即放行同账号下一 pts，而 user_update_events
// 继续保留，在线漏推由正常 difference 路径补偿。
defaultOutboxPoisonRetention = time.Minute
defaultOutboxPoisonInterval = 15 * time.Second
)
// RetentionWorker 周期性回收存储中的死数据。
//
// 注意：本 worker 刻意不清理 user_update_events —— pts log 永久保留。原因：TDesktop 不支持
// 账号级 updates.differenceTooLong（api_updates.cpp 收到该响应只打一行日志，且漏掉
// setRequesting(false)，会永久锁死整个 update 引擎），服务端因此无法让"落后超过保留期"的
// 客户端整库重置；一旦裁剪 events，落后客户端的 getDifference 会拿到不完整的事件链而静默
// 丢消息。详见 docs/performance-audit.md 与 docs/compatibility-matrix.md。user_update_events
// 长期膨胀作为已知 todo。
// 注意：TDesktop 不支持账号级 updates.differenceTooLong（api_updates.cpp 收到该响应只
// 记录日志且不清 requesting，会永久锁死 update 引擎），因此绝不能按普通 TTL 硬裁剪
// user_update_events。本 worker 只允许 store 删除“所有当前授权设备都明确确认”的连续安全
// 前缀；落后或缺 state 的任一设备都会把 floor 压回 0。客户端偶然带回已确认前的旧 pts 时，
// updates 服务通过普通 differenceSlice checkpoint 推进，不发送 differenceTooLong。
type RetentionWorker struct {
outbox DispatchOutboxRetentionStore
tempKeys TempAuthKeyRetentionStore // 可为 nil（不回收 temp key 绑定）
botAPIUpdates BotAPIUpdateRetentionStore // 可为 nil（不回收 Bot API 队列）
logger *zap.Logger
retention time.Duration
botAPIRetention time.Duration
interval time.Duration
batch int
outbox DispatchOutboxRetentionStore
tempKeys TempAuthKeyRetentionStore // 可为 nil（不回收 temp key 绑定）
botAPIUpdates BotAPIUpdateRetentionStore // 可为 nil（不回收 Bot API 队列）
userUpdates UserUpdateEventRetentionStore
channelUpdates ChannelUpdateEventRetentionStore
loginCodeDeliveries LoginCodeDeliveryRetentionStore
orphanAuthKeys OrphanAuthKeyRetentionStore
activeAuthKeys ActiveRawAuthKeyProvider
activeAuthKeyHeartbeat ActiveAuthKeyHeartbeatStore
logger *zap.Logger
retention time.Duration
botAPIRetention time.Duration
orphanRetention time.Duration
outboxPoisonRetention time.Duration
outboxPoisonInterval time.Duration
interval time.Duration
batch int
}
func NewRetentionWorker(outbox DispatchOutboxRetentionStore, tempKeys TempAuthKeyRetentionStore, logger *zap.Logger, retention, interval time.Duration, batch int) *RetentionWorker {
@ -65,15 +118,32 @@ func NewRetentionWorker(outbox DispatchOutboxRetentionStore, tempKeys TempAuthKe
batch = 10000
}
return &RetentionWorker{
outbox: outbox,
tempKeys: tempKeys,
logger: logger,
retention: retention,
interval: interval,
batch: batch,
outbox: outbox,
tempKeys: tempKeys,
logger: logger,
retention: retention,
outboxPoisonRetention: defaultOutboxPoisonRetention,
outboxPoisonInterval: defaultOutboxPoisonInterval,
interval: interval,
batch: batch,
}
}
// WithDispatchOutboxPoisonPolicy 配置 terminal failed head 的独立短隔离与清理周期。
// 该周期不能复用 durable update 的周级保留期，否则一条确定性构造错误会冻结该
// 用户整条在线 pts lane。<=0 分别回退到 1m/15s 的安全默认值。
func (w *RetentionWorker) WithDispatchOutboxPoisonPolicy(retention, interval time.Duration) *RetentionWorker {
if retention <= 0 {
retention = defaultOutboxPoisonRetention
}
if interval <= 0 {
interval = defaultOutboxPoisonInterval
}
w.outboxPoisonRetention = retention
w.outboxPoisonInterval = interval
return w
}
// WithBotAPIUpdateRetention 启用 bot_api_updates 队列回收；retention <=0 时用官方语义默认 24h。
func (w *RetentionWorker) WithBotAPIUpdateRetention(store BotAPIUpdateRetentionStore, retention time.Duration) *RetentionWorker {
if retention <= 0 {
@ -84,26 +154,102 @@ func (w *RetentionWorker) WithBotAPIUpdateRetention(store BotAPIUpdateRetentionS
return w
}
// WithUserUpdateRetention 启用账号 update 的共同确认安全前缀回收。TDesktop 不支持
// account differenceTooLong，具体 store 必须保证未确认前缀永不删除。
func (w *RetentionWorker) WithUserUpdateRetention(store UserUpdateEventRetentionStore) *RetentionWorker {
w.userUpdates = store
return w
}
// WithChannelUpdateRetention 启用 channel durable update 的有界 TTL 回收；复用 worker 的
// retention/interval/batch，并由 store 的 retained floor 保证旧 pts 不会读到静默空洞。
func (w *RetentionWorker) WithChannelUpdateRetention(store ChannelUpdateEventRetentionStore) *RetentionWorker {
w.channelUpdates = store
return w
}
// WithLoginCodeDeliveryRetention enables bounded seek cleanup for compact
// phone_code_hash receipts. Each row carries its own expiry derived from the
// code TTL, so this cleanup intentionally does not reuse update-log retention.
func (w *RetentionWorker) WithLoginCodeDeliveryRetention(store LoginCodeDeliveryRetentionStore) *RetentionWorker {
w.loginCodeDeliveries = store
return w
}
// WithOrphanAuthKeyRetention 启用未授权握手 key 的有界回收。active 必须提供 raw key，
// 不能提供 temp→perm business key；否则未登录或 PFS 连接会被误判为 orphan。
func (w *RetentionWorker) WithOrphanAuthKeyRetention(store OrphanAuthKeyRetentionStore, active ActiveRawAuthKeyProvider, retention time.Duration) *RetentionWorker {
w.orphanAuthKeys = store
w.activeAuthKeys = active
w.activeAuthKeyHeartbeat, _ = store.(ActiveAuthKeyHeartbeatStore)
w.orphanRetention = retention
return w
}
func (w *RetentionWorker) Run(ctx context.Context) {
w.runOnce(ctx)
ticker := time.NewTicker(w.interval)
defer ticker.Stop()
retentionTicker := time.NewTicker(w.interval)
defer retentionTicker.Stop()
poisonTicker := time.NewTicker(w.outboxPoisonInterval)
defer poisonTicker.Stop()
var (
heartbeatTicker *time.Ticker
heartbeatC <-chan time.Time
)
if interval := w.orphanHeartbeatInterval(); interval > 0 {
heartbeatTicker = time.NewTicker(interval)
heartbeatC = heartbeatTicker.C
defer heartbeatTicker.Stop()
}
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
w.runOnce(ctx)
case <-retentionTicker.C:
w.runRetentionOnce(ctx)
case <-poisonTicker.C:
w.runOutboxPoisonOnce(ctx)
case <-heartbeatC:
w.heartbeatActiveAuthKeys(ctx)
}
}
}
func (w *RetentionWorker) runOnce(ctx context.Context) {
outboxDeleted, err := w.outbox.DeleteFailed(ctx, w.retention, w.batch)
w.runOutboxPoisonOnce(ctx)
w.runRetentionOnce(ctx)
}
func (w *RetentionWorker) runOutboxPoisonOnce(ctx context.Context) {
if w.outbox == nil {
return
}
outboxDeleted, err := w.outbox.DeleteFailed(ctx, w.outboxPoisonRetention, w.batch)
if err != nil {
w.logger.Warn("清理 failed dispatch_outbox 失败", zap.Error(err))
w.logger.Error("清理 terminal failed dispatch_outbox 失败",
zap.String("signal", "dispatch_outbox_poison_cleanup_failed"),
zap.Duration("quarantine", w.outboxPoisonRetention),
zap.Error(err),
)
} else if outboxDeleted > 0 {
w.logger.Info("清理 failed dispatch_outbox 完成", zap.Int("deleted", outboxDeleted))
// Error 级结构化信号刻意保留：发生 terminal failed 代表确定性编码、事件缺失
// 或其它不可自动重试故障。任务删除只解冻在线 lane，不会删除 durable event。
w.logger.Error("terminal failed dispatch_outbox 已结束隔离并释放用户 lane",
zap.String("signal", "dispatch_outbox_poison_released"),
zap.Int("deleted", outboxDeleted),
zap.Duration("quarantine", w.outboxPoisonRetention),
)
}
}
func (w *RetentionWorker) runRetentionOnce(ctx context.Context) {
if w.loginCodeDeliveries != nil {
deleted, err := w.loginCodeDeliveries.DeleteExpiredLoginCodeDeliveries(ctx, time.Now(), w.batch)
if err != nil {
w.logger.Warn("回收过期 login-code delivery 回执失败", zap.Error(err))
} else if deleted > 0 {
w.logger.Info("回收过期 login-code delivery 回执完成", zap.Int("deleted", deleted))
}
}
if w.tempKeys != nil {
expiredBefore := time.Now().Add(-tempAuthKeyExpiryGrace).Unix()
@ -114,6 +260,24 @@ func (w *RetentionWorker) runOnce(ctx context.Context) {
w.logger.Info("回收过期 temp auth key 绑定完成", zap.Int("deleted", tempDeleted))
}
}
if w.orphanAuthKeys != nil && w.orphanRetention > 0 {
var protected [][8]byte
if w.activeAuthKeys != nil {
protected = w.activeAuthKeys.ActiveRawAuthKeyIDs()
}
if !w.touchActiveAuthKeys(ctx, protected) {
// Fail safe: if this instance cannot publish its own active set, deleting against a
// stale database heartbeat could evict keys used by another instance too. Keep all
// candidates for this pass and retry after the next heartbeat.
} else {
orphanDeleted, err := w.orphanAuthKeys.DeleteOrphaned(ctx, w.orphanRetention, w.batch, protected)
if err != nil {
w.logger.Warn("回收未授权 orphan auth key 失败", zap.Error(err))
} else if orphanDeleted > 0 {
w.logger.Info("回收未授权 orphan auth key 完成", zap.Int("deleted", orphanDeleted))
}
}
}
if w.botAPIUpdates != nil {
botAPIDeleted, err := w.botAPIUpdates.DeleteDeliveredOrExpired(ctx, botAPIConfirmedGrace, w.botAPIRetention, w.batch)
if err != nil {
@ -122,4 +286,63 @@ func (w *RetentionWorker) runOnce(ctx context.Context) {
w.logger.Info("回收 bot_api_updates 队列完成", zap.Int("deleted", botAPIDeleted))
}
}
if w.userUpdates != nil {
userDeleted, err := w.userUpdates.DeleteConfirmedPrefix(ctx, w.retention, w.batch)
if err != nil {
w.logger.Warn("回收已共同确认的 user_update_events 前缀失败", zap.Error(err))
} else if userDeleted > 0 {
w.logger.Info("回收已共同确认的 user_update_events 前缀完成", zap.Int("deleted", userDeleted))
}
}
if w.channelUpdates != nil {
channelDeleted, err := w.channelUpdates.DeleteExpiredChannelUpdateEvents(ctx, w.retention, w.batch)
if err != nil {
// store 会逐频道隔离坏 gap 后继续本轮；deleted 可能非零，必须同时记录，
// 既不能把全局 pass 伪装成完全失败，也不能吞掉不变量错误。
w.logger.Warn("回收过期 channel_update_events 存在隔离频道",
zap.Int("deleted", channelDeleted),
zap.Error(err),
)
} else if channelDeleted > 0 {
w.logger.Info("回收过期 channel_update_events 连续前缀完成", zap.Int("deleted", channelDeleted))
}
}
}
func (w *RetentionWorker) orphanHeartbeatInterval() time.Duration {
if w.activeAuthKeyHeartbeat == nil || w.activeAuthKeys == nil || w.orphanRetention <= 0 {
return 0
}
interval := w.orphanRetention / 3
if interval <= 0 {
interval = time.Nanosecond
}
if w.interval > 0 && w.interval < interval {
interval = w.interval
}
return interval
}
func (w *RetentionWorker) heartbeatActiveAuthKeys(ctx context.Context) {
if w.activeAuthKeys == nil {
return
}
w.touchActiveAuthKeys(ctx, w.activeAuthKeys.ActiveRawAuthKeyIDs())
}
// touchActiveAuthKeys returns false only when a configured durable heartbeat failed. A store that
// predates the optional heartbeat interface keeps single-instance behavior.
func (w *RetentionWorker) touchActiveAuthKeys(ctx context.Context, protected [][8]byte) bool {
if w.activeAuthKeyHeartbeat == nil {
return true
}
if err := w.activeAuthKeyHeartbeat.TouchActiveRawAuthKeys(ctx, protected); err != nil {
w.logger.Error("刷新 active raw auth key heartbeat 失败，本轮跳过 orphan GC",
zap.String("signal", "auth_key_heartbeat_failed"),
zap.Int("active_keys", len(protected)),
zap.Error(err),
)
return false
}
return true
}

220
internal/app/maintenance/retention_test.go
View file

@ -2,19 +2,57 @@ package maintenance
import (
"context"
"errors"
"testing"
"time"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
"go.uber.org/zap/zaptest/observer"
)
type fakeOutboxRetention struct {
calls int
calls int
olderThan time.Duration
limit int
deleted int
}
func (f *fakeOutboxRetention) DeleteFailed(context.Context, time.Duration, int) (int, error) {
func (f *fakeOutboxRetention) DeleteFailed(_ context.Context, olderThan time.Duration, limit int) (int, error) {
f.calls++
return 0, nil
f.olderThan = olderThan
f.limit = limit
return f.deleted, nil
}
func TestRetentionWorkerUsesIndependentOutboxPoisonPolicyAndSignalsRelease(t *testing.T) {
core, logs := observer.New(zapcore.ErrorLevel)
outbox := &fakeOutboxRetention{deleted: 2}
w := NewRetentionWorker(outbox, nil, zap.New(core), 7*24*time.Hour, time.Hour, 73).
WithDispatchOutboxPoisonPolicy(2*time.Minute, 7*time.Second)
w.runOnce(context.Background())
if outbox.calls != 1 || outbox.olderThan != 2*time.Minute || outbox.limit != 73 {
t.Fatalf("outbox poison calls/args = %d/%v/%d, want 1/2m/73", outbox.calls, outbox.olderThan, outbox.limit)
}
entries := logs.FilterMessage("terminal failed dispatch_outbox 已结束隔离并释放用户 lane").All()
if len(entries) != 1 {
t.Fatalf("poison release error signals = %d, want 1", len(entries))
}
if got := entries[0].ContextMap()["signal"]; got != "dispatch_outbox_poison_released" {
t.Fatalf("poison signal = %v", got)
}
}
func TestRetentionWorkerOutboxPoisonPolicyDefaultsAreShort(t *testing.T) {
outbox := &fakeOutboxRetention{}
w := NewRetentionWorker(outbox, nil, zap.NewNop(), 168*time.Hour, time.Hour, 100).
WithDispatchOutboxPoisonPolicy(0, 0)
w.runOutboxPoisonOnce(context.Background())
if outbox.olderThan != defaultOutboxPoisonRetention || w.outboxPoisonInterval != defaultOutboxPoisonInterval {
t.Fatalf("default poison policy = %v/%v, want %v/%v", outbox.olderThan, w.outboxPoisonInterval, defaultOutboxPoisonRetention, defaultOutboxPoisonInterval)
}
}
type fakeTempKeyRetention struct {
@ -65,6 +103,34 @@ type fakeBotAPIRetention struct {
limit int
}
type fakeLoginCodeDeliveryRetention struct {
calls int
expiredBefore time.Time
limit int
}
func (f *fakeLoginCodeDeliveryRetention) DeleteExpiredLoginCodeDeliveries(_ context.Context, expiredBefore time.Time, limit int) (int, error) {
f.calls++
f.expiredBefore = expiredBefore
f.limit = limit
return 4, nil
}
func TestRetentionWorkerReclaimsExpiredLoginCodeDeliveryReceipts(t *testing.T) {
loginCodes := &fakeLoginCodeDeliveryRetention{}
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), 168*time.Hour, time.Hour, 83).
WithLoginCodeDeliveryRetention(loginCodes)
before := time.Now()
w.runRetentionOnce(context.Background())
after := time.Now()
if loginCodes.calls != 1 || loginCodes.limit != 83 {
t.Fatalf("login-code retention calls/limit = %d/%d, want 1/83", loginCodes.calls, loginCodes.limit)
}
if loginCodes.expiredBefore.Before(before) || loginCodes.expiredBefore.After(after) {
t.Fatalf("login-code expiry boundary = %v, want within [%v,%v]", loginCodes.expiredBefore, before, after)
}
}
func (f *fakeBotAPIRetention) DeleteDeliveredOrExpired(_ context.Context, confirmedGrace, maxAge time.Duration, limit int) (int, error) {
f.calls++
f.confirmedGrace = confirmedGrace
@ -99,3 +165,151 @@ func TestRetentionWorkerBotAPIRetentionDefaultsTo24h(t *testing.T) {
t.Fatalf("default bot api retention = %v, want 24h", botAPI.maxAge)
}
}
type fakeUserUpdateRetention struct {
calls int
olderThan time.Duration
limit int
}
func (f *fakeUserUpdateRetention) DeleteConfirmedPrefix(_ context.Context, olderThan time.Duration, limit int) (int, error) {
f.calls++
f.olderThan = olderThan
f.limit = limit
return 9, nil
}
func TestRetentionWorkerReclaimsOnlyConfirmedUserUpdatePrefix(t *testing.T) {
const retention = 7 * 24 * time.Hour
store := &fakeUserUpdateRetention{}
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), retention, time.Hour, 91).
WithUserUpdateRetention(store)
w.runOnce(context.Background())
if store.calls != 1 || store.olderThan != retention || store.limit != 91 {
t.Fatalf("user update retention calls/args = %d/%v/%d, want 1/%v/91", store.calls, store.olderThan, store.limit, retention)
}
}
type fakeChannelUpdateRetention struct {
calls int
olderThan time.Duration
limit int
}
func (f *fakeChannelUpdateRetention) DeleteExpiredChannelUpdateEvents(_ context.Context, olderThan time.Duration, limit int) (int, error) {
f.calls++
f.olderThan = olderThan
f.limit = limit
return 7, nil
}
func TestRetentionWorkerReclaimsChannelUpdates(t *testing.T) {
const (
retention = 14 * 24 * time.Hour
batch = 321
)
channelUpdates := &fakeChannelUpdateRetention{}
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), retention, time.Hour, batch).
WithChannelUpdateRetention(channelUpdates)
w.runOnce(context.Background())
if channelUpdates.calls != 1 {
t.Fatalf("channel update retention calls = %d, want 1", channelUpdates.calls)
}
if channelUpdates.olderThan != retention || channelUpdates.limit != batch {
t.Fatalf("channel update retention args = (%v, %d), want (%v, %d)",
channelUpdates.olderThan, channelUpdates.limit, retention, batch)
}
}
type fakeOrphanAuthKeyRetention struct {
calls int
olderThan time.Duration
limit int
protected [][8]byte
}
func (f *fakeOrphanAuthKeyRetention) DeleteOrphaned(_ context.Context, olderThan time.Duration, limit int, protected [][8]byte) (int, error) {
f.calls++
f.olderThan = olderThan
f.limit = limit
f.protected = append([][8]byte(nil), protected...)
return 2, nil
}
type fakeActiveRawAuthKeys struct{ ids [][8]byte }
func (f fakeActiveRawAuthKeys) ActiveRawAuthKeyIDs() [][8]byte {
return append([][8]byte(nil), f.ids...)
}
func TestRetentionWorkerProtectsActiveRawAuthKeysFromOrphanGC(t *testing.T) {
store := &fakeOrphanAuthKeyRetention{}
active := fakeActiveRawAuthKeys{ids: [][8]byte{{1}, {2}}}
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), time.Hour, time.Hour, 73).
WithOrphanAuthKeyRetention(store, active, 24*time.Hour)
w.runOnce(context.Background())
if store.calls != 1 || store.olderThan != 24*time.Hour || store.limit != 73 {
t.Fatalf("orphan retention calls/args = %d/%v/%d, want 1/24h/73", store.calls, store.olderThan, store.limit)
}
if len(store.protected) != 2 || store.protected[0] != ([8]byte{1}) || store.protected[1] != ([8]byte{2}) {
t.Fatalf("protected raw auth keys = %v, want {1},{2}", store.protected)
}
}
type fakeHeartbeatOrphanRetention struct {
fakeOrphanAuthKeyRetention
heartbeatCalls int
heartbeatIDs [][8]byte
heartbeatErr error
}
func (f *fakeHeartbeatOrphanRetention) TouchActiveRawAuthKeys(_ context.Context, ids [][8]byte) error {
f.heartbeatCalls++
f.heartbeatIDs = append([][8]byte(nil), ids...)
return f.heartbeatErr
}
func TestRetentionWorkerHeartbeatsActiveKeysBeforeOrphanDelete(t *testing.T) {
store := &fakeHeartbeatOrphanRetention{}
active := fakeActiveRawAuthKeys{ids: [][8]byte{{3}, {4}}}
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), time.Hour, 2*time.Hour, 19).
WithOrphanAuthKeyRetention(store, active, 3*time.Hour)
w.runRetentionOnce(context.Background())
if store.heartbeatCalls != 1 || store.calls != 1 {
t.Fatalf("heartbeat/delete calls = %d/%d, want 1/1", store.heartbeatCalls, store.calls)
}
if len(store.heartbeatIDs) != 2 || store.heartbeatIDs[0] != ([8]byte{3}) || store.heartbeatIDs[1] != ([8]byte{4}) {
t.Fatalf("heartbeat ids = %v, want {3},{4}", store.heartbeatIDs)
}
// min(retention worker interval=2h, orphan retention/3=1h)
if got := w.orphanHeartbeatInterval(); got != time.Hour {
t.Fatalf("heartbeat interval = %v, want 1h", got)
}
}
func TestRetentionWorkerSkipsOrphanDeleteWhenHeartbeatFails(t *testing.T) {
core, logs := observer.New(zapcore.ErrorLevel)
store := &fakeHeartbeatOrphanRetention{heartbeatErr: errors.New("db unavailable")}
active := fakeActiveRawAuthKeys{ids: [][8]byte{{5}}}
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.New(core), time.Hour, 30*time.Minute, 11).
WithOrphanAuthKeyRetention(store, active, 24*time.Hour)
if got := w.orphanHeartbeatInterval(); got != 30*time.Minute {
t.Fatalf("heartbeat interval = %v, want worker interval 30m", got)
}
w.runRetentionOnce(context.Background())
if store.heartbeatCalls != 1 || store.calls != 0 {
t.Fatalf("heartbeat/delete calls = %d/%d, want 1/0", store.heartbeatCalls, store.calls)
}
entries := logs.FilterMessage("刷新 active raw auth key heartbeat 失败，本轮跳过 orphan GC").All()
if len(entries) != 1 || entries[0].ContextMap()["signal"] != "auth_key_heartbeat_failed" {
t.Fatalf("heartbeat failure signals = %+v", entries)
}
}

28
internal/app/messages/album_group.go Normal file
View file

@ -0,0 +1,28 @@
package messages
import (
"context"
"errors"
"telesrv/internal/domain"
"telesrv/internal/store"
)
// ReserveAlbumGroup 把 RPC 已验证的一批 album item 交给持久层原子预留。
// 该能力只传 domain DTO；上传媒体解析与 tg 类型仍停留在 RPC edge。
func (s *Service) ReserveAlbumGroup(ctx context.Context, userID int64, req domain.AlbumGroupReservationRequest) (int64, error) {
if s == nil || s.messages == nil || userID <= 0 {
return 0, domain.ErrAlbumGroupReservationInvalid
}
if req.SenderUserID == 0 {
req.SenderUserID = userID
}
if req.SenderUserID != userID {
return 0, domain.ErrAlbumGroupReservationInvalid
}
reservations, ok := s.messages.(store.AlbumGroupStore)
if !ok {
return 0, errors.New("message store does not support album group reservations")
}
return reservations.ReserveAlbumGroup(ctx, req)
}

43
internal/app/messages/service.go
View file

@ -2,6 +2,7 @@ package messages
import (
"context"
"fmt"
"telesrv/internal/app/userprojection"
"telesrv/internal/domain"
@ -96,6 +97,28 @@ func (s *Service) SendPrivateText(ctx context.Context, userID int64, req domain.
if req.SenderUserID == 0 {
req.SenderUserID = userID
}
if req.SenderUserID != userID {
return domain.SendPrivateTextResult{}, domain.ErrUserSendRestricted
}
if req.RandomID != 0 && !req.IdempotencyPreflighted {
fingerprint, err := store.PrivateSendFingerprint(req)
if err != nil {
return domain.SendPrivateTextResult{}, err
}
req.IdempotencyFingerprint = fingerprint
if replayStore, ok := s.messages.(store.PrivateSendReplayStore); ok {
replay, found, err := replayStore.LookupPrivateSendReplay(ctx, domain.PrivateSendReplayRequest{
SenderUserID: req.SenderUserID,
RecipientUserID: req.RecipientUserID,
RandomID: req.RandomID,
IdempotencyFingerprint: fingerprint,
})
if err != nil || found {
return replay, err
}
req.IdempotencyPreflighted = true
}
}
if err := s.ensureCanSend(ctx, req.SenderUserID); err != nil {
return domain.SendPrivateTextResult{}, err
}
@ -113,6 +136,26 @@ func (s *Service) SendPrivateText(ctx context.Context, userID int64, req domain.
return res, err
}
// LookupPrivateSendReplay exposes the immutable receipt to the RPC boundary without executing
// send permission checks, business automation or bot responders. Sender identity is still bound
// to the authenticated app-service caller.
func (s *Service) LookupPrivateSendReplay(ctx context.Context, userID int64, req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) {
if s == nil || s.messages == nil || userID == 0 {
return domain.SendPrivateTextResult{}, false, nil
}
if req.SenderUserID == 0 {
req.SenderUserID = userID
}
if req.SenderUserID != userID || req.RecipientUserID == 0 || req.RandomID == 0 {
return domain.SendPrivateTextResult{}, false, fmt.Errorf("private send replay: invalid authenticated scope")
}
replayStore, ok := s.messages.(store.PrivateSendReplayStore)
if !ok {
return domain.SendPrivateTextResult{}, false, nil
}
return replayStore.LookupPrivateSendReplay(ctx, req)
}
func (s *Service) ensureCanSend(ctx context.Context, userID int64) error {
if s == nil || s.sendGate == nil || userID == 0 {
return nil

32
internal/app/messages/service_test.go
View file

@ -29,6 +29,38 @@ func TestServiceSendPrivateTextHonorsSendPermissionGate(t *testing.T) {
}
}
func TestServicePrivateReplayPrecedesCurrentSendPermissionGate(t *testing.T) {
ctx := context.Background()
messages := memory.NewMessageStore()
allowed := NewService(messages, nil)
req := domain.SendPrivateTextRequest{
SenderUserID: 1001,
RecipientUserID: 1002,
RandomID: 91,
Message: "committed before restriction",
Date: 1_700_000_000,
}
first, err := allowed.SendPrivateText(ctx, 1001, req)
if err != nil {
t.Fatalf("first SendPrivateText: %v", err)
}
denied := NewService(messages, nil, WithSendPermissionChecker(denySendChecker{}))
req.Date++ // execution time is not part of the immutable send intent.
replay, err := denied.SendPrivateText(ctx, 1001, req)
if err != nil {
t.Fatalf("replay through denied gate: %v", err)
}
if !replay.Duplicate || replay.SenderMessage.ID != first.SenderMessage.ID {
t.Fatalf("replay = %+v, want committed duplicate %d", replay, first.SenderMessage.ID)
}
req.Message = "different intent"
if _, err := denied.SendPrivateText(ctx, 1001, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
t.Fatalf("conflicting replay err=%v, want ErrMessageRandomIDDuplicate before send gate", err)
}
}
func TestServiceForwardPrivateMessagesHonorsSendPermissionGate(t *testing.T) {
ctx := context.Background()
store := &gateMessageStore{}

185
internal/app/updates/service.go
View file

@ -27,6 +27,10 @@ type newMessageEventFinder interface {
FindNewMessageEvent(ctx context.Context, userID int64, messageBoxID int) (domain.UpdateEvent, bool, error)
}
type userUpdateRetentionCheckpointStore interface {
UserUpdateRetentionCheckpoint(ctx context.Context, authKeyID [8]byte, userID int64) (pts, date int, ok bool, err error)
}
// ServiceOption 调整 updates 服务的运行时依赖。
type ServiceOption func(*Service)
@ -146,6 +150,11 @@ func (s *Service) AcknowledgeCurrentState(ctx context.Context, authKeyID [8]byte
if err := s.saveConfirmedState(ctx, authKeyID, userID, st); err != nil {
return domain.UpdateState{}, err
}
// getState 明确建立“从当前快照开始同步”的 baseline；即使响应丢失，客户端也会
// 重试 getState/重新拉 snapshot，而不会依赖 baseline 之前的 durable event。
if err := s.observeClientState(ctx, authKeyID, userID, st); err != nil {
return domain.UpdateState{}, err
}
return st, nil
}
@ -162,6 +171,27 @@ func (s *Service) GetDifference(ctx context.Context, authKeyID [8]byte, userID i
if err != nil {
return domain.UpdateDifference{}, err
}
// 只把客户端在本次请求中实际带回的 cursor 记为 observed。绝不能把本次将要
// 返回的 State 当确认：响应可能在 socket/进程故障中丢失。恶意/损坏客户端带来的
// 超前 pts 钳到账号当前连续水位，避免把 retention 安全边界推过 durable truth。
observed := from
if observed.Pts < 0 {
observed.Pts = 0
}
if observed.Pts > st.Pts {
observed.Pts = st.Pts
}
if err := s.observeClientState(ctx, authKeyID, userID, observed); err != nil {
return domain.UpdateDifference{}, err
}
// TDesktop 不支持账号级 updates.differenceTooLong。retention 只能删除所有授权
// 设备都已确认的共同前缀；当前设备若仍带更旧 pts，用一个空的普通
// differenceSlice 把 IntermediateState 推进到已确认 checkpoint，再从 live tail 续拉。
if checkpoint, found, err := s.retainedPrefixCheckpoint(ctx, authKeyID, userID, from, st); err != nil {
return domain.UpdateDifference{}, err
} else if found {
return checkpoint, nil
}
if s.events == nil || from.Pts >= st.Pts {
if from.Date != 0 {
st.Date = from.Date
@ -176,6 +206,17 @@ func (s *Service) GetDifference(ctx context.Context, authKeyID [8]byte, userID i
return domain.UpdateDifference{}, err
}
contiguous, gapEvent, expectedPts := contiguousPrefixAndGap(events, from.Pts)
// Retention may advance after the pre-read checkpoint probe and before ListAfter obtains its
// statement snapshot. If it removed the whole requested prefix, the read is empty or starts at a
// gap. Re-read the checkpoint before returning a non-advancing empty difference; otherwise a
// client can believe synchronization completed while retaining a cursor below deleted history.
if len(contiguous) == 0 && from.Pts < st.Pts {
if checkpoint, found, err := s.retainedPrefixCheckpoint(ctx, authKeyID, userID, from, st); err != nil {
return domain.UpdateDifference{}, err
} else if found {
return checkpoint, nil
}
}
last := from.Pts
if len(contiguous) > 0 {
last = contiguous[len(contiguous)-1].Pts
@ -215,6 +256,32 @@ func (s *Service) GetDifference(ctx context.Context, authKeyID [8]byte, userID i
}, nil
}
func (s *Service) retainedPrefixCheckpoint(ctx context.Context, authKeyID [8]byte, userID int64, from, current domain.UpdateState) (domain.UpdateDifference, bool, error) {
checkpoints, ok := s.events.(userUpdateRetentionCheckpointStore)
if !ok {
return domain.UpdateDifference{}, false, nil
}
pts, date, found, err := checkpoints.UserUpdateRetentionCheckpoint(ctx, authKeyID, userID)
if err != nil {
return domain.UpdateDifference{}, false, err
}
if !found || from.Pts >= pts {
return domain.UpdateDifference{}, false, nil
}
checkpoint := from
checkpoint.Pts = pts
checkpoint.Seq = 0
if date > 0 {
checkpoint.Date = date
} else if checkpoint.Date == 0 {
checkpoint.Date = current.Date
}
if err := s.saveConfirmedState(ctx, authKeyID, userID, checkpoint); err != nil {
return domain.UpdateDifference{}, false, err
}
return domain.UpdateDifference{State: checkpoint, Partial: true}, true, nil
}
func (s *Service) currentState(ctx context.Context, userID int64) (domain.UpdateState, error) {
current, err := s.currentPts(ctx, userID)
if err != nil {
@ -235,6 +302,14 @@ func (s *Service) saveConfirmedState(ctx context.Context, authKeyID [8]byte, use
return s.states.Save(ctx, authKeyID, userID, st)
}
func (s *Service) observeClientState(ctx context.Context, authKeyID [8]byte, userID int64, st domain.UpdateState) error {
if s.states == nil {
return nil
}
st.Seq = 0
return s.states.ObserveClientState(ctx, authKeyID, userID, st)
}
// contiguousPrefix 返回从 from 起 pts 严格连续（from+1, from+2, ...）的事件前缀。
// 先按 pts 升序排序以兼容存储返回顺序，遇到空洞即停。
func contiguousPrefix(events []domain.UpdateEvent, from int) []domain.UpdateEvent {
@ -287,7 +362,7 @@ func (s *Service) RecordNewMessage(ctx context.Context, authKeyID [8]byte, userI
if date == 0 {
date = int(time.Now().Unix())
}
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
return s.recordEvent(ctx, authKeyID, [8]byte{}, userID, domain.UpdateEvent{
Type: domain.UpdateEventNewMessage,
Date: date,
Message: msg,
@ -318,7 +393,7 @@ func (s *Service) PublishNewMessage(ctx context.Context, userID int64, msg domai
if date == 0 {
date = int(time.Now().Unix())
}
return s.recordEventCore(ctx, [8]byte{}, userID, domain.UpdateEvent{
return s.recordEventCore(ctx, [8]byte{}, [8]byte{}, userID, domain.UpdateEvent{
Type: domain.UpdateEventNewMessage,
Date: date,
Message: msg,
@ -369,11 +444,11 @@ func (s *Service) RecordMessagePoll(ctx context.Context, authKeyID [8]byte, user
}
// RecordStory records a story snapshot change for offline difference replay.
func (s *Service) RecordStory(ctx context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordStory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, story domain.Story, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
if userID == 0 && story.Owner.Type == domain.PeerTypeUser {
userID = story.Owner.ID
}
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventStory,
Date: story.Date,
Peer: story.Owner,
@ -389,7 +464,7 @@ func (s *Service) RecordStoryFanout(ctx context.Context, userID int64, story dom
if userID == 0 {
return domain.UpdateEvent{}, domain.UpdateState{}, domain.ErrStoryPeerInvalid
}
return s.recordEventCore(ctx, [8]byte{}, userID, domain.UpdateEvent{
return s.recordEventCore(ctx, [8]byte{}, [8]byte{}, userID, domain.UpdateEvent{
Type: domain.UpdateEventStory,
Date: story.Date,
Peer: story.Owner,
@ -399,11 +474,11 @@ func (s *Service) RecordStoryFanout(ctx context.Context, userID int64, story dom
}
// RecordReadStories records a read boundary update for multi-device sync.
func (s *Service) RecordReadStories(ctx context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordReadStories(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
if userID == 0 {
userID = read.ViewerID
}
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventReadStories,
Date: read.Date,
Peer: read.Peer,
@ -413,11 +488,11 @@ func (s *Service) RecordReadStories(ctx context.Context, authKeyID [8]byte, user
}
// RecordSentStoryReaction records the current user's story reaction for multi-device sync.
func (s *Service) RecordSentStoryReaction(ctx context.Context, authKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordSentStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
if userID == 0 {
userID = reaction.ViewerID
}
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventSentStoryReaction,
Date: reaction.Date,
Peer: reaction.Peer,
@ -432,7 +507,7 @@ func (s *Service) RecordSentStoryReaction(ctx context.Context, authKeyID [8]byte
// sent by another user. It does not advance any owner device confirmation state:
// the owner did not initiate the RPC, but online outbox and offline difference
// must still see the durable event.
func (s *Service) RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordNewStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
if ownerUserID == 0 && reaction.Story.Owner.Type == domain.PeerTypeUser {
ownerUserID = reaction.Story.Owner.ID
}
@ -442,7 +517,7 @@ func (s *Service) RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte,
if ownerUserID == 0 || reaction.ViewerID == 0 || reaction.Reaction == nil {
return domain.UpdateEvent{}, domain.UpdateState{}, domain.ErrStoryPeerInvalid
}
return s.recordEventCore(ctx, authKeyID, ownerUserID, domain.UpdateEvent{
return s.recordEventCore(ctx, stateAuthKeyID, excludeAuthKeyID, ownerUserID, domain.UpdateEvent{
Type: domain.UpdateEventNewStoryReaction,
Date: reaction.Date,
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: reaction.ViewerID},
@ -456,7 +531,7 @@ func (s *Service) RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte,
// RecordQuickReplyMutation records account-local quick reply state changes for
// multi-device sync. Quick-reply TL updates do not carry pts, so outbox appends
// auxiliary pts bookkeeping just like other account settings events.
func (s *Service) RecordQuickReplyMutation(ctx context.Context, authKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordQuickReplyMutation(ctx context.Context, stateAuthKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
if userID == 0 {
userID = mutation.List.OwnerUserID
}
@ -481,16 +556,16 @@ func (s *Service) RecordQuickReplyMutation(ctx context.Context, authKeyID [8]byt
default:
event.Type = domain.UpdateEventQuickReplies
}
return s.recordEvent(ctx, authKeyID, userID, event, true, excludeSessionID)
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, event, true, excludeSessionID)
}
// RecordReadHistory 推进 update 状态并追加一条 read_history_inbox 事件。
func (s *Service) RecordReadHistory(ctx context.Context, authKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordReadHistory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
if userID == 0 {
userID = read.OwnerUserID
}
date := int(time.Now().Unix())
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventReadHistoryInbox,
Date: date,
Peer: read.Peer,
@ -503,8 +578,8 @@ func (s *Service) RecordReadHistory(ctx context.Context, authKeyID [8]byte, user
// RecordChannelState 记录当前账号与某频道成员关系变化（leave/kick），
// 离线设备经 difference 收到 updateChannel 后重拉 channel 状态。
func (s *Service) RecordChannelState(ctx context.Context, authKeyID [8]byte, userID, channelID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordChannelState(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventChannelState,
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
PtsCount: 1,
@ -512,8 +587,8 @@ func (s *Service) RecordChannelState(ctx context.Context, authKeyID [8]byte, use
}
// RecordContactsReset 记录通讯录视角变化，供离线设备通过 updates.getDifference 触发重拉。
func (s *Service) RecordContactsReset(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordContactsReset(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventContactsReset,
PtsCount: 1,
}, true, excludeSessionID)
@ -522,8 +597,8 @@ func (s *Service) RecordContactsReset(ctx context.Context, authKeyID [8]byte, us
// RecordDraftMessage 记录某会话云草稿变化（保存/清空都是同一事件——草稿是绝对
// 状态，重放时按 peer 重载当前值）。updateDraftMessage 无 pts 字段，走 LacksWirePts
// aux 簿记；topMsgID 是 forum 话题草稿键（复用 MaxID 列持久化）。
func (s *Service) RecordDraftMessage(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordDraftMessage(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventDraftMessage,
Peer: peer,
MaxID: topMsgID,
@ -533,8 +608,8 @@ func (s *Service) RecordDraftMessage(ctx context.Context, authKeyID [8]byte, use
// RecordDialogPinned 记录单个会话置顶状态变化；folderID 是会话所在 folder
// （0 主列表/1 归档），缺失会让离线设备把归档内置顶重放到主列表。
func (s *Service) RecordDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventDialogPinned,
Peer: peer,
Bool: pinned,
@ -544,8 +619,8 @@ func (s *Service) RecordDialogPinned(ctx context.Context, authKeyID [8]byte, use
}
// RecordPinnedDialogs 记录指定 folder 内置顶顺序变化，并把新顺序持久化给 getDifference/outbox。
func (s *Service) RecordPinnedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordPinnedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventPinnedDialogs,
Peers: append([]domain.Peer(nil), order...),
FolderID: folderID,
@ -554,8 +629,8 @@ func (s *Service) RecordPinnedDialogs(ctx context.Context, authKeyID [8]byte, us
}
// RecordSavedDialogPinned 记录收藏夹单个子会话置顶状态变化。
func (s *Service) RecordSavedDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordSavedDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventSavedDialogPinned,
Peer: peer,
Bool: pinned,
@ -564,8 +639,8 @@ func (s *Service) RecordSavedDialogPinned(ctx context.Context, authKeyID [8]byte
}
// RecordPinnedSavedDialogs 记录收藏夹置顶顺序变化，新顺序持久化给 getDifference/outbox。
func (s *Service) RecordPinnedSavedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordPinnedSavedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventPinnedSavedDialogs,
Peers: append([]domain.Peer(nil), order...),
PtsCount: 1,
@ -573,8 +648,8 @@ func (s *Service) RecordPinnedSavedDialogs(ctx context.Context, authKeyID [8]byt
}
// RecordDialogUnreadMark 记录手动未读标记变化。
func (s *Service) RecordDialogUnreadMark(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordDialogUnreadMark(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventDialogUnreadMark,
Peer: peer,
Bool: unread,
@ -583,8 +658,8 @@ func (s *Service) RecordDialogUnreadMark(ctx context.Context, authKeyID [8]byte,
}
// RecordChannelViewForumAsMessages records a per-account forum presentation state change.
func (s *Service) RecordChannelViewForumAsMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, enabled bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordChannelViewForumAsMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, enabled bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventChannelViewForum,
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
Bool: enabled,
@ -594,8 +669,8 @@ func (s *Service) RecordChannelViewForumAsMessages(ctx context.Context, authKeyI
// RecordChannelDiscussionInbox 记录 forum 话题级已读（updateReadChannelDiscussionInbox），
// 占一个账号 pts（LacksWirePts），供自己其它设备在线同步与离线差分恢复。
func (s *Service) RecordChannelDiscussionInbox(ctx context.Context, authKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordChannelDiscussionInbox(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventReadChannelDiscussionInbox,
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
TopMsgID: topicID,
@ -605,8 +680,8 @@ func (s *Service) RecordChannelDiscussionInbox(ctx context.Context, authKeyID [8
}
// RecordPeerSettings 记录 peer settings 变化。
func (s *Service) RecordPeerSettings(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordPeerSettings(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventPeerSettings,
Peer: peer,
Settings: settings,
@ -615,8 +690,8 @@ func (s *Service) RecordPeerSettings(ctx context.Context, authKeyID [8]byte, use
}
// RecordPeerStoryBlocked 记录当前账号 story blocklist 对某个 peer 的可见状态变化。
func (s *Service) RecordPeerStoryBlocked(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordPeerStoryBlocked(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventPeerStoryBlocked,
Peer: peer,
Bool: blocked,
@ -625,13 +700,13 @@ func (s *Service) RecordPeerStoryBlocked(ctx context.Context, authKeyID [8]byte,
}
// RecordDialogFilter 记录单个 filter 的创建、更新或删除；folder 为 nil 表示删除。
func (s *Service) RecordDialogFilter(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) RecordDialogFilter(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
var copyFolder *domain.DialogFolder
if folder != nil {
f := *folder
copyFolder = &f
}
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventDialogFilter,
FilterID: folderID,
DialogFilter: copyFolder,
@ -640,8 +715,8 @@ func (s *Service) RecordDialogFilter(ctx context.Context, authKeyID [8]byte, use
}
// RecordDialogFilterOrder 记录 filter 顺序变化。
func (s *Service) RecordDialogFilterOrder(ctx context.Context, authKeyID [8]byte, userID int64, order []int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordDialogFilterOrder(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventDialogFilterOrder,
FilterOrder: append([]int(nil), order...),
PtsCount: 1,
@ -649,16 +724,16 @@ func (s *Service) RecordDialogFilterOrder(ctx context.Context, authKeyID [8]byte
}
// RecordDialogFiltersReload 通知其他设备重新拉取 filter 列表。
func (s *Service) RecordDialogFiltersReload(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordDialogFiltersReload(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventDialogFilters,
PtsCount: 1,
}, true, excludeSessionID)
}
// RecordFolderPeers 记录归档/还原会话的 folder_id 变化。
func (s *Service) RecordFolderPeers(ctx context.Context, authKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordFolderPeers(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventFolderPeers,
FolderPeers: append([]domain.FolderPeerUpdate(nil), peers...),
PtsCount: 1,
@ -666,8 +741,8 @@ func (s *Service) RecordFolderPeers(ctx context.Context, authKeyID [8]byte, user
}
// RecordChannelAvailableMessages records a local channel history clear for multi-device sync.
func (s *Service) RecordChannelAvailableMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, availableMinID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
func (s *Service) RecordChannelAvailableMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, availableMinID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
Type: domain.UpdateEventChannelAvailable,
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
MaxID: availableMinID,
@ -675,15 +750,15 @@ func (s *Service) RecordChannelAvailableMessages(ctx context.Context, authKeyID
}, true, excludeSessionID)
}
func (s *Service) recordEvent(ctx context.Context, authKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEventCore(ctx, authKeyID, userID, event, dispatch, excludeSessionID, true)
func (s *Service) recordEvent(ctx context.Context, stateAuthKeyID, excludeAuthKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEventCore(ctx, stateAuthKeyID, excludeAuthKeyID, userID, event, dispatch, excludeSessionID, true)
}
func (s *Service) recordEventWithoutState(ctx context.Context, userID int64, event domain.UpdateEvent) (domain.UpdateEvent, domain.UpdateState, error) {
return s.recordEventCore(ctx, [8]byte{}, userID, event, false, 0, false)
return s.recordEventCore(ctx, [8]byte{}, [8]byte{}, userID, event, false, 0, false)
}
func (s *Service) recordEventCore(ctx context.Context, authKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64, saveState bool) (domain.UpdateEvent, domain.UpdateState, error) {
func (s *Service) recordEventCore(ctx context.Context, stateAuthKeyID, excludeAuthKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64, saveState bool) (domain.UpdateEvent, domain.UpdateState, error) {
date := event.Date
if date == 0 {
date = int(time.Now().Unix())
@ -698,7 +773,7 @@ func (s *Service) recordEventCore(ctx context.Context, authKeyID [8]byte, userID
var err error
if dispatch {
if appender, ok := s.events.(dispatchingEventAppender); ok {
event, err = appender.AppendAllocatedWithDispatch(ctx, userID, event, authKeyID, excludeSessionID)
event, err = appender.AppendAllocatedWithDispatch(ctx, userID, event, excludeAuthKeyID, excludeSessionID)
} else {
event, err = s.events.AppendAllocated(ctx, userID, event)
}
@ -735,7 +810,7 @@ func (s *Service) recordEventCore(ctx context.Context, authKeyID [8]byte, userID
st.Pts = event.Pts
}
if saveState && s.states != nil {
if err := s.states.Save(ctx, authKeyID, userID, st); err != nil {
if err := s.states.Save(ctx, stateAuthKeyID, userID, st); err != nil {
return domain.UpdateEvent{}, domain.UpdateState{}, err
}
}

190
internal/app/updates/service_test.go
View file

@ -100,7 +100,7 @@ func TestRecordReadHistoryFeedsGetDifference(t *testing.T) {
Peer: peer,
MaxID: 10,
Changed: true,
}, 0)
}, [8]byte{}, 0)
if err != nil {
t.Fatalf("RecordReadHistory: %v", err)
}
@ -132,7 +132,7 @@ func TestRecordChannelReadHistoryKeepsChannelPtsPayload(t *testing.T) {
StillUnreadCount: 3,
ChannelPts: 77,
Changed: true,
}, 0)
}, [8]byte{}, 0)
if err != nil {
t.Fatalf("RecordReadHistory: %v", err)
}
@ -157,24 +157,24 @@ func TestRecordSettingsEventsFeedGetDifference(t *testing.T) {
ownerUserID := int64(1000000001)
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
if _, _, err := svc.RecordContactsReset(ctx, authKeyID, ownerUserID, 0); err != nil {
if _, _, err := svc.RecordContactsReset(ctx, authKeyID, ownerUserID, [8]byte{}, 0); err != nil {
t.Fatalf("RecordContactsReset: %v", err)
}
if _, _, err := svc.RecordDialogPinned(ctx, authKeyID, ownerUserID, peer, true, 0, 0); err != nil {
if _, _, err := svc.RecordDialogPinned(ctx, authKeyID, ownerUserID, peer, true, 0, [8]byte{}, 0); err != nil {
t.Fatalf("RecordDialogPinned: %v", err)
}
order := []domain.Peer{peer}
if _, _, err := svc.RecordPinnedDialogs(ctx, authKeyID, ownerUserID, 0, order, 0); err != nil {
if _, _, err := svc.RecordPinnedDialogs(ctx, authKeyID, ownerUserID, 0, order, [8]byte{}, 0); err != nil {
t.Fatalf("RecordPinnedDialogs: %v", err)
}
if _, _, err := svc.RecordDialogUnreadMark(ctx, authKeyID, ownerUserID, peer, false, 0); err != nil {
if _, _, err := svc.RecordDialogUnreadMark(ctx, authKeyID, ownerUserID, peer, false, [8]byte{}, 0); err != nil {
t.Fatalf("RecordDialogUnreadMark: %v", err)
}
settings := domain.PeerSettings{ShareContact: true}
if _, _, err := svc.RecordPeerSettings(ctx, authKeyID, ownerUserID, peer, settings, 0); err != nil {
if _, _, err := svc.RecordPeerSettings(ctx, authKeyID, ownerUserID, peer, settings, [8]byte{}, 0); err != nil {
t.Fatalf("RecordPeerSettings: %v", err)
}
stateEvent, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, ownerUserID, peer, true, 0)
stateEvent, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, ownerUserID, peer, true, [8]byte{}, 0)
if err != nil {
t.Fatalf("RecordPeerStoryBlocked: %v", err)
}
@ -221,22 +221,29 @@ func TestRecordSettingsEventsFeedGetDifference(t *testing.T) {
func TestRecordSettingsEventUsesDispatchAppender(t *testing.T) {
ctx := context.Background()
var authKeyID [8]byte
authKeyID[0] = 4
authKeyID := [8]byte{4}
rawAuthKeyID := [8]byte{4, 9}
events := &captureDispatchAppender{UpdateEventStore: memory.NewUpdateEventStore()}
svc := NewService(memory.NewUpdateStateStore(), events)
states := &captureStateStore{}
svc := NewService(states, events)
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
event, state, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, peer, true, 0, 42)
event, state, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, peer, true, 0, rawAuthKeyID, 42)
if err != nil {
t.Fatalf("RecordDialogPinned: %v", err)
}
if event.Pts != 1 || state.Pts != 1 {
t.Fatalf("event/state = %+v / %+v, want first pts", event, state)
}
if !events.dispatched || events.excludeAuthKeyID != authKeyID || events.excludeSessionID != 42 || events.event.Type != domain.UpdateEventDialogPinned || events.event.Peer != peer {
if !events.dispatched || events.excludeAuthKeyID != rawAuthKeyID || events.excludeSessionID != 42 || events.event.Type != domain.UpdateEventDialogPinned || events.event.Peer != peer {
t.Fatalf("dispatch capture = %+v exclude_auth=%v exclude_session=%d dispatched=%v, want dialog_pinned outbox", events.event, events.excludeAuthKeyID, events.excludeSessionID, events.dispatched)
}
if states.lastSaveAuthKeyID != authKeyID {
t.Fatalf("device state auth key = %x, want business/perm %x", states.lastSaveAuthKeyID, authKeyID)
}
if _, found, err := states.Get(ctx, rawAuthKeyID, 1000000001); err != nil || found {
t.Fatalf("raw temp key unexpectedly owns device state: found=%v err=%v", found, err)
}
}
func TestRecordSettingsEventDispatchFailureDoesNotRecordEvent(t *testing.T) {
@ -246,7 +253,7 @@ func TestRecordSettingsEventDispatchFailureDoesNotRecordEvent(t *testing.T) {
events := &failingDispatchAppender{UpdateEventStore: memory.NewUpdateEventStore()}
svc := NewService(memory.NewUpdateStateStore(), events)
_, _, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}, true, 0, 42)
_, _, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}, true, 0, authKeyID, 42)
if !errors.Is(err, errDispatchFailed) {
t.Fatalf("RecordDialogPinned err = %v, want dispatch failure", err)
}
@ -267,7 +274,7 @@ func TestRecordPeerStoryBlockedUsesDispatchAppender(t *testing.T) {
svc := NewService(memory.NewUpdateStateStore(), events)
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
event, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, 1000000001, peer, true, 91)
event, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, 1000000001, peer, true, authKeyID, 91)
if err != nil {
t.Fatalf("RecordPeerStoryBlocked: %v", err)
}
@ -294,7 +301,7 @@ func TestRecordStoryUsesDispatchAppenderExcludeCurrentSession(t *testing.T) {
Caption: "owner story",
}
event, state, err := svc.RecordStory(ctx, authKeyID, owner.ID, story, 1234)
event, state, err := svc.RecordStory(ctx, authKeyID, owner.ID, story, authKeyID, 1234)
if err != nil {
t.Fatalf("RecordStory: %v", err)
}
@ -330,7 +337,7 @@ func TestRecordStoryReadAndSentReactionExcludeCurrentSession(t *testing.T) {
MaxReadID: story.ID,
Advanced: true,
Date: 1700000201,
}, 2233)
}, authKeyID, 2233)
if err != nil {
t.Fatalf("RecordReadStories: %v", err)
}
@ -350,7 +357,7 @@ func TestRecordStoryReadAndSentReactionExcludeCurrentSession(t *testing.T) {
Reaction: reaction,
Changed: true,
Date: 1700000202,
}, 2233)
}, authKeyID, 2233)
if err != nil {
t.Fatalf("RecordSentStoryReaction: %v", err)
}
@ -384,7 +391,7 @@ func TestRecordNewStoryReactionDispatchesWithoutSavingDeviceState(t *testing.T)
},
Reaction: reaction,
Date: 1700000101,
}, 0)
}, [8]byte{}, 0)
if err != nil {
t.Fatalf("RecordNewStoryReaction: %v", err)
}
@ -493,7 +500,8 @@ func TestAcknowledgeCurrentStateAdvancesConfirmedWatermark(t *testing.T) {
authKeyID[0] = 11
userID := int64(1000000001)
events := memory.NewUpdateEventStore()
svc := NewService(memory.NewUpdateStateStore(), events)
states := memory.NewUpdateStateStore()
svc := NewService(states, events)
if err := events.Append(ctx, userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNewMessage, Pts: 1, PtsCount: 1,
Date: 1700000001, Message: domain.Message{ID: 1, OwnerUserID: userID},
@ -527,6 +535,138 @@ func TestAcknowledgeCurrentStateAdvancesConfirmedWatermark(t *testing.T) {
if confirmed.Pts != 3 {
t.Fatalf("confirmed watermark = %d, want advanced to 3", confirmed.Pts)
}
observed, ok := states.ObservedClientState(authKeyID, userID)
if !ok || observed.Pts != 3 {
t.Fatalf("getState observed watermark = %+v/%v, want pts=3", observed, ok)
}
}
func TestGetDifferenceRetainsOnlyClientObservedInputCursor(t *testing.T) {
ctx := context.Background()
authKeyID := [8]byte{12}
const userID int64 = 1000000012
events := memory.NewUpdateEventStore()
states := memory.NewUpdateStateStore()
svc := NewService(states, events)
for pts := 1; pts <= 2; pts++ {
if err := events.Append(ctx, userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNewMessage, Pts: pts, PtsCount: 1,
Date: 1700000100 + pts, Message: domain.Message{ID: pts, OwnerUserID: userID},
}); err != nil {
t.Fatalf("append pts=%d: %v", pts, err)
}
}
// 服务端把 pts=1..2 放进 response，并不证明客户端收到了 response；observed 只能
// 保持在本次 request 实际携带的 pts=0。
diff, err := svc.GetDifference(ctx, authKeyID, userID, domain.UpdateState{Pts: 0, Date: 1700000100})
if err != nil {
t.Fatalf("first difference: %v", err)
}
if diff.State.Pts != 2 || len(diff.Events) != 2 {
t.Fatalf("first difference = %+v, want response through pts=2", diff)
}
observed, ok := states.ObservedClientState(authKeyID, userID)
if !ok || observed.Pts != 0 {
t.Fatalf("observed after merely sending response = %+v/%v, want pts=0", observed, ok)
}
// 客户端下一次明确带回 pts=2 后，才允许 retention 把共同安全水位推进到 2。
if _, err := svc.GetDifference(ctx, authKeyID, userID, domain.UpdateState{Pts: 2, Date: 1700000102}); err != nil {
t.Fatalf("confirming difference: %v", err)
}
observed, ok = states.ObservedClientState(authKeyID, userID)
if !ok || observed.Pts != 2 {
t.Fatalf("observed after client carried cursor = %+v/%v, want pts=2", observed, ok)
}
}
type retentionCheckpointEvents struct {
*memory.UpdateEventStore
pts int
date int
current int
missFirst bool
calls int
}
func (s *retentionCheckpointEvents) UserUpdateRetentionCheckpoint(_ context.Context, _ [8]byte, _ int64) (int, int, bool, error) {
s.calls++
if s.missFirst && s.calls == 1 {
return 0, 0, false, nil
}
return s.pts, s.date, s.pts > 0, nil
}
func (s *retentionCheckpointEvents) MaxContiguousPts(_ context.Context, _ int64) (int, error) {
return s.current, nil
}
func TestGetDifferenceBelowRetainedFloorUsesEmptySliceCheckpoint(t *testing.T) {
ctx := context.Background()
authKeyID := [8]byte{13}
const userID int64 = 1000000013
base := memory.NewUpdateEventStore()
events := &retentionCheckpointEvents{UpdateEventStore: base, pts: 2, date: 1700000202, current: 3}
states := memory.NewUpdateStateStore()
svc := NewService(states, events)
// Retention already removed pts 1..2; only the live tail remains.
if err := base.Append(ctx, userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNoop, Pts: 3, PtsCount: 1, Date: 1700000203,
}); err != nil {
t.Fatalf("append live tail: %v", err)
}
checkpoint, err := svc.GetDifference(ctx, authKeyID, userID, domain.UpdateState{Pts: 0, Date: 1700000200})
if err != nil {
t.Fatalf("difference below retained floor: %v", err)
}
if !checkpoint.Partial || len(checkpoint.Events) != 0 || checkpoint.State.Pts != 2 || checkpoint.State.Date != 1700000202 {
t.Fatalf("checkpoint difference = %+v, want empty differenceSlice at pts/date 2/1700000202", checkpoint)
}
tail, err := svc.GetDifference(ctx, authKeyID, userID, checkpoint.State)
if err != nil {
t.Fatalf("difference from retained floor: %v", err)
}
if tail.Partial || len(tail.Events) != 1 || tail.Events[0].Pts != 3 || tail.State.Pts != 3 {
t.Fatalf("tail difference = %+v, want normal event pts=3", tail)
}
}
func TestGetDifferenceRechecksCheckpointWhenRetentionRacesEventRead(t *testing.T) {
ctx := context.Background()
authKeyID := [8]byte{14}
const userID int64 = 1000000014
base := memory.NewUpdateEventStore()
events := &retentionCheckpointEvents{
UpdateEventStore: base,
pts: 2,
date: 1700000302,
current: 3,
missFirst: true,
}
if err := base.Append(ctx, userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNoop, Pts: 3, PtsCount: 1, Date: 1700000303,
}); err != nil {
t.Fatalf("append live tail: %v", err)
}
diff, err := NewService(memory.NewUpdateStateStore(), events).GetDifference(
ctx,
authKeyID,
userID,
domain.UpdateState{Pts: 0, Date: 1700000300},
)
if err != nil {
t.Fatalf("difference across retention race: %v", err)
}
if events.calls != 2 {
t.Fatalf("checkpoint probes = %d, want pre-read plus post-gap recheck", events.calls)
}
if !diff.Partial || len(diff.Events) != 0 || diff.State.Pts != 2 || diff.State.Date != 1700000302 {
t.Fatalf("race checkpoint difference = %+v, want empty differenceSlice at retained floor", diff)
}
}
type captureDispatchAppender struct {
@ -559,8 +699,9 @@ func (s *failingDispatchAppender) AppendAllocatedWithDispatch(context.Context, i
}
type captureStateStore struct {
saveCount int
states map[[16]byte]domain.UpdateState
saveCount int
lastSaveAuthKeyID [8]byte
states map[[16]byte]domain.UpdateState
}
func (s *captureStateStore) Get(_ context.Context, authKeyID [8]byte, userID int64) (domain.UpdateState, bool, error) {
@ -576,10 +717,15 @@ func (s *captureStateStore) Save(_ context.Context, authKeyID [8]byte, userID in
s.states = make(map[[16]byte]domain.UpdateState)
}
s.saveCount++
s.lastSaveAuthKeyID = authKeyID
s.states[captureStateKey(authKeyID, userID)] = state
return nil
}
func (s *captureStateStore) ObserveClientState(_ context.Context, _ [8]byte, _ int64, _ domain.UpdateState) error {
return nil
}
func (s *captureStateStore) Delete(_ context.Context, authKeyID [8]byte, userID int64) error {
if s.states != nil {
delete(s.states, captureStateKey(authKeyID, userID))

3
internal/compat/layerwire/README.md
View file

@ -20,6 +20,7 @@
| `schema/canonical-227.tl` | **embed**，运行期 walker 的 227 字段布局（= gotd `td/_schema/tdesktop.tl` 的副本） | gotd 升级时 re-sync |
| `_schema/layer-2NN.tl` | 历史层官方 schema（从 TDesktop git 抽，**仅生成期用**，下划线=不编译/不 embed） | 升级/下探 floor 时抽取 |
| `schema/client-drift.tl` | **声明式**：客户端发的旧构造器老布局（body 与 227 不同的） | 发现客户端漂移时 +1 行 |
| `schema/routable-compat.tl` | **仅结构预检**：已有 RPC fallback adapter 的非 canonical wire 布局（当前只含 4 个 DrKLO theme 构造器）；与 canonical 图合并后完整 walk，但不自动升级 | 收敛既有手写 adapter 时维护，禁止借此新增业务 fallback |
| `client_aliases.go` | 客户端漂移里 **body 与 227 字节一致**的，纯 `老CRC→227CRC` | 发现纯换 CRC 漂移时 +1 条 |
| `tables_gen.go` | **生成产物**（勿手改）：官方层降级表 + 入站升级表 + 新类型集 | 跑 `gen` 重生成 |
| `gen/main.go` | 生成器：对拍 schema、证明机械性、产 `tables_gen.go` | 升级逻辑变更时 |
@ -95,7 +96,7 @@ gofmt -w internal/compat/layerwire/ && go build ./... && go vet ./internal/...
- 绿 = 通用引擎已能自动升级（复制共享字段 + 插 flags=0 + 按 kind 补默认）。**完事**。
- `TestInboundDriftCoverage` 报 `needs converter A->B` = 有字段类型变更 → 往 `inbound.go fieldConverters` 加一条 `"A->B"`（可复用，参照 `Vector<int>->Vector<InputMessage>`）。
- 报 `field X not defaultable` 或字段**改名** → 往 `inbound.go driftFieldRenames` 加 `"<method>\x00<227字段>": "<老字段>"`（参照 `bots.exportBotToken\x00bot`）。
4. **绝不**为此写一个新的 `handleLegacyXxx` 解码 handler——那是旧做法，已全删。统一走数据 + 通用引擎。
4. **绝不**为此写一个新的 `handleLegacyXxx` 解码 handler——统一走数据 + 通用引擎。`routable-compat.tl` 只给既存 DrKLO theme fallback 补 dispatcher 前结构门禁，不是新增 adapter 的入口。
## 操作 4：出站 `TestCoverageGate` 失败

28
internal/compat/layerwire/fallback.go
View file

@ -30,8 +30,11 @@ func init() {
// replaceWithBare consumes the canonical (227-only) object and emits a
// bodyless constructor id the target layer understands.
func replaceWithBare(id uint32) fallbackFunc {
return func(cl *ctorLayout, in, out *bin.Buffer, layer int) error {
if err := canonical.skipObject(in); err != nil {
return func(cl *ctorLayout, in, out *bin.Buffer, layer, depth int, walk *walkState) error {
if err := in.ConsumeID(cl.crc); err != nil {
return err
}
if err := walk.skipCtorBody(canonical, in, cl, depth); err != nil {
return err
}
out.PutID(id)
@ -47,6 +50,8 @@ var peerVectorField = fieldLayout{
elem: &fieldLayout{kind: kindObject, typeName: "Peer", flagBit: -1},
}
var pollOptionBytesField = fieldLayout{kind: kindBytes, flagBit: -1}
// transcodePollAnswerVoters downgrades pollAnswerVoters: canonical (227) made
// voters conditional (flags.2?int) and added recent_voters (flags.2?Vector<Peer>);
// older layers carry voters as a plain int. The leading CRC is already consumed.
@ -54,34 +59,35 @@ var peerVectorField = fieldLayout{
// 227: flags:# chosen:flags.0?true correct:flags.1?true option:bytes
// voters:flags.2?int recent_voters:flags.2?Vector<Peer>
// <=226: flags:# chosen:flags.0?true correct:flags.1?true option:bytes voters:int
func transcodePollAnswerVoters(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer int) error {
func transcodePollAnswerVoters(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer, depth int, walk *walkState) error {
flags, err := in.Uint32()
if err != nil {
return err
}
option, err := in.Bytes()
if err != nil {
optionStart := in.Buf
if err := walk.skipValue(canonical, in, &pollOptionBytesField, cl, depth); err != nil {
return err
}
optionRaw := optionStart[:len(optionStart)-len(in.Buf)]
var voters int
if flags&(1<<2) != 0 {
if voters, err = in.Int(); err != nil {
return err
}
if err := canonical.skipValue(in, &peerVectorField); err != nil {
if err := walk.skipValue(canonical, in, &peerVectorField, cl, depth); err != nil {
return err
}
}
out.PutID(target)
out.PutUint32(flags & 0b11) // retain chosen/correct, clear the moved bit 2
out.PutBytes(option)
out.Put(optionRaw)
out.PutInt(voters)
return nil
}
// fallbackMessageEntity replaces any 227-only MessageEntity with
// messageEntityUnknown, preserving offset/length so text positions stay valid.
func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer int) error {
func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer, depth int, walk *walkState) error {
id, err := in.PeekID()
if err != nil {
return err
@ -89,7 +95,7 @@ func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer int) error
if err := in.ConsumeID(id); err != nil {
return err
}
offset, length, err := canonical.decodeOffsetLength(in, cl)
offset, length, err := canonical.decodeOffsetLength(in, cl, depth, walk)
if err != nil {
return err
}
@ -101,7 +107,7 @@ func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer int) error
// decodeOffsetLength walks a constructor body (no leading CRC) per the canonical
// layout, returning its offset/length int fields and discarding the rest.
func (m *schemaModel) decodeOffsetLength(in *bin.Buffer, cl *ctorLayout) (offset, length int, err error) {
func (m *schemaModel) decodeOffsetLength(in *bin.Buffer, cl *ctorLayout, depth int, walk *walkState) (offset, length int, err error) {
var flags map[string]uint32
for i := range cl.fields {
f := &cl.fields[i]
@ -129,7 +135,7 @@ func (m *schemaModel) decodeOffsetLength(in *bin.Buffer, cl *ctorLayout) (offset
return
}
default:
if err = m.skipValue(in, f); err != nil {
if err = walk.skipValue(m, in, f, cl, depth); err != nil {
return
}
}

59
internal/compat/layerwire/inbound.go
View file

@ -47,16 +47,19 @@ var driftFieldRenames = map[string]string{
// fieldConverter rewrites one field whose wire type changed between the old and
// canonical layout. Keyed by "<oldTypeSig>-><newTypeSig>"; raw is the old field's
// encoded bytes. Reusable across any method with the same type change.
type fieldConverter func(raw []byte, out *bin.Buffer) error
type fieldConverter func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error
var fieldConverters = map[string]fieldConverter{
// id:Vector<int> -> id:Vector<InputMessage> (wrap each int in inputMessageID).
"Vector<int>->Vector<InputMessage>": func(raw []byte, out *bin.Buffer) error {
"Vector<int>->Vector<InputMessage>": func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error {
in := &bin.Buffer{Buf: raw}
n, err := in.VectorHeader()
if err != nil {
return err
}
if max := walk.vectorLimit(owner, field); n > max {
return limitf("vector %s.%s length %d exceeds limit %d", ownerName(owner), fieldName(field), n, max)
}
out.PutVectorHeader(n)
for i := 0; i < n; i++ {
v, err := in.Int()
@ -66,10 +69,13 @@ var fieldConverters = map[string]fieldConverter{
out.PutID(inputMessageID)
out.PutInt(v)
}
if in.Len() != 0 {
return malformedf("%d trailing bytes in Vector<int> converter", in.Len())
}
return nil
},
// bot_id:long -> bot:InputUser{user_id, access_hash=0}.
"long->InputUser": func(raw []byte, out *bin.Buffer) error {
"long->InputUser": func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error {
in := &bin.Buffer{Buf: raw}
id, err := in.Long()
if err != nil {
@ -78,11 +84,14 @@ var fieldConverters = map[string]fieldConverter{
out.PutID(inputUserID)
out.PutLong(id)
out.PutLong(0)
if in.Len() != 0 {
return malformedf("%d trailing bytes in long converter", in.Len())
}
return nil
},
// channel:InputChannel -> peer:InputPeer for the old channels.editCreator
// Android constructor. Concrete layouts are otherwise byte-compatible.
"InputChannel->InputPeer": func(raw []byte, out *bin.Buffer) error {
"InputChannel->InputPeer": func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error {
in := &bin.Buffer{Buf: raw}
id, err := in.ID()
if err != nil {
@ -116,8 +125,12 @@ var fieldConverters = map[string]fieldConverter{
// id + body) is what to dispatch.
func UpgradeInbound(id uint32, in *bin.Buffer) (*bin.Buffer, bool, error) {
if newID, ok := UpgradeMethodCRC(id); ok {
if len(in.Buf) < 4 {
return nil, false, fmt.Errorf("layerwire: short inbound buffer for %#08x", id)
target := canonical.byCRC[newID]
if target == nil || !target.isFunc {
return nil, true, malformedf("alias %#08x targets unknown canonical method %#08x", id, newID)
}
if err := validateAliasedMethod(id, target, in.Buf); err != nil {
return nil, true, err
}
// Copy rather than rewrite in place: never mutate the caller's buffer
// (matches the body-transform path, which also returns a fresh buffer).
@ -126,15 +139,36 @@ func UpgradeInbound(id uint32, in *bin.Buffer) (*bin.Buffer, bool, error) {
return out, true, nil
}
if old := driftModel.byCRC[id]; old != nil {
out, err := upgradeFromDrift(old, in)
out, err := upgradeFromDrift(old, in, newWalkState())
if err != nil {
return nil, false, fmt.Errorf("layerwire: upgrade %s (%#08x): %w", old.name, id, err)
return nil, true, classifyWalkError(fmt.Errorf("layerwire: upgrade %s (%#08x): %w", old.name, id, err))
}
return out, true, nil
}
return nil, false, nil
}
// validateAliasedMethod validates the old-id/canonical-body shape before
// allocating the replacement buffer. The body is walked against the canonical
// target layout while the original constructor id remains untouched.
func validateAliasedMethod(oldID uint32, target *ctorLayout, raw []byte) error {
walk := newWalkState()
if err := walk.enter(1, "constructor"); err != nil {
return err
}
b := &bin.Buffer{Buf: raw}
if err := b.ConsumeID(oldID); err != nil {
return classifyWalkError(err)
}
if err := walk.skipCtorBody(canonical, b, target, 1); err != nil {
return classifyWalkError(err)
}
if b.Len() != 0 {
return malformedf("%d trailing bytes after aliased method %s", b.Len(), target.name)
}
return nil
}
// IsClientDrift reports whether id is a client-private constructor (DrKLO
// constructor drift), as opposed to official layer drift from api.tl.
func IsClientDrift(id uint32) bool {
@ -146,11 +180,14 @@ func IsClientDrift(id uint32) bool {
// upgradeFromDrift rebuilds a canonical (227) request from an old client-drift
// body, comparing the declared old layout to the canonical layout field by field.
func upgradeFromDrift(old *ctorLayout, in *bin.Buffer) (*bin.Buffer, error) {
func upgradeFromDrift(old *ctorLayout, in *bin.Buffer, walk *walkState) (*bin.Buffer, error) {
target := canonical.byName[old.name]
if target == nil {
return nil, fmt.Errorf("no canonical method %q", old.name)
}
if err := walk.enter(1, "constructor"); err != nil {
return nil, err
}
if err := in.ConsumeID(old.crc); err != nil {
return nil, err
}
@ -179,7 +216,7 @@ func upgradeFromDrift(old *ctorLayout, in *bin.Buffer) (*bin.Buffer, error) {
continue
}
pre := in.Buf
if err := canonical.skipValue(in, f); err != nil {
if err := walk.skipValue(canonical, in, f, old, 1); err != nil {
return nil, fmt.Errorf("decode old field %q: %w", f.name, err)
}
vals[f.name] = pre[:len(pre)-len(in.Buf)]
@ -208,7 +245,7 @@ func upgradeFromDrift(old *ctorLayout, in *bin.Buffer) (*bin.Buffer, error) {
if conv == nil {
return nil, fmt.Errorf("field %q: no converter %s->%s", nf.name, typeSig(of), typeSig(nf))
}
if err := conv(vals[oldName], out); err != nil {
if err := conv(vals[oldName], out, walk, old, of); err != nil {
return nil, fmt.Errorf("field %q convert: %w", nf.name, err)
}
} else {

80
internal/compat/layerwire/routable.go Normal file
View file

@ -0,0 +1,80 @@
package layerwire
import (
_ "embed"
"fmt"
"github.com/gotd/td/bin"
)
const maxOpaqueRequestBytes = 16 << 20
//go:embed schema/routable-compat.tl
var routableCompatSchema string
// routable combines the canonical Layer 227 model with the small set of
// explicitly declared compatibility-only methods. Nested objects in those
// methods are canonical Input* constructors, so one combined graph is needed
// for the same depth/vector/bytes walker to validate the complete request.
var routable = mustLoadRoutable()
func mustLoadRoutable() *schemaModel {
compat, err := parseSchemaModel(routableCompatSchema)
if err != nil {
panic("layerwire: parse routable compat schema: " + err.Error())
}
m := &schemaModel{
byCRC: make(map[uint32]*ctorLayout, len(canonical.byCRC)+len(compat.byCRC)),
byName: make(map[string]*ctorLayout, len(canonical.byName)+len(compat.byName)),
bareByT: make(map[string]*ctorLayout, len(canonical.bareByT)),
ctorsOfT: make(map[string][]*ctorLayout, len(canonical.ctorsOfT)),
}
for id, cl := range canonical.byCRC {
m.byCRC[id] = cl
}
for name, cl := range canonical.byName {
m.byName[name] = cl
}
for name, cl := range canonical.bareByT {
m.bareByT[name] = cl
}
for name, ctors := range canonical.ctorsOfT {
m.ctorsOfT[name] = ctors
}
for id, cl := range compat.byCRC {
if existing := m.byCRC[id]; existing != nil {
panic(fmt.Sprintf("layerwire: routable compat crc %#08x collides with %s", id, existing.name))
}
m.byCRC[id] = cl
m.byName[cl.name] = cl
}
return m
}
// ValidateRoutableRequest validates every request shape the router knows how to
// decode, including compatibility-only fallback methods. known=false denotes
// a genuinely unknown top-level constructor. Such a request is never decoded:
// it is treated as opaque, word-aligned TL data, bounded by both this total-size
// cap and mtprotoedge's transport/RPC budgets, and must continue to the router's
// compatibility trace rather than being mislabeled as malformed input.
func ValidateRoutableRequest(body []byte) (known bool, err error) {
b := &bin.Buffer{Buf: body}
id, err := b.PeekID()
if err != nil {
return false, classifyWalkError(err)
}
cl := routable.byCRC[id]
if cl == nil {
if len(body) > maxOpaqueRequestBytes {
return false, limitf("opaque request length %d exceeds limit %d", len(body), maxOpaqueRequestBytes)
}
if len(body)%bin.Word != 0 {
return false, malformedf("opaque request length %d is not word aligned", len(body))
}
return false, nil
}
if !cl.isFunc {
return true, malformedf("constructor %s (%#08x) is not a method", cl.name, id)
}
return true, validateRequestLayout(routable, cl, body)
}

43
internal/compat/layerwire/routable_test.go Normal file
View file

@ -0,0 +1,43 @@
package layerwire
import (
"errors"
"testing"
"github.com/gotd/td/bin"
"github.com/gotd/td/tg"
)
func TestValidateRoutableRequestCompatibilityAndUnknown(t *testing.T) {
t.Run("legacy theme is fully walked", func(t *testing.T) {
var b bin.Buffer
b.PutID(0x8d9d742b)
b.PutString("android")
(&tg.InputThemeSlug{Slug: "night"}).Encode(&b)
b.PutLong(42)
known, err := ValidateRoutableRequest(b.Buf)
if err != nil || !known {
t.Fatalf("legacy theme known=%v err=%v, want true/nil", known, err)
}
b.Buf = b.Buf[:len(b.Buf)-4]
known, err = ValidateRoutableRequest(b.Buf)
if !known || !errors.Is(err, ErrMalformed) {
t.Fatalf("truncated legacy theme known=%v err=%v, want true/malformed", known, err)
}
})
t.Run("unknown stays opaque and bounded", func(t *testing.T) {
var b bin.Buffer
b.PutID(0x12345678)
b.PutUint32(0xffffffff)
known, err := ValidateRoutableRequest(b.Buf)
if err != nil || known {
t.Fatalf("opaque unknown known=%v err=%v, want false/nil", known, err)
}
known, err = ValidateRoutableRequest(append(b.Buf, 1))
if known || !errors.Is(err, ErrMalformed) {
t.Fatalf("unaligned unknown known=%v err=%v, want false/malformed", known, err)
}
})
}

11
internal/compat/layerwire/schema/routable-compat.tl Normal file
View file

@ -0,0 +1,11 @@
// Hand-maintained request layouts that are intentionally handled by the RPC
// fallback instead of gotd's canonical ServerDispatcher. They still belong in
// the structural preflight model: fallback handlers must never become a way to
// bypass the canonical vector/depth/bytes budgets.
---functions---
compat.legacyCreateTheme#8432c21f flags:# slug:string title:string document:flags.2?InputDocument settings:flags.3?InputThemeSettings = Object;
compat.legacyUpdateTheme#5cb367d5 flags:# format:string theme:InputTheme slug:flags.0?string title:flags.1?string document:flags.2?InputDocument settings:flags.3?InputThemeSettings = Object;
compat.legacyInstallTheme#7ae43737 flags:# dark:flags.0?true format:flags.1?string theme:flags.1?InputTheme = Object;
compat.legacyGetTheme#8d9d742b format:string theme:InputTheme document_id:long = Object;

62
internal/compat/layerwire/tables.go
View file

@ -139,11 +139,11 @@ func (lt *layerTables) fieldDirty(f *fieldLayout) bool {
// field drop. The leading CRC has already been consumed from in; the transform
// reads the canonical body from in and writes the target-layer object (whose
// constructor id is target) to out.
type structuralFunc func(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer int) error
type structuralFunc func(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer, depth int, walk *walkState) error
// fallbackFunc replaces a layer-absent (227-only) constructor with an
// equivalent the target layer understands. The leading CRC is NOT yet consumed.
type fallbackFunc func(cl *ctorLayout, in, out *bin.Buffer, layer int) error
type fallbackFunc func(cl *ctorLayout, in, out *bin.Buffer, layer, depth int, walk *walkState) error
// structuralTransforms and the newType fallback registries are populated in
// fallback.go. newTypeFallbacks is keyed by canonical CRC (specific override);
@ -175,11 +175,12 @@ func Transcode(canonicalBytes []byte, layer int) ([]byte, error) {
}
in := &bin.Buffer{Buf: canonicalBytes}
out := &bin.Buffer{}
if err := lt.transcodeObject(in, out, layer); err != nil {
return nil, err
walk := newWalkState()
if err := lt.transcodeObject(in, out, layer, 1, walk); err != nil {
return nil, classifyWalkError(err)
}
if in.Len() != 0 {
return nil, fmt.Errorf("layerwire: %d trailing bytes after transcode to layer %d", in.Len(), layer)
return nil, malformedf("%d trailing bytes after transcode to layer %d", in.Len(), layer)
}
return out.Buf, nil
}
@ -198,7 +199,10 @@ func UpgradeMethodCRC(oldID uint32) (uint32, bool) {
return newID, ok
}
func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer, depth int, walk *walkState) error {
if err := walk.enter(depth, "constructor"); err != nil {
return err
}
id, err := in.PeekID()
if err != nil {
return err
@ -216,10 +220,10 @@ func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
if fn == nil {
return fmt.Errorf("layerwire: no structural transform %q for %s@%d", rule.structural, cl.name, layer)
}
return fn(cl, rule.target, in, out, layer)
return fn(cl, rule.target, in, out, layer, depth, walk)
}
out.PutID(rule.target)
return lt.transcodeBody(in, out, cl, rule.keep, layer)
return lt.transcodeBody(in, out, cl, rule.keep, layer, depth, walk)
}
if lt.newTypes[id] {
fn := newTypeFallbacks[id]
@ -229,12 +233,15 @@ func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
if fn == nil {
return fmt.Errorf("layerwire: %s (%#08x) absent at layer %d and no fallback", cl.name, id, layer)
}
return fn(cl, in, out, layer)
return fn(cl, in, out, layer, depth, walk)
}
if !lt.dirty[id] {
// Unaffected subtree: byte-for-byte copy.
pre := in.Buf
if err := canonical.skipObject(in); err != nil {
if err := in.ConsumeID(id); err != nil {
return err
}
if err := walk.skipCtorBody(canonical, in, cl, depth); err != nil {
return err
}
out.Put(pre[:len(pre)-len(in.Buf)])
@ -245,14 +252,14 @@ func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
return err
}
out.PutID(id)
return lt.transcodeBody(in, out, cl, nil, layer)
return lt.transcodeBody(in, out, cl, nil, layer, depth, walk)
}
// transcodeBody re-encodes a constructor body. keep==nil means retain every
// field (recursing into dirty descendants); otherwise only the named canonical
// fields are written, flag integers are remasked to the retained bits, and
// dropped fields are read-and-discarded.
func (lt *layerTables) transcodeBody(in, out *bin.Buffer, cl *ctorLayout, keep map[string]bool, layer int) error {
func (lt *layerTables) transcodeBody(in, out *bin.Buffer, cl *ctorLayout, keep map[string]bool, layer, depth int, walk *walkState) error {
kept := func(name string) bool { return keep == nil || keep[name] }
var flags map[string]uint32
for i := range cl.fields {
@ -276,10 +283,10 @@ func (lt *layerTables) transcodeBody(in, out *bin.Buffer, cl *ctorLayout, keep m
continue
}
if kept(f.name) {
if err := lt.transcodeValue(in, out, f, layer); err != nil {
if err := lt.transcodeValue(in, out, f, cl, layer, depth, walk); err != nil {
return fmt.Errorf("%s.%s: %w", cl.name, f.name, err)
}
} else if err := canonical.skipValue(in, f); err != nil {
} else if err := walk.skipValue(canonical, in, f, cl, depth); err != nil {
return fmt.Errorf("%s.%s (drop): %w", cl.name, f.name, err)
}
}
@ -301,10 +308,10 @@ func (lt *layerTables) keptMask(cl *ctorLayout, flagName string, kept func(strin
// transcodeValue writes one present field value, recursing only into dirty
// subtrees and byte-copying everything else.
func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, layer int) error {
func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, owner *ctorLayout, layer, depth int, walk *walkState) error {
if !lt.fieldDirty(f) {
pre := in.Buf
if err := canonical.skipValue(in, f); err != nil {
if err := walk.skipValue(canonical, in, f, owner, depth); err != nil {
return err
}
out.Put(pre[:len(pre)-len(in.Buf)])
@ -312,6 +319,10 @@ func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, layer
}
switch f.kind {
case kindVector, kindVectorBare:
vectorDepth := depth + 1
if vectorDepth <= 0 || vectorDepth > walk.limits.maxDepth {
return limitf("vector nesting depth %d exceeds limit %d", vectorDepth, walk.limits.maxDepth)
}
if f.kind == kindVector {
id, err := in.Uint32()
if err != nil {
@ -326,23 +337,36 @@ func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, layer
if err != nil {
return err
}
if n < 0 {
return malformedf("negative vector length %d", n)
}
if max := walk.vectorLimit(owner, f); n > max {
return limitf("vector %s.%s length %d exceeds limit %d", ownerName(owner), fieldName(f), n, max)
}
if err := walk.addUnits(n, "vector "+ownerName(owner)+"."+fieldName(f)); err != nil {
return err
}
out.PutInt(n)
for i := 0; i < n; i++ {
if err := lt.transcodeValue(in, out, f.elem, layer); err != nil {
if err := lt.transcodeValue(in, out, f.elem, nil, layer, vectorDepth, walk); err != nil {
return err
}
}
return nil
case kindObject:
return lt.transcodeObject(in, out, layer)
return lt.transcodeObject(in, out, layer, depth+1, walk)
case kindBareObject:
bareDepth := depth + 1
if err := walk.enter(bareDepth, "bare constructor"); err != nil {
return err
}
cl, ok := canonical.bareByT[f.typeName]
if !ok {
return fmt.Errorf("unknown bare type %q", f.typeName)
}
// Bare objects have no CRC and (within 220..227) no changed bare ctor;
// recurse all-kept to reach any dirty descendants.
return lt.transcodeBody(in, out, cl, nil, layer)
return lt.transcodeBody(in, out, cl, nil, layer, bareDepth, walk)
default:
// Primitive marked dirty should be impossible.
return fmt.Errorf("unexpected dirty primitive kind %d", f.kind)

397
internal/compat/layerwire/walk.go
View file

@ -1,32 +1,205 @@
package layerwire
import (
"errors"
"fmt"
"io"
"math"
"github.com/gotd/td/bin"
)
// ErrMalformed identifies invalid or truncated TL wire data. Callers may use
// errors.Is to distinguish it from an otherwise well-formed request which was
// rejected by a walker resource limit.
var ErrMalformed = errors.New("layerwire: malformed TL")
// ErrResourceLimit identifies structurally valid-looking TL input which would
// exceed a walker resource budget.
var ErrResourceLimit = errors.New("layerwire: resource limit")
const (
defaultMaxVectorElements = 4096
defaultMaxWalkDepth = 32
defaultMaxWalkUnits = 131072 // constructors + declared vector elements
defaultMaxFieldBytes = 16 << 20
defaultMaxTotalBytes = 32 << 20
)
// A very small number of API methods have a documented limit above the
// package-wide default. Keeping overrides keyed by constructor and field makes
// every exception explicit and prevents a large vector in an unrelated method
// from inheriting the larger allowance.
type vectorLimitKey struct {
owner string
field string
}
var vectorElementLimitOverrides = map[vectorLimitKey]int{
{owner: "contacts.editCloseFriends", field: "id"}: 5000,
{owner: "contacts.setBlocked", field: "id"}: 5000,
}
type walkLimits struct {
maxVectorElements int
maxDepth int
maxUnits uint64
maxFieldBytes uint64
maxTotalBytes uint64
}
var defaultWalkLimits = walkLimits{
maxVectorElements: defaultMaxVectorElements,
maxDepth: defaultMaxWalkDepth,
maxUnits: defaultMaxWalkUnits,
maxFieldBytes: defaultMaxFieldBytes,
maxTotalBytes: defaultMaxTotalBytes,
}
// walkState is deliberately request-scoped. Every branch of one transform
// shares it, so splitting a large value across nested constructors or vectors
// cannot reset the aggregate budgets.
type walkState struct {
limits walkLimits
units uint64
bytes uint64
}
func newWalkState() *walkState {
return &walkState{limits: defaultWalkLimits}
}
func malformedf(format string, args ...any) error {
return fmt.Errorf("%w: %s", ErrMalformed, fmt.Sprintf(format, args...))
}
func limitf(format string, args ...any) error {
return fmt.Errorf("%w: %s", ErrResourceLimit, fmt.Sprintf(format, args...))
}
// classifyWalkError makes all public walker/transform failures classifiable,
// including errors returned by the low-level gotd bin decoder.
func classifyWalkError(err error) error {
if err == nil || errors.Is(err, ErrMalformed) || errors.Is(err, ErrResourceLimit) {
return err
}
return fmt.Errorf("%w: %v", ErrMalformed, err)
}
func (s *walkState) enter(depth int, what string) error {
if depth <= 0 || depth > s.limits.maxDepth {
return limitf("%s nesting depth %d exceeds limit %d", what, depth, s.limits.maxDepth)
}
return s.addUnits(1, what)
}
func (s *walkState) addUnits(n int, what string) error {
if n < 0 {
return malformedf("negative %s count %d", what, n)
}
u := uint64(n)
// Subtraction form avoids overflow even if limits are changed later.
if s.units > s.limits.maxUnits || u > s.limits.maxUnits-s.units {
return limitf("constructor/vector element budget exceeds %d at %s", s.limits.maxUnits, what)
}
s.units += u
return nil
}
func (s *walkState) addBytes(n uint64, what string) error {
if n > s.limits.maxFieldBytes {
return limitf("%s payload length %d exceeds per-field limit %d", what, n, s.limits.maxFieldBytes)
}
if s.bytes > s.limits.maxTotalBytes || n > s.limits.maxTotalBytes-s.bytes {
return limitf("string/bytes payload budget exceeds %d at %s", s.limits.maxTotalBytes, what)
}
s.bytes += n
return nil
}
func (s *walkState) vectorLimit(owner *ctorLayout, f *fieldLayout) int {
if owner != nil && f != nil {
if n := vectorElementLimitOverrides[vectorLimitKey{owner: owner.name, field: f.name}]; n > 0 {
return n
}
}
return s.limits.maxVectorElements
}
const maxConstructorFlagWords = 8
type constructorFlagWord struct {
name string
value uint32
}
// ValidateCanonicalRequest performs a complete, allocation-free structural
// preflight of one canonical Layer 227 method request. It is intended for the
// router seam immediately before typed dispatch. A successful result means the
// walker consumed exactly one known function constructor and all of its body.
func ValidateCanonicalRequest(body []byte) error {
b := &bin.Buffer{Buf: body}
id, err := b.PeekID()
if err != nil {
return classifyWalkError(err)
}
cl := canonical.byCRC[id]
if cl == nil {
return malformedf("unknown canonical request constructor %#08x", id)
}
if !cl.isFunc {
return malformedf("constructor %s (%#08x) is not a method", cl.name, id)
}
return validateRequestLayout(canonical, cl, body)
}
func validateRequestLayout(m *schemaModel, cl *ctorLayout, body []byte) error {
b := &bin.Buffer{Buf: body}
s := newWalkState()
if err := s.skipObject(m, b, 1); err != nil {
return classifyWalkError(err)
}
if b.Len() != 0 {
return malformedf("%d trailing bytes after canonical request %s", b.Len(), cl.name)
}
return nil
}
// skipObject advances b past one boxed object (CRC + body), resolving the
// constructor from the canonical schema.
// constructor from m. This compatibility wrapper creates a fresh budget; all
// production transforms call the stateful variant directly.
func (m *schemaModel) skipObject(b *bin.Buffer) error {
return classifyWalkError(newWalkState().skipObject(m, b, 1))
}
func (s *walkState) skipObject(m *schemaModel, b *bin.Buffer, depth int) error {
if err := s.enter(depth, "constructor"); err != nil {
return err
}
id, err := b.PeekID()
if err != nil {
return err
}
cl, ok := m.byCRC[id]
if !ok {
return fmt.Errorf("layerwire: unknown constructor %#08x", id)
return malformedf("unknown constructor %#08x", id)
}
if err := b.ConsumeID(id); err != nil {
return err
}
return m.skipCtorBody(b, cl)
return s.skipCtorBody(m, b, cl, depth)
}
// skipCtorBody advances b past a constructor body (no leading CRC), evaluating
// flag integers so conditional fields are read iff present.
func (m *schemaModel) skipCtorBody(b *bin.Buffer, cl *ctorLayout) error {
var flags map[string]uint32
// flag integers so conditional fields are read iff present. The constructor's
// unit and depth have already been charged by the caller.
func (s *walkState) skipCtorBody(m *schemaModel, b *bin.Buffer, cl *ctorLayout, depth int) error {
// Layer 227 constructors currently use at most flags + flags2. Keep generous fixed stack
// storage so the allocation-free preflight remains allocation-free on the hottest flagged
// methods; the explicit bound also prevents a future malformed/generated layout from turning
// every request into an attacker-amplified map allocation.
var flags [maxConstructorFlagWords]constructorFlagWord
flagCount := 0
for i := range cl.fields {
f := &cl.fields[i]
if f.isFlags {
@ -34,59 +207,80 @@ func (m *schemaModel) skipCtorBody(b *bin.Buffer, cl *ctorLayout) error {
if err != nil {
return fmt.Errorf("%s.%s: %w", cl.name, f.name, err)
}
if flags == nil {
flags = make(map[string]uint32, 2)
if flagCount >= len(flags) {
return limitf("constructor %s has more than %d flags words", cl.name, len(flags))
}
flags[f.name] = v
flags[flagCount] = constructorFlagWord{name: f.name, value: v}
flagCount++
continue
}
if f.conditional() && flags[f.flagName]&(1<<uint(f.flagBit)) == 0 {
continue
if f.conditional() {
var (
flagValue uint32
found bool
)
for j := 0; j < flagCount; j++ {
if flags[j].name == f.flagName {
flagValue = flags[j].value
found = true
break
}
}
if !found {
return malformedf("constructor %s conditional field %s references missing flags word %s", cl.name, f.name, f.flagName)
}
if flagValue&(1<<uint(f.flagBit)) == 0 {
continue
}
}
if err := m.skipValue(b, f); err != nil {
if err := s.skipValue(m, b, f, cl, depth); err != nil {
return fmt.Errorf("%s.%s: %w", cl.name, f.name, err)
}
}
return nil
}
// skipValue advances b past one (already known-present) field value.
func (m *schemaModel) skipValue(b *bin.Buffer, f *fieldLayout) error {
// skipValue advances b past one already-known-present field value.
func (s *walkState) skipValue(m *schemaModel, b *bin.Buffer, f *fieldLayout, owner *ctorLayout, depth int) error {
switch f.kind {
case kindInt:
_, err := b.Int()
return err
case kindLong:
_, err := b.Long()
return err
case kindDouble:
_, err := b.Double()
return err
return skipFixed(b, 4)
case kindLong, kindDouble:
return skipFixed(b, 8)
case kindInt128:
_, err := b.Int128()
return err
return skipFixed(b, 16)
case kindInt256:
_, err := b.Int256()
return err
return skipFixed(b, 32)
case kindBytes:
_, err := b.Bytes()
return err
return s.skipTLBytes(b, "bytes")
case kindString:
_, err := b.String()
return err
return s.skipTLBytes(b, "string")
case kindBool:
_, err := b.Bool()
return err
if err := s.addUnits(1, "Bool constructor"); err != nil {
return err
}
id, err := b.Uint32()
if err != nil {
return err
}
if id != bin.TypeTrue && id != bin.TypeFalse {
return malformedf("invalid Bool constructor %#08x", id)
}
return nil
case kindTrue:
return nil
case kindVector, kindVectorBare:
vectorDepth := depth + 1
if vectorDepth <= 0 || vectorDepth > s.limits.maxDepth {
return limitf("vector nesting depth %d exceeds limit %d", vectorDepth, s.limits.maxDepth)
}
if f.kind == kindVector {
id, err := b.Uint32()
if err != nil {
return err
}
if id != vectorTypeID {
return fmt.Errorf("expected vector id, got %#08x", id)
return malformedf("expected vector id, got %#08x", id)
}
}
n, err := b.Int()
@ -94,23 +288,146 @@ func (m *schemaModel) skipValue(b *bin.Buffer, f *fieldLayout) error {
return err
}
if n < 0 {
return fmt.Errorf("negative vector length %d", n)
return malformedf("negative vector length %d", n)
}
if max := s.vectorLimit(owner, f); n > max {
return limitf("vector %s.%s length %d exceeds limit %d", ownerName(owner), fieldName(f), n, max)
}
if err := s.addUnits(n, "vector "+ownerName(owner)+"."+fieldName(f)); err != nil {
return err
}
if width, ok := fixedWireWidth(f.elem); ok {
total, ok := checkedMulInt(n, width)
if !ok {
return malformedf("vector byte length overflow: %d * %d", n, width)
}
return skipFixed(b, total)
}
for i := 0; i < n; i++ {
if err := m.skipValue(b, f.elem); err != nil {
return err
if err := s.skipValue(m, b, f.elem, nil, vectorDepth); err != nil {
return fmt.Errorf("vector element %d: %w", i, err)
}
}
return nil
case kindObject:
return m.skipObject(b)
return s.skipObject(m, b, depth+1)
case kindBareObject:
bareDepth := depth + 1
if err := s.enter(bareDepth, "bare constructor"); err != nil {
return err
}
cl, ok := m.bareByT[f.typeName]
if !ok {
return fmt.Errorf("unknown bare type %q", f.typeName)
return malformedf("unknown bare type %q", f.typeName)
}
return m.skipCtorBody(b, cl)
return s.skipCtorBody(m, b, cl, bareDepth)
default:
return fmt.Errorf("bad wire kind %d", f.kind)
return malformedf("bad wire kind %d", f.kind)
}
}
// skipTLBytes parses TL's 1/4-byte length prefix directly and advances the
// input slice. Unlike bin.Buffer.Bytes it never copies payload data.
func (s *walkState) skipTLBytes(b *bin.Buffer, what string) error {
if len(b.Buf) == 0 {
return io.ErrUnexpectedEOF
}
var header, payload uint64
switch b.Buf[0] {
case 254:
if len(b.Buf) < 4 {
return io.ErrUnexpectedEOF
}
header = 4
payload = uint64(b.Buf[1]) | uint64(b.Buf[2])<<8 | uint64(b.Buf[3])<<16
case 255:
return malformedf("invalid %s length prefix 255", what)
default:
header = 1
payload = uint64(b.Buf[0])
}
if err := s.addBytes(payload, what); err != nil {
return err
}
encoded, ok := checkedAddUint64(header, payload)
if !ok {
return malformedf("%s encoded length overflow", what)
}
withPadding, ok := checkedAddUint64(encoded, 3)
if !ok {
return malformedf("%s padded length overflow", what)
}
padded := withPadding &^ uint64(3)
if padded > uint64(math.MaxInt) {
return malformedf("%s padded length %d overflows int", what, padded)
}
if uint64(len(b.Buf)) < padded {
return io.ErrUnexpectedEOF
}
b.Buf = b.Buf[int(padded):]
return nil
}
func skipFixed(b *bin.Buffer, n int) error {
if n < 0 {
return malformedf("negative fixed-width skip %d", n)
}
if len(b.Buf) < n {
return io.ErrUnexpectedEOF
}
b.Buf = b.Buf[n:]
return nil
}
func fixedWireWidth(f *fieldLayout) (int, bool) {
if f == nil {
return 0, false
}
switch f.kind {
case kindInt:
return 4, true
case kindLong, kindDouble:
return 8, true
case kindInt128:
return 16, true
case kindInt256:
return 32, true
case kindTrue:
return 0, true
default:
// Bool deliberately stays on the element loop so constructor ids are
// validated and charged to the aggregate constructor budget.
return 0, false
}
}
func checkedMulInt(a, b int) (int, bool) {
if a < 0 || b < 0 {
return 0, false
}
if a != 0 && b > math.MaxInt/a {
return 0, false
}
return a * b, true
}
func checkedAddUint64(a, b uint64) (uint64, bool) {
if b > math.MaxUint64-a {
return 0, false
}
return a + b, true
}
func ownerName(cl *ctorLayout) string {
if cl == nil || cl.name == "" {
return "<nested>"
}
return cl.name
}
func fieldName(f *fieldLayout) string {
if f == nil || f.name == "" {
return "<element>"
}
return f.name
}

263
internal/compat/layerwire/walk_limits_test.go Normal file
View file

@ -0,0 +1,263 @@
package layerwire
import (
"errors"
"math"
"testing"
"github.com/gotd/td/bin"
"github.com/gotd/td/tg"
)
func TestValidateCanonicalRequestFlaggedHotPathAllocatesNothing(t *testing.T) {
var body bin.Buffer
req := &tg.MessagesSendMessageRequest{
Peer: &tg.InputPeerSelf{},
Message: "hello",
RandomID: 7,
}
if err := req.Encode(&body); err != nil {
t.Fatalf("encode request: %v", err)
}
if err := ValidateCanonicalRequest(body.Buf); err != nil {
t.Fatalf("validate request: %v", err)
}
if allocs := testing.AllocsPerRun(1000, func() {
if err := ValidateCanonicalRequest(body.Buf); err != nil {
panic(err)
}
}); allocs != 0 {
t.Fatalf("canonical request preflight allocations = %.2f, want 0", allocs)
}
}
func TestValidateCanonicalRequestVectorLimits(t *testing.T) {
editCloseFriends := canonical.byName["contacts.editCloseFriends"]
if editCloseFriends == nil {
t.Fatal("contacts.editCloseFriends missing from canonical schema")
}
t.Run("explicit_5000_override", func(t *testing.T) {
var body bin.Buffer
body.PutID(editCloseFriends.crc)
body.PutVectorHeader(5000)
for i := 0; i < 5000; i++ {
body.PutLong(int64(i))
}
if err := ValidateCanonicalRequest(body.Buf); err != nil {
t.Fatalf("validate legal 5000-element close-friends request: %v", err)
}
})
t.Run("override_stops_at_5000", func(t *testing.T) {
var body bin.Buffer
body.PutID(editCloseFriends.crc)
body.PutVectorHeader(5001)
err := ValidateCanonicalRequest(body.Buf)
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want ErrResourceLimit", err)
}
})
t.Run("default_4096", func(t *testing.T) {
getMessages := canonical.byName["messages.getMessages"]
var body bin.Buffer
body.PutID(getMessages.crc)
body.PutVectorHeader(defaultMaxVectorElements + 1)
err := ValidateCanonicalRequest(body.Buf)
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want ErrResourceLimit", err)
}
})
t.Run("max_int32_count_rejected_before_iteration", func(t *testing.T) {
var body bin.Buffer
body.PutID(editCloseFriends.crc)
body.PutID(vectorTypeID)
body.PutInt32(math.MaxInt32)
err := ValidateCanonicalRequest(body.Buf)
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want ErrResourceLimit", err)
}
})
}
func TestValidateCanonicalRequestDepthLimit(t *testing.T) {
invoke := canonical.byName["invokeWithoutUpdates"]
leaf := canonical.byName["help.getConfig"]
if invoke == nil || leaf == nil {
t.Fatal("generic wrapper methods missing from canonical schema")
}
request := func(wrappers int) []byte {
var body bin.Buffer
for i := 0; i < wrappers; i++ {
body.PutID(invoke.crc)
}
body.PutID(leaf.crc)
return body.Buf
}
if err := ValidateCanonicalRequest(request(defaultMaxWalkDepth - 1)); err != nil {
t.Fatalf("depth exactly %d rejected: %v", defaultMaxWalkDepth, err)
}
err := ValidateCanonicalRequest(request(defaultMaxWalkDepth))
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("depth %d error = %v, want ErrResourceLimit", defaultMaxWalkDepth+1, err)
}
}
func TestTLBytesSkipIsZeroCopyAndBounded(t *testing.T) {
var encoded bin.Buffer
encoded.PutBytes([]byte("payload"))
fieldLen := len(encoded.Buf)
raw := append(encoded.Copy(), 0xaa, 0xbb, 0xcc, 0xdd)
b := &bin.Buffer{Buf: raw}
walk := newWalkState()
if err := walk.skipTLBytes(b, "bytes"); err != nil {
t.Fatalf("skip bytes: %v", err)
}
if len(b.Buf) != 4 || &b.Buf[0] != &raw[fieldLen] {
t.Fatalf("walker did not retain the original backing buffer")
}
t.Run("per_field_budget", func(t *testing.T) {
limited := newWalkState()
limited.limits.maxFieldBytes = 3
probe := &bin.Buffer{Buf: encoded.Copy()}
err := limited.skipTLBytes(probe, "bytes")
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want ErrResourceLimit", err)
}
})
t.Run("aggregate_budget", func(t *testing.T) {
limited := newWalkState()
limited.limits.maxTotalBytes = 10
first := &bin.Buffer{Buf: encoded.Copy()}
if err := limited.skipTLBytes(first, "bytes"); err != nil {
t.Fatalf("first field: %v", err)
}
second := &bin.Buffer{Buf: encoded.Copy()}
err := limited.skipTLBytes(second, "bytes")
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("second error = %v, want ErrResourceLimit", err)
}
})
t.Run("truncated_payload_is_malformed", func(t *testing.T) {
importAuth := canonical.byName["auth.importAuthorization"]
var body bin.Buffer
body.PutID(importAuth.crc)
body.PutLong(1)
body.Put([]byte{5, 'a', 'b'}) // declares five bytes, lacks payload/padding
err := ValidateCanonicalRequest(body.Buf)
if !errors.Is(err, ErrMalformed) || errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want only ErrMalformed", err)
}
})
}
func TestInboundTransformsShareWalkerBudgets(t *testing.T) {
t.Run("canonical_alias", func(t *testing.T) {
var body bin.Buffer
body.PutID(0x41d41ade) // DrKLO messages.forwardMessages alias
body.PutUint32(0)
body.PutID(canonical.byName["inputPeerEmpty"].crc)
body.PutID(vectorTypeID)
body.PutInt32(math.MaxInt32)
_, ok, err := UpgradeInbound(0x41d41ade, &body)
if !ok || !errors.Is(err, ErrResourceLimit) {
t.Fatalf("ok=%v error=%v, want matched ErrResourceLimit", ok, err)
}
})
t.Run("drift_body_transform", func(t *testing.T) {
var body bin.Buffer
body.PutID(0x2e1ee318) // DrKLO langpack.getStrings body transform
body.PutString("en")
body.PutID(vectorTypeID)
body.PutInt32(math.MaxInt32)
_, ok, err := UpgradeInbound(0x2e1ee318, &body)
if !ok || !errors.Is(err, ErrResourceLimit) {
t.Fatalf("ok=%v error=%v, want matched ErrResourceLimit", ok, err)
}
})
t.Run("outbound_structural_transform", func(t *testing.T) {
poll := canonical.byName["pollAnswerVoters"]
var body bin.Buffer
body.PutID(poll.crc)
body.PutUint32(1 << 2)
body.PutBytes(nil)
body.PutInt(1)
body.PutID(vectorTypeID)
body.PutInt32(math.MaxInt32)
_, err := Transcode(body.Buf, CanonicalLayer-1)
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want ErrResourceLimit", err)
}
})
}
func TestWalkerArithmeticAndMalformedClassification(t *testing.T) {
if defaultMaxWalkUnits != 131072 {
t.Fatalf("default constructor/vector budget = %d, want 131072", defaultMaxWalkUnits)
}
if _, ok := checkedMulInt(math.MaxInt, 2); ok {
t.Fatal("checkedMulInt accepted overflow")
}
if _, ok := checkedAddUint64(math.MaxUint64, 1); ok {
t.Fatal("checkedAddUint64 accepted overflow")
}
t.Run("aggregate_constructor_and_vector_units", func(t *testing.T) {
editCloseFriends := canonical.byName["contacts.editCloseFriends"]
var body bin.Buffer
body.PutID(editCloseFriends.crc)
body.PutVectorHeader(4)
for i := 0; i < 4; i++ {
body.PutLong(int64(i))
}
walk := newWalkState()
walk.limits.maxUnits = 4 // top constructor + four elements needs five
probe := &bin.Buffer{Buf: body.Buf}
err := walk.skipObject(canonical, probe, 1)
if !errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want ErrResourceLimit", err)
}
})
tests := []struct {
name string
body []byte
}{
{name: "empty"},
{name: "unknown_constructor", body: []byte{1, 2, 3, 4}},
{name: "trailing_bytes", body: append(methodIDBytes(canonical.byName["help.getConfig"].crc), 0, 0, 0, 0)},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := ValidateCanonicalRequest(tt.body)
if !errors.Is(err, ErrMalformed) || errors.Is(err, ErrResourceLimit) {
t.Fatalf("error = %v, want only ErrMalformed", err)
}
})
}
}
func methodIDBytes(id uint32) []byte {
var b bin.Buffer
b.PutID(id)
return b.Buf
}
func FuzzValidateCanonicalRequest(f *testing.F) {
f.Add(methodIDBytes(canonical.byName["help.getConfig"].crc))
f.Add([]byte{})
f.Add([]byte{1, 2, 3, 4})
f.Fuzz(func(t *testing.T, body []byte) {
err := ValidateCanonicalRequest(body)
if err != nil && !errors.Is(err, ErrMalformed) && !errors.Is(err, ErrResourceLimit) {
t.Fatalf("unclassified walker error: %v", err)
}
})
}

104
internal/config/config.go
View file

@ -29,6 +29,28 @@ type Config struct {
RSAKeyPath string
// DC 是本 server 的 DC ID。
DC int
// MTProtoMaxConnections / PerIP 覆盖 raw Accept、codec sniff、握手到认证 session
// 的完整物理连接生命周期；负数关闭对应 admission 上限。
MTProtoMaxConnections int
MTProtoMaxConnectionsPerIP int
// MTProtoMaxConcurrentHandshakes 限制昂贵 RSA/DH exchange 并发；负数关闭。
MTProtoMaxConcurrentHandshakes int
// MTProto RPC 使用 Server 共享公平调度器；per-connection 与 global 预算共同限制
// goroutine、排队任务和 request body 内存。
MTProtoRPCMaxInflight int
MTProtoRPCQueueSize int
MTProtoRPCTimeout time.Duration
MTProtoRPCGlobalWorkers int
MTProtoRPCGlobalMaxTasks int
MTProtoRPCGlobalMaxBytes int64
// MTProtoInboundFrameGlobalMaxBytes 是 transport wire + 最大解密 plaintext 的
// 进程级在途预算；frame 长度读出后、payload 分配前预留。
MTProtoInboundFrameGlobalMaxBytes int64
// MTProto outbound mailbox 按连接有界；resend pending body 另受 Server 全局预算约束。
MTProtoOutboundQueueSize int
MTProtoOutboundControlQueueSize int
MTProtoOutboundTrackedGlobalMaxBytes int64
MTProtoOutboundWriteGlobalMaxBytes int64
// DebugAddr 是 net/http/pprof 调试端点监听地址（CPU/heap/goroutine/mutex/block 剖析）。
// telesrv 是宿主进程、不在 docker 内，docker stats 看不到它，性能定位主要靠此端点。
@ -76,6 +98,12 @@ type Config struct {
// AuthCodeMaxAttempts 是同一 phone_code_hash / email verification code 的最大错误次数。
// 达到上限后验证码立即失效，用户必须重发。
AuthCodeMaxAttempts int
// AuthCodePhoneRateLimit / AuthCodeAuthKeyRateLimit 对未授权验证码签发按规范化手机号摘要
// 与连接实际 raw auth_key 分别限流。两个维度共用 AuthCodeRateWindow；<=0 关闭对应维度。
// 手机号只以 SHA-256 摘要进入限流 key，禁止把原文写入 Redis key 或日志。
AuthCodePhoneRateLimit int
AuthCodeAuthKeyRateLimit int
AuthCodeRateWindow time.Duration
// LoginEmailEnable 启用手机号登录流程中的邮箱验证码投递。
LoginEmailEnable bool
// LoginEmailRequireSetup 为 true 时，没有登录邮箱的账号/新手机号会要求先设置邮箱。
@ -136,6 +164,9 @@ type Config struct {
AIPrivacyLogContent bool
// TempKeyResolveCacheMaxEntries 是 Router temp→perm 解析缓存容量。
TempKeyResolveCacheMaxEntries int
// TempKeyResolveCacheTTL 是 temp→perm 绑定的进程内复核周期。绑定/revoke 有精确
// 失效，TTL 作为跨进程或异常路径兜底；默认 30m 避免大连接数下每 5s 全量打 PG。
TempKeyResolveCacheTTL time.Duration
// ChannelRowCacheMaxEntries 是「共享频道行」进程内缓存容量(channelID→domain.Channel)。
// 由 channels 表 LISTEN/NOTIFY 触发器实时失效(强一致、零 TTL)。<=0 禁用缓存与监听。
@ -153,8 +184,8 @@ type Config struct {
// ChannelBoostCacheTTL 是 boost 读投影在未收到写侧通知时的最大陈旧窗口。
ChannelBoostCacheTTL time.Duration
// OutboxWorkers 是并发 claim 的 outbox worker 数。默认 1，保证同一用户 pts update
// 在线投递顺序与持久化顺序一致；后续需要吞吐时应改成按 target_user_id 分片的串行 worker。
// OutboxWorkers 是并发 outbox worker 数。用户先稳定哈希到固定 logical shard，
// 每个 shard 只归一个 worker，故提高 worker 数不会破坏同一用户 pts 顺序。
OutboxWorkers int
// OutboxBatch 是 transactional outbox worker 每次 claim 的最大条数。
// 调大提升吞吐、增大单批 PG/推送压力；调小降低延迟抖动。配套压测见 docs/message-module.md。
@ -164,6 +195,13 @@ type Config struct {
// OutboxLeaseTimeout 是 'dispatching' 行被判定为租约过期、允许其它 worker 重新 claim 的时长。
// 取值需大于单批投递耗时，否则会重复推送；过大则 worker 崩溃后积压恢复变慢。
OutboxLeaseTimeout time.Duration
// OutboxPoisonRetention 是 terminal failed outbox head 的隔离窗口。隔离期内保留
// last_error 供排障，期满只删除在线投递任务；durable user_update_events 仍保留，
// 客户端可经 updates.getDifference 恢复。
OutboxPoisonRetention time.Duration
// OutboxPoisonCleanupInterval 独立于大表 retention 周期清理 terminal failed head，
// 避免一条确定性坏事件长期冻结同账号更高 pts 的在线投递 lane。
OutboxPoisonCleanupInterval time.Duration
// OutboundPushTimeout 是 best-effort updates 推送等待 outbound 队列接受的最长时间。
OutboundPushTimeout time.Duration
// SendRateLimit 是账号级发送窗口内允许的消息条数；<=0 表示关闭发送限流。
@ -182,6 +220,9 @@ type Config struct {
// BotAPIUpdateRetention 是 bot_api_updates 投递队列的最大保留期（官方 Bot API 语义 24h）；
// 已确认的行另按固定短宽限提前回收（性能审计 H1）。
BotAPIUpdateRetention time.Duration
// OrphanAuthKeyRetention 是握手已创建、但没有 authorization/temp binding/活跃连接的
// auth key 最短保留期。过期后由有界 GC 回收；客户端收到 -404 会重建 key。
OrphanAuthKeyRetention time.Duration
// RetentionInterval 是 retention worker 的运行间隔。
RetentionInterval time.Duration
// RetentionBatch 是单次 retention 最多删除的行数。
@ -323,19 +364,33 @@ func Load() (Config, error) {
// AdvertiseIP 当前不影响 help.getConfig——getConfig 返回空 DCOptions，
// 客户端使用其写死的 static DC 地址（见 compat/tdesktop/config.go）。
// 字段与默认值保留，供未来需要显式下发 DC 地址时使用。
AdvertiseIP: envOr("TELESRV_ADVERTISE_IP", "127.0.0.1"),
RSAKeyPath: envOr("TELESRV_RSA_KEY", "data/server_rsa.pem"),
DC: envIntOr("TELESRV_DC", 2),
DebugAddr: envAllowEmptyOr("TELESRV_DEBUG_ADDR", "127.0.0.1:6060"),
BotAPIAddr: envAllowEmptyOr("TELESRV_BOT_API_ADDR", ""),
AdminAPIAddr: envAllowEmptyOr("TELESRV_ADMIN_API_ADDR", ""),
AdminAPIToken: envOr("TELESRV_ADMIN_API_TOKEN", ""),
PublicBaseURL: publicBaseURL,
PublicLinkWebAddr: envAllowEmptyOr("TELESRV_PUBLIC_LINK_WEB_ADDR", ""),
AdminUIAddr: envOr("TELESRV_ADMIN_UI_ADDR", "127.0.0.1:2600"),
AdminUIPassword: envOr("TELESRV_ADMIN_UI_PASSWORD", ""),
AdminUIToken: envOr("TELESRV_ADMIN_UI_TOKEN", ""),
AdminSessionKey: envOr("TELESRV_ADMIN_SESSION_KEY", ""),
AdvertiseIP: envOr("TELESRV_ADVERTISE_IP", "127.0.0.1"),
RSAKeyPath: envOr("TELESRV_RSA_KEY", "data/server_rsa.pem"),
DC: envIntOr("TELESRV_DC", 2),
MTProtoMaxConnections: envIntOr("TELESRV_MTPROTO_MAX_CONNECTIONS", 200000),
MTProtoMaxConnectionsPerIP: envIntOr("TELESRV_MTPROTO_MAX_CONNECTIONS_PER_IP", 4096),
MTProtoMaxConcurrentHandshakes: envIntOr("TELESRV_MTPROTO_MAX_CONCURRENT_HANDSHAKES", 256),
MTProtoRPCMaxInflight: envIntOr("TELESRV_MTPROTO_RPC_MAX_INFLIGHT", 32),
MTProtoRPCQueueSize: envIntOr("TELESRV_MTPROTO_RPC_QUEUE_SIZE", 64),
MTProtoRPCTimeout: envDurationOr("TELESRV_MTPROTO_RPC_TIMEOUT", 30*time.Second),
MTProtoRPCGlobalWorkers: envIntOr("TELESRV_MTPROTO_RPC_GLOBAL_WORKERS", 256),
MTProtoRPCGlobalMaxTasks: envIntOr("TELESRV_MTPROTO_RPC_GLOBAL_MAX_TASKS", 8192),
MTProtoRPCGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_RPC_GLOBAL_MAX_BYTES", 512<<20),
MTProtoInboundFrameGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_INBOUND_FRAME_GLOBAL_MAX_BYTES", 512<<20),
MTProtoOutboundQueueSize: envIntOr("TELESRV_MTPROTO_OUTBOUND_QUEUE_SIZE", 128),
MTProtoOutboundControlQueueSize: envIntOr("TELESRV_MTPROTO_OUTBOUND_CONTROL_QUEUE_SIZE", 32),
MTProtoOutboundTrackedGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_OUTBOUND_TRACKED_GLOBAL_MAX_BYTES", 512<<20),
MTProtoOutboundWriteGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_OUTBOUND_WRITE_GLOBAL_MAX_BYTES", 512<<20),
DebugAddr: envAllowEmptyOr("TELESRV_DEBUG_ADDR", "127.0.0.1:6060"),
BotAPIAddr: envAllowEmptyOr("TELESRV_BOT_API_ADDR", ""),
AdminAPIAddr: envAllowEmptyOr("TELESRV_ADMIN_API_ADDR", ""),
AdminAPIToken: envOr("TELESRV_ADMIN_API_TOKEN", ""),
PublicBaseURL: publicBaseURL,
PublicLinkWebAddr: envAllowEmptyOr("TELESRV_PUBLIC_LINK_WEB_ADDR", ""),
AdminUIAddr: envOr("TELESRV_ADMIN_UI_ADDR", "127.0.0.1:2600"),
AdminUIPassword: envOr("TELESRV_ADMIN_UI_PASSWORD", ""),
AdminUIToken: envOr("TELESRV_ADMIN_UI_TOKEN", ""),
AdminSessionKey: envOr("TELESRV_ADMIN_SESSION_KEY", ""),
// 用 127.0.0.1 而非 localhost：localhost 在 Windows 上会先解析到 IPv6 ::1，而 Docker
// Desktop 的端口转发只在 IPv4 监听，IPv6 连接要等 ~1s 超时才回退 IPv4（实测 localhost
@ -351,6 +406,9 @@ func Load() (Config, error) {
DevAuthCode: envOr("TELESRV_DEV_AUTH_CODE", "12345"),
AuthCodeTTL: envDurationOr("TELESRV_AUTH_CODE_TTL", 5*time.Minute),
AuthCodeMaxAttempts: envIntOr("TELESRV_AUTH_CODE_MAX_ATTEMPTS", 5),
AuthCodePhoneRateLimit: envIntOr("TELESRV_AUTH_CODE_PHONE_RATE_LIMIT", 5),
AuthCodeAuthKeyRateLimit: envIntOr("TELESRV_AUTH_CODE_AUTH_KEY_RATE_LIMIT", 20),
AuthCodeRateWindow: envDurationOr("TELESRV_AUTH_CODE_RATE_WINDOW", 10*time.Minute),
LoginEmailEnable: envBoolOr("TELESRV_LOGIN_EMAIL_ENABLE", false),
LoginEmailRequireSetup: envBoolOr("TELESRV_LOGIN_EMAIL_REQUIRE_SETUP", false),
LoginEmailCodeLength: envIntOr("TELESRV_LOGIN_EMAIL_CODE_LENGTH", 6),
@ -381,17 +439,22 @@ func Load() (Config, error) {
AIRateLimit: envIntOr("TELESRV_AI_RATE_LIMIT", 20),
AIRateWindow: envDurationOr("TELESRV_AI_RATE_WINDOW", time.Minute),
AIPrivacyLogContent: envBoolOr("TELESRV_AI_LOG_CONTENT", false),
TempKeyResolveCacheMaxEntries: envIntOr("TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES", 4096),
TempKeyResolveCacheMaxEntries: envIntOr("TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES", 262144),
TempKeyResolveCacheTTL: envDurationOr("TELESRV_TEMP_KEY_CACHE_TTL", 30*time.Minute),
ChannelRowCacheMaxEntries: envIntOr("TELESRV_CHANNEL_ROW_CACHE_MAX", 50000),
ChannelMemberCacheMaxEntries: envIntOr("TELESRV_CHANNEL_MEMBER_CACHE_MAX", 100000),
ChannelDialogCacheMaxEntries: envIntOr("TELESRV_CHANNEL_DIALOG_CACHE_MAX", 100000),
ChannelBoostCacheMaxEntries: envIntOr("TELESRV_CHANNEL_BOOST_CACHE_MAX", 100000),
ChannelBoostCacheTTL: envDurationOr("TELESRV_CHANNEL_BOOST_CACHE_TTL", 10*time.Second),
OutboxWorkers: envIntOr("TELESRV_OUTBOX_WORKERS", 1),
OutboxBatch: envIntOr("TELESRV_OUTBOX_BATCH", 100),
OutboxInterval: envDurationOr("TELESRV_OUTBOX_INTERVAL", 200*time.Millisecond),
OutboxLeaseTimeout: envDurationOr("TELESRV_OUTBOX_LEASE_TIMEOUT", 30*time.Second),
OutboxWorkers: envIntOr("TELESRV_OUTBOX_WORKERS", 4),
OutboxBatch: envIntOr("TELESRV_OUTBOX_BATCH", 100),
OutboxInterval: envDurationOr("TELESRV_OUTBOX_INTERVAL", 200*time.Millisecond),
OutboxLeaseTimeout: envDurationOr("TELESRV_OUTBOX_LEASE_TIMEOUT", 30*time.Second),
OutboxPoisonRetention: envDurationOr("TELESRV_OUTBOX_POISON_RETENTION", time.Minute),
OutboxPoisonCleanupInterval: envDurationOr(
"TELESRV_OUTBOX_POISON_CLEANUP_INTERVAL", 15*time.Second,
),
OutboundPushTimeout: envDurationOr("TELESRV_OUTBOUND_PUSH_TIMEOUT", 200*time.Millisecond),
SendRateLimit: envIntOr("TELESRV_SEND_RATE_LIMIT", 30),
SendRateWindow: envDurationOr("TELESRV_SEND_RATE_WINDOW", time.Minute),
@ -400,6 +463,7 @@ func Load() (Config, error) {
ChannelNudgeMaxTargets: envIntOr("TELESRV_CHANNEL_NUDGE_MAX_TARGETS", 0),
UpdateEventRetention: envDurationOr("TELESRV_UPDATE_EVENT_RETENTION", 168*time.Hour),
BotAPIUpdateRetention: envDurationOr("TELESRV_BOT_API_UPDATE_RETENTION", 24*time.Hour),
OrphanAuthKeyRetention: envDurationOr("TELESRV_ORPHAN_AUTH_KEY_RETENTION", 24*time.Hour),
RetentionInterval: envDurationOr("TELESRV_RETENTION_INTERVAL", time.Hour),
RetentionBatch: envIntOr("TELESRV_RETENTION_BATCH", 10000),
UploadPartTTL: envDurationOr("TELESRV_UPLOAD_PART_TTL", 24*time.Hour),

73
internal/config/config_test.go
View file

@ -37,6 +37,62 @@ func TestLoadUsesExplicitAdvertiseIP(t *testing.T) {
}
}
func TestLoadMTProtoAdmissionAndRPCBudgets(t *testing.T) {
disableDefaultConfigFile(t)
t.Setenv("TELESRV_MTPROTO_MAX_CONNECTIONS", "12345")
t.Setenv("TELESRV_MTPROTO_MAX_CONNECTIONS_PER_IP", "234")
t.Setenv("TELESRV_MTPROTO_MAX_CONCURRENT_HANDSHAKES", "45")
t.Setenv("TELESRV_MTPROTO_RPC_MAX_INFLIGHT", "7")
t.Setenv("TELESRV_MTPROTO_RPC_QUEUE_SIZE", "19")
t.Setenv("TELESRV_MTPROTO_RPC_TIMEOUT", "9s")
t.Setenv("TELESRV_MTPROTO_RPC_GLOBAL_WORKERS", "33")
t.Setenv("TELESRV_MTPROTO_RPC_GLOBAL_MAX_TASKS", "444")
t.Setenv("TELESRV_MTPROTO_RPC_GLOBAL_MAX_BYTES", "555555")
t.Setenv("TELESRV_MTPROTO_INBOUND_FRAME_GLOBAL_MAX_BYTES", "777777")
t.Setenv("TELESRV_MTPROTO_OUTBOUND_QUEUE_SIZE", "88")
t.Setenv("TELESRV_MTPROTO_OUTBOUND_CONTROL_QUEUE_SIZE", "22")
t.Setenv("TELESRV_MTPROTO_OUTBOUND_TRACKED_GLOBAL_MAX_BYTES", "888888")
t.Setenv("TELESRV_MTPROTO_OUTBOUND_WRITE_GLOBAL_MAX_BYTES", "999999")
t.Setenv("TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES", "666")
t.Setenv("TELESRV_TEMP_KEY_CACHE_TTL", "17m")
t.Setenv("TELESRV_ORPHAN_AUTH_KEY_RETENTION", "36h")
cfg, err := Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.MTProtoMaxConnections != 12345 || cfg.MTProtoMaxConnectionsPerIP != 234 || cfg.MTProtoMaxConcurrentHandshakes != 45 {
t.Fatalf("admission config = %d/%d/%d", cfg.MTProtoMaxConnections, cfg.MTProtoMaxConnectionsPerIP, cfg.MTProtoMaxConcurrentHandshakes)
}
if cfg.MTProtoRPCMaxInflight != 7 || cfg.MTProtoRPCQueueSize != 19 || cfg.MTProtoRPCTimeout != 9*time.Second ||
cfg.MTProtoRPCGlobalWorkers != 33 || cfg.MTProtoRPCGlobalMaxTasks != 444 || cfg.MTProtoRPCGlobalMaxBytes != 555555 {
t.Fatalf("rpc budget config = %d/%d/%v/%d/%d/%d", cfg.MTProtoRPCMaxInflight, cfg.MTProtoRPCQueueSize, cfg.MTProtoRPCTimeout, cfg.MTProtoRPCGlobalWorkers, cfg.MTProtoRPCGlobalMaxTasks, cfg.MTProtoRPCGlobalMaxBytes)
}
if cfg.MTProtoInboundFrameGlobalMaxBytes != 777777 {
t.Fatalf("inbound frame budget config = %d", cfg.MTProtoInboundFrameGlobalMaxBytes)
}
if cfg.MTProtoOutboundQueueSize != 88 || cfg.MTProtoOutboundControlQueueSize != 22 || cfg.MTProtoOutboundTrackedGlobalMaxBytes != 888888 || cfg.MTProtoOutboundWriteGlobalMaxBytes != 999999 {
t.Fatalf("outbound config = %d/%d/%d/%d", cfg.MTProtoOutboundQueueSize, cfg.MTProtoOutboundControlQueueSize, cfg.MTProtoOutboundTrackedGlobalMaxBytes, cfg.MTProtoOutboundWriteGlobalMaxBytes)
}
if cfg.TempKeyResolveCacheMaxEntries != 666 || cfg.TempKeyResolveCacheTTL != 17*time.Minute || cfg.OrphanAuthKeyRetention != 36*time.Hour {
t.Fatalf("auth key resource config = %d/%v/%v", cfg.TempKeyResolveCacheMaxEntries, cfg.TempKeyResolveCacheTTL, cfg.OrphanAuthKeyRetention)
}
}
func TestLoadOutboxPoisonPolicy(t *testing.T) {
disableDefaultConfigFile(t)
t.Setenv("TELESRV_OUTBOX_POISON_RETENTION", "2m")
t.Setenv("TELESRV_OUTBOX_POISON_CLEANUP_INTERVAL", "7s")
cfg, err := Load()
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.OutboxPoisonRetention != 2*time.Minute || cfg.OutboxPoisonCleanupInterval != 7*time.Second {
t.Fatalf("outbox poison policy = %v/%v, want 2m/7s", cfg.OutboxPoisonRetention, cfg.OutboxPoisonCleanupInterval)
}
}
func TestLoadBusinessAIProvider(t *testing.T) {
disableDefaultConfigFile(t)
t.Setenv("TELESRV_BUSINESS_AI_PROVIDER", "echo")
@ -76,8 +132,11 @@ func TestLoadLoginEmailDefaultsDisabled(t *testing.T) {
if cfg.LoginEmailRequireSetup {
t.Fatal("LoginEmailRequireSetup = true, want false")
}
if cfg.AuthCodeTTL != 5*time.Minute || cfg.AuthCodeMaxAttempts != 5 || cfg.LoginEmailCodeLength != 6 {
t.Fatalf("auth/login email defaults = %v/%d/%d", cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength)
if cfg.AuthCodeTTL != 5*time.Minute || cfg.AuthCodeMaxAttempts != 5 || cfg.LoginEmailCodeLength != 6 ||
cfg.AuthCodePhoneRateLimit != 5 || cfg.AuthCodeAuthKeyRateLimit != 20 || cfg.AuthCodeRateWindow != 10*time.Minute {
t.Fatalf("auth/login email defaults = ttl=%v attempts=%d length=%d phone_limit=%d key_limit=%d window=%v",
cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength,
cfg.AuthCodePhoneRateLimit, cfg.AuthCodeAuthKeyRateLimit, cfg.AuthCodeRateWindow)
}
}
@ -87,6 +146,9 @@ func TestLoadLoginEmailSMTPConfig(t *testing.T) {
t.Setenv("TELESRV_LOGIN_EMAIL_REQUIRE_SETUP", "true")
t.Setenv("TELESRV_AUTH_CODE_TTL", "3m")
t.Setenv("TELESRV_AUTH_CODE_MAX_ATTEMPTS", "4")
t.Setenv("TELESRV_AUTH_CODE_PHONE_RATE_LIMIT", "3")
t.Setenv("TELESRV_AUTH_CODE_AUTH_KEY_RATE_LIMIT", "9")
t.Setenv("TELESRV_AUTH_CODE_RATE_WINDOW", "2m")
t.Setenv("TELESRV_LOGIN_EMAIL_CODE_LENGTH", "7")
t.Setenv("TELESRV_SMTP_HOST", "smtp.example.test")
t.Setenv("TELESRV_SMTP_PORT", "2525")
@ -103,8 +165,11 @@ func TestLoadLoginEmailSMTPConfig(t *testing.T) {
if !cfg.LoginEmailEnable || !cfg.LoginEmailRequireSetup {
t.Fatalf("login email flags = %v/%v, want true/true", cfg.LoginEmailEnable, cfg.LoginEmailRequireSetup)
}
if cfg.AuthCodeTTL != 3*time.Minute || cfg.AuthCodeMaxAttempts != 4 || cfg.LoginEmailCodeLength != 7 {
t.Fatalf("auth/login email config = %v/%d/%d", cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength)
if cfg.AuthCodeTTL != 3*time.Minute || cfg.AuthCodeMaxAttempts != 4 || cfg.LoginEmailCodeLength != 7 ||
cfg.AuthCodePhoneRateLimit != 3 || cfg.AuthCodeAuthKeyRateLimit != 9 || cfg.AuthCodeRateWindow != 2*time.Minute {
t.Fatalf("auth/login email config = ttl=%v attempts=%d length=%d phone_limit=%d key_limit=%d window=%v",
cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength,
cfg.AuthCodePhoneRateLimit, cfg.AuthCodeAuthKeyRateLimit, cfg.AuthCodeRateWindow)
}
if cfg.SMTPHost != "smtp.example.test" || cfg.SMTPPort != 2525 || cfg.SMTPUsername != "smtp-user" || cfg.SMTPPassword != "smtp-pass" || cfg.SMTPFrom != "noreply@example.test" || cfg.SMTPTLSMode != "none" || cfg.SMTPTimeout != 2*time.Second {
t.Fatalf("smtp config = %#v", cfg)

51
internal/domain/album_group.go Normal file
View file

@ -0,0 +1,51 @@
package domain
import (
"crypto/sha256"
"errors"
)
// ErrAlbumGroupReservationInvalid 表示相册分组预留缺少发送者、目标、
// random_id 或 proposed grouped_id。RPC 边界通常会更早拦截这些输入；
// domain/store 仍 fail-fast，避免坏绑定进入持久层。
var ErrAlbumGroupReservationInvalid = errors.New("album group reservation invalid")
// AlbumGroupReservationRequest 在任何相册 item 落库或上传媒体解析前，原子地把
// 一组 random_id 绑定到同一个 grouped_id。Peer 是幂等作用域的一部分：同一发送者
// 可以在不同会话中复用 random_id，而不会互相污染相册分组。
type AlbumGroupReservationRequest struct {
SenderUserID int64
Peer Peer
Items []AlbumGroupReservationItem
ProposedGroupedID int64
}
// AlbumGroupReservationItem 把 random_id 与该 item 的不可变客户端意图绑定。
// IntentHash 是在媒体解析/服务端派生字段产生前计算的 SHA-256；相同 random_id
// 若携带不同意图必须报冲突，不能借旧 album reservation 绕过发送幂等校验。
type AlbumGroupReservationItem struct {
RandomID int64
IntentHash []byte
}
// Validate 校验持久层必须依赖的最小不变量。同一批内重复 random_id 与发送幂等
// 冲突同义，必须显式失败，不能静默去重后改变客户端请求的消息条数。
func (r AlbumGroupReservationRequest) Validate() error {
if r.SenderUserID <= 0 || r.Peer.ID <= 0 || r.ProposedGroupedID == 0 || len(r.Items) == 0 {
return ErrAlbumGroupReservationInvalid
}
if r.Peer.Type != PeerTypeUser && r.Peer.Type != PeerTypeChannel {
return ErrAlbumGroupReservationInvalid
}
seen := make(map[int64]struct{}, len(r.Items))
for _, item := range r.Items {
if item.RandomID == 0 || len(item.IntentHash) != sha256.Size {
return ErrAlbumGroupReservationInvalid
}
if _, exists := seen[item.RandomID]; exists {
return ErrMessageRandomIDDuplicate
}
seen[item.RandomID] = struct{}{}
}
return nil
}

75
internal/domain/channel.go
View file

@ -11,6 +11,9 @@ const (
MaxChannelDifferenceLimit = 100
// MaxChannelDifferenceTooLongMessages limits the latest message snapshot returned by channelDifferenceTooLong.
MaxChannelDifferenceTooLongMessages = 100
// MaxChannelUpdateRetentionBatch bounds one channel durable-log pruning transaction.
// Retention advances the recoverable floor only through rows actually deleted in that transaction.
MaxChannelUpdateRetentionBatch = 10000
// MaxChannelParticipantsLimit limits a single participants page.
MaxChannelParticipantsLimit = 200
// MaxChannelParticipantsOffset bounds channels.getParticipants deep OFFSET work.
@ -1184,6 +1187,23 @@ type DirtyChannel struct {
Pts int
}
// ChannelUpdateRetentionCheckpoint is the durable recovery boundary for one channel.
// Events with pts <= RetainedThroughPts may be absent; callers below that floor must receive
// channelDifferenceTooLong. LatestEventDate/LatestPts survive event pruning and keep account-level
// dirty-channel nudges reconstructable after the hot event rows have been removed.
type ChannelUpdateRetentionCheckpoint struct {
ChannelID int64
RetainedThroughPts int
LatestEventDate int
LatestPts int
}
// ChannelUpdateRetentionResult describes one bounded, atomic prune operation.
type ChannelUpdateRetentionResult struct {
Checkpoint ChannelUpdateRetentionCheckpoint
Deleted int
}
// CreateChannelRequest creates a broadcast channel or megagroup.
type CreateChannelRequest struct {
CreatorUserID int64
@ -1372,14 +1392,20 @@ type DeleteChannelResult struct {
// SendChannelMessageRequest sends one channel/supergroup message.
type SendChannelMessageRequest struct {
UserID int64
ChannelID int64
RandomID int64
Message string
Entities []MessageEntity
Media *MessageMedia
MentionUserIDs []int64
SkipDeliveryUserIDs []int64
UserID int64
ChannelID int64
RandomID int64
// IdempotencyFingerprint is the SHA-256 of the immutable client send intent. RPC callers
// provide a raw-TL per-item value; internal callers leave it empty for the store fallback.
IdempotencyFingerprint []byte
// IdempotencyPreflighted is trusted internal execution metadata; see the private-send
// equivalent. It is deliberately excluded from the durable fingerprint.
IdempotencyPreflighted bool
Message string
Entities []MessageEntity
Media *MessageMedia
MentionUserIDs []int64
SkipDeliveryUserIDs []int64
// SkipRecipientLookup lets high-level realtime fan-out use the online member
// read model instead of forcing store.SendChannelMessage to synchronously
// return an active-member recipient list after commit.
@ -1405,13 +1431,26 @@ type SendChannelMessageRequest struct {
// 虚拟频道 id;SavedPeer 是订阅者子会话分组键(订阅者发=自己,管理员回复=目标订阅者);
// SenderUserID 是实际发件人。发件权限(订阅者身份/管理员)在 RPC 层校验,store 只校验 monoforum 存在。
type SendMonoforumMessageRequest struct {
MonoforumID int64
SenderUserID int64
SavedPeer Peer
RandomID int64
Message string
Entities []MessageEntity
Date int
MonoforumID int64
SenderUserID int64
SavedPeer Peer
RandomID int64
IdempotencyFingerprint []byte
IdempotencyPreflighted bool
Message string
Entities []MessageEntity
Date int
}
// ChannelSendReplayRequest addresses either a regular channel send (SavedPeer is zero) or one
// monoforum sub-dialog send (SavedPeer is the subscriber scope). Lookup is read-only and must
// never re-run membership/permission checks or allocate pts/message ids.
type ChannelSendReplayRequest struct {
ChannelID int64
SenderUserID int64
SavedPeer Peer
RandomID int64
IdempotencyFingerprint []byte
}
// MonoforumHistoryFilter 按订阅者子会话拉取 monoforum 私信历史。
@ -1520,7 +1559,11 @@ type SendChannelMessageResult struct {
Event ChannelUpdateEvent
Recipients []int64
Duplicate bool
Discussion *SendChannelDiscussionResult
// ReplayDeleteEvent is the existing durable channel delete event paired
// with a deleted exact-random_id replay. It must be returned only to the
// caller echo and must never be fanned out as a fresh event.
ReplayDeleteEvent *ChannelUpdateEvent
Discussion *SendChannelDiscussionResult
// MentionUserIDs 是本条消息解析出的被 @ 成员；在线 fanout 按它为
// 每个接收者投影 message.mentioned/media_unread。
MentionUserIDs []int64

55
internal/domain/login_code_delivery.go Normal file
View file

@ -0,0 +1,55 @@
package domain
import (
"fmt"
"math"
"strings"
)
const officialLoginCodeMessageTemplate = `Login code: %s. Do not give this code to anyone, even if they say they are from Telegram!
This code can be used to log in to your Telegram account. We never ask it for anything else.
If you didn't request this code by trying to log in on another device, simply ignore this message.`
// LoginCodeDeliveryRequest describes one durable 777000 login-code delivery.
// PhoneCodeHash is an opaque idempotency token and must never be persisted in
// plaintext; store implementations persist only its SHA-256 digest.
type LoginCodeDeliveryRequest struct {
UserID int64
PhoneCodeHash string
Code string
Date int
// ExpiresAt is the unix second after which the compact idempotency receipt
// may be reclaimed. It must cover the corresponding code's usable lifetime.
ExpiresAt int64
}
// LoginCodeDeliveryResult returns the immutable first delivery. Created is
// false when the same phone_code_hash was already committed and replayed.
type LoginCodeDeliveryResult struct {
Message Message
Created bool
}
// OfficialLoginCodeMessage builds the account-visible incoming message from
// Telegram's official notification account. Persistence assigns ID, UID and
// Pts atomically.
func OfficialLoginCodeMessage(userID int64, code string, date int) (Message, error) {
if userID <= 0 || IsSystemUserID(userID) || strings.TrimSpace(code) == "" || len(code) > 64 || date < 0 || date > math.MaxInt32 {
return Message{}, fmt.Errorf("%w: user=%d code_length=%d date=%d", ErrLoginCodeDeliveryInvalid, userID, len(code), date)
}
body := fmt.Sprintf(officialLoginCodeMessageTemplate, code)
codeOffset := len("Login code: ")
return Message{
OwnerUserID: userID,
Peer: Peer{Type: PeerTypeUser, ID: OfficialSystemUserID},
From: Peer{Type: PeerTypeUser, ID: OfficialSystemUserID},
Date: date,
Body: body,
Entities: []MessageEntity{
{Type: MessageEntityBold, Offset: 0, Length: len("Login code:")},
{Type: MessageEntityBold, Offset: codeOffset, Length: len(code)},
},
}, nil
}

22
internal/domain/media.go
View file

@ -3,6 +3,7 @@ package domain
import (
"path/filepath"
"strings"
"time"
)
// 本文件定义媒体相关的业务值对象（文档、照片、贴纸集、可用 reaction、消息媒体）。
@ -81,6 +82,27 @@ type UploadedFileRef struct {
MD5 string // small file 客户端 md5_checksum（hex），可校验；big file 为空
}
// UploadedMediaKind identifies the durable object materialized from a one-shot upload file id.
type UploadedMediaKind string
const (
UploadedMediaPhoto UploadedMediaKind = "photo"
UploadedMediaDocument UploadedMediaKind = "document"
)
// UploadedMediaReceipt makes InputMediaUploaded* replayable after transient upload parts have
// been consumed. IntentHash binds the file id to the complete materialization intent (kind, part
// metadata and document spec); MediaID points at the immutable Photo/Document returned on every
// exact replay.
type UploadedMediaReceipt struct {
OwnerUserID int64
FileID int64
IntentHash []byte
Kind UploadedMediaKind
MediaID int64
CreatedAt time.Time
}
// DocumentSpec 描述从上传文件创建 Document 的元数据（来自 InputMediaUploadedDocument）。
type DocumentSpec struct {
MimeType string

41
internal/domain/message.go
View file

@ -260,8 +260,18 @@ type SendPrivateTextRequest struct {
OriginAuthKeyID [8]byte
OriginSessionID int64
RecipientBlocked bool
TTLPeriod int
ViaBotID int64
// IdempotencyFingerprint 是调用边界对原始、不可变发送请求计算的 SHA-256。
// RPC 层应优先填入原始 TL 请求指纹，避免链接预览、骰子结果、上传媒体
// 等服务端派生字段让合法重放看起来不同；内部调用留空时 store 会基于
// domain command 的不可变字段生成等价指纹。
IdempotencyFingerprint []byte
// IdempotencyPreflighted is internal execution metadata. A trusted caller sets it only
// after a read-only replay lookup returned absent, allowing the app/store layers to avoid
// repeating the same indexed lookup. The transactional unique-key path still fences a
// concurrent first writer; this flag is never part of the durable request fingerprint.
IdempotencyPreflighted bool
TTLPeriod int
ViaBotID int64
// GroupedID 相册分组 id（sendMultiMedia 同组共享非零值，非相册恒 0）。
GroupedID int64
// Effect 消息特效 id（私聊专属，0 表无特效；调用方已对 catalog 校验过合法性）。
@ -275,6 +285,16 @@ type SendPrivateTextRequest struct {
RichMessage *MessageRichMessage
}
// PrivateSendReplayRequest identifies one already-committed private send without carrying any
// mutable or resolver-derived message fields. The fingerprint is computed at the original
// request boundary and must be a complete SHA-256 value.
type PrivateSendReplayRequest struct {
SenderUserID int64
RecipientUserID int64
RandomID int64
IdempotencyFingerprint []byte
}
// SendPrivateTextResult 描述一次私聊文本发送的双端结果。
type SendPrivateTextResult struct {
SenderMessage Message
@ -282,6 +302,10 @@ type SendPrivateTextResult struct {
SenderEvent UpdateEvent
RecipientEvent UpdateEvent
Duplicate bool
// ReplayDeleteEvent is the already-durable sender-side deletion that must
// follow the first-send snapshot in an exact random_id replay. It never
// represents a newly allocated event.
ReplayDeleteEvent *UpdateEvent
}
// SetPrivateChatThemeRequest changes the shared theme token for a private dialog.
@ -351,12 +375,13 @@ type ForwardPrivateMessagesRequest struct {
// ForwardPrivateMessagesResult 描述一次私聊转发的 owner 维度结果。
type ForwardPrivateMessagesResult struct {
OwnerUserID int64
SenderMessages []Message
RecipientMessages []Message
SenderEvents []UpdateEvent
RecipientEvents []UpdateEvent
Duplicates []bool
OwnerUserID int64
SenderMessages []Message
RecipientMessages []Message
SenderEvents []UpdateEvent
RecipientEvents []UpdateEvent
Duplicates []bool
ReplayDeleteEvents []*UpdateEvent
}
// ReadHistoryRequest 是账号视角的 messages.readHistory 命令。

31
internal/domain/message_errors.go
View file

@ -3,13 +3,30 @@ package domain
import "errors"
var (
ErrMessageIDInvalid = errors.New("message id invalid")
ErrMessageEmpty = errors.New("message empty")
ErrMessageAuthorRequired = errors.New("message author required")
ErrMessageNotModified = errors.New("message not modified")
ErrMessageNotReadYet = errors.New("message not read yet")
ErrReplyMessageIDInvalid = errors.New("reply message id invalid")
ErrChatForwardsRestricted = errors.New("chat forwards restricted")
ErrMessageIDInvalid = errors.New("message id invalid")
ErrMessageEmpty = errors.New("message empty")
ErrMessageAuthorRequired = errors.New("message author required")
ErrMessageNotModified = errors.New("message not modified")
ErrMessageNotReadYet = errors.New("message not read yet")
// ErrMessageRandomIDDuplicate 表示同一发送者重复使用 random_id，且本次
// 不可变请求载荷与首次成功发送不一致。完全相同的重放不返回此错误，
// 而是复用首次发送结果。
ErrMessageRandomIDDuplicate = errors.New("message random id duplicate")
// ErrLoginCodeDeliveryInvalid rejects malformed durable 777000 delivery
// commands before allocating message/pts facts.
ErrLoginCodeDeliveryInvalid = errors.New("login code delivery invalid")
// ErrLoginCodeDeliveryConflict means one phone_code_hash digest was reused
// for a different account or code. It must fail closed rather than expose or
// overwrite the first account's immutable receipt.
ErrLoginCodeDeliveryConflict = errors.New("login code delivery conflict")
// ErrLoginCodeDeliveryCommitAmbiguous means PostgreSQL lost the commit
// acknowledgement and an independent receipt probe could not prove whether
// the durable 777000 transaction committed. Callers must retain the opaque
// code record until TTL expiry; deleting it could invalidate a committed but
// undisclosed delivery and make a retry impossible to reconcile.
ErrLoginCodeDeliveryCommitAmbiguous = errors.New("login code delivery commit ambiguous")
ErrReplyMessageIDInvalid = errors.New("reply message id invalid")
ErrChatForwardsRestricted = errors.New("chat forwards restricted")
// ErrPinnedSavedDialogsTooMuch 映射 PINNED_TOO_MUCH：收藏夹子会话置顶
// 数量达到 MaxPinnedSavedDialogs 上限。
ErrPinnedSavedDialogsTooMuch = errors.New("pinned saved dialogs too much")

166
internal/mtprotoedge/admission.go Normal file
View file

@ -0,0 +1,166 @@
package mtprotoedge
import (
"context"
"net"
"sync"
"time"
)
const (
defaultMaxConnections = 200_000
defaultMaxConnectionsPerIP = 4_096
defaultMaxConcurrentHandshakes = 256
acceptRetryInitialDelay = 5 * time.Millisecond
acceptRetryMaxDelay = time.Second
)
// admissionController 把 raw socket 与昂贵的 RSA/DH exchange 分开限流。
// raw 配额覆盖连接从 Accept 到物理 Close 的完整生命周期；handshake 配额只覆盖
// auth_key_id=0 的 exchange（包括 TDesktop 每条候选连接的 fake req_pq 探活）。
type admissionController struct {
mu sync.Mutex
maxConnections int
maxPerIP int
connections int
byIP map[string]int
handshakes chan struct{}
}
func newAdmissionController(maxConnections, maxPerIP, maxHandshakes int) *admissionController {
a := &admissionController{
maxConnections: maxConnections,
maxPerIP: maxPerIP,
byIP: make(map[string]int),
}
if maxHandshakes > 0 {
a.handshakes = make(chan struct{}, maxHandshakes)
}
return a
}
func (a *admissionController) wrapListener(ln net.Listener) net.Listener {
if a == nil {
return ln
}
return &admissionListener{Listener: ln, admission: a}
}
func (a *admissionController) acquireConnection(addr net.Addr) (func(), bool) {
if a == nil {
return func() {}, true
}
ip := remoteAdmissionKey(addr)
a.mu.Lock()
if (a.maxConnections > 0 && a.connections >= a.maxConnections) ||
(a.maxPerIP > 0 && a.byIP[ip] >= a.maxPerIP) {
a.mu.Unlock()
return nil, false
}
a.connections++
a.byIP[ip]++
a.mu.Unlock()
var once sync.Once
return func() {
once.Do(func() {
a.mu.Lock()
a.connections--
a.byIP[ip]--
if a.byIP[ip] == 0 {
delete(a.byIP, ip)
}
a.mu.Unlock()
})
}, true
}
func (a *admissionController) tryAcquireHandshake() (func(), bool) {
if a == nil || a.handshakes == nil {
return func() {}, true
}
select {
case a.handshakes <- struct{}{}:
var once sync.Once
return func() {
once.Do(func() { <-a.handshakes })
}, true
default:
return nil, false
}
}
func remoteAdmissionKey(addr net.Addr) string {
if addr == nil {
return "<unknown>"
}
if host, _, err := net.SplitHostPort(addr.String()); err == nil {
if ip := net.ParseIP(host); ip != nil {
return ip.String()
}
return host
}
return addr.Network() + ":" + addr.String()
}
// admissionListener 在最早的原始 Accept 边界记账，因此 mixed TCP/WebSocket 的
// sniff/upgrade 阶段也受 raw cap 保护。admittedConn.Close 负责幂等归还配额。
type admissionListener struct {
net.Listener
admission *admissionController
}
func (l *admissionListener) Accept() (net.Conn, error) {
for {
conn, err := l.Listener.Accept()
if err != nil {
return nil, err
}
release, ok := l.admission.acquireConnection(conn.RemoteAddr())
if !ok {
_ = conn.Close()
continue
}
return &admittedConn{Conn: conn, release: release}, nil
}
}
type admittedConn struct {
net.Conn
release func()
once sync.Once
}
func (c *admittedConn) Close() error {
err := c.Conn.Close()
c.once.Do(c.release)
return err
}
func isTemporaryAcceptError(err error) bool {
type temporary interface{ Temporary() bool }
e, ok := err.(temporary)
return ok && e.Temporary()
}
func nextAcceptRetryDelay(previous time.Duration) time.Duration {
if previous <= 0 {
return acceptRetryInitialDelay
}
next := previous * 2
if next > acceptRetryMaxDelay {
return acceptRetryMaxDelay
}
return next
}
func waitAcceptRetry(ctx context.Context, delay time.Duration) bool {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-timer.C:
return true
case <-ctx.Done():
return false
}
}

313
internal/mtprotoedge/admission_test.go Normal file
View file

@ -0,0 +1,313 @@
package mtprotoedge
import (
"context"
"errors"
"net"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto/codec"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
func TestAdmissionConnectionLimitsAndIdempotentRelease(t *testing.T) {
a := newAdmissionController(2, 1, 1)
ip1a := &net.TCPAddr{IP: net.ParseIP("203.0.113.1"), Port: 1000}
ip1b := &net.TCPAddr{IP: net.ParseIP("203.0.113.1"), Port: 1001}
ip2 := &net.TCPAddr{IP: net.ParseIP("203.0.113.2"), Port: 1000}
ip3 := &net.TCPAddr{IP: net.ParseIP("203.0.113.3"), Port: 1000}
release1, ok := a.acquireConnection(ip1a)
if !ok {
t.Fatal("first connection rejected")
}
if _, ok := a.acquireConnection(ip1b); ok {
t.Fatal("second connection from same IP bypassed per-IP cap")
}
release2, ok := a.acquireConnection(ip2)
if !ok {
t.Fatal("second IP connection rejected below global cap")
}
if _, ok := a.acquireConnection(ip3); ok {
t.Fatal("third connection bypassed global cap")
}
release1()
release1() // 幂等归还不得把计数减成负数。
releaseAgain, ok := a.acquireConnection(ip1b)
if !ok {
t.Fatal("released per-IP/global slot was not reusable")
}
releaseAgain()
release2()
a.mu.Lock()
defer a.mu.Unlock()
if a.connections != 0 || len(a.byIP) != 0 {
t.Fatalf("admission counters after release = %d/%v, want 0/empty", a.connections, a.byIP)
}
}
func TestAdmissionHandshakeLimitAndRelease(t *testing.T) {
a := newAdmissionController(-1, -1, 1)
release, ok := a.tryAcquireHandshake()
if !ok {
t.Fatal("first handshake rejected")
}
if _, ok := a.tryAcquireHandshake(); ok {
t.Fatal("second handshake bypassed semaphore")
}
release()
release() // 幂等
release2, ok := a.tryAcquireHandshake()
if !ok {
t.Fatal("released handshake slot was not reusable")
}
release2()
}
type oneConnListener struct {
conn net.Conn
once sync.Once
}
func (l *oneConnListener) Accept() (net.Conn, error) {
var conn net.Conn
l.once.Do(func() {
conn = l.conn
})
if conn == nil {
return nil, net.ErrClosed
}
return conn, nil
}
func (l *oneConnListener) Close() error { return l.conn.Close() }
func (l *oneConnListener) Addr() net.Addr { return l.conn.LocalAddr() }
func TestAdmissionListenerTracksUntilPhysicalClose(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
a := newAdmissionController(1, 1, 1)
ln := a.wrapListener(&oneConnListener{conn: serverSide})
conn, err := ln.Accept()
if err != nil {
t.Fatalf("Accept: %v", err)
}
a.mu.Lock()
active := a.connections
a.mu.Unlock()
if active != 1 {
t.Fatalf("active after Accept = %d, want 1", active)
}
_ = conn.Close()
_ = conn.Close()
a.mu.Lock()
active = a.connections
a.mu.Unlock()
if active != 0 {
t.Fatalf("active after physical Close = %d, want 0", active)
}
}
type temporaryAcceptTestError struct{}
func (temporaryAcceptTestError) Error() string { return "temporary accept failure" }
func (temporaryAcceptTestError) Timeout() bool { return false }
func (temporaryAcceptTestError) Temporary() bool { return true }
type temporaryThenConnListener struct {
conn net.Conn
closed chan struct{}
closeOnce sync.Once
calls atomic.Int32
}
type connThenErrorListener struct {
conn net.Conn
err error
closeOnce sync.Once
calls atomic.Int32
}
func (l *connThenErrorListener) Accept() (net.Conn, error) {
if l.calls.Add(1) == 1 {
return l.conn, nil
}
return nil, l.err
}
func (l *connThenErrorListener) Close() error {
var err error
l.closeOnce.Do(func() {
if l.conn != nil {
err = l.conn.Close()
}
})
return err
}
func (l *connThenErrorListener) Addr() net.Addr {
if l.conn != nil {
return l.conn.LocalAddr()
}
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}
}
type fixedErrorListener struct {
err error
addr net.Addr
}
func (l *fixedErrorListener) Accept() (net.Conn, error) { return nil, l.err }
func (*fixedErrorListener) Close() error { return nil }
func (l *fixedErrorListener) Addr() net.Addr { return l.addr }
func (l *temporaryThenConnListener) Accept() (net.Conn, error) {
call := l.calls.Add(1)
if call == 1 {
return nil, temporaryAcceptTestError{}
}
if call == 2 {
return l.conn, nil
}
<-l.closed
return nil, net.ErrClosed
}
func (l *temporaryThenConnListener) Close() error {
l.closeOnce.Do(func() {
close(l.closed)
_ = l.conn.Close()
})
return nil
}
func (l *temporaryThenConnListener) Addr() net.Addr {
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}
}
func TestAcceptLoopRetriesTemporaryError(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
ln := &temporaryThenConnListener{conn: serverSide, closed: make(chan struct{})}
srv := New(Options{HandshakeIdleTimeout: 100 * time.Millisecond})
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { done <- srv.acceptLoop(ctx, ln, false) }()
deadline := time.Now().Add(time.Second)
for ln.calls.Load() < 3 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if ln.calls.Load() < 3 {
cancel()
<-done
t.Fatalf("accept calls = %d, want temporary retry then next accept", ln.calls.Load())
}
cancel()
select {
case err := <-done:
if err != nil {
t.Fatalf("acceptLoop after temporary error: %v", err)
}
case <-time.After(time.Second):
t.Fatal("acceptLoop did not stop after cancel")
}
}
func TestAcceptLoopPermanentErrorCancelsAcceptedConnectionsBeforeWait(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
wantErr := errors.New("permanent accept failure")
ln := &connThenErrorListener{conn: serverSide, err: wantErr}
srv := New(Options{HandshakeIdleTimeout: time.Hour})
done := make(chan error, 1)
go func() {
done <- srv.acceptLoop(context.Background(), ln, false)
}()
select {
case err := <-done:
if !errors.Is(err, wantErr) {
t.Fatalf("acceptLoop error = %v, want %v", err, wantErr)
}
case <-time.After(time.Second):
t.Fatal("acceptLoop waited for an accepted connection before canceling it")
}
_ = clientSide.SetReadDeadline(time.Now().Add(time.Second))
var one [1]byte
if _, err := clientSide.Read(one[:]); err == nil {
t.Fatal("accepted connection remained open after permanent accept failure")
}
}
func TestServeMixedStopsAllComponentsWhenOneReturnsCleanly(t *testing.T) {
srv := New(Options{WebSocket: true})
ln := &fixedErrorListener{
err: net.ErrClosed,
addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 2398},
}
done := make(chan error, 1)
go func() {
done <- srv.serveMixed(context.Background(), ln)
}()
select {
case err := <-done:
if err != nil {
t.Fatalf("serveMixed error = %v, want nil closed-listener shutdown", err)
}
case <-time.After(time.Second):
t.Fatal("serveMixed did not stop remaining components after one clean exit")
}
}
type countingAuthKeyStore struct {
store.AuthKeyStore
gets atomic.Int32
}
func (s *countingAuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
s.gets.Add(1)
return s.AuthKeyStore.Get(ctx, id)
}
func TestUnknownAuthKeyRespondsOnceThenCloses(t *testing.T) {
keys := &countingAuthKeyStore{AuthKeyStore: memory.NewAuthKeyStore()}
addr, _, _ := startTestServer(t, Options{AuthKeys: keys})
conn := dialTransportOnly(t, addr)
var request bin.Buffer
request.PutLong(0x0102030405060708)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := conn.Send(ctx, &request); err != nil {
t.Fatalf("send unknown auth key: %v", err)
}
var response bin.Buffer
err := conn.Recv(ctx, &response)
var protocolErr *codec.ProtocolErr
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
t.Fatalf("first recv err = %T %v, want protocol -404", err, err)
}
if got := keys.gets.Load(); got != 1 {
t.Fatalf("AuthKeyStore.Get calls = %d, want 1", got)
}
response.Reset()
err = conn.Recv(ctx, &response)
if err == nil {
t.Fatal("connection remained readable after terminal -404")
}
if got := keys.gets.Load(); got != 1 {
t.Fatalf("AuthKeyStore.Get calls after close = %d, want 1", got)
}
}

56
internal/mtprotoedge/auth_key_switch_test.go Normal file
View file

@ -0,0 +1,56 @@
package mtprotoedge
import (
"testing"
"time"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
)
func TestEncryptedConnectionSwitchesAuthKeyEvenWhenSessionIDIsReused(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})
connA, authA, cipherA := dialHandshake(t, addr, dc, pub)
_, authB, cipherB := dialHandshake(t, addr, dc, pub)
msgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, connA, cipherA, authA, msgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
for range 3 { // new_session_created + pong + msgs_ack; leave no A-key frame on the socket.
readServerMessage(t, connA, cipherA, authA.AuthKey)
}
// Reuse A's session id on the same physical TCP socket, but encrypt with the independently
// established key B and B's salt. Session identity is (raw auth_key_id, session_id): comparing
// session_id alone would keep A's cached key/user identity and encrypt the reply with A.
body := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2})
sendEncryptedWithSessionSaltAndSeq(
t,
connA,
cipherB,
authB,
authA.SessionID,
authB.ServerSalt,
msgID.New(proto.MessageFromClient),
1,
body,
)
seenPong := false
for range 3 {
_, typeID, _ := readServerMessage(t, connA, cipherB, authB.AuthKey)
seenPong = seenPong || typeID == mt.PongTypeID
}
if !seenPong {
t.Fatal("new auth key did not receive pong")
}
oldKey := sessionKey{authKeyID: authA.AuthKey.ID, sessionID: authA.SessionID}
newKey := sessionKey{authKeyID: authB.AuthKey.ID, sessionID: authA.SessionID}
srv.conns.mu.RLock()
_, oldAlive := srv.conns.bySession[oldKey]
current := srv.conns.bySession[newKey]
srv.conns.mu.RUnlock()
if oldAlive || current == nil || current.authKeyID != authB.AuthKey.ID {
t.Fatalf("registry after key switch: old_alive=%v current=%v", oldAlive, current != nil)
}
}

50
internal/mtprotoedge/conn.go
View file

@ -46,27 +46,53 @@ type Conn struct {
outboundStop chan struct{}
outboundDone chan struct{}
outboundClose sync.Once
// outboundEnqueueMu orders producer registration against terminal close. Close
// flips closing under this lock before waiting, so no WaitGroup Add can race Wait.
outboundEnqueueMu sync.Mutex
outboundEnqueueWG sync.WaitGroup
outboundClosing bool
// Queue backing is intentionally small and bounded per Conn; control has a separate queue
// and strict actor priority. Server-created connections share outboundTrackedBudget.
outboundQueueSize int
outboundControlQueueSize int
outboundTrackedBudget *outboundTrackedBudget
outboundBudgetOnce sync.Once
// Encoded MTProto service frames and control vectors use independent headroom: pong,
// new_session_created, bad_msg and msgs_ack must remain admissible when the body budget is
// full. Content-related control frames keep this budget while pending for resend.
outboundControlTrackedBudget *outboundTrackedBudget
outboundControlBudgetOnce sync.Once
outboundScratchPool *outboundScratchPool
outboundScratchOnce sync.Once
// terminal 表示该 logical Conn 已停止接受新的出站操作。写失败时由
// outbound actor 置位并只发停止信号，不能在 actor 内等待自身退出。
terminal atomic.Bool
transportClose sync.Once
rpcQueue chan inboundRPC
rpcStop chan struct{}
rpcCancel context.CancelFunc
rpcClose sync.Once
rpcWG sync.WaitGroup
rpcTimeout time.Duration
rpcScheduler *inboundRPCScheduler
rpcCancel context.CancelFunc
rpcClose sync.Once
rpcMu sync.Mutex
rpcWG sync.WaitGroup
// rpcReservationWG 跟踪 Copy 前预算到 commit/abort 的短窗口，使 Close 返回时
// 全局/单连接预算都已归还或转交给明确的 queued/running task。
rpcReservationWG sync.WaitGroup
rpcTimeout time.Duration
rpcQueue []inboundRPC
rpcQueueSize int
rpcReserved int
rpcRunning int
rpcReady bool
rpcClosed bool
// inflightRPCBytes 跟踪已入队未完成的 inbound RPC body 总字节，配合 maxInflightRPCBytes
// 给 RPC 队列设字节预算（不止限条数），防对抗客户端发大请求撑内存。
inflightRPCBytes atomic.Int64
// RPC worker 懒启动：首个 RPC 入队时才起 worker（ensureInboundRPCWorkers），
// 避免握手后静默 / 纯推送目标连接白白钉住 rpcMaxInflight 个 goroutine。
// 单连接只保留并发配额；实际 worker 来自 Server 共享池，避免每连接预留 goroutine。
rpcRootCtx context.Context
rpcMaxInflight int
rpcWorkersOnce sync.Once
// sentContentMessages 只由 outbound actor 访问，用于生成 MTProto seq_no。
sentContentMessages int32
// outboundPlain/outboundWire 只由 outbound actor 访问，用于复用出站加密缓冲。
outboundPlain bin.Buffer
outboundWire bin.Buffer
// outboundRand 只由 outbound actor 访问：对 cipher 随机源的缓冲预读，
// 把每帧 padding 的 getrandom syscall 摊薄成 ~1KiB 一次。
outboundRand *bufio.Reader

428
internal/mtprotoedge/encrypted.go
View file

@ -1,6 +1,8 @@
package mtprotoedge
import (
"bytes"
"compress/gzip"
"context"
"crypto/sha256"
"encoding/binary"
@ -8,6 +10,7 @@ import (
"fmt"
"io"
"math"
"sync/atomic"
"time"
"go.uber.org/zap"
@ -59,6 +62,25 @@ func (cs *connState) reset() {
const (
maxTrackedClientMsgIDs = 400
// maxContainerMessages bounds per-frame recursive work and ack growth. Official clients batch
// far fewer messages; 1024 leaves ample headroom while preventing a 16 MiB frame of zero-body
// container entries from expanding into tens of MiB of Go objects.
maxContainerMessages = 1024
// maxDispatchDepth bounds gzip/container wrapper recursion. Normal shapes are RPC, gzip(RPC),
// container(RPC...) and gzip(container(...)); deeper nesting has no compatibility value.
maxDispatchDepth = 4
// gotd already caps each gzip expansion at 10 MiB. This cumulative cap prevents several nested
// gzip layers in one transport frame from repeatedly allocating/decompressing that allowance.
maxDispatchExpandedBytes = 32 << 20
maxSingleGZIPExpandedBytes = 10 << 20
// MTProto service vectors operate on bounded connection tracking tables. Accepting more IDs
// only burns decode/CPU and cannot improve the result.
maxServiceMessageIDs = 4096
// A decoded container descriptor is 48 bytes on 64-bit Go today. Charge 64 bytes per entry
// before allocating the exact-size slice so allocator rounding and future field growth remain
// inside the process-wide inbound budget. Message bodies stay as zero-copy views of the already
// charged plaintext frame/gzip expansion.
containerDescriptorBudgetBytes = 64
msgStateUnknown byte = 1
msgStateNotReceived byte = 2
@ -101,7 +123,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
if frame.salt != serverSalt {
c := current
temp := false
if c == nil || c.sessionID != frame.sessionID {
if c == nil || c.sessionID != frame.sessionID || c.authKeyID != key.ID {
c = s.newConn(tc, key, frame.sessionID, serverSalt)
temp = true
}
@ -113,7 +135,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
}
// 首个加密消息或 session 变化时（重新）注册连接到 SessionManager。
if current == nil || current.sessionID != frame.sessionID {
if current == nil || current.sessionID != frame.sessionID || current.authKeyID != key.ID {
if current != nil {
cs.reset()
}
@ -150,7 +172,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
)
return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code)
}
if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext); err != nil {
if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext, s.writeTimeout); err != nil {
return current, err
}
@ -224,12 +246,28 @@ func (s *Server) maybePersistSession(ctx context.Context, c *Conn, sessionID int
}
}
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, plaintext []byte) error {
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, plaintext []byte, writeTimeout time.Duration) error {
q, ok := tc.(quickAckTransport)
if !ok || !q.ConsumeQuickAckRequested() {
return nil
}
return q.SendQuickAck(ctx, clientQuickAckToken(key, plaintext))
token := clientQuickAckToken(key, plaintext)
deadline := time.Time{}
if writeTimeout > 0 {
deadline = time.Now().Add(writeTimeout)
}
if d, ok := ctx.Deadline(); ok && (deadline.IsZero() || d.Before(deadline)) {
deadline = d
}
if dq, ok := tc.(deadlineQuickAckTransport); ok {
return dq.SendQuickAckDeadline(deadline, token)
}
if deadline.IsZero() {
return q.SendQuickAck(ctx, token)
}
sendCtx, cancel := context.WithDeadline(ctx, deadline)
defer cancel()
return q.SendQuickAck(sendCtx, token)
}
// clientQuickAckToken 按 Android MTProto v2 公式计算 quick ack：SHA256(auth_key[88:120] +
@ -246,6 +284,20 @@ func clientQuickAckToken(key crypto.AuthKey, plaintext []byte) uint32 {
// dispatch 处理一条明文消息：解包 container/gzip，处理服务消息，其余转 RPC 路由。
// content-related 消息（ping、RPC）的 msg_id 会收集到 acks 以便统一确认。
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
expanded := 0
return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, b, acks, dispatchBudget{expanded: &expanded})
}
type dispatchBudget struct {
depth int
containerDepth int
expanded *int
}
func (s *Server) dispatchWithBudget(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64, budget dispatchBudget) error {
if budget.depth > maxDispatchDepth {
return fmt.Errorf("mtproto wrapper depth %d exceeds %d", budget.depth, maxDispatchDepth)
}
id, err := b.PeekID()
if err != nil {
return fmt.Errorf("peek type id: %w", err)
@ -258,20 +310,39 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
switch id {
case proto.GZIPTypeID:
var gz proto.GZIP
if err := gz.Decode(b); err != nil {
data, releaseExpansion, err := s.decodeGZIPWithGlobalBudget(b)
if err != nil {
return fmt.Errorf("decode gzip: %w", err)
}
return s.dispatch(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: gz.Data}, acks)
defer releaseExpansion()
*budget.expanded += len(data)
if *budget.expanded > maxDispatchExpandedBytes {
return fmt.Errorf("cumulative gzip expansion %d exceeds %d", *budget.expanded, maxDispatchExpandedBytes)
}
budget.depth++
return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: data}, acks, budget)
case proto.MessageContainerTypeID:
var container proto.MessageContainer
if err := container.Decode(b); err != nil {
if budget.containerDepth != 0 {
return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer)
}
count, err := containerMessageCount(b)
if err != nil {
return fmt.Errorf("decode container count: %w", err)
}
if count > maxContainerMessages {
return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer)
}
container, releaseContainer, err := s.decodeMessageContainerViews(b, count)
if err != nil {
return fmt.Errorf("decode container: %w", err)
}
defer releaseContainer()
if code := validateClientContainer(msgID, seqNo, container); code != 0 {
return s.sendBadMsg(ctx, c, msgID, seqNo, code)
}
budget.depth++
budget.containerDepth++
for i := range container.Messages {
m := container.Messages[i]
typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID()
@ -292,7 +363,7 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
return s.sendBadMsg(ctx, c, m.ID, int32(m.SeqNo), code)
}
cs.track(m.ID, int32(m.SeqNo), content, msgStateReceived)
if err := s.dispatch(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks); err != nil {
if err := s.dispatchWithBudget(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks, budget); err != nil {
return err
}
}
@ -323,6 +394,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
return s.sendFutureSalts(ctx, c, msgID, req.Num)
case mt.MsgsAckTypeID:
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return fmt.Errorf("msgs_ack vector: %w", err)
}
var ack mt.MsgsAck
if err := ack.Decode(b); err != nil {
return fmt.Errorf("decode msgs_ack: %w", err)
@ -332,6 +406,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
return nil
case mt.MsgsStateReqTypeID:
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return fmt.Errorf("msgs_state_req vector: %w", err)
}
var req mt.MsgsStateReq
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode msgs_state_req: %w", err)
@ -344,6 +421,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs)))
case mt.MsgResendReqTypeID:
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return fmt.Errorf("msg_resend_req vector: %w", err)
}
var req mt.MsgResendReq
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode msg_resend_req: %w", err)
@ -356,19 +436,22 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs)))
case mt.MsgsStateInfoTypeID:
var info mt.MsgsStateInfo
if err := info.Decode(b); err != nil {
reqMsgID, info, err := msgsStateInfoView(b)
if err != nil {
return fmt.Errorf("decode msgs_state_info: %w", err)
}
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", info.ReqMsgID), zap.Int("len", len(info.Info)))
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", reqMsgID), zap.Int("len", len(info)))
return nil
case mt.MsgsAllInfoTypeID:
var info mt.MsgsAllInfo
if err := info.Decode(b); err != nil {
count, info, err := msgsAllInfoView(b)
if err != nil {
return fmt.Errorf("decode msgs_all_info: %w", err)
}
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", len(info.MsgIDs)), zap.Int("len", len(info.Info)))
if len(info) != count {
return fmt.Errorf("decode msgs_all_info: info length %d does not match msg_ids %d", len(info), count)
}
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", count), zap.Int("len", len(info)))
return nil
case mt.DestroySessionRequestTypeID:
@ -424,11 +507,228 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
default:
ackContent()
body := b.Copy()
return s.enqueueRPC(ctx, c, msgID, id, body)
return s.enqueueRPC(ctx, c, msgID, id, b)
}
}
// decodeGZIPWithGlobalBudget reserves the maximum single-wrapper output before
// decompression starts. Once the actual size is known the excess reservation is
// returned, while the actual output remains charged through recursive dispatch.
// This closes the gap where every connection read goroutine could otherwise hold
// an unaccounted 10 MiB expansion before the shared RPC scheduler saw the body.
func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), error) {
compressed, err := gzipPackedBytesView(b)
if err != nil {
return nil, func() {}, err
}
reserved := int64(0)
release := func() {
if reserved > 0 && s.frameBudget != nil {
s.frameBudget.release(reserved)
reserved = 0
}
}
if s.frameBudget != nil {
reserved, err = s.frameBudget.reserve(maxSingleGZIPExpandedBytes, 0)
if err != nil {
return nil, func() {}, err
}
}
r, err := gzip.NewReader(bytes.NewReader(compressed))
if err != nil {
release()
return nil, func() {}, err
}
data, readErr := io.ReadAll(io.LimitReader(r, maxSingleGZIPExpandedBytes+1))
closeErr := r.Close()
if readErr != nil {
release()
return nil, func() {}, readErr
}
if closeErr != nil {
release()
return nil, func() {}, closeErr
}
if len(data) > maxSingleGZIPExpandedBytes {
release()
return nil, func() {}, fmt.Errorf("gzip expansion %d exceeds %d", len(data), maxSingleGZIPExpandedBytes)
}
if reserved > int64(len(data)) {
s.frameBudget.release(reserved - int64(len(data)))
reserved = int64(len(data))
}
return data, release, nil
}
// gzipPackedBytesView parses the TL bytes envelope without copying the compressed
// payload. proto.GZIP.Decode calls bin.Buffer.Bytes, which duplicates the compressed
// frame before allocating the decompressed result.
func gzipPackedBytesView(b *bin.Buffer) ([]byte, error) {
if b == nil || len(b.Buf) < 5 {
return nil, io.ErrUnexpectedEOF
}
if binary.LittleEndian.Uint32(b.Buf[:4]) != proto.GZIPTypeID {
return nil, fmt.Errorf("unexpected gzip constructor %#x", binary.LittleEndian.Uint32(b.Buf[:4]))
}
payload, _, err := tlBytesView(b.Buf[4:], -1)
return payload, err
}
// tlBytesView validates one TL bytes envelope and returns a view into the caller-owned buffer.
// maxPayload < 0 means that the enclosing frame budget is the only size limit. The limit is
// checked from the encoded length before touching the payload, so service messages cannot make
// generated decoders allocate an attacker-selected []byte first and validate it afterwards.
func tlBytesView(raw []byte, maxPayload int) ([]byte, int, error) {
if len(raw) < 1 {
return nil, 0, io.ErrUnexpectedEOF
}
header, size := 1, int(raw[0])
if size == 254 {
if len(raw) < 4 {
return nil, 0, io.ErrUnexpectedEOF
}
header = 4
size = int(raw[1]) | int(raw[2])<<8 | int(raw[3])<<16
} else if size == 255 {
return nil, 0, errors.New("invalid TL bytes length marker 255")
}
if maxPayload >= 0 && size > maxPayload {
return nil, 0, fmt.Errorf("TL bytes length %d exceeds %d", size, maxPayload)
}
padded := (header + size + 3) &^ 3
if size < 0 || padded < header || len(raw) < padded {
return nil, 0, io.ErrUnexpectedEOF
}
return raw[header : header+size : header+size], padded, nil
}
// decodeMessageContainerViews parses the container without proto.Message.Decode's per-body
// copies. Bodies are immutable views of b and stay alive only for this synchronous dispatch;
// enqueueRPC takes its own budgeted copy before returning. Only the exact-size descriptor slice
// is new memory, and that allocation is reserved globally first.
func (s *Server) decodeMessageContainerViews(b *bin.Buffer, count int) (proto.MessageContainer, func(), error) {
release := func() {}
if b == nil || len(b.Buf) < 8 {
return proto.MessageContainer{}, release, io.ErrUnexpectedEOF
}
if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != proto.MessageContainerTypeID {
return proto.MessageContainer{}, release, fmt.Errorf("unexpected constructor %#x", got)
}
declared := int(int32(binary.LittleEndian.Uint32(b.Buf[4:8])))
if declared != count || count < 0 || count > maxContainerMessages {
return proto.MessageContainer{}, release, fmt.Errorf("invalid message count %d", declared)
}
reserved := int64(0)
if count > 0 && s.frameBudget != nil {
var err error
reserved, err = s.frameBudget.reserve(int64(count*containerDescriptorBudgetBytes), 0)
if err != nil {
return proto.MessageContainer{}, release, err
}
release = func() {
if reserved > 0 {
s.frameBudget.release(reserved)
reserved = 0
}
}
}
messages := make([]proto.Message, count)
offset := 8
for i := range messages {
if len(b.Buf)-offset < 16 {
release()
return proto.MessageContainer{}, func() {}, io.ErrUnexpectedEOF
}
id := int64(binary.LittleEndian.Uint64(b.Buf[offset : offset+8]))
seqNo := int32(binary.LittleEndian.Uint32(b.Buf[offset+8 : offset+12]))
bodyLen := int(int32(binary.LittleEndian.Uint32(b.Buf[offset+12 : offset+16])))
offset += 16
if bodyLen < 0 || bodyLen > 1024*1024 {
release()
return proto.MessageContainer{}, func() {}, fmt.Errorf("message length %d is invalid", bodyLen)
}
if bodyLen > len(b.Buf)-offset {
release()
return proto.MessageContainer{}, func() {}, io.ErrUnexpectedEOF
}
bodyEnd := offset + bodyLen
messages[i] = proto.Message{
ID: id,
SeqNo: int(seqNo),
Bytes: bodyLen,
Body: b.Buf[offset:bodyEnd:bodyEnd],
}
offset = bodyEnd
}
return proto.MessageContainer{Messages: messages}, release, nil
}
func msgsStateInfoView(b *bin.Buffer) (int64, []byte, error) {
if b == nil || len(b.Buf) < 12 {
return 0, nil, io.ErrUnexpectedEOF
}
if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != mt.MsgsStateInfoTypeID {
return 0, nil, fmt.Errorf("unexpected constructor %#x", got)
}
info, _, err := tlBytesView(b.Buf[12:], maxServiceMessageIDs)
if err != nil {
return 0, nil, err
}
return int64(binary.LittleEndian.Uint64(b.Buf[4:12])), info, nil
}
func msgsAllInfoView(b *bin.Buffer) (int, []byte, error) {
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return 0, nil, fmt.Errorf("vector: %w", err)
}
count := int(int32(binary.LittleEndian.Uint32(b.Buf[8:12])))
// count is already non-negative and capped, but check remaining bytes before multiplying into
// an offset so malformed frames cannot produce an out-of-bounds slice.
if count > (len(b.Buf)-12)/8 {
return 0, nil, io.ErrUnexpectedEOF
}
offset := 12 + count*8
info, _, err := tlBytesView(b.Buf[offset:], maxServiceMessageIDs)
if err != nil {
return 0, nil, err
}
return count, info, nil
}
func containerMessageCount(b *bin.Buffer) (int, error) {
if b == nil || len(b.Buf) < 8 {
return 0, io.ErrUnexpectedEOF
}
if binary.LittleEndian.Uint32(b.Buf[:4]) != proto.MessageContainerTypeID {
return 0, fmt.Errorf("unexpected constructor %#x", binary.LittleEndian.Uint32(b.Buf[:4]))
}
count := int(int32(binary.LittleEndian.Uint32(b.Buf[4:8])))
if count < 0 {
return 0, fmt.Errorf("negative message count %d", count)
}
return count, nil
}
func validateFirstVectorCount(b *bin.Buffer, max int) error {
if b == nil || len(b.Buf) < 12 {
return io.ErrUnexpectedEOF
}
if got := binary.LittleEndian.Uint32(b.Buf[4:8]); got != bin.TypeVector {
return fmt.Errorf("unexpected vector constructor %#x", got)
}
count := int(int32(binary.LittleEndian.Uint32(b.Buf[8:12])))
if count < 0 {
return fmt.Errorf("negative vector count %d", count)
}
if count > max {
return fmt.Errorf("vector count %d exceeds %d", count, max)
}
return nil
}
func mergeStateInfo(primary, fallback []byte) []byte {
if len(primary) == 0 {
return fallback
@ -448,7 +748,7 @@ func mergeStateInfo(primary, fallback []byte) []byte {
// enqueueRPC 把一条 RPC 请求交给连接的 inbound 调度器。typeID 由 dispatch 传入
// （已 PeekID 过一次），method 只解析一次并随任务透传，避免同一请求三处重复 PeekID/typeName。
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, body []byte) error {
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, request *bin.Buffer) error {
method := s.typeName(typeID)
if cached, ok := s.cachedRPCResult(c, msgID); ok {
s.log.Info("RPC duplicate replay from session cache",
@ -459,13 +759,48 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
)
return c.SendEncoded(ctx, proto.MessageServerResponse, cached)
}
err := c.enqueueInboundRPC(ctx, inboundRPC{
method: method,
size: len(body),
// 两级条数/字节预算必须先于 Copy：对抗客户端不能用大量满尺寸请求在“判断队列满”
// 之前制造一轮无上限的临时 body 分配。reservation 在 commit/abort 间唯一持有预算。
reservation, err := c.reserveInboundRPC(ctx, method, request.Len())
if err != nil {
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
}
defer reservation.abort()
body := request.Copy()
responseGate := &rpcResponseGate{}
timeoutResponse := func() {
if !responseGate.tryTimeout() {
return
}
// 原 task context 已到期，使用有界的新 context 回显明确的可重试超时；
// 500 保持 TDesktop 默认重试语义，错误名区分于容量型 FLOOD_WAIT。
writeTimeout := c.writeTimeout
if writeTimeout <= 0 || writeTimeout > 5*time.Second {
writeTimeout = 5 * time.Second
}
responseCtx, cancel := context.WithTimeout(context.Background(), writeTimeout)
defer cancel()
if sendErr := s.sendResult(responseCtx, c, msgID, &mt.RPCError{
ErrorCode: 500,
ErrorMessage: "RPC_TIMEOUT",
}); sendErr != nil && !isClientDisconnect(sendErr) {
s.log.Debug("Send RPC timeout failed",
zap.String("method", method),
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID),
zap.Error(sendErr),
)
}
}
err = reservation.commit(inboundRPC{
method: method,
size: len(body),
onTimeout: timeoutResponse,
run: func(taskCtx context.Context) error {
// body 已是 enqueueRPC 入参的独立副本（dispatch 里 b.Copy()），且每个任务只 run 一次，
// body 是预算成功后生成的独立副本，且每个任务只 run 一次，
// 无需再 append 拷贝；直接复用，省掉一份 inbound 在途内存。
if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}); err != nil {
if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}, responseGate); err != nil {
fields := []zap.Field{
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", c.authKeyHex),
@ -482,8 +817,12 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
return nil
},
})
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
}
func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, msgID int64, method string, err error) error {
if errors.Is(err, ErrInboundRPCQueueFull) {
s.log.Debug("Inbound RPC queue full",
s.log.Debug("Inbound RPC capacity exhausted",
zap.String("method", method),
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", c.authKeyHex),
@ -498,7 +837,7 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
}
// handleRPC 把明文 RPC 请求交给 RPC 路由，并将结果或错误包成 rpc_result 回发。
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer) error {
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer, responseGate *rpcResponseGate) error {
if s.rpc == nil {
s.log.Warn("No RPC handler configured; dropping request", zap.String("method", method))
return nil
@ -534,12 +873,24 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
}
fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot())
if ctxErr := ctx.Err(); ctxErr != nil && err != nil {
// A canceled request context means the result cannot be delivered. Do not
// turn cancellation-derived handler errors into cacheable rpc_error replies.
s.log.Info("RPC canceled", append(fields, zap.NamedError("dispatch_error", err), zap.NamedError("context_error", ctxErr))...)
if ctxErr := ctx.Err(); ctxErr != nil {
// A canceled request context means neither a success nor an error can be delivered
// with this expired context. In particular, do not cache a late successful result and
// hand it to outbound: a past write deadline would correctly poison that transport and
// could prevent the scheduler's fresh-context RPC_TIMEOUT response from being sent.
cancelFields := append(fields, zap.NamedError("context_error", ctxErr))
if err != nil {
cancelFields = append(cancelFields, zap.NamedError("dispatch_error", err))
}
s.log.Info("RPC canceled", cancelFields...)
return ctxErr
}
// A deadline callback may have already emitted RPC_TIMEOUT while Dispatch was returning.
// Claim the single normal-response slot before serializing any success/error rpc_result.
if responseGate != nil && !responseGate.tryNormal() {
s.log.Info("RPC result suppressed after timeout", fields...)
return context.DeadlineExceeded
}
if err != nil {
var rpcErr *tgerr.Error
@ -565,6 +916,21 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
return nil
}
// rpcResponseGate guarantees exactly one terminal rpc_result per request. A running deadline
// races legitimately with a handler completing at the boundary; whichever path claims state
// first owns the response, and the other path becomes a no-op.
type rpcResponseGate struct {
state atomic.Uint32
}
func (g *rpcResponseGate) tryNormal() bool {
return g == nil || g.state.CompareAndSwap(0, 1)
}
func (g *rpcResponseGate) tryTimeout() bool {
return g != nil && g.state.CompareAndSwap(0, 2)
}
// sendResult 把 RPC 结果包成 rpc_result 并加密回发。
func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) error {
encoded, err := s.encodeRPCResult(c, reqMsgID, result)

64
internal/mtprotoedge/exchange.go
View file

@ -2,9 +2,9 @@ package mtprotoedge
import (
"context"
"encoding/binary"
"errors"
"fmt"
"sync"
"go.uber.org/zap"
@ -12,7 +12,6 @@ import (
"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/transport"
@ -31,7 +30,8 @@ func peekAuthKeyID(b *bin.Buffer) (id [8]byte, err error) {
// handleExchange 在收到 auth_key_id==0 的首帧后执行服务端 MTProto 密钥交换。
//
// first 是已读取的首帧（req_pq*），通过 bufferedConn 交还给 exchange 流程，
// 使其能从头读取握手消息。成功后将 auth key + server salt 落入 AuthKeyStore。
// 使其能从头读取握手消息。auth key + server salt 会在 DhGenOk 发出前落入
// AuthKeyStore；持久化失败时不向客户端确认握手成功。
func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first *bin.Buffer) (*bin.Buffer, error) {
if s.key.Zero() {
s.log.Error("Key exchange requested but server RSA key is not configured")
@ -62,11 +62,6 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
var encErr *exchange.UnexpectedEncryptedError
if errors.As(err, &encErr) {
replay := encErr.Frame
if len(replay) == 0 {
if lf := buffered.lastFrame(); lf != nil {
replay = lf.Buf
}
}
if len(replay) > 0 {
s.log.Debug("Key exchange interrupted by encrypted frame; replaying as existing session")
return &bin.Buffer{Buf: replay}, nil
@ -99,7 +94,7 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
zap.Duration("dur", s.clock.Now().Sub(start)),
)
return nil, s.authKeys.Save(ctx, authKeyData(res.Key, res.ServerSalt, s.clock.Now().Unix()))
return nil, nil
}
// authKeyData 把握手结果转换为 store 记录。
@ -142,9 +137,7 @@ var errTooManyHandshakeReqPQ = errors.New("too many req_pq frames in one handsha
// 用于密钥交换：serveConn 已读首帧用于 peek auth_key_id，再 push 回来交给 exchange。
type bufferedConn struct {
transport.Conn
mu sync.Mutex
pending []bin.Buffer
last bin.Buffer
reqPQCount int // 本次握手已见 req_pq(_multi) 帧数；只在握手期访问（Recv 单 goroutine）
}
@ -153,32 +146,38 @@ func newBufferedConn(conn transport.Conn) *bufferedConn {
}
func (c *bufferedConn) push(b *bin.Buffer) {
c.mu.Lock()
c.pending = append(c.pending, bin.Buffer{Buf: b.Copy()})
c.mu.Unlock()
if b == nil {
return
}
// serveConn is synchronously blocked in handleExchange, so the first frame's
// backing remains stable until the exchange returns. Keep a slice view instead
// of copying an attacker-sized transport frame.
buf := b.Buf
b.Buf = nil // transfer ownership; serveConn must not pin the frame after next Recv releases it
c.pending = append(c.pending, bin.Buffer{Buf: buf})
}
// Recv 优先返回已 push 的帧（FIFO），耗尽后读取底层连接。
func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
for {
c.mu.Lock()
if len(c.pending) > 0 {
e := c.pending[0]
c.pending[0] = bin.Buffer{}
c.pending = c.pending[1:]
c.last.ResetTo(e.Copy())
c.mu.Unlock()
b.ResetTo(e.Buf)
} else {
c.mu.Unlock()
if err := c.Conn.Recv(ctx, b); err != nil {
return err
}
c.mu.Lock()
c.last.ResetTo(b.Copy())
c.mu.Unlock()
}
if isUnencryptedMsgsAckFrame(b) {
// The ack is intentionally ignored during exchange. Drop its transport
// backing and shrink the retained high-water charge before the next Recv;
// otherwise a large trailing frame can consume global admission budget for
// the rest of a CPU-heavy key exchange even though no backing remains live.
b.Buf = nil
retainInboundFrameBackings(c.Conn, b)
continue
}
// req_pq 计数上界：仅在握手期生效（bufferedConn 只用于密钥交换），且 payload id 探测
@ -196,21 +195,17 @@ func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
// unencryptedPayloadID 返回未加密消息（auth_key_id==0）内层 TL payload 的 type id。
// 非未加密消息 / 解码失败时 ok=false。
func unencryptedPayloadID(frame *bin.Buffer) (uint32, bool) {
authKeyID, err := peekAuthKeyID(frame)
if err != nil || authKeyID != emptyAuthKeyID {
if frame == nil || len(frame.Buf) < 24 {
return 0, false
}
var msg proto.UnencryptedMessage
cp := &bin.Buffer{Buf: frame.Copy()}
if err := msg.Decode(cp); err != nil {
if binary.LittleEndian.Uint64(frame.Buf[:8]) != 0 {
return 0, false
}
payload := &bin.Buffer{Buf: msg.MessageData}
id, err := payload.PeekID()
if err != nil {
dataLen := int64(int32(binary.LittleEndian.Uint32(frame.Buf[16:20])))
if dataLen < 4 || dataLen > int64(len(frame.Buf)-20) {
return 0, false
}
return id, true
return binary.LittleEndian.Uint32(frame.Buf[20:24]), true
}
func isUnencryptedMsgsAckFrame(frame *bin.Buffer) bool {
@ -222,12 +217,3 @@ func isUnencryptedReqPQFrame(frame *bin.Buffer) bool {
id, ok := unencryptedPayloadID(frame)
return ok && (id == mt.ReqPqRequestTypeID || id == mt.ReqPqMultiRequestTypeID)
}
func (c *bufferedConn) lastFrame() *bin.Buffer {
c.mu.Lock()
defer c.mu.Unlock()
if c.last.Len() == 0 {
return nil
}
return &bin.Buffer{Buf: c.last.Copy()}
}

75
internal/mtprotoedge/exchange_compat.go
View file

@ -31,27 +31,42 @@ import (
// matches this server DC.
func (s *Server) runServerExchange(ctx context.Context, conn transport.Conn) (exchange.ServerExchangeResult, error) {
ex := serverExchangeCompat{
conn: conn,
clock: s.clock,
rand: s.rand,
timeout: exchange.DefaultTimeout,
key: s.key,
dc: s.dc,
log: s.log.Named("exchange"),
rng: compatServerRNG{rand: s.rand},
conn: conn,
clock: s.clock,
rand: s.rand,
timeout: exchange.DefaultTimeout,
key: s.key,
dc: s.dc,
log: s.log.Named("exchange"),
rng: compatServerRNG{rand: s.rand},
commitKey: s.commitExchangeAuthKey,
}
return ex.run(ctx)
}
// commitExchangeAuthKey is the durable commit point of the server exchange.
// It must complete before DhGenOk is put on the wire: after that response the
// 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 {
createdAt := s.clock.Now().Unix()
if err := s.authKeys.Save(ctx, authKeyData(result.Key, result.ServerSalt, createdAt)); err != nil {
return fmt.Errorf("persist auth key before DhGenOk: %w", err)
}
return nil
}
type serverExchangeCompat struct {
conn transport.Conn
clock clock.Clock
rand io.Reader
timeout time.Duration
key exchange.PrivateKey
dc int
log *zap.Logger
rng compatServerRNG
conn transport.Conn
clock clock.Clock
rand io.Reader
timeout time.Duration
key exchange.PrivateKey
dc int
log *zap.Logger
rng compatServerRNG
commitKey func(context.Context, exchange.ServerExchangeResult) error
}
func (s serverExchangeCompat) run(ctx context.Context) (exchange.ServerExchangeResult, error) {
@ -193,6 +208,21 @@ SendResPQ:
return exchange.ServerExchangeResult{}, wrapKeyNotFound(err)
}
serverResult := exchange.ServerExchangeResult{
Key: authKey.WithID(),
ServerSalt: crypto.ServerSalt(innerData.NewNonce, serverNonce),
}
// DhGenOk is the externally visible commit acknowledgement. Require a
// durable key commit before sending it, rather than allowing callers to
// persist after run returns. A nil hook is rejected so a future call site
// cannot accidentally reintroduce the unsafe ordering.
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 {
return exchange.ServerExchangeResult{}, err
}
s.log.Debug("Sending DhGenOk")
if err := s.writeUnencrypted(ctx, b, &mt.DhGenOk{
Nonce: req.Nonce,
@ -202,11 +232,7 @@ SendResPQ:
return exchange.ServerExchangeResult{}, err
}
serverSalt := crypto.ServerSalt(innerData.NewNonce, serverNonce)
return exchange.ServerExchangeResult{
Key: authKey.WithID(),
ServerSalt: serverSalt,
}, nil
return serverResult, nil
}
func (s serverExchangeCompat) validatePQInnerDataDC(d mt.PQInnerDataClass) error {
@ -278,9 +304,14 @@ func (s serverExchangeCompat) readUnencrypted(ctx context.Context, b *bin.Buffer
var keyID [8]byte
if err := b.PeekN(keyID[:], len(keyID)); err == nil && keyID != ([8]byte{}) {
// The exchange aborts immediately on an encrypted frame, so transfer the received backing
// to the replay error instead of making an unbudgeted near-transport-limit copy. serveConn
// keeps the existing frame reservation until replay dispatch has finished.
frame := b.Buf
b.Buf = nil
return &exchange.UnexpectedEncryptedError{
AuthKeyID: keyID,
Frame: append([]byte(nil), b.Buf...),
Frame: frame,
}
}

285
internal/mtprotoedge/exchange_test.go
View file

@ -4,6 +4,7 @@ import (
"context"
"crypto/rand"
"crypto/rsa"
"encoding/binary"
"errors"
"net"
"testing"
@ -106,6 +107,214 @@ func TestKeyExchange(t *testing.T) {
}
}
type authKeySaveContextObservation struct {
hasDeadline bool
deadline time.Time
}
type observingAuthKeyStore struct {
store.AuthKeyStore
saveContext chan authKeySaveContextObservation
}
func (s *observingAuthKeyStore) Save(ctx context.Context, key store.AuthKeyData) error {
deadline, hasDeadline := ctx.Deadline()
select {
case s.saveContext <- authKeySaveContextObservation{hasDeadline: hasDeadline, deadline: deadline}:
default:
}
return s.AuthKeyStore.Save(ctx, key)
}
type gatedAuthKeyStore struct {
store.AuthKeyStore
entered chan store.AuthKeyData
release chan struct{}
saveErr error
}
type ownershipFrameConn struct {
transport.Conn
frame []byte
}
func (c *ownershipFrameConn) Recv(_ context.Context, b *bin.Buffer) error {
b.ResetTo(c.frame)
return nil
}
func TestExchangeEncryptedReplayTransfersFrameOwnership(t *testing.T) {
backing := make([]byte, 64)
copy(backing[:8], []byte{1, 2, 3, 4, 5, 6, 7, 8})
conn := &ownershipFrameConn{frame: backing}
ex := serverExchangeCompat{conn: conn, timeout: time.Second}
var b bin.Buffer
err := ex.readUnencrypted(context.Background(), &b, &compatReqPQ{})
var encrypted *exchange.UnexpectedEncryptedError
if !errors.As(err, &encrypted) {
t.Fatalf("read encrypted frame err = %v, want UnexpectedEncryptedError", err)
}
if len(encrypted.Frame) != len(backing) || &encrypted.Frame[0] != &backing[0] {
t.Fatal("encrypted replay copied the received frame instead of transferring ownership")
}
if b.Buf != nil {
t.Fatal("exchange buffer retained transferred encrypted frame backing")
}
}
func (s *gatedAuthKeyStore) Save(ctx context.Context, key store.AuthKeyData) error {
select {
case s.entered <- key:
case <-ctx.Done():
return ctx.Err()
}
select {
case <-s.release:
case <-ctx.Done():
return ctx.Err()
}
if s.saveErr != nil {
return s.saveErr
}
return s.AuthKeyStore.Save(ctx, key)
}
// TestKeyExchangeDoesNotAcknowledgeBeforeAuthKeyCommit pins the protocol commit
// boundary: while durable Save is blocked, the client must not receive DhGenOk
// and therefore must not report a successful exchange.
func TestKeyExchangeDoesNotAcknowledgeBeforeAuthKeyCommit(t *testing.T) {
base := memory.NewAuthKeyStore()
keys := &gatedAuthKeyStore{
AuthKeyStore: base,
entered: make(chan store.AuthKeyData, 1),
release: make(chan struct{}, 1),
}
addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
conn := dialTransportOnly(t, addr)
type exchangeOutcome struct {
result exchange.ClientExchangeResult
err error
}
outcome := make(chan exchangeOutcome, 1)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
go func() {
result, err := exchange.NewExchanger(conn, 2).
WithRand(rand.Reader).
Client([]exchange.PublicKey{pub}).
Run(ctx)
outcome <- exchangeOutcome{result: result, err: err}
}()
var pending store.AuthKeyData
select {
case pending = <-keys.entered:
case <-time.After(5 * time.Second):
t.Fatal("AuthKeyStore.Save was not reached")
}
defer func() {
select {
case keys.release <- struct{}{}:
default:
}
}()
select {
case got := <-outcome:
t.Fatalf("client exchange completed before auth key commit: err=%v", got.err)
case <-time.After(150 * time.Millisecond):
}
if _, found, err := base.Get(context.Background(), pending.ID); err != nil {
t.Fatalf("Get before commit: %v", err)
} else if found {
t.Fatal("auth key became visible while durable Save was blocked")
}
keys.release <- struct{}{}
select {
case got := <-outcome:
if got.err != nil {
t.Fatalf("client exchange after commit: %v", got.err)
}
if got.result.AuthKey.ID != pending.ID {
t.Fatalf("committed auth key id = %x, client got %x", pending.ID, got.result.AuthKey.ID)
}
case <-time.After(5 * time.Second):
t.Fatal("client exchange did not finish after auth key commit")
}
if _, found, err := base.Get(context.Background(), pending.ID); err != nil {
t.Fatalf("Get after commit: %v", err)
} else if !found {
t.Fatal("auth key is not durable after successful client exchange")
}
}
// TestKeyExchangeAuthKeyCommitFailureWithholdsDhGenOk proves the failure side
// of the same invariant. The client must not observe success if storage rejects
// the key; the server closes this exchange and lets the client retry cleanly.
func TestKeyExchangeAuthKeyCommitFailureWithholdsDhGenOk(t *testing.T) {
base := memory.NewAuthKeyStore()
keys := &gatedAuthKeyStore{
AuthKeyStore: base,
entered: make(chan store.AuthKeyData, 1),
release: make(chan struct{}, 1),
saveErr: errors.New("injected auth key persistence failure"),
}
keys.release <- struct{}{}
addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
conn := dialTransportOnly(t, addr)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_, err := exchange.NewExchanger(conn, 2).
WithRand(rand.Reader).
Client([]exchange.PublicKey{pub}).
Run(ctx)
if err == nil {
t.Fatal("client exchange succeeded even though auth key commit failed")
}
select {
case attempted := <-keys.entered:
if _, found, getErr := base.Get(context.Background(), attempted.ID); getErr != nil {
t.Fatalf("Get failed key: %v", getErr)
} else if found {
t.Fatal("failed auth key commit became visible")
}
case <-time.After(time.Second):
t.Fatal("AuthKeyStore.Save was not attempted")
}
}
func TestKeyExchangeAuthKeySaveUsesHandshakeDeadline(t *testing.T) {
const handshakeMax = 10 * time.Second
observed := make(chan authKeySaveContextObservation, 1)
keys := &observingAuthKeyStore{
AuthKeyStore: memory.NewAuthKeyStore(),
saveContext: observed,
}
addr, pub, _ := startTestServer(t, Options{
DC: 2,
AuthKeys: keys,
HandshakeMaxDuration: handshakeMax,
})
_, _, _ = dialHandshake(t, addr, 2, pub)
select {
case got := <-observed:
if !got.hasDeadline {
t.Fatal("AuthKeyStore.Save context has no handshake deadline")
}
remaining := time.Until(got.deadline)
if remaining <= 0 || remaining > handshakeMax {
t.Fatalf("AuthKeyStore.Save deadline remaining = %v, want (0, %v]", remaining, handshakeMax)
}
case <-time.After(time.Second):
t.Fatal("AuthKeyStore.Save was not called")
}
}
func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})
@ -227,6 +436,82 @@ func TestKeyExchangeIgnoresUnencryptedMsgsAck(t *testing.T) {
}
}
func TestBufferedExchangePushTransfersFrameOwnershipWithoutCopy(t *testing.T) {
backing := make([]byte, 64)
for i := range backing {
backing[i] = byte(i)
}
source := &bin.Buffer{Buf: backing}
buffered := newBufferedConn(nil)
buffered.push(source)
if source.Buf != nil {
t.Fatal("push retained ownership in the source buffer")
}
var got bin.Buffer
if err := buffered.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv pending frame: %v", err)
}
if len(got.Buf) != len(backing) || &got.Buf[0] != &backing[0] {
t.Fatal("pending frame was copied instead of transferring its backing")
}
if len(buffered.pending) != 0 || cap(buffered.pending) != 0 {
t.Fatalf("consumed pending ownership retained: len=%d cap=%d", len(buffered.pending), cap(buffered.pending))
}
}
func TestBufferedExchangeLargeTrailingMsgsAckReleasesFrameBeforeNextRecv(t *testing.T) {
encodeUnencrypted := func(msg bin.Encoder, msgID int64) []byte {
var payload bin.Buffer
if err := msg.Encode(&payload); err != nil {
t.Fatalf("encode payload: %v", err)
}
var frame bin.Buffer
if err := (tgproto.UnencryptedMessage{MessageID: msgID, MessageData: payload.Raw()}).Encode(&frame); err != nil {
t.Fatalf("encode unencrypted frame: %v", err)
}
return frame.Copy()
}
intermediate := func(frame []byte) []byte {
packet := make([]byte, bin.Word+len(frame))
binary.LittleEndian.PutUint32(packet, uint32(len(frame)))
copy(packet[bin.Word:], frame)
return packet
}
// Make the ignored ack larger than the per-codec retained-buffer threshold. The following
// small req_pq frame forces bufferedConn to cross the next-Recv ownership boundary while the
// same destination bin.Buffer is reused.
ids := make([]int64, 300_000)
for i := range ids {
ids[i] = int64(i + 1)
}
ackFrame := encodeUnencrypted(&mt.MsgsAck{MsgIDs: ids}, 4)
reqFrame := encodeUnencrypted(&mt.ReqPqMultiRequest{}, 8)
packet := append(intermediate(ackFrame), intermediate(reqFrame)...)
budget := newInboundFrameBudget(2 * int64(len(ackFrame)))
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
buffered := newBufferedConn(conn)
var got bin.Buffer
if err := buffered.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv after large msgs_ack: %v", err)
}
if id, ok := unencryptedPayloadID(&got); !ok || id != mt.ReqPqMultiRequestTypeID {
t.Fatalf("returned frame type = 0x%x ok=%v, want req_pq_multi", id, ok)
}
if used, want := budget.usedBytes(), 2*int64(len(reqFrame)); used != want {
t.Fatalf("inbound budget after skipped ack = %d, want only next frame %d", used, want)
}
if cap(got.Buf) >= len(ackFrame)/2 {
t.Fatalf("large ignored ack backing retained by next frame: cap=%d ack=%d", cap(got.Buf), len(ackFrame))
}
conn.releaseInboundFrame()
if used := budget.usedBytes(); used != 0 {
t.Fatalf("inbound budget after final ownership release = %d, want 0", used)
}
}
type ackingExchangeConn struct {
transport.Conn
t *testing.T

251
internal/mtprotoedge/frame_budget.go Normal file
View file

@ -0,0 +1,251 @@
package mtprotoedge
import (
"encoding/binary"
"errors"
"fmt"
"io"
"sync/atomic"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
)
const defaultInboundFrameGlobalMaxBytes int64 = 512 << 20
var (
// ErrInboundFrameBudgetExceeded means the process-wide wire+plaintext reservation for a
// newly announced transport frame could not be acquired. The length prefix has been read,
// but the payload buffer has not been allocated and the connection must be closed.
ErrInboundFrameBudgetExceeded = errors.New("inbound frame global byte budget exceeded")
errInboundFrameCodecUnsupported = errors.New("transport codec cannot preflight inbound frame length")
errInboundFrameNotReserved = errors.New("transport codec returned a frame without reserving inbound bytes")
)
// InboundFrameBudgetedCodec is the fail-safe extension point for a custom Options.Codec.
// Implementations must parse and validate the frame length, call reserve exactly once before
// allocating or growing the payload buffer, and keep the reservation valid until Read returns.
// Built-in abridged/intermediate/padded-intermediate/full codecs are recognized directly.
type InboundFrameBudgetedCodec interface {
transport.Codec
ReadWithInboundFrameBudget(r io.Reader, b *bin.Buffer, reserve func(wireBytes, plaintextBytes int64) error) error
}
// inboundFrameBudget accounts the two per-frame buffers that can coexist while an encrypted
// request is handled: transport/wire bytes and decrypted plaintext. It deliberately charges the
// maximum plaintext size announced by framing even for an unencrypted handshake frame; that
// conservative rule makes admission independent of auth state and prevents allocation before
// auth_key_id can be inspected.
type inboundFrameBudget struct {
max int64
used atomic.Int64
}
func newInboundFrameBudget(max int64) *inboundFrameBudget {
if max <= 0 {
max = defaultInboundFrameGlobalMaxBytes
}
return &inboundFrameBudget{max: max}
}
func (b *inboundFrameBudget) reserve(wireBytes, plaintextBytes int64) (int64, error) {
return b.growReservation(0, wireBytes, plaintextBytes)
}
// growReservation atomically raises one connection's existing retained/frame reservation to
// cover a newly announced frame. Keeping the old charge until this transition is what makes a
// reused transport/plaintext backing remain accounted between frames; a small next frame cannot
// release a previously large allocation while still retaining its capacity.
func (b *inboundFrameBudget) growReservation(current, wireBytes, plaintextBytes int64) (int64, error) {
if current < 0 || wireBytes <= 0 || plaintextBytes < 0 || wireBytes > b.max || plaintextBytes > b.max-wireBytes {
return 0, fmt.Errorf("%w: wire=%d plaintext=%d limit=%d", ErrInboundFrameBudgetExceeded, wireBytes, plaintextBytes, b.max)
}
target := wireBytes + plaintextBytes
if target <= current {
return current, nil
}
n := target - current
for {
used := b.used.Load()
if n > b.max-used {
return 0, fmt.Errorf("%w: requested=%d used=%d limit=%d", ErrInboundFrameBudgetExceeded, n, used, b.max)
}
if b.used.CompareAndSwap(used, used+n) {
return target, nil
}
}
}
func (b *inboundFrameBudget) release(n int64) {
if n == 0 {
return
}
used := b.used.Add(-n)
if used < 0 {
// This is an internal ownership invariant, not recoverable input. A negative value would
// silently disable admission for subsequent frames, so fail loudly during development.
panic("mtprotoedge: inbound frame budget released more than reserved")
}
}
func (b *inboundFrameBudget) usedBytes() int64 {
return b.used.Load()
}
type inboundFrameCodecKind uint8
const (
inboundFrameCodecUnknown inboundFrameCodecKind = iota
inboundFrameCodecQuickAckAbridged
inboundFrameCodecAbridged
inboundFrameCodecIntermediate
inboundFrameCodecPaddedIntermediate
inboundFrameCodecFull
inboundFrameCodecCustom
)
func classifyInboundFrameCodec(c transport.Codec) inboundFrameCodecKind {
switch v := c.(type) {
case *quickAckAbridgedCodec:
return inboundFrameCodecQuickAckAbridged
case codec.Abridged, *codec.Abridged:
return inboundFrameCodecAbridged
case *quickAckIntermediateCodec, codec.Intermediate, *codec.Intermediate:
return inboundFrameCodecIntermediate
case *quickAckPaddedIntermediateCodec, codec.PaddedIntermediate, *codec.PaddedIntermediate:
return inboundFrameCodecPaddedIntermediate
case *codec.Full:
return inboundFrameCodecFull
case codec.NoHeader:
return classifyInboundFrameCodec(v.Codec)
case *codec.NoHeader:
if v == nil {
return inboundFrameCodecUnknown
}
return classifyInboundFrameCodec(v.Codec)
case InboundFrameBudgetedCodec:
return inboundFrameCodecCustom
default:
return inboundFrameCodecUnknown
}
}
func unwrapInboundFrameBudgetedCodec(c transport.Codec) InboundFrameBudgetedCodec {
switch v := c.(type) {
case InboundFrameBudgetedCodec:
return v
case codec.NoHeader:
return unwrapInboundFrameBudgetedCodec(v.Codec)
case *codec.NoHeader:
if v != nil {
return unwrapInboundFrameBudgetedCodec(v.Codec)
}
}
return nil
}
// inboundFramePreflightReader consumes only the framing length prefix, reserves the announced
// wire+plaintext bytes, and only then exposes the final prefix bytes to the codec. Consequently a
// budget error is observed by codec.Read before it can ResetN/Expand the payload buffer.
type inboundFramePreflightReader struct {
r io.Reader
kind inboundFrameCodecKind
reserve func(wireBytes, plaintextBytes int64) error
abridgedFirstDelivered bool
done bool
}
func (r *inboundFramePreflightReader) Read(p []byte) (int, error) {
if len(p) == 0 {
return 0, nil
}
if r.done {
return r.r.Read(p)
}
switch r.kind {
case inboundFrameCodecQuickAckAbridged:
return r.readAbridgedPrefix(p, true)
case inboundFrameCodecAbridged:
return r.readAbridgedPrefix(p, false)
case inboundFrameCodecIntermediate, inboundFrameCodecPaddedIntermediate:
return r.readWordPrefix(p, false)
case inboundFrameCodecFull:
return r.readWordPrefix(p, true)
default:
return 0, errInboundFrameCodecUnsupported
}
}
func (r *inboundFramePreflightReader) readAbridgedPrefix(p []byte, quickAck bool) (int, error) {
if !r.abridgedFirstDelivered {
var first [1]byte
if _, err := io.ReadFull(r.r, first[:]); err != nil {
return 0, err
}
lengthByte := first[0]
extended := lengthByte >= 0x7f
if quickAck {
lengthByte &= 0x7f
extended = lengthByte == 0x7f
}
if !extended {
n := int64(lengthByte) * bin.Word
if err := reserveCompatFrame(r.reserve, n, n); err != nil {
return 0, err
}
r.done = true
}
r.abridgedFirstDelivered = true
p[0] = first[0]
return 1, nil
}
var tail [3]byte
if _, err := io.ReadFull(r.r, tail[:]); err != nil {
return 0, err
}
words := uint32(tail[0]) | uint32(tail[1])<<8 | uint32(tail[2])<<16
n := int64(words) * bin.Word
if err := reserveCompatFrame(r.reserve, n, n); err != nil {
return 0, err
}
r.done = true
return copy(p, tail[:]), nil
}
func (r *inboundFramePreflightReader) readWordPrefix(p []byte, full bool) (int, error) {
var header [bin.Word]byte
if _, err := io.ReadFull(r.r, header[:]); err != nil {
return 0, err
}
raw := int64(binary.LittleEndian.Uint32(header[:]))
var wireBytes, plaintextBytes int64
if full {
// Full transport length includes length + sequence + payload + CRC.
if raw < 3*bin.Word || raw > maxTransportMessageSize {
return 0, fmt.Errorf("invalid full transport message length %d", raw)
}
wireBytes = raw
plaintextBytes = raw - 3*bin.Word
} else {
wireBytes = raw &^ int64(quickAckResponseFlag)
plaintextBytes = wireBytes
}
if err := reserveCompatFrame(r.reserve, wireBytes, plaintextBytes); err != nil {
return 0, err
}
r.done = true
return copy(p, header[:]), nil
}
func reserveCompatFrame(reserve func(wireBytes, plaintextBytes int64) error, wireBytes, plaintextBytes int64) error {
if wireBytes <= 0 || wireBytes > maxTransportMessageSize {
return fmt.Errorf("invalid transport message length %d", wireBytes)
}
return reserve(wireBytes, plaintextBytes)
}

337
internal/mtprotoedge/frame_budget_test.go Normal file
View file

@ -0,0 +1,337 @@
package mtprotoedge
import (
"bytes"
"context"
"encoding/binary"
"errors"
"io"
"net"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
)
type frameBudgetTestConn struct {
reader bytes.Reader
read int
closed bool
}
func newFrameBudgetTestConn(packet []byte) *frameBudgetTestConn {
c := &frameBudgetTestConn{}
c.reader.Reset(packet)
return c
}
func (c *frameBudgetTestConn) Read(p []byte) (int, error) {
n, err := c.reader.Read(p)
c.read += n
return n, err
}
func (*frameBudgetTestConn) Write(p []byte) (int, error) { return len(p), nil }
func (c *frameBudgetTestConn) Close() error {
c.closed = true
return nil
}
func (*frameBudgetTestConn) LocalAddr() net.Addr { return frameBudgetTestAddr("local") }
func (*frameBudgetTestConn) RemoteAddr() net.Addr { return frameBudgetTestAddr("remote") }
func (*frameBudgetTestConn) SetDeadline(time.Time) error { return nil }
func (*frameBudgetTestConn) SetReadDeadline(time.Time) error { return nil }
func (*frameBudgetTestConn) SetWriteDeadline(time.Time) error { return nil }
type frameBudgetTestAddr string
func (a frameBudgetTestAddr) Network() string { return "frame-budget-test" }
func (a frameBudgetTestAddr) String() string { return string(a) }
func newFrameBudgetTestTransport(packet []byte, c transport.Codec, budget *inboundFrameBudget) (*compatTransportConn, *frameBudgetTestConn) {
raw := newFrameBudgetTestConn(packet)
return &compatTransportConn{conn: raw, codec: c, budget: budget}, raw
}
func TestInboundFrameBudgetSupportsBuiltInCodecs(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
abridged := append([]byte{byte(len(payload) / bin.Word)}, payload...)
intermediate := make([]byte, bin.Word+len(payload))
binary.LittleEndian.PutUint32(intermediate, uint32(len(payload)))
copy(intermediate[bin.Word:], payload)
padded := make([]byte, bin.Word+len(payload)+1)
binary.LittleEndian.PutUint32(padded, uint32(len(payload)+1))
copy(padded[bin.Word:], payload)
padded[len(padded)-1] = 0xa5
var full bytes.Buffer
fullCodec := &codec.Full{}
fullPayload := &bin.Buffer{Buf: append([]byte(nil), payload...)}
if err := fullCodec.Write(&full, fullPayload); err != nil {
t.Fatalf("encode full frame: %v", err)
}
tests := []struct {
name string
packet []byte
codec transport.Codec
reservation int64
}{
{name: "abridged", packet: abridged, codec: &quickAckAbridgedCodec{}, reservation: 2 * int64(len(payload))},
{name: "intermediate", packet: intermediate, codec: &quickAckIntermediateCodec{}, reservation: 2 * int64(len(payload))},
{name: "padded_intermediate", packet: padded, codec: &quickAckPaddedIntermediateCodec{}, reservation: 2 * int64(len(payload)+1)},
{name: "full", packet: full.Bytes(), codec: &codec.Full{}, reservation: int64(full.Len() + len(payload))},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
budget := newInboundFrameBudget(tt.reservation)
conn, _ := newFrameBudgetTestTransport(tt.packet, tt.codec, budget)
var got bin.Buffer
if err := conn.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv: %v", err)
}
if !bytes.Equal(got.Raw(), payload) {
t.Fatalf("payload = %x, want %x", got.Raw(), payload)
}
if used := budget.usedBytes(); used != tt.reservation {
t.Fatalf("held budget = %d, want %d", used, tt.reservation)
}
if err := conn.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
if used := budget.usedBytes(); used != tt.reservation {
t.Fatalf("budget after concurrent Close = %d, want delivered ownership %d", used, tt.reservation)
}
conn.releaseInboundFrame()
if used := budget.usedBytes(); used != 0 {
t.Fatalf("budget after ownership release = %d, want 0", used)
}
})
}
}
func TestInboundFrameBudgetRejectsBeforePayloadAllocation(t *testing.T) {
const payloadBytes = 1 << 20
var header [bin.Word]byte
binary.LittleEndian.PutUint32(header[:], payloadBytes)
budget := newInboundFrameBudget(2*payloadBytes - 1)
conn, raw := newFrameBudgetTestTransport(header[:], &quickAckIntermediateCodec{}, budget)
var got bin.Buffer
err := conn.Recv(context.Background(), &got)
if !errors.Is(err, ErrInboundFrameBudgetExceeded) {
t.Fatalf("Recv error = %v, want ErrInboundFrameBudgetExceeded", err)
}
if raw.read != bin.Word {
t.Fatalf("wire bytes read = %d, want only %d-byte length prefix", raw.read, bin.Word)
}
if cap(got.Buf) != 0 {
t.Fatalf("payload buffer capacity = %d, want 0 before admission", cap(got.Buf))
}
if used := budget.usedBytes(); used != 0 {
t.Fatalf("budget after rejected preflight = %d, want 0", used)
}
}
func TestInboundFrameBudgetAbridgedPreflightMatchesCodecSemantics(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
quickPacket := append([]byte{0x80 | byte(len(payload)/bin.Word)}, payload...)
quickBudget := newInboundFrameBudget(int64(2 * len(payload)))
quick, _ := newFrameBudgetTestTransport(quickPacket, &quickAckAbridgedCodec{}, quickBudget)
var got bin.Buffer
if err := quick.Recv(context.Background(), &got); err != nil {
t.Fatalf("quick-ack abridged Recv: %v", err)
}
requested := quick.ConsumeQuickAckRequested()
if !bytes.Equal(got.Raw(), payload) || !requested {
t.Fatalf("quick-ack frame = %x requested=%v", got.Raw(), requested)
}
quick.releaseInboundFrame()
_ = quick.Close()
// gotd's plain codec treats every first byte >= 0x7f as the extended form (it does not
// implement the quick-ack high bit). The preflight parser must mirror that behavior; treating
// 0x82 as a short two-word frame would let the codec allocate from the following three bytes.
malicious := []byte{0x82, 0xff, 0xff, 0xff}
plainBudget := newInboundFrameBudget(defaultInboundFrameGlobalMaxBytes)
plain, raw := newFrameBudgetTestTransport(malicious, codec.Abridged{}, plainBudget)
got.Reset()
err := plain.Recv(context.Background(), &got)
if err == nil {
t.Fatal("plain abridged accepted oversized extended length")
}
if raw.read != 4 || cap(got.Buf) > 2*bin.Word {
t.Fatalf("plain abridged read=%d buffer_cap=%d, want prefix-only allocation", raw.read, cap(got.Buf))
}
_ = plain.Close()
}
func TestInboundFrameBudgetReleasedAtNextRecvAndReusable(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
frame := make([]byte, bin.Word+len(payload))
binary.LittleEndian.PutUint32(frame, uint32(len(payload)))
copy(frame[bin.Word:], payload)
packet := append(append([]byte(nil), frame...), frame...)
reservation := int64(2 * len(payload))
budget := newInboundFrameBudget(reservation)
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
for i := 0; i < 2; i++ {
var got bin.Buffer
if err := conn.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv %d: %v", i+1, err)
}
if used := budget.usedBytes(); used != reservation {
t.Fatalf("held budget after frame %d = %d, want %d", i+1, used, reservation)
}
}
conn.releaseInboundFrame()
_ = conn.Close()
}
func TestInboundFrameRetainedBackingStaysChargedAcrossSmallFrame(t *testing.T) {
const largeBytes = 1 << 20
large := make([]byte, bin.Word+largeBytes)
binary.LittleEndian.PutUint32(large, largeBytes)
smallPayload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
small := make([]byte, bin.Word+len(smallPayload))
binary.LittleEndian.PutUint32(small, uint32(len(smallPayload)))
copy(small[bin.Word:], smallPayload)
packet := append(large, small...)
budget := newInboundFrameBudget(2 * largeBytes)
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
var wire bin.Buffer
if err := conn.Recv(context.Background(), &wire); err != nil {
t.Fatalf("large Recv: %v", err)
}
// Model decryptClientFrame's exact-size plaintext reuse buffer.
plain := bin.Buffer{Buf: make([]byte, largeBytes)}
retainInboundFrameBackings(conn, &wire, &plain)
retained := int64(cap(wire.Buf) + cap(plain.Buf))
if got := budget.usedBytes(); got != retained {
t.Fatalf("retained budget after large frame = %d, want capacities %d", got, retained)
}
wire.Reset()
if err := conn.Recv(context.Background(), &wire); err != nil {
t.Fatalf("small Recv: %v", err)
}
// The small announcement must not release the large backing's charge. This was the
// warm-many-connections bypass: each socket retained MiBs while the global budget saw bytes.
if got := budget.usedBytes(); got != retained {
t.Fatalf("budget after small frame = %d, want retained high-water %d", got, retained)
}
wire.Buf = nil
plain.Buf = nil
retainInboundFrameBackings(conn, &wire, &plain)
if got := budget.usedBytes(); got != 0 {
t.Fatalf("budget after dropping reusable backings = %d, want 0", got)
}
_ = conn.Close()
}
func TestInboundFrameBudgetClosePreservesDeliveredOwnershipUntilRelease(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
frame := make([]byte, bin.Word+len(payload))
binary.LittleEndian.PutUint32(frame, uint32(len(payload)))
copy(frame[bin.Word:], payload)
budget := newInboundFrameBudget(int64(2 * len(payload)))
first, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
var got bin.Buffer
if err := first.Recv(context.Background(), &got); err != nil {
t.Fatalf("first Recv: %v", err)
}
blocked, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
var blockedPayload bin.Buffer
if err := blocked.Recv(context.Background(), &blockedPayload); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
t.Fatalf("concurrent Recv error = %v, want global budget rejection", err)
}
if cap(blockedPayload.Buf) != 0 {
t.Fatalf("blocked connection allocated payload capacity %d", cap(blockedPayload.Buf))
}
_ = blocked.Close()
if err := first.Close(); err != nil {
t.Fatalf("first Close: %v", err)
}
if used := budget.usedBytes(); used != int64(2*len(payload)) {
t.Fatalf("budget after concurrent Close = %d, want delivered frame still charged", used)
}
second, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
got.Reset()
if err := second.Recv(context.Background(), &got); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
t.Fatalf("second Recv before ownership release = %v, want budget rejection", err)
}
_ = second.Close()
first.releaseInboundFrame()
third, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
got.Reset()
if err := third.Recv(context.Background(), &got); err != nil {
t.Fatalf("third Recv after ownership release: %v", err)
}
third.releaseInboundFrame()
_ = third.Close()
}
type unsafeFrameBudgetCodec struct {
readCalled bool
}
func (*unsafeFrameBudgetCodec) WriteHeader(io.Writer) error { return nil }
func (*unsafeFrameBudgetCodec) ReadHeader(io.Reader) error { return nil }
func (*unsafeFrameBudgetCodec) Write(io.Writer, *bin.Buffer) error { return nil }
func (c *unsafeFrameBudgetCodec) Read(io.Reader, *bin.Buffer) error { c.readCalled = true; return nil }
func TestCustomCodecWithoutPreflightFailsClosed(t *testing.T) {
raw := newFrameBudgetTestConn([]byte{1, 2, 3, 4})
listener := newSingleConnListener(raw)
custom := &unsafeFrameBudgetCodec{}
budgeted := newCompatTransportListener(func() transport.Codec { return custom }, listener, newInboundFrameBudget(1024))
conn, err := budgeted.Accept()
if !errors.Is(err, errInboundFrameCodecUnsupported) {
t.Fatalf("Accept error = %v, want unsupported preflight codec", err)
}
if conn != nil {
t.Fatal("unsupported custom codec unexpectedly accepted")
}
if custom.readCalled || raw.read != 0 {
t.Fatalf("custom codec touched frame before rejection: read_called=%v wire_read=%d", custom.readCalled, raw.read)
}
}
func TestExplicitBuiltInCodecUsesBudgetedTransport(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
packet := append([]byte(nil), codec.IntermediateClientStart[:]...)
var header [bin.Word]byte
binary.LittleEndian.PutUint32(header[:], uint32(len(payload)))
packet = append(packet, header[:]...)
packet = append(packet, payload...)
raw := newFrameBudgetTestConn(packet)
budget := newInboundFrameBudget(int64(2 * len(payload)))
listener := newCompatTransportListener(
func() transport.Codec { return codec.Intermediate{} },
newSingleConnListener(raw),
budget,
)
conn, err := listener.Accept()
if err != nil {
t.Fatalf("Accept: %v", err)
}
var got bin.Buffer
if err := conn.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv: %v", err)
}
if !bytes.Equal(got.Raw(), payload) || budget.usedBytes() != int64(2*len(payload)) {
t.Fatalf("payload=%x budget=%d", got.Raw(), budget.usedBytes())
}
conn.(*compatTransportConn).releaseInboundFrame()
_ = conn.Close()
}

813
internal/mtprotoedge/inbound_rpc.go
View file

@ -1,30 +1,337 @@
package mtprotoedge
import (
"container/list"
"context"
"errors"
"sync"
"sync/atomic"
"time"
)
// ErrInboundRPCQueueFull 表示单连接 RPC 队列已满。
// ErrInboundRPCQueueFull 表示 inbound RPC 已触达单连接或进程级预算。
var ErrInboundRPCQueueFull = errors.New("inbound rpc queue full")
// maxInflightRPCBytes 是单连接已入队未完成 inbound RPC body 的总字节上限。
// 队列除按条数(queueSize)限制外，再按字节预算兜底：对抗客户端发满大请求时按字节先拒绝。
// maxInflightRPCBytes 是单连接所有已预留、排队和执行中 RPC body 的总字节上限。
// 进程级预算在 Copy 前先兜底；这里再隔离单个连接，避免一个客户端独占全局内存。
const maxInflightRPCBytes = 32 << 20 // 32 MiB
// rpcCloseWaitTimeout 是连接关闭时等待 inbound RPC worker 退出的上限。
// rpcCloseWaitTimeout 是连接/Server 关闭时等待在途 RPC 或共享 worker 退出的上限。
const rpcCloseWaitTimeout = 5 * time.Second
type inboundRPC struct {
ctx context.Context
method string
enqueuedAt time.Time
size int
run func(context.Context) error
ctx context.Context
cancel context.CancelFunc
stopRoot func() bool
stopTimeout func() bool
method string
enqueuedAt time.Time
deadline time.Time
size int
run func(context.Context) error
onTimeout func()
budget *inboundRPCGlobalReservation
ticket *inboundRPCTicket
}
func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time.Duration) {
const (
inboundRPCTicketQueued int32 = iota
inboundRPCTicketRunning
inboundRPCTicketDone
)
type inboundRPCTicket struct {
state atomic.Int32
onTimeout func()
}
// inboundRPCScheduler 是 Server 级共享调度器。ready 中每个 Conn 最多只有一个有效令牌；
// worker 每次只从该连接取一条，再把仍可运行的连接放回队尾，因此单个热点连接不能长期
// 占住共享池。worker 在首条任务到达后才创建，空闲 Server 不预起 256 个 goroutine。
type inboundRPCScheduler struct {
workers int
maxTasks int
maxBytes int64
// ready is an intrusive scheduler-owned queue rather than a bounded channel. A connection
// has at most one element, and close removes that element in O(1). This prevents closed-Conn
// stale tokens from filling a channel and making every worker block while trying to reschedule.
readyMu sync.Mutex
ready *list.List
readyIndex map[*Conn]*list.Element
readyWake chan struct{}
stopCh chan struct{}
lifecycleMu sync.Mutex
started bool
stopped bool
workersStarted bool
workerWG sync.WaitGroup
budgetMu sync.Mutex
tasks int
bytes int64
}
type inboundRPCGlobalReservation struct {
scheduler *inboundRPCScheduler
size int64
once sync.Once
}
// inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次；
// 无论 Copy 后连接关闭、入队成功还是调用方提前返回，预算都有唯一归还路径。
type inboundRPCReservation struct {
conn *Conn
global *inboundRPCGlobalReservation
ctx context.Context
method string
size int
enqueuedAt time.Time
deadline time.Time
once sync.Once
}
func newInboundRPCScheduler(workers, maxTasks int, maxBytes int64) *inboundRPCScheduler {
if workers <= 0 {
workers = 1
}
if maxTasks <= 0 {
maxTasks = 1
}
if maxBytes <= 0 {
maxBytes = 1
}
return &inboundRPCScheduler{
workers: workers,
maxTasks: maxTasks,
maxBytes: maxBytes,
ready: list.New(),
readyIndex: make(map[*Conn]*list.Element),
readyWake: make(chan struct{}, 1),
stopCh: make(chan struct{}),
}
}
// start 允许共享池开始消费。已在 start 前进入 ready 的任务会保留顺序，便于启动突发，
// 也使测试能够确定性验证轮转公平性。
func (s *inboundRPCScheduler) start() {
s.lifecycleMu.Lock()
if s.stopped {
s.lifecycleMu.Unlock()
return
}
s.started = true
shouldStart := s.readyLen() > 0
s.lifecycleMu.Unlock()
if shouldStart {
s.ensureWorkers()
}
}
func (s *inboundRPCScheduler) ensureWorkers() {
s.lifecycleMu.Lock()
defer s.lifecycleMu.Unlock()
if !s.started || s.stopped || s.workersStarted {
return
}
s.workersStarted = true
s.workerWG.Add(s.workers)
for i := 0; i < s.workers; i++ {
go s.worker()
}
}
func (s *inboundRPCScheduler) stop(timeout time.Duration) {
s.lifecycleMu.Lock()
if !s.stopped {
s.stopped = true
s.budgetMu.Lock()
// 与 reserveGlobal 在同一把锁下切断新任务；已持有 reservation 的任务仍由
// 对应 Conn 的 commit/abort/close 路径精确归还。
close(s.stopCh)
s.budgetMu.Unlock()
}
s.lifecycleMu.Unlock()
done := make(chan struct{})
go func() {
s.workerWG.Wait()
close(done)
}()
if timeout <= 0 {
<-done
return
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-done:
case <-timer.C:
}
}
func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservation, string, error) {
if size < 0 {
size = 0
}
size64 := int64(size)
s.budgetMu.Lock()
defer s.budgetMu.Unlock()
select {
case <-s.stopCh:
return nil, "scheduler_closed", ErrConnClosed
default:
}
if s.tasks >= s.maxTasks {
return nil, "global_task_budget", ErrInboundRPCQueueFull
}
// 用减法比较避免 s.bytes+size64 溢出。
if size64 > s.maxBytes-s.bytes {
return nil, "global_byte_budget", ErrInboundRPCQueueFull
}
s.tasks++
s.bytes += size64
return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil
}
func (r *inboundRPCGlobalReservation) release() {
if r == nil || r.scheduler == nil {
return
}
r.once.Do(func() {
s := r.scheduler
s.budgetMu.Lock()
s.tasks--
s.bytes -= r.size
s.budgetMu.Unlock()
})
}
func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) {
s.budgetMu.Lock()
defer s.budgetMu.Unlock()
return s.tasks, s.bytes
}
func (s *inboundRPCScheduler) schedule(c *Conn) {
if s == nil || c == nil {
return
}
// rpcReady/rpcClosed and queue membership must be tested/installed while holding rpcMu.
// Otherwise close can remove the old token between the test and enqueue, leaving a new stale
// token behind after the connection is already terminal.
c.rpcMu.Lock()
eligible := c.rpcReady && !c.rpcClosed
added := false
if eligible {
added = s.enqueueReady(c)
}
c.rpcMu.Unlock()
if !added {
return
}
s.signalReady()
s.ensureWorkers()
}
func (s *inboundRPCScheduler) worker() {
defer s.workerWG.Done()
for {
select {
case <-s.stopCh:
return
default:
}
if c := s.popReady(); c != nil {
task, ok, reschedule := c.takeInboundRPC()
if reschedule {
s.schedule(c)
}
if ok {
c.runInboundRPC(task)
}
continue
}
select {
case <-s.readyWake:
case <-s.stopCh:
return
}
}
}
func (s *inboundRPCScheduler) enqueueReady(c *Conn) bool {
select {
case <-s.stopCh:
return false
default:
}
s.readyMu.Lock()
defer s.readyMu.Unlock()
select {
case <-s.stopCh:
return false
default:
}
if _, exists := s.readyIndex[c]; exists {
return false
}
s.readyIndex[c] = s.ready.PushBack(c)
return true
}
func (s *inboundRPCScheduler) popReady() *Conn {
s.readyMu.Lock()
front := s.ready.Front()
if front == nil {
s.readyMu.Unlock()
return nil
}
c, _ := front.Value.(*Conn)
s.ready.Remove(front)
delete(s.readyIndex, c)
hasMore := s.ready.Len() > 0
s.readyMu.Unlock()
if hasMore {
// Wake another worker while this worker begins the task. A capacity-one wake channel is
// sufficient: every pop cascades another wake until the queue is drained.
s.signalReady()
}
return c
}
func (s *inboundRPCScheduler) unschedule(c *Conn) {
if s == nil || c == nil {
return
}
s.readyMu.Lock()
if el := s.readyIndex[c]; el != nil {
s.ready.Remove(el)
delete(s.readyIndex, c)
}
hasMore := s.ready.Len() > 0
s.readyMu.Unlock()
if hasMore {
s.signalReady()
}
}
func (s *inboundRPCScheduler) readyLen() int {
s.readyMu.Lock()
defer s.readyMu.Unlock()
return s.ready.Len()
}
func (s *inboundRPCScheduler) signalReady() {
select {
case s.readyWake <- struct{}{}:
default:
}
}
func (c *Conn) startInboundRPCScheduler(scheduler *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) {
if c.metrics == nil {
c.metrics = NopMetrics{}
}
@ -35,163 +342,413 @@ func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time
queueSize = 1
}
rootCtx, cancel := context.WithCancel(context.Background())
c.rpcQueue = make(chan inboundRPC, queueSize)
c.rpcStop = make(chan struct{})
c.rpcScheduler = scheduler
c.rpcCancel = cancel
c.rpcTimeout = timeout
c.rpcRootCtx = rootCtx
c.rpcMaxInflight = maxInflight
// worker 懒启动：不在此处起 worker；首个 RPC 入队时由 ensureInboundRPCWorkers 起，
// 避免握手后静默 / 纯推送目标连接白白钉住 maxInflight 个 goroutine。
c.rpcQueueSize = queueSize
// rpcQueue 保持 nil；首个成功 commit 才由 append 分配，静默连接零队列内存。
}
// ensureInboundRPCWorkers 懒启动 maxInflight 个 RPC worker（仅一次），在 enqueueInboundRPC
// 入队成功后调用。从不发 RPC 的连接（半开 / 纯推送）由此完全不起 worker。
func (c *Conn) ensureInboundRPCWorkers() {
c.rpcWorkersOnce.Do(func() {
c.rpcWG.Add(c.rpcMaxInflight)
for i := 0; i < c.rpcMaxInflight; i++ {
go c.inboundRPCWorker(c.rpcRootCtx)
// reserveInboundRPC 必须在 request body Copy 前调用。它先拿进程级条数/字节预算，
// 再预占单连接队列槽和字节预算；commit 或 abort 负责唯一释放。
func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (*inboundRPCReservation, error) {
if ctx == nil {
ctx = context.Background()
}
select {
case <-ctx.Done():
c.metrics.InboundRPCDropped(method, "context_done")
return nil, ctx.Err()
default:
}
if c.rpcScheduler == nil {
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
}
global, reason, err := c.rpcScheduler.reserveGlobal(size)
if err != nil {
c.metrics.InboundRPCDropped(method, reason)
return nil, err
}
now := time.Now()
deadline := time.Time{}
if c.rpcTimeout > 0 {
deadline = now.Add(c.rpcTimeout)
}
if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
deadline = ctxDeadline
}
if size < 0 {
size = 0
}
c.rpcMu.Lock()
if err := ctx.Err(); err != nil {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "context_done")
return nil, err
}
if c.rpcClosed {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
}
if c.rpcReserved+len(c.rpcQueue) >= c.rpcQueueSize {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "queue_full")
return nil, ErrInboundRPCQueueFull
}
if int64(size) > maxInflightRPCBytes-c.inflightRPCBytes.Load() {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "byte_budget")
return nil, ErrInboundRPCQueueFull
}
c.rpcReserved++
c.inflightRPCBytes.Add(int64(size))
// Add 与 close 的 Wait 由 rpcMu 排序：close 置 rpcClosed 后不会再发生 Add。
c.rpcReservationWG.Add(1)
c.rpcMu.Unlock()
return &inboundRPCReservation{
conn: c,
global: global,
ctx: ctx,
method: method,
size: size,
enqueuedAt: now,
deadline: deadline,
}, nil
}
// enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用
// reserveInboundRPC -> Copy -> commit，保证真正的 Copy 前预算。
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
reservation, err := c.reserveInboundRPC(ctx, task.method, task.size)
if err != nil {
return err
}
defer reservation.abort()
return reservation.commit(task)
}
func (r *inboundRPCReservation) commit(task inboundRPC) error {
result := ErrConnClosed
var (
committed bool
reschedule bool
queueLen int
queueCap int
)
r.once.Do(func() {
c := r.conn
c.rpcMu.Lock()
c.rpcReserved--
if c.rpcClosed {
c.inflightRPCBytes.Add(-int64(r.size))
} else {
// The request deadline starts when admission succeeds, not when a worker
// eventually dequeues the request. This bounds total queue + execution
// latency and lets a queued request emit its explicit timeout on time.
if r.deadline.IsZero() {
task.ctx, task.cancel = context.WithCancel(r.ctx)
} else {
task.ctx, task.cancel = context.WithDeadline(r.ctx, r.deadline)
}
task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel)
task.method = r.method
task.enqueuedAt = r.enqueuedAt
task.deadline = r.deadline
task.size = r.size
task.budget = r.global
ticket := &inboundRPCTicket{}
if task.onTimeout != nil {
onTimeout := task.onTimeout
var timeoutOnce sync.Once
ticket.onTimeout = func() {
timeoutOnce.Do(onTimeout)
}
task.onTimeout = ticket.onTimeout
}
task.ticket = ticket
if task.onTimeout != nil && !task.deadline.IsZero() {
taskCtx := task.ctx
task.stopTimeout = context.AfterFunc(taskCtx, func() {
if errors.Is(taskCtx.Err(), context.DeadlineExceeded) {
c.expireInboundRPCTicket(ticket)
}
})
}
c.rpcQueue = append(c.rpcQueue, task)
queueLen = len(c.rpcQueue)
queueCap = c.rpcQueueSize
if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
c.rpcReady = true
reschedule = true
}
committed = true
result = nil
}
c.rpcMu.Unlock()
c.rpcReservationWG.Done()
if !committed {
r.global.release()
}
})
if committed {
r.conn.metrics.InboundRPCQueued(r.method, queueLen, queueCap)
if reschedule {
r.conn.rpcScheduler.schedule(r.conn)
}
}
return result
}
func (r *inboundRPCReservation) abort() {
if r == nil {
return
}
r.once.Do(func() {
c := r.conn
c.rpcMu.Lock()
c.rpcReserved--
c.inflightRPCBytes.Add(-int64(r.size))
c.rpcMu.Unlock()
c.rpcReservationWG.Done()
r.global.release()
})
}
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
if ctx == nil {
ctx = context.Background()
func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) {
c.rpcMu.Lock()
defer c.rpcMu.Unlock()
// ready token 是可替代的：收到一个 token 就消费当前“已调度”状态。关闭后或
// 已被另一 token 抢先处理时，这只是一个无害 stale token。
if !c.rpcReady {
return inboundRPC{}, false, false
}
if c.rpcQueue == nil || c.rpcStop == nil {
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
return ErrConnClosed
c.rpcReady = false
if c.rpcClosed || len(c.rpcQueue) == 0 || c.rpcRunning >= c.rpcMaxInflight {
return inboundRPC{}, false, false
}
task.ctx = ctx
task.enqueuedAt = time.Now()
select {
case <-ctx.Done():
task = c.rpcQueue[0]
c.rpcQueue[0] = inboundRPC{}
c.rpcQueue = c.rpcQueue[1:]
if len(c.rpcQueue) == 0 {
c.rpcQueue = nil
}
c.rpcRunning++
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketRunning)
}
c.rpcWG.Add(1)
if len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight {
c.rpcReady = true
reschedule = true
}
return task, true, reschedule
}
func (c *Conn) runInboundRPC(task inboundRPC) {
defer c.finishInboundRPC(task)
now := time.Now()
ctxErr := task.ctx.Err()
if (!task.deadline.IsZero() && !now.Before(task.deadline)) || errors.Is(ctxErr, context.DeadlineExceeded) {
c.metrics.InboundRPCDropped(task.method, "queue_timeout")
if task.onTimeout != nil {
task.onTimeout()
}
return
}
if ctxErr != nil {
c.metrics.InboundRPCDropped(task.method, "context_done")
return ctx.Err()
case <-c.rpcStop:
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
return ErrConnClosed
default:
return
}
// 字节预算：先预扣 size，超 maxInflightRPCBytes 则回滚并拒绝（与条数上限并列的第二道闸）。
if task.size > 0 {
if c.inflightRPCBytes.Add(int64(task.size)) > maxInflightRPCBytes {
c.inflightRPCBytes.Add(-int64(task.size))
c.metrics.InboundRPCDropped(task.method, "byte_budget")
return ErrInboundRPCQueueFull
}
}
select {
case c.rpcQueue <- task:
c.ensureInboundRPCWorkers()
c.metrics.InboundRPCQueued(task.method, len(c.rpcQueue), cap(c.rpcQueue))
return nil
case <-ctx.Done():
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "context_done")
return ctx.Err()
case <-c.rpcStop:
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
return ErrConnClosed
default:
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "queue_full")
return ErrInboundRPCQueueFull
}
}
// releaseInflightRPCBytes 归还字节预算。与 enqueueInboundRPC 的预扣严格配对：
// 入队失败时回滚、worker 执行完(runInboundRPC)或排空丢弃(drainInboundRPCQueue)时释放。
func (c *Conn) releaseInflightRPCBytes(size int) {
if size > 0 {
c.inflightRPCBytes.Add(-int64(size))
}
}
func (c *Conn) inboundRPCWorker(rootCtx context.Context) {
defer c.rpcWG.Done()
for {
select {
case <-c.rpcStop:
return
default:
}
select {
case task := <-c.rpcQueue:
c.runInboundRPC(rootCtx, task)
case <-c.rpcStop:
return
}
}
}
func (c *Conn) runInboundRPC(rootCtx context.Context, task inboundRPC) {
defer c.releaseInflightRPCBytes(task.size)
queueWait := time.Since(task.enqueuedAt)
c.metrics.InboundRPCStarted(task.method, queueWait)
c.metrics.InboundRPCStarted(task.method, now.Sub(task.enqueuedAt))
ctx := task.ctx
if ctx == nil {
ctx = context.Background()
if task.run != nil {
_ = task.run(ctx)
}
// 合并两个取消源（task.ctx 与 rootCtx）+ 超时为最少的 context 层数：
// WithTimeout/WithCancel 的 cancel 直接作为 AfterFunc 回调，省掉单独的中间层。
var cancel context.CancelFunc
if c.rpcTimeout > 0 {
ctx, cancel = context.WithTimeout(ctx, c.rpcTimeout)
} else {
ctx, cancel = context.WithCancel(ctx)
}
func (c *Conn) finishInboundRPC(task inboundRPC) {
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketDone)
}
stopInboundRPCTask(task)
var reschedule bool
c.rpcMu.Lock()
c.rpcRunning--
c.inflightRPCBytes.Add(-int64(task.size))
if !c.rpcClosed && len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
c.rpcReady = true
reschedule = true
}
c.rpcMu.Unlock()
reservation := task.budget
// The scheduler budget may be reused immediately after release. Clear request-owned
// closures/context references first so slow metrics/rescheduling cannot overlap the old body
// with a newly admitted body under the same byte accounting.
task = inboundRPC{}
reservation.release()
c.rpcWG.Done()
if reschedule {
c.rpcScheduler.schedule(c)
}
}
// expireInboundRPCTicket removes a request that is still queued and returns its
// memory/task reservations immediately. If the worker won the dequeue race, the
// same callback only signals the running request's response gate; its body remains
// owned until the handler exits.
func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
if ticket == nil {
return
}
var (
task inboundRPC
found bool
unschedule bool
)
c.rpcMu.Lock()
for i := range c.rpcQueue {
if c.rpcQueue[i].ticket != ticket {
continue
}
task = c.rpcQueue[i]
copy(c.rpcQueue[i:], c.rpcQueue[i+1:])
last := len(c.rpcQueue) - 1
c.rpcQueue[last] = inboundRPC{}
c.rpcQueue = c.rpcQueue[:last]
if len(c.rpcQueue) == 0 {
c.rpcQueue = nil
if c.rpcReady {
c.rpcReady = false
unschedule = true
}
}
c.inflightRPCBytes.Add(-int64(task.size))
ticket.state.Store(inboundRPCTicketDone)
found = true
break
}
c.rpcMu.Unlock()
if unschedule {
c.rpcScheduler.unschedule(c)
}
if found {
method := task.method
reservation := task.budget
stopInboundRPCTask(task)
// Drop the run/context closures before returning the byte reservation. Otherwise an
// onTimeout callback that blocks or performs a slow write can keep the copied request body
// reachable after the global scheduler has advertised those bytes as available again.
task = inboundRPC{}
reservation.release()
c.metrics.InboundRPCDropped(method, "queue_timeout")
if ticket.onTimeout != nil {
ticket.onTimeout()
}
return
}
if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil {
ticket.onTimeout()
}
}
// stopInboundRPCTask disarms callbacks before canceling the context so a normal
// completion or connection close cannot manufacture an RPC_TIMEOUT response.
// A deadline callback already in flight is harmless because enqueueRPC's response
// gate makes timeout and normal rpc_result mutually exclusive.
func stopInboundRPCTask(task inboundRPC) {
if task.stopTimeout != nil {
task.stopTimeout()
}
if task.stopRoot != nil {
task.stopRoot()
}
if task.cancel != nil {
task.cancel()
}
defer cancel()
stopRoot := context.AfterFunc(rootCtx, cancel)
defer stopRoot()
_ = task.run(ctx)
}
func (c *Conn) closeInboundRPCScheduler() {
if c.rpcStop == nil {
c.beginCloseInboundRPCScheduler()
if c.rpcScheduler == nil {
return
}
c.waitInboundShutdown(rpcCloseWaitTimeout)
}
// beginCloseInboundRPCScheduler publishes closure, cancels running work and releases queued
// requests without waiting for handlers. ForceClose uses this phase before transport.Close so a
// pathological/blocking transport implementation cannot leave the RPC admission gate open.
func (c *Conn) beginCloseInboundRPCScheduler() {
if c.rpcScheduler == nil {
return
}
c.rpcClose.Do(func() {
c.rpcMu.Lock()
c.rpcClosed = true
c.rpcReady = false
queued := c.rpcQueue
c.rpcQueue = nil
for i := range queued {
c.inflightRPCBytes.Add(-int64(queued[i].size))
}
c.rpcMu.Unlock()
// Remove the scheduler-owned token after rpcClosed/rpcReady become visible. schedule()
// takes rpcMu while installing a token, so either it finishes first and is removed here,
// or it observes the closed state and cannot enqueue a new stale token afterward.
c.rpcScheduler.unschedule(c)
if c.rpcCancel != nil {
c.rpcCancel()
}
close(c.rpcStop)
// 抢占懒启动 Once：若 worker 尚未起，封住其启动，避免后续 ensureInboundRPCWorkers 的
// rpcWG.Add 与下面的 rpcWG.Wait 并发（WaitGroup 误用）。Once 互斥保证 Add happens-before Wait。
c.rpcWorkersOnce.Do(func() {})
c.drainInboundRPCQueue()
// 等 worker 退出，使关闭对 inbound 与 outbound（<-outboundDone）收敛对称；带超时防慢 handler 卡死。
c.waitInboundWorkers(rpcCloseWaitTimeout)
for i := range queued {
task := queued[i]
queued[i] = inboundRPC{}
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketDone)
}
method := task.method
reservation := task.budget
stopInboundRPCTask(task)
task = inboundRPC{}
reservation.release()
c.metrics.InboundRPCDropped(method, "connection_closed")
}
})
}
// waitInboundWorkers 等所有 inbound RPC worker 退出，最长 timeout。超时则放弃等待，
// worker 在其阻塞的底层调用返回后自行退出（rpcCancel 已发，最终收敛）。
func (c *Conn) waitInboundWorkers(timeout time.Duration) {
// waitInboundShutdown 等 Copy 前 reservation 完成 commit/abort，以及本连接已经出队的 RPC
// 完成，二者共用一个 timeout。超时后 reservation/共享 worker 会在底层调用最终返回时自行
// 收敛；连接 root context 已取消。
func (c *Conn) waitInboundShutdown(timeout time.Duration) bool {
done := make(chan struct{})
go func() {
c.rpcReservationWG.Wait()
c.rpcWG.Wait()
close(done)
}()
if timeout <= 0 {
return false
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-done:
return true
case <-timer.C:
}
}
func (c *Conn) drainInboundRPCQueue() {
for {
select {
case task := <-c.rpcQueue:
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "connection_closed")
default:
return
}
return false
}
}

422
internal/mtprotoedge/inbound_rpc_test.go
View file

@ -8,10 +8,40 @@ import (
"time"
)
func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
func newInboundTestConn(s *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) *Conn {
c := &Conn{metrics: NopMetrics{}}
c.startInboundRPCScheduler(2, 4, time.Second)
defer c.closeInboundRPCScheduler()
c.startInboundRPCScheduler(s, maxInflight, queueSize, timeout)
return c
}
func TestInboundRPCSchedulerIsLazyPerConnectionAndServer(t *testing.T) {
scheduler := newInboundRPCScheduler(4, 16, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 2, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
if c.rpcQueue != nil {
t.Fatal("new connection eagerly allocated an inbound queue")
}
scheduler.lifecycleMu.Lock()
workersStarted := scheduler.workersStarted
scheduler.lifecycleMu.Unlock()
if workersStarted {
t.Fatal("empty server eagerly started inbound RPC workers")
}
}
func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
scheduler := newInboundRPCScheduler(2, 32, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 2, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
var active atomic.Int64
var maxActive atomic.Int64
@ -73,4 +103,390 @@ func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
time.Sleep(10 * time.Millisecond)
}
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("global budget after completion = (%d tasks, %d bytes), want zero", tasks, bytes)
}
}
func TestInboundRPCSchedulerFairAcrossConnections(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
c2 := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c1.closeInboundRPCScheduler()
c2.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
order := make(chan string, 3)
enqueue := func(c *Conn, label string) {
t.Helper()
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: label,
run: func(context.Context) error {
order <- label
return nil
},
}); err != nil {
t.Fatalf("enqueue %s: %v", label, err)
}
}
// 先在 worker 启动前形成 [c1, c2] ready 顺序。c1 每次只执行一条后回到队尾，
// 因此 c2 必须在 c1 的第二条之前获得执行机会。
enqueue(c1, "c1-first")
enqueue(c1, "c1-second")
enqueue(c2, "c2-first")
scheduler.start()
want := []string{"c1-first", "c2-first", "c1-second"}
for i := range want {
select {
case got := <-order:
if got != want[i] {
t.Fatalf("execution[%d] = %q, want %q", i, got, want[i])
}
case <-time.After(time.Second):
t.Fatalf("timed out waiting for execution[%d]", i)
}
}
}
func TestInboundRPCBudgetReservedBeforeCommitAndFullyReturned(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 2, 10)
c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
c2 := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c1.closeInboundRPCScheduler()
c2.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
r1, err := c1.reserveInboundRPC(context.Background(), "one", 6)
if err != nil {
t.Fatalf("reserve first body: %v", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 6 {
t.Fatalf("budget after first pre-Copy reservation = (%d, %d), want (1, 6)", tasks, bytes)
}
if _, err := c2.reserveInboundRPC(context.Background(), "too-large", 5); !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("reserve over byte budget err = %v, want queue full", err)
}
r2, err := c2.reserveInboundRPC(context.Background(), "two", 4)
if err != nil {
t.Fatalf("reserve second body: %v", err)
}
if _, err := c1.reserveInboundRPC(context.Background(), "too-many", 0); !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("reserve over task budget err = %v, want queue full", err)
}
r1.abort()
r2.abort()
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after aborts = (%d, %d), want zero", tasks, bytes)
}
if got := c1.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("c1 inflight bytes = %d, want zero", got)
}
if got := c2.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("c2 inflight bytes = %d, want zero", got)
}
}
func TestInboundRPCPerConnectionByteBudgetRejectedBeforeCommit(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 2, int64(maxInflightRPCBytes)+1)
c := newInboundTestConn(scheduler, 1, 2, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
if _, err := c.reserveInboundRPC(context.Background(), "oversized", maxInflightRPCBytes+1); !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("reserve over per-connection byte budget err = %v, want queue full", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("global budget after per-connection rejection = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after rejection = %d, want zero", got)
}
}
func TestInboundRPCCommitRacingCloseReturnsReservation(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 4, 1<<20)
c := newInboundTestConn(scheduler, 1, 2, time.Second)
defer scheduler.stop(time.Second)
reservation, err := c.reserveInboundRPC(context.Background(), "closing", 13)
if err != nil {
t.Fatalf("reserve: %v", err)
}
closed := make(chan struct{})
go func() {
c.closeInboundRPCScheduler()
close(closed)
}()
deadline := time.Now().Add(time.Second)
for {
c.rpcMu.Lock()
isClosed := c.rpcClosed
c.rpcMu.Unlock()
if isClosed {
break
}
if time.Now().After(deadline) {
t.Fatal("connection scheduler was not marked closed")
}
time.Sleep(time.Millisecond)
}
if err := reservation.commit(inboundRPC{run: func(context.Context) error { return nil }}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("commit after close err = %v, want ErrConnClosed", err)
}
select {
case <-closed:
case <-time.After(time.Second):
t.Fatal("close did not finish after reservation commit")
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("global budget after close/commit race = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after close/commit race = %d, want zero", got)
}
}
func TestInboundRPCSchedulerCloseRemovesReadyTokenBeforeStart(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 1, 1<<20)
defer scheduler.stop(time.Second)
// A bounded ready channel used to retain one stale token per closed connection. With workers
// not started yet, the second connection then blocked forever trying to publish its token even
// though the first connection had returned every task/byte budget.
for i := 0; i < 32; i++ {
c := newInboundTestConn(scheduler, 1, 1, time.Second)
done := make(chan error, 1)
go func() {
done <- c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "close-before-start",
run: func(context.Context) error { return nil },
})
}()
select {
case err := <-done:
if err != nil {
t.Fatalf("enqueue iteration %d: %v", i, err)
}
case <-time.After(time.Second):
t.Fatalf("enqueue iteration %d blocked behind a stale ready token", i)
}
c.closeInboundRPCScheduler()
if got := scheduler.readyLen(); got != 0 {
t.Fatalf("ready tokens after close iteration %d = %d, want zero", i, got)
}
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after close churn = (%d, %d), want zero", tasks, bytes)
}
}
func TestInboundRPCExpiredInQueueNeverRunsAndSignalsTimeout(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 1, 4, 40*time.Millisecond)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
started := make(chan struct{})
release := make(chan struct{})
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "blocker",
size: 7,
run: func(context.Context) error {
close(started)
<-release // 刻意忽略 deadline，确保下一条在队列中到期。
return nil
},
}); err != nil {
t.Fatalf("enqueue blocker: %v", err)
}
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("blocker did not start")
}
var ran atomic.Bool
timedOut := make(chan struct{})
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "expires",
size: 11,
onTimeout: func() {
close(timedOut)
},
run: func(context.Context) error {
ran.Store(true)
return nil
},
}); err != nil {
t.Fatalf("enqueue expiring task: %v", err)
}
select {
case <-timedOut:
case <-time.After(time.Second):
t.Fatal("queued task did not signal timeout while the worker was still blocked")
}
deadline := time.Now().Add(time.Second)
for {
tasks, bytes := scheduler.budgetSnapshot()
if tasks == 1 && bytes == 7 {
break
}
if time.Now().After(deadline) {
t.Fatalf("budget while blocker still runs = (%d, %d), want only blocker (1, 7)", tasks, bytes)
}
time.Sleep(time.Millisecond)
}
close(release)
if ran.Load() {
t.Fatal("expired queued task entered business handler")
}
deadline = time.Now().Add(time.Second)
for {
tasks, bytes := scheduler.budgetSnapshot()
if tasks == 0 && bytes == 0 {
break
}
if time.Now().After(deadline) {
t.Fatalf("budget after timeout = (%d, %d), want zero", tasks, bytes)
}
time.Sleep(time.Millisecond)
}
}
func TestInboundRPCCloseDisarmsQueuedTimeout(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
defer scheduler.stop(time.Second)
timedOut := make(chan struct{}, 1)
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "queued",
size: 11,
onTimeout: func() {
timedOut <- struct{}{}
},
}); err != nil {
t.Fatalf("enqueue queued task: %v", err)
}
c.closeInboundRPCScheduler()
time.Sleep(60 * time.Millisecond)
select {
case <-timedOut:
t.Fatal("connection close emitted a queued RPC timeout")
default:
}
}
func TestInboundRPCRunningTimeoutSignalsWithoutReleasingBodyEarly(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
started := make(chan struct{})
release := make(chan struct{})
timedOut := make(chan struct{}, 1)
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "running",
size: 7,
onTimeout: func() {
timedOut <- struct{}{}
},
run: func(context.Context) error {
close(started)
<-release
return nil
},
}); err != nil {
t.Fatalf("enqueue running task: %v", err)
}
<-started
select {
case <-timedOut:
case <-time.After(time.Second):
t.Fatal("running task did not signal timeout while handler ignored cancellation")
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 7 {
t.Fatalf("running body budget after timeout = (%d, %d), want retained (1, 7)", tasks, bytes)
}
close(release)
deadline := time.Now().Add(time.Second)
for {
tasks, bytes := scheduler.budgetSnapshot()
if tasks == 0 && bytes == 0 {
break
}
if time.Now().After(deadline) {
t.Fatalf("running body budget after completion = (%d, %d), want zero", tasks, bytes)
}
time.Sleep(time.Millisecond)
}
}
func TestInboundRPCCloseDrainsQueueAndReturnsBudgets(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer scheduler.stop(time.Second)
started := make(chan struct{})
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "running",
size: 7,
run: func(ctx context.Context) error {
close(started)
<-ctx.Done()
return ctx.Err()
},
}); err != nil {
t.Fatalf("enqueue running task: %v", err)
}
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("running task did not start")
}
var queuedRan atomic.Bool
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "queued",
size: 11,
run: func(context.Context) error {
queuedRan.Store(true)
return nil
},
}); err != nil {
t.Fatalf("enqueue queued task: %v", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 2 || bytes != 18 {
t.Fatalf("budget before close = (%d, %d), want (2, 18)", tasks, bytes)
}
c.closeInboundRPCScheduler()
if queuedRan.Load() {
t.Fatal("queued task ran during connection close")
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after close = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection inflight bytes after close = %d, want zero", got)
}
}

9
internal/mtprotoedge/login_email_e2e_test.go
View file

@ -71,11 +71,16 @@ func TestLoginEmailEndToEnd(t *testing.T) {
passwordStore := memory.NewPasswordStore()
helpStore := memory.NewHelpStore()
codeStore := memory.NewCodeStore()
dialogStore := memory.NewDialogStore()
messageStore := memory.NewMessageStore(dialogStore)
updateEventStore := memory.NewUpdateEventStore()
emailSender := &loginEmailTestSender{}
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,
auth.WithLoginMessages(messageStore, dialogStore),
auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore)),
auth.WithPasswords(passwordStore),
auth.WithLoginEmail(auth.LoginEmailOptions{
Enabled: true,
@ -89,10 +94,10 @@ func TestLoginEmailEndToEnd(t *testing.T) {
Account: accountService,
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Updates: updates.NewService(memory.NewUpdateStateStore(), updateEventStore),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(memory.NewDialogStore()),
Dialogs: dialogs.NewService(dialogStore),
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router})

1000
internal/mtprotoedge/outbound.go
View file

File diff suppressed because it is too large Load diff

126
internal/mtprotoedge/outbound_scratch.go Normal file
View file

@ -0,0 +1,126 @@
package mtprotoedge
import (
"context"
"time"
"github.com/gotd/td/bin"
)
const (
defaultOutboundWriteMaxBytes = int64(512 << 20)
defaultOutboundScratchPool = 256
)
// outboundScratchPool bounds and reuses the encrypted wire buffer across connections. A lease
// reserves a conservative 3x wire size while writing (wire + codec/obfuscation copies), then
// shrinks to the actual retained capacity while idle in the bounded pool. Large one-off frames are
// dropped on return. This removes attacker-warmable per-Conn MiB buffers without returning to an
// unbounded allocation-per-message design.
type outboundScratchPool struct {
budget *outboundTrackedBudget
idle chan *outboundScratch
}
type outboundScratch struct {
wire bin.Buffer
reserved int
}
func newOutboundScratchPool(maxBytes int64) *outboundScratchPool {
if maxBytes <= 0 {
maxBytes = defaultOutboundWriteMaxBytes
}
return &outboundScratchPool{
budget: newOutboundTrackedBudget(maxBytes),
idle: make(chan *outboundScratch, defaultOutboundScratchPool),
}
}
func (p *outboundScratchPool) acquire(ctx context.Context, stop <-chan struct{}, wireBytes int) (*outboundScratch, error) {
return p.acquireUntil(ctx, stop, wireBytes, time.Time{})
}
func (p *outboundScratchPool) acquireUntil(ctx context.Context, stop <-chan struct{}, wireBytes int, deadline time.Time) (*outboundScratch, error) {
if p == nil || wireBytes <= 0 {
return nil, ErrOutboundMessageTooLarge
}
peak := wireBytes * 3
if peak < wireBytes { // int overflow
return nil, ErrOutboundMessageTooLarge
}
var scratch *outboundScratch
select {
case scratch = <-p.idle:
default:
}
if scratch == nil {
if err := p.budget.waitReserveUntil(ctx, stop, peak, deadline); err != nil {
return nil, err
}
return &outboundScratch{wire: bin.Buffer{Buf: make([]byte, wireBytes)}, reserved: peak}, nil
}
if cap(scratch.wire.Buf) >= wireBytes {
if extra := peak - scratch.reserved; extra > 0 {
if err := p.budget.waitReserveUntil(ctx, stop, extra, deadline); err != nil {
p.putIdle(scratch)
return nil, err
}
scratch.reserved += extra
}
scratch.wire.Buf = scratch.wire.Buf[:wireBytes]
return scratch, nil
}
// The old slice is no longer reachable after clearing it; return that retained charge before
// waiting for a larger lease, otherwise old+peak may exceed the budget and deadlock a resize
// that would fit after replacement.
old := scratch.reserved
scratch.wire.Buf = nil
scratch.reserved = 0
p.budget.release(old)
if err := p.budget.waitReserveUntil(ctx, stop, peak, deadline); err != nil {
return nil, err
}
scratch.wire.Buf = make([]byte, wireBytes)
scratch.reserved = peak
return scratch, nil
}
func (p *outboundScratchPool) release(scratch *outboundScratch) {
if p == nil || scratch == nil {
return
}
retained := cap(scratch.wire.Buf)
if retained > maxRetainedConnBuffer {
p.budget.release(scratch.reserved)
scratch.wire.Buf = nil
scratch.reserved = 0
return
}
if scratch.reserved > retained {
p.budget.release(scratch.reserved - retained)
scratch.reserved = retained
}
scratch.wire.Buf = scratch.wire.Buf[:0]
p.putIdle(scratch)
}
func (p *outboundScratchPool) putIdle(scratch *outboundScratch) {
select {
case p.idle <- scratch:
default:
p.budget.release(scratch.reserved)
scratch.wire.Buf = nil
scratch.reserved = 0
}
}
func (p *outboundScratchPool) snapshot() int64 {
if p == nil {
return 0
}
return p.budget.snapshot()
}

805
internal/mtprotoedge/outbound_test.go
View file

@ -4,7 +4,10 @@ import (
"bytes"
"context"
"crypto/rand"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
@ -13,8 +16,426 @@ import (
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
)
type failAfterTransport struct {
failAt atomic.Int32
sends atomic.Int32
stored atomic.Int32
closes atomic.Int32
mu sync.Mutex
last []byte
}
type blockingOutboundTransport struct {
started chan struct{}
release chan struct{}
once sync.Once
sends atomic.Int32
}
type blockingEncodeProbe struct {
started chan struct{}
release <-chan struct{}
active atomic.Int32
max atomic.Int32
}
func (e *blockingEncodeProbe) Encode(b *bin.Buffer) error {
active := e.active.Add(1)
for {
max := e.max.Load()
if active <= max || e.max.CompareAndSwap(max, active) {
break
}
}
e.started <- struct{}{}
<-e.release
e.active.Add(-1)
b.PutID(tg.UpdatesTooLongTypeID)
return nil
}
func newBlockingOutboundTransport() *blockingOutboundTransport {
return &blockingOutboundTransport{started: make(chan struct{}), release: make(chan struct{})}
}
func TestOutboundEncodingHasProcessWideConcurrencyBudget(t *testing.T) {
const extra = 8
total := defaultOutboundEncodeConcurrency + extra
release := make(chan struct{})
probe := &blockingEncodeProbe{
started: make(chan struct{}, total),
release: release,
}
errs := make(chan error, total)
for range total {
go func() {
_, err := encodeOutboundMessage(probe)
errs <- err
}()
}
for range defaultOutboundEncodeConcurrency {
select {
case <-probe.started:
case <-time.After(time.Second):
t.Fatal("encode workers did not fill concurrency budget")
}
}
select {
case <-probe.started:
t.Fatalf("more than %d outbound encodes ran concurrently", defaultOutboundEncodeConcurrency)
case <-time.After(50 * time.Millisecond):
}
close(release)
for range total {
if err := <-errs; err != nil {
t.Fatalf("encode: %v", err)
}
}
if got := probe.max.Load(); got != defaultOutboundEncodeConcurrency {
t.Fatalf("peak concurrent encodes = %d, want %d", got, defaultOutboundEncodeConcurrency)
}
}
func TestConnectionCloseDoesNotWaitForRunningEncoder(t *testing.T) {
release := make(chan struct{})
probe := &blockingEncodeProbe{started: make(chan struct{}, 1), release: release}
c := &Conn{metrics: NopMetrics{}}
c.startOutbound()
sendDone := make(chan error, 1)
go func() {
sendDone <- c.Send(context.Background(), proto.MessageFromServer, probe)
}()
select {
case <-probe.started:
case <-time.After(time.Second):
t.Fatal("encoder did not start")
}
closeDone := make(chan struct{})
go func() {
c.Close()
close(closeDone)
}()
select {
case <-closeDone:
case <-time.After(time.Second):
t.Fatal("Conn.Close waited for external Encoder")
}
close(release)
select {
case err := <-sendDone:
if !errors.Is(err, ErrConnClosed) {
t.Fatalf("send after concurrent close = %v, want ErrConnClosed", err)
}
case <-time.After(time.Second):
t.Fatal("send did not return after encoder release")
}
}
func TestOutboundControlVectorsUseGlobalByteBudget(t *testing.T) {
budget := newOutboundTrackedBudget(16)
c := &Conn{outboundControlTrackedBudget: budget}
op, err := c.newOutboundVectorOp(outboundAck, []int64{1, 2})
if err != nil {
t.Fatalf("reserve first vector: %v", err)
}
if got := budget.snapshot(); got != 16 {
t.Fatalf("tracked bytes after reserve = %d, want 16", got)
}
if _, err := c.newOutboundVectorOp(outboundResend, []int64{3}); !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("reserve over budget error = %v, want %v", err, ErrOutboundTrackedBudget)
}
op.releaseReservation(budget)
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after release = %d, want 0", got)
}
}
func TestEncodedControlFramesUseIndependentBudgetForQueuedAndPendingLifetime(t *testing.T) {
bodyBudget := newOutboundTrackedBudget(4)
controlBudget := newOutboundTrackedBudget(256)
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, bodyBudget)
c.outboundControlTrackedBudget = controlBudget
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
// One content frame fills the ordinary body budget and remains pending.
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("fill body budget: %v", err)
}
first, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt ordinary frame: %v", err)
}
if got := bodyBudget.snapshot(); got != 4 {
t.Fatalf("body budget = %d, want saturated 4", got)
}
created := &mt.NewSessionCreated{FirstMsgID: 1, UniqueID: 2, ServerSalt: 3}
encodedCreated, err := encodeOutboundMessageWithoutSlot(created)
if err != nil {
t.Fatalf("encode new_session_created: %v", err)
}
if err := c.SendAsync(ctx, proto.MessageFromServer, created); err != nil {
t.Fatalf("new_session_created under saturated body budget: %v", err)
}
deadline := time.Now().Add(time.Second)
for tr.stored.Load() < 2 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := tr.stored.Load(); got != 2 {
t.Fatalf("completed physical sends = %d, want 2", got)
}
second, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt control frame: %v", err)
}
if got := bodyBudget.snapshot(); got != 4 {
t.Fatalf("body budget after control send = %d, want unchanged 4", got)
}
if got := controlBudget.snapshot(); got != int64(len(encodedCreated.body)) {
t.Fatalf("control pending budget = %d, want new_session_created body %d", got, len(encodedCreated.body))
}
select {
case <-c.outboundDone:
t.Fatal("ordinary body pressure closed a healthy connection")
default:
}
// Pong is non-pending, but must also remain admissible and return its control bytes after write.
if err := c.SendAsync(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 4, PingID: 5}); err != nil {
t.Fatalf("pong under saturated body budget: %v", err)
}
deadline = time.Now().Add(time.Second)
for tr.stored.Load() < 3 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := tr.stored.Load(); got != 3 {
t.Fatalf("completed physical sends after pong = %d, want 3", got)
}
if got := controlBudget.snapshot(); got != int64(len(encodedCreated.body)) {
t.Fatalf("control budget after non-pending pong = %d, want pending %d", got, len(encodedCreated.body))
}
c.AckServerMessages([]int64{first.MessageID, second.MessageID})
deadline = time.Now().Add(time.Second)
for (bodyBudget.snapshot() != 0 || controlBudget.snapshot() != 0) && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := bodyBudget.snapshot(); got != 0 {
t.Fatalf("body budget after ACK = %d, want 0", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after ACK = %d, want 0", got)
}
}
func TestOutboundScratchPoolBoundsConcurrentWireCopies(t *testing.T) {
pool := newOutboundScratchPool(300)
first, err := pool.acquire(context.Background(), nil, 100) // 3x peak = full budget.
if err != nil {
t.Fatalf("acquire first scratch: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
if _, err := pool.acquire(ctx, nil, 100); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("second concurrent acquire = %v, want deadline backpressure", err)
}
pool.release(first)
if got := pool.snapshot(); got != 100 {
t.Fatalf("idle retained scratch = %d, want 100", got)
}
second, err := pool.acquire(context.Background(), nil, 100)
if err != nil {
t.Fatalf("reuse retained scratch: %v", err)
}
pool.release(second)
if got := pool.snapshot(); got != 100 {
t.Fatalf("scratch after reuse = %d, want one bounded idle buffer", got)
}
}
func TestOutboundScratchAdmissionUsesWriteTimeoutWithoutClosingHealthyConnection(t *testing.T) {
wireBytes := encryptedOutboundWireLen(4)
pool := newOutboundScratchPool(int64(wireBytes * 3))
blocker, err := pool.acquire(context.Background(), nil, wireBytes)
if err != nil {
t.Fatalf("occupy shared scratch budget: %v", err)
}
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundScratchPool = pool
c.writeTimeout = 25 * time.Millisecond
start := time.Now()
err = c.Send(context.Background(), proto.MessageFromServer, &tg.UpdatesTooLong{})
elapsed := time.Since(start)
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("scratch admission err = %v, want deadline exceeded", err)
}
if elapsed > 250*time.Millisecond {
t.Fatalf("scratch admission waited %v, want writeTimeout-bounded wait", elapsed)
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("writer called %d times without scratch, want 0", got)
}
if c.terminal.Load() {
t.Fatal("scratch admission timeout terminally closed a healthy connection")
}
select {
case <-c.outboundDone:
t.Fatal("outbound actor exited after pre-write scratch timeout")
default:
}
pool.release(blocker)
c.writeTimeout = time.Second
if err := c.Send(context.Background(), proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("send after scratch capacity returned: %v", err)
}
if got := tr.sends.Load(); got != 1 {
t.Fatalf("writer calls after recovery = %d, want 1", got)
}
}
func (t *blockingOutboundTransport) Send(context.Context, *bin.Buffer) error {
if t.sends.Add(1) == 1 {
close(t.started)
}
<-t.release
return io.ErrClosedPipe
}
func (t *blockingOutboundTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
func (t *blockingOutboundTransport) Close() error {
t.once.Do(func() { close(t.release) })
return nil
}
func (t *failAfterTransport) Send(_ context.Context, b *bin.Buffer) error {
n := t.sends.Add(1)
if failAt := t.failAt.Load(); failAt > 0 && n >= failAt {
return io.ErrClosedPipe
}
t.mu.Lock()
t.last = append(t.last[:0], b.Raw()...)
t.mu.Unlock()
t.stored.Add(1)
return nil
}
func (t *failAfterTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
func (t *failAfterTransport) Close() error {
t.closes.Add(1)
return nil
}
func (t *failAfterTransport) lastFrame() []byte {
t.mu.Lock()
defer t.mu.Unlock()
return append([]byte(nil), t.last...)
}
func newOutboundFailureTestConn(t *testing.T, tr transport.Conn) *Conn {
return newOutboundTestConn(t, tr, nil)
}
func newOutboundTestConn(t *testing.T, tr transport.Conn, budget *outboundTrackedBudget) *Conn {
t.Helper()
var key crypto.Key
if _, err := rand.Read(key[:]); err != nil {
t.Fatalf("rand key: %v", err)
}
c := &Conn{
transport: tr,
writer: tr,
cipher: crypto.NewServerCipher(rand.Reader),
msgID: proto.NewMessageIDGen(time.Now),
writeTimeout: time.Second,
metrics: NopMetrics{},
key: key.WithID(),
salt: 123,
sessionID: 456,
outboundTrackedBudget: budget,
}
c.startOutbound()
t.Cleanup(c.Close)
return c
}
func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) {
t.Run("defaults", func(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
c.startOutbound()
defer c.Close()
if got := cap(c.outbound); got != defaultOutboundQueueSize {
t.Fatalf("normal queue cap = %d, want %d", got, defaultOutboundQueueSize)
}
if got := cap(c.outboundControl); got != defaultOutboundControlQueueSize {
t.Fatalf("control queue cap = %d, want %d", got, defaultOutboundControlQueueSize)
}
})
t.Run("configured", func(t *testing.T) {
c := &Conn{
metrics: NopMetrics{},
outboundQueueSize: 7,
outboundControlQueueSize: 3,
}
c.startOutbound()
defer c.Close()
if got := cap(c.outbound); got != 7 {
t.Fatalf("normal queue cap = %d, want 7", got)
}
if got := cap(c.outboundControl); got != 3 {
t.Fatalf("control queue cap = %d, want 3", got)
}
})
}
func TestOutboundOptionsDefaults(t *testing.T) {
opts := Options{}
opts.setDefaults()
if opts.OutboundQueueSize != 128 || opts.OutboundControlQueueSize != 32 {
t.Fatalf("outbound queue defaults = %d/%d, want 128/32", opts.OutboundQueueSize, opts.OutboundControlQueueSize)
}
if opts.OutboundTrackedGlobalMaxBytes != 512<<20 {
t.Fatalf("outbound tracked default = %d, want %d", opts.OutboundTrackedGlobalMaxBytes, 512<<20)
}
}
func TestServerNewConnectionsShareOutboundBudgetAndQueueLimits(t *testing.T) {
srv := New(Options{
OutboundQueueSize: 7,
OutboundControlQueueSize: 3,
OutboundTrackedGlobalMaxBytes: 20,
})
var rawKey crypto.Key
key := rawKey.WithID()
c1 := srv.newConn(nil, key, 1, 1)
c2 := srv.newConn(nil, key, 2, 1)
defer c1.Close()
defer c2.Close()
if cap(c1.outbound) != 7 || cap(c1.outboundControl) != 3 || cap(c2.outbound) != 7 || cap(c2.outboundControl) != 3 {
t.Fatalf("server queue caps = %d/%d and %d/%d, want 7/3",
cap(c1.outbound), cap(c1.outboundControl), cap(c2.outbound), cap(c2.outboundControl))
}
if c1.outboundTrackedBudget != srv.outboundTrackedBudget || c2.outboundTrackedBudget != srv.outboundTrackedBudget {
t.Fatal("server connections did not receive the shared outbound tracking budget")
}
if got := srv.outboundTrackedBudget.maxBytes; got != 20 {
t.Fatalf("server outbound tracked max = %d, want 20", got)
}
}
func TestEncryptOutboundFrameDecryptsWithGotdCipher(t *testing.T) {
var key crypto.Key
if _, err := rand.Read(key[:]); err != nil {
@ -104,8 +525,375 @@ func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
}
}
func TestOutboundWriteErrorTerminallyClosesWithoutActorDeadlock(t *testing.T) {
tr := &failAfterTransport{}
tr.failAt.Store(1)
c := newOutboundFailureTestConn(t, tr)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err == nil {
t.Fatal("Send unexpectedly succeeded")
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor deadlocked while terminalizing its own write error")
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("second Send err = %v, want ErrConnClosed", err)
}
if got := tr.sends.Load(); got != 1 {
t.Fatalf("physical sends after terminal error = %d, want 1", got)
}
}
func TestOutboundResendWriteErrorTerminallyCloses(t *testing.T) {
tr := &failAfterTransport{}
c := newOutboundFailureTestConn(t, tr)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("initial Send: %v", err)
}
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt initial frame: %v", err)
}
tr.failAt.Store(2)
if _, err := c.ResendMessages(ctx, []int64{data.MessageID}); err == nil {
t.Fatal("ResendMessages unexpectedly succeeded")
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not exit after resend write error")
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
}
func TestOutboundTrackedBudgetSharedAcrossConnections(t *testing.T) {
budget := newOutboundTrackedBudget(12)
tr1 := &failAfterTransport{}
tr2 := &failAfterTransport{}
c1 := newOutboundTestConn(t, tr1, budget)
c2 := newOutboundTestConn(t, tr2, budget)
body := &encodedOutboundMessage{body: make([]byte, 8), typeID: tg.UpdatesTooLongTypeID}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := c1.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
t.Fatalf("first connection send: %v", err)
}
if got := budget.snapshot(); got != 8 {
t.Fatalf("tracked bytes after first connection = %d, want 8", got)
}
if err := c2.SendEncoded(ctx, proto.MessageFromServer, body); !errors.Is(err, ErrOutboundTrackedBudget) && !errors.Is(err, ErrConnClosed) {
t.Fatalf("second connection send err = %v, want tracked budget/closed", err)
}
select {
case <-c2.outboundDone:
case <-time.After(time.Second):
t.Fatal("budget-exhausted connection did not terminate")
}
if got := tr2.sends.Load(); got != 0 {
t.Fatalf("budget-exhausted connection wrote %d frames, want 0", got)
}
if got := budget.snapshot(); got != 8 {
t.Fatalf("tracked bytes after second rejection = %d, want first connection's 8", got)
}
c1.Close()
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after first connection close = %d, want 0", got)
}
}
func TestOutboundTrackedBudgetReleaseBroadcastsToAllWaiters(t *testing.T) {
const waiters = 8
budget := newOutboundTrackedBudget(waiters)
if !budget.reserve(waiters) {
t.Fatal("reserve initial saturated budget")
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
results := make(chan error, waiters)
for i := 0; i < waiters; i++ {
go func() {
results <- budget.waitReserve(ctx, nil, 1)
}()
}
deadline := time.Now().Add(time.Second)
for {
budget.wakeMu.Lock()
got := budget.wake.waiters
budget.wakeMu.Unlock()
if got == waiters {
break
}
if time.Now().After(deadline) {
t.Fatalf("subscribed waiters = %d, want %d", got, waiters)
}
time.Sleep(time.Millisecond)
}
// One batch release creates capacity for every waiter. A single-token notification strands
// seven of them forever because successful reservations do not produce another wake-up.
budget.release(waiters)
for i := 0; i < waiters; i++ {
if err := <-results; err != nil {
t.Fatalf("waiter %d: %v", i, err)
}
}
if got := budget.snapshot(); got != waiters {
t.Fatalf("reserved bytes after broadcast = %d, want %d", got, waiters)
}
budget.release(waiters)
}
func TestOutboundGlobalBudgetIncludesQueuedBodies(t *testing.T) {
budget := newOutboundTrackedBudget(24)
tr := newBlockingOutboundTransport()
c := newOutboundTestConn(t, tr, budget)
body := &encodedOutboundMessage{body: make([]byte, 8), typeID: tg.UpdatesTooLongTypeID}
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); err != nil {
t.Fatalf("enqueue writing body: %v", err)
}
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("outbound actor did not start blocked write")
}
for i := 0; i < 2; i++ {
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); err != nil {
t.Fatalf("enqueue queued body %d: %v", i, err)
}
}
if got := budget.snapshot(); got != 24 {
t.Fatalf("writing + queued budget = %d, want 24", got)
}
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("over-budget enqueue err = %v, want ErrOutboundTrackedBudget", err)
}
select {
case <-c.outboundDone:
t.Fatal("best-effort global pressure terminated a healthy connection")
case <-time.After(50 * time.Millisecond):
}
if got := budget.snapshot(); got != 24 {
t.Fatalf("budget after non-terminal rejection = %d, want existing 24", got)
}
if err := tr.Close(); err != nil {
t.Fatalf("close blocking transport: %v", err)
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after transport failure")
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("budget after transport close = %d, want zero", got)
}
}
func TestOutboundOversizedBodyRejectedBeforeEncryption(t *testing.T) {
budget := newOutboundTrackedBudget(64 << 20)
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, budget)
body := &encodedOutboundMessage{body: make([]byte, maxOutboundBodyBytes+1), typeID: tg.UpdatesTooLongTypeID}
err := c.SendEncoded(context.Background(), proto.MessageFromServer, body)
if !errors.Is(err, ErrOutboundMessageTooLarge) {
t.Fatalf("oversized outbound err = %v, want ErrOutboundMessageTooLarge", err)
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("oversized outbound wrote %d frames, want zero", got)
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("oversized outbound reserved %d bytes, want zero", got)
}
}
func TestOutboundCloseRaceDrainsEveryProducerReservation(t *testing.T) {
budget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, &failAfterTransport{}, budget)
body := &encodedOutboundMessage{body: make([]byte, 128), typeID: tg.UpdatesTooLongTypeID}
start := make(chan struct{})
var wg sync.WaitGroup
for i := 0; i < 128; i++ {
wg.Add(1)
go func() {
defer wg.Done()
<-start
_ = c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0)
}()
}
close(start)
c.Close()
wg.Wait()
if got := budget.snapshot(); got != 0 {
t.Fatalf("outbound budget after close/enqueue race = %d, want zero", got)
}
}
func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
t.Run("ack", func(t *testing.T) {
budget := newOutboundTrackedBudget(64)
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, budget)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
t.Fatalf("send: %v", err)
}
if got := budget.snapshot(); got != 12 {
t.Fatalf("tracked bytes after send = %d, want 12", got)
}
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt frame: %v", err)
}
c.AckServerMessages([]int64{data.MessageID})
deadline := time.Now().Add(time.Second)
for budget.snapshot() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after ack = %d, want 0", got)
}
})
t.Run("close", func(t *testing.T) {
budget := newOutboundTrackedBudget(64)
c := newOutboundTestConn(t, &failAfterTransport{}, budget)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
t.Fatalf("send: %v", err)
}
if got := budget.snapshot(); got != 12 {
t.Fatalf("tracked bytes after send = %d, want 12", got)
}
c.Close()
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after close = %d, want 0", got)
}
})
}
func TestOutboundTrackedBudgetWriteFailureReturnsReservation(t *testing.T) {
budget := newOutboundTrackedBudget(64)
tr := &failAfterTransport{}
tr.failAt.Store(1)
c := newOutboundTestConn(t, tr, budget)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err == nil {
t.Fatal("send unexpectedly succeeded")
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("write-failed connection did not terminate")
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after write failure = %d, want 0", got)
}
}
func TestOutboundStateEvictionReturnsTrackedBudget(t *testing.T) {
budget := newOutboundTrackedBudget(64)
state := newOutboundStateWithLimits(budget, 2, 8)
defer state.releaseAll()
frames := make([]*outboundFrame, 0, 3)
for id := int64(1); id <= 3; id++ {
frame := &outboundFrame{msgID: id, body: make([]byte, 4), reservedBytes: 4}
frames = append(frames, frame)
if !budget.reserve(len(frame.body)) {
t.Fatalf("reserve frame %d", id)
}
dropped := state.addReserved(frame)
if id < 3 && dropped != 0 {
t.Fatalf("frame %d dropped %d, want 0", id, dropped)
}
if id == 3 && dropped != 1 {
t.Fatalf("third frame dropped %d, want 1", dropped)
}
}
if got := budget.snapshot(); got != 8 {
t.Fatalf("tracked bytes after eviction = %d, want 8", got)
}
if frames[0].body != nil {
t.Fatal("evicted frame retained its body reference")
}
state.releaseAll()
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after state close = %d, want 0", got)
}
}
func TestOutboundStateReleasesMixedBodyAndControlBudgets(t *testing.T) {
bodyBudget := newOutboundTrackedBudget(16)
controlBudget := newOutboundTrackedBudget(16)
state := newOutboundStateWithLimits(bodyBudget, 1, 16)
if !controlBudget.reserve(4) {
t.Fatal("reserve control frame")
}
controlFrame := &outboundFrame{
msgID: 1,
body: make([]byte, 4),
reservedBytes: 4,
reservationBudget: controlBudget,
}
if dropped := state.addReserved(controlFrame); dropped != 0 {
t.Fatalf("first add dropped %d, want 0", dropped)
}
if !bodyBudget.reserve(4) {
t.Fatal("reserve body frame")
}
bodyFrame := &outboundFrame{
msgID: 2,
body: make([]byte, 4),
reservedBytes: 4,
reservationBudget: bodyBudget,
}
if dropped := state.addReserved(bodyFrame); dropped != 1 {
t.Fatalf("second add dropped %d, want control frame eviction", dropped)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after eviction = %d, want 0", got)
}
if got := bodyBudget.snapshot(); got != 4 {
t.Fatalf("body budget after eviction = %d, want 4", got)
}
if controlFrame.body != nil || controlFrame.reservationBudget != nil {
t.Fatal("evicted control frame retained body or budget ownership")
}
state.releaseAll()
if got := bodyBudget.snapshot(); got != 0 {
t.Fatalf("body budget after state close = %d, want 0", got)
}
if bodyFrame.body != nil || bodyFrame.reservationBudget != nil {
t.Fatal("closed body frame retained body or budget ownership")
}
}
func TestSendBestEffortQueueFullBehavior(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: newOutboundTrackedBudget(1 << 20)}
c.outbound = make(chan outboundOp, 1)
c.outboundControl = make(chan outboundOp, 1)
c.outboundStop = make(chan struct{})
@ -140,6 +928,21 @@ func TestSendBestEffortQueueFullBehavior(t *testing.T) {
}
}
func TestSendAsyncControlQueueBoundary(t *testing.T) {
c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: newOutboundTrackedBudget(1 << 20)}
c.outbound = make(chan outboundOp, 1)
c.outboundControl = make(chan outboundOp, 1)
c.outboundStop = make(chan struct{})
c.outboundControl <- outboundOp{kind: outboundAck}
if err := c.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); err != nil {
t.Fatalf("SendAsync on full control queue: %v", err)
}
if got := len(c.outboundControl); got != 1 {
t.Fatalf("control queue len = %d, want bounded at 1", got)
}
}
func TestFrameNeedsAckServiceExceptions(t *testing.T) {
cases := []struct {
name string

11
internal/mtprotoedge/passkey_e2e_test.go
View file

@ -96,17 +96,22 @@ func TestPasskeyEndToEnd(t *testing.T) {
userStore := memory.NewUserStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
dialogStore := memory.NewDialogStore()
messageStore := memory.NewMessageStore(dialogStore)
updateEventStore := memory.NewUpdateEventStore()
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(), code,
auth.WithLoginMessages(messageStore, dialogStore),
auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore))),
Account: account.NewService(memory.NewPasswordStore(), account.WithUsers(userStore)),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Updates: updates.NewService(memory.NewUpdateStateStore(), updateEventStore),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(memory.NewDialogStore()),
Dialogs: dialogs.NewService(dialogStore),
Passkey: passkeyService,
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)

70
internal/mtprotoedge/quick_ack_deadline_test.go Normal file
View file

@ -0,0 +1,70 @@
package mtprotoedge
import (
"context"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
)
type quickAckDeadlineProbe struct {
requested bool
deadline time.Time
token uint32
}
func (p *quickAckDeadlineProbe) ConsumeQuickAckRequested() bool {
if !p.requested {
return false
}
p.requested = false
return true
}
func (p *quickAckDeadlineProbe) SendQuickAck(ctx context.Context, token uint32) error {
p.deadline, _ = ctx.Deadline()
p.token = token
return nil
}
func (p *quickAckDeadlineProbe) SendQuickAckDeadline(deadline time.Time, token uint32) error {
p.deadline = deadline
p.token = token
return nil
}
func (*quickAckDeadlineProbe) Send(context.Context, *bin.Buffer) error { return nil }
func (*quickAckDeadlineProbe) Recv(context.Context, *bin.Buffer) error { return nil }
func (*quickAckDeadlineProbe) Close() error { return nil }
func TestQuickAckUsesServerWriteDeadline(t *testing.T) {
probe := &quickAckDeadlineProbe{requested: true}
var key crypto.Key
authKey := key.WithID()
before := time.Now()
if err := sendQuickAckIfRequested(context.Background(), probe, authKey, []byte("plain"), 50*time.Millisecond); err != nil {
t.Fatalf("send quick ack: %v", err)
}
if probe.deadline.IsZero() {
t.Fatal("quick ack did not receive a write deadline")
}
if probe.deadline.Before(before.Add(40*time.Millisecond)) || probe.deadline.After(time.Now().Add(60*time.Millisecond)) {
t.Fatalf("quick ack deadline = %v, want about server timeout from now", probe.deadline)
}
}
func TestQuickAckHonorsEarlierCallerDeadline(t *testing.T) {
probe := &quickAckDeadlineProbe{requested: true}
ctxDeadline := time.Now().Add(25 * time.Millisecond)
ctx, cancel := context.WithDeadline(context.Background(), ctxDeadline)
defer cancel()
var key crypto.Key
if err := sendQuickAckIfRequested(ctx, probe, key.WithID(), []byte("plain"), time.Second); err != nil {
t.Fatalf("send quick ack: %v", err)
}
if delta := probe.deadline.Sub(ctxDeadline); delta < -time.Millisecond || delta > time.Millisecond {
t.Fatalf("quick ack deadline = %v, want caller deadline %v", probe.deadline, ctxDeadline)
}
}

146
internal/mtprotoedge/rpc_test.go
View file

@ -107,6 +107,118 @@ func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
close(handler.release)
}
func TestInboundRPCQueuedDeadlineReturnsRPCTimeout(t *testing.T) {
const dc = 2
handler := &queueDeadlineRPC{
firstStarted: make(chan struct{}),
releaseFirst: make(chan struct{}),
}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 2,
RPCTimeout: 60 * time.Millisecond,
RPCGlobalWorkers: 1,
})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
firstReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, firstReqID, 1, &tg.HelpGetConfigRequest{})
select {
case <-handler.firstStarted:
case <-time.After(time.Second):
t.Fatal("timed out waiting for first rpc to start")
}
secondReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, secondReqID, 3, &tg.HelpGetConfigRequest{})
// 第一条故意忽略 context，使第二条越过自身从入队起计算的 deadline 后才有机会出队。
time.Sleep(120 * time.Millisecond)
close(handler.releaseFirst)
result := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, secondReqID)
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode rpc timeout: %v", err)
}
if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "RPC_TIMEOUT" {
t.Fatalf("rpc_error = %d %q, want 500 RPC_TIMEOUT", rpcErr.ErrorCode, rpcErr.ErrorMessage)
}
if calls := handler.calls.Load(); calls != 1 {
t.Fatalf("handler calls = %d, want 1 (expired queued RPC must not dispatch)", calls)
}
}
func TestInboundRPCRunningDeadlineReturnsExactlyOneTimeout(t *testing.T) {
for _, tc := range []struct {
name string
honorContext bool
}{
{name: "handler_honors_context", honorContext: true},
{name: "handler_temporarily_ignores_context", honorContext: false},
} {
t.Run(tc.name, func(t *testing.T) {
const dc = 2
handler := &runningDeadlineRPC{
started: make(chan struct{}),
release: make(chan struct{}),
honorContext: tc.honorContext,
}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCTimeout: 60 * time.Millisecond,
RPCGlobalWorkers: 1,
})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
reqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, reqID, 1, &tg.HelpGetConfigRequest{})
select {
case <-handler.started:
case <-time.After(time.Second):
t.Fatal("timed out waiting for running rpc")
}
// In the ignore-context case this result must arrive before release is closed: the
// scheduler deadline, not eventual handler return, owns the timeout response.
result := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, reqID)
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode running rpc timeout: %v", err)
}
if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "RPC_TIMEOUT" {
t.Fatalf("rpc_error = %d %q, want 500 RPC_TIMEOUT", rpcErr.ErrorCode, rpcErr.ErrorMessage)
}
close(handler.release)
})
}
}
func TestRPCResponseGateExactlyOnce(t *testing.T) {
for i := 0; i < 100; i++ {
gate := &rpcResponseGate{}
results := make(chan bool, 2)
go func() { results <- gate.tryNormal() }()
go func() { results <- gate.tryTimeout() }()
wins := 0
if <-results {
wins++
}
if <-results {
wins++
}
if wins != 1 {
t.Fatalf("iteration %d response gate winners = %d, want 1", i, wins)
}
}
}
func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) {
const dc = 2
handler := &countingConfigRPC{}
@ -208,6 +320,40 @@ func (h *blockingRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.B
func (h *blockingRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
type queueDeadlineRPC struct {
calls atomic.Int32
firstStarted chan struct{}
releaseFirst chan struct{}
}
func (h *queueDeadlineRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) {
if h.calls.Add(1) == 1 {
close(h.firstStarted)
<-h.releaseFirst
}
return &tg.Config{ThisDC: 2}, nil
}
func (h *queueDeadlineRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
type runningDeadlineRPC struct {
started chan struct{}
release chan struct{}
honorContext bool
}
func (h *runningDeadlineRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
close(h.started)
if h.honorContext {
<-ctx.Done()
return nil, ctx.Err()
}
<-h.release
return &tg.Config{ThisDC: 2}, nil
}
func (h *runningDeadlineRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
type canceledInternalRPC struct {
calls atomic.Int32
firstDone chan struct{}

126
internal/mtprotoedge/same_port_mux.go
View file

@ -40,6 +40,13 @@ type samePortMux struct {
closed chan struct{}
once sync.Once
// sniffing contains only sockets still owned by dispatch while it reads the first four
// bytes. Keeping an explicit registry lets Close interrupt every slow-loris read without a
// second cancellation goroutine per raw connection. A socket is removed under sniffMu before
// successful child-listener hand-off, establishing the ownership barrier.
sniffMu sync.Mutex
sniffing map[net.Conn]struct{}
}
func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux {
@ -51,6 +58,7 @@ func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux
addr: base.Addr(),
sniffTimeout: sniffTimeout,
closed: make(chan struct{}),
sniffing: make(map[net.Conn]struct{}),
}
m.tcp = newSamePortMuxListener(m.addr, m.closed)
m.http = newSamePortMuxListener(m.addr, m.closed)
@ -69,7 +77,15 @@ func (m *samePortMux) HTTP() net.Listener {
func (m *samePortMux) Serve(ctx context.Context) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
// Every exit path must publish cancellation and close both child listeners before waiting
// for sniff/delivery goroutines. A permanent Accept error can otherwise leave a dispatch
// blocked on a full child backlog while the old defer order waits for it before canceling.
var wg sync.WaitGroup
defer func() {
cancel()
_ = m.Close()
wg.Wait()
}()
go func() {
<-ctx.Done()
@ -77,17 +93,23 @@ func (m *samePortMux) Serve(ctx context.Context) error {
}()
// 每条连接一个窥探 goroutine：wg 让 Serve 在退出前等待在途窥探把连接交接完成。
var wg sync.WaitGroup
defer wg.Wait()
var tempDelay time.Duration
for {
conn, err := m.base.Accept()
if err != nil {
if ctx.Err() != nil || isSamePortMuxClosed(m.closed) || isNetClosed(err) {
return nil
}
if isTemporaryAcceptError(err) {
tempDelay = nextAcceptRetryDelay(tempDelay)
if !waitAcceptRetry(ctx, tempDelay) {
return nil
}
continue
}
return err
}
tempDelay = 0
wg.Add(1)
go func() {
defer wg.Done()
@ -99,6 +121,19 @@ func (m *samePortMux) Serve(ctx context.Context) error {
func (m *samePortMux) Close() error {
m.once.Do(func() {
close(m.closed)
// Snapshot under the ownership lock, then close outside it. finishSniff observes
// m.closed and refuses hand-off even after the map is cleared, so dispatch cannot race
// this snapshot and deliver a socket that Close is about to terminate.
m.sniffMu.Lock()
sniffing := make([]net.Conn, 0, len(m.sniffing))
for conn := range m.sniffing {
sniffing = append(sniffing, conn)
delete(m.sniffing, conn)
}
m.sniffMu.Unlock()
for _, conn := range sniffing {
_ = conn.Close()
}
_ = m.tcp.Close()
_ = m.http.Close()
_ = m.base.Close()
@ -109,6 +144,20 @@ func (m *samePortMux) Close() error {
// dispatch 窥探单条连接的前 4 字节并把它交给 tcp 或 http 子 listener。窥探带 sniffTimeout
// 读上界，慢/半开连接最多占用本 goroutine sniffTimeout 后即被回收。
func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
// SetReadDeadline bounds an otherwise healthy slow-loris connection, but Close only owns
// the base listener, not sockets Accept has already returned. Register temporary ownership so
// mux shutdown can close this read immediately. finishSniff removes the socket before hand-off.
if !m.beginSniff(conn) {
_ = conn.Close()
return
}
finishedSniff := false
defer func() {
if !finishedSniff {
m.finishSniff(conn)
}
}()
var header [4]byte
if err := conn.SetReadDeadline(time.Now().Add(m.sniffTimeout)); err != nil {
_ = conn.Close()
@ -118,6 +167,14 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
_ = conn.Close()
return
}
// From this point onward deliver/child-listener closure owns cancellation. Removing the
// registry entry under sniffMu is the hand-off barrier: Close either captured and closed this
// socket, or it can no longer find it. A concurrently closed mux refuses delivery.
if !m.finishSniff(conn) {
_ = conn.Close()
return
}
finishedSniff = true
if err := conn.SetReadDeadline(time.Time{}); err != nil {
_ = conn.Close()
return
@ -137,6 +194,30 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
}
}
func (m *samePortMux) beginSniff(conn net.Conn) bool {
m.sniffMu.Lock()
defer m.sniffMu.Unlock()
if isSamePortMuxClosed(m.closed) {
return false
}
if m.sniffing == nil {
m.sniffing = make(map[net.Conn]struct{})
}
m.sniffing[conn] = struct{}{}
return true
}
// finishSniff returns true only when dispatch still owned the socket and the mux remained open
// through the ownership barrier. A false result means Close captured the socket; dispatch must
// not hand it to a child listener.
func (m *samePortMux) finishSniff(conn net.Conn) bool {
m.sniffMu.Lock()
defer m.sniffMu.Unlock()
_, owned := m.sniffing[conn]
delete(m.sniffing, conn)
return owned && !isSamePortMuxClosed(m.closed)
}
// isHTTPHeaderPrefix 判断前 4 字节是否是 HTTP 请求行起始。
//
// 这里只认 GET/POST/HEAD/OPTI，与 gotd generateInit 排除的前缀集合「严格对齐」：合法的
@ -243,6 +324,10 @@ type samePortMuxListener struct {
ch chan net.Conn
closed chan struct{}
once sync.Once
deliveryMu sync.Mutex
closing bool
deliveryWG sync.WaitGroup
}
func newSamePortMuxListener(addr net.Addr, parentClosed <-chan struct{}) *samePortMuxListener {
@ -273,7 +358,25 @@ func (l *samePortMuxListener) Accept() (net.Conn, error) {
func (l *samePortMuxListener) Close() error {
l.once.Do(func() {
// Add and Wait on a WaitGroup must not race while the counter may still be zero.
// The delivery gate serializes the final Add with the transition to closing; after
// closing becomes true no producer can enter, so waiting and draining are safe.
l.deliveryMu.Lock()
l.closing = true
close(l.closed)
l.deliveryMu.Unlock()
l.deliveryWG.Wait()
for {
select {
case conn := <-l.ch:
if conn != nil {
_ = conn.Close()
}
default:
return
}
}
})
return nil
}
@ -283,6 +386,11 @@ func (l *samePortMuxListener) Addr() net.Addr {
}
func (l *samePortMuxListener) deliver(ctx context.Context, conn net.Conn) bool {
if !l.beginDelivery() {
return false
}
defer l.deliveryWG.Done()
select {
case <-l.closed:
return false
@ -292,3 +400,13 @@ func (l *samePortMuxListener) deliver(ctx context.Context, conn net.Conn) bool {
return false
}
}
func (l *samePortMuxListener) beginDelivery() bool {
l.deliveryMu.Lock()
defer l.deliveryMu.Unlock()
if l.closing {
return false
}
l.deliveryWG.Add(1)
return true
}

240
internal/mtprotoedge/same_port_mux_test.go Normal file
View file

@ -0,0 +1,240 @@
package mtprotoedge
import (
"bytes"
"context"
"errors"
"io"
"net"
"testing"
"time"
)
func TestSamePortMuxListenerCloseWaitsAndReturnsBacklogAdmission(t *testing.T) {
admission := newAdmissionController(4, 4, 1)
listener := &samePortMuxListener{
addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)},
ch: make(chan net.Conn, 1),
closed: make(chan struct{}),
}
backlog, backlogPeer := trackedMuxPipe(t, admission, 1001)
defer backlogPeer.Close()
if !listener.deliver(context.Background(), backlog) {
t.Fatal("initial backlog delivery was rejected")
}
// Deterministically model a producer that passed the delivery gate but has not yet
// completed. Close must publish closed first, then wait before draining the backlog.
if !listener.beginDelivery() {
t.Fatal("in-flight delivery gate unexpectedly closed")
}
closeDone := make(chan struct{})
go func() {
_ = listener.Close()
close(closeDone)
}()
select {
case <-listener.closed:
case <-time.After(time.Second):
t.Fatal("Close did not publish listener closure")
}
select {
case <-closeDone:
t.Fatal("Close returned before in-flight delivery completed")
default:
}
listener.deliveryWG.Done()
select {
case <-closeDone:
case <-time.After(time.Second):
t.Fatal("Close did not finish after delivery completed")
}
assertAdmissionConnections(t, admission, 0)
late, latePeer := trackedMuxPipe(t, admission, 1002)
defer latePeer.Close()
if listener.deliver(context.Background(), late) {
t.Fatal("delivery after Close unexpectedly succeeded")
}
_ = late.Close() // dispatch owns and closes a rejected delivery.
assertAdmissionConnections(t, admission, 0)
}
func TestSamePortMuxPermanentAcceptErrorCancelsBlockedDeliveryBeforeWait(t *testing.T) {
serverSide, clientSide := net.Pipe()
defer clientSide.Close()
wantErr := errors.New("same-port permanent accept failure")
base := &connThenErrorListener{conn: serverSide, err: wantErr}
closed := make(chan struct{})
mux := &samePortMux{
base: base,
addr: base.Addr(),
sniffTimeout: time.Hour,
closed: closed,
}
// An unbuffered child listener deterministically leaves dispatch blocked in deliver: no
// consumer is running, and the base listener immediately returns a permanent second error.
mux.tcp = &samePortMuxListener{addr: mux.addr, ch: make(chan net.Conn), closed: make(chan struct{})}
mux.http = &samePortMuxListener{addr: mux.addr, ch: make(chan net.Conn), closed: make(chan struct{})}
writeDone := make(chan error, 1)
go func() {
_, err := clientSide.Write([]byte{0xef, 0, 0, 0})
writeDone <- err
}()
serveDone := make(chan error, 1)
go func() {
serveDone <- mux.Serve(context.Background())
}()
select {
case err := <-serveDone:
if !errors.Is(err, wantErr) {
t.Fatalf("Serve error = %v, want %v", err, wantErr)
}
case <-time.After(time.Second):
t.Fatal("same-port Serve waited for blocked delivery before canceling it")
}
select {
case <-writeDone:
case <-time.After(time.Second):
t.Fatal("sniff writer remained blocked after same-port shutdown")
}
}
func TestSamePortMuxShutdownInterruptsSlowSniffImmediately(t *testing.T) {
tests := []struct {
name string
shutdown func(context.CancelFunc, *samePortMux)
}{
{
name: "context cancel",
shutdown: func(cancel context.CancelFunc, _ *samePortMux) {
cancel()
},
},
{
name: "mux close",
shutdown: func(_ context.CancelFunc, mux *samePortMux) {
_ = mux.Close()
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
base, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
mux := newSamePortMux(base, time.Minute)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
serveDone := make(chan error, 1)
go func() { serveDone <- mux.Serve(ctx) }()
peer, err := net.Dial("tcp", base.Addr().String())
if err != nil {
t.Fatalf("dial: %v", err)
}
defer peer.Close()
// No bytes are written: dispatch is blocked in the four-byte sniff with a one-minute
// deadline. Shutdown must close this accepted socket instead of waiting for it.
tt.shutdown(cancel, mux)
select {
case err := <-serveDone:
if err != nil {
t.Fatalf("Serve: %v", err)
}
case <-time.After(time.Second):
t.Fatal("Serve waited for the sniff deadline after shutdown")
}
if err := peer.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set peer deadline: %v", err)
}
var one [1]byte
if _, err := peer.Read(one[:]); err == nil {
t.Fatal("slow sniff socket remained open after mux shutdown")
}
})
}
}
func TestSamePortMuxSuccessfulHandoffReleasesSniffOwnership(t *testing.T) {
base, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
mux := newSamePortMux(base, time.Minute)
ctx, cancel := context.WithCancel(context.Background())
serveDone := make(chan error, 1)
go func() { serveDone <- mux.Serve(ctx) }()
peer, err := net.Dial("tcp", base.Addr().String())
if err != nil {
t.Fatalf("dial: %v", err)
}
defer peer.Close()
if _, err := peer.Write([]byte{0xef, 0, 0, 0}); err != nil {
t.Fatalf("write sniff prefix: %v", err)
}
accepted, err := mux.TCP().Accept()
if err != nil {
t.Fatalf("accept child: %v", err)
}
defer accepted.Close()
// Once dispatch has delivered the Conn, canceling the mux may close listeners/backlog but
// must not let the old sniff watcher close a socket now owned by the child consumer.
cancel()
select {
case err := <-serveDone:
if err != nil {
t.Fatalf("Serve: %v", err)
}
case <-time.After(time.Second):
t.Fatal("Serve did not stop after cancel")
}
if _, err := peer.Write([]byte{1, 2, 3, 4}); err != nil {
t.Fatalf("write after handoff/shutdown: %v", err)
}
if err := accepted.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set accepted deadline: %v", err)
}
got := make([]byte, 8)
if _, err := io.ReadFull(accepted, got); err != nil {
t.Fatalf("read handed-off connection: %v", err)
}
want := []byte{0xef, 0, 0, 0, 1, 2, 3, 4}
if !bytes.Equal(got, want) {
t.Fatalf("handed-off bytes = %x, want %x", got, want)
}
}
func trackedMuxPipe(t *testing.T, admission *admissionController, port int) (net.Conn, net.Conn) {
t.Helper()
server, peer := net.Pipe()
release, ok := admission.acquireConnection(&net.TCPAddr{
IP: net.ParseIP("203.0.113.20"),
Port: port,
})
if !ok {
_ = server.Close()
_ = peer.Close()
t.Fatal("test connection admission rejected")
}
return &admittedConn{Conn: server, release: release}, peer
}
func assertAdmissionConnections(t *testing.T, admission *admissionController, want int) {
t.Helper()
admission.mu.Lock()
got := admission.connections
byIP := len(admission.byIP)
admission.mu.Unlock()
if got != want || (want == 0 && byIP != 0) {
t.Fatalf("admission state = connections:%d by_ip:%d, want connections:%d", got, byIP, want)
}
}

315
internal/mtprotoedge/server.go
View file

@ -47,7 +47,9 @@ type RPCHandler interface {
type Options struct {
// Logger 日志器。默认 zap.NewNop()。
Logger *zap.Logger
// Codec 传输 codec 构造器。nil 表示自动探测（intermediate/abridged/full）。
// Codec 传输 codec 构造器。nil 表示自动探测（intermediate/abridged/full）。自定义
// codec 必须是 gotd 内置四种 codec（可包 NoHeader），或实现 InboundFrameBudgetedCodec；
// 无法在 payload 分配前预检长度的 codec 会 fail-closed。
Codec func() transport.Codec
// ObfuscatedTCP 先按 MTProto TCP obfuscation 解包，再自动探测 codec。
// Telegram Desktop 的 tcpo_only endpoint 会走这个 64 字节前缀流程。
@ -72,12 +74,45 @@ type Options struct {
HandshakeMaxDuration time.Duration
// WriteTimeout 单次写入超时。默认 30s。
WriteTimeout time.Duration
// MaxConnections 是进程接受的 raw 物理连接总上限，覆盖 codec sniff、握手和
// 已认证连接的完整生命周期。默认 200000；负数表示不限制。
MaxConnections int
// MaxConnectionsPerIP 是单 remote IP 的 raw 物理连接上限。默认 4096，
// 为共享 NAT 与 TDesktop 多候选连接保留足够突发；负数表示不限制。
MaxConnectionsPerIP int
// MaxConcurrentHandshakes 是同时执行 auth_key_id=0 RSA/DH exchange 的上限。
// 达限时已完成 transport framing 的连接收到 -429 后断开。默认 256；负数表示不限制。
MaxConcurrentHandshakes int
// RPCMaxInflight 是单连接同时处理的 RPC 上限。默认 32。
RPCMaxInflight int
// RPCQueueSize 是单连接等待处理的 RPC 队列长度。默认 256。
// RPCQueueSize 是单连接等待处理的 RPC 队列长度。默认 64；队列按首条请求懒分配。
RPCQueueSize int
// RPCTimeout 是单个 RPC 在连接层的最大处理时长。默认 30s。
// 超时从 Copy 前预算/入队开始计算，排队时间包含在内。
RPCTimeout time.Duration
// RPCGlobalWorkers 是 Server 共享 inbound RPC worker 数。默认 256。
RPCGlobalWorkers int
// RPCGlobalMaxTasks 是全进程已预留、排队和执行中的 RPC 条数上限。默认 8192。
RPCGlobalMaxTasks int
// RPCGlobalMaxBytes 是上述 RPC body 的总字节预算。默认 512 MiB。
RPCGlobalMaxBytes int64
// InboundFrameGlobalMaxBytes 是所有物理连接当前正在处理的 transport wire buffer
// 与最大解密 plaintext buffer 的总预算。长度前缀读取后、payload 分配前预留，默认
// 512 MiB；非正值使用默认值。
InboundFrameGlobalMaxBytes int64
// OutboundQueueSize / OutboundControlQueueSize 是每连接普通与控制 mailbox 容量。
// 默认 128/32；控制队列在 actor 中保持严格优先。
OutboundQueueSize int
OutboundControlQueueSize int
// OutboundTrackedGlobalMaxBytes 是所有连接为 msg_resend_req 保留的 RPC/update body
// 总预算。默认 512 MiB；编码后的 MTProto service frame 与控制向量另用 64 MiB
// control budget（包括需 resend tracking 的 new_session_created 等），避免 body 压力
// 阻断连接维持消息。可靠响应无法 tracking 时终止该连接，durable best-effort update
// 则只丢在线加速并由 difference 恢复。
OutboundTrackedGlobalMaxBytes int64
// OutboundWriteGlobalMaxBytes bounds concurrent encrypted wire/codec/obfuscation scratch.
// Scratch is shared and pooled across connections; default 512 MiB.
OutboundWriteGlobalMaxBytes int64
// DC 是本 server 的 DC ID。默认 2。
DC int
@ -115,15 +150,48 @@ func (o *Options) setDefaults() {
if o.WriteTimeout == 0 {
o.WriteTimeout = 30 * time.Second
}
if o.MaxConnections == 0 {
o.MaxConnections = defaultMaxConnections
}
if o.MaxConnectionsPerIP == 0 {
o.MaxConnectionsPerIP = defaultMaxConnectionsPerIP
}
if o.MaxConcurrentHandshakes == 0 {
o.MaxConcurrentHandshakes = defaultMaxConcurrentHandshakes
}
if o.RPCMaxInflight <= 0 {
o.RPCMaxInflight = 32
}
if o.RPCQueueSize <= 0 {
o.RPCQueueSize = 256
o.RPCQueueSize = 64
}
if o.RPCTimeout == 0 {
o.RPCTimeout = 30 * time.Second
}
if o.RPCGlobalWorkers <= 0 {
o.RPCGlobalWorkers = 256
}
if o.RPCGlobalMaxTasks <= 0 {
o.RPCGlobalMaxTasks = 8192
}
if o.RPCGlobalMaxBytes <= 0 {
o.RPCGlobalMaxBytes = 512 << 20
}
if o.InboundFrameGlobalMaxBytes <= 0 {
o.InboundFrameGlobalMaxBytes = defaultInboundFrameGlobalMaxBytes
}
if o.OutboundQueueSize <= 0 {
o.OutboundQueueSize = defaultOutboundQueueSize
}
if o.OutboundControlQueueSize <= 0 {
o.OutboundControlQueueSize = defaultOutboundControlQueueSize
}
if o.OutboundTrackedGlobalMaxBytes <= 0 {
o.OutboundTrackedGlobalMaxBytes = defaultOutboundTrackedMaxBytes
}
if o.OutboundWriteGlobalMaxBytes <= 0 {
o.OutboundWriteGlobalMaxBytes = defaultOutboundWriteMaxBytes
}
if o.DC == 0 {
o.DC = 2
}
@ -150,30 +218,38 @@ func (o *Options) setDefaults() {
// 接受连接、协商 codec、完成密钥交换、解密并分发加密消息到 RPC 路由，处理服务消息，
// 并把活跃连接注册到 SessionManager 以支持主动推送（updates 等）。不含业务逻辑。
type Server struct {
log *zap.Logger
codec func() transport.Codec
obfuscated bool
websocket bool
websocketOrigins []string
readTimeout time.Duration
handshakeTimeout time.Duration
handshakeMaxDur time.Duration
writeTimeout time.Duration
rpcInflight int
rpcQueueSize int
rpcTimeout time.Duration
log *zap.Logger
codec func() transport.Codec
obfuscated bool
websocket bool
websocketOrigins []string
readTimeout time.Duration
handshakeTimeout time.Duration
handshakeMaxDur time.Duration
writeTimeout time.Duration
rpcInflight int
rpcQueueSize int
rpcTimeout time.Duration
rpcScheduler *inboundRPCScheduler
frameBudget *inboundFrameBudget
outboundQueueSize int
outboundControlQueueSize int
outboundTrackedBudget *outboundTrackedBudget
outboundControlBudget *outboundTrackedBudget
outboundScratchPool *outboundScratchPool
dc int
key exchange.PrivateKey
authKeys store.AuthKeyStore
sessions store.SessionStore
conns *SessionManager
rpc RPCHandler
metrics Metrics
cipher crypto.Cipher
clock clock.Clock
rand io.Reader
types *tmap.Map
dc int
key exchange.PrivateKey
authKeys store.AuthKeyStore
sessions store.SessionStore
conns *SessionManager
rpc RPCHandler
metrics Metrics
cipher crypto.Cipher
clock clock.Clock
rand io.Reader
types *tmap.Map
admission *admissionController
rpcResults *rpcResultCache
@ -189,30 +265,38 @@ func New(opts Options) *Server {
conns = NewSessionManager(opts.Logger.Named("sessions"))
}
return &Server{
log: opts.Logger,
codec: opts.Codec,
obfuscated: opts.ObfuscatedTCP,
websocket: opts.WebSocket,
websocketOrigins: append([]string(nil), opts.WebSocketAllowedOrigins...),
readTimeout: opts.ReadTimeout,
handshakeTimeout: opts.HandshakeIdleTimeout,
handshakeMaxDur: opts.HandshakeMaxDuration,
writeTimeout: opts.WriteTimeout,
rpcInflight: opts.RPCMaxInflight,
rpcQueueSize: opts.RPCQueueSize,
rpcTimeout: opts.RPCTimeout,
dc: opts.DC,
key: exchange.PrivateKey{RSA: opts.RSAKey},
authKeys: opts.AuthKeys,
sessions: opts.Sessions,
conns: conns,
rpc: opts.RPC,
metrics: opts.Metrics,
cipher: crypto.NewServerCipher(opts.Rand),
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
rpcResults: newRPCResultCache(opts.Clock.Now),
log: opts.Logger,
codec: opts.Codec,
obfuscated: opts.ObfuscatedTCP,
websocket: opts.WebSocket,
websocketOrigins: append([]string(nil), opts.WebSocketAllowedOrigins...),
readTimeout: opts.ReadTimeout,
handshakeTimeout: opts.HandshakeIdleTimeout,
handshakeMaxDur: opts.HandshakeMaxDuration,
writeTimeout: opts.WriteTimeout,
rpcInflight: opts.RPCMaxInflight,
rpcQueueSize: opts.RPCQueueSize,
rpcTimeout: opts.RPCTimeout,
rpcScheduler: newInboundRPCScheduler(opts.RPCGlobalWorkers, opts.RPCGlobalMaxTasks, opts.RPCGlobalMaxBytes),
frameBudget: newInboundFrameBudget(opts.InboundFrameGlobalMaxBytes),
outboundQueueSize: opts.OutboundQueueSize,
outboundControlQueueSize: opts.OutboundControlQueueSize,
outboundTrackedBudget: newOutboundTrackedBudget(opts.OutboundTrackedGlobalMaxBytes),
outboundControlBudget: newOutboundTrackedBudget(defaultOutboundControlMaxBytes),
outboundScratchPool: newOutboundScratchPool(opts.OutboundWriteGlobalMaxBytes),
dc: opts.DC,
key: exchange.PrivateKey{RSA: opts.RSAKey},
authKeys: opts.AuthKeys,
sessions: opts.Sessions,
conns: conns,
rpc: opts.RPC,
metrics: opts.Metrics,
cipher: crypto.NewServerCipher(opts.Rand),
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
rpcResults: newRPCResultCache(opts.Clock.Now),
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
}
}
@ -224,27 +308,40 @@ func (s *Server) Conns() *SessionManager {
// newConn 基于一次解密结果创建一个可发送的连接对象。
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
c := &Conn{
transport: tc,
writer: tc,
cipher: s.cipher,
msgID: proto.NewMessageIDGen(s.clock.Now),
writeTimeout: s.writeTimeout,
metrics: s.metrics,
authKeyID: key.ID,
authKeyHex: hex.EncodeToString(key.ID[:]),
sessionID: sessionID,
salt: salt,
key: key,
createdAt: s.clock.Now(),
transport: tc,
writer: tc,
cipher: s.cipher,
msgID: proto.NewMessageIDGen(s.clock.Now),
writeTimeout: s.writeTimeout,
metrics: s.metrics,
authKeyID: key.ID,
authKeyHex: hex.EncodeToString(key.ID[:]),
sessionID: sessionID,
salt: salt,
key: key,
createdAt: s.clock.Now(),
outboundQueueSize: s.outboundQueueSize,
outboundControlQueueSize: s.outboundControlQueueSize,
outboundTrackedBudget: s.outboundTrackedBudget,
outboundControlTrackedBudget: s.outboundControlBudget,
outboundScratchPool: s.outboundScratchPool,
}
c.startOutbound()
c.startInboundRPCScheduler(s.rpcInflight, s.rpcQueueSize, s.rpcTimeout)
c.startInboundRPCScheduler(s.rpcScheduler, s.rpcInflight, s.rpcQueueSize, s.rpcTimeout)
return c
}
// Serve 在 ln 上运行 MTProto 连接循环，直到 ctx 取消或发生不可恢复错误。
// ctx 取消时优雅退出：关闭 listener 并等待在途连接处理结束。
func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
// 共享 worker 池只在 Server 真正 Serve 后允许消费，并在首条 RPC 到达时懒启动。
// serveTCP/serveMixed 返回前会等待连接 goroutine 收敛，各 Conn 已先排空/取消任务；
// 最后再停止全局池，避免关闭过程中留下无人消费但仍占预算的队列。
s.rpcScheduler.start()
defer s.rpcScheduler.stop(rpcCloseWaitTimeout)
// 只在最外层 listener 包一次，确保 same-port mux 的 sniff/HTTP upgrade 也计入
// raw admission，而不是等连接已经分流后才计数。
ln = s.admission.wrapListener(ln)
if s.websocket {
return s.serveMixed(ctx, ln)
}
@ -292,11 +389,15 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
)
defer s.log.Info("Stopped")
go func() {
<-ctx.Done()
stopAll := func() {
cancel()
_ = mux.Close()
_ = httpServer.Close()
_ = wsLn.Close()
}
go func() {
<-ctx.Done()
stopAll()
}()
errCh := make(chan error, 4)
@ -330,17 +431,19 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
errCh <- nil
}()
// The four services form one lifecycle: even a clean/closed-listener return from any one
// component means the remaining three can no longer make forward progress as a complete
// same-port server. Stop them immediately, then collect their terminal results.
var firstErr error
for i := 0; i < 4; i++ {
if err := <-errCh; err != nil {
firstErr = err
}
stopAll()
for i := 1; i < 4; i++ {
if err := <-errCh; err != nil && firstErr == nil {
firstErr = err
cancel()
}
}
cancel()
_ = mux.Close()
_ = httpServer.Close()
_ = wsLn.Close()
wg.Wait()
return firstErr
}
@ -351,22 +454,39 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
// 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆（裸 MTProto TCP）；WebSocket
// 连接传 false（gotd 升级处理器已完成去混淆）。
func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error {
ctx, cancel := context.WithCancel(ctx)
var wg sync.WaitGroup
defer func() {
// A permanent Accept error is itself a terminal lifecycle event. Cancel accepted
// connections and close the listener before waiting; otherwise a live connection can
// keep the WaitGroup blocked forever and prevent the accept error from being returned.
cancel()
_ = ln.Close()
wg.Wait()
}()
go func() {
<-ctx.Done()
_ = ln.Close()
}()
var wg sync.WaitGroup
defer wg.Wait()
var tempDelay time.Duration
for {
raw, err := ln.Accept()
if err != nil {
if ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
return nil
}
if isTemporaryAcceptError(err) {
tempDelay = nextAcceptRetryDelay(tempDelay)
s.log.Debug("Temporary accept error; retrying", zap.Duration("backoff", tempDelay), zap.Error(err))
if !waitAcceptRetry(ctx, tempDelay) {
return nil
}
continue
}
return fmt.Errorf("accept: %w", err)
}
tempDelay = 0
wg.Add(1)
go func() {
@ -428,7 +548,7 @@ func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, err
if obfuscated {
ln = transport.ObfuscatedListener(ln)
}
return newCompatTransportListener(s.codec, ln).Accept()
return newCompatTransportListener(s.codec, ln, s.frameBudget).Accept()
}
// serveConn 处理单个传输连接：读帧并按 auth_key_id 分流。
@ -444,6 +564,13 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
var current *Conn
defer func() {
// A successful Recv transfers the frame reservation to serveConn. Release it only after
// this stack has stopped using b/plain; transport.Close may have raced us earlier and must
// not return that memory budget prematurely.
releaseInboundFrameOwnership(conn)
// 先同步关闭物理 socket，解除可能阻塞在 writer.Send 的 outbound actor；
// 再停止 logical Conn，避免 Close 等 actor 时反过来等到 write deadline。
_ = conn.Close()
if current != nil {
s.conns.Unregister(current)
current.Close()
@ -468,7 +595,7 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
var replay *bin.Buffer
for {
if replay != nil {
b.ResetTo(replay.Copy())
b.ResetTo(replay.Buf)
replay = nil
} else {
// 建立 session 前（current==nil，握手 + 首个加密消息之前）用较短的 handshakeTimeout
@ -491,11 +618,29 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
}
if authKeyID == emptyAuthKeyID {
releaseHandshake, admitted := s.admission.tryAcquireHandshake()
if !admitted {
if err := s.sendProtoError(ctx, conn, codec.CodeTransportFlood); err != nil {
return err
}
return nil
}
next, err := s.handleExchange(ctx, conn, &b)
releaseHandshake()
if err != nil {
return err
}
replay = next
// Exchange has finished consuming the original transport frame. Drop its
// potentially near-16MiB backing immediately. A replay frame is the gotd
// encrypted-frame copy and keeps the existing frame reservation until it is
// dispatched; a completed handshake has no surviving frame and can release now.
trimOversizedInboundBuffer(&b)
if replay == nil {
releaseInboundFrameOwnership(conn)
} else {
retainInboundFrameBackings(conn, replay)
}
continue
}
@ -513,7 +658,9 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
if err := s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound); err != nil {
return err
}
continue
// -404 对 TDesktop 是 terminal key failure；继续保留 socket 只会允许
// 同一客户端反复触发 AuthKeyStore 查询。回包一次后立即断开。
return nil
}
fetchedKey = &d
}
@ -522,6 +669,20 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
if err != nil {
return err
}
trimOversizedInboundBuffer(&b)
trimOversizedInboundBuffer(&plain)
retainInboundFrameBackings(conn, &b, &plain)
}
}
// 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.
const maxRetainedConnBuffer = 2 << 20
func trimOversizedInboundBuffer(b *bin.Buffer) {
if b != nil && cap(b.Buf) > maxRetainedConnBuffer {
b.Buf = nil
}
}

2
internal/mtprotoedge/server_test.go
View file

@ -358,7 +358,7 @@ func TestSamePortWebSocketTransportRoundTrip(t *testing.T) {
serverDone := make(chan error, 1)
go func() {
l := newCompatTransportListener(nil, wsLn)
l := newCompatTransportListener(nil, wsLn, newInboundFrameBudget(defaultInboundFrameGlobalMaxBytes))
defer func() { _ = l.Close() }()
conn, err := l.Accept()

556
internal/mtprotoedge/session_manager.go
View file

@ -4,7 +4,9 @@ import (
"context"
"errors"
"fmt"
"sort"
"sync"
"sync/atomic"
"time"
"go.uber.org/zap"
@ -36,7 +38,8 @@ const (
// （typing/presence，不写 durable log）经 PushToUserTransient* 在未就绪时直接跳过、不入队，
// 因此本队列被老化/溢出/重试耗尽丢弃时，丢的一定是 durable 条目——getDifference 以
// user_update_events 兜底补齐，丢弃不丢数据。
pendingPushMaxAge = 60 * time.Second
pendingPushMaxAge = 60 * time.Second
defaultPendingPushMaxBytes = int64(256 << 20)
// maxSessionsPerAuthKey：单个 raw auth_key 允许同时在线的 session 上限。telesrv 单 DC，
// 一个客户端的全部连接（主连接 + 并发下载/上传）共享同一 auth_key、各用独立 session_id，
// 故此上限须高于真实客户端单设备的并发连接峰值，否则会误杀活跃下载/主连接：
@ -52,10 +55,51 @@ const (
maxChannelIndexPerSession = 8192
)
// forceCloseBatchTimeout is one deadline for a whole revoke/replace/eviction batch. Conn.Close
// already bounds its inbound-RPC wait, but calling ForceClose serially would multiply that bound
// by the number of sessions. The batch helper starts every close concurrently and waits at most
// this one shared interval.
const forceCloseBatchTimeout = rpcCloseWaitTimeout
// maxForceCloseParallelism caps control-plane close goroutines even if a corrupted/runtime index
// hands a revoke path far more sessions than maxSessionsPerAuthKey. Every Conn's producer/RPC gate
// is closed synchronously before these workers start, so a stuck transport.Close cannot admit more
// memory while the bounded workers continue draining physical sockets in the background.
const maxForceCloseParallelism = 64
type queuedPush struct {
t proto.MessageType
msg bin.Encoder
at time.Time
t proto.MessageType
encoded *encodedOutboundMessage
reservation *pendingPushReservation
at time.Time
}
type pendingPushReservation struct {
budget *outboundTrackedBudget
bytes int
refs atomic.Int32
}
func (r *pendingPushReservation) retain() {
if r == nil {
return
}
if refs := r.refs.Add(1); refs <= 1 {
panic("mtprotoedge: retained released pending push reservation")
}
}
func (r *pendingPushReservation) release() {
if r == nil {
return
}
refs := r.refs.Add(-1)
if refs < 0 {
panic("mtprotoedge: pending push reservation released more than retained")
}
if refs == 0 {
r.budget.release(r.bytes)
}
}
type sessionKey struct {
@ -85,6 +129,7 @@ type SessionManager struct {
bySessionMembers map[sessionKey]map[int64]struct{}
pending map[sessionKey][]queuedPush // updates-ready 前暂存的主动推送
flushing map[sessionKey]bool // 置位时暂存正在排空的 session；排空完成前推送继续进 pending 保序
pendingBudget *outboundTrackedBudget // 未就绪 session 暂存 encoded body 的进程级上限
lifecycle SessionLifecycleObserver
log *zap.Logger
@ -107,6 +152,7 @@ func NewSessionManager(log *zap.Logger) *SessionManager {
bySessionMembers: make(map[sessionKey]map[int64]struct{}),
pending: make(map[sessionKey][]queuedPush),
flushing: make(map[sessionKey]bool),
pendingBudget: newOutboundTrackedBudget(defaultPendingPushMaxBytes),
log: log,
}
}
@ -161,11 +207,13 @@ func (m *SessionManager) Register(c *Conn) {
)
m.mu.Unlock()
if replaced != nil {
replaced.Close()
}
if evicted != nil {
evicted.Close()
// 同 identity 的新物理连接已经原子接管索引；立即关闭旧 transport，不能只停
// actor 后让旧 FD/read goroutine 滞留到 read timeout。replacement 与 cap eviction
// 共用一个并发关闭批次，不能把每条 Conn 的 RPC 等待上界串行相加。
if replaced != nil || evicted != nil {
if !forceCloseConnBatch([]*Conn{replaced, evicted}, forceCloseBatchTimeout) {
m.log.Warn("Session replacement/eviction close exceeded shared deadline")
}
}
}
@ -217,7 +265,12 @@ func (m *SessionManager) DestroySession(sessionID int64) bool {
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
c.Close()
if !forceCloseConnBatch([]*Conn{c}, forceCloseBatchTimeout) {
m.log.Warn("Destroyed session close exceeded shared deadline",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", sessionID),
)
}
if observer != nil && offlineUser != 0 {
observer.SessionOffline(key.authKeyID, sessionID, offlineUser, lastForUser)
}
@ -230,7 +283,7 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
delete(m.pending, key)
m.deletePendingLocked(key)
m.mu.Unlock()
return false
}
@ -243,7 +296,12 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
c.Close()
if !forceCloseConnBatch([]*Conn{c}, forceCloseBatchTimeout) {
m.log.Warn("Destroyed session close exceeded shared deadline",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int64("session_id", sessionID),
)
}
if observer != nil && offlineUser != 0 {
observer.SessionOffline(authKeyID, sessionID, offlineUser, lastForUser)
}
@ -287,7 +345,7 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
c.membershipsSynced.Store(false)
// 身份变化即丢弃暂存推送：它们属于前一个账号，flush 给新账号是跨账号泄露。
// 同时取消进行中的排空（runFlush 还另有 owner 校验做批内兜底）。
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
}
}
@ -298,7 +356,7 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
c.membershipsSynced.Store(false)
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
}
}
@ -399,7 +457,7 @@ func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
c.membershipsSynced.Store(false)
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
c.userID.Store(0)
c.userIDResolved.Store(false)
@ -459,8 +517,11 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int
)
}
m.mu.Unlock()
for _, c := range conns {
c.ForceClose()
if !forceCloseConnBatch(conns, forceCloseBatchTimeout) {
m.log.Warn("Revoked auth-key session close exceeded shared deadline",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int("sessions", len(conns)),
)
}
if observer != nil {
for _, e := range events {
@ -493,8 +554,11 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc
}
observer := m.lifecycle
m.mu.Unlock()
for _, c := range conns {
c.ForceClose()
if !forceCloseConnBatch(conns, forceCloseBatchTimeout) {
m.log.Warn("Raw auth-key session close exceeded shared deadline",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int("sessions", len(conns)),
)
}
if observer != nil {
for _, e := range events {
@ -504,6 +568,99 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc
return len(conns)
}
// forceCloseConnBatch closes every producer/RPC gate first, then closes physical transports with a
// bounded worker set. Physical close and actor/RPC convergence share one batch deadline; the wait is
// never multiplied by the number of sessions. Workers may finish physical closes after the caller's
// deadline, but no timed-out Conn can enqueue more work in that interval. Nil/duplicate entries are
// removed so Register's replacement/eviction slots cannot close the same Conn twice.
func forceCloseConnBatch(conns []*Conn, timeout time.Duration) bool {
if len(conns) == 0 {
return true
}
unique := make([]*Conn, 0, len(conns))
seen := make(map[*Conn]struct{}, len(conns))
for _, c := range conns {
if c == nil {
continue
}
if _, ok := seen[c]; ok {
continue
}
seen[c] = struct{}{}
unique = append(unique, c)
}
if len(unique) == 0 {
return true
}
// This phase is non-blocking and must precede transport.Close: it is the safety boundary if
// an implementation of transport.Conn.Close itself blocks past the batch deadline.
for _, c := range unique {
c.beginTerminalShutdown()
}
workers := min(len(unique), maxForceCloseParallelism)
jobs := make(chan *Conn, len(unique))
for _, c := range unique {
jobs <- c
}
close(jobs)
var closeWG sync.WaitGroup
closeWG.Add(workers)
for range workers {
go func() {
defer closeWG.Done()
for c := range jobs {
c.closeTransport()
}
}()
}
physicalDone := make(chan struct{})
go func() {
closeWG.Wait()
close(physicalDone)
}()
if timeout <= 0 {
return false
}
deadline := time.Now().Add(timeout)
timer := time.NewTimer(time.Until(deadline))
defer timer.Stop()
select {
case <-physicalDone:
case <-timer.C:
return false
}
// All physical close calls returned. Wait for memory-owning actor/RPC work using the same
// deadline; the first genuinely stuck Conn consumes the remaining allowance, not a fresh 5s.
for _, c := range unique {
remaining := time.Until(deadline)
if remaining <= 0 {
return false
}
if c.rpcScheduler != nil && !c.waitInboundShutdown(remaining) {
return false
}
if c.outboundDone == nil {
continue
}
remaining = time.Until(deadline)
if remaining <= 0 {
return false
}
wait := time.NewTimer(remaining)
select {
case <-c.outboundDone:
wait.Stop()
case <-wait.C:
return false
}
}
return true
}
// UnbindAuthKey 清理某业务 auth_key 下所有活跃连接的登录用户缓存。
func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
m.mu.Lock()
@ -520,7 +677,7 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
m.clearChannelMembershipsLocked(c, key)
c.membershipsSynced.Store(false)
// 授权解除后暂存推送属于已登出的账号，不能等下一个登录者置位时 flush 出去。
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
c.userIDResolved.Store(true)
count++
@ -595,7 +752,7 @@ func (m *SessionManager) runFlush(c *Conn, key sessionKey, owner int64, attempt
}
if c.userID.Load() != owner {
// 排空期间发生登出/换号：剩余暂存属于旧账号，丢弃且不得发给新账号。
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
m.mu.Unlock()
return
@ -613,38 +770,47 @@ func (m *SessionManager) runFlush(c *Conn, key sessionKey, owner int64, attempt
// 每条发送前复查身份：登出/换号后 batch 的剩余条目不能继续发到已易主的连接。
if c.userID.Load() != owner {
m.mu.Lock()
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
m.mu.Unlock()
releaseQueuedPushes(batch[i:])
return
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
err := c.Send(ctx, item.t, item.msg)
// Pending entries are durable account updates. Shared body-budget pressure is not
// evidence that this socket is corrupt, so use the non-terminal enqueue path; after
// bounded retries, getDifference is the authoritative recovery path.
err := c.SendBestEffortEncoded(ctx, item.t, item.encoded, 5*time.Second)
cancel()
if err == nil {
item.release()
continue
}
m.mu.Lock()
if cur, ok := m.bySession[key]; !ok || cur != c || !m.flushing[key] || c.userID.Load() != owner {
// 连接换代/取消/易主：剩余 batch 不属于当前连接当前账号，丢弃。
if c.userID.Load() != owner {
delete(m.pending, key)
m.deletePendingLocked(key)
delete(m.flushing, key)
}
m.mu.Unlock()
releaseQueuedPushes(batch[i:])
return
}
rest := append(append([]queuedPush(nil), batch[i:]...), m.pending[key]...)
if len(rest) > maxPendingPushesPerSession {
// 与 queueLocked 溢出策略一致：丢最旧留最新，让 pts 空洞集中在最前端，
// flush 首条即触发客户端 gap 检测，恢复路径最短。
rest = rest[len(rest)-maxPendingPushesPerSession:]
dropped := len(rest) - maxPendingPushesPerSession
releaseQueuedPushes(rest[:dropped])
rest = rest[dropped:]
}
m.pending[key] = rest
if attempt+1 >= maxFlushAttempts {
// 重试用尽：置位激活避免 idle 客户端永久断流；剩余暂存中的 durable 更新
// 由客户端后续 pts 空洞触发 getDifference 补齐。
c.receivesUpdates.Store(true)
m.deletePendingLocked(key)
delete(m.flushing, key)
m.mu.Unlock()
m.log.Debug("Flush gave up after retries; activated with getDifference fallback",
@ -717,41 +883,61 @@ func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, session
// PushToSession 向指定 session 推送一条消息。
func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.Lock()
m.mu.RLock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous {
m.mu.Unlock()
m.mu.RUnlock()
return ErrSessionAmbiguous
}
if !ok {
m.mu.Unlock()
m.mu.RUnlock()
return ErrSessionNotFound
}
if !c.receivesUpdates.Load() {
m.queueLocked(key, t, msg)
m.mu.Unlock()
return nil
ready := c.receivesUpdates.Load()
m.mu.RUnlock()
if ready {
return c.Send(ctx, t, msg)
}
m.mu.Unlock()
return c.Send(ctx, t, msg)
return m.queueOrSendPrepared(ctx, key, t, msg)
}
// PushToSessionForAuthKey 向指定 raw auth_key_id + session_id 推送一条消息。
func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.Lock()
m.mu.RLock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
m.mu.RUnlock()
return ErrSessionNotFound
}
ready := c.receivesUpdates.Load()
m.mu.RUnlock()
if ready {
return c.Send(ctx, t, msg)
}
return m.queueOrSendPrepared(ctx, key, t, msg)
}
func (m *SessionManager) queueOrSendPrepared(ctx context.Context, key sessionKey, t proto.MessageType, msg bin.Encoder) error {
encoded, reservation, err := m.preparePendingPush(ctx, msg)
if err != nil {
return err
}
defer reservation.release()
m.mu.Lock()
c, ok := m.bySession[key]
if !ok {
m.mu.Unlock()
return ErrSessionNotFound
}
if !c.receivesUpdates.Load() {
m.queueLocked(key, t, msg)
_ = m.queuePreparedLocked(key, t, encoded, reservation)
m.mu.Unlock()
return nil
}
m.mu.Unlock()
return c.Send(ctx, t, msg)
return c.SendEncoded(ctx, t, encoded)
}
// PushToSessionForAuthKeyImmediate 向指定 raw auth_key_id + session_id 立即推送一条消息。
@ -782,7 +968,7 @@ func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, ex
return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg)
}
// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定业务 auth_key + session。
// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定 raw auth_key + session。
func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg)
}
@ -793,23 +979,28 @@ func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, use
// 漏 temp-key 设备）。未就绪连接跳过、不进 pending——密聊消息 durable 在 qts 队列，
// 离线设备靠 getDifference 补回（在线推送只是加速器）。c.userID 复查防跨账号泄露。
func (m *SessionManager) PushToUserAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder) (int, error) {
getEncoded := onceEncodedOutbound(msg)
return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, false, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
if err != nil {
return err
}
return c.SendEncoded(ctx, t, encoded)
})
// Secret-chat qts is the durable source of truth, so online delivery is an accelerator just
// like account pts fan-out. Do not synchronously wait for every PFS/raw connection's socket.
return m.pushToBusinessAuthKeyBestEffort(ctx, userID, businessAuthKeyID, t, msg, 2*time.Second)
}
// PushToUserAuthKeyTransient 是 PushToUserAuthKey 的 transient（typing）best-effort 版本。
func (m *SessionManager) PushToUserAuthKeyTransient(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(msg)
return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, true, func(c *Conn) error {
return m.pushToBusinessAuthKeyBestEffort(ctx, userID, businessAuthKeyID, t, msg, timeout)
}
func (m *SessionManager) pushToBusinessAuthKeyBestEffort(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(ctx, msg)
var deadline time.Time
if timeout > 0 {
deadline = time.Now().Add(timeout)
}
if ctx != nil {
if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
deadline = ctxDeadline
}
}
return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
@ -817,11 +1008,18 @@ func (m *SessionManager) PushToUserAuthKeyTransient(ctx context.Context, userID
if err != nil {
return err
}
return c.SendBestEffortEncoded(ctx, t, encoded, timeout)
remaining := timeout
if !deadline.IsZero() {
remaining = time.Until(deadline)
if remaining < 0 {
remaining = 0
}
}
return c.SendBestEffortEncoded(ctx, t, encoded, remaining)
})
}
func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, transient bool, send func(*Conn) error) (int, error) {
func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, send func(*Conn) error) (int, error) {
m.mu.Lock()
candidates := m.businessAuthKeyCandidatesLocked(businessAuthKeyID)
conns := make([]*Conn, 0, len(candidates))
@ -836,7 +1034,6 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64
conns = append(conns, c)
}
m.mu.Unlock()
_ = transient
var firstErr error
sent := 0
for _, c := range conns {
@ -845,6 +1042,18 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64
continue
}
if err := send(c); err != nil {
if errors.Is(err, ErrOutboundTrackedBudget) {
// Shared process pressure is not evidence that this particular socket is
// slow. Skip this online accelerator; durable qts/difference is the truth.
continue
}
if errors.Is(err, ErrOutboundQueueFull) {
c.dropSlowConsumer()
continue
}
if errors.Is(err, ErrConnClosed) {
continue
}
if firstErr == nil {
firstErr = err
}
@ -856,7 +1065,7 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64
}
func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
getEncoded := onceEncodedOutbound(msg)
getEncoded := onceEncodedOutbound(ctx, msg)
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
@ -875,7 +1084,7 @@ func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAu
// 下一次状态变化重建，囤积过期 transient 既无意义又会被 pending 的老化/溢出/重试耗尽误当
// 「durable 兜底」丢弃。走 best-effort 发送，不阻塞调用方。
func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(msg)
getEncoded := onceEncodedOutbound(ctx, msg)
return m.pushToUserWithSender(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, false, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
@ -897,7 +1106,19 @@ func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Co
}
func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(msg)
getEncoded := onceEncodedOutbound(ctx, msg)
// timeout 是整次 fan-out 的等待预算，不是每个 session 各自一份。健康连接始终先走
// SendBestEffortEncoded 的非阻塞快路径；预算耗尽后 remaining=0，仍会尝试快路径，
// 但不会再为后续慢连接串行等待。
var deadline time.Time
if timeout > 0 {
deadline = time.Now().Add(timeout)
}
if ctx != nil {
if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
deadline = ctxDeadline
}
}
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
@ -906,18 +1127,25 @@ func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64,
if err != nil {
return err
}
return c.SendBestEffortEncoded(ctx, t, encoded, timeout)
remaining := timeout
if !deadline.IsZero() {
remaining = time.Until(deadline)
if remaining < 0 {
remaining = 0
}
}
return c.SendBestEffortEncoded(ctx, t, encoded, remaining)
})
}
func onceEncodedOutbound(msg bin.Encoder) func() (*encodedOutboundMessage, error) {
func onceEncodedOutbound(ctx context.Context, msg bin.Encoder) func() (*encodedOutboundMessage, error) {
var (
encoded *encodedOutboundMessage
err error
)
return func() (*encodedOutboundMessage, error) {
if encoded == nil && err == nil {
encoded, err = encodeOutboundMessage(msg)
encoded, err = encodeOutboundMessageContext(ctx, msg)
}
return encoded, err
}
@ -957,6 +1185,10 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
}
m.mu.RUnlock()
if needQueue {
// TL encoding and the process-wide pending-byte reservation may be expensive or
// briefly block on the global encode gate. Do both before taking SessionManager.mu,
// then share the immutable body across every not-ready session found by the re-scan.
pendingEncoded, pendingReservation, pendingErr := m.preparePendingPush(ctx, msg)
// 写锁下完整重扫（读锁释放到此之间状态可能变化，以重扫结果为准）。
conns = conns[:0]
queued, dropped, excluded, skipped = 0, 0, 0, 0
@ -972,7 +1204,7 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
skipped++
continue
}
if m.queueLocked(key, t, msg) {
if pendingErr == nil && m.queuePreparedLocked(key, t, pendingEncoded, pendingReservation) {
queued++
if debug {
m.log.Debug("Push queued (session not updates-ready)",
@ -996,6 +1228,15 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
conns = append(conns, c)
}
m.mu.Unlock()
if pendingReservation != nil {
pendingReservation.release() // drop producer ref; queued entries own the body now.
}
if pendingErr != nil && debug {
m.log.Debug("Drop pending pushes outside byte budget",
zap.Int64("user_id", userID),
zap.Error(pendingErr),
)
}
}
var firstErr error
@ -1008,6 +1249,29 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
continue
}
if err := send(c); err != nil {
if errors.Is(err, ErrOutboundTrackedBudget) {
// Do not turn pressure owned by other sockets into a reconnect storm on
// healthy recipients. The durable event remains recoverable by difference.
dropped++
continue
}
// 对 durable/best-effort fan-out，队列满意味着该 socket 已成为慢消费者。
// 立即摘除并把它视为离线：不能让其错误把已经投递给健康 session 的 outbox
// 行整体重试。该 session 的 durable gap 由 getDifference 恢复。
if errors.Is(err, ErrOutboundQueueFull) {
c.dropSlowConsumer()
if debug {
m.log.Debug("Drop slow outbound consumer",
zap.Int64("user_id", userID),
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int64("session_id", c.sessionID),
)
}
continue
}
if errors.Is(err, ErrConnClosed) {
continue
}
if firstErr == nil {
firstErr = err
}
@ -1055,6 +1319,24 @@ func (m *SessionManager) Online() int {
return len(m.bySession)
}
// ActiveRawAuthKeyIDs 返回当前物理连接实际使用的 raw auth_key_id 去重快照。
// maintenance 用它保护“已建 key 但尚未登录”的长连接不被 orphan GC 删除；不能用
// business/temp→perm key 替代，否则活跃 temp 连接仍可能误删。
func (m *SessionManager) ActiveRawAuthKeyIDs() [][8]byte {
m.mu.RLock()
defer m.mu.RUnlock()
seen := make(map[[8]byte]struct{}, len(m.bySession))
out := make([][8]byte, 0, len(m.byAuthKey))
for key := range m.bySession {
if _, ok := seen[key.authKeyID]; ok {
continue
}
seen[key.authKeyID] = struct{}{}
out = append(out, key.authKeyID)
}
return out
}
// IsUserOnline returns whether userID has at least one active connection.
func (m *SessionManager) IsUserOnline(userID int64) bool {
if userID == 0 {
@ -1268,6 +1550,54 @@ func (m *SessionManager) OnlineChannelMemberUserIDsExcluding(channelID int64, ex
return out
}
// OnlineChannelIDsSnapshot returns every channel with at least one live joined-member session in
// strictly ascending order. The global SessionManager lock is held only while copying map keys;
// sorting and all recovery database work happen after unlock. The fixed saturation-recovery actor
// is the sole caller, so its exceptional-path temporary memory is one int64 slice (peak about 8*C
// bytes) rather than repeated O(C) scans under the connection/membership lock.
func (m *SessionManager) OnlineChannelIDsSnapshot() []int64 {
m.mu.RLock()
out := make([]int64, 0, len(m.byMemberChannel))
for channelID, sessions := range m.byMemberChannel {
if channelID <= 0 || len(sessions) == 0 {
continue
}
live := false
for key := range sessions {
if _, ok := m.bySession[key]; ok {
live = true
break
}
}
if !live {
continue
}
out = append(out, channelID)
}
m.mu.RUnlock()
sort.Slice(out, func(i, j int) bool { return out[i] < out[j] })
return out
}
// OnlineChannelIDsAfter is retained for bounded diagnostics/tests. Production recovery takes one
// OnlineChannelIDsSnapshot per generation and slices it into pages, avoiding repeated full scans.
func (m *SessionManager) OnlineChannelIDsAfter(afterChannelID int64, limit int) []int64 {
if limit <= 0 {
return nil
}
const maxRecoveryPage = 4096
if limit > maxRecoveryPage {
limit = maxRecoveryPage
}
all := m.OnlineChannelIDsSnapshot()
start := sort.Search(len(all), func(i int) bool { return all[i] > afterChannelID })
end := start + limit
if end > len(all) {
end = len(all)
}
return all[start:end]
}
func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64, channelID int64, limit int) []int64 {
if channelID == 0 {
return nil
@ -1314,7 +1644,7 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
if dropPending {
delete(m.pending, key)
m.deletePendingLocked(key)
}
delete(m.flushing, key)
return uid
@ -1430,8 +1760,11 @@ func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedP
now := time.Now()
pending := make([]queuedPush, 0, len(q))
dropped := 0
for _, item := range q {
for i := range q {
item := q[i]
q[i] = queuedPush{}
if now.Sub(item.at) > pendingPushMaxAge {
item.release()
dropped++
continue
}
@ -1447,9 +1780,43 @@ func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedP
return pending
}
// queueLocked 暂存一条主动推送，返回是否实际入队——stale 丢批分支会连同当前
// 这条一起丢弃，调用方据此区分 queued/dropped 计数，避免投递日志失真。
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool {
// preparePendingPush encodes outside SessionManager.mu and reserves the one physical body before
// releasing the process-wide encode slot. Multiple not-ready sessions may then share this
// immutable body via reservation refs instead of encoding/copying it once per session.
func (m *SessionManager) preparePendingPush(ctx context.Context, msg bin.Encoder) (*encodedOutboundMessage, *pendingPushReservation, error) {
var (
encoded *encodedOutboundMessage
bytes int
)
err := withOutboundEncodeSlot(ctx, nil, func() error {
var err error
encoded, err = encodeOutboundMessageWithoutSlot(msg)
if err != nil {
return err
}
if encoded == nil {
return errors.New("nil encoded pending push")
}
bytes = len(encoded.body)
if bytes > maxOutboundBodyBytes {
return fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, bytes, maxOutboundBodyBytes)
}
if !m.pendingBudget.reserve(bytes) {
return ErrOutboundTrackedBudget
}
return nil
})
if err != nil {
return nil, nil, err
}
reservation := &pendingPushReservation{budget: m.pendingBudget, bytes: bytes}
reservation.refs.Store(1) // producer ownership; queue entries retain below.
return encoded, reservation, nil
}
// queuePreparedLocked 暂存一条已编码的主动推送，返回是否实际入队。
// 调用方必须在锁外保持 reservation 的 producer ref，并在全部入队完成后 release。
func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType, encoded *encodedOutboundMessage, reservation *pendingPushReservation) bool {
q := m.pending[key]
// 过期保护：最早一条暂存已超过 pendingPushMaxAge（session 迟迟未 ready）时，丢整批并
// 不再囤这条，记 trace。避免「登录后从不 getState」的连接长期占用 pending 内存。
@ -1459,11 +1826,21 @@ func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bi
zap.Int64("session_id", key.sessionID),
zap.Int("dropped", len(q)),
)
delete(m.pending, key)
m.deletePendingLocked(key)
return false
}
push := queuedPush{t: t, msg: msg, at: time.Now()}
if encoded == nil || reservation == nil {
return false
}
reservation.retain()
push := queuedPush{
t: t,
encoded: encoded,
reservation: reservation,
at: time.Now(),
}
if len(q) >= maxPendingPushesPerSession {
q[0].release()
copy(q, q[1:])
q[len(q)-1] = push
m.pending[key] = q
@ -1473,6 +1850,22 @@ func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bi
return true
}
// queueLocked remains as a test/internal single-target convenience. Production fan-out prepares
// outside m.mu and calls queuePreparedLocked so TL encoding never serializes the session registry.
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool {
encoded, reservation, err := m.preparePendingPush(context.Background(), msg)
if err != nil {
m.log.Debug("Drop pending push outside byte budget",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Error(err),
)
return false
}
defer reservation.release()
return m.queuePreparedLocked(key, t, encoded, reservation)
}
func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey, bool, bool) {
set := m.bySessionID[sessionID]
if len(set) == 0 {
@ -1490,7 +1883,7 @@ func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey
func (m *SessionManager) dropPendingBySessionLocked(sessionID int64) {
for key := range m.pending {
if key.sessionID == sessionID {
delete(m.pending, key)
m.deletePendingLocked(key)
}
}
}
@ -1528,7 +1921,7 @@ func (m *SessionManager) sweepStalePending() {
if len(q) == 0 || now.Sub(q[0].at) <= pendingPushMaxAge {
continue
}
delete(m.pending, key)
m.deletePendingLocked(key)
dropped++
}
if dropped > 0 {
@ -1536,6 +1929,27 @@ func (m *SessionManager) sweepStalePending() {
}
}
func (q *queuedPush) release() {
if q == nil {
return
}
reservation := q.reservation
*q = queuedPush{}
reservation.release()
}
func releaseQueuedPushes(q []queuedPush) {
for i := range q {
q[i].release()
}
}
func (m *SessionManager) deletePendingLocked(key sessionKey) {
q := m.pending[key]
delete(m.pending, key)
releaseQueuedPushes(q)
}
func addConnIndex[K comparable](idx map[K]map[int64]*Conn, key K, sessionID int64, c *Conn) {
set := idx[key]
if set == nil {
@ -1630,7 +2044,7 @@ func shouldExcludeSession(c *Conn, excludeAuthKeyID *[8]byte, excludeSessionID i
if excludeAuthKeyID == nil || *excludeAuthKeyID == ([8]byte{}) {
return true
}
return connUsesBusinessAuthKey(c, *excludeAuthKeyID)
return c.authKeyID == *excludeAuthKeyID
}
func sessionKeyLog(id [8]byte) string {

519
internal/mtprotoedge/session_manager_test.go
View file

@ -3,6 +3,8 @@ package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
@ -27,6 +29,35 @@ type closeCountingTransport struct {
closes int
}
type slowCloseTransport struct {
delay time.Duration
release <-chan struct{}
done chan struct{}
once sync.Once
closes atomic.Int32
}
func newSlowCloseTransport(delay time.Duration, release <-chan struct{}) *slowCloseTransport {
return &slowCloseTransport{delay: delay, release: release, done: make(chan struct{})}
}
func (*slowCloseTransport) Send(context.Context, *bin.Buffer) error {
return errors.New("test transport send")
}
func (*slowCloseTransport) Recv(context.Context, *bin.Buffer) error {
return errors.New("test transport recv")
}
func (t *slowCloseTransport) Close() error {
t.closes.Add(1)
if t.release != nil {
<-t.release
} else if t.delay > 0 {
time.Sleep(t.delay)
}
t.once.Do(func() { close(t.done) })
return nil
}
func (t *closeCountingTransport) Send(context.Context, *bin.Buffer) error {
return errors.New("test transport send")
}
@ -78,6 +109,43 @@ func TestSessionManagerRegistry(t *testing.T) {
}
}
func TestSessionManagerReplacementClosesOldPhysicalTransport(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
oldTransport := &closeCountingTransport{}
old := &Conn{sessionID: 42, authKeyID: raw, transport: oldTransport}
replacement := &Conn{sessionID: 42, authKeyID: raw}
sm.Register(old)
sm.Register(replacement)
if oldTransport.closes != 1 {
t.Fatalf("old transport closes = %d, want 1", oldTransport.closes)
}
// 旧 serveConn 稍后退出时不得把 replacement 从索引删掉。
sm.Unregister(old)
if got, ok := sm.bySession[sessionKey{authKeyID: raw, sessionID: 42}]; !ok || got != replacement {
t.Fatal("old unregister removed the replacement connection")
}
}
func TestSessionManagerDestroyClosesPhysicalTransport(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{4, 5, 6}
physical := &closeCountingTransport{}
c := &Conn{sessionID: 77, authKeyID: raw, transport: physical}
sm.Register(c)
if !sm.DestroySessionForAuthKey(raw, 77) {
t.Fatal("DestroySessionForAuthKey returned false")
}
if physical.closes != 1 {
t.Fatalf("destroyed transport closes = %d, want 1", physical.closes)
}
if sm.Online() != 0 {
t.Fatalf("online after destroy = %d, want 0", sm.Online())
}
}
func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(100)
@ -88,6 +156,7 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
}
c.userID.Store(userID)
c.userIDResolved.Store(true)
@ -115,6 +184,117 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
}
}
func TestSessionManagerPendingFanoutSharesOneEncodedBodyAndBudget(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(102)
keys := make([]sessionKey, 0, 2)
for i := 0; i < 2; i++ {
c := &Conn{sessionID: int64(i + 1), authKeyID: [8]byte{byte(i + 1)}}
c.userID.Store(userID)
c.userIDResolved.Store(true)
sm.Register(c)
keys = append(keys, connSessionKey(c))
}
encodes := 0
msg := &countingOutboundEncoder{count: &encodes}
sent, err := sm.PushToUserExceptSession(context.Background(), userID, 0, proto.MessageFromServer, msg)
if err != nil {
t.Fatalf("push: %v", err)
}
if sent != 2 || encodes != 1 {
t.Fatalf("pending fanout = sent:%d encodes:%d, want 2/1", sent, encodes)
}
sm.mu.Lock()
first := sm.pending[keys[0]][0]
second := sm.pending[keys[1]][0]
if first.encoded != second.encoded || first.reservation != second.reservation {
sm.mu.Unlock()
t.Fatal("pending sessions did not share encoded body/reservation")
}
wantBytes := int64(len(first.encoded.body))
sm.deletePendingLocked(keys[0])
if got := sm.pendingBudget.snapshot(); got != wantBytes {
sm.mu.Unlock()
t.Fatalf("budget after first session drop = %d, want shared body %d", got, wantBytes)
}
sm.deletePendingLocked(keys[1])
sm.mu.Unlock()
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("budget after last session drop = %d, want 0", got)
}
}
func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(101)
// 三个满队列模拟三个慢设备；没有 outbound actor，确保队列在测试期间不会自行排空。
slow := make([]*Conn, 0, 3)
for i := 0; i < 3; i++ {
tr := &closeCountingTransport{}
c := &Conn{
sessionID: int64(i + 1),
authKeyID: [8]byte{byte(i + 1)},
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.outbound <- outboundOp{}
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
sm.Register(c)
slow = append(slow, c)
}
healthy := &Conn{
sessionID: 99,
authKeyID: [8]byte{99},
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
healthy.userID.Store(userID)
healthy.userIDResolved.Store(true)
healthy.receivesUpdates.Store(true)
sm.Register(healthy)
const budget = 40 * time.Millisecond
start := time.Now()
sent, err := sm.PushToUserExceptSessionBestEffort(
context.Background(), userID, 0, proto.MessageFromServer, &tg.UpdatesTooLong{}, budget,
)
elapsed := time.Since(start)
if err != nil {
t.Fatalf("push: %v", err)
}
if sent != 1 {
t.Fatalf("sent = %d, want only healthy session", sent)
}
if elapsed >= 3*budget {
t.Fatalf("fan-out waited %v; want one shared %v budget, not one per slow session", elapsed, budget)
}
if got := len(healthy.outbound); got != 1 {
t.Fatalf("healthy queued ops = %d, want 1", got)
}
if healthy.terminal.Load() {
t.Fatal("healthy session was terminalized")
}
for i, c := range slow {
if !c.terminal.Load() {
t.Fatalf("slow session %d was not terminalized", i)
}
if tr := c.transport.(*closeCountingTransport); tr.closes != 1 {
t.Fatalf("slow session %d transport closes = %d, want 1", i, tr.closes)
}
}
}
func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw1 := [8]byte{1}
@ -130,6 +310,9 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
}
sm.BindAuthKeyForSession(raw1, 42, perm1)
// 两条 PFS/raw 连接可以解析到同一业务 perm key 且复用同一个 session_id；
// 精确排除必须只匹配 raw1，不能按 business key 把 raw2 一并排除。
sm.BindAuthKeyForSession(raw2, 42, perm1)
sm.BindUserForAuthKey(raw1, 42, 100)
sm.BindUserForAuthKey(raw2, 42, 200)
@ -148,7 +331,7 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
sm.BindUserForAuthKey(raw1, 42, 300)
sm.BindUserForAuthKey(raw2, 42, 300)
sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, perm1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, raw1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push except scoped session: %v", err)
}
@ -213,6 +396,151 @@ func TestSessionManagerCloseSessionsForBusinessAuthKeyClosesBoundTempAndRaw(t *t
}
}
func TestSessionManagerCloseSessionsRunsSlowPhysicalClosesConcurrently(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
business := [8]byte{9, 9, 9}
const sessions = 8
const closeDelay = 75 * time.Millisecond
transports := make([]*slowCloseTransport, 0, sessions)
for i := 0; i < sessions; i++ {
raw := [8]byte{byte(i + 1)}
tr := newSlowCloseTransport(closeDelay, nil)
c := &Conn{sessionID: int64(i + 1), authKeyID: raw, transport: tr}
sm.Register(c)
sm.BindAuthKeyForSession(raw, c.sessionID, business)
transports = append(transports, tr)
}
started := time.Now()
if got := sm.CloseSessionsForBusinessAuthKey(business); got != sessions {
t.Fatalf("closed sessions = %d, want %d", got, sessions)
}
elapsed := time.Since(started)
// A serial implementation takes ~600ms. Leave ample Windows/CI scheduling margin while
// still proving that the per-Conn delay is not multiplied by the session count.
if elapsed >= 4*closeDelay {
t.Fatalf("batch close elapsed = %v, want concurrent closes near %v", elapsed, closeDelay)
}
for i, tr := range transports {
select {
case <-tr.done:
default:
t.Fatalf("transport %d close had not completed when batch returned", i)
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport %d closes = %d, want 1", i, got)
}
}
}
func TestSessionManagerCloseRawSessionsExceptRunsConcurrentlyAndPreservesExcluded(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{6, 6, 6}
const sessions = 7
const excludedSession = int64(4)
const closeDelay = 60 * time.Millisecond
transports := make([]*slowCloseTransport, 0, sessions)
for i := 0; i < sessions; i++ {
tr := newSlowCloseTransport(closeDelay, nil)
c := &Conn{sessionID: int64(i + 1), authKeyID: raw, transport: tr}
sm.Register(c)
transports = append(transports, tr)
}
started := time.Now()
if got, want := sm.CloseSessionsForRawAuthKeyExcept(raw, excludedSession), sessions-1; got != want {
t.Fatalf("closed sessions = %d, want %d", got, want)
}
if elapsed := time.Since(started); elapsed >= 4*closeDelay {
t.Fatalf("raw-key batch close elapsed = %v, want concurrent closes near %v", elapsed, closeDelay)
}
for i, tr := range transports {
sessionID := int64(i + 1)
if sessionID == excludedSession {
if got := tr.closes.Load(); got != 0 {
t.Fatalf("excluded transport closes = %d, want 0", got)
}
continue
}
select {
case <-tr.done:
default:
t.Fatalf("transport for session %d had not closed", sessionID)
}
}
if _, ok := sm.bySession[sessionKey{authKeyID: raw, sessionID: excludedSession}]; !ok {
t.Fatal("excluded session was removed from the registry")
}
// Clean up the deliberately preserved connection without making the assertion path depend
// on test process teardown.
if !sm.DestroySessionForAuthKey(raw, excludedSession) {
t.Fatal("cleanup destroy of excluded session failed")
}
}
func TestForceCloseBatchTimeoutStillClosesProducerAndRPCGates(t *testing.T) {
release := make(chan struct{})
const sessions = 4
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
defer scheduler.stop(time.Second)
conns := make([]*Conn, 0, sessions)
transports := make([]*slowCloseTransport, 0, sessions)
for i := 0; i < sessions; i++ {
tr := newSlowCloseTransport(0, release)
c := &Conn{
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "shutdown.budget",
size: 32,
}); err != nil {
t.Fatalf("enqueue queued RPC %d: %v", i, err)
}
conns = append(conns, c)
transports = append(transports, tr)
}
started := time.Now()
if completed := forceCloseConnBatch(conns, 40*time.Millisecond); completed {
t.Fatal("blocked transport close batch unexpectedly completed")
}
if elapsed := time.Since(started); elapsed > 250*time.Millisecond {
t.Fatalf("timed batch close blocked for %v", elapsed)
}
for i, c := range conns {
if !c.terminal.Load() {
t.Fatalf("connection %d producer gate remains open after batch timeout", i)
}
select {
case <-c.outboundStop:
default:
t.Fatalf("connection %d outbound stop was not published", i)
}
select {
case <-c.rpcRootCtx.Done():
default:
t.Fatalf("connection %d RPC root remains open after batch timeout", i)
}
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("RPC budget after batch gate close = tasks:%d bytes:%d, want zero", tasks, bytes)
}
close(release)
for i, tr := range transports {
select {
case <-tr.done:
case <-time.After(time.Second):
t.Fatalf("transport %d did not finish after release", i)
}
}
}
func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1}
@ -238,6 +566,63 @@ func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) {
}
}
func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
business := [8]byte{9, 9}
const userID = int64(100)
newConn := func(raw [8]byte, sessionID int64, queueFull bool) (*Conn, *closeCountingTransport) {
transport := &closeCountingTransport{}
c := &Conn{
authKeyID: raw,
sessionID: sessionID,
metrics: NopMetrics{},
transport: transport,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.receivesUpdates.Store(true)
if queueFull {
c.outbound <- outboundOp{}
}
sm.Register(c)
sm.BindAuthKeyForSession(raw, sessionID, business)
sm.BindUserForAuthKey(raw, sessionID, userID)
return c, transport
}
slowOne, slowOneTransport := newConn([8]byte{1}, 11, true)
slowTwo, slowTwoTransport := newConn([8]byte{2}, 12, true)
healthy, healthyTransport := newConn([8]byte{3}, 13, false)
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
defer cancel()
started := time.Now()
sent, err := sm.PushToUserAuthKey(ctx, userID, business, proto.MessageFromServer, &tg.UpdatesTooLong{})
elapsed := time.Since(started)
if err != nil {
t.Fatalf("PushToUserAuthKey: %v", err)
}
if sent != 1 {
t.Fatalf("sent = %d, want only healthy connection", sent)
}
if elapsed > 100*time.Millisecond {
t.Fatalf("elapsed = %v, want one shared deadline rather than per-session waits", elapsed)
}
if !slowOne.terminal.Load() || !slowTwo.terminal.Load() || slowOneTransport.closes != 1 || slowTwoTransport.closes != 1 {
t.Fatalf("slow connections not terminal/closed: one=%v/%d two=%v/%d",
slowOne.terminal.Load(), slowOneTransport.closes, slowTwo.terminal.Load(), slowTwoTransport.closes)
}
if healthy.terminal.Load() || healthyTransport.closes != 0 {
t.Fatalf("healthy connection was dropped: terminal=%v closes=%d", healthy.terminal.Load(), healthyTransport.closes)
}
select {
case <-healthy.outbound:
default:
t.Fatal("healthy PFS connection did not receive best-effort enqueue")
}
}
func TestSessionManagerChannelInterestIndex(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
@ -368,8 +753,16 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
select {
case op := <-c.outbound:
if op.msg != msg {
t.Fatalf("enqueued msg = %T, want original update", op.msg)
defer op.releaseReservation(c.outboundTrackedBudget)
if op.encoded == nil {
t.Fatal("immediate push did not retain its encoded body")
}
var got tg.UpdateShort
if err := got.Decode(&bin.Buffer{Buf: op.encoded.body}); err != nil {
t.Fatalf("decode enqueued update: %v", err)
}
if _, ok := got.Update.(*tg.UpdateLoginToken); !ok || got.Date != msg.Date {
t.Fatalf("enqueued update = %+v, want login token date %d", got, msg.Date)
}
case <-time.After(time.Second):
t.Fatal("immediate push was not enqueued")
@ -383,6 +776,126 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
}
}
func TestPendingPushBodiesUseGlobalByteBudgetAndReleaseOnDrop(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
encoded, err := encodeOutboundMessage(msg)
if err != nil {
t.Fatalf("encode pending fixture: %v", err)
}
sm.pendingBudget = newOutboundTrackedBudget(int64(len(encoded.body)))
key := sessionKey{authKeyID: [8]byte{9}, sessionID: 77}
sm.mu.Lock()
first := sm.queueLocked(key, proto.MessageFromServer, msg)
second := sm.queueLocked(key, proto.MessageFromServer, msg)
sm.mu.Unlock()
if !first || second {
t.Fatalf("pending queue results = first %v second %v, want true/false at byte cap", first, second)
}
if got := sm.pendingBudget.snapshot(); got != int64(len(encoded.body)) {
t.Fatalf("pending body budget = %d, want %d", got, len(encoded.body))
}
sm.mu.Lock()
sm.deletePendingLocked(key)
sm.mu.Unlock()
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("pending body budget after drop = %d, want zero", got)
}
}
func TestPendingFlushGlobalBodyPressureDoesNotTerminateHealthyConnection(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
key := sessionKey{authKeyID: [8]byte{6}, sessionID: 66}
c := &Conn{
authKeyID: key.authKeyID,
sessionID: key.sessionID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
outboundTrackedBudget: newOutboundTrackedBudget(1),
}
const userID = int64(606)
c.userID.Store(userID)
c.userIDResolved.Store(true)
sm.Register(c)
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
sm.mu.Lock()
if !sm.queueLocked(key, proto.MessageFromServer, msg) {
sm.mu.Unlock()
t.Fatal("queue pending push")
}
sm.flushing[key] = true
sm.mu.Unlock()
// Enter at the final retry so the test exercises the durable-difference fallback without
// waiting for the production backoff timer.
sm.runFlush(c, key, userID, maxFlushAttempts-1)
if c.terminal.Load() {
t.Fatal("shared body pressure terminated a healthy pending-flush connection")
}
if !c.receivesUpdates.Load() {
t.Fatal("pending flush did not activate difference fallback after bounded retries")
}
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("pending budget after fallback = %d, want zero", got)
}
}
func TestPendingPushBudgetSurvivesTakeAndReturnsAcrossOverflowAndUnregister(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
encoded, err := encodeOutboundMessage(msg)
if err != nil {
t.Fatalf("encode pending fixture: %v", err)
}
bytesPerPush := int64(len(encoded.body))
sm.pendingBudget = newOutboundTrackedBudget(bytesPerPush * (maxPendingPushesPerSession + 8))
key := sessionKey{authKeyID: [8]byte{7}, sessionID: 55}
c := &Conn{authKeyID: key.authKeyID, sessionID: key.sessionID}
sm.Register(c)
sm.mu.Lock()
for i := 0; i < maxPendingPushesPerSession+5; i++ {
if !sm.queueLocked(key, proto.MessageFromServer, msg) {
sm.mu.Unlock()
t.Fatalf("queue pending push %d unexpectedly failed", i)
}
}
if got, want := sm.pendingBudget.snapshot(), bytesPerPush*maxPendingPushesPerSession; got != want {
sm.mu.Unlock()
t.Fatalf("budget after overflow replacement = %d, want %d", got, want)
}
batch := sm.takePendingLocked(key, true)
sm.mu.Unlock()
if len(batch) != maxPendingPushesPerSession {
t.Fatalf("taken pending pushes = %d, want %d", len(batch), maxPendingPushesPerSession)
}
// take transfers ownership to runFlush; deleting the map entry must not release bodies while
// the batch still references them.
if got, want := sm.pendingBudget.snapshot(), bytesPerPush*maxPendingPushesPerSession; got != want {
t.Fatalf("budget after take = %d, want transferred ownership %d", got, want)
}
releaseQueuedPushes(batch)
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("budget after taken batch release = %d, want 0", got)
}
sm.mu.Lock()
if !sm.queueLocked(key, proto.MessageFromServer, msg) {
sm.mu.Unlock()
t.Fatal("queue before unregister failed")
}
sm.mu.Unlock()
sm.Unregister(c)
if got := sm.pendingBudget.snapshot(); got != 0 {
t.Fatalf("budget after unregister = %d, want 0", got)
}
}
// TestSessionManagerPush 验证主动推送端到端：两个 client 连接握手并建立 session 后，
// server 经 PushToSession / PushToUser 主动向其推送，client 收到。
func TestSessionManagerPush(t *testing.T) {

60
internal/mtprotoedge/session_membership_gen_test.go
View file

@ -1,12 +1,62 @@
package mtprotoedge
import (
"slices"
"testing"
"time"
"go.uber.org/zap/zaptest"
)
func TestOnlineChannelIDsSnapshotAndDiagnosticPagesStableAscending(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{4, 5, 6}
c := &Conn{sessionID: 77, authKeyID: raw}
sm.Register(c)
sm.BindUserForAuthKey(raw, 77, 100)
sm.SetSessionChannelMemberships(raw, 77, 100, []int64{50, 10, 30, 20, 40}, sm.ChannelMembershipGeneration(raw, 77))
want := []int64{10, 20, 30, 40, 50}
snapshot := sm.OnlineChannelIDsSnapshot()
if !slices.Equal(snapshot, want) {
t.Fatalf("online channel snapshot = %v, want %v", snapshot, want)
}
var got []int64
after := int64(0)
for {
page := sm.OnlineChannelIDsAfter(after, 2)
if len(page) == 0 {
break
}
for _, channelID := range page {
if channelID <= after {
t.Fatalf("page %v not strictly after cursor %d", page, after)
}
after = channelID
got = append(got, channelID)
}
}
if !slices.Equal(got, want) {
t.Fatalf("paged online channels = %v, want %v", got, want)
}
// The recovery actor owns a stable copy: later membership changes are visible to the next
// generation, not spliced into the in-flight sorted snapshot.
sm.AddUserChannelMembership(100, 5)
if !slices.Equal(snapshot, want) {
t.Fatalf("owned snapshot mutated after membership insert: %v", snapshot)
}
if current := sm.OnlineChannelIDsSnapshot(); !slices.Equal(current, []int64{5, 10, 20, 30, 40, 50}) {
t.Fatalf("next online channel snapshot = %v", current)
}
// Removing the only live session must immediately remove all channel ids from the recovery
// enumeration; stale membership map entries are never enough without a live bySession key.
sm.Unregister(c)
if got := sm.OnlineChannelIDsAfter(0, 10); len(got) != 0 {
t.Fatalf("online channels after unregister = %v, want empty", got)
}
}
// TestSetSessionChannelMembershipsDetectsConcurrentIncrementalUpdates 验证全量
// membership 同步的丢失更新防护：同步方在读持久成员列表前采样修订号，读取窗口内
// 若发生增量 join/leave（另一设备操作经 Add/RemoveUserChannelMembership 落索引），
@ -67,13 +117,18 @@ func TestRegisterEvictsOldestSessionAtCap(t *testing.T) {
base := time.Unix(1_700_000_000, 0)
const oldestSession = int64(100)
oldestTransport := &closeCountingTransport{}
for i := 0; i < maxSessionsPerAuthKey; i++ {
sid := int64(i + 1)
created := base.Add(time.Duration(i+1) * time.Second)
if sid == oldestSession {
created = base // 唯一早于所有其它连接的时间戳，且故意不在注册顺序首位。
}
sm.Register(&Conn{sessionID: sid, authKeyID: raw, createdAt: created})
c := &Conn{sessionID: sid, authKeyID: raw, createdAt: created}
if sid == oldestSession {
c.transport = oldestTransport
}
sm.Register(c)
}
sm.Register(&Conn{sessionID: 9999, authKeyID: raw, createdAt: base.Add(time.Hour)})
@ -92,4 +147,7 @@ func TestRegisterEvictsOldestSessionAtCap(t *testing.T) {
if total != maxSessionsPerAuthKey {
t.Fatalf("sessions for auth key = %d, want cap %d", total, maxSessionsPerAuthKey)
}
if oldestTransport.closes != 1 {
t.Fatalf("evicted transport closes = %d, want 1", oldestTransport.closes)
}
}

65
internal/mtprotoedge/shutdown_gate_test.go Normal file
View file

@ -0,0 +1,65 @@
package mtprotoedge
import (
"testing"
"time"
)
func TestTerminalFailurePathsCloseGatesBeforeBlockingTransportClose(t *testing.T) {
tests := []struct {
name string
run func(*Conn)
}{
{name: "write failure", run: (*Conn).failTransport},
{name: "slow consumer", run: (*Conn).dropSlowConsumer},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
release := make(chan struct{})
tr := newSlowCloseTransport(0, release)
scheduler := newInboundRPCScheduler(1, 1, 1024)
defer scheduler.stop(time.Second)
c := &Conn{
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
returned := make(chan struct{})
go func() {
tt.run(c)
close(returned)
}()
deadline := time.Now().Add(time.Second)
for tr.closes.Load() == 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if tr.closes.Load() == 0 {
t.Fatal("terminal path did not enter transport.Close")
}
if !c.terminal.Load() {
t.Fatal("producer terminal gate was not published before blocking Close")
}
select {
case <-c.outboundStop:
default:
t.Fatal("outbound stop was not published before blocking Close")
}
select {
case <-c.rpcRootCtx.Done():
default:
t.Fatal("RPC root was not canceled before blocking Close")
}
close(release)
select {
case <-returned:
case <-time.After(time.Second):
t.Fatal("terminal path did not return after transport release")
}
})
}
}

206
internal/mtprotoedge/structural_limits_test.go Normal file
View file

@ -0,0 +1,206 @@
package mtprotoedge
import (
"context"
"strings"
"testing"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"go.uber.org/zap/zaptest"
)
func TestContainerMessageCountAndServiceVectorCaps(t *testing.T) {
var container bin.Buffer
container.PutID(proto.MessageContainerTypeID)
container.PutInt(maxContainerMessages)
if got, err := containerMessageCount(&container); err != nil || got != maxContainerMessages {
t.Fatalf("container count = %d/%v, want %d/nil", got, err, maxContainerMessages)
}
container.Buf[4]++
if got, err := containerMessageCount(&container); err != nil || got != maxContainerMessages+1 {
t.Fatalf("oversized container preflight = %d/%v, want %d/nil", got, err, maxContainerMessages+1)
}
ack := mt.MsgsAck{MsgIDs: make([]int64, maxServiceMessageIDs)}
var encoded bin.Buffer
if err := ack.Encode(&encoded); err != nil {
t.Fatalf("encode msgs_ack: %v", err)
}
if err := validateFirstVectorCount(&encoded, maxServiceMessageIDs); err != nil {
t.Fatalf("service vector at cap: %v", err)
}
// Count lives after constructor + vector constructor. We only mutate the declared count: the
// preflight must reject before generated Decode attempts a long loop/allocation.
encoded.Buf[8]++
if err := validateFirstVectorCount(&encoded, maxServiceMessageIDs); err == nil {
t.Fatal("service vector above cap unexpectedly accepted")
}
}
func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) {
encoded := bin.Buffer{}
wantBody := []byte{0x11, 0x22, 0x33, 0x44}
message := proto.Message{ID: 1, SeqNo: 1, Bytes: len(wantBody), Body: wantBody}
if err := (&proto.MessageContainer{Messages: []proto.Message{message}}).Encode(&encoded); err != nil {
t.Fatalf("encode container: %v", err)
}
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes - 1)
if _, release, err := s.decodeMessageContainerViews(&encoded, 1); err == nil {
release()
t.Fatal("descriptor allocation unexpectedly bypassed process budget")
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("failed descriptor reservation leaked %d bytes", got)
}
s.frameBudget = newInboundFrameBudget(2 * containerDescriptorBudgetBytes)
container, release, err := s.decodeMessageContainerViews(&encoded, 1)
if err != nil {
t.Fatalf("decode budgeted container: %v", err)
}
if got := s.frameBudget.usedBytes(); got != containerDescriptorBudgetBytes {
t.Fatalf("descriptor budget = %d, want %d", got, containerDescriptorBudgetBytes)
}
container.Messages[0].Body[0] = 0x99
if encoded.Buf[8+16] != 0x99 {
t.Fatal("container body was copied instead of viewing the charged input frame")
}
release()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("released descriptor budget = %d, want zero", got)
}
truncated := bin.Buffer{Buf: encoded.Buf[:len(encoded.Buf)-1]}
if _, _, err := s.decodeMessageContainerViews(&truncated, 1); err == nil {
t.Fatal("truncated container unexpectedly decoded")
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("failed container decode leaked %d bytes", got)
}
}
func TestServiceInfoViewsRejectOversizedBytesWithoutDecodeCopy(t *testing.T) {
state := mt.MsgsStateInfo{ReqMsgID: 7, Info: make([]byte, maxServiceMessageIDs)}
var encodedState bin.Buffer
if err := state.Encode(&encodedState); err != nil {
t.Fatalf("encode msgs_state_info: %v", err)
}
reqMsgID, info, err := msgsStateInfoView(&encodedState)
if err != nil || reqMsgID != state.ReqMsgID || len(info) != maxServiceMessageIDs {
t.Fatalf("state info view = id %d len %d err %v", reqMsgID, len(info), err)
}
info[0] = 0x7f
if encodedState.Buf[16] != 0x7f {
t.Fatal("msgs_state_info view unexpectedly copied info")
}
state.Info = make([]byte, maxServiceMessageIDs+1)
encodedState.Reset()
if err := state.Encode(&encodedState); err != nil {
t.Fatalf("encode oversized msgs_state_info: %v", err)
}
if _, _, err := msgsStateInfoView(&encodedState); err == nil || !strings.Contains(err.Error(), "exceeds") {
t.Fatalf("oversized msgs_state_info err = %v, want capped rejection", err)
}
all := mt.MsgsAllInfo{MsgIDs: []int64{1, 2}, Info: []byte{4, 4}}
var encodedAll bin.Buffer
if err := all.Encode(&encodedAll); err != nil {
t.Fatalf("encode msgs_all_info: %v", err)
}
count, allInfo, err := msgsAllInfoView(&encodedAll)
if err != nil || count != 2 || len(allInfo) != 2 {
t.Fatalf("all info view = count %d len %d err %v", count, len(allInfo), err)
}
all.Info = make([]byte, maxServiceMessageIDs+1)
encodedAll.Reset()
if err := all.Encode(&encodedAll); err != nil {
t.Fatalf("encode oversized msgs_all_info: %v", err)
}
if _, _, err := msgsAllInfoView(&encodedAll); err == nil || !strings.Contains(err.Error(), "exceeds") {
t.Fatalf("oversized msgs_all_info err = %v, want capped rejection", err)
}
}
func TestDispatchRejectsExcessiveWrapperDepthBeforeRPC(t *testing.T) {
var body bin.Buffer
if err := (&mt.MsgsStateInfo{ReqMsgID: 1, Info: []byte{4}}).Encode(&body); err != nil {
t.Fatalf("encode leaf: %v", err)
}
encoded := body.Copy()
for i := 0; i < maxDispatchDepth+1; i++ {
var wrapped bin.Buffer
if err := (proto.GZIP{Data: encoded}).Encode(&wrapped); err != nil {
t.Fatalf("encode gzip depth %d: %v", i+1, err)
}
encoded = wrapped.Copy()
}
s := New(Options{Logger: zaptest.NewLogger(t)})
var acks []int64
err := s.dispatch(context.Background(), newConnState(), nil, 4, 0, &bin.Buffer{Buf: encoded}, &acks)
if err == nil || !strings.Contains(err.Error(), "wrapper depth") {
t.Fatalf("deep wrapper err = %v, want wrapper depth rejection", err)
}
}
func TestOversizedConnectionBuffersAreReleasedAfterFrame(t *testing.T) {
inbound := &bin.Buffer{Buf: make([]byte, 1, maxRetainedConnBuffer+1)}
trimOversizedInboundBuffer(inbound)
if inbound.Buf != nil {
t.Fatalf("oversized inbound buffer cap=%d, want released", cap(inbound.Buf))
}
regular := &bin.Buffer{Buf: make([]byte, 1, maxRetainedConnBuffer)}
trimOversizedInboundBuffer(regular)
if cap(regular.Buf) != maxRetainedConnBuffer {
t.Fatalf("regular inbound buffer cap=%d, want retained", cap(regular.Buf))
}
pool := newOutboundScratchPool(16 << 20)
scratch, err := pool.acquire(context.Background(), nil, maxRetainedConnBuffer+1)
if err != nil {
t.Fatalf("acquire oversized outbound scratch: %v", err)
}
pool.release(scratch)
if got := pool.snapshot(); got != 0 {
t.Fatalf("oversized outbound scratch retained %d bytes, want 0", got)
}
}
func TestGZIPExpansionUsesProcessBudgetBeforeDecode(t *testing.T) {
payload := make([]byte, 1<<20)
var wrapped bin.Buffer
if err := (proto.GZIP{Data: payload}).Encode(&wrapped); err != nil {
t.Fatalf("encode gzip: %v", err)
}
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes - 1)
if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
release()
t.Fatal("gzip decode unexpectedly bypassed saturated process budget")
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("failed gzip reservation leaked %d bytes", got)
}
s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
decoded, release, err := s.decodeGZIPWithGlobalBudget(&wrapped)
if err != nil {
t.Fatalf("budgeted gzip decode: %v", err)
}
if len(decoded) != len(payload) {
t.Fatalf("decoded bytes = %d, want %d", len(decoded), len(payload))
}
if got := s.frameBudget.usedBytes(); got != int64(len(payload)) {
t.Fatalf("held expansion budget = %d, want %d", got, len(payload))
}
release()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("released expansion budget = %d, want zero", got)
}
}

257
internal/mtprotoedge/transport_compat.go
View file

@ -34,16 +34,21 @@ type quickAckTransport interface {
SendQuickAck(ctx context.Context, token uint32) error
}
type deadlineQuickAckTransport interface {
SendQuickAckDeadline(deadline time.Time, token uint32) error
}
type compatTransportListener struct {
codec func() transport.Codec
listener net.Listener
budget *inboundFrameBudget
}
func newCompatTransportListener(codec func() transport.Codec, listener net.Listener) transportListener {
if codec != nil {
return transport.ListenCodec(codec, listener)
func newCompatTransportListener(codec func() transport.Codec, listener net.Listener, budget *inboundFrameBudget) transportListener {
if budget == nil {
panic("mtprotoedge: nil inbound frame budget")
}
return &compatTransportListener{listener: listener}
return &compatTransportListener{codec: codec, listener: listener, budget: budget}
}
// singleConnListener 是一个只产出一条「已接受」连接、随后阻塞到关闭的 net.Listener。
@ -89,9 +94,27 @@ func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
}
}()
connCodec, reader, err := detectCompatCodec(conn)
if err != nil {
return nil, errors.Wrap(err, "detect codec")
var (
connCodec transport.Codec
reader io.Reader = conn
)
if l.codec != nil {
connCodec = l.codec()
if classifyInboundFrameCodec(connCodec) == inboundFrameCodecUnknown {
// Unknown codecs are rejected before their header or first frame is read. Without an
// explicit preflight contract, calling Codec.Read could allocate from an attacker-
// controlled length before the process-wide budget can be reserved.
return nil, errInboundFrameCodecUnsupported
}
if err := connCodec.ReadHeader(conn); err != nil {
return nil, errors.Wrap(err, "read codec header")
}
} else {
var err error
connCodec, reader, err = detectCompatCodec(conn)
if err != nil {
return nil, errors.Wrap(err, "detect codec")
}
}
return &compatTransportConn{
@ -99,7 +122,8 @@ func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
reader: reader,
Conn: conn,
},
codec: connCodec,
codec: connCodec,
budget: l.budget,
}, nil
}
@ -121,11 +145,17 @@ func (w wrappedCompatConn) Read(p []byte) (int, error) {
}
type compatTransportConn struct {
conn net.Conn
codec transport.Codec
conn net.Conn
codec transport.Codec
budget *inboundFrameBudget
readMux sync.Mutex
writeMux sync.Mutex
frameMu sync.Mutex
heldFrameBytes int64
frameDelivered bool
closed bool
}
func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error {
@ -157,6 +187,11 @@ func (c *compatTransportConn) ConsumeQuickAckRequested() bool {
}
func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) error {
deadline, _ := ctx.Deadline()
return c.SendQuickAckDeadline(deadline, token)
}
func (c *compatTransportConn) SendQuickAckDeadline(deadline time.Time, token uint32) error {
q, ok := c.codec.(quickAckCodec)
if !ok {
return nil
@ -165,7 +200,6 @@ func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) er
c.writeMux.Lock()
defer c.writeMux.Unlock()
deadline, _ := ctx.Deadline()
if err := c.conn.SetWriteDeadline(deadline); err != nil {
return errors.Wrap(err, "set write deadline")
}
@ -189,19 +223,185 @@ func (c *compatTransportConn) RecvDeadline(deadline time.Time, b *bin.Buffer) er
c.readMux.Lock()
defer c.readMux.Unlock()
// Starting the next Recv proves the previous frame slices are no longer consumed, but its
// reusable backing remains live. Keep the high-water reservation until the new length prefix
// atomically grows/reuses it; serveConn later shrinks it to the actually retained capacities.
c.beginInboundFrameRead()
if err := c.conn.SetReadDeadline(deadline); err != nil {
c.releaseInboundFrame()
return errors.Wrap(err, "set read deadline")
}
if err := c.codec.Read(c.conn, b); err != nil {
if err := c.readInboundFrame(b); err != nil {
// A short payload or protocol error cannot escape while retaining a reservation.
c.releaseInboundFrame()
return errors.Wrap(err, "read")
}
return nil
}
func (c *compatTransportConn) Close() error {
c.frameMu.Lock()
c.closed = true
// Never release here: Close can race both a delivered frame owned by serveConn and a codec read
// still writing into b. Recv's error path or serveConn's deferred ownership release is the
// unique point where those backings become dead.
c.frameMu.Unlock()
return c.conn.Close()
}
func (c *compatTransportConn) readInboundFrame(b *bin.Buffer) error {
kind := classifyInboundFrameCodec(c.codec)
if kind == inboundFrameCodecUnknown {
return errInboundFrameCodecUnsupported
}
reserveCalls := 0
reserved := false
var reserveErr error
reserve := func(wireBytes, plaintextBytes int64) error {
reserveCalls++
if reserveCalls != 1 {
reserveErr = errors.New("inbound frame codec reserved more than once")
return reserveErr
}
reserveErr = c.reserveInboundFrame(wireBytes, plaintextBytes)
reserved = reserveErr == nil
return reserveErr
}
var err error
if kind == inboundFrameCodecCustom {
custom := unwrapInboundFrameBudgetedCodec(c.codec)
if custom == nil {
return errInboundFrameCodecUnsupported
}
err = custom.ReadWithInboundFrameBudget(c.conn, b, reserve)
} else {
preflight := &inboundFramePreflightReader{r: c.conn, kind: kind, reserve: reserve}
err = c.codec.Read(preflight, b)
}
if err != nil {
return err
}
if reserveErr != nil {
return reserveErr
}
if reserveCalls != 1 || !reserved {
return errInboundFrameNotReserved
}
return c.markInboundFrameDelivered()
}
func (c *compatTransportConn) reserveInboundFrame(wireBytes, plaintextBytes int64) error {
c.frameMu.Lock()
defer c.frameMu.Unlock()
if c.closed {
return net.ErrClosed
}
n, err := c.budget.growReservation(c.heldFrameBytes, wireBytes, plaintextBytes)
if err != nil {
return err
}
c.heldFrameBytes = n
return nil
}
func (c *compatTransportConn) beginInboundFrameRead() {
c.frameMu.Lock()
c.frameDelivered = false
c.frameMu.Unlock()
}
// retainInboundFrameBytes shrinks the high-water frame charge to the capacities that serveConn
// intentionally keeps for reuse after dispatch. It may grow only to account allocator rounding;
// callers drop both buffers and retry with zero when that extra admission is unavailable.
func (c *compatTransportConn) retainInboundFrameBytes(n int64) bool {
if n < 0 {
return false
}
c.frameMu.Lock()
old := c.heldFrameBytes
if n > old {
grown, err := c.budget.growReservation(old, n, 0)
if err != nil {
c.frameMu.Unlock()
return false
}
c.heldFrameBytes = grown
c.frameMu.Unlock()
return true
}
c.heldFrameBytes = n
c.frameMu.Unlock()
c.budget.release(old - n)
return true
}
func (c *compatTransportConn) releaseInboundFrame() {
c.frameMu.Lock()
n := c.heldFrameBytes
c.heldFrameBytes = 0
c.frameDelivered = false
c.frameMu.Unlock()
c.budget.release(n)
}
func (c *compatTransportConn) markInboundFrameDelivered() error {
c.frameMu.Lock()
defer c.frameMu.Unlock()
if c.closed {
// Do not hand a frame to the consumer after Close. The read error path keeps ownership
// accounting until the codec has stopped touching its backing, then releases it.
return net.ErrClosed
}
if c.heldFrameBytes == 0 {
return errInboundFrameNotReserved
}
c.frameDelivered = true
return nil
}
type inboundFrameOwnershipReleaser interface {
releaseInboundFrame()
}
type inboundFrameBackingRetainer interface {
retainInboundFrameBytes(int64) bool
}
func releaseInboundFrameOwnership(conn transport.Conn) {
if releaser, ok := conn.(inboundFrameOwnershipReleaser); ok {
releaser.releaseInboundFrame()
}
}
// retainInboundFrameBackings transfers the current-frame reservation into a persistent charge
// for reusable buffer capacities. If allocator rounding would exceed the available budget, drop
// both backings and release the reservation rather than retaining unaccounted memory.
func retainInboundFrameBackings(conn transport.Conn, buffers ...*bin.Buffer) {
retainer, ok := conn.(inboundFrameBackingRetainer)
if !ok {
return
}
var retained int64
for _, b := range buffers {
if b == nil {
continue
}
retained += int64(cap(b.Buf))
}
if retainer.retainInboundFrameBytes(retained) {
return
}
for _, b := range buffers {
if b != nil {
b.Buf = nil
}
}
if !retainer.retainInboundFrameBytes(0) {
panic("mtprotoedge: failed to release inbound frame backing reservation")
}
}
func detectCompatCodec(c io.Reader) (transport.Codec, io.Reader, error) {
var buf [4]byte
if _, err := io.ReadFull(c, buf[:1]); err != nil {
@ -233,7 +433,6 @@ type quickAckCodec interface {
type quickAckAbridgedCodec struct {
quickAckRequested bool
wbuf []byte
}
func (*quickAckAbridgedCodec) WriteHeader(w io.Writer) error {
@ -261,7 +460,7 @@ func (q *quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error {
header[3] = byte(words >> 16)
headerLen = 4
}
return writeCompatPacket(w, &q.wbuf, header[:headerLen], b.Raw())
return writeCompatPacket(w, header[:headerLen], b.Raw())
}
func (q *quickAckAbridgedCodec) Read(r io.Reader, b *bin.Buffer) error {
@ -287,7 +486,6 @@ func (*quickAckAbridgedCodec) quickAckResponse(token uint32) [4]byte {
type quickAckIntermediateCodec struct {
quickAckRequested bool
wbuf []byte
}
func (*quickAckIntermediateCodec) WriteHeader(w io.Writer) error {
@ -304,7 +502,7 @@ func (q *quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
}
var header [4]byte
binary.LittleEndian.PutUint32(header[:], uint32(b.Len()))
return writeCompatPacket(w, &q.wbuf, header[:], b.Raw())
return writeCompatPacket(w, header[:], b.Raw())
}
func (q *quickAckIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
@ -330,7 +528,6 @@ func (*quickAckIntermediateCodec) quickAckResponse(token uint32) [4]byte {
type quickAckPaddedIntermediateCodec struct {
quickAckRequested bool
wbuf []byte
rand *bufio.Reader
}
@ -356,13 +553,11 @@ func (q *quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) erro
return err
}
n := int(padding[0] % 4)
// header(4B) + payload + padding 一次拼进复用缓冲，单次 Write 出站。
buf := append(q.wbuf[:0], 0, 0, 0, 0)
binary.LittleEndian.PutUint32(buf[:4], uint32(b.Len()+n))
buf = append(buf, b.Raw()...)
buf = append(buf, padding[:n]...)
q.wbuf = buf
return writeAll(w, buf)
var header [4]byte
binary.LittleEndian.PutUint32(header[:], uint32(b.Len()+n))
buffers := net.Buffers{header[:], b.Raw(), padding[:n]}
_, err := buffers.WriteTo(w)
return err
}
func (q *quickAckPaddedIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
@ -449,13 +644,13 @@ func validateOutgoingCompatMessage(b *bin.Buffer) error {
return nil
}
// writeCompatPacket 把 header+payload 拼进调用方持有的复用缓冲后单次写出：
// 保持 MTProto 帧单包出站（quick ack 尾延迟契约），同时避免每帧分配拼包缓冲。
func writeCompatPacket(w io.Writer, scratch *[]byte, header, payload []byte) error {
buf := append((*scratch)[:0], header...)
buf = append(buf, payload...)
*scratch = buf
return writeAll(w, buf)
// writeCompatPacket avoids a full-frame codec copy. net.Buffers uses vectored I/O for raw TCP
// (one syscall); wrapped writers may receive ordered writes, still serialized by writeMux. The
// outbound scratch lease keeps the encrypted payload alive until all segments finish.
func writeCompatPacket(w io.Writer, header, payload []byte) error {
buffers := net.Buffers{header, payload}
_, err := buffers.WriteTo(w)
return err
}
func writeAll(w io.Writer, p []byte) error {

14
internal/mtprotoedge/transport_compat_test.go
View file

@ -75,8 +75,8 @@ func TestCompatPaddedIntermediateWriteRoundTrip(t *testing.T) {
if err := codec.Write(&out, &payload); err != nil {
t.Fatalf("write %d: %v", i, err)
}
if out.writes != 1 {
t.Fatalf("write %d: writes = %d, want 1", i, out.writes)
if out.writes < 2 || out.writes > 3 {
t.Fatalf("write %d: writes = %d, want header/payload[/padding] segments", i, out.writes)
}
total := binary.LittleEndian.Uint32(out.Bytes()[:4])
if int(total) != len(out.Bytes())-4 {
@ -97,7 +97,7 @@ func TestCompatPaddedIntermediateWriteRoundTrip(t *testing.T) {
}
}
func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) {
func TestCompatTransportCodecsWriteSegmentedPacketWithoutFullCopy(t *testing.T) {
var payload bin.Buffer
payload.PutInt32(0x01020304)
payload.PutInt32(0x05060708)
@ -107,8 +107,8 @@ func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) {
if err := (&quickAckAbridgedCodec{}).Write(&out, &payload); err != nil {
t.Fatalf("write: %v", err)
}
if out.writes != 1 {
t.Fatalf("writes = %d, want 1", out.writes)
if out.writes != 2 {
t.Fatalf("generic writer calls = %d, want header+payload; raw TCP uses vectored I/O", out.writes)
}
if got, want := out.Bytes()[0], byte(payload.Len()/bin.Word); got != want {
t.Fatalf("abridged header = %#x, want %#x", got, want)
@ -120,8 +120,8 @@ func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) {
if err := (&quickAckIntermediateCodec{}).Write(&out, &payload); err != nil {
t.Fatalf("write: %v", err)
}
if out.writes != 1 {
t.Fatalf("writes = %d, want 1", out.writes)
if out.writes != 2 {
t.Fatalf("generic writer calls = %d, want header+payload; raw TCP uses vectored I/O", out.writes)
}
if got, want := binary.LittleEndian.Uint32(out.Bytes()[:4]), uint32(payload.Len()); got != want {
t.Fatalf("intermediate length = %d, want %d", got, want)

4
internal/rpc/account.go
View file

@ -490,8 +490,8 @@ func (r *Router) onAccountVerifyEmail(ctx context.Context, req *tg.AccountVerify
SentCode: tgEmailSentCode(p.PhoneCodeHash, domain.MaskEmail(email), len(strings.TrimSpace(code))),
}, nil
}
u, loginMessage, needSignUp, signInErr := r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), p.PhoneNumber, p.PhoneCodeHash, code)
authorization, err := r.finishAuthSignIn(ctx, u, loginMessage, needSignUp, signInErr)
u, _, needSignUp, signInErr := r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), p.PhoneNumber, p.PhoneCodeHash, code)
authorization, err := r.finishAuthSignIn(ctx, u, needSignUp, signInErr)
if err != nil {
return nil, err
}

2
internal/rpc/account_business.go
View file

@ -490,7 +490,7 @@ func (r *Router) recordConnectedBusinessPeerSettings(ctx context.Context, userID
}
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
event, _, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, settings, sessionID)
event, _, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, settings, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return err
}

2
internal/rpc/account_phone.go
View file

@ -51,10 +51,12 @@ func (r *Router) onAccountChangePhone(ctx context.Context, req *tg.AccountChange
return nil, authKeyUnregisteredErr()
}
sessionID, _ := SessionIDFrom(ctx)
originRawAuthKeyID := rawAuthKeyIDForOrigin(ctx)
result, err := r.deps.Account.ChangePhone(
ctx,
userID,
authKeyID,
originRawAuthKeyID,
sessionID,
req.PhoneNumber,
req.PhoneCodeHash,

94
internal/rpc/account_themes_compat_test.go
View file

@ -6,31 +6,39 @@ import (
"github.com/gotd/td/bin"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"go.uber.org/zap"
"go.uber.org/zap/zaptest/observer"
"telesrv/internal/domain"
)
// TestLegacyThemeWireDecode 验证按 DrKLO 12.8.1 的 theme 构造器(比 gotd schema 新)
// 手写解码后能正确复用现有 handler。直接构造 DrKLO 的 wire 字节喂给 fallback compat 层。
func TestLegacyThemeWireDecode(t *testing.T) {
// TestLegacyThemeWireDispatch 验证 DrKLO theme 构造器经过 Router.Dispatch 的完整链路：
// layerwire 结构预检 -> gotd dispatcher fallback -> compat 解码 -> 现有 handler。
// 不能只直调 tryLegacyThemeRPC，否则会掩盖预检早于 fallback 的路由回归。
func TestLegacyThemeWireDispatch(t *testing.T) {
const userID = 1000010
ctx := WithUserID(context.Background(), userID)
var authKeyID [8]byte
authKeyID[0] = 1
const sessionID = 99
files := &fakeFiles{docs: map[int64]domain.Document{
777: {ID: 777, AccessHash: 7, DCID: 2, MimeType: "application/x-tgtheme-android", Size: 4096},
}}
r := newThemeRouter(t, files)
// createTheme 0x8432c21f:flags(=4,document) + slug + title + InputDocument。
// createTheme 0x8432c21f:flags(document+single settings) + slug + title。
var cb bin.Buffer
cb.PutID(legacyCreateThemeID)
cb.PutInt32(1 << 2) // document present
cb.PutString("") // empty slug → auto
cb.PutInt32((1 << 2) | (1 << 3))
cb.PutString("") // empty slug → auto
cb.PutString("Legacy Theme")
(&tg.InputDocument{ID: 777, AccessHash: 7}).Encode(&cb)
(&tg.InputThemeSettings{BaseTheme: &tg.BaseThemeDay{}, AccentColor: 0x3997d3}).Encode(&cb)
enc, handled, err := r.tryLegacyThemeRPC(ctx, &cb)
if !handled || err != nil {
t.Fatalf("createTheme legacy = handled %v err %v", handled, err)
enc, err := r.Dispatch(ctx, authKeyID, sessionID, &cb)
if err != nil {
t.Fatalf("createTheme legacy dispatch: %v", err)
}
th, ok := enc.(*tg.Theme)
if !ok {
@ -44,9 +52,29 @@ func TestLegacyThemeWireDecode(t *testing.T) {
} else if d, _ := doc.(*tg.Document); d == nil || d.ID != 777 {
t.Fatalf("created theme document = %#v, want id 777", doc)
}
if settings, ok := th.GetSettings(); !ok || len(settings) != 1 || settings[0].AccentColor != 0x3997d3 {
t.Fatalf("created theme settings = %#v ok=%v, want one legacy setting", settings, ok)
}
mustEncodeTheme(t, th)
slug := th.Slug
// updateTheme 0x5cb367d5:flags(=2,title) + format + InputTheme + title。
var ub bin.Buffer
ub.PutID(legacyUpdateThemeID)
ub.PutInt32(1 << 1)
ub.PutString("android")
(&tg.InputTheme{ID: th.ID, AccessHash: th.AccessHash}).Encode(&ub)
ub.PutString("Legacy Theme Updated")
enc, err = r.Dispatch(ctx, authKeyID, sessionID, &ub)
if err != nil {
t.Fatalf("updateTheme legacy dispatch: %v", err)
}
updated, ok := enc.(*tg.Theme)
if !ok || updated.Title != "Legacy Theme Updated" {
t.Fatalf("updateTheme legacy result = %#v, want updated title", enc)
}
// getTheme 0x8d9d742b:format + InputThemeSlug + document_id(被忽略)。
var gb bin.Buffer
gb.PutID(legacyGetThemeID)
@ -54,9 +82,9 @@ func TestLegacyThemeWireDecode(t *testing.T) {
(&tg.InputThemeSlug{Slug: slug}).Encode(&gb)
gb.PutLong(12345) // document_id ignored
enc, handled, err = r.tryLegacyThemeRPC(ctx, &gb)
if !handled || err != nil {
t.Fatalf("getTheme legacy = handled %v err %v", handled, err)
enc, err = r.Dispatch(ctx, authKeyID, sessionID, &gb)
if err != nil {
t.Fatalf("getTheme legacy dispatch: %v", err)
}
got, ok := enc.(*tg.Theme)
if !ok || got.Slug != slug {
@ -73,18 +101,44 @@ func TestLegacyThemeWireDecode(t *testing.T) {
ib.PutString("android")
(&tg.InputThemeSlug{Slug: slug}).Encode(&ib)
enc, handled, err = r.tryLegacyThemeRPC(ctx, &ib)
if !handled || err != nil {
t.Fatalf("installTheme legacy = handled %v err %v", handled, err)
enc, err = r.Dispatch(ctx, authKeyID, sessionID, &ib)
if err != nil {
t.Fatalf("installTheme legacy dispatch: %v", err)
}
if _, ok := enc.(*tg.BoolTrue); !ok {
t.Fatalf("installTheme legacy result = %T, want *tg.BoolTrue", enc)
}
// 非 theme 构造器 → 不处理。
var ob bin.Buffer
ob.PutID(0x12345678)
if _, handled, _ := r.tryLegacyThemeRPC(ctx, &ob); handled {
t.Fatalf("unrelated ctor should not be handled")
// 已声明 legacy 方法仍必须精确消费完整结构，截断字段不能到手写 decoder。
var malformed bin.Buffer
malformed.PutID(legacyCreateThemeID)
malformed.PutInt32(0)
malformed.PutString("slug") // missing title
if _, err := r.Dispatch(ctx, authKeyID, sessionID, &malformed); !tgerr.Is(err, "INPUT_REQUEST_INVALID") {
t.Fatalf("malformed legacy theme err = %v, want INPUT_REQUEST_INVALID", err)
}
}
func TestUnknownRPCReachesCompatibilityTraceAfterOpaquePreflight(t *testing.T) {
const unknownID = uint32(0x12345678)
r := newThemeRouter(t, &fakeFiles{})
r.deps.Auth = &captureAuthService{}
core, logs := observer.New(zap.WarnLevel)
r.log = zap.New(core)
var b bin.Buffer
b.PutID(unknownID)
// Deliberately resembles a forged vector count. Because the constructor is unknown, the
// body remains opaque and is never decoded or allocated from; total frame/RPC budgets bound it.
b.PutUint32(0xffffffff)
if _, err := r.Dispatch(context.Background(), [8]byte{1}, 101, &b); !tgerr.Is(err, "NOT_IMPLEMENTED") {
t.Fatalf("unknown dispatch err = %v, want NOT_IMPLEMENTED", err)
}
entries := logs.FilterMessage("Unhandled RPC (compatibility trace)").All()
if len(entries) != 1 {
t.Fatalf("compatibility trace entries = %d, want 1", len(entries))
}
if got, ok := entries[0].ContextMap()["type_id"]; !ok || got != "0x12345678" {
t.Fatalf("trace type_id = %#v, want %#x", got, unknownID)
}
}

39
internal/rpc/album_group.go Normal file
View file

@ -0,0 +1,39 @@
package rpc
import (
"context"
"errors"
"go.uber.org/zap"
"telesrv/internal/domain"
)
func (r *Router) reserveAlbumGroup(ctx context.Context, userID int64, peer domain.Peer, items []domain.AlbumGroupReservationItem) (int64, error) {
reservations, ok := r.deps.Messages.(AlbumGroupService)
if !ok {
r.log.Error("messages.sendMultiMedia album reservation capability missing",
append(r.contextLogFields(ctx), zap.Int64("user_id", userID), zap.String("peer_type", string(peer.Type)), zap.Int64("peer_id", peer.ID))...)
return 0, internalErr()
}
groupedID, err := reservations.ReserveAlbumGroup(ctx, userID, domain.AlbumGroupReservationRequest{
SenderUserID: userID,
Peer: peer,
Items: items,
ProposedGroupedID: randomNonZeroInt64(),
})
if errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
return 0, randomIDDuplicateErr()
}
if err != nil {
r.log.Error("messages.sendMultiMedia album reservation failed",
append(r.contextLogFields(ctx), zap.Error(err), zap.Int64("user_id", userID), zap.String("peer_type", string(peer.Type)), zap.Int64("peer_id", peer.ID), zap.Int("items", len(items)))...)
return 0, internalErr()
}
if groupedID == 0 {
r.log.Error("messages.sendMultiMedia album reservation returned zero grouped_id",
append(r.contextLogFields(ctx), zap.Int64("user_id", userID), zap.String("peer_type", string(peer.Type)), zap.Int64("peer_id", peer.ID))...)
return 0, internalErr()
}
return groupedID, nil
}

84
internal/rpc/auth.go
View file

@ -241,6 +241,9 @@ func (r *Router) pushLoginTokenAccepted(ctx context.Context, target loginTokenTa
// 若该手机号账号设置了登录邮箱，验证码改投递到邮箱，返回 sentCodeTypeEmailCode
// （客户端据此进入"输入邮箱验证码"界面，随后用 auth.signIn 的 email_verification 完成登录）。
func (r *Router) onAuthSendCode(ctx context.Context, req *tg.AuthSendCodeRequest) (tg.AuthSentCodeClass, error) {
if err := r.checkAuthCodeRateLimit(ctx, req.PhoneNumber); err != nil {
return nil, err
}
r.rememberClientAPIID(ctx, req.APIID)
hash, err := r.deps.Auth.SendCode(ctx, req.PhoneNumber)
if err != nil {
@ -328,25 +331,21 @@ func (r *Router) tgSentCodeForHash(ctx context.Context, hash string) (tg.AuthSen
// 带 email_verification 时走登录邮箱路径（验证码来自邮箱而非短信）。
func (r *Router) onAuthSignIn(ctx context.Context, req *tg.AuthSignInRequest) (tg.AuthAuthorizationClass, error) {
var (
u domain.User
loginMessage domain.Message
needSignUp bool
err error
u domain.User
needSignUp bool
err error
)
if verification, ok := req.GetEmailVerification(); ok {
u, loginMessage, needSignUp, err = r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, emailVerificationCode(verification))
u, _, needSignUp, err = r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, emailVerificationCode(verification))
} else {
u, loginMessage, needSignUp, err = r.deps.Auth.SignIn(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, req.PhoneCode)
u, _, needSignUp, err = r.deps.Auth.SignIn(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, req.PhoneCode)
}
return r.finishAuthSignIn(ctx, u, loginMessage, needSignUp, err)
return r.finishAuthSignIn(ctx, u, needSignUp, err)
}
func (r *Router) finishAuthSignIn(ctx context.Context, u domain.User, loginMessage domain.Message, needSignUp bool, err error) (tg.AuthAuthorizationClass, error) {
func (r *Router) finishAuthSignIn(ctx context.Context, u domain.User, needSignUp bool, err error) (tg.AuthAuthorizationClass, error) {
if err != nil {
if errors.Is(err, domain.ErrSessionPasswordNeeded) && u.ID != 0 {
if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil {
return nil, internalErr()
}
// 两步验证未完成：绝不能把 auth_key/session 标记为已登录，否则客户端忽略
// SESSION_PASSWORD_NEEDED、直接调用业务 RPC 即可绕过 2FA。失效缓存并把 session
// 置为未授权，让后续鉴权重新读到 password_pending 并拒绝；待 checkPassword 通过后再授权。
@ -360,19 +359,18 @@ func (r *Router) finishAuthSignIn(ctx context.Context, u domain.User, loginMessa
if needSignUp {
return &tg.AuthAuthorizationSignUpRequired{}, nil
}
if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil {
return nil, internalErr()
}
if id, ok := AuthKeyIDFrom(ctx); ok {
r.setAuthUserCache(id, u.ID, true)
}
r.bindSessionUser(ctx, u.ID)
r.enqueueLoginMessageBootstrap(ctx, loginMessage)
r.pushSignInServiceNotificationToOthers(ctx, u)
return &tg.AuthAuthorization{User: r.tgSelfUser(u)}, nil
}
func (r *Router) onAuthResendCode(ctx context.Context, req *tg.AuthResendCodeRequest) (tg.AuthSentCodeClass, error) {
if err := r.checkAuthCodeRateLimit(ctx, req.PhoneNumber); err != nil {
return nil, err
}
var hash string
var err error
if scoped, ok := r.deps.Auth.(interface {
@ -545,15 +543,34 @@ func (r *Router) completePendingPasswordSignIn(ctx context.Context, authKeyID [8
// onAuthResetLoginEmail 处理 auth.resetLoginEmail：用户登录设备时无法访问登录邮箱时
// 清除登录邮箱，改回手机验证码登录，返回一个新的手机 sentCode 供其继续。
type loginEmailResetConsumer interface {
ConsumeLoginEmailReset(ctx context.Context, phone, phoneCodeHash string) (userID int64, err error)
SendPhoneCodeAfterLoginEmailReset(ctx context.Context, phone string, expectedUserID int64) (string, error)
}
func (r *Router) onAuthResetLoginEmail(ctx context.Context, req *tg.AuthResetLoginEmailRequest) (tg.AuthSentCodeClass, error) {
if r.deps.Account == nil || r.deps.Auth == nil {
return nil, internalErr()
}
if err := r.deps.Account.ClearLoginEmailByPhone(ctx, req.PhoneNumber); err != nil {
if err := r.checkAuthCodeRateLimit(ctx, req.PhoneNumber); err != nil {
return nil, err
}
resetConsumer, ok := r.deps.Auth.(loginEmailResetConsumer)
if !ok {
return nil, internalErr()
}
hash, err := r.deps.Auth.SendCode(ctx, req.PhoneNumber)
resetUserID, err := resetConsumer.ConsumeLoginEmailReset(ctx, req.PhoneNumber, req.PhoneCodeHash)
if err != nil {
return nil, signInErr(err)
}
if err := r.deps.Account.ClearLoginEmail(ctx, resetUserID); err != nil {
return nil, internalErr()
}
hash, err := resetConsumer.SendPhoneCodeAfterLoginEmailReset(ctx, req.PhoneNumber, resetUserID)
if err != nil {
if errors.Is(err, auth.ErrCodeExpired) || errors.Is(err, auth.ErrCodeInvalid) {
return nil, signInErr(err)
}
if errors.Is(err, auth.ErrPhoneNumberInvalid) ||
errors.Is(err, auth.ErrSystemUserLoginForbidden) {
return nil, phoneNumberInvalidErr()
@ -564,7 +581,7 @@ func (r *Router) onAuthResetLoginEmail(ctx context.Context, req *tg.AuthResetLog
}
// emailVerificationCode 从 emailVerification 取出可校验的字符串（验证码 / Google·Apple
// 令牌）。开发环境一律按"任意非空即通过"处理，故三者等价取值。
// 令牌）；最终必须由 auth service 对签发记录精确校验。
// onAuthInitPasskeyLogin 处理 auth.initPasskeyLogin：生成一次性断言挑战（discoverable），
// 以 DataJSON（顶层含 publicKey）返回。免授权（登录前）。
func (r *Router) onAuthInitPasskeyLogin(ctx context.Context, req *tg.AuthInitPasskeyLoginRequest) (*tg.AuthPasskeyLoginOptions, error) {
@ -579,7 +596,8 @@ func (r *Router) onAuthInitPasskeyLogin(ctx context.Context, req *tg.AuthInitPas
}
// onAuthFinishPasskeyLogin 处理 auth.finishPasskeyLogin：验证登录断言并绑定 auth_key。
// 收尾与 signIn 同构（清水位→授权缓存→session 绑定）；passkey 是强因子，直接完全授权
// 收尾与 signIn 同构（Bind 原子切换 update baseline → 授权缓存 → session 绑定）；
// passkey 是强因子，直接完全授权
// （不走 SESSION_PASSWORD_NEEDED）。FromDCID/FromAuthKeyID 为多 DC 重路由用，本单 DC 忽略。
func (r *Router) onAuthFinishPasskeyLogin(ctx context.Context, req *tg.AuthFinishPasskeyLoginRequest) (tg.AuthAuthorizationClass, error) {
if r.deps.Passkey == nil || r.deps.Auth == nil {
@ -597,9 +615,6 @@ func (r *Router) onAuthFinishPasskeyLogin(ctx context.Context, req *tg.AuthFinis
if err != nil {
return nil, passkeyErr(err)
}
if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil {
return nil, internalErr()
}
if id, ok := AuthKeyIDFrom(ctx); ok {
r.setAuthUserCache(id, u.ID, true)
}
@ -621,7 +636,7 @@ func emailVerificationCode(v tg.EmailVerificationClass) string {
// onAuthImportBotAuthorization 处理 auth.importBotAuthorization：bot 程序凭 token
// 登录为 bot 账号。api_id/api_hash 与现有 sendCode 行为一致不校验（无 app 注册表）。
// 收尾与 signIn 同构（清水位→授权缓存→session 绑定），但不写登录消息、不推
// 收尾与 signIn 同构（Bind 原子切换 update baseline → 授权缓存 → session 绑定），但不写登录消息、不推
// signIn 服务通知——那是手机登录语义。
func (r *Router) onAuthImportBotAuthorization(ctx context.Context, req *tg.AuthImportBotAuthorizationRequest) (tg.AuthAuthorizationClass, error) {
if r.deps.Auth == nil {
@ -631,9 +646,6 @@ func (r *Router) onAuthImportBotAuthorization(ctx context.Context, req *tg.AuthI
if err != nil {
return nil, importBotAuthorizationErr(err)
}
if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil {
return nil, internalErr()
}
if id, ok := AuthKeyIDFrom(ctx); ok {
r.setAuthUserCache(id, u.ID, true)
}
@ -647,9 +659,6 @@ func (r *Router) onAuthSignUp(ctx context.Context, req *tg.AuthSignUpRequest) (t
if err != nil {
return nil, signInErr(err)
}
if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil {
return nil, internalErr()
}
if id, ok := AuthKeyIDFrom(ctx); ok {
r.setAuthUserCache(id, u.ID, true)
}
@ -689,18 +698,6 @@ func (r *Router) onAuthLogOut(ctx context.Context) (*tg.AuthLoggedOut, error) {
return &tg.AuthLoggedOut{}, nil
}
func (r *Router) clearAuthKeyStateOnUserChange(ctx context.Context, newUserID int64) error {
oldUserID, ok := UserIDFrom(ctx)
if !ok || oldUserID == 0 || oldUserID == newUserID {
return nil
}
id, ok := AuthKeyIDFrom(ctx)
if !ok {
return nil
}
return r.clearAuthKeyState(ctx, id)
}
func (r *Router) clearAuthKeyState(ctx context.Context, authKeyID [8]byte) error {
if r.deps.Updates == nil {
return nil
@ -770,8 +767,9 @@ func (r *Router) pushSignInServiceNotificationToOthers(ctx context.Context, u do
return
}
authKeyID, hasAuthKeyID := AuthKeyIDFrom(ctx)
rawAuthKeyID, hasRawAuthKeyID := RawAuthKeyIDFrom(ctx)
sessionID, hasSessionID := SessionIDFrom(ctx)
if !hasAuthKeyID || !hasSessionID {
if !hasAuthKeyID || !hasRawAuthKeyID || !hasSessionID {
return
}
notification := r.tgSignInServiceNotification(ctx, u, authKeyID)
@ -779,7 +777,7 @@ func (r *Router) pushSignInServiceNotificationToOthers(ctx context.Context, u do
pushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if scoped, ok := r.scopedSessions(); ok {
if sent, err := scoped.PushToUserExceptAuthKeySession(pushCtx, u.ID, authKeyID, sessionID, proto.MessageFromServer, notification); err != nil {
if sent, err := scoped.PushToUserExceptAuthKeySession(pushCtx, u.ID, rawAuthKeyID, sessionID, proto.MessageFromServer, notification); err != nil {
r.log.Debug("push sign-in service notification", zap.Int64("user_id", u.ID), zap.Int("sent", sent), zap.Error(err))
}
return

248
internal/rpc/auth_code_rate_limit_test.go Normal file
View file

@ -0,0 +1,248 @@
package rpc
import (
"context"
"crypto/sha256"
"encoding/hex"
"strings"
"testing"
"time"
"github.com/gotd/td/clock"
"github.com/gotd/td/tg"
"go.uber.org/zap/zaptest"
"telesrv/internal/app/auth"
"telesrv/internal/domain"
)
type authCodeRateTestService struct {
*captureAuthService
sendCalls int
resendCalls int
resetCalls int
resetPhone string
resetHash string
resetUserID int64
resetErr error
}
func (s *authCodeRateTestService) SendCode(context.Context, string) (string, error) {
s.sendCalls++
return "send-hash", nil
}
func (s *authCodeRateTestService) ResendCode(context.Context, string, string) (string, error) {
s.resendCalls++
return "resend-hash", nil
}
func (s *authCodeRateTestService) ConsumeLoginEmailReset(_ context.Context, phone, hash string) (int64, error) {
s.resetCalls++
s.resetPhone = phone
s.resetHash = hash
return s.resetUserID, s.resetErr
}
func (s *authCodeRateTestService) SendPhoneCodeAfterLoginEmailReset(_ context.Context, _ string, expectedUserID int64) (string, error) {
s.sendCalls++
if expectedUserID != s.resetUserID {
return "", auth.ErrCodeInvalid
}
return "send-hash", nil
}
type authCodeRateTestAccount struct {
AccountService
clearCalls int
clearUserID int64
}
func (s *authCodeRateTestAccount) ClearLoginEmail(_ context.Context, userID int64) error {
s.clearCalls++
s.clearUserID = userID
return nil
}
func TestAuthSendCodeRateLimitUsesOpaquePhoneAndRawAuthKeyKeys(t *testing.T) {
phone := "+1 (555) 123-4567"
rawAuthKeyID := [8]byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef}
limiter := &captureRateLimiter{}
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}}
r := New(Config{
AuthCodePhoneRateLimit: 5,
AuthCodeAuthKeyRateLimit: 20,
AuthCodeRateWindow: 10 * time.Minute,
}, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System)
ctx := WithRawAuthKeyID(context.Background(), rawAuthKeyID)
if _, err := r.onAuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 2040}); err != nil {
t.Fatalf("onAuthSendCode: %v", err)
}
if authService.sendCalls != 1 {
t.Fatalf("SendCode calls = %d, want 1", authService.sendCalls)
}
if len(limiter.calls) != 2 {
t.Fatalf("limiter calls = %d, want phone + raw auth key", len(limiter.calls))
}
digest := sha256.Sum256([]byte(domain.NormalizePhone(phone)))
wantPhoneKey := authCodePhoneRateLimitKeyPrefix + hex.EncodeToString(digest[:])
wantAuthKey := authCodeAuthKeyRateLimitKeyPrefix + hex.EncodeToString(rawAuthKeyID[:])
if got := limiter.calls[0]; got.key != wantAuthKey || got.cost != 1 || got.limit != 20 || got.window != 10*time.Minute {
t.Fatalf("auth-key limiter call = %+v", got)
}
if got := limiter.calls[1]; got.key != wantPhoneKey || got.cost != 1 || got.limit != 5 || got.window != 10*time.Minute {
t.Fatalf("phone limiter call = %+v", got)
}
for _, call := range limiter.calls {
if strings.Contains(call.key, domain.NormalizePhone(phone)) || strings.Contains(call.key, phone) {
t.Fatalf("limiter key leaked phone: %q", call.key)
}
}
}
func TestAuthSendCodePhoneRateLimitPrecedesBusinessLookupAndWrite(t *testing.T) {
limiter := &captureRateLimiter{block: true, retryAfter: 17}
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}}
r := New(Config{
AuthCodePhoneRateLimit: 5,
AuthCodeRateWindow: time.Minute,
}, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System)
_, err := r.onAuthSendCode(WithRawAuthKeyID(context.Background(), [8]byte{1}), &tg.AuthSendCodeRequest{PhoneNumber: "+86 188 0000 0000", APIID: 2040})
if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") || !strings.Contains(err.Error(), "(17)") {
t.Fatalf("sendCode err = %v, want FLOOD_WAIT 17", err)
}
if authService.sendCalls != 0 {
t.Fatalf("SendCode calls = %d, want 0", authService.sendCalls)
}
if len(authService.authKeyClientInfos) != 0 {
t.Fatalf("blocked sendCode persisted client info: %+v", authService.authKeyClientInfos)
}
if len(limiter.calls) != 1 || !strings.HasPrefix(limiter.calls[0].key, authCodePhoneRateLimitKeyPrefix) {
t.Fatalf("limiter calls = %+v, want only phone dimension", limiter.calls)
}
}
func TestAuthSendCodeRawAuthKeyBlockDoesNotCreatePhoneDimension(t *testing.T) {
limiter := &captureRateLimiter{block: true, retryAfter: 31}
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}}
r := New(Config{
AuthCodePhoneRateLimit: 5,
AuthCodeAuthKeyRateLimit: 20,
AuthCodeRateWindow: time.Minute,
}, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System)
rawAuthKeyID := [8]byte{7, 7, 7, 7, 7, 7, 7, 7}
_, err := r.onAuthSendCode(WithRawAuthKeyID(context.Background(), rawAuthKeyID), &tg.AuthSendCodeRequest{PhoneNumber: "15550000001"})
if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") {
t.Fatalf("sendCode err = %v, want FLOOD_WAIT", err)
}
wantKey := authCodeAuthKeyRateLimitKeyPrefix + hex.EncodeToString(rawAuthKeyID[:])
if len(limiter.calls) != 1 || limiter.calls[0].key != wantKey {
t.Fatalf("limiter calls = %+v, want only raw auth-key %q", limiter.calls, wantKey)
}
if authService.sendCalls != 0 {
t.Fatalf("SendCode calls = %d, want 0", authService.sendCalls)
}
}
func TestAuthSendCodeInvalidPhoneCreatesNoLimiterKey(t *testing.T) {
limiter := &captureRateLimiter{}
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}}
r := New(Config{
AuthCodePhoneRateLimit: 5,
AuthCodeAuthKeyRateLimit: 20,
AuthCodeRateWindow: time.Minute,
}, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System)
_, err := r.onAuthSendCode(WithRawAuthKeyID(context.Background(), [8]byte{1}), &tg.AuthSendCodeRequest{PhoneNumber: "not-a-phone"})
if err == nil || !strings.Contains(err.Error(), "PHONE_NUMBER_INVALID") {
t.Fatalf("sendCode err = %v, want PHONE_NUMBER_INVALID", err)
}
if len(limiter.calls) != 0 || authService.sendCalls != 0 {
t.Fatalf("invalid phone limiter/service calls = %d/%d, want 0/0", len(limiter.calls), authService.sendCalls)
}
}
func TestAuthResendCodeRawAuthKeyRateLimitPrecedesRotation(t *testing.T) {
limiter := &captureRateLimiter{block: true, retryAfter: 23}
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}}
r := New(Config{
AuthCodeAuthKeyRateLimit: 20,
AuthCodeRateWindow: 2 * time.Minute,
}, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System)
rawAuthKeyID := [8]byte{9, 8, 7, 6, 5, 4, 3, 2}
_, err := r.onAuthResendCode(WithRawAuthKeyID(context.Background(), rawAuthKeyID), &tg.AuthResendCodeRequest{
PhoneNumber: "8618800000000",
PhoneCodeHash: "old-hash",
})
if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") || !strings.Contains(err.Error(), "(23)") {
t.Fatalf("resendCode err = %v, want FLOOD_WAIT 23", err)
}
if authService.resendCalls != 0 {
t.Fatalf("ResendCode calls = %d, want 0", authService.resendCalls)
}
wantKey := authCodeAuthKeyRateLimitKeyPrefix + hex.EncodeToString(rawAuthKeyID[:])
if len(limiter.calls) != 1 || limiter.calls[0].key != wantKey {
t.Fatalf("limiter calls = %+v, want %q", limiter.calls, wantKey)
}
}
func TestAuthResetLoginEmailRateLimitPrecedesEmailClear(t *testing.T) {
limiter := &captureRateLimiter{block: true, retryAfter: 29}
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}}
accountService := &authCodeRateTestAccount{}
r := New(Config{AuthCodePhoneRateLimit: 5, AuthCodeRateWindow: time.Minute}, Deps{
Auth: authService, Account: accountService, Limiter: limiter,
}, zaptest.NewLogger(t), clock.System)
_, err := r.onAuthResetLoginEmail(context.Background(), &tg.AuthResetLoginEmailRequest{
PhoneNumber: "8618800000000",
PhoneCodeHash: "email-hash",
})
if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") || !strings.Contains(err.Error(), "(29)") {
t.Fatalf("resetLoginEmail err = %v, want FLOOD_WAIT 29", err)
}
if accountService.clearCalls != 0 || authService.resetCalls != 0 || authService.sendCalls != 0 {
t.Fatalf("side effects reset=%d clear=%d send=%d, want 0/0/0", authService.resetCalls, accountService.clearCalls, authService.sendCalls)
}
}
func TestAuthResetLoginEmailConsumesHashBeforeClearAndResend(t *testing.T) {
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}, resetUserID: 4242}
accountService := &authCodeRateTestAccount{}
r := New(Config{}, Deps{Auth: authService, Account: accountService}, zaptest.NewLogger(t), clock.System)
req := &tg.AuthResetLoginEmailRequest{PhoneNumber: "+1 555 123 9999", PhoneCodeHash: "email-login-hash"}
result, err := r.onAuthResetLoginEmail(context.Background(), req)
if err != nil {
t.Fatalf("onAuthResetLoginEmail: %v", err)
}
if authService.resetCalls != 1 || authService.resetPhone != req.PhoneNumber || authService.resetHash != req.PhoneCodeHash ||
accountService.clearCalls != 1 || accountService.clearUserID != authService.resetUserID || authService.sendCalls != 1 {
t.Fatalf("calls reset=%d(%q,%q uid=%d) clear=%d(uid=%d) send=%d", authService.resetCalls, authService.resetPhone, authService.resetHash, authService.resetUserID, accountService.clearCalls, accountService.clearUserID, authService.sendCalls)
}
sent, ok := result.(*tg.AuthSentCode)
if !ok || sent.PhoneCodeHash != "send-hash" {
t.Fatalf("result=%T %+v, want sentCode/send-hash", result, result)
}
}
func TestAuthResetLoginEmailInvalidHashNeverClearsFactor(t *testing.T) {
authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}, resetErr: auth.ErrCodeExpired}
accountService := &authCodeRateTestAccount{}
r := New(Config{}, Deps{Auth: authService, Account: accountService}, zaptest.NewLogger(t), clock.System)
_, err := r.onAuthResetLoginEmail(context.Background(), &tg.AuthResetLoginEmailRequest{
PhoneNumber: "15551239999",
PhoneCodeHash: "expired-email-hash",
})
if err == nil || !strings.Contains(err.Error(), "PHONE_CODE_EXPIRED") {
t.Fatalf("onAuthResetLoginEmail err=%v, want PHONE_CODE_EXPIRED", err)
}
if authService.resetCalls != 1 || accountService.clearCalls != 0 || authService.sendCalls != 0 {
t.Fatalf("calls reset=%d clear=%d send=%d, want 1/0/0", authService.resetCalls, accountService.clearCalls, authService.sendCalls)
}
}

11
internal/rpc/auth_qr_login_test.go
View file

@ -107,15 +107,16 @@ func TestAuthLoginTokenAcceptedByAndroidBindsTargetSession(t *testing.T) {
if snap.sessionID != targetSession || snap.userID != scannerUserID || !snap.userResolved {
t.Fatalf("target session snapshot = %+v, want session/user/resolved %d/%d/true", snap, targetSession, scannerUserID)
}
if snap.messageType != proto.MessageFromServer {
t.Fatalf("push message type = %v, want MessageFromServer", snap.messageType)
}
if !sessions.immediatePushSeen() {
t.Fatal("login token update was not pushed through the immediate pre-auth path")
}
short, ok := snap.message.(*tg.UpdateShort)
immediateType, immediateMessage := sessions.immediatePushSnapshot()
if immediateType != proto.MessageFromServer {
t.Fatalf("immediate push message type = %v, want MessageFromServer", immediateType)
}
short, ok := immediateMessage.(*tg.UpdateShort)
if !ok {
t.Fatalf("push message = %T, want *tg.UpdateShort", snap.message)
t.Fatalf("immediate push message = %T, want *tg.UpdateShort", immediateMessage)
}
if _, ok := short.Update.(*tg.UpdateLoginToken); !ok {
t.Fatalf("pushed update = %T, want *tg.UpdateLoginToken", short.Update)

3
internal/rpc/bots_longtail.go
View file

@ -151,7 +151,6 @@ func (r *Router) sendBotAllowedServiceMessageWith(ctx context.Context, userID, b
if err != nil {
return domain.SendPrivateTextResult{}, err
}
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
return r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{
SenderUserID: userID,
@ -165,7 +164,7 @@ func (r *Router) sendBotAllowedServiceMessageWith(ctx context.Context, userID, b
},
},
Date: int(r.clock.Now().Unix()),
OriginAuthKeyID: authKeyID,
OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx),
OriginSessionID: sessionID,
RecipientBlocked: recipientBlocked,
})

788
internal/rpc/channel_fanout_dispatcher.go
View file

@ -2,8 +2,10 @@ package rpc
import (
"context"
"fmt"
"sync"
"sync/atomic"
"time"
"github.com/gotd/td/tg"
"go.uber.org/zap"
@ -23,12 +25,28 @@ import (
// DrKLO Android ~1.5s 乱序窗口与 TDesktop PtsWaiter 的连续性期望（设计 §10.2）。
// - 单实例 + 无 durable 重投队列 + 同 channel 串行 → 无乱序、无自重复，故 v1 不需要
// per-session at-most-once 双水位（那是 Phase 3 跨实例的事，设计 §9/§10.1）。
// - 有界队列满时丢弃当前 job 并告警：被丢 recipient 会在该 channel 下一条成功投递的
// pts 跳变时经 getChannelDifference 收敛（设计约束 B）。
// - 有界队列满时不静默丢弃恢复触发：真实 payload job 降级为按 channel 合并、只保留
// 最高 pts 的 UpdateChannelTooLong nudge。每个 shard 独立公平 drain；即使该频道随后
// 静默，也不依赖“下一条消息”才能触发 getChannelDifference（设计约束 B）。
const (
defaultChannelFanoutShards = 64
defaultChannelFanoutBuffer = 2048
// The old 64x2048 channel buffers eagerly retained up to 131k closures (each may capture a
// message batch and ~2k recipients). Keep one small FIFO per ordering shard and enforce a
// process-wide retained-byte budget below.
defaultChannelFanoutBuffer = 64
defaultChannelFanoutMaxQueuedJobs = 4096
defaultChannelFanoutMaxQueuedBytes = 256 << 20
defaultChannelFanoutOverflowPerShard = 256
defaultChannelNudgeWorkers = 8
defaultChannelNudgeQueue = 4096
defaultChannelFanoutRecoverySweepPage = 256
channelFanoutMinRetainedBytes int64 = 64 << 10
channelFanoutNudgeRetryMin = time.Millisecond
channelFanoutNudgeRetryMax = 50 * time.Millisecond
channelFanoutRecoveryRetryMin = 10 * time.Millisecond
channelFanoutRecoveryRetryMax = time.Second
defaultChannelFanoutNudgeDeadline = 5 * time.Second
)
// channelFanoutBuilder 按 viewer 构建该 viewer 视角的 channel updates。与同步
@ -37,7 +55,7 @@ const (
type channelFanoutBuilder func(ctx context.Context, viewerUserID int64) *tg.Updates
// channelFanoutJob 是一条频道 payload fan-out 任务。Pts 仅用于日志/折叠语义；真值仍是
// channel_update_events，worker 只做在线投递。originAuthKeyID 是业务视角 auth key
// channel_update_events，worker 只做在线投递。originAuthKeyID 是物理 raw auth key
// （与 SessionManager.shouldExcludeSession 的比较侧一致），用于显式排除发起设备——异步
// 执行时请求 ctx 已失效，不能再靠 ctx 派生排除。
type channelFanoutJob struct {
@ -50,6 +68,270 @@ type channelFanoutJob struct {
originSessionID int64
prefetch channelFanoutPrefetch
build channelFanoutBuilder
// retainedBytes is a conservative reservation for the request-derived closure, result
// snapshots and explicit recipient slice. It is charged before the job enters any queue.
retainedBytes int64
// queueSeq 是 dispatcher shard 内部的 FIFO 序号。仅成功进入正常 payload queue 的
// job 占用序号；overflow watermark 记录入队失败时已经接受的最大序号，等这些更早
// payload 处理完后才发 nudge，避免 nudge 越过其之前的正常 FIFO payload。
queueSeq uint64
}
// channelFanoutOverflow 是 queue full 时的 nudge-only 恢复水位。同一 channel 只保留
// 最大 pts；barrier 是该次 overflow 之前已经进入正常 FIFO 的最后一个 shard 序号。
type channelFanoutOverflow struct {
pts int
barrier uint64
}
// channelFanoutShard 把正常 payload FIFO 与 overflow nudge mailbox 放在同一个 worker
// 下。overflowOrder 每个 channel 最多出现一次；热点 channel 只更新 map 水位，不会占满
// order，从而不能把其它 channel 的唯一恢复 nudge 永久饿死。
type channelFanoutShard struct {
jobs chan channelFanoutJob
overflowWake chan struct{}
// overflowSpace is a generation channel, not a one-token notification. A slot
// release closes the current generation and installs a fresh channel while mu is
// held, waking every waiter that observed the old full mailbox. Each waiter then
// competes under mu for the actually available slots; losers observe the new
// generation and sleep again. This avoids losing N-1 wakeups when several slots
// are released before any of N waiters gets scheduled.
overflowSpace chan struct{}
mu sync.Mutex
nextSeq uint64
processedSeq uint64
overflow map[int64]channelFanoutOverflow
overflowOrder []int64
overflowLimit int
// overflowWaiters is guarded by mu and only covers the distinct-channel
// saturation slow path. Besides making the wait lifecycle explicit, it avoids
// allocating a fresh generation channel when no goroutine is subscribed.
overflowWaiters int
}
func newChannelFanoutShard(buffer int) *channelFanoutShard {
return &channelFanoutShard{
jobs: make(chan channelFanoutJob, buffer),
overflowWake: make(chan struct{}, 1),
overflowSpace: make(chan struct{}),
overflow: make(map[int64]channelFanoutOverflow),
overflowLimit: defaultChannelFanoutOverflowPerShard,
}
}
// enqueue 尝试把真实 payload 放入正常 FIFO；满时按 channel 合并最高 pts 的 nudge-only
// watermark。返回 true 表示正常入队，false 表示已经安全降级为 overflow watermark。
func (s *channelFanoutShard) enqueue(job channelFanoutJob) bool {
s.mu.Lock()
job.queueSeq = s.nextSeq + 1
select {
case s.jobs <- job:
s.nextSeq = job.queueSeq
s.mu.Unlock()
return true
default:
s.mu.Unlock()
return false
}
}
func (s *channelFanoutShard) enqueueOverflow(channelID int64, pts int) bool {
s.mu.Lock()
accepted := s.addOverflowLocked(channelID, pts)
s.mu.Unlock()
if accepted {
s.signalOverflow()
}
return accepted
}
func (s *channelFanoutShard) enqueueOverflowWait(ctx context.Context, channelID int64, pts int, stop <-chan struct{}) bool {
if ctx == nil {
ctx = context.Background()
}
select {
case <-ctx.Done():
return false
case <-stop:
return false
default:
}
// Try and capture the current space generation under the same lock. Separating
// these operations creates a classic missed-wakeup window: a drain may free the
// mailbox after the failed try but before the waiter starts observing the signal.
s.mu.Lock()
if s.addOverflowLocked(channelID, pts) {
s.mu.Unlock()
s.signalOverflow()
return true
}
space := s.overflowSpace
s.overflowWaiters++
s.mu.Unlock()
defer func() {
s.mu.Lock()
s.overflowWaiters--
s.mu.Unlock()
}()
// Same-channel overflow is the hot saturation path and only updates an existing map item.
// Distinct-channel saturation may wait here only from the dispatcher's fixed recovery sweep
// actor. RPC producers never call this method: they publish an O(1) global recovery generation
// when the bounded mailbox is full, so request goroutines cannot be exhausted by fan-out
// admission pressure.
for {
// Cardinality is full with distinct channels. Apply bounded-memory backpressure instead of
// allocating an unbounded recovery map or dropping the recovery watermark.
select {
case <-space:
case <-ctx.Done():
return false
case <-stop:
return false
}
// Do not turn a release racing with cancellation into admission after the caller ended.
select {
case <-ctx.Done():
return false
case <-stop:
return false
default:
}
// Retrying and subscribing to the next generation must also be atomic with
// respect to a release. Broadcast wakeups can be spurious for a particular
// waiter (another waiter may win the sole slot), so loop until accepted or stopped.
s.mu.Lock()
if s.addOverflowLocked(channelID, pts) {
s.mu.Unlock()
s.signalOverflow()
return true
}
space = s.overflowSpace
s.mu.Unlock()
}
}
func (s *channelFanoutShard) addOverflowLocked(channelID int64, pts int) bool {
item, exists := s.overflow[channelID]
if !exists {
if len(s.overflow) >= s.overflowLimit {
return false
}
s.overflowOrder = append(s.overflowOrder, channelID)
// The first overflow fixes the FIFO barrier. Later same-channel losses only raise the
// durable pts watermark: moving the barrier on every merge lets a continuously full
// payload queue keep the recovery nudge one slot behind forever. UpdateChannelTooLong is
// an idempotent catch-up trigger, so it is safe for its newest pts to overtake payloads
// accepted after the first loss; those payloads become harmless duplicates after
// getChannelDifference converges the client.
item.barrier = s.nextSeq
}
if pts > item.pts {
item.pts = pts
}
s.overflow[channelID] = item
return true
}
func (s *channelFanoutShard) markProcessed(seq uint64) {
s.mu.Lock()
if seq > s.processedSeq {
s.processedSeq = seq
}
s.mu.Unlock()
}
// signalOverflowSpaceLocked announces a mailbox-cardinality decrease. Callers
// must hold s.mu. Close-and-replace provides broadcast generations without an
// unbounded waiter list, goroutine-per-waiter, or lossy fixed-capacity token queue.
func (s *channelFanoutShard) signalOverflowSpaceLocked() {
if s.overflowWaiters == 0 {
return
}
close(s.overflowSpace)
s.overflowSpace = make(chan struct{})
}
// popOverflow 仅供 mailbox/cardinality 单元测试直接释放一个 overflow；生产 drain 必须走
// tryQueueOverflow，确保 nudgeJobs 真正接收成功前不删除水位。
func (s *channelFanoutShard) popOverflow() (channelID int64, pts int, ok bool) {
s.mu.Lock()
defer s.mu.Unlock()
for remaining := len(s.overflowOrder); remaining > 0; remaining-- {
channelID = s.overflowOrder[0]
s.overflowOrder = s.overflowOrder[1:]
item, exists := s.overflow[channelID]
if !exists {
continue
}
if item.barrier > s.processedSeq {
s.overflowOrder = append(s.overflowOrder, channelID)
continue
}
delete(s.overflow, channelID)
s.signalOverflowSpaceLocked()
return channelID, item.pts, true
}
return 0, 0, false
}
// tryQueueOverflow 尝试把一个 barrier 已满足的 overflow 水位非阻塞提交给共享 nudge queue。
// 只有 channel send 成功才从 mailbox 删除；queue 满时保留原 item（包含并发合并后的最高 pts）
// 和原 order 位置。整个操作在 shard.mu 下完成，因此不会出现“读到旧 pts 后删除新 pts”的竞态。
func (s *channelFanoutShard) tryQueueOverflow(offer func(channelFanoutNudge) bool) (queued, blocked bool) {
s.mu.Lock()
defer s.mu.Unlock()
for i := 0; i < len(s.overflowOrder); {
channelID := s.overflowOrder[i]
item, exists := s.overflow[channelID]
if !exists {
copy(s.overflowOrder[i:], s.overflowOrder[i+1:])
s.overflowOrder = s.overflowOrder[:len(s.overflowOrder)-1]
continue
}
if item.barrier > s.processedSeq {
i++
continue
}
if offer(channelFanoutNudge{channelID: channelID, pts: item.pts}) {
delete(s.overflow, channelID)
copy(s.overflowOrder[i:], s.overflowOrder[i+1:])
s.overflowOrder = s.overflowOrder[:len(s.overflowOrder)-1]
s.signalOverflowSpaceLocked()
return true, false
} else {
// Shared nudge workers are saturated. Keep the exact watermark and retry from
// the shard's bounded timer; do not remove or advance the mailbox entry, and
// never park this payload worker on the nudge queue.
return false, true
}
}
return false, false
}
func (s *channelFanoutShard) signalOverflow() {
select {
case s.overflowWake <- struct{}{}:
default:
}
}
func (s *channelFanoutShard) signalEligibleOverflow() {
s.mu.Lock()
eligible := false
for _, channelID := range s.overflowOrder {
if item, ok := s.overflow[channelID]; ok && item.barrier <= s.processedSeq {
eligible = true
break
}
}
s.mu.Unlock()
if eligible {
s.signalOverflow()
}
}
// channelFanoutPrefetch 在 worker 解析出最终 recipient 集合后、逐 viewer build 之前调用一次，
@ -60,30 +342,72 @@ type channelFanoutPrefetch func(ctx context.Context, viewers []int64)
// channelFanoutDispatcher 把频道 payload fan-out 移出发送者 RPC，按 channelID 分片串行处理。
type channelFanoutDispatcher struct {
r *Router
log *zap.Logger
shards []chan channelFanoutJob
started atomic.Bool
r *Router
log *zap.Logger
shards []*channelFanoutShard
started atomic.Bool
stopped atomic.Bool
stopCh chan struct{}
stopOnce sync.Once
enqueueMu sync.RWMutex
budgetMu sync.Mutex
queuedJobs int
queuedBytes int64
maxQueuedJobs int
maxQueuedBytes int64
nudgeJobs chan int64
nudgeWorkers int
nudgeTimeout time.Duration
nudgeMu sync.Mutex
// nudgePending and nudgeJobs form one bounded coalescing mailbox. nudgeJobs contains only
// channel ids; the mutable map value always holds the highest pts observed before a worker
// takes that id. A hot channel therefore occupies one slot rather than filling the queue.
nudgePending map[int64]int
nudgeLimit int
// recoveryGeneration is the terminal in-memory saturation fallback. It deliberately carries
// no channel id: the fixed recovery actor enumerates the online membership index and reloads
// each channel's durable max pts. Thus even when every key-bearing mailbox is full, publishing
// one constant-size generation cannot fail or block an RPC producer.
recoveryGeneration atomic.Uint64
recoveryCompleted atomic.Uint64
recoveryWake chan struct{}
// dropped 保留旧字段名供既有统计兼容；现在表示“真实 payload 因 queue full 被折叠为
// nudge-only overflow watermark”的次数，不再表示恢复触发也被静默丢弃。
dropped atomic.Int64
}
type channelFanoutNudge struct {
channelID int64
pts int
}
// enqueueChannelFanout 把一条 channel-payload-pts 的 fan-out 投入异步 dispatcher。
// 从请求 ctx 抓取发起设备的业务 auth key + session_id 显式带入 job，使异步 worker 仍能
// 从请求 ctx 抓取发起设备的 raw auth key + session_id 显式带入 job，使异步 worker 仍能
// 排除发起设备回显（请求 ctx 异步时已失效）。仅用于会推进客户端 channel PtsWaiter 的真实
// payload（新消息/编辑/删除/pin）；reaction/poll（viewer-only 零 pts）、participant/TTL/
// channel state（无 channel pts）、typing（transient）不走此路径（设计 §2.1/§5 分类）。
func (r *Router) enqueueChannelFanout(ctx context.Context, scope channelFanoutScope, originUserID, channelID int64, pts int, recipients []int64, build channelFanoutBuilder) {
r.enqueueChannelFanoutWithPrefetch(ctx, scope, originUserID, channelID, pts, recipients, nil, build)
r.enqueueChannelFanoutWithPrefetch(ctx, scope, originUserID, channelID, pts, recipients, 0, nil, build)
}
// enqueueChannelFanoutWithPrefetch 同 enqueueChannelFanout，但额外带一个跨 viewer 用户投影预热钩子
// （fan-out 模板化把每 recipient 的逐 viewer 投影折叠成一次 O(owner) 投影；见 prefetchChannelFanoutUsers）。
func (r *Router) enqueueChannelFanoutWithPrefetch(ctx context.Context, scope channelFanoutScope, originUserID, channelID int64, pts int, recipients []int64, prefetch channelFanoutPrefetch, build channelFanoutBuilder) {
func (r *Router) enqueueChannelFanoutWithPrefetch(ctx context.Context, scope channelFanoutScope, originUserID, channelID int64, pts int, recipients []int64, retainedFloor int64, prefetch channelFanoutPrefetch, build channelFanoutBuilder) {
if r.channelFanout == nil || build == nil {
return
}
originAuthKeyID, _ := AuthKeyIDFrom(ctx)
originAuthKeyID := rawAuthKeyIDForOrigin(ctx)
originSessionID, _ := SessionIDFrom(ctx)
retainedBytes := int64(inboundRPCBytesFrom(ctx)) + int64(len(recipients))*8 + 4096
if retainedFloor < channelFanoutMinRetainedBytes {
retainedFloor = channelFanoutMinRetainedBytes
}
if retainedBytes < retainedFloor {
retainedBytes = retainedFloor
}
r.channelFanout.Enqueue(ctx, channelFanoutJob{
scope: scope,
originUserID: originUserID,
@ -94,6 +418,7 @@ func (r *Router) enqueueChannelFanoutWithPrefetch(ctx context.Context, scope cha
originSessionID: originSessionID,
prefetch: prefetch,
build: build,
retainedBytes: retainedBytes,
})
}
@ -110,9 +435,22 @@ func newChannelFanoutDispatcher(r *Router, shards, buffer int) *channelFanoutDis
if buffer <= 0 {
buffer = defaultChannelFanoutBuffer
}
d := &channelFanoutDispatcher{r: r, log: r.log.Named("channel-fanout"), shards: make([]chan channelFanoutJob, shards)}
d := &channelFanoutDispatcher{
r: r,
log: r.log.Named("channel-fanout"),
shards: make([]*channelFanoutShard, shards),
stopCh: make(chan struct{}),
maxQueuedJobs: defaultChannelFanoutMaxQueuedJobs,
maxQueuedBytes: defaultChannelFanoutMaxQueuedBytes,
nudgeJobs: make(chan int64, defaultChannelNudgeQueue),
nudgeWorkers: defaultChannelNudgeWorkers,
nudgeTimeout: defaultChannelFanoutNudgeDeadline,
nudgePending: make(map[int64]int),
nudgeLimit: defaultChannelNudgeQueue,
recoveryWake: make(chan struct{}, 1),
}
for i := range d.shards {
d.shards[i] = make(chan channelFanoutJob, buffer)
d.shards[i] = newChannelFanoutShard(buffer)
}
return d
}
@ -124,22 +462,138 @@ func (d *channelFanoutDispatcher) Run(ctx context.Context) {
return
}
var wg sync.WaitGroup
for i := range d.shards {
wg.Add(1)
go func() {
defer wg.Done()
<-ctx.Done()
d.enqueueMu.Lock()
d.stopped.Store(true)
d.stopOnce.Do(func() { close(d.stopCh) })
d.enqueueMu.Unlock()
}()
wg.Add(1)
go func() {
defer wg.Done()
d.runRecoverySweeps(ctx)
}()
for range d.nudgeWorkers {
wg.Add(1)
ch := d.shards[i]
go func() {
defer wg.Done()
for {
select {
case <-ctx.Done():
return
case job := <-ch:
case channelID := <-d.nudgeJobs:
nudge, ok := d.takeNudge(channelID)
if !ok {
continue
}
timeout := d.nudgeTimeout
if timeout <= 0 {
timeout = defaultChannelFanoutNudgeDeadline
}
nudgeCtx, cancel := context.WithTimeout(ctx, timeout)
complete := d.r.runChannelFanoutOverflowNudge(nudgeCtx, nudge.channelID, nudge.pts)
cancel()
if !complete && ctx.Err() == nil {
// A deadline may leave only a prefix of online members nudged. Do not try to
// remember that recipient subset: request a durable max-pts sweep instead.
d.requestRecoverySweep("nudge deadline")
}
}
}
}()
}
for i := range d.shards {
wg.Add(1)
shard := d.shards[i]
go func() {
defer wg.Done()
var retryTimer *time.Timer
var retryC <-chan time.Time
retryDelay := channelFanoutNudgeRetryMin
stopRetryTimer := func() {
if retryTimer != nil {
retryTimer.Stop()
}
}
defer stopRetryTimer()
scheduleRetry := func() {
if retryC != nil {
return
}
if retryTimer == nil {
retryTimer = time.NewTimer(retryDelay)
} else {
retryTimer.Reset(retryDelay)
}
retryC = retryTimer.C
if retryDelay < channelFanoutNudgeRetryMax {
retryDelay *= 2
if retryDelay > channelFanoutNudgeRetryMax {
retryDelay = channelFanoutNudgeRetryMax
}
}
}
drain := func() {
// While a retry is armed, payload completions and coalescing wakeups must not
// defeat backoff and spin on a full shared queue.
if retryC != nil {
return
}
queued, blocked := d.drainOneOverflow(shard)
if queued {
retryDelay = channelFanoutNudgeRetryMin
shard.signalEligibleOverflow()
return
}
if blocked {
scheduleRetry()
}
}
for {
select {
case <-ctx.Done():
return
case job := <-shard.jobs:
d.r.runChannelFanoutJob(ctx, job)
d.releaseQueuedJob(job)
shard.markProcessed(job.queueSeq)
// 每处理一条正常 FIFO payload，主动尝试 drain 一条已经越过
// barrier 的 overflow。这样持续灌满正常队列的热点频道也不能
// 永久饿死其它频道的恢复 nudge。
drain()
case <-shard.overflowWake:
drain()
case <-retryC:
retryC = nil
drain()
}
}
}()
}
wg.Wait()
// Workers may choose ctx.Done while jobs remain buffered. Release every reservation and
// drop closure references so tests/restarts do not retain the global budget after shutdown.
for _, shard := range d.shards {
for {
select {
case job := <-shard.jobs:
d.releaseQueuedJob(job)
default:
goto drained
}
}
drained:
shard.mu.Lock()
clear(shard.overflow)
shard.overflowOrder = nil
shard.mu.Unlock()
}
d.nudgeMu.Lock()
clear(d.nudgePending)
d.nudgeMu.Unlock()
}
func (d *channelFanoutDispatcher) shardIndex(channelID int64) int {
@ -152,8 +606,10 @@ func (d *channelFanoutDispatcher) shardIndex(channelID int64) int {
}
// Enqueue 投递一条 fan-out 任务。dispatcher 未启动时同步执行（用请求 ctx，保持旧行为）；
// 已启动时投入对应分片，满则丢弃 + 告警（该 channel 下一条消息的 pts 跳变会经
// getChannelDifference 兜底）。
// 已启动时投入对应分片。满时正常 payload 不阻塞请求路径，而是按 channel 合并为最高 pts
// 的 nudge-only overflow watermark，由同 shard worker 在更早的 FIFO payload 后公平 drain。
// 若 overflow cardinality 也已满，只发布一个常量大小的全局 recovery generation；固定后台
// actor 随后从 durable channel pts 重建全部在线 channel 的 nudge。RPC goroutine 永不等待 slot。
func (d *channelFanoutDispatcher) Enqueue(reqCtx context.Context, job channelFanoutJob) {
if d == nil || job.build == nil {
return
@ -162,14 +618,256 @@ func (d *channelFanoutDispatcher) Enqueue(reqCtx context.Context, job channelFan
d.r.runChannelFanoutJob(reqCtx, job)
return
}
shard := d.shards[d.shardIndex(job.channelID)]
select {
case shard <- job:
default:
d.dropped.Add(1)
d.log.Warn("channel fanout queue full, dropped realtime push (recovered via next pts gap / getChannelDifference)",
zap.Int64("channel_id", job.channelID), zap.Int("pts", job.pts))
d.enqueueMu.RLock()
if d.stopped.Load() {
d.enqueueMu.RUnlock()
return
}
shard := d.shards[d.shardIndex(job.channelID)]
queued := false
if d.reserveQueuedJob(job) {
queued = shard.enqueue(job)
if queued {
d.enqueueMu.RUnlock()
return
}
d.releaseQueuedJob(job)
}
d.dropped.Add(1)
if job.scope != channelFanoutMembers || job.pts <= 0 {
// 当前所有 enqueue 入口均为 members + durable pts；若未来新增其它 scope，必须先
// 定义其 overflow 恢复面，不能误把 viewer-only/no-pts 更新伪装成 channel nudge。
d.log.Error("channel fanout queue full for non-coalescible job; overflow contract violated",
zap.Int64("channel_id", job.channelID), zap.Int("pts", job.pts), zap.Int("scope", int(job.scope)))
d.enqueueMu.RUnlock()
return
}
channelID, pts := job.channelID, job.pts
// The payload closure and recipient snapshot are no longer needed after normal queue
// admission failed. Make them unreachable before applying overflow-cardinality backpressure;
// otherwise blocked producers would retain unbudgeted request bodies while waiting for one of
// the fixed mailbox slots. Inbound RPC concurrency remains the producer-count bound.
job.recipients = nil
job.prefetch = nil
job.build = nil
if !shard.enqueueOverflow(channelID, pts) {
// Every key-bearing in-memory structure is bounded. Once the shard mailbox has no distinct
// channel slot, do not add another queue and do not park this RPC worker. A generation bit is
// enough because channels.pts/channel_update_events are already the durable truth: the fixed
// recovery actor can enumerate all online channel ids and reconstruct the highest watermark.
d.requestRecoverySweep("overflow cardinality full")
d.log.Warn("channel fanout overflow cardinality exhausted; scheduled durable max-pts recovery sweep",
zap.Int64("channel_id", channelID), zap.Int("pts", pts))
d.enqueueMu.RUnlock()
return
}
d.log.Warn("channel fanout capacity exhausted, coalesced realtime payload into highest-pts overflow nudge",
zap.Int64("channel_id", channelID), zap.Int("pts", pts))
d.enqueueMu.RUnlock()
}
func (d *channelFanoutDispatcher) reserveQueuedJob(job channelFanoutJob) bool {
size := job.retainedBytes
if size < channelFanoutMinRetainedBytes {
size = channelFanoutMinRetainedBytes
}
d.budgetMu.Lock()
defer d.budgetMu.Unlock()
if d.queuedJobs >= d.maxQueuedJobs || size > d.maxQueuedBytes-d.queuedBytes {
return false
}
d.queuedJobs++
d.queuedBytes += size
return true
}
func (d *channelFanoutDispatcher) releaseQueuedJob(job channelFanoutJob) {
size := job.retainedBytes
if size < channelFanoutMinRetainedBytes {
size = channelFanoutMinRetainedBytes
}
d.budgetMu.Lock()
d.queuedJobs--
d.queuedBytes -= size
if d.queuedJobs < 0 || d.queuedBytes < 0 {
panic("channel fanout queue budget underflow")
}
d.budgetMu.Unlock()
}
func (d *channelFanoutDispatcher) queuedBudgetSnapshot() (jobs int, bytes int64) {
d.budgetMu.Lock()
defer d.budgetMu.Unlock()
return d.queuedJobs, d.queuedBytes
}
func (d *channelFanoutDispatcher) drainOneOverflow(shard *channelFanoutShard) (queued, blocked bool) {
return shard.tryQueueOverflow(d.offerNudge)
}
// offerNudge inserts one channel id into the bounded shared queue and stores its mutable highest
// pts in nudgePending. It never blocks. A same-channel update succeeds even when cardinality is
// full because it consumes no additional queue slot.
func (d *channelFanoutDispatcher) offerNudge(nudge channelFanoutNudge) bool {
if nudge.channelID == 0 || nudge.pts <= 0 {
return true
}
d.nudgeMu.Lock()
if current, exists := d.nudgePending[nudge.channelID]; exists {
if nudge.pts > current {
d.nudgePending[nudge.channelID] = nudge.pts
}
d.nudgeMu.Unlock()
return true
}
if len(d.nudgePending) >= d.nudgeLimit {
d.nudgeMu.Unlock()
return false
}
d.nudgePending[nudge.channelID] = nudge.pts
select {
case d.nudgeJobs <- nudge.channelID:
d.nudgeMu.Unlock()
return true
default:
// nudgeJobs has the same cardinality bound as nudgePending. This branch is reachable only
// while a test overrides one without the other or an invariant regresses; roll back rather
// than retain an unreachable map entry.
delete(d.nudgePending, nudge.channelID)
d.nudgeMu.Unlock()
return false
}
}
func (d *channelFanoutDispatcher) takeNudge(channelID int64) (channelFanoutNudge, bool) {
d.nudgeMu.Lock()
pts, ok := d.nudgePending[channelID]
if ok {
delete(d.nudgePending, channelID)
}
d.nudgeMu.Unlock()
return channelFanoutNudge{channelID: channelID, pts: pts}, ok
}
func (d *channelFanoutDispatcher) requestRecoverySweep(reason string) {
generation := d.recoveryGeneration.Add(1)
select {
case d.recoveryWake <- struct{}{}:
default:
}
d.log.Debug("channel fanout durable recovery sweep requested",
zap.Uint64("generation", generation), zap.String("reason", reason))
}
// runRecoverySweeps owns the only potentially waiting overflow admission path. Producers publish
// generations and return; this fixed actor reconstructs channel ids from the live membership index
// and watermarks from durable channels.pts. A generation is marked complete only after every page
// and every channel in that page has successfully entered its shard's barrier-preserving overflow
// mailbox. Errors retain the generation and retry with bounded backoff.
func (d *channelFanoutDispatcher) runRecoverySweeps(ctx context.Context) {
completed := d.recoveryCompleted.Load()
retryDelay := channelFanoutRecoveryRetryMin
var retryTimer *time.Timer
defer func() {
if retryTimer != nil {
retryTimer.Stop()
}
}()
for {
target := d.recoveryGeneration.Load()
if target <= completed {
select {
case <-ctx.Done():
return
case <-d.recoveryWake:
continue
}
}
if err := d.sweepOnlineChannelRecovery(ctx); err != nil {
if ctx.Err() != nil {
return
}
d.log.Warn("channel fanout durable recovery sweep failed; retaining generation",
zap.Uint64("generation", target), zap.Duration("retry_in", retryDelay), zap.Error(err))
if retryTimer == nil {
retryTimer = time.NewTimer(retryDelay)
} else {
retryTimer.Reset(retryDelay)
}
// recoveryGeneration already records every concurrent request. A wake may shorten the
// idle wait before a healthy sweep, but it must never bypass failure backoff: otherwise
// sustained saturation plus a persistent DB error retries at producer rate.
select {
case <-ctx.Done():
return
case <-retryTimer.C:
}
if retryDelay < channelFanoutRecoveryRetryMax {
retryDelay *= 2
if retryDelay > channelFanoutRecoveryRetryMax {
retryDelay = channelFanoutRecoveryRetryMax
}
}
continue
}
completed = target
d.recoveryCompleted.Store(completed)
retryDelay = channelFanoutRecoveryRetryMin
d.log.Info("channel fanout durable recovery sweep completed", zap.Uint64("generation", completed))
// If a producer saturated after its channel had already been visited, generation is now
// greater than completed and the next loop immediately performs a fresh full pass.
}
}
func (d *channelFanoutDispatcher) sweepOnlineChannelRecovery(ctx context.Context) error {
sessions, ok := d.r.deps.Sessions.(ChannelFanoutRecoverySessionProvider)
if !ok {
return fmt.Errorf("sessions dependency lacks online channel recovery enumeration")
}
channels, ok := d.r.deps.Channels.(ChannelFanoutRecoveryPtsProvider)
if !ok {
return fmt.Errorf("channels dependency lacks durable max pts lookup")
}
channelIDs := sessions.OnlineChannelIDsSnapshot()
for i, channelID := range channelIDs {
if channelID <= 0 || (i > 0 && channelID <= channelIDs[i-1]) {
return fmt.Errorf("online channel recovery snapshot is not strictly ascending: index=%d got=%d", i, channelID)
}
}
for start := 0; start < len(channelIDs); start += defaultChannelFanoutRecoverySweepPage {
end := start + defaultChannelFanoutRecoverySweepPage
if end > len(channelIDs) {
end = len(channelIDs)
}
page := channelIDs[start:end]
ptsByChannel, err := channels.MaxChannelPtsBatch(ctx, page)
if err != nil {
return fmt.Errorf("load durable max pts for online channel page [%d:%d]: %w", start, end, err)
}
for _, channelID := range page {
pts := ptsByChannel[channelID]
if pts > 0 {
shard := d.shards[d.shardIndex(channelID)]
if !shard.enqueueOverflowWait(ctx, channelID, pts, d.stopCh) {
if err := ctx.Err(); err != nil {
return err
}
return fmt.Errorf("dispatcher stopped while admitting recovery for channel %d", channelID)
}
}
}
}
return nil
}
// runChannelFanoutOverflowNudge 是 queue-full 的 nudge-only 降级路径。不能复用原 job 的
// origin exclude：同一 channel 水位可能合并多个不同发起 session；向全部在线成员发最高 pts
// nudge 是保守且幂等的，已追上 pts 的 TDesktop 会直接忽略。
func (r *Router) runChannelFanoutOverflowNudge(ctx context.Context, channelID int64, pts int) bool {
if r.deps.Sessions == nil || channelID == 0 || pts <= 0 {
return true
}
return r.nudgeBeyondCapChannelMembers(ctx, channelID, pts, nil)
}
// runChannelFanoutJob 执行一条 fan-out：与同步 pushChannelUpdatesWithScope 等价，区别是
@ -179,7 +877,7 @@ func (r *Router) runChannelFanoutJob(ctx context.Context, job channelFanoutJob)
if r.deps.Sessions == nil || job.build == nil {
return
}
pushCtx := WithSessionID(WithAuthKeyID(ctx, job.originAuthKeyID), job.originSessionID)
pushCtx := WithSessionID(WithRawAuthKeyID(ctx, job.originAuthKeyID), job.originSessionID)
recipients := r.channelFanoutRecipients(ctx, job.scope, job.channelID, job.recipients)
// 预热跨 viewer 用户投影（fan-out 模板化）：在逐 viewer build 之前一次性算好每 recipient 的
// 投影并预热共享 cache，使 build 只命中缓存、不再 O(viewer) 逐个 ForViewer。覆盖 recipients +
@ -278,6 +976,7 @@ func (r *Router) enqueueChannelMessageFanout(ctx context.Context, originUserID i
ownerIDs := channelMessageFanoutOwnerIDs(res, extraUserIDs)
skip := skipDeliverySet(res.SkipDeliveryUserIDs)
r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, res.Channel.ID, res.Event.Pts, res.Recipients,
0,
func(bgCtx context.Context, viewers []int64) {
r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs)
},
@ -342,6 +1041,7 @@ func (r *Router) enqueueChannelEditMessageFanout(ctx context.Context, originUser
ownerIDs := channelEditMessageFanoutOwnerIDs(res)
nudgePts := max(res.Event.Pts, res.ServiceEvent.Pts)
r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, res.Channel.ID, nudgePts, res.Recipients,
0,
func(bgCtx context.Context, viewers []int64) {
r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs)
},
@ -358,6 +1058,7 @@ func (r *Router) enqueueChannelMessagesFanout(ctx context.Context, originUserID,
fanoutCache := newViewerPeerCache(r)
ownerIDs := channelMessagesFanoutOwnerIDs(results, extraUserIDs)
r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, channelID, pts, recipients,
int64(len(results))*(64<<10),
func(bgCtx context.Context, viewers []int64) {
r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs)
},
@ -385,28 +1086,37 @@ func (r *Router) channelNudgeMaxTargets() int {
// getChannelDifference（设计 §10.3）。走 pushUserUpdates（best-effort、未就绪入 pending、非
// transient），符合设计 §决策4 的 nudge 投递可靠性要求。SessionManager 未实现 ChannelNudgeProvider
// 时（测试/未装配）静默跳过，不影响完整 payload 投递。
func (r *Router) nudgeBeyondCapChannelMembers(ctx context.Context, channelID int64, pts int, delivered map[int64]struct{}) {
func (r *Router) nudgeBeyondCapChannelMembers(ctx context.Context, channelID int64, pts int, delivered map[int64]struct{}) bool {
provider, ok := r.deps.Sessions.(ChannelNudgeProvider)
if !ok || channelID == 0 || pts <= 0 {
return
return true
}
targets := provider.OnlineChannelMemberUserIDsExcluding(channelID, delivered, r.channelNudgeMaxTargets())
if len(targets) == 0 {
return
return true
}
date := int(r.clock.Now().Unix())
tooLong := &tg.UpdateChannelTooLong{ChannelID: channelID}
tooLong.SetPts(pts)
updates := &tg.Updates{
Updates: []tg.UpdateClass{tooLong},
Users: []tg.UserClass{},
Chats: []tg.ChatClass{},
Date: date,
Seq: 0,
}
for _, userID := range targets {
select {
case <-ctx.Done():
return false
default:
}
if userID == 0 {
continue
}
tooLong := &tg.UpdateChannelTooLong{ChannelID: channelID}
tooLong.SetPts(pts)
r.pushUserUpdates(ctx, userID, &tg.Updates{
Updates: []tg.UpdateClass{tooLong},
Users: []tg.UserClass{},
Chats: []tg.ChatClass{},
Date: date,
Seq: 0,
})
// The nudge is viewer-independent and immutable. Reuse the TL object across
// recipients; SessionManager encodes before enqueue and never mutates it.
r.pushUserUpdates(ctx, userID, updates)
}
return ctx.Err() == nil
}

1145
internal/rpc/channel_fanout_dispatcher_test.go
View file

File diff suppressed because it is too large Load diff

4
internal/rpc/channels_legacy_chat.go
View file

@ -445,14 +445,13 @@ func (r *Router) onMessagesSetChatTheme(ctx context.Context, req *tg.MessagesSet
if err != nil {
return nil, err
}
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
res, err := r.deps.Messages.SetChatTheme(ctx, userID, domain.SetPrivateChatThemeRequest{
OwnerUserID: userID,
Peer: peer,
Emoticon: emoticon,
Date: int(r.clock.Now().Unix()),
OriginAuthKeyID: authKeyID,
OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx),
OriginSessionID: sessionID,
RecipientBlocked: recipientBlocked,
})
@ -578,6 +577,7 @@ func (r *Router) enqueueChannelWallpaperFanout(ctx context.Context, originUserID
fanoutCache := newViewerPeerCache(r)
ownerIDs := channelMessageFanoutOwnerIDs(sendRes, nil)
r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, res.Channel.ID, res.Event.Pts, res.Recipients,
0,
func(bgCtx context.Context, viewers []int64) {
r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs)
},

6
internal/rpc/channels_members.go
View file

@ -369,7 +369,11 @@ func (r *Router) recordChannelStateForUser(ctx context.Context, userID, channelI
authKeyID, _ = AuthKeyIDFrom(ctx)
excludeSessionID, _ = SessionIDFrom(ctx)
}
event, _, err := r.deps.Updates.RecordChannelState(ctx, authKeyID, userID, channelID, excludeSessionID)
excludeAuthKeyID := [8]byte{}
if excludeCurrent {
excludeAuthKeyID = rawAuthKeyIDForOrigin(ctx)
}
event, _, err := r.deps.Updates.RecordChannelState(ctx, authKeyID, userID, channelID, excludeAuthKeyID, excludeSessionID)
if err != nil {
return
}

6
internal/rpc/channels_stubs.go
View file

@ -264,7 +264,7 @@ func (r *Router) recordChannelAvailableMessages(ctx context.Context, userID, cha
}
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
recorded, _, err := r.deps.Updates.RecordChannelAvailableMessages(ctx, authKeyID, userID, channelID, availableMinID, sessionID)
recorded, _, err := r.deps.Updates.RecordChannelAvailableMessages(ctx, authKeyID, userID, channelID, availableMinID, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return event
}
@ -300,7 +300,7 @@ func (r *Router) recordChannelReadInbox(ctx context.Context, userID int64, read
StillUnreadCount: read.StillUnreadCount,
ChannelPts: read.Pts,
Changed: read.Changed,
}, sessionID)
}, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return domain.UpdateEvent{}, internalErr()
}
@ -673,6 +673,8 @@ func channelInvalidErr(err error) error {
return tgerr400("USER_ALREADY_PARTICIPANT")
case errors.Is(err, domain.ErrReplyMessageIDInvalid):
return replyMessageIDInvalidErr()
case errors.Is(err, domain.ErrMessageRandomIDDuplicate):
return randomIDDuplicateErr()
default:
if seconds, ok := domain.SlowModeWaitSeconds(err); ok {
return tgerr.New(420, fmt.Sprintf("SLOWMODE_WAIT_%d", seconds))

2
internal/rpc/channels_topics.go
View file

@ -42,7 +42,7 @@ func (r *Router) onChannelsToggleViewForumAsMessages(ctx context.Context, req *t
if r.deps.Updates != nil {
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
event, _, err = r.deps.Updates.RecordChannelViewForumAsMessages(ctx, authKeyID, userID, channelID, req.Enabled, sessionID)
event, _, err = r.deps.Updates.RecordChannelViewForumAsMessages(ctx, authKeyID, userID, channelID, req.Enabled, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return nil, internalErr()
}

8
internal/rpc/channels_updates.go
View file

@ -206,6 +206,14 @@ func (r *Router) channelMessagesUpdatesWithPeerCache(ctx context.Context, viewer
updates = append(updates, update)
}
}
if res.Duplicate && res.ReplayDeleteEvent != nil {
if update := tgChannelUpdate(viewerUserID, *res.ReplayDeleteEvent); update != nil {
updates = append(updates, update)
}
if res.ReplayDeleteEvent.Date > date {
date = res.ReplayDeleteEvent.Date
}
}
collectChannelUpdatePeerRefs(res.Event, res.Channel.ID, userIDs, channelIDs)
collectChannelMessagePeerRefs(res.Message, res.Channel.ID, userIDs, channelIDs)
if date == 0 {

4
internal/rpc/chatlists.go
View file

@ -379,7 +379,7 @@ func (r *Router) recordChatlistFilterUpdate(ctx context.Context, userID int64, f
if r.deps.Updates != nil {
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, sessionID)
event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return internalErr()
}
@ -400,7 +400,7 @@ func (r *Router) chatlistFilterUpdates(ctx context.Context, userID int64, filter
if r.deps.Updates != nil {
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, sessionID)
event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return nil, internalErr()
}

18
internal/rpc/contacts.go
View file

@ -984,10 +984,10 @@ func (r *Router) recordAcceptedContactTargetUpdates(ctx context.Context, userID,
return internalErr()
}
var zeroAuthKeyID [8]byte
if err := r.recordPeerSettingsForUser(ctx, zeroAuthKeyID, targetUserID, peer, settings, 0); err != nil {
if err := r.recordPeerSettingsForUser(ctx, zeroAuthKeyID, targetUserID, peer, settings, zeroAuthKeyID, 0); err != nil {
return internalErr()
}
if err := r.recordContactsResetForUser(ctx, zeroAuthKeyID, targetUserID, 0); err != nil {
if err := r.recordContactsResetForUser(ctx, zeroAuthKeyID, targetUserID, zeroAuthKeyID, 0); err != nil {
return internalErr()
}
peerUser := domain.User{ID: userID}
@ -1017,14 +1017,14 @@ func (r *Router) pushContactsReset(ctx context.Context, userID int64) {
func (r *Router) recordContactsReset(ctx context.Context, userID int64) error {
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
return r.recordContactsResetForUser(ctx, authKeyID, userID, sessionID)
return r.recordContactsResetForUser(ctx, authKeyID, userID, rawAuthKeyIDForOrigin(ctx), sessionID)
}
func (r *Router) recordContactsResetForUser(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) error {
func (r *Router) recordContactsResetForUser(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) error {
if r.deps.Updates == nil || userID == 0 {
return nil
}
event, _, err := r.deps.Updates.RecordContactsReset(ctx, authKeyID, userID, excludeSessionID)
event, _, err := r.deps.Updates.RecordContactsReset(ctx, stateAuthKeyID, userID, excludeAuthKeyID, excludeSessionID)
if err == nil && excludeSessionID != 0 {
r.bookkeepAuxPtsForCurrentSession(ctx, event)
}
@ -1034,7 +1034,7 @@ func (r *Router) recordContactsResetForUser(ctx context.Context, authKeyID [8]by
func (r *Router) recordPeerSettings(ctx context.Context, userID int64, peer domain.Peer, settings domain.PeerSettings) error {
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
return r.recordPeerSettingsForUser(ctx, authKeyID, userID, peer, settings, sessionID)
return r.recordPeerSettingsForUser(ctx, authKeyID, userID, peer, settings, rawAuthKeyIDForOrigin(ctx), sessionID)
}
func (r *Router) recordPeerStoryBlocked(ctx context.Context, userID int64, peer domain.Peer, blocked bool) error {
@ -1043,18 +1043,18 @@ func (r *Router) recordPeerStoryBlocked(ctx context.Context, userID int64, peer
if r.deps.Updates == nil || userID == 0 {
return nil
}
event, _, err := r.deps.Updates.RecordPeerStoryBlocked(ctx, authKeyID, userID, peer, blocked, sessionID)
event, _, err := r.deps.Updates.RecordPeerStoryBlocked(ctx, authKeyID, userID, peer, blocked, rawAuthKeyIDForOrigin(ctx), sessionID)
if err == nil && sessionID != 0 {
r.bookkeepAuxPtsForCurrentSession(ctx, event)
}
return err
}
func (r *Router) recordPeerSettingsForUser(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) error {
func (r *Router) recordPeerSettingsForUser(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeAuthKeyID [8]byte, excludeSessionID int64) error {
if r.deps.Updates == nil || userID == 0 {
return nil
}
event, _, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, settings, excludeSessionID)
event, _, err := r.deps.Updates.RecordPeerSettings(ctx, stateAuthKeyID, userID, peer, settings, excludeAuthKeyID, excludeSessionID)
if err == nil && excludeSessionID != 0 {
r.bookkeepAuxPtsForCurrentSession(ctx, event)
}

23
internal/rpc/context.go
View file

@ -16,8 +16,21 @@ const (
sessionIDKey
userIDKey
invokeWithoutUpdatesKey
inboundRPCBytesKey
)
func withInboundRPCBytes(ctx context.Context, n int) context.Context {
if n < 0 {
n = 0
}
return context.WithValue(ctx, inboundRPCBytesKey, n)
}
func inboundRPCBytesFrom(ctx context.Context) int {
v, _ := ctx.Value(inboundRPCBytesKey).(int)
return v
}
const currentClientLayer = 227
var androidSDKVersionRE = regexp.MustCompile(`\bsdk\s+\d+\b`)
@ -153,6 +166,16 @@ func AuthKeyIDFrom(ctx context.Context) ([8]byte, bool) {
return v, ok
}
// rawAuthKeyIDForOrigin 返回用于 update/outbox 当前 session 排除的物理 raw key。
// 单测/非 edge 调用若没有注入 raw key，才回退业务 key；生产 Router context 两者都有。
func rawAuthKeyIDForOrigin(ctx context.Context) [8]byte {
if id, ok := RawAuthKeyIDFrom(ctx); ok {
return id
}
id, _ := AuthKeyIDFrom(ctx)
return id
}
// WithSessionID 在 ctx 注入调用方的 MTProto session_id。
func WithSessionID(ctx context.Context, id int64) context.Context {
return context.WithValue(ctx, sessionIDKey, id)

74
internal/rpc/deps.go
View file

@ -71,6 +71,7 @@ type ScopedSessionBinder interface {
UserIDResolvedForAuthKey(rawAuthKeyID [8]byte, sessionID int64) (userID int64, resolved bool)
SetReceivesUpdatesForAuthKey(rawAuthKeyID [8]byte, sessionID int64, receives bool)
PushToSessionForAuthKey(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error
// excludeAuthKeyID is the physical/raw auth key, paired with session_id.
PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error)
}
@ -158,6 +159,21 @@ type ChannelNudgeProvider interface {
OnlineChannelMemberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64
}
// ChannelFanoutRecoverySessionProvider snapshots the process-local online joined-channel index in
// stable channel-id order. It is used only after every keyed fan-out mailbox is saturated. The
// fixed recovery actor accepts the temporary 8*C id slice so one sweep never repeatedly scans the
// SessionManager index or holds its global lock while sorting/database work runs.
type ChannelFanoutRecoverySessionProvider interface {
OnlineChannelIDsSnapshot() []int64
}
// ChannelFanoutRecoveryPtsProvider reloads the authoritative channel pts after in-memory fan-out
// saturation. Production channels.Service implements it through the channel store; keeping this
// separate from ChannelsService avoids burdening lightweight RPC fakes that never run the worker.
type ChannelFanoutRecoveryPtsProvider interface {
MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error)
}
// RateLimiter 抽象 RPC 高频写操作限流。
type RateLimiter interface {
Allow(ctx context.Context, key string, limit int, window time.Duration) (allowed bool, retryAfterSeconds int, err error)
@ -255,7 +271,7 @@ type UserPremiumStatusService interface {
// AccountService 抽象账号设置查询。
type AccountService interface {
SendChangePhoneCode(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64, phone string) (string, domain.AuthCodeDelivery, error)
ChangePhone(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error)
ChangePhone(ctx context.Context, userID int64, authKeyID, originRawAuthKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error)
GetPassword(ctx context.Context, userID int64) (domain.PasswordSettings, error)
GetPasswordSettings(ctx context.Context, userID int64, check domain.PasswordCheck) (domain.PrivatePasswordSettings, error)
UpdatePasswordSettings(ctx context.Context, userID int64, check domain.PasswordCheck, input domain.PasswordInputSettings) error
@ -270,10 +286,9 @@ type AccountService interface {
SendLoginEmailCode(ctx context.Context, userID int64, phone, phoneCodeHash, email string, setup bool) (string, int, error)
VerifyLoginEmail(ctx context.Context, userID int64, phone, phoneCodeHash, code string, setup bool) (string, error)
SetLoginEmail(ctx context.Context, userID int64, email string) error
SetLoginEmailByPhone(ctx context.Context, phone, email string) error
LoginEmail(ctx context.Context, userID int64) (string, bool, error)
LoginEmailByPhone(ctx context.Context, phone string) (string, bool, error)
ClearLoginEmailByPhone(ctx context.Context, phone string) error
ClearLoginEmail(ctx context.Context, userID int64) error
ResetPassword(ctx context.Context, userID int64) (domain.PasswordResetResult, error)
DeclinePasswordReset(ctx context.Context, userID int64) error
SaveMusic(ctx context.Context, userID int64, req domain.SaveMusicRequest) (bool, error)
@ -336,30 +351,30 @@ type UpdatesService interface {
ClearAuthKey(ctx context.Context, authKeyID [8]byte) error
RecordNewMessage(ctx context.Context, authKeyID [8]byte, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error)
PublishNewMessage(ctx context.Context, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error)
RecordStory(ctx context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordStory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, story domain.Story, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordStoryFanout(ctx context.Context, userID int64, story domain.Story) (domain.UpdateEvent, domain.UpdateState, error)
RecordReadStories(ctx context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordSentStoryReaction(ctx context.Context, authKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordQuickReplyMutation(ctx context.Context, authKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordReadHistory(ctx context.Context, authKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordContactsReset(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelState(ctx context.Context, authKeyID [8]byte, userID, channelID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPinnedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordSavedDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPinnedSavedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogUnreadMark(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPeerSettings(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPeerStoryBlocked(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogFilter(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogFilterOrder(ctx context.Context, authKeyID [8]byte, userID int64, order []int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogFiltersReload(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordFolderPeers(ctx context.Context, authKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelAvailableMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, availableMinID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelViewForumAsMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, enabled bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelDiscussionInbox(ctx context.Context, authKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDraftMessage(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordReadStories(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordSentStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordNewStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordQuickReplyMutation(ctx context.Context, stateAuthKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordReadHistory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordContactsReset(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelState(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPinnedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordSavedDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPinnedSavedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogUnreadMark(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPeerSettings(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordPeerStoryBlocked(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogFilter(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogFilterOrder(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDialogFiltersReload(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordFolderPeers(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelAvailableMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, availableMinID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelViewForumAsMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, enabled bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordChannelDiscussionInbox(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
RecordDraftMessage(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error)
}
// ContactsService 抽象通讯录查询。
@ -455,6 +470,13 @@ type MessagesService interface {
DeleteSavedHistory(ctx context.Context, userID int64, req domain.DeleteSavedHistoryRequest) (domain.DeleteSavedHistoryResult, error)
}
// AlbumGroupService 是 MessagesService 的可选、生产必备能力：sendMultiMedia 在
// 解析任何媒体或落第一条消息前，持久预留整批 random_id 的 grouped_id。
// 单独定义可避免让不触发 sendMultiMedia 的轻量测试替身实现无关方法。
type AlbumGroupService interface {
ReserveAlbumGroup(ctx context.Context, userID int64, req domain.AlbumGroupReservationRequest) (int64, error)
}
// StoriesService 抽象 story 读取、已读、观看与 reaction 状态。
type StoriesService interface {
CreateStory(ctx context.Context, userID int64, req domain.StoryCreateRequest) (domain.StoryCreateResult, error)

9
internal/rpc/dialogs_rpc_test.go
View file

@ -172,7 +172,11 @@ func TestMessagesGetPeerDialogsReturnsRequestedDialogsAndState(t *testing.T) {
func TestDialogSettingRPCsRecordDurableUpdates(t *testing.T) {
var authKeyID [8]byte
authKeyID[0] = 9
ctx := WithSessionID(WithAuthKeyID(WithUserID(context.Background(), 1000000001), authKeyID), 55)
rawAuthKeyID := [8]byte{9, 7}
ctx := WithRawAuthKeyID(
WithSessionID(WithAuthKeyID(WithUserID(context.Background(), 1000000001), authKeyID), 55),
rawAuthKeyID,
)
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
dialogPeer := &tg.InputDialogPeer{Peer: &tg.InputPeerUser{UserID: peer.ID}}
dialogs := &captureDialogs{}
@ -187,6 +191,9 @@ func TestDialogSettingRPCsRecordDurableUpdates(t *testing.T) {
if len(updates.events) != 1 || updates.events[0].Type != domain.UpdateEventDialogPinned || updates.events[0].Peer != peer || !updates.events[0].Bool || updates.excludeSessionID != 55 {
t.Fatalf("pin event = %+v, want durable dialog_pinned", updates.events)
}
if updates.authKeyID != authKeyID || updates.excludeAuthKeyID != rawAuthKeyID {
t.Fatalf("durable update keys = state:%x exclude:%x, want business:%x raw:%x", updates.authKeyID, updates.excludeAuthKeyID, authKeyID, rawAuthKeyID)
}
if ok, err := r.onMessagesReorderPinnedDialogs(ctx, &tg.MessagesReorderPinnedDialogsRequest{Order: []tg.InputDialogPeerClass{dialogPeer}}); err != nil || !ok {
t.Fatalf("reorder pinned = %v, %v", ok, err)

2
internal/rpc/encrypted_chats.go
View file

@ -88,7 +88,7 @@ func (r *Router) recordEncryptionEventBestEffort(ctx context.Context, chatID int
// 全部活跃密聊，并向对端推送 encryptedChatDiscarded（在线）+ 写 durable 事件（离线 getDifference
// 补偿）。ownerUserID 是被销毁设备的所有者，用于定位对端。best-effort：失败仅记日志，绝不阻断
// 登出/撤销。修复 P1：此前 onAuthLogOut 等不级联 discard，对端继续往死 auth_key 投递成静默死链
//（消息 acked=f / qts 永久积压，对端永看不到 discarded）。
// （消息 acked=f / qts 永久积压，对端永看不到 discarded）。
func (r *Router) discardSecretChatsForAuthKey(ctx context.Context, businessAuthKeyID, ownerUserID int64) {
if r.deps.SecretChats == nil || businessAuthKeyID == 0 || ownerUserID == 0 {
return

6
internal/rpc/errors.go
View file

@ -271,6 +271,8 @@ func chatForwardsRestrictedErr() error { return tgerr.New(400, "CHAT_FORWARDS_RE
func inputRequestInvalidErr() error { return tgerr.New(400, "INPUT_REQUEST_INVALID") }
func inputRequestTooLongErr() error { return tgerr.New(400, "INPUT_REQUEST_TOO_LONG") }
func persistentTimestampInvalidErr() error { return tgerr.New(400, "PERSISTENT_TIMESTAMP_INVALID") }
func channelForumMissingErr() error { return tgerr.New(400, "CHANNEL_FORUM_MISSING") }
@ -281,6 +283,10 @@ func topicIDInvalidErr() error { return tgerr.New(400, "TOPIC_ID_INVALID") }
// randomIDEmptyErr 表示发送消息缺少 random_id。
func randomIDEmptyErr() error { return tgerr.New(400, "RANDOM_ID_EMPTY") }
// randomIDDuplicateErr 表示同一发送者重复使用 random_id，但请求载荷与首次
// 成功发送不一致。Layer 227 为该错误定义的 code 是 500。
func randomIDDuplicateErr() error { return tgerr.New(500, "RANDOM_ID_DUPLICATE") }
// scheduleDateInvalidErr 表示当前阶段不支持定时消息。
func scheduleDateInvalidErr() error { return tgerr.New(400, "SCHEDULE_DATE_INVALID") }

2
internal/rpc/folders.go
View file

@ -53,7 +53,7 @@ func (r *Router) onFoldersEditPeerFolders(ctx context.Context, folderPeers []tg.
if r.deps.Updates != nil {
authKeyID, _ := AuthKeyIDFrom(ctx)
sessionID, _ := SessionIDFrom(ctx)
event, _, err = r.deps.Updates.RecordFolderPeers(ctx, authKeyID, userID, updates, sessionID)
event, _, err = r.deps.Updates.RecordFolderPeers(ctx, authKeyID, userID, updates, rawAuthKeyIDForOrigin(ctx), sessionID)
if err != nil {
return nil, internalErr()
}

5
internal/rpc/idle_backoff.go
View file

@ -57,7 +57,10 @@ func runIdleBackoffLoop(ctx context.Context, interval, maxIdleInterval time.Dura
case <-timer.C:
}
if dispatch(ctx) {
timer.Reset(backoff.ActiveDelay())
// 有积压时立即继续 drain；interval 只用于空闲轮询。旧逻辑每个非空
// batch 也固定等待 base，形成 batch/base 的人工吞吐上限。
_ = backoff.ActiveDelay()
timer.Reset(0)
continue
}
timer.Reset(backoff.IdleDelay())

31
internal/rpc/idle_backoff_test.go
View file

@ -1,6 +1,8 @@
package rpc
import (
"context"
"sync/atomic"
"testing"
"time"
)
@ -27,6 +29,35 @@ func TestIdleBackoffSequenceAndReset(t *testing.T) {
}
}
func TestIdleBackoffLoopDrainsActiveWorkImmediately(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
var calls atomic.Int32
done := make(chan struct{})
started := time.Now()
go func() {
defer close(done)
runIdleBackoffLoop(ctx, time.Second, time.Second, func(context.Context) bool {
if calls.Add(1) < 4 {
return true
}
cancel()
return false
})
}()
select {
case <-done:
case <-time.After(300 * time.Millisecond):
t.Fatal("active drain waited for idle interval")
}
if got := calls.Load(); got != 4 {
t.Fatalf("dispatch calls = %d, want 4 consecutive active drains", got)
}
if elapsed := time.Since(started); elapsed >= 300*time.Millisecond {
t.Fatalf("active drain elapsed = %v, want no 1s base delay", elapsed)
}
}
func TestIdleBackoffSanitizesMaxBelowBase(t *testing.T) {
backoff := newIdleBackoff(2*time.Second, time.Second)
if got := backoff.IdleDelay(); got != 2*time.Second {

84
internal/rpc/message_idempotency.go Normal file
View file

@ -0,0 +1,84 @@
package rpc
import (
"crypto/sha256"
"fmt"
"github.com/gotd/td/bin"
"github.com/gotd/td/tg"
)
// rpcRequestFingerprint 在任何自动实体补全、链接预览解析或上传媒体落库之前，
// 对客户端原始 TL request 取稳定指纹。这样 lost-response 重放不会因服务端派生
// photo/document id、pending webpage 状态等变化而被误判为另一条消息。
func rpcRequestFingerprint(req bin.Encoder) ([]byte, error) {
if req == nil {
return nil, fmt.Errorf("fingerprint rpc request: nil request")
}
var b bin.Buffer
if err := req.Encode(&b); err != nil {
return nil, fmt.Errorf("fingerprint rpc request: %w", err)
}
sum := sha256.Sum256(b.Raw())
return sum[:], nil
}
// sendMessageIdempotencyFingerprint fingerprints only the durable message intent.
// clear_draft/background/update_stickersets_order are one-shot client-side delivery
// hints: after a lost response DrKLO/TDesktop may legitimately retry the same
// random_id without them. They must not turn an exact send replay into
// RANDOM_ID_DUPLICATE.
func sendMessageIdempotencyFingerprint(req *tg.MessagesSendMessageRequest) ([]byte, error) {
if req == nil {
return nil, fmt.Errorf("fingerprint messages.sendMessage: nil request")
}
clone := *req
clone.Flags = 0
clone.ClearDraft = false
clone.Background = false
clone.UpdateStickersetsOrder = false
return rpcRequestFingerprint(&clone)
}
func sendMediaIdempotencyFingerprint(req *tg.MessagesSendMediaRequest) ([]byte, error) {
if req == nil {
return nil, fmt.Errorf("fingerprint messages.sendMedia: nil request")
}
clone := *req
clone.Flags = 0
clone.ClearDraft = false
clone.Background = false
clone.UpdateStickersetsOrder = false
return rpcRequestFingerprint(&clone)
}
// sendMultiMediaItemIdempotencyFingerprint deliberately reduces a batch to one
// InputSingleMedia. A retry containing only the failed subset therefore produces
// the same fingerprint for every surviving random_id as the original batch.
func sendMultiMediaItemIdempotencyFingerprint(req *tg.MessagesSendMultiMediaRequest, item tg.InputSingleMedia) ([]byte, error) {
if req == nil {
return nil, fmt.Errorf("fingerprint messages.sendMultiMedia item: nil request")
}
clone := *req
clone.Flags = 0
clone.ClearDraft = false
clone.Background = false
clone.UpdateStickersetsOrder = false
clone.MultiMedia = []tg.InputSingleMedia{item}
return rpcRequestFingerprint(&clone)
}
// forwardMessagesItemIdempotencyFingerprint makes the source message id and its
// paired random_id the unit of idempotency. Hashing the full ID/RandomID vectors
// incorrectly rejects a legal retry that contains only a failed subset.
func forwardMessagesItemIdempotencyFingerprint(req *tg.MessagesForwardMessagesRequest, messageID int, randomID int64) ([]byte, error) {
if req == nil {
return nil, fmt.Errorf("fingerprint messages.forwardMessages item: nil request")
}
clone := *req
clone.Flags = 0
clone.Background = false
clone.ID = []int{messageID}
clone.RandomID = []int64{randomID}
return rpcRequestFingerprint(&clone)
}

Some files were not shown because too many files have changed in this diff Show more