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

Initial open source release

Browse source
This commit is contained in:
A 2026-06-04 01:37:39 +08:00
commit 74992e893f
377 changed files with 118084 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

114
internal/mtprotoedge/conn.go Normal file
View file

@ -0,0 +1,114 @@
package mtprotoedge
import (
"context"
"sync"
"sync/atomic"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
)
// Conn 是一个已识别 session 的客户端连接，持有向其加密发送消息所需的全部上下文。
// 由 SessionManager 管理，供请求响应与主动 push 共用。
//
// Send 并发安全：所有出站消息先进 per-Conn outbound actor，由它串行分配 msg_id/seq_no、
// 加密并写 transport，避免高并发 RPC 响应与 push 交错造成 MTProto 顺序错误。
type outboundWriter interface {
Send(context.Context, *bin.Buffer) error
}
type Conn struct {
transport transport.Conn
writer outboundWriter
cipher crypto.Cipher
msgID *proto.MessageIDGen
writeTimeout time.Duration
metrics Metrics
authKeyID [8]byte
sessionID int64
salt int64
key crypto.AuthKey
outbound chan outboundOp
outboundControl chan outboundOp
outboundStop chan struct{}
outboundDone chan struct{}
outboundClose sync.Once
rpcQueue chan inboundRPC
rpcStop chan struct{}
rpcCancel context.CancelFunc
rpcClose sync.Once
rpcWG sync.WaitGroup
rpcTimeout time.Duration
// inflightRPCBytes 跟踪已入队未完成的 inbound RPC body 总字节，配合 maxInflightRPCBytes
// 给 RPC 队列设字节预算（不止限条数），防对抗客户端发大请求撑内存。
inflightRPCBytes atomic.Int64
// RPC worker 懒启动：首个 RPC 入队时才起 worker（ensureInboundRPCWorkers），
// 避免握手后静默 / 纯推送目标连接白白钉住 rpcMaxInflight 个 goroutine。
rpcRootCtx context.Context
rpcMaxInflight int
rpcWorkersOnce sync.Once
// sentContentMessages 只由 outbound actor 访问，用于生成 MTProto seq_no。
sentContentMessages int32
identityMu sync.RWMutex
businessAuthKeyID [8]byte
businessAuthKeyResolved bool
userID atomic.Int64
userIDResolved atomic.Bool
receivesUpdates atomic.Bool
}
// AuthKeyID 返回连接的 auth_key_id。
func (c *Conn) AuthKeyID() [8]byte { return c.authKeyID }
// BusinessAuthKeyID 返回业务视角的 auth_key_id。
//
// temp auth_key 绑定后解析为 perm auth_key；第二个返回值表示本连接是否已完成解析，
// 即便解析结果等于原始 auth_key_id 也会返回 true，以避免每个 RPC 重复查绑定表。
func (c *Conn) BusinessAuthKeyID() ([8]byte, bool) {
c.identityMu.RLock()
defer c.identityMu.RUnlock()
return c.businessAuthKeyID, c.businessAuthKeyResolved
}
// SetBusinessAuthKeyID 缓存业务视角 auth_key_id。
func (c *Conn) SetBusinessAuthKeyID(id [8]byte) {
c.identityMu.Lock()
changed := !c.businessAuthKeyResolved || c.businessAuthKeyID != id
c.businessAuthKeyID = id
c.businessAuthKeyResolved = true
c.identityMu.Unlock()
if changed {
c.userID.Store(0)
c.userIDResolved.Store(false)
}
}
// SessionID 返回连接的 session_id。
func (c *Conn) SessionID() int64 { return c.sessionID }
// UserID 返回绑定的用户 id；未登录为 0。
func (c *Conn) UserID() int64 { return c.userID.Load() }
// UserIDResolved 返回 user_id 授权状态是否已为当前连接解析过。
//
// resolved=true 且 userID=0 表示该 auth_key 当前未登录；这样登录前的多次 RPC
// 不会反复查询授权表，后续登录成功会由 BindUser 覆盖为真实用户。
func (c *Conn) UserIDResolved() (userID int64, resolved bool) {
return c.userID.Load(), c.userIDResolved.Load()
}
// ReceivesUpdates 报告该连接是否接收主动推送的 updates。
func (c *Conn) ReceivesUpdates() bool { return c.receivesUpdates.Load() }
// SetReceivesUpdates 设置该连接是否接收主动推送的 updates。
// 登录后的主连接在 updates.getState/getDifference 建立同步基线后置为 true。
func (c *Conn) SetReceivesUpdates(v bool) { c.receivesUpdates.Store(v) }

33
internal/mtprotoedge/destroy_auth_key.go Normal file
View file

@ -0,0 +1,33 @@
package mtprotoedge
import (
"fmt"
"github.com/gotd/td/bin"
)
const (
destroyAuthKeyRequestTypeID = 0xd1435160
destroyAuthKeyOkTypeID = 0xf660e1d4
)
type destroyAuthKeyRequest struct{}
func (*destroyAuthKeyRequest) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyRequestTypeID)
return nil
}
func (*destroyAuthKeyRequest) Decode(b *bin.Buffer) error {
if err := b.ConsumeID(destroyAuthKeyRequestTypeID); err != nil {
return fmt.Errorf("decode destroy_auth_key: %w", err)
}
return nil
}
type destroyAuthKeyOk struct{}
func (*destroyAuthKeyOk) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyOkTypeID)
return nil
}

6
internal/mtprotoedge/doc.go Normal file
View file

@ -0,0 +1,6 @@
// Package mtprotoedge 是 MTProto 连接层：TCP/WS listener、密钥交换、auth key 查找与持久化、
// 消息加解密、session/server-salt/msg-id/ack/container/gzip，以及 invokeWithLayer、initConnection、
// invokeWithoutUpdates 等 wrapper 的 unwrap。
//
// 它只把 MTProto 世界转换成「已解密、已识别 session 的 RPC 请求」，不得包含业务逻辑。
package mtprotoedge

90
internal/mtprotoedge/e2e_test.go Normal file
View file

@ -0,0 +1,90 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"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/auth"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
// TestTelegramClientEndToEnd 是连接层的最强端到端验证：用 gotd/td 的完整
// telegram.Client（而非底层 cipher）连本地 mtprotoedge，client 自动经
// invokeWithLayer(initConnection(help.getConfig)) 完成初始化，并取得含本地 DC 的 Config。
func TestTelegramClientEndToEnd(t *testing.T) {
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()
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), "12345"),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
}
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)
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
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: zaptest.NewLogger(t).Named("client"),
SessionStorage: &session.StorageMemory{},
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
client := telegram.NewClient(1, "hash", opts)
if err := client.Run(ctx, func(ctx context.Context) error {
cfg, err := tg.NewClient(client).HelpGetConfig(ctx)
if err != nil {
return err
}
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")
}
return nil
}); err != nil {
t.Fatalf("telegram client run: %v", err)
}
cancel()
if err := <-serveErr; err != nil {
t.Errorf("serve: %v", err)
}
}

702
internal/mtprotoedge/encrypted.go Normal file
View file

@ -0,0 +1,702 @@
package mtprotoedge
import (
"context"
"encoding/hex"
"errors"
"fmt"
"math"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tgerr"
"github.com/gotd/td/transport"
"telesrv/internal/store"
)
// connState 是单连接的 MTProto 运行态。
type connState struct {
sentCreated bool
seen map[int64]clientMsgRecord // 已处理的 client msg_id，用于幂等和 msgs_state_req
order []int64
minSeen int64
maxSeen int64
}
type clientMsgRecord struct {
state byte
seqNo int32
content bool
}
func newConnState() *connState {
return &connState{
seen: make(map[int64]clientMsgRecord),
minSeen: math.MaxInt64,
}
}
const (
maxTrackedClientMsgIDs = 400
msgStateUnknown byte = 1
msgStateNotReceived byte = 2
msgStateNotReceivedHigh byte = 3
msgStateReceived byte = 4
badMsgIDTooLow = 16
badMsgIDTooHigh = 17
badMsgIDInvalidBits = 18
badMsgSeqTooLow = 32
badMsgSeqTooHigh = 33
badMsgSeqNotEven = 34
badMsgSeqNotOdd = 35
badMsgContainer = 64
)
// 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}
data, err := s.cipher.DecryptFromBuffer(key, b)
if err != nil {
return current, fmt.Errorf("decrypt: %w", err)
}
if data.Salt != keyData.ServerSalt {
c := current
temp := false
if c == nil || c.sessionID != data.SessionID {
c = s.newConn(tc, key, data.SessionID, keyData.ServerSalt)
temp = true
}
err := s.sendBadServerSalt(ctx, c, data.MessageID, data.SeqNo, keyData.ServerSalt)
if temp {
c.Close()
}
return current, err
}
// 首个加密消息或 session 变化时（重新）注册连接到 SessionManager。
if current == nil || current.sessionID != data.SessionID {
if current != nil {
s.conns.Unregister(current)
current.Close()
}
current = s.newConn(tc, key, data.SessionID, keyData.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)
}
body := data.Data()
typeID, err := (&bin.Buffer{Buf: body}).PeekID()
if err != nil {
return current, fmt.Errorf("peek encrypted payload type id: %w", err)
}
if code := validateClientEnvelope(s.clock.Now(), data.MessageID, data.SeqNo, typeID); code != 0 {
s.log.Debug("Sending bad_msg_notification",
zap.Int64("msg_id", data.MessageID),
zap.Int32("seq_no", data.SeqNo),
zap.Uint32("type_id", typeID),
zap.Int("code", code),
)
return current, s.sendBadMsg(ctx, current, data.MessageID, data.SeqNo, code)
}
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 {
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
}
return current, s.sendAck(ctx, current, data.MessageID)
}
if code := cs.validateSeq(data.MessageID, data.SeqNo, content); code != 0 {
s.log.Debug("Sending bad_msg_notification",
zap.Int64("msg_id", data.MessageID),
zap.Int32("seq_no", data.SeqNo),
zap.Uint32("type_id", typeID),
zap.Int("code", code),
)
return current, s.sendBadMsg(ctx, current, data.MessageID, data.SeqNo, code)
}
cs.track(data.MessageID, data.SeqNo, content, msgStateReceived)
if !cs.sentCreated {
cs.sentCreated = true
s.log.Debug("Sending new_session_created", zap.Int64("msg_id", data.MessageID), zap.Int32("seq_no", data.SeqNo))
if err := s.sendNewSessionCreated(ctx, current, data.MessageID); err != nil {
return current, err
}
}
var acks []int64
if err := s.dispatch(ctx, cs, current, data.MessageID, data.SeqNo, &bin.Buffer{Buf: body}, &acks); err != nil {
return current, err
}
if len(acks) > 0 {
if err := s.sendAck(ctx, current, acks...); err != nil {
return current, err
}
}
return current, nil
}
// dispatch 处理一条明文消息：解包 container/gzip，处理服务消息，其余转 RPC 路由。
// content-related 消息（ping、RPC）的 msg_id 会收集到 acks 以便统一确认。
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
id, err := b.PeekID()
if err != nil {
return fmt.Errorf("peek type id: %w", err)
}
ackContent := func() {
if clientMessageNeedsAck(id) {
*acks = append(*acks, msgID)
}
}
switch id {
case proto.GZIPTypeID:
var gz proto.GZIP
if err := gz.Decode(b); err != nil {
return fmt.Errorf("decode gzip: %w", err)
}
return s.dispatch(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: gz.Data}, acks)
case proto.MessageContainerTypeID:
var container proto.MessageContainer
if err := container.Decode(b); err != nil {
return fmt.Errorf("decode container: %w", err)
}
if code := validateClientContainer(msgID, seqNo, container); code != 0 {
return s.sendBadMsg(ctx, c, msgID, seqNo, code)
}
for i := range container.Messages {
m := container.Messages[i]
typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID()
if err != nil {
return fmt.Errorf("peek container message type id: %w", err)
}
content := clientMessageNeedsAck(typeID)
if record, ok := cs.seenRecord(m.ID); ok {
if record.content {
*acks = append(*acks, m.ID)
}
continue
}
if code := cs.validateSeq(m.ID, int32(m.SeqNo), content); code != 0 {
return s.sendBadMsg(ctx, c, m.ID, int32(m.SeqNo), code)
}
cs.track(m.ID, int32(m.SeqNo), content, msgStateReceived)
if err := s.dispatch(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks); err != nil {
return err
}
}
return nil
case mt.PingRequestTypeID:
var ping mt.PingRequest
if err := ping.Decode(b); err != nil {
return fmt.Errorf("decode ping: %w", err)
}
ackContent()
return s.sendPong(ctx, c, msgID, ping.PingID)
case mt.PingDelayDisconnectRequestTypeID:
var ping mt.PingDelayDisconnectRequest
if err := ping.Decode(b); err != nil {
return fmt.Errorf("decode ping_delay_disconnect: %w", err)
}
ackContent()
return s.sendPong(ctx, c, msgID, ping.PingID)
case mt.GetFutureSaltsRequestTypeID:
var req mt.GetFutureSaltsRequest
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode get_future_salts: %w", err)
}
ackContent()
return s.sendFutureSalts(ctx, c, msgID, req.Num)
case mt.MsgsAckTypeID:
var ack mt.MsgsAck
if err := ack.Decode(b); err != nil {
return fmt.Errorf("decode msgs_ack: %w", err)
}
c.AckServerMessages(ack.MsgIDs)
s.log.Debug("Received msgs_ack", zap.Int64s("msg_ids", ack.MsgIDs))
return nil
case mt.MsgsStateReqTypeID:
var req mt.MsgsStateReq
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode msgs_state_req: %w", err)
}
ackContent()
outgoing, err := c.OutgoingStateInfo(ctx, req.MsgIDs)
if err != nil {
return err
}
return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs)))
case mt.MsgResendReqTypeID:
var req mt.MsgResendReq
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode msg_resend_req: %w", err)
}
ackContent()
outgoing, err := c.ResendMessages(ctx, req.MsgIDs)
if err != nil {
return err
}
return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs)))
case mt.MsgsStateInfoTypeID:
var info mt.MsgsStateInfo
if err := info.Decode(b); err != nil {
return fmt.Errorf("decode msgs_state_info: %w", err)
}
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", info.ReqMsgID), zap.Int("len", len(info.Info)))
return nil
case mt.MsgsAllInfoTypeID:
var info mt.MsgsAllInfo
if err := info.Decode(b); err != nil {
return fmt.Errorf("decode msgs_all_info: %w", err)
}
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", len(info.MsgIDs)), zap.Int("len", len(info.Info)))
return nil
case mt.DestroySessionRequestTypeID:
var req mt.DestroySessionRequest
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode destroy_session: %w", err)
}
ackContent()
return s.sendDestroySession(ctx, c, req.SessionID)
case mt.HTTPWaitRequestTypeID:
var req mt.HTTPWaitRequest
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode http_wait: %w", err)
}
s.log.Debug("Received http_wait",
zap.Int("max_delay", req.MaxDelay),
zap.Int("wait_after", req.WaitAfter),
zap.Int("max_wait", req.MaxWait),
)
return nil
case mt.RPCDropAnswerRequestTypeID:
var req mt.RPCDropAnswerRequest
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode rpc_drop_answer: %w", err)
}
ackContent()
s.log.Debug("Received rpc_drop_answer", zap.Int64("req_msg_id", req.ReqMsgID))
return s.sendResult(ctx, c, msgID, &mt.RPCAnswerUnknown{})
case destroyAuthKeyRequestTypeID:
var req destroyAuthKeyRequest
if err := req.Decode(b); err != nil {
return err
}
ackContent()
s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])))
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{})
default:
ackContent()
body := b.Copy()
return s.enqueueRPC(ctx, c, msgID, body)
}
}
func mergeStateInfo(primary, fallback []byte) []byte {
if len(primary) == 0 {
return fallback
}
info := make([]byte, len(fallback))
copy(info, fallback)
for i, state := range primary {
if i >= len(info) {
break
}
if state != 0 {
info[i] = state
}
}
return info
}
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, body []byte) error {
id, _ := (&bin.Buffer{Buf: body}).PeekID()
method := s.typeName(id)
err := c.enqueueInboundRPC(ctx, inboundRPC{
method: method,
size: len(body),
run: func(taskCtx context.Context) error {
// body 已是 enqueueRPC 入参的独立副本（dispatch 里 b.Copy()），且每个任务只 run 一次，
// 无需再 append 拷贝；直接复用，省掉一份 inbound 在途内存。
if err := s.handleRPC(taskCtx, c, msgID, &bin.Buffer{Buf: body}); err != nil {
s.log.Info("RPC async handler failed",
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])),
zap.Int64("session_id", c.sessionID),
zap.Error(err),
)
return err
}
return nil
},
})
if errors.Is(err, ErrInboundRPCQueueFull) {
s.log.Debug("Inbound RPC queue full",
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 s.sendResult(ctx, c, msgID, &mt.RPCError{
ErrorCode: 420,
ErrorMessage: "FLOOD_WAIT_1",
})
}
return err
}
// handleRPC 把明文 RPC 请求交给 RPC 路由，并将结果或错误包成 rpc_result 回发。
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buffer) error {
id, _ := b.PeekID()
method := s.typeName(id)
if s.rpc == nil {
s.log.Warn("No RPC handler configured; dropping request", zap.String("method", method))
return nil
}
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)
fields := []zap.Field{
zap.String("method", method),
zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])),
zap.Int64("session_id", c.sessionID),
zap.Int64("msg_id", msgID),
zap.Duration("dur", dur),
}
if businessAuthKeyID, ok := c.BusinessAuthKeyID(); ok {
fields = append(fields, zap.String("business_auth_key_id", hex.EncodeToString(businessAuthKeyID[:])))
}
if userID := c.UserID(); userID != 0 {
fields = append(fields, zap.Int64("user_id", userID))
}
if err != nil {
var rpcErr *tgerr.Error
if errors.As(err, &rpcErr) {
s.log.Info("RPC error", append(fields, zap.Int("code", rpcErr.Code), zap.String("error", rpcErr.Message))...)
return s.sendResult(ctx, c, msgID, &mt.RPCError{
ErrorCode: rpcErr.Code,
ErrorMessage: rpcErr.Message,
})
}
s.log.Info("RPC internal error", append(fields, zap.Error(err))...)
return s.sendResult(ctx, c, msgID, &mt.RPCError{
ErrorCode: 500,
ErrorMessage: "INTERNAL",
})
}
s.log.Info("RPC handled", fields...)
return s.sendResult(ctx, c, msgID, result)
}
// sendResult 把 RPC 结果包成 rpc_result 并加密回发。
func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) error {
var buf bin.Buffer
if err := result.Encode(&buf); err != nil {
return fmt.Errorf("encode rpc result: %w", err)
}
return c.Send(ctx, proto.MessageServerResponse, &proto.Result{
RequestMessageID: reqMsgID,
Result: buf.Raw(),
})
}
// sendPong 回复 mt.PingRequest / mt.PingDelayDisconnectRequest。
func (s *Server) sendPong(ctx context.Context, c *Conn, reqMsgID, pingID int64) error {
return c.SendAsync(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: reqMsgID, PingID: pingID})
}
// sendFutureSalts 回复 MTProto get_future_salts。
//
// 第一阶段只维护当前 auth key 的权威 server_salt，因此返回当前 salt 的有效窗口。
// 后续如引入 salt rotation，可在这里扩展为多条未来 salt。
func (s *Server) sendFutureSalts(ctx context.Context, c *Conn, reqMsgID int64, num int) error {
if num < 0 {
num = 0
}
if num > 1 {
num = 1
}
now := int(s.clock.Now().Unix())
salts := make([]mt.FutureSalt, 0, num)
if num == 1 {
salts = append(salts, mt.FutureSalt{
ValidSince: now - 300,
ValidUntil: now + 24*60*60,
Salt: c.salt,
})
}
return c.SendAsync(ctx, proto.MessageServerResponse, &mt.FutureSalts{
ReqMsgID: reqMsgID,
Now: now,
Salts: salts,
})
}
// sendNewSessionCreated 在连接首个加密消息后通知客户端新 session 已建立。
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,
ServerSalt: c.salt,
})
}
// 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})
}
// sendMsgsStateInfo 回复 msgs_state_req/msg_resend_req。
func (s *Server) sendMsgsStateInfo(ctx context.Context, c *Conn, reqMsgID int64, info []byte) error {
return c.SendAsync(ctx, proto.MessageServerResponse, &mt.MsgsStateInfo{ReqMsgID: reqMsgID, Info: info})
}
func (s *Server) sendDestroySession(ctx context.Context, c *Conn, sessionID int64) error {
removed := false
if sessionID != c.sessionID {
removed = s.conns.DestroySessionForAuthKey(c.authKeyID, sessionID)
if err := s.sessions.Delete(ctx, sessionID); err != nil {
s.log.Debug("Delete session record failed",
zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])),
zap.Int64("session_id", sessionID),
zap.Error(err),
)
}
}
if removed {
return c.Send(ctx, proto.MessageServerResponse, &mt.DestroySessionOk{SessionID: sessionID})
}
return c.Send(ctx, proto.MessageServerResponse, &mt.DestroySessionNone{SessionID: sessionID})
}
// sendBadMsg 通知客户端消息存在协议层错误（msg_id/seqno 非法）。
func (s *Server) sendBadMsg(ctx context.Context, c *Conn, badMsgID int64, badSeqno int32, code int) error {
return c.SendAsync(ctx, proto.MessageFromServer, &mt.BadMsgNotification{
BadMsgID: badMsgID,
BadMsgSeqno: int(badSeqno),
ErrorCode: code,
})
}
// sendBadServerSalt 通知客户端修正 server_salt（error_code 48）。
func (s *Server) sendBadServerSalt(ctx context.Context, c *Conn, badMsgID int64, badSeqno int32, newSalt int64) error {
return c.SendPriority(ctx, proto.MessageFromServer, &mt.BadServerSalt{
BadMsgID: badMsgID,
BadMsgSeqno: int(badSeqno),
ErrorCode: 48,
NewServerSalt: newSalt,
})
}
// typeName 返回 TL TypeID 的可读名称，未知时回退到 hex。
func (s *Server) typeName(id uint32) string {
if name := s.types.Get(id); name != "" {
return name
}
return fmt.Sprintf("%#x", id)
}
func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint32) int {
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
return badMsgIDInvalidBits
}
msgTime := proto.MessageID(msgID).Time()
if msgTime.Before(now.Add(-300 * time.Second)) {
return badMsgIDTooLow
}
if msgTime.After(now.Add(30 * time.Second)) {
return badMsgIDTooHigh
}
if clientMessageNeedsAck(typeID) {
if seqNo%2 == 0 {
return badMsgSeqNotOdd
}
} else if seqNo%2 != 0 {
return badMsgSeqNotEven
}
return 0
}
func validateClientContainer(containerMsgID int64, containerSeqNo int32, container proto.MessageContainer) int {
for _, m := range container.Messages {
if m.ID >= containerMsgID || int32(m.SeqNo) > containerSeqNo {
return badMsgContainer
}
typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID()
if err != nil {
return badMsgContainer
}
if typeID == proto.MessageContainerTypeID {
return badMsgContainer
}
if code := validateClientContainerEnvelope(m.ID, int32(m.SeqNo), typeID); code != 0 {
return badMsgContainer
}
}
return 0
}
func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) int {
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
return badMsgIDInvalidBits
}
if clientMessageNeedsAck(typeID) {
if seqNo%2 == 0 {
return badMsgSeqNotOdd
}
} else if seqNo%2 != 0 {
return badMsgSeqNotEven
}
return 0
}
func clientMessageNeedsAck(typeID uint32) bool {
switch typeID {
case proto.MessageContainerTypeID,
mt.MsgsAckTypeID,
mt.HTTPWaitRequestTypeID,
mt.BadMsgNotificationTypeID,
mt.BadServerSaltTypeID,
mt.MsgsAllInfoTypeID,
mt.MsgsStateInfoTypeID,
mt.MsgDetailedInfoTypeID,
mt.MsgNewDetailedInfoTypeID:
return false
default:
return true
}
}
func (cs *connState) seenRecord(msgID int64) (clientMsgRecord, bool) {
record, ok := cs.seen[msgID]
return record, ok
}
func (cs *connState) validateSeq(msgID int64, seqNo int32, content bool) int {
if !content {
return 0
}
for seenMsgID, record := range cs.seen {
if !record.content {
continue
}
if seenMsgID < msgID && record.seqNo >= seqNo {
return badMsgSeqTooLow
}
if seenMsgID > msgID && record.seqNo <= seqNo {
return badMsgSeqTooHigh
}
}
return 0
}
func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) {
cs.seen[msgID] = clientMsgRecord{
state: state,
seqNo: seqNo,
content: content,
}
cs.order = append(cs.order, msgID)
if msgID < cs.minSeen {
cs.minSeen = msgID
}
if msgID > cs.maxSeen {
cs.maxSeen = msgID
}
if len(cs.order) > maxTrackedClientMsgIDs {
oldest := cs.order[0]
cs.order = cs.order[1:]
delete(cs.seen, oldest)
if oldest == cs.minSeen || oldest == cs.maxSeen {
cs.recomputeRange()
}
}
}
func (cs *connState) stateInfo(msgIDs []int64) []byte {
info := make([]byte, len(msgIDs))
if len(cs.seen) == 0 {
for i := range info {
info[i] = msgStateUnknown
}
return info
}
for i, id := range msgIDs {
if id < cs.minSeen {
info[i] = msgStateUnknown
continue
}
if id > cs.maxSeen {
info[i] = msgStateNotReceivedHigh
continue
}
record, ok := cs.seen[id]
if !ok {
info[i] = msgStateNotReceived
continue
}
info[i] = record.state
}
return info
}
func (cs *connState) recomputeRange() {
cs.minSeen = math.MaxInt64
cs.maxSeen = 0
for id := range cs.seen {
if id < cs.minSeen {
cs.minSeen = id
}
if id > cs.maxSeen {
cs.maxSeen = id
}
}
if len(cs.seen) == 0 {
cs.minSeen = math.MaxInt64
}
}

