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

chore: refresh gramsrv public release

Browse source
This commit is contained in:
A 2026-06-30 14:37:43 +08:00
parent 75cebe8dbf
commit 70b6820474
1274 changed files with 378751 additions and 59919 deletions
Download patch file Download diff file Expand all files Collapse all files

428
internal/mtprotoedge/bot_callback_e2e_test.go Normal file
View file

@ -0,0 +1,428 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"fmt"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
botsapp "telesrv/internal/app/bots"
"telesrv/internal/app/contacts"
"telesrv/internal/app/dialogs"
"telesrv/internal/app/help"
"telesrv/internal/app/langpack"
messageapp "telesrv/internal/app/messages"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
// botCallbackEnv 搭建一套内存 server，供 P3 callback / startBot / markup e2e 复用。
type botCallbackEnv struct {
rsaKey *rsa.PrivateKey
addr *net.TCPAddr
bots *botsapp.Service
newCli func(*session.StorageMemory, telegram.UpdateHandler) *telegram.Client
newCliH func(*session.StorageMemory) *telegram.Client
}
func newBotCallbackEnv(t *testing.T, ctx context.Context) *botCallbackEnv {
t.Helper()
const dc = 2
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tcpAddr := ln.Addr().(*net.TCPAddr)
userStore := memory.NewUserStore()
authzStore := memory.NewAuthorizationStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
langPackStore := memory.NewLangPackStore()
dialogStore := memory.NewDialogStore()
messageStore := memory.NewMessageStore(dialogStore)
botStore := memory.NewBotStore(userStore)
botsService := botsapp.NewService(userStore, botStore, messageStore,
botsapp.WithLogger(zaptest.NewLogger(t).Named("bots")))
activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions"))
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore,
memory.NewTempAuthKeyBindingStore(), "12345", auth.WithBotLogin(botStore)),
Account: account.NewService(memory.NewPasswordStore()),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(dialogStore),
Messages: messageapp.NewService(messageStore, dialogStore, messageapp.WithBotResponder(botsService)),
Bots: botsService,
LangPack: langpack.NewService(langPackStore),
Sessions: activeSessions,
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
botsService.SetRouterHooks(router)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
go func() { _ = srv.Serve(ctx, ln) }()
newCli := func(storage *session.StorageMemory, handler telegram.UpdateHandler) *telegram.Client {
if handler == nil {
handler = telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil })
}
return telegram.NewClient(1, "hash", telegram.Options{
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: storage,
UpdateHandler: handler,
})
}
return &botCallbackEnv{
rsaKey: rsaKey, addr: tcpAddr, bots: botsService,
newCli: newCli,
newCliH: func(s *session.StorageMemory) *telegram.Client { return newCli(s, nil) },
}
}
func registerOwner(t *testing.T, ctx context.Context, env *botCallbackEnv, phone, name string) (tg.User, *session.StorageMemory) {
t.Helper()
var owner tg.User
storage := &session.StorageMemory{}
client := env.newCliH(storage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
hash := sent.(*tg.AuthSentCode).PhoneCodeHash
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: hash, PhoneCode: "12345"}); err != nil {
return err
}
res, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{PhoneNumber: phone, PhoneCodeHash: hash, FirstName: name})
if err != nil {
return err
}
owner = *res.(*tg.AuthAuthorization).User.(*tg.User)
return nil
}); err != nil {
t.Fatalf("owner signUp: %v", err)
}
return owner, storage
}
func historyMessageList(history tg.MessagesMessagesClass) []tg.MessageClass {
switch v := history.(type) {
case *tg.MessagesMessages:
return v.Messages
case *tg.MessagesMessagesSlice:
return v.Messages
case *tg.MessagesChannelMessages:
return v.Messages
default:
return nil
}
}
func updatesFromClass(u tg.UpdatesClass) []tg.UpdateClass {
switch v := u.(type) {
case *tg.Updates:
return v.Updates
case *tg.UpdatesCombined:
return v.Updates
case *tg.UpdateShort:
return []tg.UpdateClass{v.Update}
default:
return nil
}
}
// TestBotInlineKeyboardCallbackFlow 验证 P3 callback 全链路（真实 gotd 双客户端并发）：
// bot 发带 inline callback markup 的消息 → owner getHistory 见 markup → owner
// getBotCallbackAnswer 挂起 → bot 经 updateBotCallbackQuery 收到 query → setBotCallbackAnswer
// → owner 收到 answer；并验证 bot 不应答 → BOT_RESPONSE_TIMEOUT；data 字节级保真。
func TestBotInlineKeyboardCallbackFlow(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
defer cancel()
env := newBotCallbackEnv(t, ctx)
owner, ownerStorage := registerOwner(t, ctx, env, "+15550004001", "Owner")
botUser, botToken, err := env.bots.CreateBot(context.Background(), owner.ID, "CB Bot", "cb_test_bot")
if err != nil {
t.Fatalf("create bot: %v", err)
}
// 含 0x00 与高位字节的 callback data：验证字节级 round-trip（I2）。
callbackData := []byte{0x00, 0x01, 0xFF, 0x80, 'a', 'b'}
const noAnswerText = "ignore"
botOnline := make(chan struct{})
botStop := make(chan struct{})
botErr := make(chan error, 1)
// bot 在后台保持在线，update handler 把收到的 callback query 转交应答 goroutine。
cbCh := make(chan *tg.UpdateBotCallbackQuery, 8)
botHandler := telegram.UpdateHandlerFunc(func(_ context.Context, u tg.UpdatesClass) error {
for _, upd := range updatesFromClass(u) {
if q, ok := upd.(*tg.UpdateBotCallbackQuery); ok {
select {
case cbCh <- q:
default:
}
}
}
return nil
})
go func() {
botClient := env.newCli(&session.StorageMemory{}, botHandler)
botErr <- botClient.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(botClient)
if _, err := raw.AuthImportBotAuthorization(ctx, &tg.AuthImportBotAuthorizationRequest{APIID: 1, APIHash: "hash", BotAuthToken: botToken}); err != nil {
return err
}
// 裸 RPC 置 receivesUpdates，使 push 可达（热恢复同步修复后语义）。
if _, err := raw.UpdatesGetState(ctx); err != nil {
return err
}
// 解析 owner 资料拿到 bot 视角的 access_hash（access_hash=0 跳过校验，
// 与既有 bot e2e 同款），再发带 inline callback+url keyboard 的消息。
got, err := raw.UsersGetUsers(ctx, []tg.InputUserClass{&tg.InputUser{UserID: owner.ID}})
if err != nil {
return err
}
if len(got) != 1 {
// 该回调由 go func() 启动并经 botErr channel 上报；在此 goroutine 直接
// t.Fatalf 会在错误的 goroutine 上 Goexit（go1.26 vet testinggoroutine 亦报），
// 故返回 error 让测试主 goroutine 经 botErr 失败。
return fmt.Errorf("bot getUsers(owner) = %d, want 1", len(got))
}
ownerSeen := got[0].(*tg.User)
markup := &tg.ReplyInlineMarkup{Rows: []tg.KeyboardButtonRow{{Buttons: []tg.KeyboardButtonClass{
&tg.KeyboardButtonCallback{Text: "Press", Data: callbackData},
&tg.KeyboardButtonURL{Text: "Site", URL: "https://example.com/x"},
}}}}
req := &tg.MessagesSendMessageRequest{
Peer: &tg.InputPeerUser{UserID: ownerSeen.ID, AccessHash: ownerSeen.AccessHash},
Message: "tap a button",
RandomID: 778801,
}
req.SetReplyMarkup(markup)
if _, err := raw.MessagesSendMessage(ctx, req); err != nil {
return err
}
// 应答 goroutine：除 noAnswerText 外，对每个 query 回 setBotCallbackAnswer。
go func() {
for {
select {
case q := <-cbCh:
data, _ := q.GetData()
if string(data) == noAnswerText {
continue // 制造超时分支
}
_, _ = raw.MessagesSetBotCallbackAnswer(ctx, &tg.MessagesSetBotCallbackAnswerRequest{
QueryID: q.QueryID,
Alert: true,
Message: "ok:" + string(data),
})
case <-botStop:
return
}
}
}()
close(botOnline)
select {
case <-botStop:
case <-ctx.Done():
}
return nil
})
}()
select {
case <-botOnline:
case err := <-botErr:
t.Fatalf("bot client exited early: %v", err)
case <-ctx.Done():
t.Fatal("bot did not come online")
}
ownerClient := env.newCliH(ownerStorage)
if err := ownerClient.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(ownerClient)
// owner 解析 bot 拿到自己视角的 access_hash（公开 username 冷启动解析）。
resolved, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: "cb_test_bot"})
if err != nil {
return err
}
seenBot := resolved.Users[0].(*tg.User)
botPeer := &tg.InputPeerUser{UserID: seenBot.ID, AccessHash: seenBot.AccessHash}
_ = botUser
// 轮询 getHistory 直到 bot 的 markup 消息到达。
var msgID int
deadline := time.Now().Add(15 * time.Second)
for {
h, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{Peer: botPeer, Limit: 5})
if err != nil {
return err
}
for _, m := range historyMessageList(h) {
msg, ok := m.(*tg.Message)
if !ok {
continue
}
rm, ok := msg.GetReplyMarkup()
if !ok {
continue
}
inline, ok := rm.(*tg.ReplyInlineMarkup)
if !ok || len(inline.Rows) != 1 || len(inline.Rows[0].Buttons) != 2 {
t.Fatalf("unexpected markup shape: %#v", rm)
}
cbBtn, ok := inline.Rows[0].Buttons[0].(*tg.KeyboardButtonCallback)
if !ok {
t.Fatalf("first button not callback: %#v", inline.Rows[0].Buttons[0])
}
if string(cbBtn.Data) != string(callbackData) {
t.Fatalf("callback data round-trip mismatch: got %v want %v", cbBtn.Data, callbackData)
}
if _, ok := inline.Rows[0].Buttons[1].(*tg.KeyboardButtonURL); !ok {
t.Fatalf("second button not url: %#v", inline.Rows[0].Buttons[1])
}
msgID = msg.ID
}
if msgID != 0 {
break
}
if time.Now().After(deadline) {
t.Fatal("markup message did not arrive in owner history")
}
time.Sleep(50 * time.Millisecond)
}
// 按下 callback 按钮：getBotCallbackAnswer 挂起直到 bot 应答。
ans, err := raw.MessagesGetBotCallbackAnswer(ctx, &tg.MessagesGetBotCallbackAnswerRequest{
Peer: botPeer,
MsgID: msgID,
Data: callbackData,
})
if err != nil {
return err
}
if !ans.Alert || ans.Message != "ok:"+string(callbackData) {
t.Fatalf("callback answer = alert:%v msg:%q, want alert:true msg:%q", ans.Alert, ans.Message, "ok:"+string(callbackData))
}
// 超时分支：bot 收到但不应答（noAnswerText）→ BOT_RESPONSE_TIMEOUT。
// 用短 ctx 触发客户端侧超时，避免等满 25s 服务端窗口。
toCtx, toCancel := context.WithTimeout(ctx, 3*time.Second)
defer toCancel()
_, err = raw.MessagesGetBotCallbackAnswer(toCtx, &tg.MessagesGetBotCallbackAnswerRequest{
Peer: botPeer,
MsgID: msgID,
Data: []byte(noAnswerText),
})
if err == nil {
t.Fatal("expected timeout error for unanswered callback")
}
return nil
}); err != nil {
t.Fatalf("owner callback flow: %v", err)
}
close(botStop)
select {
case <-botErr:
case <-time.After(5 * time.Second):
}
}
// TestBotStartBotFlow 验证 messages.startBot 产生可见 "/start <param>" 消息（I7）。
func TestBotStartBotFlow(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
env := newBotCallbackEnv(t, ctx)
owner, ownerStorage := registerOwner(t, ctx, env, "+15550004002", "Owner")
_, botToken, err := env.bots.CreateBot(context.Background(), owner.ID, "Start Bot", "start_test_bot")
if err != nil {
t.Fatalf("create bot: %v", err)
}
// owner 调 startBot(bot, payload)。先经公开 username 解析 bot 拿 access_hash。
ownerClient := env.newCliH(ownerStorage)
if err := ownerClient.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(ownerClient)
resolved, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: "start_test_bot"})
if err != nil {
return err
}
seenBot := resolved.Users[0].(*tg.User)
_, err = raw.MessagesStartBot(ctx, &tg.MessagesStartBotRequest{
Bot: &tg.InputUser{UserID: seenBot.ID, AccessHash: seenBot.AccessHash},
Peer: &tg.InputPeerUser{UserID: seenBot.ID, AccessHash: seenBot.AccessHash},
RandomID: 994401,
StartParam: "ref123",
})
return err
}); err != nil {
t.Fatalf("owner startBot: %v", err)
}
// bot 登录后 getHistory(peer=owner) 应见到 "/start ref123"。
botClient := env.newCliH(&session.StorageMemory{})
if err := botClient.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(botClient)
if _, err := raw.AuthImportBotAuthorization(ctx, &tg.AuthImportBotAuthorizationRequest{APIID: 1, APIHash: "hash", BotAuthToken: botToken}); err != nil {
return err
}
got, err := raw.UsersGetUsers(ctx, []tg.InputUserClass{&tg.InputUser{UserID: owner.ID}})
if err != nil {
return err
}
ownerSeen := got[0].(*tg.User)
h, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerUser{UserID: ownerSeen.ID, AccessHash: ownerSeen.AccessHash},
Limit: 5,
})
if err != nil {
return err
}
found := false
var bodies []string
for _, m := range historyMessageList(h) {
if msg, ok := m.(*tg.Message); ok {
bodies = append(bodies, msg.Message)
if msg.Message == "/start ref123" {
found = true
}
}
}
if !found {
t.Fatalf("bot did not receive '/start ref123' message; history bodies=%v", bodies)
}
return nil
}); err != nil {
t.Fatalf("bot startBot receive: %v", err)
}
}

835
internal/mtprotoedge/bot_e2e_test.go Normal file
View file