435
internal/mtprotoedge/encrypted_test.go Normal file
View file

@ -0,0 +1,435 @@
package mtprotoedge
import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
)
// TestEncryptedPingPong 验证 M2/M4：握手后 client 加密 ping，
// server 回 new_session_created + pong + msgs_ack。
func TestEncryptedPingPong(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)
const pingID int64 = 0x1234beef
pingMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, pingMsgID, &mt.PingRequest{PingID: pingID})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created")
pongBuf := mustHave(t, replies, mt.PongTypeID, "pong")
var pong mt.Pong
if err := pong.Decode(pongBuf); err != nil {
t.Fatalf("decode pong: %v", err)
}
if pong.PingID != pingID {
t.Fatalf("pong.PingID = %#x, want %#x", pong.PingID, pingID)
}
if pong.MsgID != pingMsgID {
t.Fatalf("pong.MsgID = %d, want %d (req msg id)", pong.MsgID, pingMsgID)
}
}
// TestDuplicateMsgIDIdempotent 验证 M4：相同 msg_id 的重复 content 请求被幂等处理，
// server 重发已缓存的 rpc_result，并重新 ack，不重复执行业务。
func TestDuplicateMsgIDIdempotent(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)
msgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, msgID, &mt.RPCDropAnswerRequest{ReqMsgID: msgID - 4})
first := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
mustHave(t, first, proto.ResultTypeID, "first rpc_result")
// 相同 msg_id —— 幂等：重发已有 rpc_result，并重新 ack。
sendEncrypted(t, conn, cipher, auth, msgID, &mt.RPCDropAnswerRequest{ReqMsgID: msgID - 4})
second := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
mustHave(t, second, proto.ResultTypeID, "resent rpc_result")
mustHave(t, second, mt.MsgsAckTypeID, "second ack")
}
// TestGetFutureSalts 验证 MTProto service message get_future_salts 由连接层直接响应，
// 不再落到业务 RPC fallback。
func TestGetFutureSalts(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)
sendEncrypted(t, conn, cipher, auth, reqMsgID, &mt.GetFutureSaltsRequest{Num: 32})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.FutureSaltsTypeID)
buf := mustHave(t, replies, mt.FutureSaltsTypeID, "future_salts")
var salts mt.FutureSalts
if err := salts.Decode(buf); err != nil {
t.Fatalf("decode future_salts: %v", err)
}
if salts.ReqMsgID != reqMsgID {
t.Fatalf("future_salts.req_msg_id = %d, want %d", salts.ReqMsgID, reqMsgID)
}
if len(salts.Salts) != 1 {
t.Fatalf("future_salts len = %d, want 1", len(salts.Salts))
}
if got := salts.Salts[0].Salt; got != auth.ServerSalt {
t.Fatalf("future salt = %#x, want server salt %#x", got, auth.ServerSalt)
}
if salts.Salts[0].ValidSince > salts.Now || salts.Salts[0].ValidUntil <= salts.Now {
t.Fatalf("future salt validity = [%d,%d], now %d", salts.Salts[0].ValidSince, salts.Salts[0].ValidUntil, salts.Now)
}
}
// TestMsgsStateReq 验证 MTProto service message msgs_state_req 由连接层直接响应，
// 不再落到业务 RPC fallback。
func TestMsgsStateReq(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)
asked := []int64{reqMsgID, reqMsgID - 4, reqMsgID + 4}
sendEncrypted(t, conn, cipher, auth, reqMsgID, &mt.MsgsStateReq{MsgIDs: asked})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsStateInfoTypeID)
buf := mustHave(t, replies, mt.MsgsStateInfoTypeID, "msgs_state_info")
var info mt.MsgsStateInfo
if err := info.Decode(buf); err != nil {
t.Fatalf("decode msgs_state_info: %v", err)
}
if info.ReqMsgID != reqMsgID {
t.Fatalf("msgs_state_info.req_msg_id = %d, want %d", info.ReqMsgID, reqMsgID)
}
if len(info.Info) != len(asked) {
t.Fatalf("msgs_state_info len = %d, want %d", len(info.Info), len(asked))
}
want := []byte{4, 1, 3}
for i, b := range info.Info {
if b != want[i] {
t.Fatalf("msgs_state_info[%d] = %d, want %d", i, b, want[i])
}
}
}
// TestMsgResendReq 验证 MTProto msg_resend_req 由连接层按状态查询兜底响应，
// 不会落入业务 RPC fallback。
func TestMsgResendReq(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)
asked := []int64{reqMsgID, reqMsgID - 4, reqMsgID + 4}
sendEncrypted(t, conn, cipher, auth, reqMsgID, &mt.MsgResendReq{MsgIDs: asked})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsStateInfoTypeID)
buf := mustHave(t, replies, mt.MsgsStateInfoTypeID, "msgs_state_info")
var info mt.MsgsStateInfo
if err := info.Decode(buf); err != nil {
t.Fatalf("decode msgs_state_info: %v", err)
}
if info.ReqMsgID != reqMsgID {
t.Fatalf("msgs_state_info.req_msg_id = %d, want %d", info.ReqMsgID, reqMsgID)
}
if len(info.Info) != len(asked) {
t.Fatalf("msgs_state_info len = %d, want %d", len(info.Info), len(asked))
}
want := []byte{4, 1, 3}
for i, b := range info.Info {
if b != want[i] {
t.Fatalf("msgs_state_info[%d] = %d, want %d", i, b, want[i])
}
}
}
// TestDestroySession 验证 destroy_session 返回 raw DestroySessionRes，
// 避免客户端清理旧 session 时掉到 RPC fallback。
func TestDestroySession(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)
targetSessionID := auth.SessionID + 4
sendEncrypted(t, conn, cipher, auth, reqMsgID, &mt.DestroySessionRequest{SessionID: targetSessionID})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.DestroySessionNoneTypeID)
buf := mustHave(t, replies, mt.DestroySessionNoneTypeID, "destroy_session_none")
var res mt.DestroySessionNone
if err := res.Decode(buf); err != nil {
t.Fatalf("decode destroy_session_none: %v", err)
}
if res.SessionID != targetSessionID {
t.Fatalf("destroy_session_none.session_id = %d, want %d", res.SessionID, targetSessionID)
}
}
// TestRPCDropAnswer 验证 rpc_drop_answer 以 rpc_result 包装 RpcDropAnswer 返回，
// 与 gotd/td 和 TDesktop 的请求/响应模型对齐。
func TestRPCDropAnswer(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)
droppedReqID := reqMsgID - 4
sendEncrypted(t, conn, cipher, auth, reqMsgID, &mt.RPCDropAnswerRequest{ReqMsgID: droppedReqID})
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
buf := mustHave(t, replies, proto.ResultTypeID, "rpc_result")
var result proto.Result
if err := result.Decode(buf); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if result.RequestMessageID != reqMsgID {
t.Fatalf("rpc_result.req_msg_id = %d, want %d", result.RequestMessageID, reqMsgID)
}
answer, err := mt.DecodeRPCDropAnswer(&bin.Buffer{Buf: result.Result})
if err != nil {
t.Fatalf("decode RpcDropAnswer: %v", err)
}
if _, ok := answer.(*mt.RPCAnswerUnknown); !ok {
t.Fatalf("RpcDropAnswer = %T, want *mt.RPCAnswerUnknown", answer)
}
}
// TestHTTPWaitInContainerDoesNotNeedAck 验证 http_wait 在 container 中被协议层吞掉，
// 但同 container 内的 ping 仍按 content-related service request 回 ack。
func TestHTTPWaitInContainerDoesNotNeedAck(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)
waitMsgID := clientMsgID.New(proto.MessageFromClient)
pingMsgID := clientMsgID.New(proto.MessageFromClient)
containerMsgID := clientMsgID.New(proto.MessageFromClient)
waitBody := mustEncodeTL(t, &mt.HTTPWaitRequest{MaxDelay: 0, WaitAfter: 0, MaxWait: 25_000})
pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 7})
sendEncrypted(t, conn, cipher, auth, containerMsgID, &proto.MessageContainer{
Messages: []proto.Message{
{ID: waitMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody},
{ID: pingMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody},
},
})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
mustHave(t, replies, mt.PongTypeID, "pong")
ackBuf := mustHave(t, replies, mt.MsgsAckTypeID, "msgs_ack")
var ack mt.MsgsAck
if err := ack.Decode(ackBuf); err != nil {
t.Fatalf("decode msgs_ack: %v", err)
}
if len(ack.MsgIDs) != 1 || ack.MsgIDs[0] != pingMsgID {
t.Fatalf("msgs_ack = %+v, want only ping msg_id %d", ack.MsgIDs, pingMsgID)
}
}
// TestOldMessageInFreshContainerAccepted verifies TDesktop's bad_msg recovery
// path: an old request can be resent inside a fresh container msg_id.
func TestOldMessageInFreshContainerAccepted(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
oldMsgIDGen := proto.NewMessageIDGen(func() time.Time {
return time.Now().Add(-10 * time.Minute)
})
freshMsgIDGen := proto.NewMessageIDGen(time.Now)
oldPingMsgID := oldMsgIDGen.New(proto.MessageFromClient)
containerMsgID := freshMsgIDGen.New(proto.MessageFromClient)
pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 42})
sendEncrypted(t, conn, cipher, auth, containerMsgID, &proto.MessageContainer{
Messages: []proto.Message{
{ID: oldPingMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody},
},
})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
buf := mustHave(t, replies, mt.PongTypeID, "pong")
var pong mt.Pong
if err := pong.Decode(buf); err != nil {
t.Fatalf("decode pong: %v", err)
}
if pong.MsgID != oldPingMsgID || pong.PingID != 42 {
t.Fatalf("pong = %+v, want msg_id=%d ping_id=42", pong, oldPingMsgID)
}
}
func TestPingDelayDisconnectOddSeqAccepted(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, 1, &mt.PingDelayDisconnectRequest{
PingID: 9,
DisconnectDelay: 60,
})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
buf := mustHave(t, replies, mt.PongTypeID, "pong")
var pong mt.Pong
if err := pong.Decode(buf); err != nil {
t.Fatalf("decode pong: %v", err)
}
if pong.MsgID != reqMsgID || pong.PingID != 9 {
t.Fatalf("pong = %+v, want msg_id=%d ping_id=9", pong, reqMsgID)
}
}
// TestDestroyAuthKey 验证 MTProto service message destroy_auth_key 由连接层直接响应，
// 避免 TDesktop 清理旧 key 时落到业务 RPC fallback。
func TestDestroyAuthKey(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)
sendEncrypted(t, conn, cipher, auth, reqMsgID, &destroyAuthKeyRequest{})
replies := collectReplies(t, conn, cipher, auth.AuthKey, destroyAuthKeyOkTypeID)
mustHave(t, replies, destroyAuthKeyOkTypeID, "destroy_auth_key_ok")
}
// TestBadServerSalt 验证客户端带错 server_salt 时 server 返回 bad_server_salt，
// 并携带当前 auth key 的权威 salt。
func TestBadServerSalt(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)
wrongSalt := auth.ServerSalt + 1
sendEncryptedWithSalt(t, conn, cipher, auth, wrongSalt, reqMsgID, &mt.PingRequest{PingID: 1})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.BadServerSaltTypeID)
buf := mustHave(t, replies, mt.BadServerSaltTypeID, "bad_server_salt")
var bad mt.BadServerSalt
if err := bad.Decode(buf); err != nil {
t.Fatalf("decode bad_server_salt: %v", err)
}
if bad.BadMsgID != reqMsgID {
t.Fatalf("bad_server_salt.bad_msg_id = %d, want %d", bad.BadMsgID, reqMsgID)
}
if bad.ErrorCode != 48 {
t.Fatalf("bad_server_salt.error_code = %d, want 48", bad.ErrorCode)
}
if bad.NewServerSalt != auth.ServerSalt {
t.Fatalf("bad_server_salt.new_server_salt = %#x, want %#x", bad.NewServerSalt, auth.ServerSalt)
}
}
func TestBadMsgSeqOddExpected(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, &tg.HelpGetConfigRequest{})
bad := readBadMsgNotification(t, conn, cipher, auth.AuthKey)
if bad.BadMsgID != reqMsgID || bad.BadMsgSeqno != 0 || bad.ErrorCode != badMsgSeqNotOdd {
t.Fatalf("bad_msg = %+v, want msg_id=%d seq=0 code=%d", bad, reqMsgID, badMsgSeqNotOdd)
}
}
func TestBadMsgSeqEvenExpected(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, 1, &mt.MsgsAck{MsgIDs: []int64{reqMsgID}})
bad := readBadMsgNotification(t, conn, cipher, auth.AuthKey)
if bad.BadMsgID != reqMsgID || bad.BadMsgSeqno != 1 || bad.ErrorCode != badMsgSeqNotEven {
t.Fatalf("bad_msg = %+v, want msg_id=%d seq=1 code=%d", bad, reqMsgID, badMsgSeqNotEven)
}
}
func TestBadMsgSeqTooLow(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, 3, &tg.HelpGetConfigRequest{})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
secondMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, secondMsgID, 1, &tg.HelpGetConfigRequest{})
bad := readBadMsgNotification(t, conn, cipher, auth.AuthKey)
if bad.BadMsgID != secondMsgID || bad.BadMsgSeqno != 1 || bad.ErrorCode != badMsgSeqTooLow {
t.Fatalf("bad_msg = %+v, want msg_id=%d seq=1 code=%d", bad, secondMsgID, badMsgSeqTooLow)
}
}
func TestBadMsgSeqTooHigh(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)
lowMsgID := clientMsgID.New(proto.MessageFromClient)
highMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, highMsgID, 1, &tg.HelpGetConfigRequest{})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
sendEncryptedWithSeq(t, conn, cipher, auth, lowMsgID, 3, &tg.HelpGetConfigRequest{})
bad := readBadMsgNotification(t, conn, cipher, auth.AuthKey)
if bad.BadMsgID != lowMsgID || bad.BadMsgSeqno != 3 || bad.ErrorCode != badMsgSeqTooHigh {
t.Fatalf("bad_msg = %+v, want msg_id=%d seq=3 code=%d", bad, lowMsgID, badMsgSeqTooHigh)
}
}
func readBadMsgNotification(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) mt.BadMsgNotification {
t.Helper()
replies := collectReplies(t, conn, cipher, key, mt.BadMsgNotificationTypeID)
buf := mustHave(t, replies, mt.BadMsgNotificationTypeID, "bad_msg_notification")
var bad mt.BadMsgNotification
if err := bad.Decode(buf); err != nil {
t.Fatalf("decode bad_msg_notification: %v", err)
}
return bad
}
func mustEncodeTL(t *testing.T, msg bin.Encoder) []byte {
t.Helper()
var b bin.Buffer
if err := msg.Encode(&b); err != nil {
t.Fatalf("encode TL: %v", err)
}
return b.Copy()
}