@ -0,0 +1,835 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"net"
"regexp"
"strings"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"github.com/gotd/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
botsapp "telesrv/internal/app/bots"
"telesrv/internal/app/contacts"
"telesrv/internal/app/dialogs"
"telesrv/internal/app/help"
"telesrv/internal/app/langpack"
messageapp "telesrv/internal/app/messages"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/domain"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
var botTokenRe = regexp.MustCompile(`(\d+):([A-Za-z0-9_-]{35})`)
// TestBotFatherCreateAndBotLoginFlow 是 bot 主链路端到端验证：
// TestBotManagementRPCFlow 验证 P2 bots.* 管理 RPC 端到端：
// bot 自己调 setBotCommands/setBotInfo/setBotMenuButton 后 getFullUser 反映新值且
// bot_info_version 单调递增；owner 视角 getFullUser(bot) 带 bot_can_edit 且 owner 可
// 经 bot:InputUser 代改 bot name。
func TestBotManagementRPCFlow(t *testing.T) {
const (
dc = 2
code = "12345"
)
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tcpAddr := ln.Addr().(*net.TCPAddr)
userStore := memory.NewUserStore()
authzStore := memory.NewAuthorizationStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
langPackStore := memory.NewLangPackStore()
dialogStore := memory.NewDialogStore()
messageStore := memory.NewMessageStore(dialogStore)
botStore := memory.NewBotStore(userStore)
botsService := botsapp.NewService(userStore, botStore, messageStore,
botsapp.WithLogger(zaptest.NewLogger(t).Named("bots")))
activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions"))
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore,
memory.NewTempAuthKeyBindingStore(), code, auth.WithBotLogin(botStore)),
Account: account.NewService(memory.NewPasswordStore()),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(dialogStore),
Messages: messageapp.NewService(messageStore, dialogStore, messageapp.WithBotResponder(botsService)),
Bots: botsService,
LangPack: langpack.NewService(langPackStore),
Sessions: activeSessions,
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
botsService.SetRouterHooks(router)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
newClient := func(storage *session.StorageMemory) *telegram.Client {
opts := telegram.Options{
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: storage,
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
return telegram.NewClient(1, "hash", opts)
}
// owner 注册。
ownerStorage := &session.StorageMemory{}
var owner tg.User
{
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: "+15550003001", APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
hash := sent.(*tg.AuthSentCode).PhoneCodeHash
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: "+15550003001", PhoneCodeHash: hash, PhoneCode: code}); err != nil {
return err
}
res, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{PhoneNumber: "+15550003001", PhoneCodeHash: hash, FirstName: "Owner"})
if err != nil {
return err
}
owner = *res.(*tg.AuthAuthorization).User.(*tg.User)
return nil
}); err != nil {
t.Fatalf("owner signUp: %v", err)
}
}
// 直接经 service 建 bot（绕过 BotFather 对话，聚焦管理 RPC）。
botUser, botToken, err := botsService.CreateBot(context.Background(), owner.ID, "Mgmt Bot", "mgmt_test_bot")
if err != nil {
t.Fatalf("create bot: %v", err)
}
getFullSelfBotInfo := func(raw *tg.Client) (tg.BotInfo, *tg.User) {
t.Helper()
full, err := raw.UsersGetFullUser(ctx, &tg.InputUserSelf{})
if err != nil {
t.Fatalf("getFullUser self: %v", err)
}
bi, ok := full.FullUser.GetBotInfo()
if !ok {
t.Fatalf("self userFull lacks bot_info")
}
return bi, full.Users[0].(*tg.User)
}
// bot 登录并调管理 RPC。
botStorage := &session.StorageMemory{}
var verAfterCommands int
{
client := newClient(botStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
if _, err := raw.AuthImportBotAuthorization(ctx, &tg.AuthImportBotAuthorizationRequest{APIID: 1, APIHash: "hash", BotAuthToken: botToken}); err != nil {
return err
}
_, self0 := getFullSelfBotInfo(raw)
v0, _ := self0.GetBotInfoVersion()
// setBotCommands（default scope）。
ok, err := raw.BotsSetBotCommands(ctx, &tg.BotsSetBotCommandsRequest{
Scope: &tg.BotCommandScopeDefault{},
LangCode: "",
Commands: []tg.BotCommand{{Command: "start", Description: "begin"}, {Command: "help", Description: "show help"}},
})
if err != nil || !ok {
t.Fatalf("setBotCommands = %v,%v", ok, err)
}
bi, self1 := getFullSelfBotInfo(raw)
cmds, _ := bi.GetCommands()
if len(cmds) != 2 || cmds[0].Command != "start" {
t.Fatalf("bot_info commands = %+v, want [start,help]", cmds)
}
v1, _ := self1.GetBotInfoVersion()
if v1 <= v0 {
t.Fatalf("bot_info_version not bumped after setBotCommands: %d -> %d", v0, v1)
}
verAfterCommands = v1
// getBotCommands 回读。
got, err := raw.BotsGetBotCommands(ctx, &tg.BotsGetBotCommandsRequest{Scope: &tg.BotCommandScopeDefault{}})
if err != nil || len(got) != 2 {
t.Fatalf("getBotCommands = %+v, %v", got, err)
}
// setBotMenuButton(webview)。
if ok, err := raw.BotsSetBotMenuButton(ctx, &tg.BotsSetBotMenuButtonRequest{
UserID: &tg.InputUserSelf{},
Button: &tg.BotMenuButton{Text: "Open", URL: "https://example.com/app"},
}); err != nil || !ok {
t.Fatalf("setBotMenuButton = %v,%v", ok, err)
}
bi2, _ := getFullSelfBotInfo(raw)
mb, _ := bi2.GetMenuButton()
if btn, ok := mb.(*tg.BotMenuButton); !ok || btn.URL != "https://example.com/app" {
t.Fatalf("menu button = %#v, want webview", mb)
}
// setBotInfo(description) by bot self（不带 bot 参数）；Description 是 flag.1，须 SetDescription 置位。
infoReq := &tg.BotsSetBotInfoRequest{LangCode: ""}
infoReq.SetDescription("what I do")
if ok, err := raw.BotsSetBotInfo(ctx, infoReq); err != nil || !ok {
t.Fatalf("setBotInfo(description) = %v,%v", ok, err)
}
info, err := raw.BotsGetBotInfo(ctx, &tg.BotsGetBotInfoRequest{LangCode: ""})
if err != nil || info.Description != "what I do" {
t.Fatalf("getBotInfo = %#v, %v", info, err)
}
// bot self 不应带 bot_can_edit（owner 视角才有）。
_, selfU := getFullSelfBotInfo(raw)
if selfU.GetBotCanEdit() {
t.Fatalf("bot self carries bot_can_edit, want only owner view")
}
return nil
}); err != nil {
t.Fatalf("bot management flow: %v", err)
}
}
// owner 视角：getFullUser(bot) 带 bot_can_edit + commands；owner 代改 name。
{
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
resolved, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: "mgmt_test_bot"})
if err != nil {
return err
}
seen := resolved.Users[0].(*tg.User)
botInput := &tg.InputUser{UserID: seen.ID, AccessHash: seen.AccessHash}
full, err := raw.UsersGetFullUser(ctx, botInput)
if err != nil {
return err
}
ownerSeesBot := full.Users[0].(*tg.User)
if !ownerSeesBot.GetBotCanEdit() {
t.Fatalf("owner view bot_can_edit = false, want true")
}
if bi, ok := full.FullUser.GetBotInfo(); ok {
if cmds, _ := bi.GetCommands(); len(cmds) != 2 {
t.Fatalf("owner view bot commands = %d, want 2", len(cmds))
}
}
// owner 代改 bot name（带 bot:InputUser）。
req := &tg.BotsSetBotInfoRequest{LangCode: ""}
req.SetBot(botInput)
req.SetName("Owner Renamed")
if ok, err := raw.BotsSetBotInfo(ctx, req); err != nil || !ok {
t.Fatalf("owner setBotInfo(name) = %v,%v", ok, err)
}
got, err := raw.UsersGetUsers(ctx, []tg.InputUserClass{botInput})
if err != nil || len(got) != 1 {
t.Fatalf("getUsers(bot) = %+v, %v", got, err)
}
if u := got[0].(*tg.User); u.FirstName != "Owner Renamed" {
t.Fatalf("bot first_name = %q, want 'Owner Renamed'", u.FirstName)
}
return nil
}); err != nil {
t.Fatalf("owner management flow: %v", err)
}
}
_ = botUser
_ = verAfterCommands
cancel()
if err := <-serveErr; err != nil {
t.Errorf("serve: %v", err)
}
}
// 用户注册 → resolveUsername(BotFather) → /newbot 对话拿 token →
// 外部客户端 auth.importBotAuthorization 登录为 bot →
// 用户与 bot 互发消息 → 错误 token 拿 ACCESS_TOKEN_INVALID。
func TestBotFatherCreateAndBotLoginFlow(t *testing.T) {
const (
dc = 2
code = "12345"
)
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tcpAddr := ln.Addr().(*net.TCPAddr)
userStore := memory.NewUserStore()
authzStore := memory.NewAuthorizationStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
langPackStore := memory.NewLangPackStore()
dialogStore := memory.NewDialogStore()
messageStore := memory.NewMessageStore(dialogStore)
botStore := memory.NewBotStore(userStore)
botsService := botsapp.NewService(userStore, botStore, messageStore,
botsapp.WithLogger(zaptest.NewLogger(t).Named("bots")))
activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions"))
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore,
memory.NewTempAuthKeyBindingStore(), code, auth.WithBotLogin(botStore)),
Account: account.NewService(memory.NewPasswordStore()),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(dialogStore),
Messages: messageapp.NewService(messageStore, dialogStore, messageapp.WithBotResponder(botsService)),
Bots: botsService,
LangPack: langpack.NewService(langPackStore),
Sessions: activeSessions,
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
botsService.SetRouterHooks(router)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
newClient := func(storage *session.StorageMemory) *telegram.Client {
opts := telegram.Options{
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: storage,
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
return telegram.NewClient(1, "hash", opts)
}
messagesOf := func(history tg.MessagesMessagesClass) []tg.MessageClass {
t.Helper()
switch v := history.(type) {
case *tg.MessagesMessages:
return v.Messages
case *tg.MessagesMessagesSlice:
return v.Messages
default:
t.Fatalf("history = %T %+v, want messages", history, history)
return nil
}
}
// 1) owner 注册。
ownerStorage := &session.StorageMemory{}
var owner tg.User
{
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{
PhoneNumber: "+15550002001", APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{},
})
if err != nil {
return err
}
hash := sent.(*tg.AuthSentCode).PhoneCodeHash
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{
PhoneNumber: "+15550002001", PhoneCodeHash: hash, PhoneCode: code,
}); err != nil {
return err
}
res, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{
PhoneNumber: "+15550002001", PhoneCodeHash: hash, FirstName: "Owner",
})
if err != nil {
return err
}
owner = *res.(*tg.AuthAuthorization).User.(*tg.User)
return nil
}); err != nil {
t.Fatalf("owner signUp: %v", err)
}
}
// 2) owner 与 BotFather 对话创建 bot，拿 token。
var (
botToken string
botUsername = "owner_e2e_bot"
botUserID int64
)
{
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
resolved, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: "BotFather"})
if err != nil {
return err
}
if len(resolved.Users) != 1 {
t.Fatalf("resolve BotFather users = %d, want 1", len(resolved.Users))
}
botFather := resolved.Users[0].(*tg.User)
if botFather.ID != domain.BotFatherUserID || !botFather.Bot {
t.Fatalf("resolved BotFather = %+v, want bot flag with id %d", botFather, domain.BotFatherUserID)
}
if v, ok := botFather.GetBotInfoVersion(); !ok || v < 1 {
t.Fatalf("BotFather bot_info_version = %d,%v, want >=1", v, ok)
}
if _, hasStatus := botFather.GetStatus(); hasStatus {
t.Fatalf("BotFather carries status %+v, want none for bots", botFather.Status)
}
peer := &tg.InputPeerUser{UserID: botFather.ID, AccessHash: botFather.AccessHash}
// userFull.bot_info 必须存在且 user_id 匹配（TDesktop P0 兼容点）。
full, err := raw.UsersGetFullUser(ctx, &tg.InputUser{UserID: botFather.ID, AccessHash: botFather.AccessHash})
if err != nil {
return err
}
botInfo, ok := full.FullUser.GetBotInfo()
if !ok {
t.Fatalf("BotFather userFull lacks bot_info")
}
if id, ok := botInfo.GetUserID(); !ok || id != domain.BotFatherUserID {
t.Fatalf("bot_info.user_id = %d,%v, want %d", id, ok, domain.BotFatherUserID)
}
if cmds, ok := botInfo.GetCommands(); !ok || len(cmds) == 0 {
t.Fatalf("BotFather bot_info commands empty, want seeded commands")
}
randomID := int64(31001)
// BotFather 回复异步到达（go routine），轮询历史顶部直到出现新的 incoming 回复。
sendAndReply := func(text string) string {
t.Helper()
randomID++
beforeTop := 0
if h, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{Peer: peer, Limit: 1}); err == nil {
if m := messagesOf(h); len(m) > 0 {
if msg, ok := m[0].(*tg.Message); ok {
beforeTop = msg.ID
}
}
}
if _, err := raw.MessagesSendMessage(ctx, &tg.MessagesSendMessageRequest{
Peer: peer, Message: text, RandomID: randomID,
}); err != nil {
t.Fatalf("send %q: %v", text, err)
}
deadline := time.Now().Add(10 * time.Second)
for {
history, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{Peer: peer, Limit: 5})
if err != nil {
t.Fatalf("history after %q: %v", text, err)
}
msgs := messagesOf(history)
if len(msgs) > 0 {
if top, ok := msgs[0].(*tg.Message); ok && !top.Out && top.ID > beforeTop {
return top.Message
}
}
if time.Now().After(deadline) {
t.Fatalf("timed out waiting for BotFather reply to %q", text)
}
time.Sleep(50 * time.Millisecond)
}
}
if reply := sendAndReply("/newbot"); !strings.Contains(reply, "choose a name") {
t.Fatalf("/newbot reply = %q", reply)
}
if reply := sendAndReply("Owner E2E Bot"); !strings.Contains(reply, "username") {
t.Fatalf("name reply = %q", reply)
}
reply := sendAndReply(botUsername)
match := botTokenRe.FindStringSubmatch(reply)
if match == nil {
t.Fatalf("done reply = %q, want token", reply)
}
botToken = match[0]
// 新 bot 可被 resolve，且带 bot flags。
resolvedBot, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: botUsername})
if err != nil {
return err
}
created := resolvedBot.Users[0].(*tg.User)
if !created.Bot {
t.Fatalf("created bot user lacks bot flag: %+v", created)
}
if v, ok := created.GetBotInfoVersion(); !ok || v < 1 {
t.Fatalf("created bot bot_info_version = %d,%v, want >=1", v, ok)
}
botUserID = created.ID
return nil
}); err != nil {
t.Fatalf("newbot flow: %v", err)
}
}
// 3) 错误 token 必须拿 ACCESS_TOKEN_INVALID。
{
client := newClient(&session.StorageMemory{})
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
_, err := raw.AuthImportBotAuthorization(ctx, &tg.AuthImportBotAuthorizationRequest{
APIID: 1, APIHash: "hash", BotAuthToken: "12345:notarealtokennotarealtokennotareal_",
})
if !tgerr.Is(err, "ACCESS_TOKEN_INVALID") {
t.Fatalf("bad token err = %v, want ACCESS_TOKEN_INVALID", err)
}
return nil
}); err != nil {
t.Fatalf("bad token flow: %v", err)
}
}
// 4) bot 客户端凭 token 登录；self 必须是 bot；userFull.bot_info 自洽。
botStorage := &session.StorageMemory{}
{
client := newClient(botStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
res, err := raw.AuthImportBotAuthorization(ctx, &tg.AuthImportBotAuthorizationRequest{
APIID: 1, APIHash: "hash", BotAuthToken: botToken,
})
if err != nil {
return err
}
authz, ok := res.(*tg.AuthAuthorization)
if !ok {
t.Fatalf("importBotAuthorization = %T, want *tg.AuthAuthorization", res)
}
self := authz.User.(*tg.User)
if self.ID != botUserID || !self.Bot || !self.Self {
t.Fatalf("bot self = %+v, want self bot id %d", self, botUserID)
}
if v, ok := self.GetBotInfoVersion(); !ok || v < 1 {
t.Fatalf("bot self bot_info_version = %d,%v, want >=1", v, ok)
}
// 登录态生效：getUsers(self) 与 getState 正常。
got, err := raw.UsersGetUsers(ctx, []tg.InputUserClass{&tg.InputUserSelf{}})
if err != nil {
return err
}
if len(got) != 1 || got[0].(*tg.User).ID != botUserID {
t.Fatalf("bot getUsers(self) = %+v, want id %d", got, botUserID)
}
if _, err := raw.UpdatesGetState(ctx); err != nil {
return err
}
full, err := raw.UsersGetFullUser(ctx, &tg.InputUserSelf{})
if err != nil {
return err
}
botInfo, ok := full.FullUser.GetBotInfo()
if !ok {
t.Fatalf("bot self userFull lacks bot_info")
}
if id, ok := botInfo.GetUserID(); !ok || id != botUserID {
t.Fatalf("bot self bot_info.user_id = %d,%v, want %d", id, ok, botUserID)
}
return nil
}); err != nil {
t.Fatalf("bot login flow: %v", err)
}
}
// 5) owner 给 bot 发消息。
{
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
resolvedBot, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: botUsername})
if err != nil {
return err
}
created := resolvedBot.Users[0].(*tg.User)
if _, err := raw.MessagesSendMessage(ctx, &tg.MessagesSendMessageRequest{
Peer: &tg.InputPeerUser{UserID: created.ID, AccessHash: created.AccessHash},
Message: "hi bot",
RandomID: 41001,
}); err != nil {
return err
}
return nil
}); err != nil {
t.Fatalf("owner -> bot send: %v", err)
}
}
// 6) bot 读到消息并回复。
{
client := newClient(botStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
// access_hash=0 跳过校验（与现有 users.getUsers 语义一致），bot 据此拿 owner 资料。
got, err := raw.UsersGetUsers(ctx, []tg.InputUserClass{&tg.InputUser{UserID: owner.ID}})
if err != nil {
return err
}
if len(got) != 1 {
t.Fatalf("bot getUsers(owner) = %d users, want 1", len(got))
}
ownerSeen := got[0].(*tg.User)
peer := &tg.InputPeerUser{UserID: ownerSeen.ID, AccessHash: ownerSeen.AccessHash}
history, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{Peer: peer, Limit: 5})
if err != nil {
return err
}
msgs := messagesOf(history)
if len(msgs) == 0 {
t.Fatalf("bot history empty, want incoming message")
}
top, ok := msgs[0].(*tg.Message)
if !ok || top.Message != "hi bot" || top.Out {
t.Fatalf("bot latest = %#v, want incoming 'hi bot'", msgs[0])
}
if _, err := raw.MessagesSendMessage(ctx, &tg.MessagesSendMessageRequest{
Peer: peer,
Message: "hello human",
RandomID: 51001,
}); err != nil {
return err
}
// P2: bot 自管理（bot-only RPC）：setBotCommands / setBotMenuButton。
ok2, err := raw.BotsSetBotCommands(ctx, &tg.BotsSetBotCommandsRequest{
Scope: &tg.BotCommandScopeDefault{},
LangCode: "",
Commands: []tg.BotCommand{
{Command: "start", Description: "start the bot"},
{Command: "ping", Description: "check liveness"},
},
})
if err != nil || !ok2 {
t.Fatalf("setBotCommands = %v err=%v, want true", ok2, err)
}
cmds, err := raw.BotsGetBotCommands(ctx, &tg.BotsGetBotCommandsRequest{
Scope: &tg.BotCommandScopeDefault{}, LangCode: "",
})
if err != nil || len(cmds) != 2 || cmds[0].Command != "start" {
t.Fatalf("getBotCommands = %+v err=%v, want [start ping]", cmds, err)
}
okBtn, err := raw.BotsSetBotMenuButton(ctx, &tg.BotsSetBotMenuButtonRequest{
UserID: &tg.InputUserEmpty{},
Button: &tg.BotMenuButton{Text: "Open", URL: "https://example.test/app"},
})
if err != nil || !okBtn {
t.Fatalf("setBotMenuButton = %v err=%v, want true", okBtn, err)
}
gotBtn, err := raw.BotsGetBotMenuButton(ctx, &tg.InputUserEmpty{})
if err != nil {
return err
}
if web, ok := gotBtn.(*tg.BotMenuButton); !ok || web.URL != "https://example.test/app" {
t.Fatalf("getBotMenuButton = %#v, want webview button", gotBtn)
}
return nil
}); err != nil {
t.Fatalf("bot read/reply: %v", err)
}
}
// 7) owner 看到 bot 回复。
{
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
resolvedBot, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: botUsername})
if err != nil {
return err
}
created := resolvedBot.Users[0].(*tg.User)
history, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerUser{UserID: created.ID, AccessHash: created.AccessHash},
Limit: 5,
})
if err != nil {
return err
}
msgs := messagesOf(history)
if len(msgs) == 0 {
t.Fatalf("owner history with bot empty")
}
top, ok := msgs[0].(*tg.Message)
if !ok || top.Message != "hello human" || top.Out {
t.Fatalf("owner latest = %#v, want incoming 'hello human'", msgs[0])
}
// P2: owner 视角的 bot 元数据闭环。
// (a) bot 设置命令+菜单按钮后 bot_info_version 已 bump（>1）。
if v, ok := created.GetBotInfoVersion(); !ok || v <= 1 {
t.Fatalf("bot_info_version = %d,%v after metadata changes, want > 1", v, ok)
}
// (b) getFullUser：bot_info 带新命令与 webview 菜单按钮；owner 看到 bot_can_edit。
full, err := raw.UsersGetFullUser(ctx, &tg.InputUser{UserID: created.ID, AccessHash: created.AccessHash})
if err != nil {
return err
}
botInfo, ok := full.FullUser.GetBotInfo()
if !ok {
t.Fatal("owner getFullUser(bot) lacks bot_info")
}
if cmds, ok := botInfo.GetCommands(); !ok || len(cmds) != 2 || cmds[1].Command != "ping" {
t.Fatalf("bot_info commands = %+v,%v, want [start ping]", cmds, ok)
}
if btn, ok := botInfo.GetMenuButton(); !ok {
t.Fatal("bot_info lacks menu_button")
} else if web, isWeb := btn.(*tg.BotMenuButton); !isWeb || web.URL != "https://example.test/app" {
t.Fatalf("menu_button = %#v, want webview", btn)
}
fullUser := full.Users[0].(*tg.User)
if !fullUser.BotCanEdit {
t.Fatalf("owner getFullUser(bot) user lacks bot_can_edit: %+v", fullUser)
}
// (c) owner 经 bot 参数代设 setBotInfo（about+description）→ getFullUser 反映。
setReq := &tg.BotsSetBotInfoRequest{LangCode: ""}
setReq.SetBot(&tg.InputUser{UserID: created.ID, AccessHash: created.AccessHash})
setReq.SetAbout("e2e about")
setReq.SetDescription("e2e description")
if okSet, err := raw.BotsSetBotInfo(ctx, setReq); err != nil || !okSet {
t.Fatalf("owner setBotInfo = %v err=%v, want true", okSet, err)
}
full2, err := raw.UsersGetFullUser(ctx, &tg.InputUser{UserID: created.ID, AccessHash: created.AccessHash})
if err != nil {
return err
}
if full2.FullUser.About != "e2e about" {
t.Fatalf("about after setBotInfo = %q, want 'e2e about'", full2.FullUser.About)
}
if bi2, ok := full2.FullUser.GetBotInfo(); !ok {
t.Fatal("getFullUser after setBotInfo lacks bot_info")
} else if desc, _ := bi2.GetDescription(); desc != "e2e description" {
t.Fatalf("description = %q, want 'e2e description'", desc)
}
// (d) owner 对非自己的 bot（BotFather）代设 → BOT_INVALID。
badReq := &tg.BotsSetBotInfoRequest{LangCode: ""}
badReq.SetBot(&tg.InputUser{UserID: domain.BotFatherUserID, AccessHash: domain.BotFatherAccessHash})
badReq.SetAbout("nope")
if _, err := raw.BotsSetBotInfo(ctx, badReq); !tgerr.Is(err, "BOT_INVALID") {
t.Fatalf("setBotInfo on BotFather err = %v, want BOT_INVALID", err)
}
// (e) 非 bot 用户调 bot-only RPC → USER_BOT_REQUIRED。
if _, err := raw.BotsGetBotCommands(ctx, &tg.BotsGetBotCommandsRequest{
Scope: &tg.BotCommandScopeDefault{}, LangCode: "",
}); !tgerr.Is(err, "USER_BOT_REQUIRED") {
t.Fatalf("owner getBotCommands err = %v, want USER_BOT_REQUIRED", err)
}
return nil
}); err != nil {
t.Fatalf("owner read bot reply: %v", err)
}
}
// 8) owner 走 BotFather /revoke：旧 token 失效 + bot 全部 authorization 被撤销
//（已登录 session 失效闭环；bot 凭旧 auth_key 重连后将得 401）。
{
if auths, err := authzStore.ListByUser(ctx, botUserID); err != nil || len(auths) == 0 {
t.Fatalf("bot authorizations before revoke = %d err=%v, want >=1", len(auths), err)
}
client := newClient(ownerStorage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
resolved, err := raw.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{Username: "BotFather"})
if err != nil {
return err
}
bf := resolved.Users[0].(*tg.User)
peer := &tg.InputPeerUser{UserID: bf.ID, AccessHash: bf.AccessHash}
randomID := int64(61001)
// 等到 BotFather 对该条的回复出现再发下一条——回复异步处理，裸 sleep 在
// 负载下会乱序（见 P2 审查的 sleep flaky 项）。轮询顶部 incoming 回复。
sendAwait := func(text string) {
t.Helper()
before := 0
if h, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{Peer: peer, Limit: 1}); err == nil {
if m := messagesOf(h); len(m) > 0 {
if msg, ok := m[0].(*tg.Message); ok {
before = msg.ID
}
}
}
randomID++
if _, err := raw.MessagesSendMessage(ctx, &tg.MessagesSendMessageRequest{
Peer: peer, Message: text, RandomID: randomID,
}); err != nil {
t.Fatalf("send %q: %v", text, err)
}
deadline := time.Now().Add(10 * time.Second)
for {
h, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{Peer: peer, Limit: 3})
if err != nil {
t.Fatalf("history after %q: %v", text, err)
}
msgs := messagesOf(h)
if len(msgs) > 0 {
if top, ok := msgs[0].(*tg.Message); ok && !top.Out && top.ID > before {
return
}
}
if time.Now().After(deadline) {
t.Fatalf("timed out waiting for BotFather reply to %q", text)
}
time.Sleep(40 * time.Millisecond)
}
}
sendAwait("/revoke")
sendAwait("@" + botUsername)
// revoke 在 BotFather goroutine 内执行；轮询 authorization 清空。
deadline := time.Now().Add(10 * time.Second)
for {
auths, err := authzStore.ListByUser(ctx, botUserID)
if err != nil {
return err
}
if len(auths) == 0 {
return nil
}
if time.Now().After(deadline) {
t.Fatalf("bot authorizations not revoked: %d rows remain", len(auths))
}
time.Sleep(50 * time.Millisecond)
}
}); err != nil {
t.Fatalf("revoke flow: %v", err)
}
}
cancel()
if err := <-serveErr; err != nil {
t.Errorf("serve: %v", err)
}
}

30
internal/mtprotoedge/conn.go
View file

@ -10,6 +10,8 @@ import (
"github.com/gotd/td/crypto"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
"telesrv/internal/compat/layerwire"
)
// Conn 是一个已识别 session 的客户端连接，持有向其加密发送消息所需的全部上下文。
@ -67,8 +69,36 @@ type Conn struct {
userID atomic.Int64
userIDResolved atomic.Bool
receivesUpdates atomic.Bool
// membershipsSynced 表示该连接的 channel membership 推送路由（byMemberChannel）
// 已成功建立。它与 receivesUpdates 共同构成「session 完全就绪」：membership
// 同步失败时保持 false，让置位短路放行、下一条 RPC 重试同步，避免
// 「已置位但 channel 路由缺失」的 session 静默漏收超级群推送。
membershipsSynced atomic.Bool
// keyDestroyed 标记本连接的 auth_key 已被 destroy_auth_key 删除。serveConn 对已建立
// 连接复用缓存密钥跳过每帧 AuthKeyStore 回查；置位后强制回落到 Get→AuthKeyNotFound，
// 维持「destroy_auth_key 发起连接下一帧自然失效」契约。只由 destroy_auth_key 处理器置位。
keyDestroyed atomic.Bool
// lastSessionSaveUnix 是上次把本连接 session 持久化到 SessionStore 的 unix 秒，用于把
// 每帧 Save 去抖到固定间隔——session 持久化是软状态（生产无热读路径，仅观测/未来用）。
// 只由单连接的读循环 goroutine 访问。
lastSessionSaveUnix atomic.Int64
// clientLayer 是本连接协商的 TL layer（invokeWithLayer/initConnection），由 handleRPC
// 在每次 Dispatch 后从 RPC 注册表刷新。出站(rpc_result/push)按此把 227 对象降级给老客户端；
// 0 表示尚未协商，按 canonical(227) 处理=不降级。
clientLayer atomic.Int32
}
// ClientLayer 返回连接协商的 TL layer；未协商时返回 canonical layer（227，不降级）。
func (c *Conn) ClientLayer() int {
if l := c.clientLayer.Load(); l != 0 {
return int(l)
}
return layerwire.CanonicalLayer
}
// SetClientLayer 记录连接协商的 TL layer。
func (c *Conn) SetClientLayer(layer int) { c.clientLayer.Store(int32(layer)) }
// AuthKeyID 返回连接的 auth_key_id。
func (c *Conn) AuthKeyID() [8]byte { return c.authKeyID }

72
internal/mtprotoedge/conn_mgmt_fixes_test.go Normal file
View file

@ -0,0 +1,72 @@
package mtprotoedge
import (
"context"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestPushSkipsConnReboundToOtherUser 锁定跨账号投递窗口的修复：pushToUserWithSender 在锁外
// 发送前复查 c.userID。模拟「收集 conns 后、send 前」连接被并发换绑（atomic userID 改了但
// byUser 索引尚未更新）的窗口，验证本属于 userA 的 update 不会投递到已易主为 userB 的连接。
func TestPushSkipsConnReboundToOtherUser(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userA, userB = int64(100), int64(200)
mk := func(sid int64, authKey byte) *Conn {
c := &Conn{
sessionID: sid,
authKeyID: [8]byte{authKey},
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outboundStop: make(chan struct{}),
}
c.userID.Store(userA)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
sm.Register(c)
return c
}
stale := mk(1, 1) // 注册为 userA 后被换绑到 userB
live := mk(2, 2) // 始终 userA
// 直接改 atomic userID、不动 byUser 索引：复现锁释放后到 send 前的换绑窗口。
stale.userID.Store(userB)
// 走 best-effort 推送：与普通推送共用 pushToUserWithSender（含 userID 复查），但只入队
// 不等 op.done，故无需起 outbound actor 即可观察「投递 vs 跳过」（op 进 buffered c.outbound）。
if _, err := sm.PushToUserExceptSessionBestEffort(context.Background(), userA, 0, proto.MessageFromServer, &tg.UpdatesTooLong{}, time.Second); err != nil {
t.Fatalf("push: %v", err)
}
if n := len(stale.outbound); n != 0 {
t.Fatalf("rebound conn received %d ops, want 0 (must skip cross-account delivery)", n)
}
if n := len(live.outbound); n != 1 {
t.Fatalf("live conn received %d ops, want 1", n)
}
}
// TestRegisterEvictsAtSessionCap 锁定单 auth_key 的 session 数上限：超出 maxSessionsPerAuthKey
// 的新 session 注册会驱逐一个现有 session，防对抗客户端用海量 session_id 撑爆索引。
func TestRegisterEvictsAtSessionCap(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
authKey := [8]byte{7}
for i := 1; i <= maxSessionsPerAuthKey+5; i++ {
// 裸 Conn（无 outbound/rpc 通道）：被驱逐时 Close() 为安全 no-op。
sm.Register(&Conn{sessionID: int64(i), authKeyID: authKey})
}
sm.mu.RLock()
perKey := len(sm.byAuthKey[authKey])
total := len(sm.bySession)
sm.mu.RUnlock()
if perKey != maxSessionsPerAuthKey {
t.Fatalf("sessions for auth key = %d, want cap %d", perKey, maxSessionsPerAuthKey)
}
if total != maxSessionsPerAuthKey {
t.Fatalf("total online = %d, want %d", total, maxSessionsPerAuthKey)
}
}

8
internal/mtprotoedge/destroy_auth_key.go
View file

@ -9,6 +9,7 @@ import (
const (
destroyAuthKeyRequestTypeID = 0xd1435160
destroyAuthKeyOkTypeID = 0xf660e1d4
destroyAuthKeyFailTypeID = 0xea109b13
)
type destroyAuthKeyRequest struct{}
@ -31,3 +32,10 @@ func (*destroyAuthKeyOk) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyOkTypeID)
return nil
}
type destroyAuthKeyFail struct{}
func (*destroyAuthKeyFail) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyFailTypeID)
return nil
}

8
internal/mtprotoedge/e2e_test.go
View file