152
internal/mtprotoedge/exchange.go Normal file
View file

@ -0,0 +1,152 @@
package mtprotoedge
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"telesrv/internal/store"
)
// emptyAuthKeyID 是未加密消息（密钥交换）的 auth_key_id（全零）。
var emptyAuthKeyID [8]byte
// peekAuthKeyID 读取消息前 8 字节的 auth_key_id，不消费 buffer。
func peekAuthKeyID(b *bin.Buffer) (id [8]byte, err error) {
err = b.PeekN(id[:], len(id))
return id, err
}
// handleExchange 在收到 auth_key_id==0 的首帧后执行服务端 MTProto 密钥交换。
//
// first 是已读取的首帧（req_pq*），通过 bufferedConn 交还给 exchange 流程，
// 使其能从头读取握手消息。成功后将 auth key + server salt 落入 AuthKeyStore。
func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first *bin.Buffer) (*bin.Buffer, error) {
if s.key.Zero() {
s.log.Error("Key exchange requested but server RSA key is not configured")
return nil, s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound)
}
buffered := newBufferedConn(conn)
buffered.push(first)
start := s.clock.Now()
res, err := exchange.NewExchanger(buffered, s.dc).
WithClock(s.clock).
WithRand(s.rand).
WithLogger(s.log.Named("exchange")).
Server(s.key).
Run(ctx)
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
}
}
var exErr *exchange.ServerExchangeError
if errors.As(err, &exErr) {
s.log.Info("Key exchange rejected", zap.Int32("code", exErr.Code), zap.Error(err))
return nil, s.sendProtoError(ctx, conn, exErr.Code)
}
return nil, fmt.Errorf("key exchange: %w", err)
}
s.metrics.HandshakeDone(s.clock.Now().Sub(start))
s.log.Info("Key exchange completed",
zap.Object("auth_key", res.Key),
zap.Int64("server_salt", res.ServerSalt),
zap.Duration("dur", s.clock.Now().Sub(start)),
)
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{
ID: key.ID,
Value: [256]byte(key.Value),
ServerSalt: salt,
CreatedAt: createdAt,
}
}
// sendProtoError 向客户端发送 transport 级协议错误（-code）。
func (s *Server) sendProtoError(ctx context.Context, conn transport.Conn, code int32) error {
var buf bin.Buffer
buf.PutInt32(-code)
ctx, cancel := context.WithTimeout(ctx, s.writeTimeout)
defer cancel()
if err := conn.Send(ctx, &buf); err != nil {
return fmt.Errorf("send proto error %d: %w", code, err)
}
return nil
}
// 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
}
func newBufferedConn(conn transport.Conn) *bufferedConn {
return &bufferedConn{Conn: conn}
}
func (c *bufferedConn) push(b *bin.Buffer) {
c.mu.Lock()
c.pending = append(c.pending, bin.Buffer{Buf: b.Copy()})
c.mu.Unlock()
}
// Recv 优先返回已 push 的帧（FIFO），耗尽后读取底层连接。
func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
c.mu.Lock()
if len(c.pending) > 0 {
e := c.pending[0]
c.pending = c.pending[1:]
c.last.ResetTo(e.Copy())
c.mu.Unlock()
b.ResetTo(e.Buf)
return nil
}
c.mu.Unlock()
if err := c.Conn.Recv(ctx, b); err != nil {
return err
}
c.mu.Lock()
c.last.ResetTo(b.Copy())
c.mu.Unlock()
return nil
}
func (c *bufferedConn) lastFrame() *bin.Buffer {
c.mu.Lock()
defer c.mu.Unlock()
if c.last.Len() == 0 {
return nil
}
return &bin.Buffer{Buf: c.last.Copy()}
}

171
internal/mtprotoedge/exchange_test.go Normal file
View file

@ -0,0 +1,171 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
tgproto "github.com/gotd/td/proto"
"github.com/gotd/td/transport"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
// TestKeyExchange 验证 M1：client 用 server 公钥完成 MTProto 密钥交换，
// 双方得到一致的 auth key 与 server salt，且 server 将其存入 AuthKeyStore。
func TestKeyExchange(t *testing.T) {
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)
}
keys := memory.NewAuthKeyStore()
srv := New(Options{
Logger: zaptest.NewLogger(t),
DC: dc,
RSAKey: rsaKey,
AuthKeys: keys,
})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
// client：TCP 拨号 + intermediate 握手，跑 client 端密钥交换。
raw, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := transport.Intermediate.Handshake(raw)
if err != nil {
t.Fatalf("transport handshake: %v", err)
}
pub := exchange.PublicKey{RSA: &rsaKey.PublicKey}
exchCtx, ec := context.WithTimeout(context.Background(), 10*time.Second)
defer ec()
res, err := exchange.NewExchanger(conn, dc).
WithRand(rand.Reader).
WithLogger(zaptest.NewLogger(t).Named("client")).
Client([]exchange.PublicKey{pub}).
Run(exchCtx)
if err != nil {
t.Fatalf("client exchange: %v", err)
}
// server 在 Run 返回后落库，轮询等待。
var saved store.AuthKeyData
found := false
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
saved, found, _ = keys.Get(context.Background(), res.AuthKey.ID)
if found {
break
}
time.Sleep(20 * time.Millisecond)
}
if !found {
t.Fatalf("server did not store auth key %x", res.AuthKey.ID)
}
if saved.Value != [256]byte(res.AuthKey.Value) {
t.Fatal("server auth key value mismatch")
}
if saved.ServerSalt != res.ServerSalt {
t.Fatalf("server salt mismatch: server=%d client=%d", saved.ServerSalt, res.ServerSalt)
}
cancel()
select {
case err := <-serveErr:
if err != nil {
t.Fatalf("serve: %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("server did not stop after ctx cancel")
}
}
func TestReconnectFakeReqPQThenEncryptedFrame(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
firstConn, auth, cipher := dialHandshake(t, addr, dc, pub)
_ = firstConn.Close()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial reconnect: %v", err)
}
conn, err := transport.Intermediate.Handshake(raw)
if err != nil {
t.Fatalf("transport reconnect: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
var reqPayload bin.Buffer
nonce, err := randInt128ForTest()
if err != nil {
t.Fatalf("nonce: %v", err)
}
if err := (&mt.ReqPqMultiRequest{Nonce: nonce}).Encode(&reqPayload); err != nil {
t.Fatalf("encode req_pq_multi: %v", err)
}
var fakeReq bin.Buffer
if err := (tgproto.UnencryptedMessage{
MessageID: int64(tgproto.NewMessageID(time.Now(), tgproto.MessageFromClient)),
MessageData: reqPayload.Raw(),
}).Encode(&fakeReq); err != nil {
t.Fatalf("encode fake req_pq: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
if err := conn.Send(ctx, &fakeReq); err != nil {
cancel()
t.Fatalf("send fake req_pq: %v", err)
}
cancel()
msgGen := tgproto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, msgGen.New(tgproto.MessageFromClient), &mt.PingRequest{PingID: 7})
var resPQFrame bin.Buffer
ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
err = conn.Recv(ctx, &resPQFrame)
cancel()
if err != nil {
t.Fatalf("recv resPQ: %v", err)
}
var plain tgproto.UnencryptedMessage
if err := plain.Decode(&resPQFrame); err != nil {
t.Fatalf("decode resPQ frame: %v", err)
}
if id, err := (&bin.Buffer{Buf: plain.MessageData}).PeekID(); err != nil || id != mt.ResPQTypeID {
t.Fatalf("resPQ payload id = %#x err=%v, want %#x", id, err, mt.ResPQTypeID)
}
got := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
mustHave(t, got, mt.PongTypeID, "pong after fake req_pq reconnect")
}
func randInt128ForTest() (v bin.Int128, err error) {
_, err = rand.Read(v[:])
return v, err
}

209
internal/mtprotoedge/helpers_test.go Normal file
View file

@ -0,0 +1,209 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
)
// startTestServer 生成 RSA key、监听随机端口并启动 Server，返回监听地址与公钥。
// 通过 t.Cleanup 自动取消并校验优雅退出。opts 的 RSAKey/Logger/DC 会被补默认。
func startTestServer(t *testing.T, opts Options) (addr string, pub exchange.PublicKey, srv *Server) {
t.Helper()
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
opts.RSAKey = rsaKey
if opts.Logger == nil {
opts.Logger = zaptest.NewLogger(t)
}
if opts.DC == 0 {
opts.DC = 2
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
srv = New(opts)
ctx, cancel := context.WithCancel(context.Background())
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
t.Cleanup(func() {
cancel()
select {
case err := <-serveErr:
if err != nil {
t.Errorf("serve: %v", err)
}
case <-time.After(5 * time.Second):
t.Error("server did not stop after ctx cancel")
}
})
return ln.Addr().String(), exchange.PublicKey{RSA: &rsaKey.PublicKey}, srv
}
// dialHandshake 建立 TCP 连接、完成 intermediate 协商与 MTProto 密钥交换，
// 返回连接、握手结果与 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() })
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")).
Client([]exchange.PublicKey{pub}).
Run(ctx)
if err != nil {
t.Fatalf("client exchange: %v", err)
}
return conn, auth, crypto.NewClientCipher(rand.Reader)
}
// 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()
sendEncryptedWithSalt(t, conn, cipher, auth, auth.ServerSalt, msgID, msg)
}
// sendEncryptedWithSalt 用指定 salt 加密并发送一条消息。
func sendEncryptedWithSalt(t *testing.T, conn transport.Conn, cipher crypto.Cipher, auth exchange.ClientExchangeResult, salt, msgID int64, msg bin.Encoder) {
t.Helper()
body, seqNo := encodeClientMessageForTest(t, msg)
sendEncryptedWithSaltAndSeq(t, conn, cipher, auth, salt, msgID, seqNo, body)
}
func sendEncryptedWithSeq(t *testing.T, conn transport.Conn, cipher crypto.Cipher, auth exchange.ClientExchangeResult, msgID int64, seqNo int32, msg bin.Encoder) {
t.Helper()
body := encodeClientMessageBodyForTest(t, msg)
sendEncryptedWithSaltAndSeq(t, conn, cipher, auth, auth.ServerSalt, msgID, seqNo, body)
}
func sendEncryptedWithSaltAndSeq(t *testing.T, conn transport.Conn, cipher crypto.Cipher, auth exchange.ClientExchangeResult, salt, msgID int64, seqNo int32, body []byte) {
t.Helper()
var buf bin.Buffer
if err := cipher.Encrypt(auth.AuthKey, crypto.EncryptedMessageData{
Salt: salt,
SessionID: auth.SessionID,
MessageID: msgID,
SeqNo: seqNo,
MessageDataLen: int32(len(body)),
MessageDataWithPadding: body,
}, &buf); err != nil {
t.Fatalf("encrypt: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &buf); err != nil {
t.Fatalf("send: %v", err)
}
}
func encodeClientMessageForTest(t *testing.T, msg bin.Encoder) ([]byte, int32) {
t.Helper()
raw := encodeClientMessageBodyForTest(t, msg)
typeID, err := (&bin.Buffer{Buf: raw}).PeekID()
if err != nil {
t.Fatalf("peek encrypted message type: %v", err)
}
if container, ok := msg.(*proto.MessageContainer); ok {
return raw, clientContainerSeqNoForTest(container)
}
if clientMessageNeedsAck(typeID) {
return raw, 1
}
return raw, 0
}
func encodeClientMessageBodyForTest(t *testing.T, msg bin.Encoder) []byte {
t.Helper()
var body bin.Buffer
if err := msg.Encode(&body); err != nil {
t.Fatalf("encode encrypted message: %v", err)
}
return body.Copy()
}
func clientContainerSeqNoForTest(container *proto.MessageContainer) int32 {
var maxSeq int32
for _, msg := range container.Messages {
if seq := int32(msg.SeqNo); seq > maxSeq {
maxSeq = seq
}
}
if maxSeq%2 != 0 {
maxSeq++
}
return maxSeq
}
// collectReplies 读取并解密 server 回发的消息，按 TypeID 收集明文 buffer，
// 直到见到 wantID（含）或达到上限。用于断言一次请求触发的多条响应
// （new_session_created / 业务响应 / msgs_ack）。
func collectReplies(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey, wantID uint32) map[uint32]*bin.Buffer {
t.Helper()
got := make(map[uint32]*bin.Buffer)
for i := 0; i < 8; i++ {
_, id, plain := readServerMessage(t, conn, cipher, key)
got[id] = plain
if id == wantID {
break
}
}
return got
}
func readServerMessage(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) (*crypto.EncryptedMessageData, uint32, *bin.Buffer) {
t.Helper()
var buf bin.Buffer
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
err := conn.Recv(ctx, &buf)
cancel()
if err != nil {
t.Fatalf("recv server message: %v", err)
}
data, err := cipher.DecryptFromBuffer(key, &buf)
if err != nil {
t.Fatalf("decrypt server message: %v", err)
}
plain := append([]byte(nil), data.Data()...)
id, err := (&bin.Buffer{Buf: plain}).PeekID()
if err != nil {
t.Fatalf("peek server message: %v", err)
}
return data, id, &bin.Buffer{Buf: plain}
}
// mustHave 断言 replies 含指定 TypeID 的消息并返回其 buffer。
func mustHave(t *testing.T, replies map[uint32]*bin.Buffer, id uint32, name string) *bin.Buffer {
t.Helper()
b, ok := replies[id]
if !ok {
t.Fatalf("missing %s (%#x)", name, id)
}
return b
}

195
internal/mtprotoedge/inbound_rpc.go Normal file
View file

@ -0,0 +1,195 @@
package mtprotoedge
import (
"context"
"errors"
"time"
)
// ErrInboundRPCQueueFull 表示单连接 RPC 队列已满。
var ErrInboundRPCQueueFull = errors.New("inbound rpc queue full")
// maxInflightRPCBytes 是单连接已入队未完成 inbound RPC body 的总字节上限。
// 队列除按条数(queueSize)限制外，再按字节预算兜底：对抗客户端发满大请求时按字节先拒绝。
const maxInflightRPCBytes = 32 << 20 // 32 MiB
// rpcCloseWaitTimeout 是连接关闭时等待 inbound RPC worker 退出的上限。
const rpcCloseWaitTimeout = 5 * time.Second
type inboundRPC struct {
ctx context.Context
method string
enqueuedAt time.Time
size int
run func(context.Context) error
}
func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time.Duration) {
if c.metrics == nil {
c.metrics = NopMetrics{}
}
if maxInflight <= 0 {
maxInflight = 1
}
if queueSize <= 0 {
queueSize = 1
}
rootCtx, cancel := context.WithCancel(context.Background())
c.rpcQueue = make(chan inboundRPC, queueSize)
c.rpcStop = make(chan struct{})
c.rpcCancel = cancel
c.rpcTimeout = timeout
c.rpcRootCtx = rootCtx
c.rpcMaxInflight = maxInflight
// worker 懒启动：不在此处起 worker；首个 RPC 入队时由 ensureInboundRPCWorkers 起，
// 避免握手后静默 / 纯推送目标连接白白钉住 maxInflight 个 goroutine。
}
// ensureInboundRPCWorkers 懒启动 maxInflight 个 RPC worker（仅一次），在 enqueueInboundRPC
// 入队成功后调用。从不发 RPC 的连接（半开 / 纯推送）由此完全不起 worker。
func (c *Conn) ensureInboundRPCWorkers() {
c.rpcWorkersOnce.Do(func() {
c.rpcWG.Add(c.rpcMaxInflight)
for i := 0; i < c.rpcMaxInflight; i++ {
go c.inboundRPCWorker(c.rpcRootCtx)
}
})
}
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
if ctx == nil {
ctx = context.Background()
}
if c.rpcQueue == nil || c.rpcStop == nil {
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
return ErrConnClosed
}
task.ctx = ctx
task.enqueuedAt = time.Now()
select {
case <-ctx.Done():
c.metrics.InboundRPCDropped(task.method, "context_done")
return ctx.Err()
case <-c.rpcStop:
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
return ErrConnClosed
default:
}
// 字节预算：先预扣 size，超 maxInflightRPCBytes 则回滚并拒绝（与条数上限并列的第二道闸）。
if task.size > 0 {
if c.inflightRPCBytes.Add(int64(task.size)) > maxInflightRPCBytes {
c.inflightRPCBytes.Add(-int64(task.size))
c.metrics.InboundRPCDropped(task.method, "byte_budget")
return ErrInboundRPCQueueFull
}
}
select {
case c.rpcQueue <- task:
c.ensureInboundRPCWorkers()
c.metrics.InboundRPCQueued(task.method, len(c.rpcQueue), cap(c.rpcQueue))
return nil
case <-ctx.Done():
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "context_done")
return ctx.Err()
case <-c.rpcStop:
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
return ErrConnClosed
default:
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "queue_full")
return ErrInboundRPCQueueFull
}
}
// releaseInflightRPCBytes 归还字节预算。与 enqueueInboundRPC 的预扣严格配对：
// 入队失败时回滚、worker 执行完(runInboundRPC)或排空丢弃(drainInboundRPCQueue)时释放。
func (c *Conn) releaseInflightRPCBytes(size int) {
if size > 0 {
c.inflightRPCBytes.Add(-int64(size))
}
}
func (c *Conn) inboundRPCWorker(rootCtx context.Context) {
defer c.rpcWG.Done()
for {
select {
case <-c.rpcStop:
return
default:
}
select {
case task := <-c.rpcQueue:
c.runInboundRPC(rootCtx, task)
case <-c.rpcStop:
return
}
}
}
func (c *Conn) runInboundRPC(rootCtx context.Context, task inboundRPC) {
defer c.releaseInflightRPCBytes(task.size)
queueWait := time.Since(task.enqueuedAt)
c.metrics.InboundRPCStarted(task.method, queueWait)
ctx := task.ctx
if ctx == nil {
ctx = context.Background()
}
ctx, cancel := context.WithCancel(ctx)
defer cancel()
stopRoot := context.AfterFunc(rootCtx, cancel)
defer stopRoot()
if c.rpcTimeout > 0 {
var timeoutCancel context.CancelFunc
ctx, timeoutCancel = context.WithTimeout(ctx, c.rpcTimeout)
defer timeoutCancel()
}
_ = task.run(ctx)
}
func (c *Conn) closeInboundRPCScheduler() {
if c.rpcStop == nil {
return
}
c.rpcClose.Do(func() {
if c.rpcCancel != nil {
c.rpcCancel()
}
close(c.rpcStop)
// 抢占懒启动 Once：若 worker 尚未起，封住其启动，避免后续 ensureInboundRPCWorkers 的
// rpcWG.Add 与下面的 rpcWG.Wait 并发（WaitGroup 误用）。Once 互斥保证 Add happens-before Wait。
c.rpcWorkersOnce.Do(func() {})
c.drainInboundRPCQueue()
// 等 worker 退出，使关闭对 inbound 与 outbound（<-outboundDone）收敛对称；带超时防慢 handler 卡死。
c.waitInboundWorkers(rpcCloseWaitTimeout)
})
}
// waitInboundWorkers 等所有 inbound RPC worker 退出，最长 timeout。超时则放弃等待，
// worker 在其阻塞的底层调用返回后自行退出（rpcCancel 已发，最终收敛）。
func (c *Conn) waitInboundWorkers(timeout time.Duration) {
done := make(chan struct{})
go func() {
c.rpcWG.Wait()
close(done)
}()
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-done:
case <-timer.C:
}
}
func (c *Conn) drainInboundRPCQueue() {
for {
select {
case task := <-c.rpcQueue:
c.releaseInflightRPCBytes(task.size)
c.metrics.InboundRPCDropped(task.method, "connection_closed")
default:
return
}
}
}

76
internal/mtprotoedge/inbound_rpc_test.go Normal file
View file

@ -0,0 +1,76 @@
package mtprotoedge
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
)
func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
c.startInboundRPCScheduler(2, 4, time.Second)
defer c.closeInboundRPCScheduler()
var active atomic.Int64
var maxActive atomic.Int64
var done atomic.Int64
started := make(chan struct{}, 6)
release := make(chan struct{})
task := inboundRPC{
method: "test.method",
run: func(ctx context.Context) error {
cur := active.Add(1)
for {
old := maxActive.Load()
if cur <= old || maxActive.CompareAndSwap(old, cur) {
break
}
}
started <- struct{}{}
select {
case <-release:
case <-ctx.Done():
}
active.Add(-1)
done.Add(1)
return nil
},
}
for i := 0; i < 2; i++ {
if err := c.enqueueInboundRPC(context.Background(), task); err != nil {
t.Fatalf("enqueue active task %d: %v", i, err)
}
}
for i := 0; i < 2; i++ {
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("timed out waiting for active rpc workers")
}
}
for i := 0; i < 4; i++ {
if err := c.enqueueInboundRPC(context.Background(), task); err != nil {
t.Fatalf("enqueue queued task %d: %v", i, err)
}
}
if err := c.enqueueInboundRPC(context.Background(), task); !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("enqueue over capacity err = %v, want ErrInboundRPCQueueFull", err)
}
if got := maxActive.Load(); got != 2 {
t.Fatalf("max active = %d, want 2", got)
}
close(release)
deadline := time.After(2 * time.Second)
for done.Load() != 6 {
select {
case <-deadline:
t.Fatalf("done = %d, want 6", done.Load())
default:
time.Sleep(10 * time.Millisecond)
}
}
}

457
internal/mtprotoedge/login_e2e_test.go Normal file
View file

@ -0,0 +1,457 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"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/langpack"
messageapp "telesrv/internal/app/messages"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/domain"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
// TestLoginRegisterFlow 是登录注册闭环的端到端验证：telegram.Client 连本地 server，
// 依次 sendCode → signIn(需注册) → signUp → getUsers(self)，验证注册后能用 self 查回自己。
func TestLoginRegisterFlow(t *testing.T) {
const (
dc = 2
phone = "+8613800138000"
wantPhone = "8613800138000"
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()
if err := helpStore.UpsertAppConfig(context.Background(), domain.AppConfig{
Client: "tdesktop",
Hash: 4,
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)
}
if err := helpStore.UpsertCountries(context.Background(), []domain.Country{
{ISO2: "US", DefaultName: "United States", CountryCodes: []domain.CountryCode{{CountryCode: "1", Prefixes: []string{"1"}}}},
}); err != nil {
t.Fatalf("seed countries: %v", err)
}
langPackStore := memory.NewLangPackStore()
if err := langPackStore.UpsertPack(context.Background(), domain.LangPack{
LangPack: "tdesktop",
LangCode: "en",
Version: 1,
Strings: []domain.LangPackString{{Key: "lng_language_name", Value: "English"}},
}); err != nil {
t.Fatalf("seed langpack: %v", err)
}
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code),
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(memory.NewDialogStore()),
LangPack: langpack.NewService(langPackStore),
}
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)
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
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: zaptest.NewLogger(t).Named("client"),
SessionStorage: &session.StorageMemory{},
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
client := telegram.NewClient(1, "hash", opts)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
// 1) sendCode → phone_code_hash
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 {
t.Fatalf("sendCode result = %T, want *tg.AuthSentCode", sent)
}
hash := sentCode.PhoneCodeHash
// 2) signIn → 新用户应得 SignUpRequired
signInRes, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{
PhoneNumber: phone,
PhoneCodeHash: hash,
PhoneCode: code,
})
if err != nil {
return err
}
if _, ok := signInRes.(*tg.AuthAuthorizationSignUpRequired); !ok {
t.Fatalf("signIn result = %T, want *tg.AuthAuthorizationSignUpRequired", signInRes)
}
// 3) signUp → 创建用户并返回授权
signUpRes, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{
PhoneNumber: phone,
PhoneCodeHash: hash,
FirstName: "Test",
LastName: "User",
})
if err != nil {
return err
}
authz, ok := signUpRes.(*tg.AuthAuthorization)
if !ok {
t.Fatalf("signUp result = %T, want *tg.AuthAuthorization", signUpRes)
}
newUser, ok := authz.User.(*tg.User)
if !ok {
t.Fatalf("signUp user = %T, want *tg.User", authz.User)
}
if !newUser.Self || newUser.FirstName != "Test" || newUser.Phone != wantPhone {
t.Fatalf("signUp user = %+v, want self FirstName=Test Phone=%s", newUser, wantPhone)
}
// 4) getUsers(self) → 注册后能查回自己
got, err := raw.UsersGetUsers(ctx, []tg.InputUserClass{&tg.InputUserSelf{}})
if err != nil {
return err
}
if len(got) != 1 {
t.Fatalf("getUsers returned %d users, want 1", len(got))
}
self, ok := got[0].(*tg.User)
if !ok {
t.Fatalf("getUsers[0] = %T, want *tg.User", got[0])
}
if self.ID != newUser.ID || self.FirstName != "Test" || self.Phone != wantPhone {
t.Fatalf("getUsers self = %+v, want id=%d FirstName=Test Phone=%s", self, newUser.ID, wantPhone)
}
// 5) 启动配置、账号安全和登录后的空账号 RPC 走业务服务并可编码。
appConfig, err := raw.HelpGetAppConfig(ctx, 0)
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)
}
countriesRes, err := raw.HelpGetCountriesList(ctx, &tg.HelpGetCountriesListRequest{LangCode: "en"})
if err != nil {
return err
}
if countries, ok := countriesRes.(*tg.HelpCountriesList); !ok || len(countries.Countries) != 1 {
t.Fatalf("help.getCountriesList = %T %+v, want 1 country", countriesRes, countriesRes)
}
password, err := raw.AccountGetPassword(ctx)
if err != nil {
return err
}
if password.HasPassword || len(password.SecureRandom) == 0 {
t.Fatalf("account.getPassword = %+v, want no password with secure random", password)
}
state, err := raw.UpdatesGetState(ctx)
if err != nil {
return err
}
if state.Date == 0 {
t.Fatal("updates.getState Date is zero")
}
diff, err := raw.UpdatesGetDifference(ctx, &tg.UpdatesGetDifferenceRequest{
Pts: state.Pts,
Date: state.Date,
Qts: state.Qts,
})
if err != nil {
return err
}
if _, ok := diff.(*tg.UpdatesDifferenceEmpty); !ok {
t.Fatalf("updates.getDifference = %T, want *tg.UpdatesDifferenceEmpty", diff)
}
contactsRes, err := raw.ContactsGetContacts(ctx, 0)
if err != nil {
return err
}
if contacts, ok := contactsRes.(*tg.ContactsContacts); !ok || len(contacts.Contacts) != 0 {
t.Fatalf("contacts.getContacts = %T %+v, want empty *tg.ContactsContacts", contactsRes, contactsRes)
}
dialogsRes, err := raw.MessagesGetDialogs(ctx, &tg.MessagesGetDialogsRequest{
OffsetPeer: &tg.InputPeerEmpty{},
Limit: 20,
})
if err != nil {
return err
}
if dialogs, ok := dialogsRes.(*tg.MessagesDialogs); !ok || len(dialogs.Dialogs) != 0 {
t.Fatalf("messages.getDialogs = %T %+v, want empty *tg.MessagesDialogs", dialogsRes, dialogsRes)
}
pinned, err := raw.MessagesGetPinnedDialogs(ctx, 0)
if err != nil {
return err
}
if len(pinned.Dialogs) != 0 || pinned.State.Date == 0 {
t.Fatalf("messages.getPinnedDialogs = %+v, want empty dialogs with state", pinned)
}
pack, err := raw.LangpackGetLangPack(ctx, &tg.LangpackGetLangPackRequest{
LangPack: "tdesktop",
LangCode: "en",
})
if err != nil {
return err
}
if pack.Version != 1 || len(pack.Strings) != 1 {
t.Fatalf("langpack.getLangPack = %+v, want version 1 with 1 string", pack)
}
strings, err := raw.LangpackGetStrings(ctx, &tg.LangpackGetStringsRequest{
LangPack: "tdesktop",
LangCode: "en",
Keys: []string{"lng_language_name"},
})
if err != nil {
return err
}
if len(strings) != 1 {
t.Fatalf("langpack.getStrings returned %d strings, want 1", len(strings))
}
return nil
}); err != nil {
t.Fatalf("login/register flow: %v", err)
}
cancel()
if err := <-serveErr; err != nil {
t.Errorf("serve: %v", err)
}
}
func TestPrivateMessageRoundTripFlow(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)
activeSessions := NewSessionManager(zaptest.NewLogger(t).Named("sessions"))
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code),
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),
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)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
ctx, cancel := context.WithTimeout(context.Background(), 30*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: zaptest.NewLogger(t).Named("client"),
SessionStorage: storage,
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
return telegram.NewClient(1, "hash", opts)
}
storageA := &session.StorageMemory{}
storageB := &session.StorageMemory{}
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
}
}
signUp := func(storage *session.StorageMemory, phone, firstName string) tg.User {
t.Helper()
client := newClient(storage)
var out tg.User
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: code,
}); err != nil {
return err
}
res, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{
PhoneNumber: phone,
PhoneCodeHash: hash,
FirstName: firstName,
})
if err != nil {
return err
}
authz := res.(*tg.AuthAuthorization)
u := authz.User.(*tg.User)
out = *u
return nil
}); err != nil {
t.Fatalf("signUp %s: %v", firstName, err)
}
return out
}
userA := signUp(storageA, "+15550001001", "Alice")
userB := signUp(storageB, "+15550001002", "Bob")
sendAndRead := func(storage *session.StorageMemory, to tg.User, body string, randomID int64) {
t.Helper()
client := newClient(storage)
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
updates, err := raw.MessagesSendMessage(ctx, &tg.MessagesSendMessageRequest{
Peer: &tg.InputPeerUser{UserID: to.ID, AccessHash: to.AccessHash},
Message: body,
RandomID: randomID,
})
if err != nil {
return err
}
gotUpdates, ok := updates.(*tg.Updates)
if !ok || len(gotUpdates.Updates) < 2 {
t.Fatalf("send updates = %T %+v, want message id + new message", updates, updates)
}
history, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerUser{UserID: to.ID, AccessHash: to.AccessHash},
Limit: 10,
})
if err != nil {
return err
}
msgs := messagesOf(history)
if len(msgs) == 0 {
t.Fatalf("history = %T %+v, want messages", history, history)
}
msg, ok := msgs[0].(*tg.Message)
if !ok || msg.Message != body || !msg.Out {
t.Fatalf("latest history message = %#v, want outgoing %q", msgs[0], body)
}
return nil
}); err != nil {
t.Fatalf("send %q: %v", body, err)
}
}
sendAndRead(storageA, userB, "hello bob", 1001)
clientB := newClient(storageB)
if err := clientB.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(clientB)
history, err := raw.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerUser{UserID: userA.ID, AccessHash: userA.AccessHash},
Limit: 10,
})
if err != nil {
return err
}
msgs := messagesOf(history)
if len(msgs) == 0 {
t.Fatalf("bob history = %T %+v, want incoming message", history, history)
}
msg, ok := msgs[0].(*tg.Message)
if !ok || msg.Message != "hello bob" || msg.Out {
t.Fatalf("bob latest message = %#v, want incoming hello bob", msgs[0])
}
return nil
}); err != nil {
t.Fatalf("bob read incoming: %v", err)
}
sendAndRead(storageB, userA, "hi alice", 2001)
cancel()
if err := <-serveErr; err != nil {
t.Errorf("serve: %v", err)
}
}

67
internal/mtprotoedge/metrics.go Normal file
View file

@ -0,0 +1,67 @@
package mtprotoedge
import "time"
// Metrics 接收连接层运行指标。实现可对接 Prometheus 等监控系统；
// 默认 NopMetrics（零开销）。第一阶段仅预留钩子，正式指标后续接入。
type Metrics interface {
// ConnOpened 在接受一个连接时调用。
ConnOpened()
// ConnClosed 在一个连接结束时调用。
ConnClosed()
// HandshakeDone 在一次密钥交换成功完成时调用，d 为握手耗时。
HandshakeDone(d time.Duration)
// RPCHandled 在一次 RPC 处理完成时调用：method 为 TL 方法名，
// d 为耗时，err 非 nil 表示失败。
RPCHandled(method string, d time.Duration, err error)
// InboundRPCQueued 在 RPC 成功进入单连接 bounded queue 时调用。
InboundRPCQueued(method string, len, cap int)
// InboundRPCStarted 在 RPC 从 bounded queue 取出开始执行时调用。
InboundRPCStarted(method string, queueWait time.Duration)
// InboundRPCDropped 在 RPC 因队列满、连接关闭或调度错误被丢弃时调用。
InboundRPCDropped(method, reason string)
// OutboundSend 在一条 server 出站消息完成写入或失败时调用。
OutboundSend(typeID uint32, queueWait time.Duration, bytes int, err error)
// OutboundResend 在一次 msg_resend_req/重复 RPC 触发重发后调用。
OutboundResend(count int, err error)
// OutboundDropped 在出站队列或状态跟踪因背压丢弃时调用。
OutboundDropped(reason string)
// OutboundQueueWait 在出站入队等待超过阈值时调用。
OutboundQueueWait(len, cap int)
}
// NopMetrics 是 Metrics 的空实现。
type NopMetrics struct{}
// ConnOpened 实现 Metrics。
func (NopMetrics) ConnOpened() {}
// ConnClosed 实现 Metrics。
func (NopMetrics) ConnClosed() {}
// HandshakeDone 实现 Metrics。
func (NopMetrics) HandshakeDone(time.Duration) {}
// RPCHandled 实现 Metrics。
func (NopMetrics) RPCHandled(string, time.Duration, error) {}
// InboundRPCQueued 实现 Metrics。
func (NopMetrics) InboundRPCQueued(string, int, int) {}
// InboundRPCStarted 实现 Metrics。
func (NopMetrics) InboundRPCStarted(string, time.Duration) {}
// InboundRPCDropped 实现 Metrics。
func (NopMetrics) InboundRPCDropped(string, string) {}
// OutboundSend 实现 Metrics。
func (NopMetrics) OutboundSend(uint32, time.Duration, int, error) {}
// OutboundResend 实现 Metrics。
func (NopMetrics) OutboundResend(int, error) {}
// OutboundDropped 实现 Metrics。
func (NopMetrics) OutboundDropped(string) {}
// OutboundQueueWait 实现 Metrics。
func (NopMetrics) OutboundQueueWait(int, int) {}

73
internal/mtprotoedge/metrics_test.go Normal file
View file

@ -0,0 +1,73 @@
package mtprotoedge
import (
"sync/atomic"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/clock"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"telesrv/internal/rpc"
)
type countingMetrics struct {
connOpened atomic.Int64
connClosed atomic.Int64
handshakes atomic.Int64
rpcs atomic.Int64
inbound atomic.Int64
outbound atomic.Int64
}
func (m *countingMetrics) ConnOpened() { m.connOpened.Add(1) }
func (m *countingMetrics) ConnClosed() { m.connClosed.Add(1) }
func (m *countingMetrics) HandshakeDone(time.Duration) { m.handshakes.Add(1) }
func (m *countingMetrics) RPCHandled(string, time.Duration, error) { m.rpcs.Add(1) }
func (m *countingMetrics) InboundRPCQueued(string, int, int) {}
func (m *countingMetrics) InboundRPCStarted(string, time.Duration) { m.inbound.Add(1) }
func (m *countingMetrics) InboundRPCDropped(string, string) {}
func (m *countingMetrics) OutboundSend(uint32, time.Duration, int, error) {
m.outbound.Add(1)
}
func (m *countingMetrics) OutboundResend(int, error) {}
func (m *countingMetrics) OutboundDropped(string) {}
func (m *countingMetrics) OutboundQueueWait(int, int) {}
// TestMetricsHooks 验证 M5：连接、握手、RPC 的 metrics 钩子被正确调用。
func TestMetricsHooks(t *testing.T) {
const dc = 2
m := &countingMetrics{}
router := rpc.New(rpc.Config{DC: dc, IP: "127.0.0.1", Port: 2398}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: router, Metrics: m})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &tg.HelpGetConfigRequest{})
collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
if got := m.connOpened.Load(); got < 1 {
t.Errorf("ConnOpened called %d times, want >= 1", got)
}
if got := m.handshakes.Load(); got != 1 {
t.Errorf("HandshakeDone called %d times, want 1", got)
}
if got := m.rpcs.Load(); got != 1 {
t.Errorf("RPCHandled called %d times, want 1", got)
}
if got := m.inbound.Load(); got != 1 {
t.Errorf("InboundRPCStarted called %d times, want 1", got)
}
// new_session_created / ack 走 fire-and-forget（异步），可能在 client 收到 rpc_result 后
// 才被 outbound actor 处理；轮询等其最终发送完成。M5 验证发送计数，不约束同步时序。
deadline := time.Now().Add(2 * time.Second)
for m.outbound.Load() < 3 && time.Now().Before(deadline) {
time.Sleep(5 * time.Millisecond)
}
if got := m.outbound.Load(); got < 3 {
t.Errorf("OutboundSend called %d times, want >= 3", got)
}
}

647
internal/mtprotoedge/outbound.go Normal file
View file