@ -10,6 +10,7 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
@ -61,7 +62,7 @@ func TestTelegramClientEndToEnd(t *testing.T) {
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: zaptest.NewLogger(t).Named("client"),
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: &session.StorageMemory{},
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
@ -75,8 +76,9 @@ func TestTelegramClientEndToEnd(t *testing.T) {
if cfg.ThisDC != dc {
t.Errorf("config.ThisDC = %d, want %d", cfg.ThisDC, dc)
}
if len(cfg.DCOptions) == 0 {
t.Error("config.DCOptions is empty")
// 不下发 DCOptions：客户端使用自己的 DCList / 写死 static 地址。
if len(cfg.DCOptions) != 0 {
t.Errorf("config.DCOptions = %+v, want empty", cfg.DCOptions)
}
return nil
}); err != nil {

187
internal/mtprotoedge/encrypted.go
View file

@ -7,6 +7,7 @@ import (
"encoding/hex"
"errors"
"fmt"
"io"
"math"
"time"
@ -19,6 +20,8 @@ import (
"github.com/gotd/td/tgerr"
"github.com/gotd/td/transport"
"telesrv/internal/compat/layerwire"
"telesrv/internal/observability/dbtrace"
"telesrv/internal/store"
)
@ -69,22 +72,34 @@ const (
// handleEncrypted 解密加密消息，按需注册连接，处理服务消息并分发明文 payload。
// 返回（可能新建/更新的）当前连接对象，供 serveConn 维护生命周期。
func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, keyData store.AuthKeyData, b *bin.Buffer) (*Conn, error) {
key := crypto.AuthKey{Value: crypto.Key(keyData.Value), ID: keyData.ID}
// fetchedKey 非 nil 表示本帧的 auth key 是刚从 AuthKeyStore 查出的（首帧/换 auth key/被销毁
// 后回落）；为 nil 表示走快路径——serveConn 判定 current 仍持同一未销毁的 auth key，直接复用
// current.key/current.salt 解密，既不回查 AuthKeyStore 也不重建 store.AuthKeyData。
func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b *bin.Buffer) (*Conn, error) {
var key crypto.AuthKey
var serverSalt int64
if fetchedKey != nil {
key = crypto.AuthKey{Value: crypto.Key(fetchedKey.Value), ID: fetchedKey.ID}
serverSalt = fetchedKey.ServerSalt
} else {
// 快路径：复用已建立连接缓存的密钥与盐（同一 auth key 的后续帧，含同连接换 session）。
key = current.key
serverSalt = current.salt
}
data, err := s.cipher.DecryptFromBuffer(key, b)
if err != nil {
return current, fmt.Errorf("decrypt: %w", err)
}
if data.Salt != keyData.ServerSalt {
if data.Salt != serverSalt {
c := current
temp := false
if c == nil || c.sessionID != data.SessionID {
c = s.newConn(tc, key, data.SessionID, keyData.ServerSalt)
c = s.newConn(tc, key, data.SessionID, serverSalt)
temp = true
}
err := s.sendBadServerSalt(ctx, c, data.MessageID, data.SeqNo, keyData.ServerSalt)
err := s.sendBadServerSalt(ctx, c, data.MessageID, data.SeqNo, serverSalt)
if temp {
c.Close()
}
@ -100,18 +115,11 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
s.conns.Unregister(current)
current.Close()
}
current = s.newConn(tc, key, data.SessionID, keyData.ServerSalt)
current = s.newConn(tc, key, data.SessionID, serverSalt)
s.conns.Register(current)
}
if err := s.sessions.Save(ctx, store.SessionData{
ID: data.SessionID,
AuthKeyID: key.ID,
Salt: keyData.ServerSalt,
LastSeen: s.clock.Now().Unix(),
}); err != nil {
return current, fmt.Errorf("save session: %w", err)
}
s.maybePersistSession(ctx, current, data.SessionID, key.ID, serverSalt)
body := data.Data()
typeID, err := (&bin.Buffer{Buf: body}).PeekID()
@ -133,11 +141,9 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
content := clientMessageNeedsAck(typeID)
if record, ok := cs.seenRecord(data.MessageID); ok {
s.log.Debug("Duplicate msg_id; re-ack only", zap.Int64("msg_id", data.MessageID))
if resent, err := current.ResendByRequest(ctx, data.MessageID); err != nil {
s.log.Debug("Duplicate msg_id; replay cached result if available", zap.Int64("msg_id", data.MessageID))
if err := s.replayRPCResultByRequest(ctx, current, data.MessageID); err != nil {
return current, err
} else if resent {
s.log.Debug("Resent cached rpc_result for duplicate msg_id", zap.Int64("msg_id", data.MessageID))
}
if !record.content {
return current, nil
@ -175,6 +181,34 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
return current, nil
}
// sessionSaveMinInterval 是单连接持久化 session 记录的最小间隔。把原本「每帧一次 Redis SET」
// 去抖到固定间隔——session 是软状态（生产无热读路径），只需周期刷新 last_seen/续 TTL。
const sessionSaveMinInterval = 30 * time.Second
// maybePersistSession 按 sessionSaveMinInterval 去抖持久化 session，失败只告警不断连。
// 原实现每帧同步 Save 且失败即断连：N 连接×帧率的 Redis 写放大 + Redis 抖动级联断连。
func (s *Server) maybePersistSession(ctx context.Context, c *Conn, sessionID int64, authKeyID [8]byte, salt int64) {
if c == nil {
return
}
now := s.clock.Now().Unix()
if last := c.lastSessionSaveUnix.Load(); last != 0 && now-last < int64(sessionSaveMinInterval/time.Second) {
return
}
c.lastSessionSaveUnix.Store(now)
if err := s.sessions.Save(ctx, store.SessionData{
ID: sessionID,
AuthKeyID: authKeyID,
Salt: salt,
LastSeen: now,
}); err != nil {
s.log.Warn("Persist session failed (non-fatal)",
zap.Int64("session_id", sessionID),
zap.Error(err),
)
}
}
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, data *crypto.EncryptedMessageData) error {
q, ok := tc.(quickAckTransport)
if !ok || !q.ConsumeQuickAckRequested() {
@ -236,6 +270,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
}
content := clientMessageNeedsAck(typeID)
if record, ok := cs.seenRecord(m.ID); ok {
if err := s.replayRPCResultByRequest(ctx, c, m.ID); err != nil {
return err
}
if record.content {
*acks = append(*acks, m.ID)
}
@ -360,6 +397,19 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
}
ackContent()
s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])))
// 真正销毁：删密钥库记录（每帧回查，删除后该 key 的入站帧立即失效）并主动
// 断开同 key 的其他连接——出站推送用连接持有的密钥副本加密、不回查密钥库，
// 不断开的话被销毁 key 的空闲连接仍能持续收到推送。发起连接除外：响应要
// 先送达，它的下一帧会因密钥缺失自然断开。授权（authorizations）不在此清理，
// destroy_auth_key 是 PFS 密钥轮换的清理动作，不等于登出。
if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil {
s.log.Warn("Delete auth key failed", zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), zap.Error(err))
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{})
}
// 标记密钥已销毁：发起连接被 CloseSessionsForRawAuthKeyExcept 排除（响应需先送达），
// 它下一帧不能再走 serveConn 的密钥复用快路径，须回落到 Get→AuthKeyNotFound 自然失效。
c.keyDestroyed.Store(true)
s.conns.CloseSessionsForRawAuthKeyExcept(c.authKeyID, c.sessionID)
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{})
default:
@ -389,6 +439,15 @@ func mergeStateInfo(primary, fallback []byte) []byte {
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, body []byte) error {
id, _ := (&bin.Buffer{Buf: body}).PeekID()
method := s.typeName(id)
if cached, ok := s.cachedRPCResult(c, msgID); ok {
s.log.Info("RPC duplicate replay from session cache",
zap.String("method", method),
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])),
zap.Int64("session_id", c.sessionID),
)
return c.SendEncoded(ctx, proto.MessageServerResponse, cached)
}
err := c.enqueueInboundRPC(ctx, inboundRPC{
method: method,
size: len(body),
@ -431,10 +490,18 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf
return nil
}
ctx, dbStats := dbtrace.WithStats(ctx)
start := s.clock.Now()
result, err := s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b)
dur := s.clock.Now().Sub(start)
s.metrics.RPCHandled(method, dur, err)
// 刷新本连接协商 layer（invokeWithLayer/initConnection 已被 Dispatch 处理并登记），
// 供 rpc_result 与后续 push 出站降级使用。仅在确实观测到 layer 时更新——缓存被驱逐
// 时 NegotiatedLayer 返回 ok=false，此时必须保留连接已记住的 layer，绝不覆盖成默认值，
// 否则长连接老客户端的条目被驱逐后会被误降回 227。
if layer, ok := s.rpc.NegotiatedLayer(c.authKeyID, c.sessionID); ok {
c.SetClientLayer(layer)
}
fields := []zap.Field{
zap.String("method", method),
@ -449,6 +516,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf
if userID := c.UserID(); userID != 0 {
fields = append(fields, zap.Int64("user_id", userID))
}
fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot())
if err != nil {
var rpcErr *tgerr.Error
@ -472,14 +540,72 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf
// 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)
if err != nil {
return err
}
s.storeRPCResult(c, reqMsgID, encoded)
return c.SendEncoded(ctx, proto.MessageServerResponse, encoded)
}
// encodeRPCResult 编码 rpc_result。proto.Result.Result 是裸 boxed 对象字节，故在包入
// rpc_result 之前对其按连接协商 layer 降级（layer==227 直通，零开销）。降级失败 fail-safe：
// 记日志并发送 canonical 字节——宁可老客户端对个别长尾对象渲染异常，也不让连接/流崩。
func (s *Server) encodeRPCResult(c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) {
var buf bin.Buffer
if err := result.Encode(&buf); err != nil {
return fmt.Errorf("encode rpc result: %w", err)
return nil, fmt.Errorf("encode rpc result: %w", err)
}
return c.Send(ctx, proto.MessageServerResponse, &proto.Result{
inner := buf.Raw()
if layer := c.ClientLayer(); layer < layerwire.CanonicalLayer {
if down, err := layerwire.Transcode(inner, layer); err != nil {
s.log.Warn("layerwire downgrade failed; sending canonical rpc_result",
zap.Int("layer", layer), zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
} else {
inner = down
}
}
encoded, err := encodeOutboundMessage(&proto.Result{
RequestMessageID: reqMsgID,
Result: buf.Raw(),
Result: inner,
})
if err != nil {
return nil, err
}
return encoded, nil
}
func (s *Server) cachedRPCResult(c *Conn, reqMsgID int64) (*encodedOutboundMessage, bool) {
if s == nil || s.rpcResults == nil || c == nil {
return nil, false
}
return s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID)
}
func (s *Server) replayRPCResultByRequest(ctx context.Context, c *Conn, reqMsgID int64) error {
if c == nil {
return nil
}
if resent, err := c.ResendByRequest(ctx, reqMsgID); err != nil {
return err
} else if resent {
s.log.Debug("Resent connection cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
return nil
}
if cached, ok := s.cachedRPCResult(c, reqMsgID); ok {
if err := c.SendEncoded(ctx, proto.MessageServerResponse, cached); err != nil {
return err
}
s.log.Debug("Resent session cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
}
return nil
}
func (s *Server) storeRPCResult(c *Conn, reqMsgID int64, encoded *encodedOutboundMessage) {
if s == nil || s.rpcResults == nil || c == nil {
return
}
s.rpcResults.Put(c.authKeyID, c.sessionID, reqMsgID, encoded)
}
// sendPong 回复 mt.PingRequest / mt.PingDelayDisconnectRequest。
@ -515,14 +641,25 @@ func (s *Server) sendFutureSalts(ctx context.Context, c *Conn, reqMsgID int64, n
}
// sendNewSessionCreated 在连接首个加密消息后通知客户端新 session 已建立。
// unique_id 必须每个 server session 实例独立：客户端按 unique_id 去重，
// 复用同一值会让断线重连后的 new_session_created 被吞掉，错过的差分补拉
// （Android 收到后才调 getDifference）随之丢失。
func (s *Server) sendNewSessionCreated(ctx context.Context, c *Conn, firstMsgID int64) error {
return c.SendAsync(ctx, proto.MessageFromServer, &mt.NewSessionCreated{
FirstMsgID: firstMsgID,
UniqueID: s.sessionUID,
UniqueID: s.newServerSessionUID(),
ServerSalt: c.salt,
})
}
func (s *Server) newServerSessionUID() int64 {
var b [8]byte
if _, err := io.ReadFull(s.rand, b[:]); err == nil {
return int64(binary.LittleEndian.Uint64(b[:]))
}
return s.clock.Now().UnixNano()
}
// sendAck 确认收到客户端 content-related 消息。
func (s *Server) sendAck(ctx context.Context, c *Conn, ids ...int64) error {
return c.SendAsync(ctx, proto.MessageFromServer, &mt.MsgsAck{MsgIDs: ids})
@ -640,7 +777,11 @@ func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) in
func clientMessageAllowsEitherSeqParity(typeID uint32) bool {
switch typeID {
case mt.PingDelayDisconnectRequestTypeID:
case mt.PingDelayDisconnectRequestTypeID,
// get_future_salts 的 seqno 奇偶在客户端间不一致：部分客户端按内容消息发奇数，
// gotd 按服务消息发偶数。两者都合法（官方服务器都接受），故不在此卡奇偶，避免
// 误判 bad_msg 触发客户端重连风暴。ack/content 行为仍由 clientMessageNeedsAck 决定。
mt.GetFutureSaltsRequestTypeID:
return true
default:
return false

28
internal/mtprotoedge/encrypted_test.go
View file

@ -301,6 +301,34 @@ func TestPingDelayDisconnectEvenSeqAccepted(t *testing.T) {
}
}
func TestPingDelayDisconnectPongUsesEvenSeqNo(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
reqMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, reqMsgID, 0, &mt.PingDelayDisconnectRequest{
PingID: 11,
DisconnectDelay: 10,
})
for i := 0; i < 4; i++ {
data, id, _ := readServerMessage(t, conn, cipher, auth.AuthKey)
if id == mt.BadMsgNotificationTypeID {
t.Fatal("ping_delay_disconnect produced bad_msg_notification")
}
if id != mt.PongTypeID {
continue
}
if data.SeqNo%2 != 0 {
t.Fatalf("pong seq_no = %d, want even non-content-related seq_no", data.SeqNo)
}
return
}
t.Fatal("pong was not returned")
}
func TestPingDelayDisconnectOddSeqAccepted(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})

110
internal/mtprotoedge/exchange.go
View file

@ -4,11 +4,11 @@ import (
"context"
"errors"
"fmt"
"strings"
"sync"
"go.uber.org/zap"
"github.com/gotd/log/logzap"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
@ -42,20 +42,53 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
buffered := newBufferedConn(conn)
buffered.push(first)
// 给整个密钥交换设总时长上界。HandshakeIdleTimeout 只约束单次读 idle，对一个持续发包的
// 客户端无效——若客户端陷入「ResPQ→nonce 失步→重发 req_pq」的握手重启死循环，无界的
// serverExchange 会对每个 req_pq 盲回 ResPQ、永不收敛地空转刷日志/占 CPU。超时即放弃本次
// 握手并断开，客户端重连发起全新握手（无残留相位差）即恢复。
runCtx := ctx
if s.handshakeMaxDur > 0 {
var cancel context.CancelFunc
runCtx, cancel = context.WithTimeout(ctx, s.handshakeMaxDur)
defer cancel()
}
start := s.clock.Now()
res, err := exchange.NewExchanger(buffered, s.dc).
WithClock(s.clock).
WithRand(s.rand).
WithLogger(s.log.Named("exchange")).
WithLogger(logzap.New(s.log.Named("exchange"))).
Server(s.key).
Run(ctx)
Run(runCtx)
if err != nil {
if isEncryptedFrameDuringExchange(err) {
replay := buffered.lastFrame()
if replay != nil {
s.log.Debug("Key exchange interrupted by encrypted frame; replaying as existing session")
return replay, nil
// gotd v0.158：握手中读到非零 auth_key_id 帧（客户端用既有 auth key 而非重新交换）
// 经类型化 UnexpectedEncryptedError 暴露并随附原始帧（旧版仅靠错误文案匹配，升级后失效）。
// 把该帧当既有会话首帧 replay，切勿回 -404——TDesktop 会判定 temp key 被销毁、丢弃并
// 重跑密钥交换，引发重连/重交换风暴。
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
}
}
// req_pq 帧数超界（客户端握手重启死循环）：瞬断，促客户端重连发起全新握手。
if errors.Is(err, errTooManyHandshakeReqPQ) {
s.log.Info("Key exchange aborted: too many req_pq retries (client handshake restart loop)",
zap.Int("max", maxHandshakeReqPQ))
return nil, err
}
// 仅本握手的总时长上界到点（ctx 自身未取消）：放弃并断开，促客户端重连。
if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil {
s.log.Info("Key exchange aborted: exceeded max duration (possible client req_pq restart loop)",
zap.Duration("max", s.handshakeMaxDur))
return nil, err
}
var exErr *exchange.ServerExchangeError
if errors.As(err, &exErr) {
@ -67,7 +100,7 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
s.metrics.HandshakeDone(s.clock.Now().Sub(start))
s.log.Info("Key exchange completed",
zap.Object("auth_key", res.Key),
zap.Int64("auth_key_id", res.Key.IntID()),
zap.Int64("server_salt", res.ServerSalt),
zap.Duration("dur", s.clock.Now().Sub(start)),
)
@ -75,11 +108,6 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
return nil, s.authKeys.Save(ctx, authKeyData(res.Key, res.ServerSalt, s.clock.Now().Unix()))
}
func isEncryptedFrameDuringExchange(err error) bool {
msg := err.Error()
return strings.Contains(msg, "unexpected auth_key_id") && strings.Contains(msg, "plaintext message")
}
// authKeyData 把握手结果转换为 store 记录。
func authKeyData(key crypto.AuthKey, salt, createdAt int64) store.AuthKeyData {
return store.AuthKeyData{
@ -103,14 +131,27 @@ func (s *Server) sendProtoError(ctx context.Context, conn transport.Conn, code i
return nil
}
// maxHandshakeReqPQ 是一次密钥交换内允许的 req_pq(_multi) 帧数上界。正常握手只发 1 个
// req_pq（含个别客户端的「fake+真」也就 2 个）；客户端因 nonce 失步陷入「收到 ResPQ→立刻
// 重启握手换 nonce 重发 req_pq」死循环时，会在同一连接上无限发 req_pq，而委托给 gotd 的
// serverExchange 会对每个都盲回 ResPQ、永不收敛（见 docs/client-compat-notes.md 的握手风暴）。
// 超过此上界即在 telesrv 传输层瞬断该连接，促客户端重连发起全新握手（无残留相位差）即恢复。
// 留足余量（8）容纳少量正常重连重启。它与 HandshakeMaxDuration 总时长上界互补（按次/按时）。
const maxHandshakeReqPQ = 8
// errTooManyHandshakeReqPQ 表示一次握手内 req_pq 帧数超过 maxHandshakeReqPQ（疑似客户端握手
// 重启死循环）。从 bufferedConn.Recv 抛出，经 serverExchange 透传回 handleExchange 断开连接。
var errTooManyHandshakeReqPQ = errors.New("too many req_pq frames in one handshake (client restart loop)")
// bufferedConn 包装 transport.Conn，可把已读取的帧重新交给后续 Recv。
//
// 用于密钥交换：serveConn 已读首帧用于 peek auth_key_id，再 push 回来交给 exchange。
type bufferedConn struct {
transport.Conn
mu sync.Mutex
pending []bin.Buffer
last bin.Buffer
mu sync.Mutex
pending []bin.Buffer
last bin.Buffer
reqPQCount int // 本次握手已见 req_pq(_multi) 帧数；只在握手期访问（Recv 单 goroutine）
}
func newBufferedConn(conn transport.Conn) *bufferedConn {
@ -146,27 +187,46 @@ func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
if isUnencryptedMsgsAckFrame(b) {
continue
}
// req_pq 计数上界：仅在握手期生效（bufferedConn 只用于密钥交换），且 payload id 探测
// 与上面的 msgs_ack 跳过同量级开销，不碰加密消息热路径。超界即瞬断，止住握手死循环。
if isUnencryptedReqPQFrame(b) {
c.reqPQCount++
if c.reqPQCount > maxHandshakeReqPQ {
return errTooManyHandshakeReqPQ
}
}
return nil
}
}
func isUnencryptedMsgsAckFrame(frame *bin.Buffer) bool {
// 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 {
return false
return 0, false
}
var msg proto.UnencryptedMessage
copy := &bin.Buffer{Buf: frame.Copy()}
if err := msg.Decode(copy); err != nil {
return false
cp := &bin.Buffer{Buf: frame.Copy()}
if err := msg.Decode(cp); err != nil {
return 0, false
}
payload := &bin.Buffer{Buf: msg.MessageData}
id, err := payload.PeekID()
if err != nil {
return false
return 0, false
}
return id == mt.MsgsAckTypeID
return id, true
}
func isUnencryptedMsgsAckFrame(frame *bin.Buffer) bool {
id, ok := unencryptedPayloadID(frame)
return ok && id == mt.MsgsAckTypeID
}
func isUnencryptedReqPQFrame(frame *bin.Buffer) bool {
id, ok := unencryptedPayloadID(frame)
return ok && (id == mt.ReqPqRequestTypeID || id == mt.ReqPqMultiRequestTypeID)
}
func (c *bufferedConn) lastFrame() *bin.Buffer {

5
internal/mtprotoedge/exchange_test.go
View file

@ -10,6 +10,7 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/bin"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
@ -64,7 +65,7 @@ func TestKeyExchange(t *testing.T) {
defer ec()
res, err := exchange.NewExchanger(conn, dc).
WithRand(rand.Reader).
WithLogger(zaptest.NewLogger(t).Named("client")).
WithLogger(logzap.New(zaptest.NewLogger(t).Named("client"))).
Client([]exchange.PublicKey{pub}).
Run(exchCtx)
if err != nil {
@ -144,7 +145,7 @@ func TestKeyExchangeIgnoresUnencryptedMsgsAck(t *testing.T) {
defer ec()
res, err := exchange.NewExchanger(&ackingExchangeConn{Conn: conn, t: t}, dc).
WithRand(rand.Reader).
WithLogger(zaptest.NewLogger(t).Named("client")).
WithLogger(logzap.New(zaptest.NewLogger(t).Named("client"))).
Client([]exchange.PublicKey{pub}).
Run(exchCtx)
if err != nil {

66
internal/mtprotoedge/flush_identity_test.go Normal file
View file

@ -0,0 +1,66 @@
package mtprotoedge
import (
"context"
"testing"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestRunFlushDiscardsBatchOnIdentitySwitch 验证排空进行中连接易主（登出/换号致
// c.userID != owner）时，属于旧账号的暂存被丢弃、不发给新账号、也不回排进 pending。
// 这是对抗审查发现的 P0：在飞 batch 逃过 bind/unbind 的 pending 清理。
func TestRunFlushDiscardsBatchOnIdentitySwitch(t *testing.T) {
sm := NewSessionManager(nil)
raw := [8]byte{7}
const sessionID = int64(99)
key := sessionKey{authKeyID: raw, sessionID: sessionID}
c := &Conn{
sessionID: sessionID,
authKeyID: raw,
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outboundStop: make(chan struct{}),
}
sm.Register(c)
sm.BindUserForAuthKey(raw, sessionID, 100) // 当前账号 A=100
// session 未就绪：两条推送进 pending。
for i := 0; i < 2; i++ {
if err := sm.PushToSessionForAuthKey(context.Background(), raw, sessionID, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("queue pending: %v", err)
}
}
sm.mu.RLock()
pending := len(sm.pending[key])
sm.mu.RUnlock()
if pending != 2 {
t.Fatalf("pending = %d, want 2", pending)
}
// 模拟「排空已启动（flushing=true，owner=旧账号 A），但运行到时连接已换号成 B」。
sm.mu.Lock()
sm.flushing[key] = true
sm.mu.Unlock()
c.userID.Store(200) // 换号后的新账号 B
sm.runFlush(c, key, 100, 0) // owner=旧账号 A，当前 userID=B → 必须丢弃
if n := len(c.outbound); n != 0 {
t.Fatalf("sent %d pushes to new owner, want 0 (discarded)", n)
}
sm.mu.RLock()
pendingAfter := len(sm.pending[key])
flushing := sm.flushing[key]
sm.mu.RUnlock()
if pendingAfter != 0 {
t.Fatalf("pending = %d after identity switch, want 0 (discarded, not requeued)", pendingAfter)
}
if flushing {
t.Fatal("flushing flag not cleared after discard")
}
if c.receivesUpdates.Load() {
t.Fatal("receivesUpdates set despite identity switch")
}
}

61
internal/mtprotoedge/handshake_reqpq_cap_test.go Normal file
View file

@ -0,0 +1,61 @@
package mtprotoedge
import (
"context"
"errors"
"testing"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
)
// reqPQConn 是只会不断返回同一个 req_pq_multi 帧的假 transport.Conn，用于驱动 bufferedConn
// 的 req_pq 计数上界。除 Recv 外的方法不会被 bufferedConn.Recv 调用。
type reqPQConn struct {
transport.Conn
frame []byte
}
func (c *reqPQConn) Recv(_ context.Context, b *bin.Buffer) error {
b.ResetTo(append([]byte(nil), c.frame...))
return nil
}
func buildReqPQFrame(t *testing.T) []byte {
t.Helper()
var payload bin.Buffer
if err := (&mt.ReqPqMultiRequest{}).Encode(&payload); err != nil {
t.Fatalf("encode req_pq_multi: %v", err)
}
msg := proto.UnencryptedMessage{MessageID: 1, MessageData: payload.Raw()}
var frame bin.Buffer
if err := msg.Encode(&frame); err != nil {
t.Fatalf("encode unencrypted message: %v", err)
}
return append([]byte(nil), frame.Raw()...)
}
// TestBufferedConnReqPQCapAborts 锁定握手 req_pq 计数上界：连续 req_pq 超过 maxHandshakeReqPQ
// 后，bufferedConn.Recv 返回 errTooManyHandshakeReqPQ，止住握手重启死循环（不依赖 20s 总超时）。
func TestBufferedConnReqPQCapAborts(t *testing.T) {
frame := buildReqPQFrame(t)
bc := newBufferedConn(&reqPQConn{frame: frame})
ctx := context.Background()
var b bin.Buffer
// 前 maxHandshakeReqPQ 个 req_pq 正常返回（容纳「fake+真」与少量正常重连重启）。
for i := 0; i < maxHandshakeReqPQ; i++ {
if err := bc.Recv(ctx, &b); err != nil {
t.Fatalf("req_pq %d: unexpected err %v (cap should not trip yet)", i+1, err)
}
if !isUnencryptedReqPQFrame(&b) {
t.Fatalf("frame %d not recognized as req_pq", i+1)
}
}
// 第 maxHandshakeReqPQ+1 个触发上界，瞬断。
if err := bc.Recv(ctx, &b); !errors.Is(err, errTooManyHandshakeReqPQ) {
t.Fatalf("after cap: err = %v, want errTooManyHandshakeReqPQ", err)
}
}

28
internal/mtprotoedge/helpers_test.go
View file

@ -10,6 +10,7 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
@ -61,21 +62,13 @@ func startTestServer(t *testing.T, opts Options) (addr string, pub exchange.Publ
// 返回连接、握手结果与 client 端 cipher。连接通过 t.Cleanup 自动关闭。
func dialHandshake(t *testing.T, addr string, dc int, pub exchange.PublicKey) (transport.Conn, exchange.ClientExchangeResult, crypto.Cipher) {
t.Helper()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := transport.Intermediate.Handshake(raw)
if err != nil {
t.Fatalf("transport handshake: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
conn := dialTransportOnly(t, addr)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
auth, err := exchange.NewExchanger(conn, dc).
WithRand(rand.Reader).
WithLogger(zaptest.NewLogger(t).Named("client")).
WithLogger(logzap.New(zaptest.NewLogger(t).Named("client"))).
Client([]exchange.PublicKey{pub}).
Run(ctx)
if err != nil {
@ -84,6 +77,21 @@ func dialHandshake(t *testing.T, addr string, dc int, pub exchange.PublicKey) (t
return conn, auth, crypto.NewClientCipher(rand.Reader)
}
func dialTransportOnly(t *testing.T, addr string) transport.Conn {
t.Helper()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := transport.Intermediate.Handshake(raw)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
return conn
}
// sendEncrypted 用 client cipher 加密并发送一条带 msgID 的消息。
func sendEncrypted(t *testing.T, conn transport.Conn, cipher crypto.Cipher, auth exchange.ClientExchangeResult, msgID int64, msg bin.Encoder) {
t.Helper()

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

@ -0,0 +1,65 @@
package mtprotoedge
import (
"bytes"
"testing"
"github.com/gotd/td/bin"
"github.com/gotd/td/tg"
)
// TestConnDowngradedClone verifies the outbound seam downgrades a canonical
// (227) object to the connection's negotiated layer, is a no-op for 227, and
// — critically for push fan-out — never mutates the shared input message (one
// pre-encoded update is reused across many connections of differing layers).
func TestConnDowngradedClone(t *testing.T) {
const (
message227CRC = 0x7600b9d3
message220CRC = 0xb92f76cf
)
msg := &tg.Message{
ID: 2,
FromID: &tg.PeerUser{UserID: 3},
PeerID: &tg.PeerUser{UserID: 3},
Date: 1,
Message: "hi",
}
// layer 220: returns a NEW message rewritten to the 220 constructor id,
// leaving the shared input untouched (227).
enc, err := encodeOutboundMessage(msg)
if err != nil {
t.Fatalf("encode: %v", err)
}
c := &Conn{metrics: NopMetrics{}}
c.SetClientLayer(220)
out := c.downgradedClone(enc)
if id, _ := (&bin.Buffer{Buf: out.body}).PeekID(); id != message220CRC {
t.Fatalf("downgraded id = %#08x, want %#08x", id, message220CRC)
}
if out.typeID != message220CRC {
t.Fatalf("downgraded typeID = %#08x, want %#08x", out.typeID, message220CRC)
}
// Input must be unmodified — this is what makes shared push fan-out safe.
if id, _ := (&bin.Buffer{Buf: enc.body}).PeekID(); id != message227CRC {
t.Fatalf("input message was mutated: id now %#08x, want 227 %#08x", id, message227CRC)
}
// Two connections sharing one pre-encoded message get independent results.
encShared, _ := encodeOutboundMessage(msg)
c220 := &Conn{metrics: NopMetrics{}}
c220.SetClientLayer(220)
c227 := &Conn{metrics: NopMetrics{}} // ClientLayer() defaults to 227
out220 := c220.downgradedClone(encShared)
out227 := c227.downgradedClone(encShared)
if id, _ := (&bin.Buffer{Buf: out220.body}).PeekID(); id != message220CRC {
t.Fatalf("shared->220 id = %#08x, want %#08x", id, message220CRC)
}
if out227 != encShared {
t.Errorf("227 connection should pass the shared message through unchanged (same pointer)")
}
if !bytes.Equal(encShared.body, out227.body) {
t.Errorf("227 passthrough altered bytes")
}
}

17
internal/mtprotoedge/login_e2e_test.go
View file

@ -4,12 +4,14 @@ import (
"context"
"crypto/rand"
"crypto/rsa"
"fmt"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
@ -56,9 +58,12 @@ func TestLoginRegisterFlow(t *testing.T) {
authzStore := memory.NewAuthorizationStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
// seed hash 必须高于 help service 的代码默认 hash（低于默认值的 store 行会被
// 视为陈旧 seed 残留而被默认 config 覆盖），否则断言拿到的是默认 config。
const seedAppConfigHash = 1_000_000
if err := helpStore.UpsertAppConfig(context.Background(), domain.AppConfig{
Client: "tdesktop",
Hash: 4,
Hash: seedAppConfigHash,
JSON: []byte(`{"chat_read_mark_expire_period":604800,"chat_read_mark_size_threshold":50,"pm_read_date_expire_period":604800,"quote_length_max":1024,"telegram_antispam_group_size_min":200,"telegram_antispam_user_id":"5434988373"}`),
}); err != nil {
t.Fatalf("seed app config: %v", err)
@ -99,7 +104,7 @@ func TestLoginRegisterFlow(t *testing.T) {
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: zaptest.NewLogger(t).Named("client"),
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: &session.StorageMemory{},
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
@ -180,8 +185,10 @@ func TestLoginRegisterFlow(t *testing.T) {
if err != nil {
return err
}
if cfg, ok := appConfig.(*tg.HelpAppConfig); !ok || cfg.Hash != 4 {
t.Fatalf("help.getAppConfig = %T %+v, want hash=4 config", appConfig, appConfig)
// 注意：client.Run 回调里 t.Fatalf 只会杀当前 goroutine、测试主协程
// 会等到 ctx 超时——断言失败用 return fmt.Errorf 让 Run 立即返回。
if cfg, ok := appConfig.(*tg.HelpAppConfig); !ok || cfg.Hash != seedAppConfigHash {
return fmt.Errorf("help.getAppConfig = %T %+v, want seeded hash=%d config", appConfig, appConfig, seedAppConfigHash)
}
countriesRes, err := raw.HelpGetCountriesList(ctx, &tg.HelpGetCountriesListRequest{LangCode: "en"})
if err != nil {
@ -319,7 +326,7 @@ func TestPrivateMessageRoundTripFlow(t *testing.T) {
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: zaptest.NewLogger(t).Named("client"),
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: storage,
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}

236
internal/mtprotoedge/login_email_e2e_test.go Normal file
View file

@ -0,0 +1,236 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"fmt"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
"telesrv/internal/app/contacts"
"telesrv/internal/app/dialogs"
"telesrv/internal/app/help"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
// TestLoginEmailEndToEnd 端到端验证登录邮箱：设备 A 注册并设置登录邮箱（loginChange），
// 一个全新设备 B 调 sendCode 收到 sentCodeTypeEmailCode，凭任意邮箱验证码经 signIn
// (email_verification) 完成登录。
func TestLoginEmailEndToEnd(t *testing.T) {
const (
dc = 2
phone = "+8613800138777"
wantPhone = "8613800138777"
code = "12345"
email = "owner@example.com"
wantMask = "o***r@example.com"
)
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tcpAddr := ln.Addr().(*net.TCPAddr)
userStore := memory.NewUserStore()
authzStore := memory.NewAuthorizationStore()
authKeyStore := memory.NewAuthKeyStore()
passwordStore := memory.NewPasswordStore()
helpStore := memory.NewHelpStore()
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code, auth.WithPasswords(passwordStore)),
Account: account.NewService(passwordStore, account.WithUsers(userStore)),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(memory.NewDialogStore()),
}
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})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
// 关停顺序：cancel 后等 Serve 真正返回（含其末尾日志）再让测试结束，避免
// 服务端 goroutine 在测试完成后写 zaptest logger 触发 panic。
defer func() {
cancel()
<-serveErr
}()
newClient := func() *telegram.Client {
return telegram.NewClient(1, "hash", telegram.Options{
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: &session.StorageMemory{}, // 每个 client 独立 session = 独立 auth key = 独立"设备"
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
})
}
// 设备 A：注册并设置登录邮箱。client.Run 回调里断言失败一律 return error（回调跑在
// 独立 goroutine，t.Fatalf 只会杀该 goroutine 并让测试空等到超时）。
deviceA := newClient()
if err := deviceA.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(deviceA)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
hash := sent.(*tg.AuthSentCode).PhoneCodeHash
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: hash, PhoneCode: code}); err != nil {
return err
}
if _, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{PhoneNumber: phone, PhoneCodeHash: hash, FirstName: "Owner"}); err != nil {
return err
}
// 设置登录邮箱（loginChange，已登录）。
sentEmail, err := raw.AccountSendVerifyEmailCode(ctx, &tg.AccountSendVerifyEmailCodeRequest{
Purpose: &tg.EmailVerifyPurposeLoginChange{},
Email: email,
})
if err != nil {
return err
}
if sentEmail.EmailPattern != wantMask {
return fmt.Errorf("sentEmailCode pattern = %q, want %q", sentEmail.EmailPattern, wantMask)
}
verified, err := raw.AccountVerifyEmail(ctx, &tg.AccountVerifyEmailRequest{
Purpose: &tg.EmailVerifyPurposeLoginChange{},
Verification: &tg.EmailVerificationCode{Code: "whatever"},
})
if err != nil {
return err
}
ev, ok := verified.(*tg.AccountEmailVerified)
if !ok {
return fmt.Errorf("verifyEmail result = %T, want *tg.AccountEmailVerified", verified)
}
if ev.Email != email {
return fmt.Errorf("verifyEmail email = %q, want %q", ev.Email, email)
}
// getPassword 下发登录邮箱掩码。
pwd, err := raw.AccountGetPassword(ctx)
if err != nil {
return err
}
gotMask, ok := pwd.GetLoginEmailPattern()
if !ok || gotMask != wantMask {
return fmt.Errorf("getPassword login_email_pattern = %q ok=%v, want %q", gotMask, ok, wantMask)
}
return nil
}); err != nil {
t.Fatalf("device A: %v", err)
}
// 设备 B（全新 auth key）：sendCode 应改投邮箱，凭任意邮箱码登录。
deviceB := newClient()
if err := deviceB.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(deviceB)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
sentCode, ok := sent.(*tg.AuthSentCode)
if !ok {
return fmt.Errorf("sendCode result = %T, want *tg.AuthSentCode", sent)
}
emailType, ok := sentCode.Type.(*tg.AuthSentCodeTypeEmailCode)
if !ok {
return fmt.Errorf("sendCode type = %T, want *tg.AuthSentCodeTypeEmailCode (login email should switch delivery)", sentCode.Type)
}
if emailType.EmailPattern != wantMask {
return fmt.Errorf("sentCodeTypeEmailCode pattern = %q, want %q", emailType.EmailPattern, wantMask)
}
signInRes, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{
PhoneNumber: phone,
PhoneCodeHash: sentCode.PhoneCodeHash,
EmailVerification: &tg.EmailVerificationCode{Code: "any-email-code"},
})
if err != nil {
return err
}
authz, ok := signInRes.(*tg.AuthAuthorization)
if !ok {
return fmt.Errorf("signIn result = %T, want *tg.AuthAuthorization", signInRes)
}
self, ok := authz.User.(*tg.User)
if !ok || !self.Self || self.Phone != wantPhone {
return fmt.Errorf("signIn user = %+v, want self phone=%s", authz.User, wantPhone)
}
return nil
}); err != nil {
t.Fatalf("device B: %v", err)
}
// 设备 C：无法访问登录邮箱 → resetLoginEmail 清除登录邮箱、改回手机验证码登录。
deviceC := newClient()
if err := deviceC.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(deviceC)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
sentCode := sent.(*tg.AuthSentCode)
if _, ok := sentCode.Type.(*tg.AuthSentCodeTypeEmailCode); !ok {
return fmt.Errorf("pre-reset sendCode type = %T, want email code", sentCode.Type)
}
// 重置登录邮箱：返回一个新的手机验证码 sentCode（sentCodeTypeApp）。
resetRes, err := raw.AuthResetLoginEmail(ctx, &tg.AuthResetLoginEmailRequest{PhoneNumber: phone, PhoneCodeHash: sentCode.PhoneCodeHash})
if err != nil {
return err
}
resetSent, ok := resetRes.(*tg.AuthSentCode)
if !ok {
return fmt.Errorf("resetLoginEmail result = %T, want *tg.AuthSentCode", resetRes)
}
if _, ok := resetSent.Type.(*tg.AuthSentCodeTypeApp); !ok {
return fmt.Errorf("resetLoginEmail sentCode type = %T, want *tg.AuthSentCodeTypeApp (back to phone)", resetSent.Type)
}
// 用手机验证码完成登录。
signInRes, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: resetSent.PhoneCodeHash, PhoneCode: code})
if err != nil {
return err
}
if _, ok := signInRes.(*tg.AuthAuthorization); !ok {
return fmt.Errorf("post-reset signIn result = %T, want *tg.AuthAuthorization", signInRes)
}
return nil
}); err != nil {
t.Fatalf("device C: %v", err)
}
}

45
internal/mtprotoedge/new_session_uid_test.go Normal file
View file

@ -0,0 +1,45 @@
package mtprotoedge
import (
"testing"
"time"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestNewSessionCreatedUniqueIDPerSession 验证两次 session 建立收到的
// new_session_created.unique_id 互不相同。客户端按 unique_id 去重，复用同一值
// 会让断线重连后的 new_session_created 被吞掉，依赖它触发的差分补拉随之丢失。
func TestNewSessionCreatedUniqueIDPerSession(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
firstMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, firstMsgID, 1, &tg.HelpGetConfigRequest{})
first := collectReplies(t, conn, cipher, auth.AuthKey, mt.NewSessionCreatedTypeID)
var created1 mt.NewSessionCreated
if err := created1.Decode(mustHave(t, first, mt.NewSessionCreatedTypeID, "first new_session_created")); err != nil {
t.Fatalf("decode first new_session_created: %v", err)
}
nextSessionID := auth.SessionID + 1
if nextSessionID == 0 {
nextSessionID++
}
secondMsgID := clientMsgID.New(proto.MessageFromClient)
body := encodeClientMessageBodyForTest(t, &tg.HelpGetConfigRequest{})
sendEncryptedWithSessionSaltAndSeq(t, conn, cipher, auth, nextSessionID, auth.ServerSalt, secondMsgID, 1, body)
second := collectReplies(t, conn, cipher, auth.AuthKey, mt.NewSessionCreatedTypeID)
var created2 mt.NewSessionCreated
if err := created2.Decode(mustHave(t, second, mt.NewSessionCreatedTypeID, "second new_session_created")); err != nil {
t.Fatalf("decode second new_session_created: %v", err)
}
if created1.UniqueID == created2.UniqueID {
t.Fatalf("new_session_created.unique_id reused across sessions: %d", created1.UniqueID)
}
}

69
internal/mtprotoedge/outbound.go
View file

@ -14,6 +14,8 @@ import (
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"telesrv/internal/compat/layerwire"
)
var (
@ -117,6 +119,16 @@ func (c *Conn) Close() {
})
}
// ForceClose 停止连接并关闭底层 transport。
// 仅用于授权撤销 / destroy_auth_key 这类“必须让对端立即断线”的路径；普通生命周期仍由
// serveConn 统一关闭 transport，避免正常 push/索引清理把长连接误伤成硬断。
func (c *Conn) ForceClose() {
if c.transport != nil {
_ = c.transport.Close()
}
c.Close()
}
// Send 加密并发送一条 server 消息。
func (c *Conn) Send(ctx context.Context, t proto.MessageType, msg bin.Encoder) error {
return c.send(ctx, t, msg, false)
@ -424,7 +436,9 @@ func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
if err == nil {
err = c.writeFrame(op.ctx, frame)
}
if err == nil && frameNeedsAck(frame.typeID) {
if err == nil && frame != nil && frameNeedsAck(frame.typeID) {
// 写成功后才提交 content seq_no 递增（peekSeqNo 已按当前计数算好本帧 seq_no）。
c.commitContentSeqNo()
if dropped := state.add(frame); dropped > 0 {
for i := 0; i < dropped; i++ {
c.metrics.OutboundDropped("tracked_queue_overflow")
@ -502,17 +516,55 @@ func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder, encoded *encoded
return nil, err
}
}
// 出站统一在此按本连接协商 layer 降级：
// - push fan-out 用 onceEncodedOutbound 把更新编码一次(canonical)再 SendEncoded 给多条
// 连接共享，故必须在此**逐连接**降级，且**绝不改共享 encoded**(downgradedClone 拷贝)。
// - rpc_result 的内层对象已在 encodeRPCResult 按 layer 降级，其 mt.* 外壳在此为顶层直通(no-op)。
// - 控制消息(mt.*)顶层直通。layer>=227 整条零开销。
encoded = c.downgradedClone(encoded)
content := frameNeedsAck(encoded.typeID)
msgID := c.msgID.New(t)
return &outboundFrame{
msgID: msgID,
seqNo: c.nextSeqNo(content),
seqNo: c.peekSeqNo(content),
typeID: encoded.typeID,
body: encoded.body,
reqMsgID: encoded.reqMsgID,
}, nil
}
// downgradedClone 返回按本连接协商 layer 降级后的消息，**绝不修改入参**——push fan-out
// 多条连接共享同一 encoded，逐连接降级必须各自拷贝，否则会污染其他连接的字节。
// layer>=227 或 Transcode 直通(mt.* / 无变化)时原样返回入参，零拷贝。降级失败 fail-safe：
// 返回 canonical 并计 metrics（宁可老客户端对个别长尾对象渲染异常，也不让连接/流崩）。
func (c *Conn) downgradedClone(encoded *encodedOutboundMessage) *encodedOutboundMessage {
if encoded == nil {
return nil
}
if c.ClientLayer() >= layerwire.CanonicalLayer {
return encoded
}
down, err := layerwire.Transcode(encoded.body, c.ClientLayer())
if err != nil {
c.metrics.OutboundDropped("layerwire_downgrade_failed")
return encoded
}
if sameBacking(down, encoded.body) {
return encoded // 直通：未变(mt.*/顶层未知)，无需拷贝或重算 typeID
}
out := &encodedOutboundMessage{body: down, typeID: encoded.typeID, reqMsgID: encoded.reqMsgID}
if id, e := (&bin.Buffer{Buf: down}).PeekID(); e == nil {
out.typeID = id
}
return out
}
// sameBacking reports whether a and b share the same backing array (Transcode
// returns its input unchanged for passthrough cases).
func sameBacking(a, b []byte) bool {
return len(a) == len(b) && (len(a) == 0 || &a[0] == &b[0])
}
func encodeOutboundMessage(msg bin.Encoder) (*encodedOutboundMessage, error) {
if msg == nil {
return nil, errors.New("nil outbound message")
@ -532,15 +584,22 @@ func encodeOutboundMessage(msg bin.Encoder) (*encodedOutboundMessage, error) {
}, nil
}
func (c *Conn) nextSeqNo(content bool) int32 {
// peekSeqNo 计算本帧的 seq_no，但不提交 content 计数递增——递增延到 writeFrame 成功后
// （commitContentSeqNo）。这样写失败（超时/连接关）但连接存活时，下一条 content 帧会复用
// 同一 seq_no 而非留下间隙，避免严格校验的客户端把间隙误判为丢帧。只由 outbound actor 调用。
func (c *Conn) peekSeqNo(content bool) int32 {
seqNo := c.sentContentMessages * 2
if content {
seqNo++
c.sentContentMessages++
}
return seqNo
}
// commitContentSeqNo 在一条 content 帧成功写出后提交 seq_no 递增。只由 outbound actor 调用。
func (c *Conn) commitContentSeqNo() {
c.sentContentMessages++
}
func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
if ctx == nil {
ctx = context.Background()
@ -628,6 +687,8 @@ func growBinBufferLen(b *bin.Buffer, n int) {
func frameNeedsAck(typeID uint32) bool {
switch typeID {
case mt.MsgsAckTypeID,
mt.PongTypeID,
mt.FutureSaltsTypeID,
mt.BadMsgNotificationTypeID,
mt.BadServerSaltTypeID,
mt.MsgsStateInfoTypeID,

20
internal/mtprotoedge/outbound_test.go
View file

@ -104,6 +104,26 @@ func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
}
}
func TestFrameNeedsAckServiceExceptions(t *testing.T) {
cases := []struct {
name string
id uint32
want bool
}{
{name: "pong", id: mt.PongTypeID, want: false},
{name: "future_salts", id: mt.FutureSaltsTypeID, want: false},
{name: "msgs_ack", id: mt.MsgsAckTypeID, want: false},
{name: "updatesTooLong", id: tg.UpdatesTooLongTypeID, want: true},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := frameNeedsAck(tc.id); got != tc.want {
t.Fatalf("frameNeedsAck(%s) = %v, want %v", tc.name, got, tc.want)
}
})
}
}
func TestOutboundResendAndAckState(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})

297
internal/mtprotoedge/passkey_e2e_test.go Normal file
View file

@ -0,0 +1,297 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"encoding/base64"
"encoding/json"
"fmt"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
"telesrv/internal/app/contacts"
"telesrv/internal/app/dialogs"
"telesrv/internal/app/help"
passkeyapp "telesrv/internal/app/passkey"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
"telesrv/internal/webauthn/webauthntest"
)
const passkeyTestOrigin = "android:apk-key-hash:e2e-test"
// publicKeyOptions 提取 DataJSON(顶层 publicKey)里的 challenge / rpId / user.id。
type publicKeyOptions struct {
PublicKey struct {
Challenge string `json:"challenge"`
RPID string `json:"rpId"`
RP struct {
ID string `json:"id"`
} `json:"rp"`
User struct {
ID string `json:"id"`
} `json:"user"`
} `json:"publicKey"`
}
func parsePasskeyOptions(t *testing.T, data string) publicKeyOptions {
t.Helper()
var opts publicKeyOptions
if err := json.Unmarshal([]byte(data), &opts); err != nil {
t.Fatalf("parse passkey options %q: %v", data, err)
}
return opts
}
func decodeB64URL(t *testing.T, s string) []byte {
t.Helper()
b, err := base64.RawURLEncoding.DecodeString(s)
if err != nil {
if b, err = base64.URLEncoding.DecodeString(s); err != nil {
t.Fatalf("base64url decode %q: %v", s, err)
}
}
return b
}
// TestPasskeyEndToEnd 端到端验证 passkey:设备 A(已登录)注册 passkey,设备 B(全新
// auth key)经 initPasskeyLogin/finishPasskeyLogin 用软件 authenticator 的断言登录;再
// getPasskeys/deletePasskey 往返。整条链路验签真实(软件 authenticator 用真私钥签名)。
func TestPasskeyEndToEnd(t *testing.T) {
const (
dc = 2
phone = "+8613800139000"
wantPhone = "8613800139000"
code = "12345"
rpID = "telesrv.test"
)
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tcpAddr := ln.Addr().(*net.TCPAddr)
userStore := memory.NewUserStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
passkeyService := passkeyapp.NewService(memory.NewPasskeyStore(), memory.NewPasskeyChallengeStore(), rpID, dc)
deps := rpc.Deps{
Auth: auth.NewService(userStore, memory.NewAuthorizationStore(), memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code),
Account: account.NewService(memory.NewPasswordStore(), account.WithUsers(userStore)),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(memory.NewDialogStore()),
Passkey: passkeyService,
}
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})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
defer func() {
cancel()
<-serveErr
}()
newClient := func() *telegram.Client {
return telegram.NewClient(1, "hash", telegram.Options{
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: &session.StorageMemory{},
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
})
}
authn, err := webauthntest.New()
if err != nil {
t.Fatalf("new authenticator: %v", err)
}
var userHandle string
// 设备 A:注册 + 注册 passkey + getPasskeys。
deviceA := newClient()
if err := deviceA.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(deviceA)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
hash := sent.(*tg.AuthSentCode).PhoneCodeHash
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: hash, PhoneCode: code}); err != nil {
return err
}
if _, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{PhoneNumber: phone, PhoneCodeHash: hash, FirstName: "Pass"}); err != nil {
return err
}
// 1) 注册选项 → 软件 authenticator 造 attestation。
regOpts, err := raw.AccountInitPasskeyRegistration(ctx)
if err != nil {
return err
}
opts := parsePasskeyOptions(t, regOpts.Options.Data)
if opts.PublicKey.RP.ID != rpID {
return fmt.Errorf("registration rp.id = %q, want %q", opts.PublicKey.RP.ID, rpID)
}
userHandle = string(decodeB64URL(t, opts.PublicKey.User.ID)) // "2:<userId>"
challenge := decodeB64URL(t, opts.PublicKey.Challenge)
clientData, attObj, err := authn.Register(rpID, passkeyTestOrigin, challenge)
if err != nil {
return err
}
credB64 := authn.CredentialIDB64()
cred := &tg.InputPasskeyCredentialPublicKey{
ID: credB64,
RawID: credB64,
Response: &tg.InputPasskeyResponseRegister{
ClientData: tg.DataJSON{Data: string(clientData)},
AttestationData: attObj,
},
}
pk, err := raw.AccountRegisterPasskey(ctx, cred)
if err != nil {
return err
}
if pk.ID != credB64 {
return fmt.Errorf("registered passkey id = %q, want %q", pk.ID, credB64)
}
// 2) getPasskeys 应含该 passkey。
list, err := raw.AccountGetPasskeys(ctx)
if err != nil {
return err
}
if len(list.Passkeys) != 1 || list.Passkeys[0].ID != credB64 {
return fmt.Errorf("getPasskeys = %+v, want 1 passkey id=%q", list.Passkeys, credB64)
}
return nil
}); err != nil {
t.Fatalf("device A: %v", err)
}
// 设备 B:全新 auth key,用 passkey 登录。
deviceB := newClient()
if err := deviceB.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(deviceB)
loginOpts, err := raw.AuthInitPasskeyLogin(ctx, &tg.AuthInitPasskeyLoginRequest{APIID: 1, APIHash: "hash"})
if err != nil {
return err
}
opts := parsePasskeyOptions(t, loginOpts.Options.Data)
if opts.PublicKey.RPID != rpID {
return fmt.Errorf("login rpId = %q, want %q", opts.PublicKey.RPID, rpID)
}
challenge := decodeB64URL(t, opts.PublicKey.Challenge)
clientData, authData, sig, err := authn.Assert(rpID, passkeyTestOrigin, challenge)
if err != nil {
return err
}
credB64 := authn.CredentialIDB64()
res, err := raw.AuthFinishPasskeyLogin(ctx, &tg.AuthFinishPasskeyLoginRequest{
Credential: &tg.InputPasskeyCredentialPublicKey{
ID: credB64,
RawID: credB64,
Response: &tg.InputPasskeyResponseLogin{
ClientData: tg.DataJSON{Data: string(clientData)},
AuthenticatorData: authData,
Signature: sig,
UserHandle: userHandle,
},
},
})
if err != nil {
return err
}
authz, ok := res.(*tg.AuthAuthorization)
if !ok {
return fmt.Errorf("finishPasskeyLogin result = %T, want *tg.AuthAuthorization", res)
}
self, ok := authz.User.(*tg.User)
if !ok || !self.Self || self.Phone != wantPhone {
return fmt.Errorf("passkey login user = %+v, want self phone=%s", authz.User, wantPhone)
}
// 第二次断言(计数器递增)仍应通过。
loginOpts2, err := raw.AuthInitPasskeyLogin(ctx, &tg.AuthInitPasskeyLoginRequest{APIID: 1, APIHash: "hash"})
if err != nil {
return err
}
ch2 := decodeB64URL(t, parsePasskeyOptions(t, loginOpts2.Options.Data).PublicKey.Challenge)
cd2, ad2, sig2, err := authn.Assert(rpID, passkeyTestOrigin, ch2)
if err != nil {
return err
}
if _, err := raw.AuthFinishPasskeyLogin(ctx, &tg.AuthFinishPasskeyLoginRequest{
Credential: &tg.InputPasskeyCredentialPublicKey{
ID: credB64, RawID: credB64,
Response: &tg.InputPasskeyResponseLogin{
ClientData: tg.DataJSON{Data: string(cd2)}, AuthenticatorData: ad2, Signature: sig2, UserHandle: userHandle,
},
},
}); err != nil {
return fmt.Errorf("second passkey login: %w", err)
}
return nil
}); err != nil {
t.Fatalf("device B: %v", err)
}
// 设备 A 再次连接:删除 passkey → getPasskeys 空。
deviceC := newClient()
if err := deviceC.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(deviceC)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: sent.(*tg.AuthSentCode).PhoneCodeHash, PhoneCode: code}); err != nil {
return err
}
deleted, err := raw.AccountDeletePasskey(ctx, authn.CredentialIDB64())
if err != nil {
return err
}
if !deleted {
return fmt.Errorf("deletePasskey returned false")
}
list, err := raw.AccountGetPasskeys(ctx)
if err != nil {
return err
}
if len(list.Passkeys) != 0 {
return fmt.Errorf("getPasskeys after delete = %+v, want empty", list.Passkeys)
}
return nil
}); err != nil {
t.Fatalf("device C: %v", err)
}
}

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