@ -0,0 +1,647 @@
package mtprotoedge
import (
"context"
"errors"
"fmt"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
)
var (
// ErrConnClosed 表示连接的出站 actor 已关闭。
ErrConnClosed = errors.New("mtproto connection closed")
// ErrOutboundQueueFull 表示 best-effort update push 未能在预算内进入出站队列。
ErrOutboundQueueFull = errors.New("mtproto outbound queue full")
)
const (
maxOutboundQueue = 1024
maxTrackedServerMsgIDs = 4096
maxTrackedAckedMsgIDs = 1024
// maxTrackedServerBytes 是 pending（已发送待 ack、用于 resend）总 body 字节上限。
// 与 maxTrackedServerMsgIDs 并列：客户端从不 ack 时，大响应体按字节滚动丢弃，
// 防 pending 被「4096 条 × 大 body」撑爆。
maxTrackedServerBytes = 64 << 20 // 64 MiB
)
type outboundOpKind byte
const (
outboundSend outboundOpKind = iota + 1
outboundAck
outboundQueryState
outboundResend
outboundResendByRequest
)
type outboundOp struct {
kind outboundOpKind
control bool
ctx context.Context
msgType proto.MessageType
msg bin.Encoder
ids []int64
reqMsgID int64
enqueuedAt time.Time
done chan outboundResult
}
type outboundResult struct {
info []byte
resent bool
err error
}
type outboundFrame struct {
msgID int64
seqNo int32
typeID uint32
body []byte
reqMsgID int64
sentAt time.Time
sends int
}
type outboundState struct {
pending map[int64]*outboundFrame
order []int64
byRequest map[int64]int64
acked map[int64]struct{}
ackOrder []int64
totalBytes int
}
func newOutboundState() *outboundState {
return &outboundState{
pending: make(map[int64]*outboundFrame),
byRequest: make(map[int64]int64),
acked: make(map[int64]struct{}),
}
}
func (c *Conn) startOutbound() {
if c.metrics == nil {
c.metrics = NopMetrics{}
}
c.outbound = make(chan outboundOp, maxOutboundQueue)
c.outboundControl = make(chan outboundOp, maxOutboundQueue/4)
c.outboundStop = make(chan struct{})
c.outboundDone = make(chan struct{})
go c.outboundLoop()
}
// Close 停止连接的出站 actor。它不关闭底层 transport；transport 生命周期仍由 serveConn 管理。
func (c *Conn) Close() {
c.closeInboundRPCScheduler()
c.outboundClose.Do(func() {
if c.outboundStop != nil {
close(c.outboundStop)
<-c.outboundDone
}
})
}
// Send 加密并发送一条 server 消息。
func (c *Conn) Send(ctx context.Context, t proto.MessageType, msg bin.Encoder) error {
return c.send(ctx, t, msg, false)
}
// SendPriority 加密并优先发送一条 server 控制消息。
func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.Encoder) error {
return c.send(ctx, t, msg, true)
}
// SendBestEffort 只等待消息进入普通 outbound 队列，不等待网络写完成。
// 用于 updates fanout：队列拥塞时返回 ErrOutboundQueueFull，durable outbox/getDifference 负责兜底。
func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.Encoder, timeout time.Duration) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
writeCtx := context.Background()
if ctx != nil {
writeCtx = context.WithoutCancel(ctx)
}
op := outboundOp{
kind: outboundSend,
ctx: writeCtx,
msgType: t,
msg: msg,
enqueuedAt: time.Now(),
}
if timeout == 0 {
select {
case c.outbound <- op:
return nil
case <-c.outboundStop:
return ErrConnClosed
default:
c.metrics.OutboundDropped("push_queue_full")
return ErrOutboundQueueFull
}
}
enqueueCtx := ctx
if enqueueCtx == nil {
enqueueCtx = context.Background()
}
var cancel context.CancelFunc
if timeout > 0 {
enqueueCtx, cancel = context.WithTimeout(enqueueCtx, timeout)
defer cancel()
}
if err := c.enqueueOutbound(enqueueCtx, op); err != nil {
if errors.Is(err, context.DeadlineExceeded) && timeout > 0 {
c.metrics.OutboundDropped("push_queue_timeout")
return ErrOutboundQueueFull
}
return err
}
return nil
}
func (c *Conn) send(ctx context.Context, t proto.MessageType, msg bin.Encoder, control bool) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
op := outboundOp{
kind: outboundSend,
control: control,
ctx: ctx,
msgType: t,
msg: msg,
enqueuedAt: time.Now(),
done: make(chan outboundResult, 1),
}
if err := c.enqueueOutbound(ctx, op); err != nil {
return err
}
select {
case res := <-op.done:
return res.err
case <-ctx.Done():
return ctx.Err()
case <-c.outboundStop:
return ErrConnClosed
}
}
// SendAsync 入队一条 server 消息但不等待发送结果（fire-and-forget），用于读循环里的控制消息
// （ack/pong/new_session_created/bad_msg/future_salts/state_info）：避免读循环被 outbound 写
// 阻塞而连带卡死。走优先(control)队列保证不被普通 push 拖后；队列满时丢弃并记 metrics——此时
// 连接多已严重拥塞，控制消息丢失由客户端重传 / 读写超时兜底。返回非 nil 仅表示连接已关闭。
func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encoder) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
op := outboundOp{
kind: outboundSend,
control: true,
ctx: ctx,
msgType: t,
msg: msg,
enqueuedAt: time.Now(),
// done 为 nil：fire-and-forget，handleOutboundSend 的 finish 对 nil done 安全跳过。
}
select {
case c.outboundControl <- op:
return nil
case <-c.outboundStop:
return ErrConnClosed
default:
c.metrics.OutboundDropped("control_queue_full")
return nil
}
}
// AckServerMessages 接收客户端 msgs_ack，释放已确认的 server 出站消息。
func (c *Conn) AckServerMessages(ids []int64) {
if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil {
return
}
copied := append([]int64(nil), ids...)
op := outboundOp{kind: outboundAck, control: true, ids: copied}
select {
case c.outboundControl <- op:
case <-c.outboundStop:
default:
c.metrics.OutboundDropped("ack_queue_full")
}
}
// OutgoingStateInfo 返回本连接出站消息的状态。返回值中 0 表示无出站侧意见，
// 调用方可继续用入站 connState 兜底。
func (c *Conn) OutgoingStateInfo(ctx context.Context, ids []int64) ([]byte, error) {
if c.outbound == nil {
return nil, ErrConnClosed
}
op := outboundOp{
kind: outboundQueryState,
control: true,
ctx: ctx,
ids: append([]int64(nil), ids...),
done: make(chan outboundResult, 1),
}
if err := c.enqueueOutbound(ctx, op); err != nil {
return nil, err
}
select {
case res := <-op.done:
return res.info, res.err
case <-ctx.Done():
return nil, ctx.Err()
case <-c.outboundStop:
return nil, ErrConnClosed
}
}
// ResendMessages 重发仍在 outgoing queue 中的 server 消息，并返回对应状态。
func (c *Conn) ResendMessages(ctx context.Context, ids []int64) ([]byte, error) {
if c.outbound == nil {
return nil, ErrConnClosed
}
op := outboundOp{
kind: outboundResend,
control: true,
ctx: ctx,
ids: append([]int64(nil), ids...),
done: make(chan outboundResult, 1),
}
if err := c.enqueueOutbound(ctx, op); err != nil {
return nil, err
}
select {
case res := <-op.done:
return res.info, res.err
case <-ctx.Done():
return nil, ctx.Err()
case <-c.outboundStop:
return nil, ErrConnClosed
}
}
// ResendByRequest 在重复 RPC 请求到达时，按原 client msg_id 找到并重发已有 rpc_result。
func (c *Conn) ResendByRequest(ctx context.Context, reqMsgID int64) (bool, error) {
if c.outbound == nil {
return false, ErrConnClosed
}
op := outboundOp{
kind: outboundResendByRequest,
control: true,
ctx: ctx,
reqMsgID: reqMsgID,
done: make(chan outboundResult, 1),
}
if err := c.enqueueOutbound(ctx, op); err != nil {
return false, err
}
select {
case res := <-op.done:
return res.resent, res.err
case <-ctx.Done():
return false, ctx.Err()
case <-c.outboundStop:
return false, ErrConnClosed
}
}
func (c *Conn) enqueueOutbound(ctx context.Context, op outboundOp) error {
if ctx == nil {
ctx = context.Background()
}
q := c.outbound
if op.control {
q = c.outboundControl
}
select {
case q <- op:
return nil
case <-ctx.Done():
return ctx.Err()
case <-c.outboundStop:
return ErrConnClosed
default:
}
c.metrics.OutboundQueueWait(len(q), cap(q))
select {
case q <- op:
return nil
case <-ctx.Done():
return ctx.Err()
case <-c.outboundStop:
return ErrConnClosed
}
}
func (c *Conn) outboundLoop() {
defer close(c.outboundDone)
state := newOutboundState()
for {
select {
case op := <-c.outboundControl:
c.handleOutboundOp(state, op)
continue
default:
}
select {
case <-c.outboundStop:
c.drainOutbound()
return
case op := <-c.outboundControl:
c.handleOutboundOp(state, op)
case op := <-c.outbound:
c.handleOutboundOp(state, op)
}
}
}
func (c *Conn) drainOutbound() {
for {
select {
case op := <-c.outboundControl:
op.finish(outboundResult{err: ErrConnClosed})
case op := <-c.outbound:
op.finish(outboundResult{err: ErrConnClosed})
default:
return
}
}
}
func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) {
switch op.kind {
case outboundSend:
c.handleOutboundSend(state, op)
case outboundAck:
state.ack(op.ids)
case outboundQueryState:
op.finish(outboundResult{info: state.stateInfo(op.ids)})
case outboundResend:
info, err := c.handleOutboundResend(state, op.ctx, op.ids)
op.finish(outboundResult{info: info, err: err})
case outboundResendByRequest:
resent, err := c.handleOutboundResendByRequest(state, op.ctx, op.reqMsgID)
op.finish(outboundResult{resent: resent, err: err})
default:
op.finish(outboundResult{err: fmt.Errorf("unknown outbound op %d", op.kind)})
}
}
func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
frame, err := c.buildFrame(op.msgType, op.msg)
if err == nil {
err = c.writeFrame(op.ctx, frame)
}
if err == nil && frameNeedsAck(frame.typeID) {
if dropped := state.add(frame); dropped > 0 {
for i := 0; i < dropped; i++ {
c.metrics.OutboundDropped("tracked_queue_overflow")
}
}
}
queueWait := time.Since(op.enqueuedAt)
bytes := 0
typeID := uint32(0)
if frame != nil {
bytes = len(frame.body)
typeID = frame.typeID
}
c.metrics.OutboundSend(typeID, queueWait, bytes, err)
op.finish(outboundResult{err: err})
}
func (c *Conn) handleOutboundResend(state *outboundState, ctx context.Context, ids []int64) ([]byte, error) {
info := make([]byte, len(ids))
resent := 0
for i, id := range ids {
if state.isKnown(id) {
info[i] = msgStateReceived
}
frame, ok := state.pending[id]
if !ok {
continue
}
if err := c.writeFrame(ctx, frame); err != nil {
c.metrics.OutboundResend(resent, err)
return info, err
}
frame.sentAt = time.Now()
frame.sends++
resent++
}
c.metrics.OutboundResend(resent, nil)
return info, nil
}
func (c *Conn) handleOutboundResendByRequest(state *outboundState, ctx context.Context, reqMsgID int64) (bool, error) {
msgID, ok := state.byRequest[reqMsgID]
if !ok {
return false, nil
}
frame, ok := state.pending[msgID]
if !ok {
return false, nil
}
if err := c.writeFrame(ctx, frame); err != nil {
c.metrics.OutboundResend(0, err)
return false, err
}
frame.sentAt = time.Now()
frame.sends++
c.metrics.OutboundResend(1, nil)
return true, nil
}
func (op outboundOp) finish(res outboundResult) {
if op.done == nil {
return
}
select {
case op.done <- res:
default:
}
}
func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder) (*outboundFrame, error) {
if msg == nil {
return nil, errors.New("nil outbound message")
}
var body bin.Buffer
if err := msg.Encode(&body); err != nil {
return nil, fmt.Errorf("encode outbound: %w", err)
}
typeID, err := (&bin.Buffer{Buf: body.Raw()}).PeekID()
if err != nil {
return nil, fmt.Errorf("peek outbound type id: %w", err)
}
content := frameNeedsAck(typeID)
msgID := c.msgID.New(t)
return &outboundFrame{
msgID: msgID,
seqNo: c.nextSeqNo(content),
typeID: typeID,
body: body.Copy(),
reqMsgID: outboundRequestMsgID(msg),
}, nil
}
func (c *Conn) nextSeqNo(content bool) int32 {
seqNo := c.sentContentMessages * 2
if content {
seqNo++
c.sentContentMessages++
}
return seqNo
}
func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
if ctx == nil {
ctx = context.Background()
}
var out bin.Buffer
if err := c.cipher.Encrypt(c.key, crypto.EncryptedMessageData{
Salt: c.salt,
SessionID: c.sessionID,
MessageID: frame.msgID,
SeqNo: frame.seqNo,
MessageDataLen: int32(len(frame.body)),
MessageDataWithPadding: frame.body,
}, &out); err != nil {
return fmt.Errorf("encrypt: %w", err)
}
sendCtx := ctx
cancel := func() {}
if c.writeTimeout > 0 {
sendCtx, cancel = context.WithTimeout(ctx, c.writeTimeout)
}
defer cancel()
writer := c.writer
if writer == nil {
writer = c.transport
}
if err := writer.Send(sendCtx, &out); err != nil {
return fmt.Errorf("send: %w", err)
}
if frame.sentAt.IsZero() {
frame.sentAt = time.Now()
frame.sends = 1
}
return nil
}
func frameNeedsAck(typeID uint32) bool {
switch typeID {
case mt.MsgsAckTypeID,
mt.BadMsgNotificationTypeID,
mt.BadServerSaltTypeID,
mt.MsgsStateInfoTypeID,
mt.MsgsAllInfoTypeID,
mt.MsgDetailedInfoTypeID,
mt.MsgNewDetailedInfoTypeID,
proto.MessageContainerTypeID:
return false
default:
return true
}
}
func outboundRequestMsgID(msg bin.Encoder) int64 {
switch v := msg.(type) {
case *proto.Result:
return v.RequestMessageID
default:
return 0
}
}
func (s *outboundState) add(frame *outboundFrame) int {
s.pending[frame.msgID] = frame
s.order = append(s.order, frame.msgID)
s.totalBytes += len(frame.body)
if frame.reqMsgID != 0 {
s.byRequest[frame.reqMsgID] = frame.msgID
}
return s.shrinkPending()
}
func (s *outboundState) ack(ids []int64) {
for _, id := range ids {
frame, ok := s.pending[id]
if !ok {
continue
}
delete(s.pending, id)
s.totalBytes -= len(frame.body)
if frame.reqMsgID != 0 {
delete(s.byRequest, frame.reqMsgID)
}
s.markAcked(id)
}
if len(s.order) > maxTrackedServerMsgIDs*2 {
s.compactOrder()
}
}
func (s *outboundState) stateInfo(ids []int64) []byte {
info := make([]byte, len(ids))
for i, id := range ids {
if s.isKnown(id) {
info[i] = msgStateReceived
}
}
return info
}
func (s *outboundState) isKnown(id int64) bool {
if _, ok := s.pending[id]; ok {
return true
}
_, ok := s.acked[id]
return ok
}
func (s *outboundState) markAcked(id int64) {
if _, ok := s.acked[id]; ok {
return
}
s.acked[id] = struct{}{}
s.ackOrder = append(s.ackOrder, id)
for len(s.ackOrder) > maxTrackedAckedMsgIDs {
oldest := s.ackOrder[0]
s.ackOrder = s.ackOrder[1:]
delete(s.acked, oldest)
}
}
func (s *outboundState) shrinkPending() int {
dropped := 0
for (len(s.pending) > maxTrackedServerMsgIDs || s.totalBytes > maxTrackedServerBytes) && len(s.order) > 0 {
oldest := s.order[0]
s.order = s.order[1:]
frame, ok := s.pending[oldest]
if !ok {
continue
}
delete(s.pending, oldest)
s.totalBytes -= len(frame.body)
if frame.reqMsgID != 0 {
delete(s.byRequest, frame.reqMsgID)
}
dropped++
}
return dropped
}
func (s *outboundState) compactOrder() {
filtered := s.order[:0]
for _, id := range s.order {
if _, ok := s.pending[id]; ok {
filtered = append(filtered, id)
}
}
s.order = filtered
}

130
internal/mtprotoedge/outbound_test.go Normal file
View file

@ -0,0 +1,130 @@
package mtprotoedge
import (
"context"
"sync"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
func TestOutboundActorSerializesConcurrentSends(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)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
srv.Conns().SetReceivesUpdates(auth.SessionID, true)
const sends = 64
var wg sync.WaitGroup
errs := make(chan error, sends)
for i := 0; i < sends; i++ {
wg.Add(1)
go func() {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
errs <- srv.Conns().PushToSession(ctx, auth.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{})
}()
}
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatalf("push: %v", err)
}
}
var prevMsgID int64
var prevSeqNo int32 = -1
for i := 0; i < sends; i++ {
data, id, _ := readServerMessage(t, conn, cipher, auth.AuthKey)
if id != tg.UpdatesTooLongTypeID {
t.Fatalf("message %d type = %#x, want updatesTooLong", i, id)
}
if i > 0 && data.MessageID <= prevMsgID {
t.Fatalf("message %d msg_id = %d after %d, want strictly increasing", i, data.MessageID, prevMsgID)
}
if data.SeqNo%2 != 1 {
t.Fatalf("message %d seq_no = %d, want odd content-related seq_no", i, data.SeqNo)
}
if i > 0 && data.SeqNo <= prevSeqNo {
t.Fatalf("message %d seq_no = %d after %d, want increasing", i, data.SeqNo, prevSeqNo)
}
prevMsgID = data.MessageID
prevSeqNo = data.SeqNo
}
}
func TestOutboundResendAndAckState(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)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
srv.Conns().SetReceivesUpdates(auth.SessionID, true)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
if err := srv.Conns().PushToSession(ctx, auth.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
cancel()
t.Fatalf("push: %v", err)
}
cancel()
original, id, _ := readServerMessage(t, conn, cipher, auth.AuthKey)
if id != tg.UpdatesTooLongTypeID {
t.Fatalf("pushed type = %#x, want updatesTooLong", id)
}
resendReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, resendReqID, 3, &mt.MsgResendReq{MsgIDs: []int64{original.MessageID}})
resent, resentType, _ := readServerMessage(t, conn, cipher, auth.AuthKey)
if resentType != tg.UpdatesTooLongTypeID {
t.Fatalf("resent type = %#x, want updatesTooLong", resentType)
}
if resent.MessageID != original.MessageID || resent.SeqNo != original.SeqNo {
t.Fatalf("resent frame = (msg_id=%d seq=%d), want original (msg_id=%d seq=%d)",
resent.MessageID, resent.SeqNo, original.MessageID, original.SeqNo)
}
_, stateType, stateBuf := readServerMessage(t, conn, cipher, auth.AuthKey)
if stateType != mt.MsgsStateInfoTypeID {
t.Fatalf("state type = %#x, want msgs_state_info", stateType)
}
assertStateInfo(t, stateBuf, resendReqID, []byte{msgStateReceived})
_, ackType, _ := readServerMessage(t, conn, cipher, auth.AuthKey)
if ackType != mt.MsgsAckTypeID {
t.Fatalf("ack type = %#x, want msgs_ack", ackType)
}
sendEncryptedWithSeq(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), 4, &mt.MsgsAck{MsgIDs: []int64{original.MessageID}})
ackedResendReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, ackedResendReqID, 5, &mt.MsgResendReq{MsgIDs: []int64{original.MessageID}})
_, ackedStateType, ackedStateBuf := readServerMessage(t, conn, cipher, auth.AuthKey)
if ackedStateType != mt.MsgsStateInfoTypeID {
t.Fatalf("after ack type = %#x, want msgs_state_info without resend", ackedStateType)
}
assertStateInfo(t, ackedStateBuf, ackedResendReqID, []byte{msgStateReceived})
}
func assertStateInfo(t *testing.T, b *bin.Buffer, reqMsgID int64, want []byte) {
t.Helper()
var info mt.MsgsStateInfo
if err := info.Decode(b); err != nil {
t.Fatalf("decode msgs_state_info: %v", err)
}
if info.ReqMsgID != reqMsgID {
t.Fatalf("msgs_state_info.req_msg_id = %d, want %d", info.ReqMsgID, reqMsgID)
}
if string(info.Info) != string(want) {
t.Fatalf("msgs_state_info.info = %v, want %v", []byte(info.Info), want)
}
}

145
internal/mtprotoedge/rpc_test.go Normal file
View file

@ -0,0 +1,145 @@
package mtprotoedge
import (
"context"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/clock"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"telesrv/internal/rpc"
)
// TestRPCGetConfig 验证 M3：握手后 client 加密 help.getConfig，
// server 经 tg.ServerDispatcher 路由并回 rpc_result（含本地 DC），外加 new_session_created + ack。
func TestRPCGetConfig(t *testing.T) {
const (
dc = 2
advIP = "127.0.0.1"
advPort = 12345
)
router := rpc.New(rpc.Config{DC: dc, IP: advIP, Port: advPort}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: router})
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{})
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
if _, ok := replies[mt.MsgsAckTypeID]; !ok {
for id, b := range collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID) {
replies[id] = b
}
}
mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created")
mustHave(t, replies, mt.MsgsAckTypeID, "msgs_ack")
resultBuf := mustHave(t, replies, proto.ResultTypeID, "rpc_result")
var res proto.Result
if err := res.Decode(resultBuf); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if res.RequestMessageID != reqMsgID {
t.Fatalf("rpc_result req_msg_id = %d, want %d", res.RequestMessageID, reqMsgID)
}
var cfg tg.Config
if err := cfg.Decode(&bin.Buffer{Buf: res.Result}); err != nil {
t.Fatalf("decode config: %v", err)
}
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)
}
}
func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
const dc = 2
handler := &blockingRPC{
started: make(chan struct{}, 1),
release: make(chan struct{}),
}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCTimeout: 5 * time.Second,
})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
firstReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, firstReqID, 1, &tg.HelpGetConfigRequest{})
select {
case <-handler.started:
case <-time.After(time.Second):
t.Fatal("timed out waiting for first rpc to start")
}
secondReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, secondReqID, 3, &tg.HelpGetConfigRequest{})
thirdReqID := clientMsgID.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, thirdReqID, 5, &tg.HelpGetConfigRequest{})
result := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, thirdReqID)
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode rpc_error: %v", err)
}
if rpcErr.ErrorCode != 420 || rpcErr.ErrorMessage != "FLOOD_WAIT_1" {
t.Fatalf("rpc_error = %d %q, want 420 FLOOD_WAIT_1", rpcErr.ErrorCode, rpcErr.ErrorMessage)
}
close(handler.release)
}
type blockingRPC struct {
started chan struct{}
release chan struct{}
}
func (h *blockingRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
select {
case h.started <- struct{}{}:
default:
}
select {
case <-h.release:
return &tg.Config{ThisDC: 2}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
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++ {
_, id, plain := readServerMessage(t, conn, cipher, key)
if id != proto.ResultTypeID {
continue
}
var result proto.Result
if err := result.Decode(plain); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if result.RequestMessageID == reqMsgID {
return result
}
}
t.Fatalf("missing rpc_result for req_msg_id %d", reqMsgID)
return proto.Result{}
}

63
internal/mtprotoedge/rsakey.go Normal file
View file

@ -0,0 +1,63 @@
package mtprotoedge
import (
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/pem"
"errors"
"fmt"
"os"
"path/filepath"
)
// rsaKeyBits 是 server RSA 私钥位数。MTProto 要求 2048-bit。
const rsaKeyBits = 2048
// LoadOrGenerateRSAKey 从 path 加载 PEM 编码的 server RSA 私钥；
// 不存在则生成 2048-bit 新密钥并持久化（含父目录）。
//
// server RSA 私钥用于 MTProto 密钥交换；其公钥 fingerprint 需 patch 进 TDesktop
// （记录于 docs/tdesktop-patch-notes.md）。
func LoadOrGenerateRSAKey(path string) (*rsa.PrivateKey, error) {
data, err := os.ReadFile(path)
switch {
case err == nil:
key, perr := parseRSAKeyPEM(data)
if perr != nil {
return nil, fmt.Errorf("parse %q: %w", path, perr)
}
return key, nil
case errors.Is(err, os.ErrNotExist):
// 继续生成。
default:
return nil, fmt.Errorf("read %q: %w", path, err)
}
key, err := rsa.GenerateKey(rand.Reader, rsaKeyBits)
if err != nil {
return nil, fmt.Errorf("generate rsa key: %w", err)
}
if dir := filepath.Dir(path); dir != "" && dir != "." {
if err := os.MkdirAll(dir, 0o700); err != nil {
return nil, fmt.Errorf("create key dir %q: %w", dir, err)
}
}
pemBytes := pem.EncodeToMemory(&pem.Block{
Type: "RSA PRIVATE KEY",
Bytes: x509.MarshalPKCS1PrivateKey(key),
})
if err := os.WriteFile(path, pemBytes, 0o600); err != nil {
return nil, fmt.Errorf("write key %q: %w", path, err)
}
return key, nil
}
func parseRSAKeyPEM(data []byte) (*rsa.PrivateKey, error) {
block, _ := pem.Decode(data)
if block == nil {
return nil, errors.New("no PEM block found")
}
return x509.ParsePKCS1PrivateKey(block.Bytes)
}

364
internal/mtprotoedge/server.go Normal file
View file

@ -0,0 +1,364 @@
package mtprotoedge
import (
"context"
"crypto/rsa"
"errors"
"fmt"
"io"
"net"
"sync"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/clock"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/tg"
"github.com/gotd/td/tmap"
"github.com/gotd/td/transport"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
// RPCHandler 把解密后的 RPC 请求体路由到响应。由 internal/rpc 实现。
//
// b 是明文 RPC 请求（已剥离 MTProto 外壳）；返回的 bin.Encoder 会被包成 rpc_result。
// 返回 *tgerr.Error 时连接层将其转为 rpc_error 回发；其他 error 视为连接级故障。
type RPCHandler interface {
Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error)
}
// Options 配置 Server。
type Options struct {
// Logger 日志器。默认 zap.NewNop()。
Logger *zap.Logger
// Codec 传输 codec 构造器。nil 表示自动探测（intermediate/abridged/full）。
Codec func() transport.Codec
// ObfuscatedTCP 先按 MTProto TCP obfuscation 解包，再自动探测 codec。
// Telegram Desktop 的 tcpo_only endpoint 会走这个 64 字节前缀流程。
ObfuscatedTCP bool
// ReadTimeout 单次读取超时。默认 5m。
ReadTimeout time.Duration
// HandshakeIdleTimeout 是连接「建立 session 前」（握手 + 首个加密消息之前）的读超时，
// 比 ReadTimeout 短，用于快速回收握手后静默的半开 / 异常连接。默认 60s。
HandshakeIdleTimeout time.Duration
// WriteTimeout 单次写入超时。默认 30s。
WriteTimeout time.Duration
// RPCMaxInflight 是单连接同时处理的 RPC 上限。默认 32。
RPCMaxInflight int
// RPCQueueSize 是单连接等待处理的 RPC 队列长度。默认 256。
RPCQueueSize int
// RPCTimeout 是单个 RPC 在连接层的最大处理时长。默认 30s。
RPCTimeout time.Duration
// DC 是本 server 的 DC ID。默认 2。
DC int
// RSAKey 是 server RSA 私钥，用于密钥交换。nil 时无法完成握手。
RSAKey *rsa.PrivateKey
// AuthKeys 持久化 auth key。默认内存实现。
AuthKeys store.AuthKeyStore
// Sessions 记录在线 MTProto session（持久化数据）。默认内存实现。
Sessions store.SessionStore
// ActiveSessions 管理活跃连接。默认新建；传入时可让 RPC 层共享同一注册表。
ActiveSessions *SessionManager
// RPC 是 typed RPC 路由。nil 时加密 RPC 被丢弃并记录。
RPC RPCHandler
// Metrics 接收连接层指标。默认 NopMetrics。
Metrics Metrics
// Clock 用于消息 ID 与时间戳。默认 clock.System。
Clock clock.Clock
// Rand 随机源。默认 crypto.DefaultRand()。
Rand io.Reader
}
func (o *Options) setDefaults() {
if o.Logger == nil {
o.Logger = zap.NewNop()
}
if o.ReadTimeout == 0 {
o.ReadTimeout = 5 * time.Minute
}
if o.HandshakeIdleTimeout == 0 {
o.HandshakeIdleTimeout = 60 * time.Second
}
if o.WriteTimeout == 0 {
o.WriteTimeout = 30 * time.Second
}
if o.RPCMaxInflight <= 0 {
o.RPCMaxInflight = 32
}
if o.RPCQueueSize <= 0 {
o.RPCQueueSize = 256
}
if o.RPCTimeout == 0 {
o.RPCTimeout = 30 * time.Second
}
if o.DC == 0 {
o.DC = 2
}
if o.AuthKeys == nil {
o.AuthKeys = memory.NewAuthKeyStore()
}
if o.Sessions == nil {
o.Sessions = memory.NewSessionStore()
}
if o.Metrics == nil {
o.Metrics = NopMetrics{}
}
if o.Clock == nil {
o.Clock = clock.System
}
if o.Rand == nil {
o.Rand = crypto.DefaultRand()
}
}
// Server 是 MTProto 连接层（mtprotoedge）。
//
// 职责见 doc.go。它把原始 TCP 字节流转换为「已解密、已识别 session 的 RPC 请求」：
// 接受连接、协商 codec、完成密钥交换、解密并分发加密消息到 RPC 路由，处理服务消息，
// 并把活跃连接注册到 SessionManager 以支持主动推送（updates 等）。不含业务逻辑。
type Server struct {
log *zap.Logger
codec func() transport.Codec
obfuscated bool
readTimeout time.Duration
handshakeTimeout time.Duration
writeTimeout time.Duration
rpcInflight int
rpcQueueSize int
rpcTimeout time.Duration
dc int
key exchange.PrivateKey
authKeys store.AuthKeyStore
sessions store.SessionStore
conns *SessionManager
rpc RPCHandler
metrics Metrics
cipher crypto.Cipher
clock clock.Clock
rand io.Reader
types *tmap.Map
// sessionUID 是本进程 server session 唯一标识，写入 new_session_created。
sessionUID int64
// onFrame 是测试钩子：收到一帧时回调其字节数；生产为 nil。
onFrame func(n int)
}
// New 创建 Server。
func New(opts Options) *Server {
opts.setDefaults()
conns := opts.ActiveSessions
if conns == nil {
conns = NewSessionManager(opts.Logger.Named("sessions"))
}
return &Server{
log: opts.Logger,
codec: opts.Codec,
obfuscated: opts.ObfuscatedTCP,
readTimeout: opts.ReadTimeout,
handshakeTimeout: opts.HandshakeIdleTimeout,
writeTimeout: opts.WriteTimeout,
rpcInflight: opts.RPCMaxInflight,
rpcQueueSize: opts.RPCQueueSize,
rpcTimeout: opts.RPCTimeout,
dc: opts.DC,
key: exchange.PrivateKey{RSA: opts.RSAKey},
authKeys: opts.AuthKeys,
sessions: opts.Sessions,
conns: conns,
rpc: opts.RPC,
metrics: opts.Metrics,
cipher: crypto.NewServerCipher(opts.Rand),
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
sessionUID: opts.Clock.Now().UnixNano(),
}
}
// Conns 返回活跃连接注册表，供业务层主动推送（updates 等）。
func (s *Server) Conns() *SessionManager {
return s.conns
}
// newConn 基于一次解密结果创建一个可发送的连接对象。
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
c := &Conn{
transport: tc,
writer: tc,
cipher: s.cipher,
msgID: proto.NewMessageIDGen(s.clock.Now),
writeTimeout: s.writeTimeout,
metrics: s.metrics,
authKeyID: key.ID,
sessionID: sessionID,
salt: salt,
key: key,
}
c.startOutbound()
c.startInboundRPCScheduler(s.rpcInflight, s.rpcQueueSize, s.rpcTimeout)
return c
}
// Serve 在 ln 上运行 MTProto 连接循环，直到 ctx 取消或发生不可恢复错误。
// ctx 取消时优雅退出：关闭 listener 并等待在途连接处理结束。
func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
transportListener := ln
if s.obfuscated {
transportListener = transport.ObfuscatedListener(ln)
}
l := transport.ListenCodec(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 阻塞。
go func() {
<-ctx.Done()
_ = l.Close()
}()
var wg sync.WaitGroup
defer wg.Wait()
for {
conn, err := l.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))
}
}()
}
}
// serveConn 处理单个传输连接：读帧并按 auth_key_id 分流。
//
// - auth_key_id == 0：未加密的密钥交换起始消息，执行握手并落地 auth key。
// - auth_key_id 已注册：加密消息，解密、注册连接并分发到 RPC 路由。
// - auth_key_id 未注册：回 AuthKeyNotFound，促使客户端重新握手。
//
// 连接建立 session 后注册到 SessionManager，结束时注销。
func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) {
s.metrics.ConnOpened()
s.log.Debug("Connection accepted")
var current *Conn
defer func() {
if current != nil {
s.conns.Unregister(current)
current.Close()
}
s.metrics.ConnClosed()
s.log.Debug("Connection closed", zap.Error(err))
}()
// ctx 取消或处理结束时关闭连接，解除 Recv 阻塞。
ctx, cancel := context.WithCancel(ctx)
defer cancel()
go func() {
<-ctx.Done()
_ = conn.Close()
}()
cs := newConnState()
var b bin.Buffer
var replay *bin.Buffer
for {
if replay != nil {
b.ResetTo(replay.Copy())
replay = nil
} else {
// 建立 session 前（current==nil，握手 + 首个加密消息之前）用较短的 handshakeTimeout
// 快速回收静默的半开 / 异常连接；建立 session 后用 readTimeout（客户端有 ping 心跳）。
timeout := s.readTimeout
if current == nil {
timeout = s.handshakeTimeout
}
if err := s.recv(ctx, conn, &b, timeout); err != nil {
return err
}
if s.onFrame != nil {
s.onFrame(b.Len())
}
}
authKeyID, err := peekAuthKeyID(&b)
if err != nil {
return fmt.Errorf("peek auth key id: %w", err)
}
if authKeyID == emptyAuthKeyID {
next, err := s.handleExchange(ctx, conn, &b)
if err != nil {
return err
}
replay = next
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
}
continue
}
current, err = s.handleEncrypted(ctx, conn, cs, current, data, &b)
if err != nil {
return err
}
}
}
func (s *Server) recv(ctx context.Context, conn transport.Conn, b *bin.Buffer, timeout time.Duration) error {
b.Reset()
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
return conn.Recv(ctx, b)
}
// isClientDisconnect 判断错误是否为正常的客户端断开/服务关闭，不应作为异常记录。
func isClientDisconnect(err error) bool {
switch {
case errors.Is(err, io.EOF),
errors.Is(err, net.ErrClosed),
errors.Is(err, context.Canceled),
errors.Is(err, context.DeadlineExceeded):
return true
}
var nerr *net.OpError
if errors.As(err, &nerr) && (nerr.Op == "read" || nerr.Op == "write") {
return true
}
return false
}

165
internal/mtprotoedge/server_test.go Normal file
View file

@ -0,0 +1,165 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/mtproxy"
"github.com/gotd/td/mtproxy/obfuscator"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
)
// TestServerAcceptAndCodec 验证 M0：
// server 能接受连接、自动协商 codec、读到客户端帧，并在 ctx 取消时优雅退出。
func TestServerAcceptAndCodec(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)})
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) }()
// 客户端：TCP 拨号 + intermediate 协议握手 + 发送一帧。
raw, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := transport.Intermediate.Handshake(raw)
if err != nil {
t.Fatalf("handshake: %v", err)
}
// payload 必须 ≠ 4 字节：codec 把恰好 4 字节的帧当作 transport 协议错误码（checkProtocolError）。
// 真实 MTProto 帧远大于 4 字节，这里发 8 字节模拟一个普通帧。
var b bin.Buffer
b.PutInt32(0x12345678)
b.PutInt32(0x0badf00d)
sendCtx, sc := context.WithTimeout(context.Background(), 5*time.Second)
defer sc()
if err := conn.Send(sendCtx, &b); err != nil {
t.Fatalf("send: %v", err)
}
select {
case n := <-frames:
if n <= 0 {
t.Fatalf("received empty frame, len = %d", n)
}
case <-time.After(5 * time.Second):
t.Fatal("server did not receive frame in time")
}
_ = conn.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")
}
}
// TestServerAcceptObfuscatedAbridged 验证 TDesktop tcpo_only 连接形态：
// 先做 MTProto TCP obfuscation，再在解密后的流上使用 abridged codec。
func TestServerAcceptObfuscatedAbridged(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})
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) }()
bad, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("bad dial: %v", err)
}
_ = bad.Close()
time.Sleep(50 * time.Millisecond)
select {
case err := <-serveErr:
t.Fatalf("server stopped after bad obfuscated accept: %v", err)
default:
}
raw, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("dial: %v", err)
}
obfs := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfs.Handshake((codec.Abridged{}).ObfuscatedTag(), 2, mtproxy.Secret{}); err != nil {
t.Fatalf("obfuscated handshake: %v", err)
}
conn, err := transport.NewProtocol(func() transport.Codec {
return transport.Abridged.CodecNoHeader()
}).Handshake(obfs)
if err != nil {
t.Fatalf("transport handshake: %v", err)
}
var b bin.Buffer
b.PutInt32(0x12345678)
b.PutInt32(0x0badf00d)
sendCtx, sc := context.WithTimeout(context.Background(), 5*time.Second)
defer sc()
if err := conn.Send(sendCtx, &b); err != nil {
t.Fatalf("send: %v", err)
}
select {
case n := <-frames:
if n <= 0 {
t.Fatalf("received empty frame, len = %d", n)
}
case <-time.After(5 * time.Second):
t.Fatal("server did not receive frame in time")
}
_ = conn.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")
}
}

978
internal/mtprotoedge/session_manager.go Normal file
View file