@ -0,0 +1,56 @@
package mtprotoedge
import (
"context"
"testing"
"time"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestSetReceivesUpdatesFlushesPendingBeforeActivation 验证置位时先排空暂存推送
// 再激活 receivesUpdates：客户端必须真实收到暂存消息，且激活最终完成。
// 排空失败时 receivesUpdates 保持 false、暂存回排，由下一次置位重试。
func TestSetReceivesUpdatesFlushesPendingBeforeActivation(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncryptedWithSeq(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), 1, &tg.HelpGetConfigRequest{})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
raw := auth.AuthKey.ID
ctx := context.Background()
// 完全就绪还要求 membership 路由建立（ReceivesUpdatesForAuthKey 的另一半条件）。
srv.Conns().BindUserForAuthKey(raw, auth.SessionID, 100)
srv.Conns().SetSessionChannelMemberships(raw, auth.SessionID, 100, nil)
// 未就绪：推送进 pending 而非直发。
for i := 0; i < 2; i++ {
if err := srv.Conns().PushToSessionForAuthKey(ctx, raw, auth.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("queue pending push: %v", err)
}
}
if srv.Conns().ReceivesUpdatesForAuthKey(raw, auth.SessionID) {
t.Fatal("session ready before activation")
}
srv.Conns().SetReceivesUpdatesForAuthKey(raw, auth.SessionID, true)
// 客户端应收到暂存的 updatesTooLong（flush 直发）。
replies := collectReplies(t, conn, cipher, auth.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, replies, tg.UpdatesTooLongTypeID, "flushed pending push")
// 排空完成后 receivesUpdates 才置位（异步，轮询等待）。
deadline := time.Now().Add(5 * time.Second)
for !srv.Conns().ReceivesUpdatesForAuthKey(raw, auth.SessionID) {
if time.Now().After(deadline) {
t.Fatal("session never became ready after flush")
}
time.Sleep(10 * time.Millisecond)
}
}