@ -0,0 +1,978 @@
package mtprotoedge
import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
)
// ErrSessionNotFound 表示目标 session 当前无活跃连接。
var ErrSessionNotFound = errors.New("session not found")
// ErrSessionAmbiguous 表示仅用 session_id 无法唯一定位连接。
var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys")
const (
maxPendingPushesPerSession = 32
// pendingPushMaxAge：session 注册后迟迟不调 updates.getState（receivesUpdates 恒 false）时，
// 其暂存的主动推送最长保留时长。超过即丢整批并不再囤——正常 TDesktop 登录后秒级就会
// getState 建立同步基线；长期不 ready 多为异常/对抗连接。丢弃不丢消息：getDifference 以
// user_update_events durable log 兜底补齐。
pendingPushMaxAge = 60 * time.Second
)
type queuedPush struct {
t proto.MessageType
msg bin.Encoder
at time.Time
}
type sessionKey struct {
authKeyID [8]byte
sessionID int64
}
// SessionLifecycleObserver receives active connection lifecycle events.
type SessionLifecycleObserver interface {
SessionOffline(rawAuthKeyID [8]byte, sessionID, userID int64, lastForUser bool)
}
// SessionManager 是活跃连接注册表，支持按 session / auth-key / user 查找并主动 push。
//
// 它管理运行态的在线连接，与持久化的 store.SessionStore 互补：后者记录 session 数据，
// 前者持有可发送的活跃连接。所有方法并发安全。
type SessionManager struct {
mu sync.RWMutex
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
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 前暂存的主动推送
lifecycle SessionLifecycleObserver
log *zap.Logger
}
// NewSessionManager 创建空的连接注册表。
func NewSessionManager(log *zap.Logger) *SessionManager {
if log == nil {
log = zap.NewNop()
}
return &SessionManager{
bySession: make(map[sessionKey]*Conn),
bySessionID: make(map[int64]map[[8]byte]*Conn),
byAuthKey: make(map[[8]byte]map[int64]*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),
log: log,
}
}
// SetLifecycleObserver installs a best-effort active session lifecycle observer.
func (m *SessionManager) SetLifecycleObserver(observer SessionLifecycleObserver) {
m.mu.Lock()
m.lifecycle = observer
m.mu.Unlock()
}
// Register 注册一个活跃连接。若同 raw auth_key_id + session_id 已存在（重连），旧连接被替换并移除索引。
func (m *SessionManager) Register(c *Conn) {
m.mu.Lock()
key := connSessionKey(c)
var replaced *Conn
if old, ok := m.bySession[key]; ok && old != c {
replaced = old
m.removeLocked(old, false)
}
m.bySession[key] = c
addSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID, c)
addConnIndex(m.byAuthKey, c.authKeyID, c.sessionID, c)
if uid := c.userID.Load(); uid != 0 {
c.userIDResolved.Store(true)
addUserIndex(m.byUser, uid, key, c)
}
m.log.Debug("Session registered",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", c.sessionID),
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
if replaced != nil {
replaced.Close()
}
}
// Unregister 注销一个连接（仅当它仍是当前注册的同一对象，避免误删重连后的新连接）。
func (m *SessionManager) Unregister(c *Conn) {
m.mu.Lock()
var (
observer SessionLifecycleObserver
offlineUser int64
lastForUser bool
)
if cur, ok := m.bySession[connSessionKey(c)]; ok && cur == c {
offlineUser = m.removeLocked(c, true)
if offlineUser != 0 {
lastForUser = len(m.byUser[offlineUser]) == 0
observer = m.lifecycle
}
m.log.Debug("Session unregistered",
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int64("session_id", c.sessionID),
zap.Int("online", len(m.bySession)),
)
}
m.mu.Unlock()
if observer != nil && offlineUser != 0 {
observer.SessionOffline(c.authKeyID, c.sessionID, offlineUser, lastForUser)
}
}
// DestroySession 移除指定 session 的运行态索引，供 MTProto destroy_session 使用。
func (m *SessionManager) DestroySession(sessionID int64) bool {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if !ambiguous {
m.dropPendingBySessionLocked(sessionID)
}
m.mu.Unlock()
return false
}
offlineUser := m.removeLocked(c, true)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
observer := m.lifecycle
m.log.Debug("Session destroyed",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", sessionID),
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
c.Close()
if observer != nil && offlineUser != 0 {
observer.SessionOffline(key.authKeyID, sessionID, offlineUser, lastForUser)
}
return true
}
// DestroySessionForAuthKey 精确移除某个 raw auth_key_id 下的 session。
func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID int64) bool {
m.mu.Lock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
delete(m.pending, key)
m.mu.Unlock()
return false
}
offlineUser := m.removeLocked(c, true)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
observer := m.lifecycle
m.log.Debug("Session destroyed",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int64("session_id", sessionID),
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
c.Close()
if observer != nil && offlineUser != 0 {
observer.SessionOffline(authKeyID, sessionID, offlineUser, lastForUser)
}
return true
}
// BindUser 缓存 session 的授权用户。userID=0 表示当前 auth_key 已确认未登录。
// 登录后绑定非 0 userID，使其可经 PushToUser 收到推送。
func (m *SessionManager) BindUser(sessionID, userID int64) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if ambiguous {
m.log.Warn("Skip BindUser for ambiguous session_id", zap.Int64("session_id", sessionID))
}
m.mu.Unlock()
return
}
m.bindUserLocked(c, key, userID)
m.mu.Unlock()
}
// BindUserForAuthKey 缓存指定 raw auth_key_id + session_id 的授权用户。
func (m *SessionManager) BindUserForAuthKey(authKeyID [8]byte, sessionID, userID int64) {
m.mu.Lock()
defer m.mu.Unlock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
return
}
m.bindUserLocked(c, key, userID)
}
func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
if old := c.userID.Swap(userID); old != 0 {
removeUserIndex(m.byUser, old, key)
if old != userID {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
}
}
c.userIDResolved.Store(true)
if userID != 0 {
addUserIndex(m.byUser, userID, key, c)
} else {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
}
}
// UserID 返回 session 当前缓存的登录用户 id。未绑定或离线时 ok=false。
func (m *SessionManager) UserID(sessionID int64) (int64, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return 0, false
}
userID := c.userID.Load()
if userID == 0 {
return 0, false
}
return userID, true
}
// UserIDForAuthKey 返回指定 raw auth_key_id + session_id 当前缓存的登录用户 id。
func (m *SessionManager) UserIDForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) {
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
if !ok {
return 0, false
}
userID := c.userID.Load()
if userID == 0 {
return 0, false
}
return userID, true
}
// UserIDResolved 返回 session 的 user_id 授权状态是否已经查过。
// resolved=true 且 userID=0 表示该 session 当前未登录。
func (m *SessionManager) UserIDResolved(sessionID int64) (int64, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return 0, false
}
return c.UserIDResolved()
}
// UserIDResolvedForAuthKey 返回指定 raw auth_key_id + session_id 的 user_id 缓存状态。
func (m *SessionManager) UserIDResolvedForAuthKey(authKeyID [8]byte, sessionID int64) (int64, bool) {
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
if !ok {
return 0, false
}
return c.UserIDResolved()
}
// BindAuthKey 缓存业务视角 auth_key_id（temp auth_key 解析后的 perm auth_key）。
func (m *SessionManager) BindAuthKey(sessionID int64, authKeyID [8]byte) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if ambiguous {
m.log.Warn("Skip BindAuthKey for ambiguous session_id", zap.Int64("session_id", sessionID))
}
m.mu.Unlock()
return
}
m.bindAuthKeyLocked(c, key, authKeyID)
m.mu.Unlock()
}
// BindAuthKeyForSession 缓存指定 raw auth_key_id + session_id 的业务 auth_key_id。
func (m *SessionManager) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte) {
m.mu.Lock()
defer m.mu.Unlock()
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
return
}
m.bindAuthKeyLocked(c, key, authKeyID)
}
func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8]byte) {
oldAuthKeyID, resolved := c.BusinessAuthKeyID()
changed := !resolved || oldAuthKeyID != authKeyID
oldUserID := c.userID.Load()
c.SetBusinessAuthKeyID(authKeyID)
if changed {
if oldUserID != 0 {
removeUserIndex(m.byUser, oldUserID, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.userID.Store(0)
c.userIDResolved.Store(false)
}
}
// AuthKeyID 返回 session 缓存的业务视角 auth_key_id。
// ok=false 表示该连接尚未完成 temp→perm 解析。
func (m *SessionManager) AuthKeyID(sessionID int64) ([8]byte, bool) {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return [8]byte{}, false
}
return c.BusinessAuthKeyID()
}
// AuthKeyIDForSession 返回指定 raw auth_key_id + session_id 缓存的业务 auth_key_id。
func (m *SessionManager) AuthKeyIDForSession(rawAuthKeyID [8]byte, sessionID int64) ([8]byte, bool) {
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}]
m.mu.RUnlock()
if !ok {
return [8]byte{}, false
}
return c.BusinessAuthKeyID()
}
// 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 {
if !connUsesBusinessAuthKey(c, authKeyID) {
continue
}
if old := c.userID.Swap(0); old != 0 {
removeUserIndex(m.byUser, old, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
c.userIDResolved.Store(true)
count++
}
return count
}
// SetReceivesUpdates 标记 session 是否已完成 updates 同步入口。
//
// TDesktop 登录后会先调用 updates.getState/getDifference 建立本地同步基线。
// 在此之前收到的主动 updates 先暂存，待 session 可接收后再异步下发。
func (m *SessionManager) SetReceivesUpdates(sessionID int64, receives bool) {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous || !ok {
if ambiguous {
m.log.Warn("Skip SetReceivesUpdates for ambiguous session_id", zap.Int64("session_id", sessionID))
}
m.mu.Unlock()
return
}
c.receivesUpdates.Store(receives)
if !receives {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
}
pending := m.takePendingLocked(key, receives)
m.mu.Unlock()
if len(pending) > 0 {
go m.flushPending(key, pending)
}
}
// SetReceivesUpdatesForAuthKey 标记指定 raw auth_key_id + session_id 是否接收主动 updates。
func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, sessionID int64, receives bool) {
m.mu.Lock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
m.mu.Unlock()
return
}
c.receivesUpdates.Store(receives)
if !receives {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
}
pending := m.takePendingLocked(key, receives)
m.mu.Unlock()
if len(pending) > 0 {
go m.flushPending(key, pending)
}
}
// PushToSession 向指定 session 推送一条消息。
func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.Lock()
c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID)
if ambiguous {
m.mu.Unlock()
return ErrSessionAmbiguous
}
if !ok {
m.mu.Unlock()
return ErrSessionNotFound
}
if !c.receivesUpdates.Load() {
m.queueLocked(key, t, msg)
m.mu.Unlock()
return nil
}
m.mu.Unlock()
return c.Send(ctx, t, msg)
}
// PushToSessionForAuthKey 向指定 raw auth_key_id + session_id 推送一条消息。
func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
m.mu.Lock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
m.mu.Unlock()
return ErrSessionNotFound
}
if !c.receivesUpdates.Load() {
m.queueLocked(key, t, msg)
m.mu.Unlock()
return nil
}
m.mu.Unlock()
return c.Send(ctx, t, msg)
}
// PushToUser 向某 user 所有活跃连接推送，返回已发送或已暂存的连接数。
// 发送在释放锁后进行，避免持锁阻塞于网络 IO。
func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.PushToUserExceptAuthKeySession(ctx, userID, [8]byte{}, 0, t, msg)
}
// PushToUserExceptSession 向某 user 所有活跃连接推送，但跳过指定 session。
// 未完成 updates 同步入口的 session 会先暂存，等 SetReceivesUpdates(true) 后再发。
func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg)
}
// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定业务 auth_key + session。
func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg)
}
func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
return c.Send(ctx, t, msg)
})
}
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)
}
func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, timeout)
}
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) {
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
return c.SendBestEffort(ctx, t, msg, timeout)
})
}
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) {
m.mu.Lock()
conns := make([]*Conn, 0, len(m.byUser[userID]))
queued := 0
for key, c := range m.byUser[userID] {
if shouldExcludeSession(c, excludeAuthKeyID, excludeSessionID) {
continue
}
if !c.receivesUpdates.Load() {
m.queueLocked(key, t, msg)
queued++
continue
}
conns = append(conns, c)
}
m.mu.Unlock()
var firstErr error
sent := 0
for _, c := range conns {
if err := send(c); err != nil {
if firstErr == nil {
firstErr = err
}
continue
}
sent++
}
return sent + queued, firstErr
}
// Online 返回当前活跃连接数。
func (m *SessionManager) Online() int {
m.mu.RLock()
defer m.mu.RUnlock()
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 {
return false
}
m.mu.RLock()
defer m.mu.RUnlock()
return len(m.byUser[userID]) > 0
}
// OnlineUserIDsForCandidates filters an explicit candidate set against the
// active user index. It avoids exporting or sorting the whole online map.
func (m *SessionManager) OnlineUserIDsForCandidates(candidateUserIDs []int64, limit int) []int64 {
if len(candidateUserIDs) == 0 {
return nil
}
m.mu.RLock()
defer m.mu.RUnlock()
out := make([]int64, 0, minInt(len(candidateUserIDs), positiveLimitOrLen(limit, len(candidateUserIDs))))
seen := make(map[int64]struct{}, len(candidateUserIDs))
for _, userID := range candidateUserIDs {
if userID == 0 {
continue
}
if _, ok := seen[userID]; ok {
continue
}
seen[userID] = struct{}{}
if len(m.byUser[userID]) == 0 {
continue
}
out = append(out, userID)
if limit > 0 && len(out) >= limit {
break
}
}
return out
}
// TrackChannelInterest replaces the channel viewer set for one live session.
// Realtime transient fan-out uses this as the current active-viewer candidate
// set; durable channel updates use the broader membership index instead.
func (m *SessionManager) TrackChannelInterest(rawAuthKeyID [8]byte, sessionID, userID int64, channelIDs []int64) {
if userID == 0 {
return
}
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
m.mu.Lock()
defer m.mu.Unlock()
c, ok := m.bySession[key]
if !ok || c.userID.Load() != userID {
return
}
m.clearChannelInterestsLocked(key)
if len(channelIDs) == 0 {
return
}
m.trackChannelIndexLocked(m.byChannel, m.bySessionChannels, key, userID, channelIDs)
}
// ClearChannelInterest removes the active-viewer channel set for one live
// session while leaving its joined-channel membership index intact.
func (m *SessionManager) ClearChannelInterest(rawAuthKeyID [8]byte, sessionID, userID int64) {
if userID == 0 {
return
}
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
m.mu.Lock()
defer m.mu.Unlock()
c, ok := m.bySession[key]
if !ok || c.userID.Load() != userID {
return
}
m.clearChannelInterestsLocked(key)
}
// OnlineChannelUserIDs returns users with active sessions that have recently
// proven current interest in channelID. The result is intentionally unsorted and bounded.
func (m *SessionManager) OnlineChannelUserIDs(channelID int64, limit int) []int64 {
return m.onlineChannelUsers(m.byChannel, channelID, limit)
}
// SetSessionChannelMemberships replaces the joined-channel index for one
// updates-ready session. This index is broader than TrackChannelInterest and is
// used for durable channel updates such as new/edit/delete message.
func (m *SessionManager) SetSessionChannelMemberships(rawAuthKeyID [8]byte, sessionID, userID int64, channelIDs []int64) {
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
m.mu.Lock()
defer m.mu.Unlock()
c, ok := m.bySession[key]
if !ok {
return
}
m.clearChannelMembershipsLocked(key)
if userID == 0 || c.userID.Load() != userID {
return
}
m.trackChannelIndexLocked(m.byMemberChannel, m.bySessionMembers, key, userID, channelIDs)
}
// AddUserChannelMembership adds channelID to every live session for userID.
// It is called after successful join/invite approval paths.
func (m *SessionManager) AddUserChannelMembership(userID, channelID int64) {
if userID == 0 || channelID == 0 {
return
}
m.mu.Lock()
defer m.mu.Unlock()
for key, c := range m.byUser[userID] {
if c == nil || c.userID.Load() != userID {
continue
}
m.trackChannelIndexLocked(m.byMemberChannel, m.bySessionMembers, key, userID, []int64{channelID})
}
}
// RemoveUserChannelMembership removes channelID from every live session for userID.
// It is called after leave/kick/ban/delete paths.
func (m *SessionManager) RemoveUserChannelMembership(userID, channelID int64) {
if userID == 0 || channelID == 0 {
return
}
m.mu.Lock()
defer m.mu.Unlock()
for key := range m.byUser[userID] {
m.removeChannelIndexLocked(m.byMemberChannel, m.bySessionMembers, key, channelID)
}
}
// OnlineChannelMemberUserIDs returns users with active sessions that are indexed
// as joined members of channelID. The result is intentionally unsorted; callers
// still verify business membership before pushing.
func (m *SessionManager) OnlineChannelMemberUserIDs(channelID int64, limit int) []int64 {
return m.onlineChannelUsers(m.byMemberChannel, channelID, limit)
}
func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64, channelID int64, limit int) []int64 {
if channelID == 0 {
return nil
}
m.mu.RLock()
defer m.mu.RUnlock()
sessions := index[channelID]
if len(sessions) == 0 {
return nil
}
out := make([]int64, 0, positiveLimitOrLen(limit, len(sessions)))
seen := make(map[int64]struct{}, len(sessions))
for key, userID := range sessions {
if userID == 0 {
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) removeLocked(c *Conn, dropPending bool) int64 {
key := connSessionKey(c)
delete(m.bySession, key)
removeSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID)
removeConnIndex(m.byAuthKey, c.authKeyID, c.sessionID)
uid := c.userID.Load()
if uid != 0 {
removeUserIndex(m.byUser, uid, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(key)
if dropPending {
delete(m.pending, key)
}
return uid
}
func (m *SessionManager) clearChannelInterestsLocked(key sessionKey) {
m.clearChannelIndexLocked(m.byChannel, m.bySessionChannels, key)
}
func (m *SessionManager) clearChannelMembershipsLocked(key sessionKey) {
m.clearChannelIndexLocked(m.byMemberChannel, m.bySessionMembers, key)
}
func (m *SessionManager) trackChannelIndexLocked(index map[int64]map[sessionKey]int64, reverse map[sessionKey]map[int64]struct{}, key sessionKey, userID int64, channelIDs []int64) {
channels := reverse[key]
if channels == nil {
channels = make(map[int64]struct{}, len(channelIDs))
reverse[key] = channels
}
for _, channelID := range channelIDs {
if channelID == 0 {
continue
}
channels[channelID] = struct{}{}
sessions := index[channelID]
if sessions == nil {
sessions = make(map[sessionKey]int64)
index[channelID] = sessions
}
sessions[key] = userID
}
}
func (m *SessionManager) clearChannelIndexLocked(index map[int64]map[sessionKey]int64, reverse map[sessionKey]map[int64]struct{}, key sessionKey) {
channels := reverse[key]
if len(channels) == 0 {
delete(reverse, key)
return
}
for channelID := range channels {
sessions := index[channelID]
delete(sessions, key)
if len(sessions) == 0 {
delete(index, channelID)
}
}
delete(reverse, key)
}
func (m *SessionManager) removeChannelIndexLocked(index map[int64]map[sessionKey]int64, reverse map[sessionKey]map[int64]struct{}, key sessionKey, channelID int64) {
channels := reverse[key]
delete(channels, channelID)
if len(channels) == 0 {
delete(reverse, key)
}
sessions := index[channelID]
delete(sessions, key)
if len(sessions) == 0 {
delete(index, channelID)
}
}
func positiveLimitOrLen(limit, length int) int {
if limit > 0 && limit < length {
return limit
}
return length
}
func minInt(a, b int) int {
if a < b {
return a
}
return b
}
func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedPush {
if !ready || len(m.pending[key]) == 0 {
return nil
}
pending := append([]queuedPush(nil), m.pending[key]...)
delete(m.pending, key)
return pending
}
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) {
q := m.pending[key]
// 过期保护：最早一条暂存已超过 pendingPushMaxAge（session 迟迟未 ready）时，丢整批并
// 不再囤这条，记 trace。避免「登录后从不 getState」的连接长期占用 pending 内存。
if len(q) > 0 && time.Since(q[0].at) > pendingPushMaxAge {
m.log.Debug("Drop stale pending pushes (session not ready in time)",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", key.sessionID),
zap.Int("dropped", len(q)),
)
delete(m.pending, key)
return
}
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
}
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),
)
}
}
}
func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey, bool, bool) {
set := m.bySessionID[sessionID]
if len(set) == 0 {
return nil, sessionKey{}, false, false
}
if len(set) > 1 {
return nil, sessionKey{}, false, true
}
for authKeyID, c := range set {
return c, sessionKey{authKeyID: authKeyID, sessionID: sessionID}, true, false
}
return nil, sessionKey{}, false, false
}
func (m *SessionManager) dropPendingBySessionLocked(sessionID int64) {
for key := range m.pending {
if key.sessionID == sessionID {
delete(m.pending, key)
}
}
}
func addConnIndex[K comparable](idx map[K]map[int64]*Conn, key K, sessionID int64, c *Conn) {
set := idx[key]
if set == nil {
set = make(map[int64]*Conn)
idx[key] = set
}
set[sessionID] = c
}
func removeConnIndex[K comparable](idx map[K]map[int64]*Conn, key K, sessionID int64) {
if set := idx[key]; set != nil {
delete(set, sessionID)
if len(set) == 0 {
delete(idx, key)
}
}
}
func addSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte, c *Conn) {
set := idx[sessionID]
if set == nil {
set = make(map[[8]byte]*Conn)
idx[sessionID] = set
}
set[authKeyID] = c
}
func removeSessionIDIndex(idx map[int64]map[[8]byte]*Conn, sessionID int64, authKeyID [8]byte) {
if set := idx[sessionID]; set != nil {
delete(set, authKeyID)
if len(set) == 0 {
delete(idx, sessionID)
}
}
}
func addUserIndex(idx map[int64]map[sessionKey]*Conn, userID int64, key sessionKey, c *Conn) {
set := idx[userID]
if set == nil {
set = make(map[sessionKey]*Conn)
idx[userID] = set
}
set[key] = c
}
func removeUserIndex(idx map[int64]map[sessionKey]*Conn, userID int64, key sessionKey) {
if set := idx[userID]; set != nil {
delete(set, key)
if len(set) == 0 {
delete(idx, userID)
}
}
}
func connSessionKey(c *Conn) sessionKey {
return sessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID}
}
func connUsesBusinessAuthKey(c *Conn, authKeyID [8]byte) bool {
id, resolved := c.BusinessAuthKeyID()
if resolved {
return id == authKeyID
}
return c.authKeyID == authKeyID
}
func shouldExcludeSession(c *Conn, excludeAuthKeyID *[8]byte, excludeSessionID int64) bool {
if excludeSessionID == 0 {
return false
}
if c.sessionID != excludeSessionID {
return false
}
if excludeAuthKeyID == nil || *excludeAuthKeyID == ([8]byte{}) {
return true
}
return connUsesBusinessAuthKey(c, *excludeAuthKeyID)
}
func sessionKeyLog(id [8]byte) string {
return fmt.Sprintf("%x", id)
}