156
internal/mtprotoedge/rpc_result_cache.go Normal file
View file

@ -0,0 +1,156 @@
package mtprotoedge
import (
"container/list"
"sync"
"time"
)
const (
rpcResultCacheTTL = 3 * time.Minute
rpcResultCacheMaxEntries = 4096
rpcResultCacheMaxBytes = 64 << 20
)
type rpcResultCacheKey struct {
authKeyID [8]byte
sessionID int64
reqMsgID int64
}
type rpcResultCacheEntry struct {
key rpcResultCacheKey
encoded *encodedOutboundMessage
size int
expiresAt time.Time
}
type rpcResultCache struct {
mu sync.Mutex
now func() time.Time
ttl time.Duration
maxEntries int
maxBytes int
bytes int
order *list.List
byKey map[rpcResultCacheKey]*list.Element
}
func newRPCResultCache(now func() time.Time) *rpcResultCache {
if now == nil {
now = time.Now
}
return &rpcResultCache{
now: now,
ttl: rpcResultCacheTTL,
maxEntries: rpcResultCacheMaxEntries,
maxBytes: rpcResultCacheMaxBytes,
order: list.New(),
byKey: make(map[rpcResultCacheKey]*list.Element),
}
}
func (c *rpcResultCache) Get(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool) {
if c == nil || reqMsgID == 0 {
return nil, false
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
now := c.now()
c.mu.Lock()
defer c.mu.Unlock()
elem, ok := c.byKey[key]
if !ok {
return nil, false
}
entry := elem.Value.(*rpcResultCacheEntry)
if !entry.expiresAt.After(now) {
c.removeElement(elem)
return nil, false
}
return cloneEncodedOutboundMessage(entry.encoded), true
}
func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encoded *encodedOutboundMessage) {
if c == nil || reqMsgID == 0 || encoded == nil {
return
}
copied := cloneEncodedOutboundMessage(encoded)
if copied == nil {
return
}
size := len(copied.body)
if c.maxBytes > 0 && size > c.maxBytes {
return
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
now := c.now()
c.mu.Lock()
defer c.mu.Unlock()
c.expireLocked(now)
if elem, ok := c.byKey[key]; ok {
c.removeElement(elem)
}
entry := &rpcResultCacheEntry{
key: key,
encoded: copied,
size: size,
expiresAt: now.Add(c.ttl),
}
elem := c.order.PushBack(entry)
c.byKey[key] = elem
c.bytes += size
c.trimLocked()
}
func (c *rpcResultCache) expireLocked(now time.Time) {
for elem := c.order.Front(); elem != nil; {
next := elem.Next()
entry := elem.Value.(*rpcResultCacheEntry)
if entry.expiresAt.After(now) {
return
}
c.removeElement(elem)
elem = next
}
}
func (c *rpcResultCache) trimLocked() {
for c.order.Len() > 0 {
tooManyEntries := c.maxEntries > 0 && c.order.Len() > c.maxEntries
tooManyBytes := c.maxBytes > 0 && c.bytes > c.maxBytes
if !tooManyEntries && !tooManyBytes {
return
}
c.removeElement(c.order.Front())
}
}
func (c *rpcResultCache) removeElement(elem *list.Element) {
if elem == nil {
return
}
entry := elem.Value.(*rpcResultCacheEntry)
delete(c.byKey, entry.key)
c.bytes -= entry.size
if c.bytes < 0 {
c.bytes = 0
}
c.order.Remove(elem)
}
func cloneEncodedOutboundMessage(src *encodedOutboundMessage) *encodedOutboundMessage {
if src == nil {
return nil
}
body := append([]byte(nil), src.body...)
return &encodedOutboundMessage{
body: body,
typeID: src.typeID,
reqMsgID: src.reqMsgID,
}
}

57
internal/mtprotoedge/rpc_test.go
View file

@ -2,6 +2,7 @@ package mtprotoedge
import (
"context"
"sync/atomic"
"testing"
"time"
@ -59,11 +60,9 @@ func TestRPCGetConfig(t *testing.T) {
if cfg.ThisDC != dc {
t.Fatalf("config.ThisDC = %d, want %d", cfg.ThisDC, dc)
}
if len(cfg.DCOptions) != 1 {
t.Fatalf("config.DCOptions count = %d, want 1", len(cfg.DCOptions))
}
if got := cfg.DCOptions[0]; got.ID != dc || got.IPAddress != advIP || got.Port != advPort {
t.Fatalf("DCOption = %+v, want id=%d ip=%s port=%d", got, dc, advIP, advPort)
// 不下发 DCOptions：客户端使用写死的 static DC 地址（空列表令其保留本地地址）。
if len(cfg.DCOptions) != 0 {
t.Fatalf("config.DCOptions = %+v, want empty (client uses pinned static address)", cfg.DCOptions)
}
}
@ -107,6 +106,52 @@ func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
close(handler.release)
}
func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) {
const dc = 2
handler := &countingConfigRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
reqMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, reqMsgID, &tg.HelpGetConfigRequest{})
first := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, reqMsgID)
if calls := handler.calls.Load(); calls != 1 {
t.Fatalf("handler calls after first request = %d, want 1", calls)
}
var firstConfig tg.Config
if err := firstConfig.Decode(&bin.Buffer{Buf: first.Result}); err != nil {
t.Fatalf("decode first config: %v", err)
}
if firstConfig.ThisDC != dc {
t.Fatalf("first config.ThisDC = %d, want %d", firstConfig.ThisDC, dc)
}
_ = conn.Close()
replayConn := dialTransportOnly(t, addr)
sendEncrypted(t, replayConn, cipher, auth, reqMsgID, &tg.HelpGetConfigRequest{})
second := readRPCResultForRequest(t, replayConn, cipher, auth.AuthKey, reqMsgID)
if calls := handler.calls.Load(); calls != 1 {
t.Fatalf("handler calls after replay = %d, want 1", calls)
}
if string(second.Result) != string(first.Result) {
t.Fatalf("replayed rpc_result payload changed")
}
}
type countingConfigRPC struct {
calls atomic.Int32
}
func (h *countingConfigRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) {
h.calls.Add(1)
return &tg.Config{ThisDC: 2}, nil
}
func (h *countingConfigRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
type blockingRPC struct {
started chan struct{}
release chan struct{}
@ -125,6 +170,8 @@ func (h *blockingRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.B
}
}
func (h *blockingRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func readRPCResultForRequest(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey, reqMsgID int64) proto.Result {
t.Helper()
for i := 0; i < 12; i++ {

294
internal/mtprotoedge/same_port_mux.go Normal file
View file

@ -0,0 +1,294 @@
package mtprotoedge
import (
"bytes"
"context"
"errors"
"io"
"net"
"net/http"
"strings"
"sync"
"time"
)
// samePortMuxBacklog 是 tcp/http 两个子 listener 的握手缓冲深度，吸收接入突发。
const samePortMuxBacklog = 1024
// websocketAllowedPaths 是允许升级为 WebSocket 的本地路径白名单。
//
// telegram-tt(WebA)按 `/apiws{_test}{_premium}` 拼 URL，故四种组合都要放行；
// 其中 `/apiws_test_premium` 是「测试服 + 会员」组合，缺它会让会员账号在测试服 404。
var websocketAllowedPaths = map[string]struct{}{
"/apiws": {},
"/apiws_test": {},
"/apiws_premium": {},
"/apiws_test_premium": {},
}
// samePortMux 在同一个 listener 上把「HTTP(WebSocket 升级请求)」与「裸 MTProto TCP」
// 两类连接拆开：每条新连接只窥探前 4 字节即可判定走向。每条连接的窥探都在各自的
// goroutine 里完成（带 sniffTimeout 上界），慢连接只占用自己的 goroutine，绝不阻塞其他
// 连接的接入与分流——这避免了固定 worker 池被 slow-loris 占满导致的接入饥饿。
type samePortMux struct {
base net.Listener
addr net.Addr
sniffTimeout time.Duration
tcp *samePortMuxListener
http *samePortMuxListener
closed chan struct{}
once sync.Once
}
func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux {
if sniffTimeout <= 0 {
sniffTimeout = 5 * time.Second
}
m := &samePortMux{
base: base,
addr: base.Addr(),
sniffTimeout: sniffTimeout,
closed: make(chan struct{}),
}
m.tcp = newSamePortMuxListener(m.addr, m.closed)
m.http = newSamePortMuxListener(m.addr, m.closed)
return m
}
// TCP 返回裸 MTProto TCP 连接的 listener（连接已窥探，前 4 字节会被回放）。
func (m *samePortMux) TCP() net.Listener {
return m.tcp
}
// HTTP 返回 WebSocket 升级请求的 listener，交给 http.Server.Serve。
func (m *samePortMux) HTTP() net.Listener {
return m.http
}
func (m *samePortMux) Serve(ctx context.Context) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
go func() {
<-ctx.Done()
_ = m.Close()
}()
// 每条连接一个窥探 goroutine：wg 让 Serve 在退出前等待在途窥探把连接交接完成。
var wg sync.WaitGroup
defer wg.Wait()
for {
conn, err := m.base.Accept()
if err != nil {
if ctx.Err() != nil || isSamePortMuxClosed(m.closed) || isNetClosed(err) {
return nil
}
return err
}
wg.Add(1)
go func() {
defer wg.Done()
m.dispatch(ctx, conn)
}()
}
}
func (m *samePortMux) Close() error {
m.once.Do(func() {
close(m.closed)
_ = m.tcp.Close()
_ = m.http.Close()
_ = m.base.Close()
})
return nil
}
// dispatch 窥探单条连接的前 4 字节并把它交给 tcp 或 http 子 listener。窥探带 sniffTimeout
// 读上界，慢/半开连接最多占用本 goroutine sniffTimeout 后即被回收。
func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
var header [4]byte
if err := conn.SetReadDeadline(time.Now().Add(m.sniffTimeout)); err != nil {
_ = conn.Close()
return
}
if _, err := io.ReadFull(conn, header[:]); err != nil {
_ = conn.Close()
return
}
if err := conn.SetReadDeadline(time.Time{}); err != nil {
_ = conn.Close()
return
}
wrapped := &prefixedNetConn{
Conn: conn,
reader: io.MultiReader(bytes.NewReader(header[:]), conn),
}
target := m.tcp
if isHTTPHeaderPrefix(header) {
target = m.http
}
if !target.deliver(ctx, wrapped) {
_ = conn.Close()
}
}
// isHTTPHeaderPrefix 判断前 4 字节是否是 HTTP 请求行起始。
//
// 这里只认 GET/POST/HEAD/OPTI，与 gotd generateInit 排除的前缀集合「严格对齐」：合法的
// obfuscated2 init 头被保证不会以这四个前缀开头（见 mtproxy/obfuscated2/keys_util.go），
// 故裸 MTProto 永不会被误判为 HTTP。刻意不扩展到 PUT/DELETE 等其他方法——generateInit
// 并未排除它们，扩展白名单反而会让随机 init 偶发（2^-32）被误分流。真实 WebSocket 升级一律
// 是 GET，浏览器不会用其他方法,因此当前集合既安全又完备。
func isHTTPHeaderPrefix(header [4]byte) bool {
switch string(header[:]) {
case "GET ", "POST", "HEAD", "OPTI":
return true
default:
return false
}
}
func websocketRouteHandler(handler http.Handler, allowedOrigins []string) http.Handler {
origins := websocketOriginSet(allowedOrigins)
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if _, ok := websocketAllowedPaths[r.URL.Path]; !ok {
http.NotFound(w, r)
return
}
// gotd 的 WebsocketListener 把 websocket.Accept 的 AcceptOptions 写死且不带
// InsecureSkipVerify/OriginPatterns，coder/websocket 因此会对「Origin.Host != Host」
// 的握手返回 403。浏览器(WebA/telegram-tt)发起的 ws 连接必然带 Origin(=页面来源，
// ≠ 本服务监听地址)，握手会被无条件拒绝；而不带 Origin 的非浏览器客户端(gotd 测试
// 客户端)却能通过——故单测全绿、真浏览器全挂。
//
// 白名单确认后再把 Origin 改写成与 Host 同源，让 Accept 放行且无需 fork gotd。
// 无 Origin 的非浏览器客户端允许通过；浏览器来源必须显式配置，"*" 仅用于临时调试。
if !websocketOriginAllowed(origins, r.Header.Get("Origin")) {
http.Error(w, "websocket origin forbidden", http.StatusForbidden)
return
}
if r.Header.Get("Origin") != "" {
r.Header.Set("Origin", "http://"+r.Host)
}
handler.ServeHTTP(w, r)
})
}
func websocketOriginSet(origins []string) map[string]struct{} {
out := make(map[string]struct{}, len(origins))
for _, origin := range origins {
origin = strings.TrimRight(strings.TrimSpace(origin), "/")
if origin != "" {
out[strings.ToLower(origin)] = struct{}{}
}
}
return out
}
func websocketOriginAllowed(allowed map[string]struct{}, origin string) bool {
origin = strings.TrimRight(strings.TrimSpace(origin), "/")
if origin == "" {
return true
}
if _, ok := allowed["*"]; ok {
return true
}
_, ok := allowed[strings.ToLower(origin)]
return ok
}
func minDuration(a, b time.Duration) time.Duration {
if a <= 0 {
return b
}
if b <= 0 || a < b {
return a
}
return b
}
func isNetClosed(err error) bool {
return errors.Is(err, net.ErrClosed)
}
func isSamePortMuxClosed(ch <-chan struct{}) bool {
select {
case <-ch:
return true
default:
return false
}
}
// prefixedNetConn 把被窥探掉的前缀字节回放在数据流最前面，使下游(去混淆/codec 探测/
// http.Server)看到完整原始字节流。
type prefixedNetConn struct {
reader io.Reader
net.Conn
}
func (p *prefixedNetConn) Read(b []byte) (int, error) {
return p.reader.Read(b)
}
// samePortMuxListener 是一个内存 listener：dispatch 把分流后的连接投递进来，下游
// (serveMixed 的 accept 循环 / http.Server) 从这里 Accept。
type samePortMuxListener struct {
addr net.Addr
ch chan net.Conn
closed chan struct{}
once sync.Once
}
func newSamePortMuxListener(addr net.Addr, parentClosed <-chan struct{}) *samePortMuxListener {
closed := make(chan struct{})
l := &samePortMuxListener{
addr: addr,
ch: make(chan net.Conn, samePortMuxBacklog),
closed: closed,
}
go func() {
select {
case <-parentClosed:
_ = l.Close()
case <-closed:
}
}()
return l
}
func (l *samePortMuxListener) Accept() (net.Conn, error) {
select {
case <-l.closed:
return nil, net.ErrClosed
case conn := <-l.ch:
return conn, nil
}
}
func (l *samePortMuxListener) Close() error {
l.once.Do(func() {
close(l.closed)
})
return nil
}
func (l *samePortMuxListener) Addr() net.Addr {
return l.addr
}
func (l *samePortMuxListener) deliver(ctx context.Context, conn net.Conn) bool {
select {
case <-l.closed:
return false
case l.ch <- conn:
return true
case <-ctx.Done():
return false
}
}

233
internal/mtprotoedge/server.go
View file

@ -7,6 +7,7 @@ import (
"fmt"
"io"
"net"
"net/http"
"sync"
"time"
@ -33,6 +34,12 @@ import (
// 返回 *tgerr.Error 时连接层将其转为 rpc_error 回发；其他 error 视为连接级故障。
type RPCHandler interface {
Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error)
// NegotiatedLayer returns the TL layer the session negotiated via
// invokeWithLayer and whether one was ever observed. Used to downgrade
// outbound objects for clients compiled on an older layer. ok=false means
// unknown (cold/evicted) — the caller must keep the connection's last-known
// layer rather than overwrite it.
NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool)
}
// Options 配置 Server。
@ -44,11 +51,24 @@ type Options struct {
// ObfuscatedTCP 先按 MTProto TCP obfuscation 解包，再自动探测 codec。
// Telegram Desktop 的 tcpo_only endpoint 会走这个 64 字节前缀流程。
ObfuscatedTCP bool
// WebSocket 在同一个 listener 上接受 MTProto over WebSocket(/apiws*)。
// 开启后仅在连接建立时读取前 4 字节做 HTTP/TCP 分流；MTProto TCP
// 后续仍走原 ObfuscatedTCP + codec 热路径。
WebSocket bool
// WebSocketAllowedOrigins 是允许浏览器发起 WebSocket upgrade 的页面 origin。
// 空列表表示只接受无 Origin 的非浏览器客户端；"*" 表示允许所有来源（仅调试）。
WebSocketAllowedOrigins []string
// ReadTimeout 单次读取超时。默认 5m。
ReadTimeout time.Duration
// HandshakeIdleTimeout 是连接「建立 session 前」（握手 + 首个加密消息之前）的读超时，
// 比 ReadTimeout 短，用于快速回收握手后静默的半开 / 异常连接。默认 60s。
HandshakeIdleTimeout time.Duration
// HandshakeMaxDuration 是单次密钥交换（serverExchange）的总时长上界。HandshakeIdleTimeout
// 只约束「单次读 idle」，对一个持续发包的客户端无效——若客户端陷入「收到 ResPQ→nonce 失步
// →重发 req_pq」的握手重启死循环（见 docs/client-compat-notes.md），无界的 serverExchange 会
// 对每个 req_pq 盲回 ResPQ、永不收敛地空转刷日志/占 CPU。本上界给整个握手设总预算，超时即
// 放弃并断开，客户端重连发起全新握手（无残留相位差）即恢复。正常握手 <1s。默认 20s。
HandshakeMaxDuration time.Duration
// WriteTimeout 单次写入超时。默认 30s。
WriteTimeout time.Duration
// RPCMaxInflight 是单连接同时处理的 RPC 上限。默认 32。
@ -88,6 +108,9 @@ func (o *Options) setDefaults() {
if o.HandshakeIdleTimeout == 0 {
o.HandshakeIdleTimeout = 60 * time.Second
}
if o.HandshakeMaxDuration == 0 {
o.HandshakeMaxDuration = 20 * time.Second
}
if o.WriteTimeout == 0 {
o.WriteTimeout = 30 * time.Second
}
@ -129,8 +152,11 @@ 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
@ -148,8 +174,7 @@ type Server struct {
rand io.Reader
types *tmap.Map
// sessionUID 是本进程 server session 唯一标识，写入 new_session_created。
sessionUID int64
rpcResults *rpcResultCache
// onFrame 是测试钩子：收到一帧时回调其字节数；生产为 nil。
onFrame func(n int)
@ -166,8 +191,11 @@ func New(opts Options) *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,
@ -183,7 +211,7 @@ func New(opts Options) *Server {
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
sessionUID: opts.Clock.Now().UnixNano(),
rpcResults: newRPCResultCache(opts.Clock.Now),
}
}
@ -214,49 +242,192 @@ func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt
// Serve 在 ln 上运行 MTProto 连接循环，直到 ctx 取消或发生不可恢复错误。
// ctx 取消时优雅退出：关闭 listener 并等待在途连接处理结束。
func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
if s.websocket {
return s.serveMixed(ctx, ln)
}
return s.serveTCP(ctx, ln)
}
func (s *Server) serveTCP(ctx context.Context, ln net.Listener) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
transportListener := ln
if s.obfuscated {
transportListener = transport.ObfuscatedListener(ln)
}
l := newCompatTransportListener(s.codec, transportListener)
s.log.Info("Serving", zap.String("addr", ln.Addr().String()), zap.Int("dc", s.dc), zap.Bool("obfuscated_tcp", s.obfuscated))
defer s.log.Info("Stopped")
// ctx 取消时关闭 listener，解除 Accept 阻塞。
return s.acceptLoop(ctx, ln, s.obfuscated)
}
func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
// 嗅探(读首 4 字节做 HTTP/TCP 分流)的读超时必须对齐「建立 session 前」的读超时
// s.handshakeTimeout(默认 60s)——与 serveDetectedConn/serveConn 的 pre-session 读
// deadline 完全一致。一条尚未发出首帧的连接正处于「pre-session idle」状态：合法 MTProto
// 客户端(如 DrKLO)会预开「暖」连接、在有请求前并不立即发送 obfuscated2 init。此前用
// minDuration(5s,...) 把嗅探压到 5s，比非 mux 路径激进 12 倍，会把这些暖连接在 5s 误杀，
// 触发客户端 6s 重连风暴并误判「后端不健康」回退到外部 DNS。per-conn goroutine 模型已消解
// slow-loris 接入饥饿，故嗅探用满 handshakeTimeout 是安全的。
mux := newSamePortMux(ln, s.handshakeTimeout)
wsLn, wsHandler := transport.WebsocketListener(ln.Addr())
httpServer := &http.Server{
Handler: websocketRouteHandler(wsHandler, s.websocketOrigins),
ReadHeaderTimeout: minDuration(10*time.Second, s.handshakeTimeout),
BaseContext: func(net.Listener) context.Context {
return ctx
},
}
s.log.Info("Serving",
zap.String("addr", ln.Addr().String()),
zap.Int("dc", s.dc),
zap.Bool("obfuscated_tcp", s.obfuscated),
zap.Bool("websocket", true),
zap.Strings("websocket_origins", s.websocketOrigins),
)
defer s.log.Info("Stopped")
go func() {
<-ctx.Done()
_ = l.Close()
_ = mux.Close()
_ = httpServer.Close()
_ = wsLn.Close()
}()
errCh := make(chan error, 4)
var wg sync.WaitGroup
wg.Add(4)
// 分流器：窥探前 4 字节把 HTTP(WebSocket 升级) 与裸 MTProto TCP 拆开。
go func() {
defer wg.Done()
errCh <- mux.Serve(ctx)
}()
// 裸 MTProto TCP：每条连接在自己的 goroutine 里完成去混淆 + codec 探测。
go func() {
defer wg.Done()
errCh <- s.acceptLoop(ctx, mux.TCP(), s.obfuscated)
}()
// WebSocket：gotd 升级处理器已剥离 obfuscated2 并补回 codec tag，这里只需探测 codec。
go func() {
defer wg.Done()
errCh <- s.acceptLoop(ctx, wsLn, false)
}()
go func() {
defer wg.Done()
if err := httpServer.Serve(mux.HTTP()); err != nil {
if errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
errCh <- nil
return
}
errCh <- fmt.Errorf("websocket http serve: %w", err)
return
}
errCh <- nil
}()
var firstErr error
for i := 0; i < 4; i++ {
if err := <-errCh; err != nil && firstErr == nil {
firstErr = err
cancel()
}
}
cancel()
_ = mux.Close()
_ = httpServer.Close()
_ = wsLn.Close()
wg.Wait()
return firstErr
}
// acceptLoop 接受裸连接，并为每条连接单独起 goroutine 完成「去混淆 + codec 探测 +
// serveConn」。探测在 accept 循环之外、带握手超时进行——慢/半开/坏 init 的客户端只占用
// 自己的 goroutine，绝不阻塞其他连接的接入；单条连接的握手失败也只关闭该连接，不会拖垮
// 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆（裸 MTProto TCP）；WebSocket
// 连接传 false（gotd 升级处理器已完成去混淆）。
func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error {
go func() {
<-ctx.Done()
_ = ln.Close()
}()
var wg sync.WaitGroup
defer wg.Wait()
for {
conn, err := l.Accept()
raw, err := ln.Accept()
if err != nil {
if ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
return nil
}
if s.obfuscated && isClientDisconnect(err) {
s.log.Debug("Ignoring failed obfuscated accept", zap.Error(err))
continue
}
return fmt.Errorf("accept: %w", err)
}
wg.Add(1)
go func() {
defer wg.Done()
if err := s.serveConn(ctx, conn); err != nil && !isClientDisconnect(err) {
s.log.Info("Connection closed with error", zap.Error(err))
}
s.serveDetectedConn(ctx, raw, obfuscated)
}()
}
}
// serveDetectedConn 把一条裸连接提升为 transport.Conn（去混淆 + codec 探测）后运行 MTProto
// 连接循环。提升过程的读取放在本 goroutine、且受握手读超时约束，而非塞在 accept 循环里，
// 这样慢连接不会阻塞其他连接接入，去混淆/codec 握手本身也有时间上界。
func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated bool) {
// 握手读超时只覆盖去混淆 + codec 探测这一小段；用真实墙钟时间（SetReadDeadline 语义），
// 不走可能被测试注入的逻辑 clock。
if err := raw.SetReadDeadline(time.Now().Add(s.handshakeTimeout)); err != nil {
_ = raw.Close()
return
}
// 探测阶段若 ctx 取消，主动关闭 raw 解除阻塞读取——否则去混淆读会一直挂到握手超时，
// 把半开连接拖进优雅退出的等待里。探测结束即停掉该 watcher，连接服务期由 serveConn
// 自己的 ctx watcher 接管。
promoted := make(chan struct{})
go func() {
select {
case <-ctx.Done():
_ = raw.Close()
case <-promoted:
}
}()
conn, err := s.promoteConn(raw, obfuscated)
close(promoted)
if err != nil {
// 去混淆/codec 探测失败（读超时、客户端中途断开、坏 init 等）只影响这一条连接，
// 记 debug 即可。
if !isClientDisconnect(err) {
s.log.Debug("Transport handshake failed", zap.Error(err))
}
_ = raw.Close()
return
}
// 探测完成，撤掉握手读超时；后续每帧读写由 serveConn / 传输层各自管理超时。
if err := raw.SetReadDeadline(time.Time{}); err != nil {
_ = conn.Close()
return
}
if err := s.serveConn(ctx, conn); err != nil && !isClientDisconnect(err) {
s.log.Info("Connection closed with error", zap.Error(err))
}
}
// promoteConn 复用与 listener 组合完全一致的「obfuscated2 去混淆 + codec 探测」管线，但针对
// 单条连接，使其可在 accept 循环之外执行。obfuscated 对 WebSocket 连接必须为 false（gotd
// 升级处理器已剥离 obfuscated2 并补回 codec tag）。
func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, error) {
var ln net.Listener = newSingleConnListener(raw)
if obfuscated {
ln = transport.ObfuscatedListener(ln)
}
return newCompatTransportListener(s.codec, ln).Accept()
}
// serveConn 处理单个传输连接：读帧并按 auth_key_id 分流。
//
// - auth_key_id == 0：未加密的密钥交换起始消息，执行握手并落地 auth key。
@ -322,18 +493,26 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
continue
}
data, found, err := s.authKeys.Get(ctx, authKeyID)
if err != nil {
return fmt.Errorf("lookup auth key: %w", err)
}
if !found {
if err := s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound); err != nil {
return err
// 已建立连接复用缓存密钥走快路径（fetchedKey=nil）：避开每帧回查 AuthKeyStore——
// 这是 mtprotoedge 层最热的库访问点。密钥材料创建后不可变；销毁(destroy_auth_key)/
// 撤销由 SessionManager 主动 Close 连接保证失效，不依赖此被动回查。仅 destroy_auth_key
// 的发起连接置 keyDestroyed，使其下一帧回落到 Get→AuthKeyNotFound，维持原契约。
var fetchedKey *store.AuthKeyData
if current == nil || current.authKeyID != authKeyID || current.keyDestroyed.Load() {
d, found, err := s.authKeys.Get(ctx, authKeyID)
if err != nil {
return fmt.Errorf("lookup auth key: %w", err)
}
continue
if !found {
if err := s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound); err != nil {
return err
}
continue
}
fetchedKey = &d
}
current, err = s.handleEncrypted(ctx, conn, cs, current, data, &b)
current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b)
if err != nil {
return err
}

400
internal/mtprotoedge/server_test.go
View file