215
internal/mtprotoedge/session_manager_bench_test.go Normal file
View file

@ -0,0 +1,215 @@
package mtprotoedge
import (
"context"
"fmt"
"os"
"runtime"
"sort"
"strconv"
"sync"
"sync/atomic"
"testing"
"time"
"go.uber.org/zap"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// 连接层 fan-out / churn 压测：聚焦 SessionManager 的锁争用，不走真实 socket / 加密。
//
// 构造的 Conn 故意不 startOutbound：pushToUser 持锁快照 byUser 后，锁外对每个 conn 调 c.Send，
// 此时 c.outbound==nil 立即返回 ErrConnClosed（见 outbound.go），因此测量集中在「持锁段 + 分发开销」，
// 即分片要消除的全局锁热点。每条连接仅结构体内存（无 1024 容量的 outbound channel、无 goroutine），
// 故可注册到 20 万规模。
//
// 用法：
//
// go test ./internal/mtprotoedge/ -run '^$' -bench BenchmarkSessionManager -benchmem -cpu 1,4,8
// go test ./internal/mtprotoedge/ -run '^$' -bench BenchmarkSessionManagerPushConcurrent -mutexprofile mu.out
// TELESRV_LOAD_CONNS=200000 go test ./internal/mtprotoedge/ -run TestSessionManagerFanoutThroughput -v -timeout 300s
func benchConn(sessionID int64, authKeyID [8]byte, userID int64) *Conn {
c := &Conn{sessionID: sessionID, authKeyID: authKeyID}
if userID != 0 {
c.userID.Store(userID)
c.userIDResolved.Store(true)
}
c.receivesUpdates.Store(true) // 走 fanout 的「收集 conns→锁外 Send」分支，而非 pending 暂存
return c
}
func authKeyIDFromInt(v uint64) [8]byte {
var id [8]byte
for i := 0; i < 8; i++ {
id[i] = byte(v >> (8 * i))
}
return id
}
// seedSessions 注册 conns 个连接，每个 user 绑定 connsPerUser 个连接（模拟多设备）。
// 返回注册的 userID 列表（去重、有序范围 [1, userCount]）。
func seedSessions(sm *SessionManager, conns, connsPerUser int) (userCount int) {
if connsPerUser < 1 {
connsPerUser = 1
}
for i := 0; i < conns; i++ {
userID := int64(i/connsPerUser) + 1
sm.Register(benchConn(int64(i)+1, authKeyIDFromInt(uint64(i)+1), userID))
}
return (conns + connsPerUser - 1) / connsPerUser
}
// BenchmarkSessionManagerPushConcurrent 模拟 20 万在线下的真实热点：大量 goroutine 并发对
// 不同 user pushToUser，全部抢同一把全局锁。-mutexprofile 会把 SessionManager.mu 顶上来。
func BenchmarkSessionManagerPushConcurrent(b *testing.B) {
const conns = 200_000
const connsPerUser = 2
sm := NewSessionManager(zap.NewNop())
userCount := seedSessions(sm, conns, connsPerUser)
msg := &tg.UpdatesTooLong{}
ctx := context.Background()
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
var n uint64
for pb.Next() {
n++
userID := int64(n%uint64(userCount)) + 1
_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
}
})
}
// BenchmarkSessionManagerRegisterChurn 测连接建立/断开的锁成本：并发 Register+Unregister。
// 20 万在线意味着持续的 connect/disconnect churn，每次都抢全局写锁。
func BenchmarkSessionManagerRegisterChurn(b *testing.B) {
sm := NewSessionManager(zap.NewNop())
var seq atomic.Uint64
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
id := seq.Add(1)
c := benchConn(int64(id), authKeyIDFromInt(id), int64(id))
sm.Register(c)
sm.Unregister(c)
}
})
}
// BenchmarkSessionManagerPushFanoutWidth 测单次 push 的 fanout 广度成本：一个 user 绑定很多连接，
// 单次 PushToUser 要持锁遍历全部。真实私聊 user 设备数少（2-4），此为上界参考。
func BenchmarkSessionManagerPushFanoutWidth(b *testing.B) {
for _, width := range []int{1, 4, 16, 64} {
b.Run(fmt.Sprintf("width=%d", width), func(b *testing.B) {
sm := NewSessionManager(zap.NewNop())
const userID = 1
for i := 0; i < width; i++ {
sm.Register(benchConn(int64(i)+1, authKeyIDFromInt(uint64(i)+1), userID))
}
msg := &tg.UpdatesTooLong{}
ctx := context.Background()
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
}
})
}
}
// TestSessionManagerFanoutThroughput 是数据驱动吞吐测：注册 N 连接后，P 个 goroutine 持续并发
// push，测全局锁下的实际 push 吞吐与 p99。默认小规模冒烟；设 TELESRV_LOAD_CONNS 放大到 20 万。
func TestSessionManagerFanoutThroughput(t *testing.T) {
conns := envIntDefault("TELESRV_LOAD_CONNS", 20_000)
connsPerUser := envIntDefault("TELESRV_LOAD_CONNS_PER_USER", 2)
workers := envIntDefault("TELESRV_LOAD_PUSH_WORKERS", 0) // 0 → GOMAXPROCS
duration := time.Duration(envIntDefault("TELESRV_LOAD_SECONDS", 3)) * time.Second
if workers <= 0 {
workers = runtime.GOMAXPROCS(0)
}
sm := NewSessionManager(zap.NewNop())
t0 := time.Now()
userCount := seedSessions(sm, conns, connsPerUser)
seedWall := time.Since(t0)
if got := sm.Online(); got != conns {
t.Fatalf("online = %d, want %d", got, conns)
}
msg := &tg.UpdatesTooLong{}
ctx := context.Background()
var ops atomic.Int64
perWorkerLat := make([][]time.Duration, workers)
deadline := time.Now().Add(duration)
var wg sync.WaitGroup
start := time.Now()
for w := 0; w < workers; w++ {
wg.Add(1)
go func(w int) {
defer wg.Done()
lat := make([]time.Duration, 0, 1<<16)
var n uint64
for time.Now().Before(deadline) {
// 批量 256 次再查一次时钟，降低 time.Now 占比。
for j := 0; j < 256; j++ {
n++
userID := int64(n%uint64(userCount)) + 1
s := time.Now()
_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
lat = append(lat, time.Since(s))
}
ops.Add(256)
}
perWorkerLat[w] = lat
}(w)
}
wg.Wait()
wall := time.Since(start)
all := make([]time.Duration, 0, ops.Load())
for _, l := range perWorkerLat {
all = append(all, l...)
}
sortDurations(all)
total := ops.Load()
thr := float64(total) / wall.Seconds()
t.Logf("==== session_manager fan-out throughput ====")
t.Logf("config: conns=%d connsPerUser=%d users=%d pushWorkers=%d dur=%s seed=%s",
conns, connsPerUser, userCount, workers, duration, seedWall.Round(time.Millisecond))
t.Logf("push: %d ops in %s -> %.0f push/s", total, wall.Round(time.Millisecond), thr)
t.Logf("push.lat p50=%s p90=%s p99=%s max=%s",
pct(all, 50), pct(all, 90), pct(all, 99), pct(all, 100))
t.Logf("=============================================")
}
func pct(sorted []time.Duration, p int) time.Duration {
if len(sorted) == 0 {
return 0
}
idx := (p*len(sorted))/100 - 1
if idx < 0 {
idx = 0
}
if idx >= len(sorted) {
idx = len(sorted) - 1
}
return sorted[idx]
}
func sortDurations(d []time.Duration) {
sort.Slice(d, func(i, j int) bool { return d[i] < d[j] })
}
func envIntDefault(key string, def int) int {
if v := os.Getenv(key); v != "" {
if n, err := strconv.Atoi(v); err == nil {
return n
}
}
return def
}

296
internal/mtprotoedge/session_manager_test.go Normal file
View file

@ -0,0 +1,296 @@
package mtprotoedge
import (
"context"
"errors"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestSessionManagerRegistry 验证注册表的注册/注销/查找语义（不涉及网络发送）。
func TestSessionManagerRegistry(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
c := &Conn{sessionID: 42, authKeyID: [8]byte{1, 2, 3}}
c.receivesUpdates.Store(true)
sm.Register(c)
if got := sm.Online(); got != 1 {
t.Fatalf("online = %d, want 1", got)
}
sm.BindAuthKey(42, [8]byte{1, 2, 3})
sm.BindUser(42, 100)
if userID, ok := sm.UserID(42); !ok || userID != 100 {
t.Fatalf("cached user = %d ok %v, want 100/true", userID, ok)
}
sm.BindAuthKey(42, [8]byte{9})
if userID, ok := sm.UserID(42); ok || userID != 0 {
t.Fatalf("cached user after auth key switch = %d ok %v, want 0/false", userID, ok)
}
if userID, resolved := sm.UserIDResolved(42); resolved || userID != 0 {
t.Fatalf("resolved user after auth key switch = %d resolved %v, want unresolved", userID, resolved)
}
sm.BindUser(42, 0)
if userID, resolved := sm.UserIDResolved(42); !resolved || userID != 0 {
t.Fatalf("negative user cache = %d resolved %v, want 0/true", userID, resolved)
}
sm.Unregister(c)
if got := sm.Online(); got != 0 {
t.Fatalf("online after unregister = %d, want 0", got)
}
err := sm.PushToSession(context.Background(), 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
if !errors.Is(err, ErrSessionNotFound) {
t.Fatalf("push to missing session err = %v, want ErrSessionNotFound", err)
}
}
func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw1 := [8]byte{1}
raw2 := [8]byte{2}
perm1 := [8]byte{9}
c1 := &Conn{sessionID: 42, authKeyID: raw1}
c2 := &Conn{sessionID: 42, authKeyID: raw2}
sm.Register(c1)
sm.Register(c2)
if got := sm.Online(); got != 2 {
t.Fatalf("online = %d, want 2", got)
}
sm.BindAuthKeyForSession(raw1, 42, perm1)
sm.BindUserForAuthKey(raw1, 42, 100)
sm.BindUserForAuthKey(raw2, 42, 200)
if userID, ok := sm.UserIDForAuthKey(raw1, 42); !ok || userID != 100 {
t.Fatalf("scoped user raw1 = %d ok %v, want 100/true", userID, ok)
}
if userID, ok := sm.UserIDForAuthKey(raw2, 42); !ok || userID != 200 {
t.Fatalf("scoped user raw2 = %d ok %v, want 200/true", userID, ok)
}
if _, ok := sm.UserID(42); ok {
t.Fatal("legacy UserID unexpectedly resolved ambiguous session_id")
}
if err := sm.PushToSession(context.Background(), 42, proto.MessageFromServer, &tg.UpdatesTooLong{}); !errors.Is(err, ErrSessionAmbiguous) {
t.Fatalf("ambiguous push err = %v, want ErrSessionAmbiguous", err)
}
sm.BindUserForAuthKey(raw1, 42, 300)
sm.BindUserForAuthKey(raw2, 42, 300)
sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, perm1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push except scoped session: %v", err)
}
if sent != 1 {
t.Fatalf("pushed to %d sessions, want 1", sent)
}
if _, ok := sm.pending[sessionKey{authKeyID: raw1, sessionID: 42}]; ok {
t.Fatal("excluded session received pending push")
}
if got := len(sm.pending[sessionKey{authKeyID: raw2, sessionID: 42}]); got != 1 {
t.Fatalf("raw2 pending pushes = %d, want 1", got)
}
}
func TestSessionManagerChannelInterestIndex(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.TrackChannelInterest(raw, 42, 100, []int64{10, 10, 20})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel 10 online users = %v, want [100]", got)
}
sm.TrackChannelInterest(raw, 42, 100, []int64{20})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel 10 after viewer switch = %v, want empty", got)
}
if got := sm.OnlineChannelUserIDs(20, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel 20 after viewer switch = %v, want [100]", got)
}
sm.TrackChannelInterest(raw, 42, 100, []int64{10})
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel 10 online members before membership sync = %v, want empty", got)
}
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{10, 30})
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel 10 online members = %v, want [100]", got)
}
if got := sm.OnlineChannelUserIDs(30, 10); len(got) != 0 {
t.Fatalf("channel 30 viewers = %v, want empty", got)
}
if got := sm.OnlineUserIDsForCandidates([]int64{0, 200, 100, 100}, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("candidate online users = %v, want [100]", got)
}
sm.BindUserForAuthKey(raw, 42, 200)
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel interest after user switch = %v, want empty", got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel membership after user switch = %v, want empty", got)
}
sm.TrackChannelInterest(raw, 42, 200, []int64{10})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 200 {
t.Fatalf("channel 10 after re-track = %v, want [200]", got)
}
sm.AddUserChannelMembership(200, 10)
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 200 {
t.Fatalf("channel 10 membership after add = %v, want [200]", got)
}
sm.RemoveUserChannelMembership(200, 10)
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel membership after remove = %v, want empty", got)
}
sm.ClearChannelInterest(raw, 42, 200)
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel interest after explicit clear = %v, want empty", got)
}
sm.Unregister(c)
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel interest after unregister = %v, want empty", got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("channel membership after unregister = %v, want empty", got)
}
}
func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
business := [8]byte{8}
c := &Conn{sessionID: 42, authKeyID: raw}
sm.Register(c)
sm.BindAuthKeyForSession(raw, 42, business)
sm.BindUserForAuthKey(raw, 42, 100)
track := func() {
sm.TrackChannelInterest(raw, 42, 100, []int64{10})
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{10})
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel viewers before cleanup = %v, want [100]", got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel members before cleanup = %v, want [100]", got)
}
}
assertCleared := func(label string) {
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
t.Fatalf("%s viewers = %v, want empty", label, got)
}
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("%s members = %v, want empty", label, got)
}
}
track()
sm.SetReceivesUpdatesForAuthKey(raw, 42, false)
assertCleared("after receivesUpdates=false")
track()
sm.BindAuthKeyForSession(raw, 42, [8]byte{9})
assertCleared("after business auth key change")
sm.BindAuthKeyForSession(raw, 42, business)
sm.BindUserForAuthKey(raw, 42, 100)
track()
if n := sm.UnbindAuthKey(business); n != 1 {
t.Fatalf("UnbindAuthKey count = %d, want 1", n)
}
assertCleared("after unbind auth key")
}
// TestSessionManagerPush 验证主动推送端到端：两个 client 连接握手并建立 session 后，
// server 经 PushToSession / PushToUser 主动向其推送，client 收到。
func TestSessionManagerPush(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})
conn1, auth1, cipher1 := dialHandshake(t, addr, dc, pub)
conn2, auth2, cipher2 := dialHandshake(t, addr, dc, pub)
// 各发一个 ping 建立 session，触发注册（并清掉 new_session_created/pong/ack）。
msgGen := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn1, cipher1, auth1, msgGen.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
collectReplies(t, conn1, cipher1, auth1.AuthKey, mt.PongTypeID)
sendEncrypted(t, conn2, cipher2, auth2, msgGen.New(proto.MessageFromClient), &mt.PingRequest{PingID: 2})
collectReplies(t, conn2, cipher2, auth2.AuthKey, mt.PongTypeID)
if got := srv.Conns().Online(); got != 2 {
t.Fatalf("online = %d, want 2", got)
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// 1) PushToSession：session2 尚未进入 updates 同步入口时先暂存，ready 后下发。
if err := srv.Conns().PushToSession(ctx, auth2.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
t.Fatalf("push to session: %v", err)
}
srv.Conns().SetReceivesUpdates(auth2.SessionID, true)
r2 := collectReplies(t, conn2, cipher2, auth2.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, r2, tg.UpdatesTooLongTypeID, "pushed updates on conn2")
// 2) BindUser + PushToUser：按 user 维度推送给 conn1。
srv.Conns().BindUser(auth1.SessionID, 100)
srv.Conns().SetReceivesUpdates(auth1.SessionID, true)
sent, err := srv.Conns().PushToUser(ctx, 100, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push to user: %v", err)
}
if sent != 1 {
t.Fatalf("pushed to %d conns, want 1", sent)
}
r1 := collectReplies(t, conn1, cipher1, auth1.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, r1, tg.UpdatesTooLongTypeID, "pushed updates on conn1")
// 3) PushToUserExceptSession：模拟 SyncUpdatesNotMe，跳过当前 session。
srv.Conns().BindUser(auth1.SessionID, 200)
srv.Conns().BindUser(auth2.SessionID, 200)
sent, err = srv.Conns().PushToUserExceptSession(ctx, 200, auth2.SessionID, proto.MessageFromServer, &tg.UpdatesTooLong{})
if err != nil {
t.Fatalf("push to user except session: %v", err)
}
if sent != 1 {
t.Fatalf("pushed to %d conns, want 1 after excluding current session", sent)
}
r1 = collectReplies(t, conn1, cipher1, auth1.AuthKey, tg.UpdatesTooLongTypeID)
mustHave(t, r1, tg.UpdatesTooLongTypeID, "pushed not-me updates on conn1")
}
func BenchmarkSessionManagerOnlineCandidateFilter(b *testing.B) {
sm := NewSessionManager(zaptest.NewLogger(b))
const online = 200_000
rawPrefix := [8]byte{9}
for i := 1; i <= online; i++ {
raw := rawPrefix
raw[1] = byte(i)
raw[2] = byte(i >> 8)
raw[3] = byte(i >> 16)
raw[4] = byte(i >> 24)
c := &Conn{sessionID: int64(i), authKeyID: raw}
sm.Register(c)
sm.BindUserForAuthKey(raw, int64(i), int64(i))
}
candidates := make([]int64, 0, 500)
for i := 0; i < 500; i++ {
candidates = append(candidates, int64(i*97+1))
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
got := sm.OnlineUserIDsForCandidates(candidates, 500)
if len(got) == 0 {
b.Fatal("no candidates matched")
}
}
}