@ -1,9 +1,16 @@
package mtprotoedge
import (
"bufio"
"bytes"
"context"
"crypto/rand"
"encoding/base64"
"errors"
"fmt"
"net"
"net/http"
"strings"
"testing"
"time"
@ -13,6 +20,7 @@ import (
"github.com/gotd/td/mtproxy"
"github.com/gotd/td/mtproxy/obfuscator"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/transport"
)
@ -223,3 +231,395 @@ func TestServerAcceptObfuscatedAbridgedQuickAckFrame(t *testing.T) {
t.Fatal("server did not stop after ctx cancel")
}
}
func TestServerSamePortWebSocketAndObfuscatedTCP(t *testing.T) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
frames := make(chan int, 2)
srv := New(Options{Logger: zaptest.NewLogger(t), ObfuscatedTCP: true, WebSocket: true})
srv.onFrame = func(n int) {
select {
case frames <- n:
default:
}
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
var wsPayload bin.Buffer
wsPayload.PutInt32(0x11223344)
wsPayload.PutInt32(0x55667788)
wsResolver := dcs.Websocket(dcs.WebsocketOptions{})
wsConn, err := wsResolver.Primary(context.Background(), 2, dcs.List{
Domains: map[int]string{
2: "ws://" + ln.Addr().String() + "/apiws",
},
})
if err != nil {
t.Fatalf("websocket dial: %v", err)
}
sendCtx, sc := context.WithTimeout(context.Background(), 5*time.Second)
if err := wsConn.Send(sendCtx, &wsPayload); err != nil {
sc()
t.Fatalf("websocket send: %v", err)
}
sc()
expectFrameLen(t, frames, wsPayload.Len())
_ = wsConn.Close()
raw, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("tcp dial: %v", err)
}
obfs := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfs.Handshake((codec.Abridged{}).ObfuscatedTag(), 2, mtproxy.Secret{}); err != nil {
t.Fatalf("tcp obfuscated handshake: %v", err)
}
tcpConn, err := transport.NewProtocol(func() transport.Codec {
return transport.Abridged.CodecNoHeader()
}).Handshake(obfs)
if err != nil {
t.Fatalf("tcp transport handshake: %v", err)
}
var tcpPayload bin.Buffer
tcpPayload.PutInt32(0x12345678)
tcpPayload.PutInt32(0x0badf00d)
sendCtx, sc = context.WithTimeout(context.Background(), 5*time.Second)
if err := tcpConn.Send(sendCtx, &tcpPayload); err != nil {
sc()
t.Fatalf("tcp send: %v", err)
}
sc()
expectFrameLen(t, frames, tcpPayload.Len())
_ = tcpConn.Close()
cancel()
select {
case err := <-serveErr:
if err != nil {
t.Fatalf("serve returned error: %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("server did not stop after ctx cancel")
}
}
func TestSamePortWebSocketTransportRoundTrip(t *testing.T) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
mux := newSamePortMux(ln, 5*time.Second)
wsLn, wsHandler := transport.WebsocketListener(ln.Addr())
httpServer := &http.Server{
Handler: websocketRouteHandler(wsHandler, []string{"http://localhost:1234"}),
ReadHeaderTimeout: 5 * time.Second,
}
serveErr := make(chan error, 2)
go func() { serveErr <- mux.Serve(ctx) }()
go func() {
err := httpServer.Serve(mux.HTTP())
if err != nil && !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, net.ErrClosed) {
serveErr <- err
return
}
serveErr <- nil
}()
defer func() {
cancel()
_ = mux.Close()
_ = httpServer.Close()
_ = wsLn.Close()
for i := 0; i < 2; i++ {
select {
case err := <-serveErr:
if err != nil {
t.Fatalf("serve returned error: %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("same-port websocket transport did not stop")
}
}
}()
serverDone := make(chan error, 1)
go func() {
l := newCompatTransportListener(nil, wsLn)
defer func() { _ = l.Close() }()
conn, err := l.Accept()
if err != nil {
serverDone <- err
return
}
defer func() { _ = conn.Close() }()
recvCtx, rc := context.WithTimeout(ctx, 5*time.Second)
defer rc()
var got bin.Buffer
if err := conn.Recv(recvCtx, &got); err != nil {
serverDone <- err
return
}
var reply bin.Buffer
reply.PutInt32(0x10203040)
reply.PutInt32(0x50607080)
sendCtx, sc := context.WithTimeout(ctx, 5*time.Second)
defer sc()
if err := conn.Send(sendCtx, &reply); err != nil {
serverDone <- err
return
}
serverDone <- nil
}()
wsResolver := dcs.Websocket(dcs.WebsocketOptions{})
wsConn, err := wsResolver.Primary(context.Background(), 2, dcs.List{
Domains: map[int]string{
2: "ws://" + ln.Addr().String() + "/apiws",
},
})
if err != nil {
t.Fatalf("websocket dial: %v", err)
}
defer func() { _ = wsConn.Close() }()
var request bin.Buffer
request.PutInt32(0x11223344)
request.PutInt32(0x55667788)
sendCtx, sc := context.WithTimeout(context.Background(), 5*time.Second)
if err := wsConn.Send(sendCtx, &request); err != nil {
sc()
t.Fatalf("websocket send: %v", err)
}
sc()
var want bin.Buffer
want.PutInt32(0x10203040)
want.PutInt32(0x50607080)
recvCtx, rc := context.WithTimeout(context.Background(), 5*time.Second)
var got bin.Buffer
if err := wsConn.Recv(recvCtx, &got); err != nil {
rc()
t.Fatalf("websocket recv: %v", err)
}
rc()
if !bytes.Equal(got.Raw(), want.Raw()) {
t.Fatalf("websocket recv = %x, want %x", got.Raw(), want.Raw())
}
select {
case err := <-serverDone:
if err != nil {
t.Fatalf("server transport: %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("server transport did not finish")
}
}
// TestSamePortMuxIdleConnNotReapedBeforeHandshakeTimeout 回归：WebSocket 同端口复用的嗅探
// 读超时（读首 4 字节做 HTTP/TCP 分流）必须对齐 HandshakeIdleTimeout，而不是旧的硬上限 5s。
// 合法 MTProto 客户端（DrKLO）会预开「暖」连接、在有请求前并不立即发 obfuscated2 init；旧实现
// 用 minDuration(5s, handshakeTimeout) 把嗅探压到 5s，比非 mux 路径（serveDetectedConn 用满
// handshakeTimeout）激进 12 倍，会把这些暖连接在 5s 误杀，触发 DrKLO 6s 重连风暴 + EPOLLRDHUP
// + 误判后端不健康回退外部 DNS（见 docs/client-compat-notes.md）。
//
// HandshakeIdleTimeout 必须 >5s 才能暴露旧的截断：连一条裸 TCP、不发任何字节，本地用 6s 读
// deadline。新实现下连接在 8s 嗅探超时前一直存活，故本地 Read 因自身 deadline 超时（net timeout）；
// 旧实现下服务端 5s 即 FIN，本地 Read 会在 6s 前拿到 EOF/reset —— 据此判定回归。
func TestSamePortMuxIdleConnNotReapedBeforeHandshakeTimeout(t *testing.T) {
addr, _, _ := startTestServer(t, Options{
WebSocket: true,
ObfuscatedTCP: true,
HandshakeIdleTimeout: 8 * time.Second, // 必须 >5s 才能区分「对齐 handshakeTimeout」与旧的 5s 截断
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer func() { _ = raw.Close() }()
// 不发任何字节，模拟客户端预开、暂未发首帧的暖连接。
if err := raw.SetReadDeadline(time.Now().Add(6 * time.Second)); err != nil {
t.Fatalf("set read deadline: %v", err)
}
n, err := raw.Read(make([]byte, 1))
if err == nil {
t.Fatalf("unexpected %d bytes on idle pre-handshake conn (server should send nothing)", n)
}
// 只有「本地读 deadline 超时」才说明连接在 6s 时仍存活（嗅探超时已对齐 8s）。
// 任何由对端关闭导致的 EOF/reset 都意味着服务端在 handshakeTimeout 前过早回收了连接。
var nerr net.Error
if errors.As(err, &nerr) && nerr.Timeout() {
return
}
t.Fatalf("server closed idle pre-handshake conn before handshake idle timeout (err=%v); "+
"same-port mux sniff timeout must align with HandshakeIdleTimeout, not a 5s cap", err)
}
func expectFrameLen(t *testing.T, frames <-chan int, want int) {
t.Helper()
select {
case n := <-frames:
if n != want {
t.Fatalf("received frame len = %d, want %d", n, want)
}
case <-time.After(5 * time.Second):
t.Fatal("server did not receive frame in time")
}
}
// TestWebSocketRouteHandlerChecksBrowserOrigin 回归保护 Origin 修复：浏览器发起的 WS 升级
// 必带 Origin（≠ Host），白名单来源需要改写 Origin 通过 coder/websocket，同名单外来源必须 403。
func TestWebSocketRouteHandlerChecksBrowserOrigin(t *testing.T) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
defer func() { _ = ln.Close() }()
_, wsHandler := transport.WebsocketListener(ln.Addr())
httpServer := &http.Server{Handler: websocketRouteHandler(wsHandler, []string{"http://localhost:1234"})}
go func() { _ = httpServer.Serve(ln) }()
defer func() { _ = httpServer.Close() }()
host := ln.Addr().String()
// 白名单跨源升级：Origin 指向页面来源（端口/主机不同于 Host）。
status := wsUpgradeStatus(t, host, "/apiws", "http://localhost:1234")
if !strings.Contains(status, "101") {
t.Fatalf("allowed cross-origin /apiws upgrade: got status %q, want 101 Switching Protocols", status)
}
status = wsUpgradeStatus(t, host, "/apiws", "http://evil.example")
if !strings.Contains(status, "403") {
t.Fatalf("disallowed origin: got status %q, want 403", status)
}
// 非白名单路径必须 404，不得升级。
status = wsUpgradeStatus(t, host, "/nope", "http://localhost:1234")
if !strings.Contains(status, "404") {
t.Fatalf("disallowed path: got status %q, want 404", status)
}
}
// wsUpgradeStatus 用裸连接发一个合法的 WebSocket 升级请求并返回状态行。
func wsUpgradeStatus(t *testing.T, host, path, origin string) string {
t.Helper()
conn, err := net.Dial("tcp", host)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer func() { _ = conn.Close() }()
var keyBytes [16]byte
if _, err := rand.Read(keyBytes[:]); err != nil {
t.Fatalf("rand: %v", err)
}
key := base64.StdEncoding.EncodeToString(keyBytes[:])
req := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"+
"Sec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Protocol: binary\r\nOrigin: %s\r\n\r\n",
path, host, key, origin)
if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
t.Fatalf("set deadline: %v", err)
}
if _, err := conn.Write([]byte(req)); err != nil {
t.Fatalf("write upgrade: %v", err)
}
statusLine, err := bufio.NewReader(conn).ReadString('\n')
if err != nil {
t.Fatalf("read status: %v", err)
}
return statusLine
}
// TestServerObfuscatedTCPNotBlockedByStalledClient 回归保护「去 worker 池 / 握手移出 accept
// 循环」的修复：一个发了几字节就挂起的连接，过去会卡死串行 accept 循环、阻塞所有后续接入；
// 现在它只占用自己的 goroutine，正常客户端仍能在握手超时内被服务。
func TestServerObfuscatedTCPNotBlockedByStalledClient(t *testing.T) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
frames := make(chan int, 1)
srv := New(Options{Logger: zaptest.NewLogger(t), ObfuscatedTCP: true, WebSocket: true})
srv.onFrame = func(n int) {
select {
case frames <- n:
default:
}
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
// 挂起连接：发 4 个非 HTTP 字节通过分流（路由到 TCP），但不补满 obfuscated2 的 64 字节
// init，使其卡在去混淆读取上。修复前这会拖死整个 TCP accept 循环。
stalled, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("stalled dial: %v", err)
}
defer func() { _ = stalled.Close() }()
if _, err := stalled.Write([]byte{0x01, 0x02, 0x03, 0x04}); err != nil {
t.Fatalf("stalled write: %v", err)
}
// 正常的 obfuscated abridged 客户端：应当照常被服务（onFrame 触发）。
raw, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("tcp dial: %v", err)
}
obfs := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfs.Handshake((codec.Abridged{}).ObfuscatedTag(), 2, mtproxy.Secret{}); err != nil {
t.Fatalf("tcp obfuscated handshake: %v", err)
}
tcpConn, err := transport.NewProtocol(func() transport.Codec {
return transport.Abridged.CodecNoHeader()
}).Handshake(obfs)
if err != nil {
t.Fatalf("tcp transport handshake: %v", err)
}
var payload bin.Buffer
payload.PutInt32(0x12345678)
payload.PutInt32(0x0badf00d)
sendCtx, sc := context.WithTimeout(context.Background(), 5*time.Second)
if err := tcpConn.Send(sendCtx, &payload); err != nil {
sc()
t.Fatalf("tcp send: %v", err)
}
sc()
expectFrameLen(t, frames, payload.Len())
_ = tcpConn.Close()
cancel()
select {
case err := <-serveErr:
if err != nil {
t.Fatalf("serve returned error: %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("server did not stop after ctx cancel")
}
}

686
internal/mtprotoedge/session_manager.go
View file

@ -21,11 +21,34 @@ var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys
const (
maxPendingPushesPerSession = 32
// maxFlushAttempts / flushRetryBackoff：排空暂存推送时 c.Send 失败（出站拥塞 5s 超时
// 或瞬时错误）后的退避重试上界。连接真死时 serveConn 会 Unregister 清理状态、提前止损；
// 这里只为「连接存活但出站暂时拥塞」做有限重试。用尽仍失败则置位激活并接受 getDifference
// 兜底——避免 idle 客户端（只发 ping、不触发置位重试）永久停在未激活态而静默断流。
maxFlushAttempts = 5
flushRetryBackoff = 2 * time.Second
// pendingPushMaxAge：session 注册后迟迟不调 updates.getState（receivesUpdates 恒 false）时，
// 其暂存的主动推送最长保留时长。超过即丢整批并不再囤——正常 TDesktop 登录后秒级就会
// getState 建立同步基线；长期不 ready 多为异常/对抗连接。丢弃不丢消息：getDifference 以
// user_update_events durable log 兜底补齐。
// getState 建立同步基线；长期不 ready 多为异常/对抗连接。
//
// 不变量：只有 durable update（写 user_update_events）才会进 pending。transient update
// （typing/presence，不写 durable log）经 PushToUserTransient* 在未就绪时直接跳过、不入队，
// 因此本队列被老化/溢出/重试耗尽丢弃时，丢的一定是 durable 条目——getDifference 以
// user_update_events 兜底补齐，丢弃不丢数据。
pendingPushMaxAge = 60 * time.Second
// maxSessionsPerAuthKey：单个 raw auth_key 允许同时在线的 session 上限。telesrv 单 DC，
// 一个客户端的全部连接（主连接 + 并发下载/上传）共享同一 auth_key、各用独立 session_id，
// 故此上限须高于真实客户端单设备的并发连接峰值，否则会误杀活跃下载/主连接：
// - TDesktop：kMaxMediaDcCount=0x10，单 DC 最多 16 路下载 + 16 路上传 + 1 主 ≈ 33；
// - DrKLO：DOWNLOAD_CONNECTIONS_COUNT=2 + UPLOAD_CONNECTIONS_COUNT=4 + 主/push ≈ 10。
// 叠加重连 churn（旧 session 在 readTimeout 内滞留）峰值约 ~70，故设 256（~3.5x 余量）。
// 它只防「单 auth_key 累积海量连接」的病态（使 CloseSessionsForRawAuthKey/pushToUser 遍历
// 退化 O(N)），超限驱逐的也只是同一设备凭据自身的连接，不会误伤别的账号。
maxSessionsPerAuthKey = 256
// maxChannelIndexPerSession：单 session 在 channel 路由索引（interest / membership）中
// 登记的 channel 数上限。membership 源于真实成员关系（大账号可能很多），interest 受客户端
// 直接控制；两者都设一个宽松上界防内存放大，超出即截断并记日志。
maxChannelIndexPerSession = 8192
)
type queuedPush struct {
@ -53,12 +76,14 @@ type SessionManager struct {
bySession map[sessionKey]*Conn
bySessionID map[int64]map[[8]byte]*Conn // sessionID → raw authKeyID → Conn，用于兼容旧 API 的唯一性检查
byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn
byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn
byUser map[int64]map[sessionKey]*Conn
byChannel map[int64]map[sessionKey]int64 // channelID → session → userID，用于频道 active-viewer 临时推送
bySessionChannels map[sessionKey]map[int64]struct{}
byMemberChannel map[int64]map[sessionKey]int64 // channelID → session → userID，用于已上线成员持久 update 推送
bySessionMembers map[sessionKey]map[int64]struct{}
pending map[sessionKey][]queuedPush // updates-ready 前暂存的主动推送
flushing map[sessionKey]bool // 置位时暂存正在排空的 session；排空完成前推送继续进 pending 保序
lifecycle SessionLifecycleObserver
log *zap.Logger
@ -73,12 +98,14 @@ func NewSessionManager(log *zap.Logger) *SessionManager {
bySession: make(map[sessionKey]*Conn),
bySessionID: make(map[int64]map[[8]byte]*Conn),
byAuthKey: make(map[[8]byte]map[int64]*Conn),
byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn),
byUser: make(map[int64]map[sessionKey]*Conn),
byChannel: make(map[int64]map[sessionKey]int64),
bySessionChannels: make(map[sessionKey]map[int64]struct{}),
byMemberChannel: make(map[int64]map[sessionKey]int64),
bySessionMembers: make(map[sessionKey]map[int64]struct{}),
pending: make(map[sessionKey][]queuedPush),
flushing: make(map[sessionKey]bool),
log: log,
}
}
@ -96,13 +123,30 @@ func (m *SessionManager) Register(c *Conn) {
key := connSessionKey(c)
var replaced *Conn
var evicted *Conn
if old, ok := m.bySession[key]; ok && old != c {
replaced = old
m.removeLocked(old, false)
} else if existing := m.byAuthKey[c.authKeyID]; len(existing) >= maxSessionsPerAuthKey {
// 同 raw auth_key 的 session 数达上限且本次是新 session：驱逐一个现有 session 让位，
// 防对抗客户端用海量 session_id 撑爆索引。驱逐对象与新连接同属一个设备凭据，
// 触顶基本是该凭据自身异常。被驱逐连接的 serveConn 会在下一帧因 actor 已关而退出。
for _, ec := range existing {
evicted = ec
m.removeLocked(ec, true)
break
}
m.log.Debug("Evicted oldest session for auth key at cap",
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int("cap", maxSessionsPerAuthKey),
)
}
m.bySession[key] = c
addSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID, c)
addConnIndex(m.byAuthKey, c.authKeyID, c.sessionID, c)
if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved {
addBusinessAuthKeyIndex(m.byBusinessAuthKey, businessAuthKeyID, key, c)
}
if uid := c.userID.Load(); uid != 0 {
c.userIDResolved.Store(true)
addUserIndex(m.byUser, uid, key, c)
@ -117,9 +161,14 @@ func (m *SessionManager) Register(c *Conn) {
if replaced != nil {
replaced.Close()
}
if evicted != nil {
evicted.Close()
}
}
// Unregister 注销一个连接（仅当它仍是当前注册的同一对象，避免误删重连后的新连接）。
// 观察者对未登录连接（userID=0）也回调：业务层据此清理按 session 维度的缓存条目，
// 否则未登录连接的元数据只能等容量上限驱逐。
func (m *SessionManager) Unregister(c *Conn) {
m.mu.Lock()
var (
@ -131,8 +180,8 @@ func (m *SessionManager) Unregister(c *Conn) {
offlineUser = m.removeLocked(c, true)
if offlineUser != 0 {
lastForUser = len(m.byUser[offlineUser]) == 0
observer = m.lifecycle
}
observer = m.lifecycle
m.log.Debug("Session unregistered",
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int64("session_id", c.sessionID),
@ -140,7 +189,7 @@ func (m *SessionManager) Unregister(c *Conn) {
)
}
m.mu.Unlock()
if observer != nil && offlineUser != 0 {
if observer != nil {
observer.SessionOffline(c.authKeyID, c.sessionID, offlineUser, lastForUser)
}
}
@ -232,6 +281,11 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
if old != userID {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.membershipsSynced.Store(false)
// 身份变化即丢弃暂存推送：它们属于前一个账号，flush 给新账号是跨账号泄露。
// 同时取消进行中的排空（runFlush 还另有 owner 校验做批内兜底）。
delete(m.pending, key)
delete(m.flushing, key)
}
}
c.userIDResolved.Store(true)
@ -240,6 +294,9 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
} else {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.membershipsSynced.Store(false)
delete(m.pending, key)
delete(m.flushing, key)
}
}
@ -327,13 +384,20 @@ func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8
oldAuthKeyID, resolved := c.BusinessAuthKeyID()
changed := !resolved || oldAuthKeyID != authKeyID
oldUserID := c.userID.Load()
if resolved {
removeBusinessAuthKeyIndex(m.byBusinessAuthKey, oldAuthKeyID, key)
}
c.SetBusinessAuthKeyID(authKeyID)
addBusinessAuthKeyIndex(m.byBusinessAuthKey, authKeyID, key, c)
if changed {
if oldUserID != 0 {
removeUserIndex(m.byUser, oldUserID, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.membershipsSynced.Store(false)
delete(m.pending, key)
delete(m.flushing, key)
c.userID.Store(0)
c.userIDResolved.Store(false)
}
@ -362,12 +426,87 @@ func (m *SessionManager) AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int
return c.BusinessAuthKeyID()
}
// CloseSessionsForBusinessAuthKey 强制断开指定业务 auth_key 的全部活跃连接，
// 供授权撤销（被踢设备）使用：出站推送用连接持有的密钥加密、不回查密钥库，
// 不断开的话被撤销的设备会继续收到推送直至自然断线；perm-key 连接的授权
// 缓存也只有断开重连才会重新回查授权表。这里必须关闭底层 transport，
// 让 WebSocket/TCP 对端马上看到断线，而不是只从在线索引摘除。
func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int {
type offlineEvent struct {
key sessionKey
userID int64
last bool
}
m.mu.Lock()
var conns []*Conn
var events []offlineEvent
for key, c := range m.businessAuthKeyCandidatesLocked(authKeyID) {
if !connUsesBusinessAuthKey(c, authKeyID) {
continue
}
uid := m.removeLocked(c, true)
conns = append(conns, c)
events = append(events, offlineEvent{key: key, userID: uid, last: uid != 0 && len(m.byUser[uid]) == 0})
}
observer := m.lifecycle
if len(conns) > 0 {
m.log.Debug("Force close sessions for revoked auth key",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int("closed", len(conns)),
)
}
m.mu.Unlock()
for _, c := range conns {
c.ForceClose()
}
if observer != nil {
for _, e := range events {
observer.SessionOffline(e.key.authKeyID, e.key.sessionID, e.userID, e.last)
}
}
return len(conns)
}
// CloseSessionsForRawAuthKeyExcept 强制断开指定 raw auth_key 的活跃连接，可排除
// 一个 session（destroy_auth_key 的发起连接：响应要先送达，它的密钥已删，下一帧
// 自然失效）。出站推送不回查密钥库，必须主动断开底层 transport 才能让销毁立即生效。
func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exceptSessionID int64) int {
type offlineEvent struct {
key sessionKey
userID int64
last bool
}
m.mu.Lock()
var conns []*Conn
var events []offlineEvent
for sessionID, c := range m.byAuthKey[authKeyID] {
if sessionID == exceptSessionID {
continue
}
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
uid := m.removeLocked(c, true)
conns = append(conns, c)
events = append(events, offlineEvent{key: key, userID: uid, last: uid != 0 && len(m.byUser[uid]) == 0})
}
observer := m.lifecycle
m.mu.Unlock()
for _, c := range conns {
c.ForceClose()
}
if observer != nil {
for _, e := range events {
observer.SessionOffline(e.key.authKeyID, e.key.sessionID, e.userID, e.last)
}
}
return len(conns)
}
// UnbindAuthKey 清理某业务 auth_key 下所有活跃连接的登录用户缓存。
func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
m.mu.Lock()
defer m.mu.Unlock()
count := 0
for key, c := range m.bySession {
for key, c := range m.businessAuthKeyCandidatesLocked(authKeyID) {
if !connUsesBusinessAuthKey(c, authKeyID) {
continue
}
@ -376,6 +515,10 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.membershipsSynced.Store(false)
// 授权解除后暂存推送属于已登出的账号，不能等下一个登录者置位时 flush 出去。
delete(m.pending, key)
delete(m.flushing, key)
c.userIDResolved.Store(true)
count++
}
@ -396,19 +539,146 @@ func (m *SessionManager) SetReceivesUpdates(sessionID int64, receives bool) {
m.mu.Unlock()
return
}
c.receivesUpdates.Store(receives)
if !receives {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
}
pending := m.takePendingLocked(key, receives)
owner, start := m.setReceivesUpdatesLocked(c, key, receives)
m.mu.Unlock()
if len(pending) > 0 {
go m.flushPending(key, pending)
if start {
go m.runFlush(c, key, owner, 0)
}
}
// setReceivesUpdatesLocked 是置位/复位的共同内核，调用方须持有 m.mu。
// 置位且有暂存时不立即置 receivesUpdates：标记 flushing 并返回该批暂存所属的 userID，
// 交由 runFlush 排空后原子置位，期间新到推送继续进 pending，保证暂存与实时推送的
// 相对顺序（否则实时直发可能先于更早 pts 的暂存条目落线）。返回的 owner 让 runFlush
// 能识别排空期间的身份切换（登出/换号），丢弃属于旧账号的剩余暂存而不发给新账号。
func (m *SessionManager) setReceivesUpdatesLocked(c *Conn, key sessionKey, receives bool) (int64, bool) {
if !receives {
c.receivesUpdates.Store(false)
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.membershipsSynced.Store(false)
// 取消进行中的排空激活：runFlush 在置位前会复查该标志，标志已删则放弃置位，
// 避免把刚置 false 的开关翻回 true。
delete(m.flushing, key)
return 0, false
}
if c.receivesUpdates.Load() || m.flushing[key] {
// 已就绪，或已有排空协程在跑（完成时会自行取走新增暂存并置位）。
return 0, false
}
if len(m.pending[key]) == 0 {
c.receivesUpdates.Store(true)
return 0, false
}
m.flushing[key] = true
return c.userID.Load(), true
}
// runFlush 把暂存推送按序直发到连接，排空（含排空期间新增）后才置位 receivesUpdates。
// 直发用 c.Send 绕过 ready 检查——此刻必然未就绪，走 PushToSessionForAuthKey 会被
// 重新暂存形成死循环。三类终止：
// - 身份切换（登出/换号致 c.userID != owner）：丢弃剩余暂存与回排数据，不发给新账号；
// - 发送失败：回排剩余并退避重试，attempt 用尽则置位激活、靠 getDifference 兜底，
// 避免 idle 客户端永久停在未激活态；
// - 排空完毕：原子置位 receivesUpdates。
func (m *SessionManager) runFlush(c *Conn, key sessionKey, owner int64, attempt int) {
for {
m.mu.Lock()
if cur, ok := m.bySession[key]; !ok || cur != c || !m.flushing[key] {
// 连接已换代（removeLocked 已清 flushing）或激活被取消（SetReceivesUpdates(false)）。
m.mu.Unlock()
return
}
if c.userID.Load() != owner {
// 排空期间发生登出/换号：剩余暂存属于旧账号，丢弃且不得发给新账号。
delete(m.pending, key)
delete(m.flushing, key)
m.mu.Unlock()
return
}
batch := m.takePendingLocked(key, true)
if len(batch) == 0 {
c.receivesUpdates.Store(true)
delete(m.flushing, key)
m.mu.Unlock()
return
}
m.mu.Unlock()
for i, item := range batch {
// 每条发送前复查身份：登出/换号后 batch 的剩余条目不能继续发到已易主的连接。
if c.userID.Load() != owner {
m.mu.Lock()
delete(m.pending, key)
delete(m.flushing, key)
m.mu.Unlock()
return
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
err := c.Send(ctx, item.t, item.msg)
cancel()
if err == nil {
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)
delete(m.flushing, key)
}
m.mu.Unlock()
return
}
rest := append(append([]queuedPush(nil), batch[i:]...), m.pending[key]...)
if len(rest) > maxPendingPushesPerSession {
// 与 queueLocked 溢出策略一致：丢最旧留最新，让 pts 空洞集中在最前端，
// flush 首条即触发客户端 gap 检测，恢复路径最短。
rest = rest[len(rest)-maxPendingPushesPerSession:]
}
m.pending[key] = rest
if attempt+1 >= maxFlushAttempts {
// 重试用尽：置位激活避免 idle 客户端永久断流；剩余暂存中的 durable 更新
// 由客户端后续 pts 空洞触发 getDifference 补齐。
c.receivesUpdates.Store(true)
delete(m.flushing, key)
m.mu.Unlock()
m.log.Debug("Flush gave up after retries; activated with getDifference fallback",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Int("requeued", len(rest)),
)
return
}
m.mu.Unlock()
m.log.Debug("Flush pending push failed; backoff retry",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Int("attempt", attempt+1),
zap.Int("requeued", len(rest)),
zap.Error(err),
)
time.AfterFunc(flushRetryBackoff*time.Duration(attempt+1), func() {
m.runFlush(c, key, owner, attempt+1)
})
return
}
// 本批发完，循环回去 re-take 排空期间新增的暂存。
}
}
// ReceivesUpdatesForAuthKey 报告指定 raw auth_key_id + session_id 的连接是否已完全就绪：
// 既接收主动 updates，channel membership 推送路由也已成功建立。无活跃连接时返回 false。
// 返回 false 会让按 RPC 置位的短路放行，下一条 RPC 重试 membership 同步——
// 否则同步失败的 session 会以「已置位但 byMemberChannel 缺失」的状态静默漏收超级群推送。
func (m *SessionManager) ReceivesUpdatesForAuthKey(authKeyID [8]byte, sessionID int64) bool {
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
return ok && c.receivesUpdates.Load() && c.membershipsSynced.Load()
}
// SetReceivesUpdatesForAuthKey 标记指定 raw auth_key_id + session_id 是否接收主动 updates。
func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, sessionID int64, receives bool) {
m.mu.Lock()
@ -418,16 +688,11 @@ func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, session
m.mu.Unlock()
return
}
c.receivesUpdates.Store(receives)
if !receives {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
}
pending := m.takePendingLocked(key, receives)
owner, start := m.setReceivesUpdatesLocked(c, key, receives)
m.mu.Unlock()
if len(pending) > 0 {
go m.flushPending(key, pending)
if start {
go m.runFlush(c, key, owner, 0)
}
}
@ -470,6 +735,22 @@ func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID
return c.Send(ctx, t, msg)
}
// PushToSessionForAuthKeyImmediate 向指定 raw auth_key_id + session_id 立即推送一条消息。
//
// 它不等待该 session 进入 updates-ready，也不写 pending 队列。仅用于登录前的握手信号
// （例如 updateLoginToken）：这类消息本身就是让客户端继续完成登录的触发器，若走普通
// durable update 队列会卡在客户端尚未调用 updates.getState 的阶段。
func (m *SessionManager) PushToSessionForAuthKeyImmediate(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.RLock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
m.mu.RUnlock()
if !ok {
return ErrSessionNotFound
}
return c.SendBestEffort(ctx, t, msg, 2*time.Second)
}
// PushToUser 向某 user 所有活跃连接推送，返回已发送或已暂存的连接数。
// 发送在释放锁后进行，避免持锁阻塞于网络 IO。
func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
@ -487,9 +768,14 @@ func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, use
return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg)
}
func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
// PushToUserAuthKey 把 msg 定向投递给【绑定到 businessAuthKeyID 这台具体设备】且属于
// userID 的就绪连接（密聊设备级投递的锚点）。索引走 byBusinessAuthKey（经
// businessAuthKeyCandidatesLocked，兼容 temp-key/PFS 连接），不是 byAuthKey（raw 索引会
// 漏 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.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, false, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
@ -501,6 +787,88 @@ func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAu
})
}
// 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 {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
if err != nil {
return err
}
return c.SendBestEffortEncoded(ctx, t, encoded, timeout)
})
}
func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, transient bool, send func(*Conn) error) (int, error) {
m.mu.Lock()
candidates := m.businessAuthKeyCandidatesLocked(businessAuthKeyID)
conns := make([]*Conn, 0, len(candidates))
for _, c := range candidates {
if c.userID.Load() != userID {
continue
}
if !c.receivesUpdates.Load() {
// 未就绪：密聊消息靠 getDifference 补，typing 直接丢——都不进 pending。
continue
}
conns = append(conns, c)
}
m.mu.Unlock()
_ = transient
var firstErr error
sent := 0
for _, c := range conns {
// 锁外发送前复查身份，防收集后并发换绑导致跨账号泄露。
if c.userID.Load() != userID {
continue
}
if err := send(c); err != nil {
if firstErr == nil {
firstErr = err
}
continue
}
sent++
}
return sent, firstErr
}
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)
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, 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)
})
}
// PushToUserTransientExceptAuthKeySession 推送 transient（短命、不写 durable log）update，
// 如 typing / presence。与普通推送的关键区别：session 未就绪（receivesUpdates=false）时直接
// 跳过该连接、不进 pending——transient 数据 getDifference 无法补，就绪后由 getState 快照 /
// 下一次状态变化重建，囤积过期 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)
return m.pushToUserWithSender(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, false, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
if err != nil {
return err
}
return c.SendBestEffortEncoded(ctx, t, encoded, timeout)
})
}
func (m *SessionManager) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, nil, excludeSessionID, t, msg, timeout)
}
@ -511,7 +879,7 @@ 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)
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
@ -536,17 +904,41 @@ func onceEncodedOutbound(msg bin.Encoder) func() (*encodedOutboundMessage, error
}
}
func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, send func(*Conn) error) (int, error) {
func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, queueWhenNotReady bool, send func(*Conn) error) (int, error) {
m.mu.Lock()
conns := make([]*Conn, 0, len(m.byUser[userID]))
total := len(m.byUser[userID])
conns := make([]*Conn, 0, total)
queued := 0
dropped := 0
excluded := 0
skipped := 0
for key, c := range m.byUser[userID] {
if shouldExcludeSession(c, excludeAuthKeyID, excludeSessionID) {
excluded++
continue
}
if !c.receivesUpdates.Load() {
m.queueLocked(key, t, msg)
queued++
if !queueWhenNotReady {
// transient（typing/presence）：未就绪即丢，不进 pending。这些 update 不写
// durable log，getDifference 无法补；就绪后由 getState 快照/下次状态变化重建。
skipped++
continue
}
if m.queueLocked(key, t, msg) {
queued++
m.log.Debug("Push queued (session not updates-ready)",
zap.Int64("user_id", userID),
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
)
} else {
dropped++
m.log.Debug("Push dropped (stale pending; durable log covers)",
zap.Int64("user_id", userID),
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
)
}
continue
}
conns = append(conns, c)
@ -556,13 +948,43 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
var firstErr error
sent := 0
for _, c := range conns {
// 锁外发送前复查身份：收集 conns 到此刻之间，连接可能被并发换绑（登出/换号，
// bindUserLocked 的 c.userID.Swap）。不复查会把本属于 userID 的 update 投递到
// 已易主的连接，构成跨账号泄露。与 AddUserChannelMembership 的同款防御一致。
if c.userID.Load() != userID {
continue
}
if err := send(c); err != nil {
if firstErr == nil {
firstErr = err
}
m.log.Debug("Push to conn failed",
zap.Int64("user_id", userID),
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int64("session_id", c.sessionID),
zap.Error(err),
)
continue
}
sent++
m.log.Debug("Push to conn ok",
zap.Int64("user_id", userID),
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int64("session_id", c.sessionID),
)
}
if total == 0 {
m.log.Debug("Push to user: no active conns", zap.Int64("user_id", userID))
} else if excluded > 0 || queued > 0 || dropped > 0 || skipped > 0 || sent < len(conns) {
m.log.Debug("Push to user summary",
zap.Int64("user_id", userID),
zap.Int("conns", total),
zap.Int("sent", sent),
zap.Int("queued", queued),
zap.Int("dropped", dropped),
zap.Int("skipped_transient", skipped),
zap.Int("excluded", excluded),
)
}
return sent + queued, firstErr
}
@ -574,31 +996,6 @@ func (m *SessionManager) Online() int {
return len(m.bySession)
}
// OnlineUserIDs returns a bounded snapshot of users that currently have active
// sessions. Callers still need to verify business visibility before pushing.
func (m *SessionManager) OnlineUserIDs(limit int) []int64 {
m.mu.RLock()
defer m.mu.RUnlock()
if len(m.byUser) == 0 {
return nil
}
capHint := len(m.byUser)
if limit > 0 && capHint > limit {
capHint = limit
}
ids := make([]int64, 0, capHint)
for userID, conns := range m.byUser {
if userID == 0 || len(conns) == 0 {
continue
}
ids = append(ids, userID)
if limit > 0 && len(ids) >= limit {
break
}
}
return ids
}
// IsUserOnline returns whether userID has at least one active connection.
func (m *SessionManager) IsUserOnline(userID int64) bool {
if userID == 0 {
@ -693,10 +1090,12 @@ func (m *SessionManager) SetSessionChannelMemberships(rawAuthKeyID [8]byte, sess
return
}
m.clearChannelMembershipsLocked(key)
c.membershipsSynced.Store(false)
if userID == 0 || c.userID.Load() != userID {
return
}
m.trackChannelIndexLocked(m.byMemberChannel, m.bySessionMembers, key, userID, channelIDs)
c.membershipsSynced.Store(true)
}
// AddUserChannelMembership adds channelID to every live session for userID.
@ -735,6 +1134,47 @@ func (m *SessionManager) OnlineChannelMemberUserIDs(channelID int64, limit int)
return m.onlineChannelUsers(m.byMemberChannel, channelID, limit)
}
// OnlineChannelMemberUserIDsExcluding 返回频道在线成员中不在 exclude 集合内的 user id，
// 用于 >cap 在线成员的 UpdateChannelTooLong nudge（P0-8）：完整 payload 已投递给 exclude
// 集合（cap 内成员），其余在线成员只发廉价 nudge 促其 getChannelDifference。单次 RLock 快照；
// 由调用方用「已收完整 payload 的 recipients」构造 exclude，使同一 user 不会既收 payload 又收
// nudge——天然规避两次独立 cap 调用的边界双投/漏投（设计 §8-D3/D32）。limit 防一次无界 nudge 风暴。
// 不做 PG active 复核：byMemberChannel 已在 join/leave/kick 维护；nudge 廉价且幂等，对刚离开成员
// 的多余 nudge 无害（其 getChannelDifference 自带访问校验）。
func (m *SessionManager) OnlineChannelMemberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64 {
if channelID == 0 {
return nil
}
m.mu.RLock()
defer m.mu.RUnlock()
sessions := m.byMemberChannel[channelID]
if len(sessions) == 0 {
return nil
}
out := make([]int64, 0)
seen := make(map[int64]struct{}, len(sessions))
for key, userID := range sessions {
if userID == 0 {
continue
}
if _, ok := exclude[userID]; ok {
continue
}
if _, ok := m.bySession[key]; !ok {
continue
}
if _, ok := seen[userID]; ok {
continue
}
seen[userID] = struct{}{}
out = append(out, userID)
if limit > 0 && len(out) >= limit {
break
}
}
return out
}
func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64, channelID int64, limit int) []int64 {
if channelID == 0 {
return nil
@ -771,6 +1211,9 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
delete(m.bySession, key)
removeSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID)
removeConnIndex(m.byAuthKey, c.authKeyID, c.sessionID)
if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved {
removeBusinessAuthKeyIndex(m.byBusinessAuthKey, businessAuthKeyID, key)
}
uid := c.userID.Load()
if uid != 0 {
removeUserIndex(m.byUser, uid, key)
@ -780,9 +1223,26 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
if dropPending {
delete(m.pending, key)
}
delete(m.flushing, key)
return uid
}
func (m *SessionManager) businessAuthKeyCandidatesLocked(authKeyID [8]byte) map[sessionKey]*Conn {
out := make(map[sessionKey]*Conn, len(m.byBusinessAuthKey[authKeyID])+len(m.byAuthKey[authKeyID]))
for key, c := range m.byBusinessAuthKey[authKeyID] {
if cur := m.bySession[key]; cur == c {
out[key] = c
}
}
for sessionID, c := range m.byAuthKey[authKeyID] {
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
if cur := m.bySession[key]; cur == c {
out[key] = c
}
}
return out
}
func (m *SessionManager) clearChannelInterestsLocked(key sessionKey) {
m.clearChannelIndexLocked(m.byChannel, m.bySessionChannels, key)
}
@ -797,10 +1257,17 @@ func (m *SessionManager) trackChannelIndexLocked(index map[int64]map[sessionKey]
channels = make(map[int64]struct{}, len(channelIDs))
reverse[key] = channels
}
truncated := 0
for _, channelID := range channelIDs {
if channelID == 0 {
continue
}
if _, exists := channels[channelID]; !exists && len(channels) >= maxChannelIndexPerSession {
// 达 per-session 上限：丢弃多出的 channel 登记（仅影响该 channel 的实时/成员
// 推送路由，durable update 仍由 getDifference/getChannelDifference 兜底）。
truncated++
continue
}
channels[channelID] = struct{}{}
sessions := index[channelID]
if sessions == nil {
@ -809,6 +1276,14 @@ func (m *SessionManager) trackChannelIndexLocked(index map[int64]map[sessionKey]
}
sessions[key] = userID
}
if truncated > 0 {
m.log.Warn("Channel index truncated for session at per-session cap",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Int("cap", maxChannelIndexPerSession),
zap.Int("truncated", truncated),
)
}
}
func (m *SessionManager) clearChannelIndexLocked(index map[int64]map[sessionKey]int64, reverse map[sessionKey]map[int64]struct{}, key sessionKey) {
@ -858,12 +1333,34 @@ func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedP
if !ready || len(m.pending[key]) == 0 {
return nil
}
pending := append([]queuedPush(nil), m.pending[key]...)
q := m.pending[key]
delete(m.pending, key)
// 取出时过滤超龄条目：暂存只为弥合「注册到就绪」的窗口，迟迟未就绪期间
// 囤下的过时 update（含 transient 类）不应在多分钟后原样下发；durable 事件
// 由 user_update_events + getDifference 兜底，丢弃不丢数据。
now := time.Now()
pending := make([]queuedPush, 0, len(q))
dropped := 0
for _, item := range q {
if now.Sub(item.at) > pendingPushMaxAge {
dropped++
continue
}
pending = append(pending, item)
}
if dropped > 0 {
m.log.Debug("Drop stale pending pushes on take",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Int("dropped", dropped),
)
}
return pending
}
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) {
// queueLocked 暂存一条主动推送，返回是否实际入队——stale 丢批分支会连同当前
// 这条一起丢弃，调用方据此区分 queued/dropped 计数，避免投递日志失真。
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool {
q := m.pending[key]
// 过期保护：最早一条暂存已超过 pendingPushMaxAge（session 迟迟未 ready）时，丢整批并
// 不再囤这条，记 trace。避免「登录后从不 getState」的连接长期占用 pending 内存。
@ -874,31 +1371,17 @@ func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bi
zap.Int("dropped", len(q)),
)
delete(m.pending, key)
return
return false
}
push := queuedPush{t: t, msg: msg, at: time.Now()}
if len(q) >= maxPendingPushesPerSession {
copy(q, q[1:])
q[len(q)-1] = push
m.pending[key] = q
return
return true
}
m.pending[key] = append(q, push)
}
func (m *SessionManager) flushPending(key sessionKey, pending []queuedPush) {
for _, item := range pending {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
err := m.PushToSessionForAuthKey(ctx, key.authKeyID, key.sessionID, item.t, item.msg)
cancel()
if err != nil {
m.log.Debug("Flush pending push failed",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Error(err),
)
}
}
return true
}
func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey, bool, bool) {
@ -923,6 +1406,47 @@ func (m *SessionManager) dropPendingBySessionLocked(sessionID int64) {
}
}
// RunPendingSweeper 周期回收长期滞留的 pending 暂存：被动老化（queueLocked/takePendingLocked）
// 只在「有新推送」或「就绪后取出」时触发，对「已注册但迟迟不调 getState、又恰好没有新推送、
// 也不断连（持续 ping 保活）」的连接无法回收其超龄 pending。本 sweeper 给出一个主动兜底，
// 与 pendingPushMaxAge 阈值一致，仅丢整批超龄、不触碰正在排空（flushing）的 session。
func (m *SessionManager) RunPendingSweeper(ctx context.Context, interval time.Duration) {
if interval <= 0 {
interval = time.Minute
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
}
m.sweepStalePending()
}
}
func (m *SessionManager) sweepStalePending() {
now := time.Now()
m.mu.Lock()
defer m.mu.Unlock()
dropped := 0
for key, q := range m.pending {
if m.flushing[key] {
// 排空协程拥有该批，回收交给 runFlush，避免与其竞态。
continue
}
if len(q) == 0 || now.Sub(q[0].at) <= pendingPushMaxAge {
continue
}
delete(m.pending, key)
dropped++
}
if dropped > 0 {
m.log.Debug("Swept stale pending sessions", zap.Int("dropped_sessions", dropped))
}
}
func addConnIndex[K comparable](idx map[K]map[int64]*Conn, key K, sessionID int64, c *Conn) {
set := idx[key]
if set == nil {
@ -941,6 +1465,24 @@ func removeConnIndex[K comparable](idx map[K]map[int64]*Conn, key K, sessionID i
}
}
func addBusinessAuthKeyIndex(idx map[[8]byte]map[sessionKey]*Conn, authKeyID [8]byte, key sessionKey, c *Conn) {
set := idx[authKeyID]
if set == nil {
set = make(map[sessionKey]*Conn)
idx[authKeyID] = set
}
set[key] = c
}
func removeBusinessAuthKeyIndex(idx map[[8]byte]map[sessionKey]*Conn, authKeyID [8]byte, key sessionKey) {
if set := idx[authKeyID]; set != nil {
delete(set, key)
if len(set) == 0 {
delete(idx, authKeyID)
}
}
}
func addSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte, c *Conn) {
set := idx[sessionID]
if set == nil {

126
internal/mtprotoedge/session_manager_test.go
View file

@ -23,6 +23,23 @@ func (e *countingOutboundEncoder) Encode(b *bin.Buffer) error {
return (&tg.UpdatesTooLong{}).Encode(b)
}
type closeCountingTransport struct {
closes int
}
func (t *closeCountingTransport) Send(context.Context, *bin.Buffer) error {
return errors.New("test transport send")
}
func (t *closeCountingTransport) Recv(context.Context, *bin.Buffer) error {
return errors.New("test transport recv")
}
func (t *closeCountingTransport) Close() error {
t.closes++
return nil
}
// TestSessionManagerRegistry 验证注册表的注册/注销/查找语义（不涉及网络发送）。
func TestSessionManagerRegistry(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
@ -146,6 +163,81 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
}
}
func TestSessionManagerCloseSessionsForBusinessAuthKeyClosesBoundTempAndRaw(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
rawTemp := [8]byte{1}
perm := [8]byte{9}
otherRaw := [8]byte{2}
otherPerm := [8]byte{8}
tempTransport := &closeCountingTransport{}
permTransport := &closeCountingTransport{}
otherTransport := &closeCountingTransport{}
cTemp := &Conn{sessionID: 11, authKeyID: rawTemp, transport: tempTransport}
cPerm := &Conn{sessionID: 12, authKeyID: perm, transport: permTransport}
cOther := &Conn{sessionID: 13, authKeyID: otherRaw}
cOther.transport = otherTransport
sm.Register(cTemp)
sm.Register(cPerm)
sm.Register(cOther)
sm.BindAuthKeyForSession(rawTemp, 11, perm)
sm.BindAuthKeyForSession(perm, 12, perm)
sm.BindAuthKeyForSession(otherRaw, 13, otherPerm)
sm.BindUserForAuthKey(rawTemp, 11, 100)
sm.BindUserForAuthKey(perm, 12, 100)
sm.BindUserForAuthKey(otherRaw, 13, 200)
if closed := sm.CloseSessionsForBusinessAuthKey(perm); closed != 2 {
t.Fatalf("closed sessions = %d, want 2", closed)
}
if tempTransport.closes != 1 || permTransport.closes != 1 {
t.Fatalf("transport closes temp=%d perm=%d, want 1/1", tempTransport.closes, permTransport.closes)
}
if otherTransport.closes != 0 {
t.Fatalf("other transport closes = %d, want 0", otherTransport.closes)
}
if got := sm.Online(); got != 1 {
t.Fatalf("online after close = %d, want 1", got)
}
if _, ok := sm.AuthKeyIDForSession(rawTemp, 11); ok {
t.Fatal("temp session still indexed after business auth key close")
}
if _, ok := sm.AuthKeyIDForSession(perm, 12); ok {
t.Fatal("raw perm session still indexed after business auth key close")
}
if userID, ok := sm.UserIDForAuthKey(otherRaw, 13); !ok || userID != 200 {
t.Fatalf("other session user = %d ok %v, want 200/true", userID, ok)
}
if closed := sm.CloseSessionsForBusinessAuthKey(perm); closed != 0 {
t.Fatalf("second close = %d, want 0", closed)
}
}
func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1}
oldPerm := [8]byte{7}
newPerm := [8]byte{8}
c := &Conn{sessionID: 21, authKeyID: raw}
sm.Register(c)
sm.BindAuthKeyForSession(raw, 21, oldPerm)
sm.BindAuthKeyForSession(raw, 21, newPerm)
if closed := sm.CloseSessionsForBusinessAuthKey(oldPerm); closed != 0 {
t.Fatalf("close old business auth key = %d, want 0", closed)
}
if got := sm.Online(); got != 1 {
t.Fatalf("online after closing old key = %d, want 1", got)
}
if closed := sm.CloseSessionsForBusinessAuthKey(newPerm); closed != 1 {
t.Fatalf("close new business auth key = %d, want 1", closed)
}
if got := sm.Online(); got != 0 {
t.Fatalf("online after closing new key = %d, want 0", got)
}
}
func TestSessionManagerChannelInterestIndex(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
@ -257,6 +349,40 @@ func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.
assertCleared("after unbind auth key")
}
func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
c := &Conn{
sessionID: 42,
authKeyID: raw,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
sm.Register(c)
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
if err := sm.PushToSessionForAuthKeyImmediate(context.Background(), raw, 42, proto.MessageFromServer, msg); err != nil {
t.Fatalf("immediate push: %v", err)
}
select {
case op := <-c.outbound:
if op.msg != msg {
t.Fatalf("enqueued msg = %T, want original update", op.msg)
}
case <-time.After(time.Second):
t.Fatal("immediate push was not enqueued")
}
sm.mu.RLock()
pending := len(sm.pending[sessionKey{authKeyID: raw, sessionID: 42}])
sm.mu.RUnlock()
if pending != 0 {
t.Fatalf("pending pushes = %d, want 0", pending)
}
}
// TestSessionManagerPush 验证主动推送端到端：两个 client 连接握手并建立 session 后，
// server 经 PushToSession / PushToUser 主动向其推送，client 收到。
func TestSessionManagerPush(t *testing.T) {

49
internal/mtprotoedge/session_ready_test.go Normal file
View file

@ -0,0 +1,49 @@
package mtprotoedge
import (
"testing"
"go.uber.org/zap/zaptest"
)
// TestReceivesUpdatesForAuthKeyRequiresMembershipSync 验证「完全就绪」查询同时要求
// receivesUpdates 与 channel membership 路由建立成功。membership 同步失败时该查询
// 必须返回 false，让按 RPC 置位的短路放行重试——否则该 session 会以「已置位但
// byMemberChannel 缺失」的状态静默漏收超级群推送，且 channel 维度没有 pts 缺口
// 信号可供客户端自愈。
func TestReceivesUpdatesForAuthKeyRequiresMembershipSync(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
c := &Conn{sessionID: 42, authKeyID: raw}
sm.Register(c)
sm.BindUserForAuthKey(raw, 42, 100)
sm.SetReceivesUpdatesForAuthKey(raw, 42, true)
if sm.ReceivesUpdatesForAuthKey(raw, 42) {
t.Fatal("ready before membership sync — a failed sync would never be retried")
}
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{7})
if !sm.ReceivesUpdatesForAuthKey(raw, 42) {
t.Fatal("not ready after successful membership sync")
}
// 没有任何频道的账号：空列表的成功同步同样算就绪。
sm.SetSessionChannelMemberships(raw, 42, 100, nil)
if !sm.ReceivesUpdatesForAuthKey(raw, 42) {
t.Fatal("not ready after successful empty membership sync")
}
// userID 与连接当前绑定不一致（换号竞态）时不算就绪，等正确身份重试。
sm.SetSessionChannelMemberships(raw, 42, 999, []int64{7})
if sm.ReceivesUpdatesForAuthKey(raw, 42) {
t.Fatal("ready after membership sync for mismatched user")
}
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{7})
// 登出清除就绪标志。
sm.BindUserForAuthKey(raw, 42, 0)
if sm.ReceivesUpdatesForAuthKey(raw, 42) {
t.Fatal("still ready after user unbind")
}
}

53
internal/mtprotoedge/transient_push_test.go Normal file
View file

@ -0,0 +1,53 @@
package mtprotoedge
import (
"context"
"testing"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestPushTransientSkipsNotReadySession 锁定不变量：transient 推送（typing/presence）对
// 未就绪 session 直接跳过、不进 pending；而普通 durable 推送会进 pending。回归 transient
// updates 与 durable 共用 pending 队列、被老化/溢出/重试耗尽误丢且 getDifference 无法补的问题。
func TestPushTransientSkipsNotReadySession(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(100)
c := &Conn{
sessionID: 7,
authKeyID: [8]byte{7},
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outboundStop: make(chan struct{}),
}
c.userID.Store(userID)
c.userIDResolved.Store(true)
// receivesUpdates 保持 false：session 未就绪（尚未 getState 建立同步基线）。
sm.Register(c)
key := connSessionKey(c)
// transient：未就绪 → 跳过、不入队。
if _, err := sm.PushToUserTransientExceptAuthKeySession(context.Background(), userID, [8]byte{}, 0, proto.MessageFromServer, &tg.UpdatesTooLong{}, 0); err != nil {
t.Fatalf("transient push: %v", err)
}
sm.mu.RLock()
n := len(sm.pending[key])
sm.mu.RUnlock()
if n != 0 {
t.Fatalf("transient push queued %d pending, want 0 (must skip not-ready session)", n)
}
// durable（普通）：未就绪 → 入 pending（就绪后排空，丢弃时由 getDifference 兜底）。
if _, err := sm.PushToUserExceptSession(context.Background(), userID, 0, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("durable push: %v", err)
}
sm.mu.RLock()
n = len(sm.pending[key])
sm.mu.RUnlock()
if n != 1 {
t.Fatalf("durable push queued %d pending, want 1", n)
}
}

38
internal/mtprotoedge/transport_compat.go
View file

@ -45,6 +45,38 @@ func newCompatTransportListener(codec func() transport.Codec, listener net.Liste
return &compatTransportListener{listener: listener}
}
// singleConnListener 是一个只产出一条「已接受」连接、随后阻塞到关闭的 net.Listener。
// 它让单条裸连接可以走 listener 形态的去混淆/codec 管线（ObfuscatedListener +
// compatTransportListener），从而把这部分阻塞读取从 accept 循环挪到每连接 goroutine。
type singleConnListener struct {
addr net.Addr
ch chan net.Conn
once sync.Once
}
func newSingleConnListener(c net.Conn) *singleConnListener {
ch := make(chan net.Conn, 1)
ch <- c
return &singleConnListener{addr: c.LocalAddr(), ch: ch}
}
func (l *singleConnListener) Accept() (net.Conn, error) {
c, ok := <-l.ch
if !ok {
return nil, net.ErrClosed
}
return c, nil
}
func (l *singleConnListener) Close() error {
l.once.Do(func() { close(l.ch) })
return nil
}
func (l *singleConnListener) Addr() net.Addr {
return l.addr
}
func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
conn, err := l.listener.Accept()
if err != nil {
@ -409,9 +441,9 @@ func validateOutgoingCompatMessage(b *bin.Buffer) error {
}
func writeCompatPacket(w io.Writer, header, payload []byte) error {
packet := make([]byte, len(header)+len(payload))
copy(packet, header)
copy(packet[len(header):], payload)
packet := make([]byte, 0, len(header)+len(payload))
packet = append(packet, header...)
packet = append(packet, payload...)
return writeAll(w, packet)
}