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

fix: sync protocol and discussion stability fixes

Browse source
This commit is contained in:
A 2026-07-12 07:05:02 +08:00
parent 9f73dc20da
commit aa21bd04e1
43 changed files with 7258 additions and 503 deletions
Download patch file Download diff file Expand all files Collapse all files

93
internal/mtprotoedge/conn.go
View file

@ -25,13 +25,31 @@ type outboundWriter interface {
Send(context.Context, *bin.Buffer) error
}
type connLifecycle uint32
const (
// The zero value is deliberately provisional so test/embedded Conn values start
// outside every SessionManager index until they complete activation.
connLifecycleProvisional connLifecycle = iota
connLifecycleClaiming
connLifecycleActive
// retired is terminal and irreversible. A physical connection that lost an
// activation claim must never become visible again, even if its read goroutine
// was already between preflight and publish when a replacement arrived.
connLifecycleRetired
)
type Conn struct {
transport transport.Conn
writer outboundWriter
cipher crypto.Cipher
msgID *proto.MessageIDGen
writeTimeout time.Duration
metrics Metrics
transport transport.Conn
// transportLease owns exactly one generation of the physical transport.
// It is nil only for directly constructed test/embedded Conns that retain
// the legacy raw-transport close fallback.
transportLease *physicalTransportLease
writer outboundWriter
cipher crypto.Cipher
msgID *proto.MessageIDGen
writeTimeout time.Duration
metrics Metrics
authKeyID [8]byte
// authKeyHex 是 authKeyID 的 hex 缓存：每条 RPC 的结构化日志都会带它，
@ -66,7 +84,11 @@ type Conn struct {
outboundScratchOnce sync.Once
// terminal 表示该 logical Conn 已停止接受新的出站操作。写失败时由
// outbound actor 置位并只发停止信号，不能在 actor 内等待自身退出。
terminal atomic.Bool
terminal atomic.Bool
// lifecycle is a monotonic activation state machine. In particular, retired
// never transitions back to claiming/active; this closes the stale-read-loop
// ABA where an evicted Conn observed "not active" and registered itself again.
lifecycle atomic.Uint32
transportClose sync.Once
rpcScheduler *inboundRPCScheduler
@ -130,6 +152,63 @@ type Conn struct {
clientLayer atomic.Int32
}
func (c *Conn) lifecycleState() connLifecycle {
if c == nil {
return connLifecycleRetired
}
return connLifecycle(c.lifecycle.Load())
}
func (c *Conn) beginActivationClaim() bool {
if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
return false
}
if !c.lifecycle.CompareAndSwap(uint32(connLifecycleProvisional), uint32(connLifecycleClaiming)) {
return false
}
// Physical close can win after the pre-check but before the lifecycle CAS.
// Do not let a doomed claimant enter SessionManager and retire a healthy old
// owner for the same logical session.
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
c.lifecycle.Store(uint32(connLifecycleRetired))
return false
}
return true
}
func (c *Conn) publishActivation() bool {
if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
return false
}
if !c.lifecycle.CompareAndSwap(uint32(connLifecycleClaiming), uint32(connLifecycleActive)) {
return false
}
// A concurrent transport failure can retire the Conn between the first
// terminal check and the CAS. Never let that intermediate active value escape.
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
c.lifecycle.Store(uint32(connLifecycleRetired))
return false
}
return true
}
func (c *Conn) isActive() bool {
return c != nil && !c.terminal.Load() && c.lifecycleState() == connLifecycleActive
}
// transferTransportOwnership hands this Conn's physical socket to the next
// logical generation. The caller must have fenced and drained the old writer.
func (c *Conn) transferTransportOwnership() (*physicalTransportLease, bool) {
if c == nil || c.transportLease == nil {
return nil, false
}
return c.transportLease.Transfer()
}
func (c *Conn) isPhysicalTransportCurrentOpen() bool {
return c != nil && (c.transportLease == nil || c.transportLease.IsCurrentOpen())
}
// ClientLayer 返回连接协商的 TL layer；未协商时返回 canonical layer（227，不降级）。
func (c *Conn) ClientLayer() int {
if l := c.clientLayer.Load(); l != 0 {

116
internal/mtprotoedge/duplicate_admission_test.go Normal file
View file

@ -0,0 +1,116 @@
package mtprotoedge
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
type blockingDuplicateRPC struct {
started chan struct{}
release chan struct{}
calls atomic.Int32
once sync.Once
}
func newBlockingDuplicateRPC() *blockingDuplicateRPC {
return &blockingDuplicateRPC{
started: make(chan struct{}, 4),
release: make(chan struct{}),
}
}
func (h *blockingDuplicateRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
h.calls.Add(1)
h.started <- struct{}{}
select {
case <-h.release:
return &tg.Config{ThisDC: 2}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
func (*blockingDuplicateRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func (h *blockingDuplicateRPC) unblock() { h.once.Do(func() { close(h.release) }) }
// TestPendingSameConnectionDuplicateDoesNotBlockFreshRequest models the core
// Android startup failure: an RPC is executing, salt correction makes the client
// resend the same msg_id, then initConnection assigns a fresh msg_id. A local
// duplicate must not synchronously join the old owner on the socket read loop,
// otherwise the fresh request remains unread until the old result/replay storm
// has drained.
func TestPendingSameConnectionDuplicateDoesNotBlockFreshRequest(t *testing.T) {
const dc = 2
handler := newBlockingDuplicateRPC()
defer handler.unblock()
addr, pub, server := startTestServer(t, Options{
DC: dc,
RPC: handler,
RPCMaxInflight: 2,
RPCGlobalWorkers: 2,
RPCQueueSize: 8,
})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
oldID := ids.New(proto.MessageFromClient)
newID := ids.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, oldID, 1, &tg.HelpGetConfigRequest{})
waitDuplicateHandlerStarts(t, handler.started, "original request")
// Same ID/seq is a retransmission, not a second business operation.
sendEncryptedWithSeq(t, conn, cipher, auth, oldID, 1, &tg.HelpGetConfigRequest{})
// A client-side init/layer rewrap legitimately assigns a new ID and next seq.
sendEncryptedWithSeq(t, conn, cipher, auth, newID, 3, &tg.HelpGetConfigRequest{})
waitDuplicateHandlerStarts(t, handler.started, "fresh request behind duplicate")
if got := handler.calls.Load(); got != 2 {
t.Fatalf("handler calls before release = %d, want original + fresh only", got)
}
handler.unblock()
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
proto.ResultTypeID: 2,
})
results := make(map[int64]int)
for _, frame := range frames {
if frame.TypeID != proto.ResultTypeID {
continue
}
var result proto.Result
if err := result.Decode(frame.Plain); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
results[result.RequestMessageID]++
}
if results[oldID] != 1 || results[newID] != 1 || len(results) != 2 {
t.Fatalf("rpc_result counts = %+v, want one for old and one for fresh id", results)
}
if got := handler.calls.Load(); got != 2 {
t.Fatalf("final handler calls = %d, want 2", got)
}
deadline := time.Now().Add(2 * time.Second)
for server.rpcResults.flightLimit.snapshot() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := server.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("duplicate admission leaked flight slots: %d", got)
}
}
func waitDuplicateHandlerStarts(t *testing.T, started <-chan struct{}, what string) {
t.Helper()
select {
case <-started:
case <-time.After(2 * time.Second):
t.Fatalf("%s did not reach handler; socket read loop is likely blocked on a duplicate", what)
}
}

708
internal/mtprotoedge/encrypted.go
View file

@ -10,6 +10,7 @@ import (
"fmt"
"io"
"math"
"sync"
"sync/atomic"
"time"
@ -19,6 +20,7 @@ import (
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/tgerr"
"github.com/gotd/td/transport"
@ -30,11 +32,16 @@ import (
// connState 是单连接的 MTProto 运行态。
type connState struct {
sentCreated bool
seen map[int64]clientMsgRecord // 已处理的 client msg_id，用于幂等和 msgs_state_req
order []int64
minSeen int64
maxSeen int64
// createdFloor is the smallest client msg_id covered by the latest
// new_session_created notification for this server-side session generation.
// It only moves down: official clients resend every request below first_msg_id,
// so advertising an outer container id while accepting smaller inner ids would
// orphan the original rpc_result messages.
createdFloor int64
seen map[int64]clientMsgRecord // 已处理的 client msg_id，用于幂等和 msgs_state_req
order []int64
minSeen int64
maxSeen int64
// maxContentMsgID/maxContentSeqNo 是已接受 content 消息的 msg_id / seq_no 高水位，
// 供 validateSeq 的 O(1) 快路径使用（客户端正常发送严格递增）。二者只增不减、
// 不随 seen 淘汰回退——快路径只接受「全扫描也必然接受」的子集，其余回落全扫描。
@ -46,6 +53,9 @@ type clientMsgRecord struct {
state byte
seqNo int32
content bool
// service is the constructor class admitted with this msg_id. A duplicate
// uses this committed class and never decodes/executes its replacement body.
service bool
}
func newConnState() *connState {
@ -76,11 +86,11 @@ const (
// MTProto service vectors operate on bounded connection tracking tables. Accepting more IDs
// only burns decode/CPU and cannot improve the result.
maxServiceMessageIDs = 4096
// A decoded container descriptor is 48 bytes on 64-bit Go today. Charge 64 bytes per entry
// before allocating the exact-size slice so allocator rounding and future field growth remain
// inside the process-wide inbound budget. Message bodies stay as zero-copy views of the already
// charged plaintext frame/gzip expansion.
containerDescriptorBudgetBytes = 64
// Charge each container entry for the decoded proto.Message view plus the staged connState,
// action and ACK descriptors retained by the single-pass inbound plan. Bodies remain zero-copy
// views of the already charged plaintext frame/gzip expansion; RPC copies have a separate batch
// admission budget.
containerDescriptorBudgetBytes = 192
msgStateUnknown byte = 1
msgStateNotReceived byte = 2
@ -97,11 +107,14 @@ const (
badMsgContainer = 64
)
var errActivationAuthKeyRejected = errors.New("activation auth key no longer exists")
// handleEncrypted 解密加密消息，按需注册连接，处理服务消息并分发明文 payload。
// 返回（可能新建/更新的）当前连接对象，供 serveConn 维护生命周期。
// fetchedKey 非 nil 表示本帧的 auth key 是刚从 AuthKeyStore 查出的（首帧/换 auth key/被销毁
// 后回落）；为 nil 表示走快路径——serveConn 判定 current 仍持同一未销毁的 auth key，直接复用
// current.key/current.salt 解密，既不回查 AuthKeyStore 也不重建 store.AuthKeyData。
// 后回落）；为 nil 表示走连接缓存快路径——serveConn 判定 current 仍持同一未销毁的 auth key，
// 直接复用 current.key/current.salt 解密。任何 provisional 在 claim 建立后、发 required
// control 前都会最终回查 AuthKeyStore，使外部撤销与 activation 线性化。
// plain 是 serveConn 持有的复用明文缓冲，frame 的 slice 仅在下一帧解密前有效。
func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b, plain *bin.Buffer) (*Conn, error) {
var key crypto.AuthKey
@ -120,30 +133,27 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
return current, fmt.Errorf("decrypt: %w", err)
}
if frame.salt != serverSalt {
c := current
temp := false
if c == nil || c.sessionID != frame.sessionID || c.authKeyID != key.ID {
c = s.newConn(tc, key, frame.sessionID, serverSalt)
temp = true
}
err := s.sendBadServerSalt(ctx, c, frame.messageID, frame.seqNo, serverSalt)
if temp {
c.Close()
}
return current, err
}
// 首个加密消息或 session 变化时（重新）注册连接到 SessionManager。
// 首个加密消息（即使 salt 尚未修正）或 session 变化时创建并保留唯一的
// provisional Conn。同一物理 transport 换 session 必须先不可逆地 fence/drain
// 旧 writer，再把物理 lease 原子转交给新 generation。为每个 bad_server_salt
// 临时创建 Conn 会在同一 socket 上启动多个 outbound actor，Android 的启动重试
// 风暴随即变成并发写和重复结果放大。
if current == nil || current.sessionID != frame.sessionID || current.authKeyID != key.ID {
if current != nil {
cs.reset()
}
if current != nil {
current.beginTerminalShutdown()
s.conns.Unregister(current)
current.Close()
if !current.waitOutboundShutdownUntil(forceCloseBatchTimeout) {
return current, errors.New("previous session outbound writer did not stop")
}
nextLease, ok := current.transferTransportOwnership()
if !ok {
return current, ErrConnClosed
}
current = s.newConnWithLease(nextLease, key, frame.sessionID, serverSalt)
} else {
current = s.newConn(tc, key, frame.sessionID, serverSalt)
}
current = s.newConn(tc, key, frame.sessionID, serverSalt)
// 注册即播种协商 layer：新 Conn 的 clientLayer 为 0（=canonical 227），若等到
// 首条 RPC 的 Dispatch 返回后才刷新，重连老客户端在首条 RPC handler 执行期间
// 收到的 pending flush / 并发 push 会漏降级。进程内重连时 rpc 层留有
@ -153,68 +163,120 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
current.SetClientLayer(layer)
}
}
s.conns.Register(current)
}
s.maybePersistSession(ctx, current, frame.sessionID, key.ID, serverSalt)
if frame.salt != serverSalt {
// bad_server_salt 是修正后重试的物理屏障：payload 与加密 envelope 都必须携带
// 同一个权威 salt，写失败则该 provisional/active Conn 不得继续接收状态。
return current, s.sendBadServerSalt(ctx, current, frame.messageID, frame.seqNo, serverSalt)
}
body := frame.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(), frame.messageID, frame.seqNo, typeID); code != 0 {
s.log.Debug("Sending bad_msg_notification",
zap.Int64("msg_id", frame.messageID),
zap.Int32("seq_no", frame.seqNo),
zap.Uint32("type_id", typeID),
zap.Int("code", code),
)
return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code)
plan, err := s.preflightInbound(cs, frame.messageID, frame.seqNo, body)
if err != nil {
var bad *dispatchBadMsgError
if errors.As(err, &bad) {
s.log.Debug("Sending bad_msg_notification",
zap.Int64("msg_id", bad.msgID),
zap.Int32("seq_no", bad.seqNo),
zap.Uint32("type_id", typeID),
zap.Int("code", bad.code),
)
return current, s.sendBadMsg(ctx, current, bad.msgID, bad.seqNo, bad.code)
}
return current, err
}
if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext, s.writeTimeout); err != nil {
defer plan.close()
if err := s.prepareInboundRPCBatch(ctx, current, plan); err != nil {
return current, err
}
if err := sendQuickAckIfRequested(ctx, current.transport, key, frame.plaintext, s.writeTimeout); err != nil {
return current, err
}
content := clientMessageNeedsAck(typeID)
if record, ok := cs.seenRecord(frame.messageID); ok {
s.log.Debug("Duplicate msg_id; replay cached result if available", zap.Int64("msg_id", frame.messageID))
if err := s.replayRPCResultByRequest(ctx, current, frame.messageID); err != nil {
moveCreatedFloor := cs.createdFloor == 0 || plan.logicalMin < cs.createdFloor
claimPending := false
if current.lifecycleState() == connLifecycleProvisional {
if !moveCreatedFloor {
return current, errors.New("provisional session has no new_session_created boundary")
}
if err := s.conns.BeginActivation(current); err != nil {
return current, err
}
if !record.content {
return current, nil
claimPending = true
defer func() {
if claimPending {
s.conns.AbortActivation(current)
}
}()
// BeginActivation has installed current in claimsByAuth, which is the shared
// linearization domain with auth-key revocation. A delete that completed before
// the claim is visible here as !found; a delete after this read must observe and
// fence the claim. This final check intentionally covers every activation path:
// first correct-salt frame, retained bad-salt provisional and session transfer.
fresh, found, getErr := s.authKeys.Get(ctx, current.authKeyID)
if getErr != nil {
return current, fmt.Errorf("revalidate activation auth key: %w", getErr)
}
if !found || fresh.ID != current.authKeyID || fresh.Value != [256]byte(current.key.Value) {
// Send the terminal protocol error while the claim still owns a live writer;
// the deferred abort then fences and removes it before serveConn returns.
if sendErr := s.sendProtoError(ctx, current.transport, codec.CodeAuthKeyNotFound); sendErr != nil {
return current, sendErr
}
return current, errActivationAuthKeyRejected
}
if current.terminal.Load() || !current.isPhysicalTransportCurrentOpen() {
return current, ErrConnClosed
}
return current, s.sendAck(ctx, current, frame.messageID)
}
if code := cs.validateSeq(frame.messageID, frame.seqNo, content); code != 0 {
s.log.Debug("Sending bad_msg_notification",
zap.Int64("msg_id", frame.messageID),
if moveCreatedFloor {
s.log.Debug("Sending new_session_created",
zap.Int64("first_msg_id", plan.logicalMin),
zap.Int64("outer_msg_id", frame.messageID),
zap.Int32("seq_no", frame.seqNo),
zap.Uint32("type_id", typeID),
zap.Int("code", code),
)
return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code)
}
cs.track(frame.messageID, frame.seqNo, content, msgStateReceived)
if !cs.sentCreated {
cs.sentCreated = true
s.log.Debug("Sending new_session_created", zap.Int64("msg_id", frame.messageID), zap.Int32("seq_no", frame.seqNo))
if err := s.sendNewSessionCreated(ctx, current, frame.messageID); err != nil {
if err := s.sendNewSessionCreated(ctx, current, plan.logicalMin); err != nil {
return current, err
}
}
if !current.isPhysicalTransportCurrentOpen() {
return current, ErrConnClosed
}
if claimPending {
if err := s.conns.PublishActivation(current); err != nil {
return current, err
}
claimPending = false
}
if moveCreatedFloor {
cs.createdFloor = plan.logicalMin
}
plan.commitState(cs)
s.maybePersistSession(ctx, current, frame.sessionID, key.ID, serverSalt)
var acks []int64
if err := s.dispatch(ctx, cs, current, frame.messageID, frame.seqNo, &bin.Buffer{Buf: body}, &acks); err != nil {
if err := s.executeInboundPlan(ctx, cs, current, plan); err != nil {
return current, err
}
if len(acks) > 0 {
if err := s.sendAck(ctx, current, acks...); err != nil {
if err := plan.commitRPCBatch(); err != nil {
return current, err
}
if len(plan.ackIDs) > 0 {
if err := s.sendAck(ctx, current, plan.ackIDs...); err != nil {
return current, err
}
}
// An overlapping cross-connection owner may run until RPCTimeout. ACK the
// accepted duplicate before joining it so the new client does not build a
// retransmit storm while its read loop waits for the shared result.
if err := s.executePendingRPCReplays(ctx, current, plan); err != nil {
return current, err
}
return current, nil
}
@ -284,236 +346,37 @@ func clientQuickAckToken(key crypto.AuthKey, plaintext []byte) uint32 {
// dispatch 处理一条明文消息：解包 container/gzip，处理服务消息，其余转 RPC 路由。
// content-related 消息（ping、RPC）的 msg_id 会收集到 acks 以便统一确认。
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
expanded := 0
return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, b, acks, dispatchBudget{expanded: &expanded})
}
type dispatchBudget struct {
depth int
containerDepth int
expanded *int
}
func (s *Server) dispatchWithBudget(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64, budget dispatchBudget) error {
if budget.depth > maxDispatchDepth {
return fmt.Errorf("mtproto wrapper depth %d exceeds %d", budget.depth, maxDispatchDepth)
}
id, err := b.PeekID()
plan, err := s.preflightInbound(cs, msgID, seqNo, b.Buf)
if err != nil {
return fmt.Errorf("peek type id: %w", err)
}
ackContent := func() {
if clientMessageNeedsAck(id) {
*acks = append(*acks, msgID)
}
var bad *dispatchBadMsgError
if errors.As(err, &bad) && c != nil {
return s.sendBadMsg(ctx, c, bad.msgID, bad.seqNo, bad.code)
}
return err
}
defer plan.close()
plan.commitState(cs)
*acks = append(*acks, plan.ackIDs...)
return s.executeInboundPlan(ctx, cs, c, plan)
}
switch id {
case proto.GZIPTypeID:
data, releaseExpansion, err := s.decodeGZIPWithGlobalBudget(b)
if err != nil {
return fmt.Errorf("decode gzip: %w", err)
}
defer releaseExpansion()
*budget.expanded += len(data)
if *budget.expanded > maxDispatchExpandedBytes {
return fmt.Errorf("cumulative gzip expansion %d exceeds %d", *budget.expanded, maxDispatchExpandedBytes)
}
budget.depth++
return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: data}, acks, budget)
// dispatchBadMsgError carries a protocol-level rejection discovered during the
// side-effect-free wrapper/container preflight. The caller emits the single
// bad_msg_notification only after the whole container has been inspected.
type dispatchBadMsgError struct {
msgID int64
seqNo int32
code int
}
case proto.MessageContainerTypeID:
if budget.containerDepth != 0 {
return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer)
}
count, err := containerMessageCount(b)
if err != nil {
return fmt.Errorf("decode container count: %w", err)
}
if count > maxContainerMessages {
return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer)
}
container, releaseContainer, err := s.decodeMessageContainerViews(b, count)
if err != nil {
return fmt.Errorf("decode container: %w", err)
}
defer releaseContainer()
if code := validateClientContainer(msgID, seqNo, container); code != 0 {
return s.sendBadMsg(ctx, c, msgID, seqNo, code)
}
budget.depth++
budget.containerDepth++
for i := range container.Messages {
m := container.Messages[i]
typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID()
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 err := s.replayRPCResultByRequest(ctx, c, m.ID); err != nil {
return err
}
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.dispatchWithBudget(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks, budget); 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:
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return fmt.Errorf("msgs_ack vector: %w", err)
}
var ack mt.MsgsAck
if err := ack.Decode(b); err != nil {
return fmt.Errorf("decode msgs_ack: %w", err)
}
c.AckServerMessages(ack.MsgIDs)
s.log.Debug("Received msgs_ack", zap.Int64s("msg_ids", ack.MsgIDs))
return nil
case mt.MsgsStateReqTypeID:
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return fmt.Errorf("msgs_state_req vector: %w", err)
}
var req mt.MsgsStateReq
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode msgs_state_req: %w", err)
}
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:
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
return fmt.Errorf("msg_resend_req vector: %w", err)
}
var req mt.MsgResendReq
if err := req.Decode(b); err != nil {
return fmt.Errorf("decode msg_resend_req: %w", err)
}
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:
reqMsgID, info, err := msgsStateInfoView(b)
if err != nil {
return fmt.Errorf("decode msgs_state_info: %w", err)
}
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", reqMsgID), zap.Int("len", len(info)))
return nil
case mt.MsgsAllInfoTypeID:
count, info, err := msgsAllInfoView(b)
if err != nil {
return fmt.Errorf("decode msgs_all_info: %w", err)
}
if len(info) != count {
return fmt.Errorf("decode msgs_all_info: info length %d does not match msg_ids %d", len(info), count)
}
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", count), zap.Int("len", len(info)))
return nil
case mt.DestroySessionRequestTypeID:
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", c.authKeyHex))
// 真正销毁：删密钥库记录（每帧回查，删除后该 key 的入站帧立即失效）并主动
// 断开同 key 的其他连接——出站推送用连接持有的密钥副本加密、不回查密钥库，
// 不断开的话被销毁 key 的空闲连接仍能持续收到推送。发起连接除外：响应要
// 先送达，它的下一帧会因密钥缺失自然断开。授权（authorizations）不在此清理，
// destroy_auth_key 是 PFS 密钥轮换的清理动作，不等于登出。
if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil {
s.log.Warn("Delete auth key failed", zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{})
}
// 标记密钥已销毁：发起连接被 CloseSessionsForRawAuthKeyExcept 排除（响应需先送达），
// 它下一帧不能再走 serveConn 的密钥复用快路径，须回落到 Get→AuthKeyNotFound 自然失效。
c.keyDestroyed.Store(true)
s.conns.CloseSessionsForRawAuthKeyExcept(c.authKeyID, c.sessionID)
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{})
default:
ackContent()
return s.enqueueRPC(ctx, c, msgID, id, b)
}
func (e *dispatchBadMsgError) Error() string {
return fmt.Sprintf("bad client message %d/%d: code %d", e.msgID, e.seqNo, e.code)
}
// decodeGZIPWithGlobalBudget reserves the maximum single-wrapper output before
// decompression starts. Once the actual size is known the excess reservation is
// returned, while the actual output remains charged through recursive dispatch.
// returned, while the actual output remains charged until the inbound plan is
// executed or aborted.
// This closes the gap where every connection read goroutine could otherwise hold
// an unaccounted 10 MiB expansion before the shared RPC scheduler saw the body.
func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), error) {
@ -571,8 +434,14 @@ func gzipPackedBytesView(b *bin.Buffer) ([]byte, error) {
if binary.LittleEndian.Uint32(b.Buf[:4]) != proto.GZIPTypeID {
return nil, fmt.Errorf("unexpected gzip constructor %#x", binary.LittleEndian.Uint32(b.Buf[:4]))
}
payload, _, err := tlBytesView(b.Buf[4:], -1)
return payload, err
payload, consumed, err := tlBytesView(b.Buf[4:], -1)
if err != nil {
return nil, err
}
if 4+consumed != len(b.Buf) {
return nil, fmt.Errorf("gzip_packed has %d trailing bytes", len(b.Buf)-(4+consumed))
}
return payload, nil
}
// tlBytesView validates one TL bytes envelope and returns a view into the caller-owned buffer.
@ -604,9 +473,9 @@ func tlBytesView(raw []byte, maxPayload int) ([]byte, int, error) {
}
// decodeMessageContainerViews parses the container without proto.Message.Decode's per-body
// copies. Bodies are immutable views of b and stay alive only for this synchronous dispatch;
// enqueueRPC takes its own budgeted copy before returning. Only the exact-size descriptor slice
// is new memory, and that allocation is reserved globally first.
// copies. Bodies stay as immutable views through the single-pass inbound plan; batch admission
// takes independent RPC copies before the backing frame/expansion is released. The descriptor
// reservation also covers staged state, actions and ACK metadata retained by that plan.
func (s *Server) decodeMessageContainerViews(b *bin.Buffer, count int) (proto.MessageContainer, func(), error) {
release := func() {}
if b == nil || len(b.Buf) < 8 {
@ -663,6 +532,10 @@ func (s *Server) decodeMessageContainerViews(b *bin.Buffer, count int) (proto.Me
}
offset = bodyEnd
}
if offset != len(b.Buf) {
release()
return proto.MessageContainer{}, func() {}, fmt.Errorf("message container has %d trailing bytes", len(b.Buf)-offset)
}
return proto.MessageContainer{Messages: messages}, release, nil
}
@ -673,10 +546,13 @@ func msgsStateInfoView(b *bin.Buffer) (int64, []byte, error) {
if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != mt.MsgsStateInfoTypeID {
return 0, nil, fmt.Errorf("unexpected constructor %#x", got)
}
info, _, err := tlBytesView(b.Buf[12:], maxServiceMessageIDs)
info, consumed, err := tlBytesView(b.Buf[12:], maxServiceMessageIDs)
if err != nil {
return 0, nil, err
}
if 12+consumed != len(b.Buf) {
return 0, nil, fmt.Errorf("msgs_state_info has %d trailing bytes", len(b.Buf)-(12+consumed))
}
return int64(binary.LittleEndian.Uint64(b.Buf[4:12])), info, nil
}
@ -691,10 +567,13 @@ func msgsAllInfoView(b *bin.Buffer) (int, []byte, error) {
return 0, nil, io.ErrUnexpectedEOF
}
offset := 12 + count*8
info, _, err := tlBytesView(b.Buf[offset:], maxServiceMessageIDs)
info, consumed, err := tlBytesView(b.Buf[offset:], maxServiceMessageIDs)
if err != nil {
return 0, nil, err
}
if offset+consumed != len(b.Buf) {
return 0, nil, fmt.Errorf("msgs_all_info has %d trailing bytes", len(b.Buf)-(offset+consumed))
}
return count, info, nil
}
@ -750,15 +629,41 @@ func mergeStateInfo(primary, fallback []byte) []byte {
// （已 PeekID 过一次），method 只解析一次并随任务透传，避免同一请求三处重复 PeekID/typeName。
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, request *bin.Buffer) error {
method := s.typeName(typeID)
if cached, ok := s.cachedRPCResult(c, msgID); ok {
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, msgID)
if err != nil {
if errors.Is(err, ErrRPCResultFlightCapacity) {
c.metrics.InboundRPCDropped(method, "flight_capacity")
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, ErrInboundRPCQueueFull)
}
return err
}
switch claim.state {
case rpcResultAcquireCompleted:
s.log.Info("RPC duplicate replay from session cache",
zap.String("method", method),
zap.Int64("msg_id", msgID),
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID),
)
return c.SendEncoded(ctx, proto.MessageServerResponse, cached)
return s.sendCachedRPCResult(ctx, c, claim.encoded)
case rpcResultAcquirePending:
encoded, ok, waitErr := claim.waiter.Wait(ctx)
if waitErr != nil || !ok || encoded == nil {
return waitErr
}
return s.sendCachedRPCResult(ctx, c, encoded)
case rpcResultAcquireOwner:
// Ownership transfers to the queued task only after commit succeeds.
default:
return ErrRPCResultFlightInvalid
}
owner := claim.owner
transferred := false
defer func() {
if !transferred {
owner.Abort()
}
}()
// 两级条数/字节预算必须先于 Copy：对抗客户端不能用大量满尺寸请求在“判断队列满”
// 之前制造一轮无上限的临时 body 分配。reservation 在 commit/abort 间唯一持有预算。
reservation, err := c.reserveInboundRPC(ctx, method, request.Len())
@ -767,11 +672,21 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
}
defer reservation.abort()
body := request.Copy()
responseGate := &rpcResponseGate{}
err = reservation.commit(s.newInboundRPCTask(c, msgID, method, body, owner))
transferred = err == nil
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
}
// newInboundRPCTask builds the exactly-once timeout/result gate shared by the
// single-message and atomic container-batch admission paths. body must already
// be an independently owned, budgeted copy.
func (s *Server) newInboundRPCTask(c *Conn, msgID int64, method string, body []byte, owner *rpcResultOwnerLease) inboundRPC {
responseGate := newRPCResponseGate()
timeoutResponse := func() {
if !responseGate.tryTimeout() {
return
}
defer responseGate.finish()
// 原 task context 已到期，使用有界的新 context 回显明确的可重试超时；
// 500 保持 TDesktop 默认重试语义，错误名区分于容量型 FLOOD_WAIT。
writeTimeout := c.writeTimeout
@ -793,10 +708,29 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
)
}
}
err = reservation.commit(inboundRPC{
return inboundRPC{
method: method,
size: len(body),
onTimeout: timeoutResponse,
release: func() {
if owner == nil {
return
}
// A running deadline callback owns the same terminal flight but executes
// on context.AfterFunc's goroutine. Let its bounded delivery complete before
// deciding that the task produced no result; otherwise release could Abort
// the owner while RPC_TIMEOUT is already inside the physical writer.
if !responseGate.wait(requiredControlMaxWait + time.Second) {
c.fenceUndeliveredRPCResult()
}
if owner.Abort() {
// connState already remembers this request. If a committed task exits
// without publishing any terminal rpc_result, a same-Conn retransmit
// would otherwise be ACKed forever. Force a fresh physical generation
// where the request can be admitted again.
c.fenceUndeliveredRPCResult()
}
},
run: func(taskCtx context.Context) error {
// body 是预算成功后生成的独立副本，且每个任务只 run 一次，
// 无需再 append 拷贝；直接复用，省掉一份 inbound 在途内存。
@ -816,8 +750,7 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
}
return nil
},
})
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
}
}
func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, msgID int64, method string, err error) error {
@ -839,8 +772,18 @@ func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, ms
// handleRPC 把明文 RPC 请求交给 RPC 路由，并将结果或错误包成 rpc_result 回发。
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer, responseGate *rpcResponseGate) error {
if s.rpc == nil {
s.log.Warn("No RPC handler configured; dropping request", zap.String("method", method))
return nil
if err := ctx.Err(); err != nil {
return err
}
if responseGate != nil && !responseGate.tryNormal() {
return context.DeadlineExceeded
}
defer responseGate.finish()
s.log.Warn("No RPC handler configured", zap.String("method", method))
return s.sendResult(ctx, c, msgID, &mt.RPCError{
ErrorCode: 500,
ErrorMessage: "NOT_IMPLEMENTED",
})
}
ctx = postresponse.WithCallbacks(ctx)
@ -874,10 +817,48 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot())
if ctxErr := ctx.Err(); ctxErr != nil {
// A canceled request context means neither a success nor an error can be delivered
// with this expired context. In particular, do not cache a late successful result and
// hand it to outbound: a past write deadline would correctly poison that transport and
// could prevent the scheduler's fresh-context RPC_TIMEOUT response from being sent.
// The old physical Conn may have been replaced after the business transaction
// committed. Never write with the expired context itself: a fenced generation
// publishes cache-only for its replacement, while a still-live generation uses
// a fresh bounded delivery context. A deadline callback that already won
// responseGate has published RPC_TIMEOUT and prevents this late result from
// overwriting it.
// Only a successful business result proves useful work completed. Errors
// observed after cancellation may themselves be cancellation-derived or
// transient and must remain retryable rather than poisoning the replay cache.
var terminal bin.Encoder
if err == nil && result != nil {
terminal = result
}
if terminal != nil && (responseGate == nil || responseGate.tryNormal()) {
defer responseGate.finish()
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() {
// Replacement/shutdown already fenced this logical generation. Cache-only
// publication is safe and lets the replacement join the completed flight.
if encoded, encodeErr := s.encodeRPCResult(c, msgID, terminal); encodeErr != nil {
s.log.Warn("Encode canceled RPC result for replay failed", append(fields, zap.Error(encodeErr))...)
} else {
s.storeRPCResult(c, msgID, encoded)
postresponse.Run(context.WithoutCancel(ctx))
}
} else {
// An individual RPC deadline can expire while the physical connection is
// still healthy. Use a fresh bounded delivery context; publishing cache-only
// here would strand same-Conn duplicates behind an ACK with no result.
writeTimeout := c.writeTimeout
if writeTimeout <= 0 || writeTimeout > 5*time.Second {
writeTimeout = 5 * time.Second
}
responseCtx, cancel := context.WithTimeout(context.Background(), writeTimeout)
sendErr := s.sendResult(responseCtx, c, msgID, terminal)
cancel()
if sendErr != nil {
s.log.Debug("Send canceled RPC result failed", append(fields, zap.Error(sendErr))...)
} else {
postresponse.Run(context.WithoutCancel(ctx))
}
}
}
cancelFields := append(fields, zap.NamedError("context_error", ctxErr))
if err != nil {
cancelFields = append(cancelFields, zap.NamedError("dispatch_error", err))
@ -891,6 +872,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
s.log.Info("RPC result suppressed after timeout", fields...)
return context.DeadlineExceeded
}
defer responseGate.finish()
if err != nil {
var rpcErr *tgerr.Error
@ -921,6 +903,12 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
// first owns the response, and the other path becomes a no-op.
type rpcResponseGate struct {
state atomic.Uint32
done chan struct{}
once sync.Once
}
func newRPCResponseGate() *rpcResponseGate {
return &rpcResponseGate{done: make(chan struct{})}
}
func (g *rpcResponseGate) tryNormal() bool {
@ -931,14 +919,84 @@ func (g *rpcResponseGate) tryTimeout() bool {
return g != nil && g.state.CompareAndSwap(0, 2)
}
func (g *rpcResponseGate) finish() {
if g == nil || g.done == nil {
return
}
g.once.Do(func() { close(g.done) })
}
// wait returns once there is no response owner, the winning response has
// completed deliver-or-fence publication, or the bounded safety deadline wins.
func (g *rpcResponseGate) wait(timeout time.Duration) bool {
if g == nil || g.state.Load() == 0 {
return true
}
if g.done == nil {
return false
}
if timeout <= 0 {
<-g.done
return true
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-g.done:
return true
case <-timer.C:
return false
}
}
// sendResult 把 RPC 结果包成 rpc_result 并加密回发。
func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) error {
if result == nil {
result = &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"}
}
encoded, err := s.encodeRPCResult(c, reqMsgID, result)
if err != nil {
// The business operation has already crossed atomic admission. Convert an
// invalid result encoder into one deterministic terminal RPC error instead of
// aborting the flight and allowing a reconnect to execute the operation again.
s.log.Warn("Encode RPC result failed; sending INTERNAL", zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
encoded, err = s.encodeRPCResult(c, reqMsgID, &mt.RPCError{
ErrorCode: 500,
ErrorMessage: "INTERNAL",
})
if err != nil {
c.fenceUndeliveredRPCResult()
return err
}
}
if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil {
// A completed result may be published before delivery only after this logical
// Conn is irreversibly fenced. SendEncoded has non-writing failure paths
// (queue/context/scratch deadline); without this terminal barrier a later
// same-Conn duplicate would be ACKed while no result can ever arrive.
c.fenceUndeliveredRPCResult()
s.storeRPCResult(c, reqMsgID, encoded)
return err
}
// On a live Conn, completed means the rpc_result has reached the reliable byte
// stream. Same-physical duplicates can therefore be ACK-only without data loss.
s.storeRPCResult(c, reqMsgID, encoded)
return c.SendEncoded(ctx, proto.MessageServerResponse, encoded)
return nil
}
// sendCachedRPCResult preserves the delivery half of the rpc_result invariant
// for completed-flight replays: either the cached result reaches this physical
// byte stream, or this logical Conn is fenced so a replacement may retry it.
func (s *Server) sendCachedRPCResult(ctx context.Context, c *Conn, encoded *encodedOutboundMessage) error {
if encoded == nil {
c.fenceUndeliveredRPCResult()
return errors.New("nil cached rpc_result")
}
if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil {
c.fenceUndeliveredRPCResult()
return err
}
return nil
}
// encodeRPCResult 编码 rpc_result。内层对象与 rpc_result 头（type_id + req_msg_id）
@ -986,13 +1044,14 @@ func (s *Server) replayRPCResultByRequest(ctx context.Context, c *Conn, reqMsgID
return nil
}
if resent, err := c.ResendByRequest(ctx, reqMsgID); err != nil {
c.fenceUndeliveredRPCResult()
return err
} else if resent {
s.log.Debug("Resent connection cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
return nil
}
if cached, ok := s.cachedRPCResult(c, reqMsgID); ok {
if err := c.SendEncoded(ctx, proto.MessageServerResponse, cached); err != nil {
if err := s.sendCachedRPCResult(ctx, c, cached); err != nil {
return err
}
s.log.Debug("Resent session cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
@ -1044,7 +1103,10 @@ func (s *Server) sendFutureSalts(ctx context.Context, c *Conn, reqMsgID int64, n
// 复用同一值会让断线重连后的 new_session_created 被吞掉，错过的差分补拉
// （Android 收到后才调 getDifference）随之丢失。
func (s *Server) sendNewSessionCreated(ctx context.Context, c *Conn, firstMsgID int64) error {
return c.SendAsync(ctx, proto.MessageFromServer, &mt.NewSessionCreated{
// This notification changes the client's request map and update recovery
// state. Unlike best-effort ack/pong traffic, it must be written successfully
// before the corresponding RPC batch starts executing.
return c.SendRequiredControl(ctx, proto.MessageFromServer, &mt.NewSessionCreated{
FirstMsgID: firstMsgID,
UniqueID: s.newServerSessionUID(),
ServerSalt: c.salt,
@ -1098,7 +1160,7 @@ func (s *Server) sendBadMsg(ctx context.Context, c *Conn, badMsgID int64, badSeq
// 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{
return c.SendRequiredControl(ctx, proto.MessageFromServer, &mt.BadServerSalt{
BadMsgID: badMsgID,
BadMsgSeqno: int(badSeqno),
ErrorCode: 48,
@ -1138,25 +1200,6 @@ func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint
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
@ -1235,10 +1278,15 @@ func (cs *connState) validateSeq(msgID int64, seqNo int32, content bool) int {
}
func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) {
cs.trackInbound(msgID, seqNo, content, false, state)
}
func (cs *connState) trackInbound(msgID int64, seqNo int32, content, service bool, state byte) {
cs.seen[msgID] = clientMsgRecord{
state: state,
seqNo: seqNo,
content: content,
service: service,
}
if content {
if msgID > cs.maxContentMsgID {

103
internal/mtprotoedge/encrypted_test.go
View file

@ -40,25 +40,72 @@ func TestEncryptedPingPong(t *testing.T) {
}
}
// TestDuplicateMsgIDIdempotent 验证 M4：相同 msg_id 的重复 content 请求被幂等处理，
// server 重发已缓存的 rpc_result，并重新 ack，不重复执行业务。
// TestDuplicateMsgIDIdempotent 验证相同物理连接上的重复 content 请求只重新 ACK，
// 原 owner 仍是唯一 rpc_result 发送者。若每次重复都重放完整结果，Android 的
// bad_server_salt 全量重试会把一个启动批次放大成 N 轮孤儿结果并饿死新 request id。
func TestDuplicateMsgIDIdempotent(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
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")
sendEncrypted(t, conn, cipher, auth, msgID, &tg.HelpGetConfigRequest{})
collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
proto.ResultTypeID: 1,
mt.MsgsAckTypeID: 1,
})
waitForAtomicCalls(t, &handler.calls, 1)
// 相同 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")
// 用一个小型重试风暴覆盖完成后的 duplicate 路径。TCP 仍存活时原结果已在同一
// 可靠字节流上；每个 duplicate 只需 ACK，不应产生第二个 rpc_result。
const duplicateCount = 16
for i := 0; i < duplicateCount; i++ {
sendEncrypted(t, conn, cipher, auth, msgID, &tg.HelpGetConfigRequest{})
}
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.MsgsAckTypeID: duplicateCount,
})
for _, frame := range frames {
if frame.TypeID == proto.ResultTypeID {
t.Fatalf("same-connection duplicate emitted an extra rpc_result")
}
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("same-connection duplicate business calls = %d, want 1", got)
}
}
func TestServiceDuplicateCannotReplaceOriginallyAdmittedPayload(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
msgID := ids.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &mt.PingRequest{PingID: 11})
collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.PongTypeID: 1,
mt.MsgsAckTypeID: 1,
})
// Same id/seq/content parity but a destructive replacement body. Duplicate
// handling must use the original committed request class and never execute it.
sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &destroyAuthKeyRequest{})
collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.MsgsAckTypeID: 1})
freshID := ids.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, freshID, 3, &mt.PingRequest{PingID: 22})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
var pong mt.Pong
if err := pong.Decode(mustHave(t, replies, mt.PongTypeID, "pong after replacement attempt")); err != nil {
t.Fatalf("decode pong: %v", err)
}
if pong.MsgID != freshID || pong.PingID != 22 {
t.Fatalf("pong after replacement = %+v, want msg=%d ping=22", pong, freshID)
}
}
// TestGetFutureSalts 验证 MTProto service message get_future_salts 由连接层直接响应，
@ -268,6 +315,14 @@ func TestOldMessageInFreshContainerAccepted(t *testing.T) {
})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
createdBuf := mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created")
var created mt.NewSessionCreated
if err := created.Decode(createdBuf); err != nil {
t.Fatalf("decode new_session_created: %v", err)
}
if created.FirstMsgID != oldPingMsgID {
t.Fatalf("new_session_created.first_msg_id = %d, want accepted inner msg_id %d", created.FirstMsgID, oldPingMsgID)
}
buf := mustHave(t, replies, mt.PongTypeID, "pong")
var pong mt.Pong
if err := pong.Decode(buf); err != nil {
@ -382,8 +437,10 @@ func TestBadServerSalt(t *testing.T) {
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")
envelope, typeID, buf := readServerMessage(t, conn, cipher, auth.AuthKey)
if typeID != mt.BadServerSaltTypeID {
t.Fatalf("bad salt reply type = %#x, want %#x", typeID, mt.BadServerSaltTypeID)
}
var bad mt.BadServerSalt
if err := bad.Decode(buf); err != nil {
@ -398,6 +455,11 @@ func TestBadServerSalt(t *testing.T) {
if bad.NewServerSalt != auth.ServerSalt {
t.Fatalf("bad_server_salt.new_server_salt = %#x, want %#x", bad.NewServerSalt, auth.ServerSalt)
}
// DrKLO stores the salt from the encrypted envelope, not only the TL payload.
// A mismatch makes every correction ineffective and re-enters the resend storm.
if envelope.Salt != bad.NewServerSalt {
t.Fatalf("bad_server_salt envelope salt = %#x, payload = %#x", envelope.Salt, bad.NewServerSalt)
}
}
func TestBadMsgSeqOddExpected(t *testing.T) {
@ -470,7 +532,7 @@ func TestBadMsgSeqTooHigh(t *testing.T) {
func TestSessionChangeResetsClientSeqState(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
addr, pub, srv := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
@ -492,6 +554,19 @@ func TestSessionChangeResetsClientSeqState(t *testing.T) {
}
mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created after session change")
mustHave(t, replies, mt.MsgsAckTypeID, "msgs_ack after session change")
oldKey := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID}
newKey := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: nextSessionID}
srv.conns.mu.RLock()
_, oldVisible := srv.conns.bySession[oldKey]
newConn := srv.conns.bySession[newKey]
claims := len(srv.conns.claims)
online := len(srv.conns.bySession)
srv.conns.mu.RUnlock()
if oldVisible || newConn == nil || !newConn.isActive() || claims != 0 || online != 1 {
t.Fatalf("same-transport switch state: old=%v new=%p active=%v claims=%d online=%d",
oldVisible, newConn, newConn != nil && newConn.isActive(), claims, online)
}
}
func readBadMsgNotification(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) mt.BadMsgNotification {

74
internal/mtprotoedge/helpers_test.go
View file

@ -190,12 +190,82 @@ func collectReplies(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key
return got
}
// serverReplyFrame preserves the wire order and encrypted envelope of a server
// reply. Tests which exercise session boundaries must not collapse replies into
// a TypeID-keyed map: both duplicate response types and their order are part of
// the observable protocol behavior.
type serverReplyFrame struct {
Message *crypto.EncryptedMessageData
TypeID uint32
Plain *bin.Buffer
}
// collectReplyFrames reads ordered server replies until every requested TypeID
// has been observed the requested number of times. Unrequested frames are kept
// in the returned slice so callers can assert ordering around control messages.
func collectReplyFrames(
t *testing.T,
conn transport.Conn,
cipher crypto.Cipher,
key crypto.AuthKey,
wantCounts map[uint32]int,
) []serverReplyFrame {
t.Helper()
remaining := make(map[uint32]int, len(wantCounts))
required := 0
for typeID, count := range wantCounts {
if count <= 0 {
continue
}
remaining[typeID] = count
required += count
}
if required == 0 {
return nil
}
// Keep the helper bounded while allowing unrelated control replies (notably
// msgs_ack) to be interleaved with the frames under test.
// One shared deadline bounds the whole collection. A per-frame deadline would
// multiply a missing-result failure by the maximum number of unrelated frames.
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
frames := make([]serverReplyFrame, 0, required)
for i := 0; i < required+16; i++ {
message, typeID, plain := readServerMessageContext(t, ctx, conn, cipher, key)
frames = append(frames, serverReplyFrame{
Message: message,
TypeID: typeID,
Plain: plain,
})
if count, ok := remaining[typeID]; ok {
if count == 1 {
delete(remaining, typeID)
} else {
remaining[typeID] = count - 1
}
}
if len(remaining) == 0 {
return frames
}
}
t.Fatalf("missing reply counts after %d frames: %+v", len(frames), remaining)
return nil
}
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)
defer cancel()
return readServerMessageContext(t, ctx, conn, cipher, key)
}
func readServerMessageContext(t *testing.T, ctx context.Context, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) (*crypto.EncryptedMessageData, uint32, *bin.Buffer) {
t.Helper()
var buf bin.Buffer
err := conn.Recv(ctx, &buf)
cancel()
if err != nil {
t.Fatalf("recv server message: %v", err)
}

880
internal/mtprotoedge/inbound_preflight.go Normal file
View file

@ -0,0 +1,880 @@
package mtprotoedge
import (
"context"
"encoding/binary"
"errors"
"fmt"
"io"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
)
type inboundItemKind uint8
const (
inboundItemDuplicate inboundItemKind = iota + 1
inboundItemServiceDuplicate
inboundItemPing
inboundItemPingDelay
inboundItemFutureSalts
inboundItemMsgsAck
inboundItemStateReq
inboundItemResendReq
inboundItemStateInfo
inboundItemAllInfo
inboundItemDestroySession
inboundItemHTTPWait
inboundItemDropAnswer
inboundItemDestroyAuthKey
inboundItemRPC
inboundItemCapacityError
inboundItemPendingRPC
// inboundItemReplayRPC is a request first observed by this physical Conn whose
// terminal result already exists in the cross-connection cache. It is distinct
// from inboundItemDuplicate: a duplicate already present in this Conn's seen
// table must only be ACKed. The original owner/result is already using the same
// reliable TCP stream, so replaying it once per retransmit wave amplifies a
// client retry burst and can starve newer request IDs behind orphan results.
inboundItemReplayRPC
)
type inboundItem struct {
kind inboundItemKind
msgID int64
seqNo int32
typeID uint32
content bool
body []byte
payload any
}
type stagedClientMessage struct {
msgID int64
seqNo int32
content bool
service bool
}
type inboundPlan struct {
items []inboundItem
staged []stagedClientMessage
ackIDs []int64
logicalMin int64
releases []func()
rpcPrepared bool
rpcReservation *inboundRPCBatchReservation
rpcTasks []inboundRPC
rpcOwners []*rpcResultOwnerLease
}
func (p *inboundPlan) close() {
if p == nil {
return
}
if p.rpcReservation != nil {
p.rpcReservation.abort()
p.rpcReservation = nil
}
for _, owner := range p.rpcOwners {
owner.Abort()
}
p.rpcOwners = nil
for i := len(p.releases) - 1; i >= 0; i-- {
p.releases[i]()
}
p.releases = nil
}
func (p *inboundPlan) commitRPCBatch() error {
if p == nil || p.rpcReservation == nil {
return nil
}
// handleEncrypted calls this only after ownership/session-control barriers,
// connState commit and every synchronous service action have completed. The
// batch is runnable immediately; using the old deferred scheduler token here
// was not a real barrier on a busy Conn because an existing ready token could
// dequeue newly appended tasks before activateRPCBatch ran.
_, err := p.rpcReservation.commit(p.rpcTasks, false)
if err != nil {
return err
}
p.rpcReservation = nil
p.rpcTasks = nil
// Every owner is now attached to exactly one queued task. Its task.release
// path aborts if no terminal rpc_result was published.
p.rpcOwners = nil
return nil
}
func (p *inboundPlan) includeLogicalID(msgID int64) {
if p.logicalMin == 0 || msgID < p.logicalMin {
p.logicalMin = msgID
}
}
func (p *inboundPlan) commitState(cs *connState) {
for _, m := range p.staged {
cs.trackInbound(m.msgID, m.seqNo, m.content, m.service, msgStateReceived)
}
}
type connStateOverlay struct {
base *connState
staged []stagedClientMessage
maxContentMsgID int64
maxContentSeqNo int32
}
func newConnStateOverlay(base *connState) connStateOverlay {
return connStateOverlay{
base: base,
maxContentMsgID: base.maxContentMsgID,
maxContentSeqNo: base.maxContentSeqNo,
}
}
func (o *connStateOverlay) seenRecord(msgID int64) (clientMsgRecord, bool) {
for i := len(o.staged) - 1; i >= 0; i-- {
m := o.staged[i]
if m.msgID == msgID {
return clientMsgRecord{state: msgStateReceived, seqNo: m.seqNo, content: m.content, service: m.service}, true
}
}
return o.base.seenRecord(msgID)
}
func (o *connStateOverlay) validateSeq(msgID int64, seqNo int32, content bool) int {
if !content {
return 0
}
if msgID > o.maxContentMsgID && seqNo > o.maxContentSeqNo {
return 0
}
for seenMsgID, record := range o.base.seen {
if !record.content {
continue
}
if seenMsgID < msgID && record.seqNo >= seqNo {
return badMsgSeqTooLow
}
if seenMsgID > msgID && record.seqNo <= seqNo {
return badMsgSeqTooHigh
}
}
for _, record := range o.staged {
if !record.content {
continue
}
if record.msgID < msgID && record.seqNo >= seqNo {
return badMsgSeqTooLow
}
if record.msgID > msgID && record.seqNo <= seqNo {
return badMsgSeqTooHigh
}
}
return 0
}
func (o *connStateOverlay) stage(msgID int64, seqNo int32, content, service bool) {
o.staged = append(o.staged, stagedClientMessage{msgID: msgID, seqNo: seqNo, content: content, service: service})
if content {
if msgID > o.maxContentMsgID {
o.maxContentMsgID = msgID
}
if seqNo > o.maxContentSeqNo {
o.maxContentSeqNo = seqNo
}
}
}
type inboundScope struct {
insideContainer bool
mustBeDuplicate bool
}
type inboundPreflightBudget struct {
depth int
containerDepth int
expanded int
now time.Time
}
func (s *Server) preflightInbound(cs *connState, msgID int64, seqNo int32, body []byte) (*inboundPlan, error) {
plan := &inboundPlan{logicalMin: 0}
overlay := newConnStateOverlay(cs)
// All envelope checks in one transport frame use the same clock sample. Apart
// from making boundary behavior deterministic, this lets walkInbound reject
// an invalid outer msg_id before spending CPU or memory on gzip expansion.
budget := &inboundPreflightBudget{now: s.clock.Now()}
if err := s.walkInbound(plan, &overlay, msgID, seqNo, body, inboundScope{}, budget); err != nil {
plan.close()
return nil, err
}
plan.staged = overlay.staged
if plan.logicalMin == 0 {
plan.close()
return nil, fmt.Errorf("inbound plan accepted no logical message")
}
for _, item := range plan.items {
if item.kind == inboundItemDestroyAuthKey && len(plan.items) != 1 {
plan.close()
return nil, &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
}
return plan, nil
}
func (s *Server) walkInbound(
plan *inboundPlan,
overlay *connStateOverlay,
msgID int64,
seqNo int32,
body []byte,
scope inboundScope,
budget *inboundPreflightBudget,
) error {
if budget.depth > maxDispatchDepth {
return fmt.Errorf("mtproto wrapper depth %d exceeds %d", budget.depth, maxDispatchDepth)
}
b := &bin.Buffer{Buf: body}
typeID, err := b.PeekID()
if err != nil {
return fmt.Errorf("peek type id: %w", err)
}
if code := validateInboundMessageID(budget.now, msgID, scope.insideContainer); code != 0 {
if scope.insideContainer {
code = badMsgContainer
}
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code}
}
// A repeated outer container still has to be decoded far enough to enumerate
// its inner request ids, but none of those already-accepted inner bodies needs
// decoding again. In particular, never reinflate a gzip body merely to discover
// the content bit we already retained with its msg_id. If an outer duplicate
// introduces an unseen inner id, reject the container before touching its body.
// A seen top-level content gzip is also safe to short-circuit: a container is
// non-content, so it cannot be hidden behind that retained record. Ambiguous
// non-content top-level gzip still expands to distinguish a container wrapper.
if record, seen := overlay.seenRecord(msgID); scope.mustBeDuplicate ||
(typeID == proto.GZIPTypeID && seen && (scope.insideContainer || record.content)) {
if !seen || record.seqNo != seqNo {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
return appendInboundDuplicate(plan, msgID, seqNo, typeID, record)
}
if typeID == proto.GZIPTypeID {
data, release, err := s.decodeGZIPWithGlobalBudget(b)
if err != nil {
return fmt.Errorf("decode gzip: %w", err)
}
plan.releases = append(plan.releases, release)
budget.expanded += len(data)
if budget.expanded > maxDispatchExpandedBytes {
return fmt.Errorf("cumulative gzip expansion %d exceeds %d", budget.expanded, maxDispatchExpandedBytes)
}
budget.depth++
err = s.walkInbound(plan, overlay, msgID, seqNo, data, scope, budget)
budget.depth--
return err
}
if typeID == proto.MessageContainerTypeID {
if scope.insideContainer || budget.containerDepth != 0 {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
if code := validateClientEnvelope(budget.now, msgID, seqNo, typeID); code != 0 {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code}
}
outerRecord, outerSeen := overlay.seenRecord(msgID)
if outerSeen {
if outerRecord.content || outerRecord.seqNo != seqNo {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
} else {
if code := overlay.validateSeq(msgID, seqNo, false); code != 0 {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code}
}
overlay.stage(msgID, seqNo, false, false)
}
count, err := containerMessageCount(b)
if err != nil {
return fmt.Errorf("decode container count: %w", err)
}
if count > maxContainerMessages {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
container, release, err := s.decodeMessageContainerViews(b, count)
if err != nil {
return fmt.Errorf("decode container: %w", err)
}
plan.releases = append(plan.releases, release)
plan.items = growInboundItems(plan.items, count)
plan.ackIDs = growInt64s(plan.ackIDs, count)
overlay.staged = growStagedMessages(overlay.staged, count+1)
if len(container.Messages) == 0 {
plan.includeLogicalID(msgID)
return nil
}
budget.depth++
budget.containerDepth++
for _, m := range container.Messages {
if m.ID >= msgID || int32(m.SeqNo) > seqNo {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
if err := s.walkInbound(plan, overlay, m.ID, int32(m.SeqNo), m.Body, inboundScope{
insideContainer: true,
mustBeDuplicate: outerSeen,
}, budget); err != nil {
return err
}
}
budget.containerDepth--
budget.depth--
return nil
}
if scope.insideContainer {
if code := validateClientContainerEnvelope(msgID, seqNo, typeID); code != 0 {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
} else if code := validateClientEnvelope(budget.now, msgID, seqNo, typeID); code != 0 {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code}
}
content := clientMessageNeedsAck(typeID)
if record, seen := overlay.seenRecord(msgID); seen {
if record.seqNo != seqNo || record.content != content {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
return appendInboundDuplicate(plan, msgID, seqNo, typeID, record)
}
if scope.mustBeDuplicate {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
}
if code := overlay.validateSeq(msgID, seqNo, content); code != 0 {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code}
}
overlay.stage(msgID, seqNo, content, inboundTypeIsService(typeID))
plan.includeLogicalID(msgID)
item, err := preflightInboundItem(msgID, seqNo, typeID, content, body)
if err != nil {
return err
}
plan.items = append(plan.items, item)
if content {
plan.ackIDs = append(plan.ackIDs, msgID)
}
return nil
}
func validateInboundMessageID(now time.Time, msgID int64, insideContainer bool) int {
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
return badMsgIDInvalidBits
}
// A container's outer envelope supplies the wall-clock admission boundary for
// its inner messages. Inner ids still need the client low bits checked before
// wrapper expansion, but intentionally keep the established no-time-check rule.
if insideContainer {
return 0
}
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
}
return 0
}
func appendInboundDuplicate(plan *inboundPlan, msgID int64, seqNo int32, typeID uint32, record clientMsgRecord) error {
plan.includeLogicalID(msgID)
kind := inboundItemDuplicate
if record.service {
kind = inboundItemServiceDuplicate
}
plan.items = append(plan.items, inboundItem{
kind: kind, msgID: msgID, seqNo: seqNo, typeID: typeID, content: record.content,
})
if record.content {
plan.ackIDs = append(plan.ackIDs, msgID)
}
return nil
}
func inboundTypeIsService(typeID uint32) bool {
switch typeID {
case mt.PingRequestTypeID,
mt.PingDelayDisconnectRequestTypeID,
mt.GetFutureSaltsRequestTypeID,
mt.MsgsAckTypeID,
mt.MsgsStateReqTypeID,
mt.MsgResendReqTypeID,
mt.MsgsStateInfoTypeID,
mt.MsgsAllInfoTypeID,
mt.DestroySessionRequestTypeID,
mt.HTTPWaitRequestTypeID,
mt.RPCDropAnswerRequestTypeID,
destroyAuthKeyRequestTypeID:
return true
default:
return false
}
}
func growInboundItems(items []inboundItem, extra int) []inboundItem {
if extra <= cap(items)-len(items) {
return items
}
grown := make([]inboundItem, len(items), len(items)+extra)
copy(grown, items)
return grown
}
func growStagedMessages(items []stagedClientMessage, extra int) []stagedClientMessage {
if extra <= cap(items)-len(items) {
return items
}
grown := make([]stagedClientMessage, len(items), len(items)+extra)
copy(grown, items)
return grown
}
func growInt64s(items []int64, extra int) []int64 {
if extra <= cap(items)-len(items) {
return items
}
grown := make([]int64, len(items), len(items)+extra)
copy(grown, items)
return grown
}
type stateInfoPayload struct {
reqMsgID int64
info []byte
}
type allInfoPayload struct {
count int
info []byte
}
// int64VectorView is a bounded, immutable view of a TL Vector<long>. Keeping the
// encoded bytes avoids one retained []int64 allocation per service message while
// a whole container is in preflight. Execution materializes the small bounded
// slice only for the existing Conn APIs that require []int64, after all wrappers
// and every sibling message have already passed structural validation.
type int64VectorView struct {
raw []byte
count int
}
func decodeInt64VectorView(b *bin.Buffer, expectedTypeID uint32, max int) (int64VectorView, error) {
if b == nil || len(b.Buf) < 12 {
return int64VectorView{}, io.ErrUnexpectedEOF
}
if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != expectedTypeID {
return int64VectorView{}, fmt.Errorf("unexpected constructor %#x", got)
}
if got := binary.LittleEndian.Uint32(b.Buf[4:8]); got != bin.TypeVector {
return int64VectorView{}, fmt.Errorf("unexpected vector constructor %#x", got)
}
count := int(int32(binary.LittleEndian.Uint32(b.Buf[8:12])))
if count < 0 {
return int64VectorView{}, fmt.Errorf("negative vector count %d", count)
}
if count > max {
return int64VectorView{}, fmt.Errorf("vector count %d exceeds %d", count, max)
}
if count > (len(b.Buf)-12)/8 {
return int64VectorView{}, io.ErrUnexpectedEOF
}
end := 12 + count*8
if end != len(b.Buf) {
return int64VectorView{}, fmt.Errorf("long vector has %d trailing bytes", len(b.Buf)-end)
}
return int64VectorView{raw: b.Buf[12:end:end], count: count}, nil
}
func (v int64VectorView) materialize() []int64 {
if v.count == 0 {
return nil
}
ids := make([]int64, v.count)
for i := range ids {
offset := i * 8
ids[i] = int64(binary.LittleEndian.Uint64(v.raw[offset : offset+8]))
}
return ids
}
func decodeInboundServiceExact(b *bin.Buffer, decoder bin.Decoder) error {
if err := decoder.Decode(b); err != nil {
return err
}
if b.Len() != 0 {
return fmt.Errorf("service message has %d trailing bytes", b.Len())
}
return nil
}
func preflightInboundItem(msgID int64, seqNo int32, typeID uint32, content bool, body []byte) (inboundItem, error) {
item := inboundItem{msgID: msgID, seqNo: seqNo, typeID: typeID, content: content, body: body}
b := &bin.Buffer{Buf: body}
switch typeID {
case mt.PingRequestTypeID:
var value mt.PingRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, fmt.Errorf("decode ping: %w", err)
}
item.kind, item.payload = inboundItemPing, value
case mt.PingDelayDisconnectRequestTypeID:
var value mt.PingDelayDisconnectRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, fmt.Errorf("decode ping_delay_disconnect: %w", err)
}
item.kind, item.payload = inboundItemPingDelay, value
case mt.GetFutureSaltsRequestTypeID:
var value mt.GetFutureSaltsRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, fmt.Errorf("decode get_future_salts: %w", err)
}
item.kind, item.payload = inboundItemFutureSalts, value
case mt.MsgsAckTypeID:
value, err := decodeInt64VectorView(b, mt.MsgsAckTypeID, maxServiceMessageIDs)
if err != nil {
return item, fmt.Errorf("decode msgs_ack: %w", err)
}
item.kind, item.payload = inboundItemMsgsAck, value
case mt.MsgsStateReqTypeID:
value, err := decodeInt64VectorView(b, mt.MsgsStateReqTypeID, maxServiceMessageIDs)
if err != nil {
return item, fmt.Errorf("decode msgs_state_req: %w", err)
}
item.kind, item.payload = inboundItemStateReq, value
case mt.MsgResendReqTypeID:
value, err := decodeInt64VectorView(b, mt.MsgResendReqTypeID, maxServiceMessageIDs)
if err != nil {
return item, fmt.Errorf("decode msg_resend_req: %w", err)
}
item.kind, item.payload = inboundItemResendReq, value
case mt.MsgsStateInfoTypeID:
reqMsgID, info, err := msgsStateInfoView(b)
if err != nil {
return item, fmt.Errorf("decode msgs_state_info: %w", err)
}
item.kind, item.payload = inboundItemStateInfo, stateInfoPayload{reqMsgID: reqMsgID, info: info}
case mt.MsgsAllInfoTypeID:
count, info, err := msgsAllInfoView(b)
if err != nil {
return item, fmt.Errorf("decode msgs_all_info: %w", err)
}
if len(info) != count {
return item, fmt.Errorf("decode msgs_all_info: info length %d does not match msg_ids %d", len(info), count)
}
item.kind, item.payload = inboundItemAllInfo, allInfoPayload{count: count, info: info}
case mt.DestroySessionRequestTypeID:
var value mt.DestroySessionRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, fmt.Errorf("decode destroy_session: %w", err)
}
item.kind, item.payload = inboundItemDestroySession, value
case mt.HTTPWaitRequestTypeID:
var value mt.HTTPWaitRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, fmt.Errorf("decode http_wait: %w", err)
}
item.kind, item.payload = inboundItemHTTPWait, value
case mt.RPCDropAnswerRequestTypeID:
var value mt.RPCDropAnswerRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, fmt.Errorf("decode rpc_drop_answer: %w", err)
}
item.kind, item.payload = inboundItemDropAnswer, value
case destroyAuthKeyRequestTypeID:
var value destroyAuthKeyRequest
if err := decodeInboundServiceExact(b, &value); err != nil {
return item, err
}
item.kind, item.payload = inboundItemDestroyAuthKey, value
default:
item.kind = inboundItemRPC
}
return item, nil
}
// prepareInboundRPCBatch performs the whole container's count/byte admission
// before copying or scheduling any API RPC. Capacity exhaustion is converted
// into one consistent terminal FLOOD_WAIT result per uncached RPC; no business
// handler from the batch is allowed to start in that case.
func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inboundPlan) error {
plan.rpcPrepared = true
// Keep service-only frames (ping/ack/http_wait) allocation-free here. These
// collections are needed only after the first real API RPC acquires ownership.
var indices []int
var specs []inboundRPCSpec
var ownersInPlan map[int64]*rpcResultOwnerLease
flightCapacity := false
for i := range plan.items {
item := &plan.items[i]
if item.kind != inboundItemRPC && item.kind != inboundItemDuplicate {
continue
}
localDuplicate := item.kind == inboundItemDuplicate
if localDuplicate {
// connState already proves this msg_id was admitted on this physical
// generation. Its original owner either still holds the flight or has
// published a result after physical write; do not consume a global flight
// slot merely to ACK the duplicate.
continue
}
method := s.typeName(item.typeID)
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, item.msgID)
if errors.Is(err, ErrRPCResultFlightCapacity) {
if item.kind == inboundItemRPC {
c.metrics.InboundRPCDropped(method, "flight_capacity")
flightCapacity = true
item.kind = inboundItemCapacityError
}
continue
}
if err != nil {
return err
}
switch claim.state {
case rpcResultAcquireCompleted:
s.log.Info("RPC duplicate replay from session cache",
zap.String("method", method),
zap.Int64("msg_id", item.msgID),
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID),
)
item.kind = inboundItemReplayRPC
item.payload = claim.encoded
case rpcResultAcquirePending:
// A malformed/replayed container may repeat the same msg_id after this
// very plan installed its owner. More generally, any request already in
// this Conn's seen table shares the owner's reliable response path. Only
// a fresh physical replacement may wait for and replay the old result.
if ownersInPlan[item.msgID] != nil {
item.kind = inboundItemDuplicate
item.payload = nil
} else {
item.kind = inboundItemPendingRPC
item.payload = claim.waiter
}
case rpcResultAcquireOwner:
if item.kind == inboundItemDuplicate {
// connState says this request was already accepted, so absence from
// both completed and in-flight tables is not authority to execute it.
claim.owner.Abort()
continue
}
if ownersInPlan == nil {
ownersInPlan = make(map[int64]*rpcResultOwnerLease)
indices = make([]int, 0, len(plan.items))
specs = make([]inboundRPCSpec, 0, len(plan.items))
}
ownersInPlan[item.msgID] = claim.owner
plan.rpcOwners = append(plan.rpcOwners, claim.owner)
item.payload = claim.owner
indices = append(indices, i)
specs = append(specs, inboundRPCSpec{method: method, size: len(item.body)})
default:
return ErrRPCResultFlightInvalid
}
}
if flightCapacity {
// One container is one API admission unit. If the cross-connection
// exactly-once table cannot claim every fresh request, none of this
// batch may reach a business handler.
for _, index := range indices {
plan.items[index].kind = inboundItemCapacityError
}
return nil
}
if len(specs) == 0 {
return nil
}
reservation, err := c.reserveInboundRPCBatch(ctx, specs)
if err != nil {
if errors.Is(err, ErrInboundRPCQueueFull) {
for _, index := range indices {
plan.items[index].kind = inboundItemCapacityError
}
return nil
}
return err
}
plan.rpcReservation = reservation
plan.rpcTasks = make([]inboundRPC, len(indices))
for i, index := range indices {
item := &plan.items[index]
body := append([]byte(nil), item.body...)
owner, _ := item.payload.(*rpcResultOwnerLease)
plan.rpcTasks[i] = s.newInboundRPCTask(c, item.msgID, specs[i].method, body, owner)
}
return nil
}
func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn, plan *inboundPlan) error {
for _, item := range plan.items {
switch item.kind {
case inboundItemDuplicate:
// The preflight plan already stages a content ACK for a locally seen
// duplicate. Do not wait for or replay its result on the same reliable
// stream; the original owner is solely responsible for that response.
continue
case inboundItemServiceDuplicate:
// Classify from the originally committed connState record, never from the
// retransmitted body. This prevents same-id payload replacement. Only an
// already cached answer is eligible for resend (rpc_drop_answer today);
// other best-effort service traffic uses a later fresh request.
if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil {
return err
}
case inboundItemReplayRPC:
if encoded, _ := item.payload.(*encodedOutboundMessage); encoded != nil {
if err := s.sendCachedRPCResult(ctx, c, encoded); err != nil {
return err
}
} else if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil {
return err
}
case inboundItemPendingRPC:
// Wait only after this plan's fresh RPC batch has become runnable; see
// executePendingRPCReplays. Blocking here could otherwise deadlock on
// an owner appended by the same container.
continue
case inboundItemPing:
if err := s.sendPong(ctx, c, item.msgID, item.payload.(mt.PingRequest).PingID); err != nil {
return err
}
case inboundItemPingDelay:
if err := s.sendPong(ctx, c, item.msgID, item.payload.(mt.PingDelayDisconnectRequest).PingID); err != nil {
return err
}
case inboundItemFutureSalts:
if err := s.sendFutureSalts(ctx, c, item.msgID, item.payload.(mt.GetFutureSaltsRequest).Num); err != nil {
return err
}
case inboundItemMsgsAck:
ids := item.payload.(int64VectorView).materialize()
c.AckServerMessages(ids)
s.log.Debug("Received msgs_ack", zap.Int64s("msg_ids", ids))
case inboundItemStateReq:
ids := item.payload.(int64VectorView).materialize()
outgoing, err := c.OutgoingStateInfo(ctx, ids)
if err != nil {
return err
}
if err := s.sendMsgsStateInfo(ctx, c, item.msgID, mergeStateInfo(outgoing, cs.stateInfo(ids))); err != nil {
return err
}
case inboundItemResendReq:
ids := item.payload.(int64VectorView).materialize()
outgoing, err := c.ResendMessages(ctx, ids)
if err != nil {
return err
}
if err := s.sendMsgsStateInfo(ctx, c, item.msgID, mergeStateInfo(outgoing, cs.stateInfo(ids))); err != nil {
return err
}
case inboundItemStateInfo:
value := item.payload.(stateInfoPayload)
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", value.reqMsgID), zap.Int("len", len(value.info)))
case inboundItemAllInfo:
value := item.payload.(allInfoPayload)
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", value.count), zap.Int("len", len(value.info)))
case inboundItemDestroySession:
if err := s.sendDestroySession(ctx, c, item.payload.(mt.DestroySessionRequest).SessionID); err != nil {
return err
}
case inboundItemHTTPWait:
value := item.payload.(mt.HTTPWaitRequest)
s.log.Debug("Received http_wait", zap.Int("max_delay", value.MaxDelay), zap.Int("wait_after", value.WaitAfter), zap.Int("max_wait", value.MaxWait))
case inboundItemDropAnswer:
value := item.payload.(mt.RPCDropAnswerRequest)
s.log.Debug("Received rpc_drop_answer", zap.Int64("req_msg_id", value.ReqMsgID))
if err := s.sendResult(ctx, c, item.msgID, &mt.RPCAnswerUnknown{}); err != nil {
return err
}
case inboundItemDestroyAuthKey:
s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", c.authKeyHex))
if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil {
s.log.Warn("Delete auth key failed", zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{})
}
c.keyDestroyed.Store(true)
s.conns.CloseSessionsForRawAuthKeyExceptConn(c.authKeyID, c)
return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{})
case inboundItemRPC:
if plan.rpcPrepared {
continue
}
if err := s.enqueueRPC(ctx, c, item.msgID, item.typeID, &bin.Buffer{Buf: item.body}); err != nil {
return err
}
case inboundItemCapacityError:
if err := s.sendResult(ctx, c, item.msgID, &mt.RPCError{
ErrorCode: 420,
ErrorMessage: "FLOOD_WAIT_1",
}); err != nil {
return err
}
default:
return fmt.Errorf("unknown inbound item kind %d", item.kind)
}
}
return nil
}
// executePendingRPCReplays joins owners already running on another physical
// connection for the same MTProto session. It never dispatches business code:
// owner Put publishes the immutable result to every waiter; owner Abort leaves
// the client free to retry after the old execution has definitely stopped.
func (s *Server) executePendingRPCReplays(ctx context.Context, c *Conn, plan *inboundPlan) error {
for _, item := range plan.items {
if item.kind != inboundItemPendingRPC {
continue
}
waiter, _ := item.payload.(*rpcResultWaiter)
if waiter == nil {
continue
}
encoded, ok, err := waiter.Wait(ctx)
if err != nil {
return err
}
if !ok || encoded == nil {
// The old owner stopped without publishing a result (normally because
// replacement cancellation reached the handler before it committed).
// A fresh-connection item still owns the decoded request body, so reacquire
// only after the prior flight is definitively gone. This is sequential
// retry, never concurrent business execution. Same-Conn seen duplicates
// have no body here and rely on the client's ordinary resend path.
if len(item.body) > 0 {
if err := s.enqueueRPC(ctx, c, item.msgID, item.typeID, &bin.Buffer{Buf: item.body}); err != nil {
return err
}
}
continue
}
if err := s.sendCachedRPCResult(ctx, c, encoded); err != nil {
return err
}
}
return nil
}

44
internal/mtprotoedge/inbound_preflight_bench_test.go Normal file
View file

@ -0,0 +1,44 @@
package mtprotoedge
import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
func BenchmarkInboundPlan32RPCContainer(b *testing.B) {
ids := proto.NewMessageIDGen(time.Now)
messages := make([]proto.Message, 32)
var request bin.Buffer
if err := (&tg.HelpGetConfigRequest{}).Encode(&request); err != nil {
b.Fatal(err)
}
requestBody := request.Copy()
for i := range messages {
messages[i] = proto.Message{
ID: ids.New(proto.MessageFromClient), SeqNo: 1 + i*2, Bytes: len(requestBody), Body: requestBody,
}
}
outerMsgID := ids.New(proto.MessageFromClient)
var container bin.Buffer
if err := (&proto.MessageContainer{Messages: messages}).Encode(&container); err != nil {
b.Fatal(err)
}
body := container.Copy()
s := New(Options{})
cs := newConnState()
b.ReportAllocs()
b.SetBytes(int64(len(body)))
b.ResetTimer()
for i := 0; i < b.N; i++ {
plan, err := s.preflightInbound(cs, outerMsgID, 64, body)
if err != nil {
b.Fatal(err)
}
plan.close()
}
}

406
internal/mtprotoedge/inbound_rpc.go
View file

@ -30,8 +30,13 @@ type inboundRPC struct {
size int
run func(context.Context) error
onTimeout func()
budget *inboundRPCGlobalReservation
ticket *inboundRPCTicket
// release drops request-scoped ownership which is independent from the body
// budget (for example a cross-connection in-flight RPC claim). It runs once
// after the request has either published a terminal result or become
// impossible to run; callers should still make the callback idempotent.
release func()
budget *inboundRPCGlobalReservation
ticket *inboundRPCTicket
}
const (
@ -76,7 +81,7 @@ type inboundRPCScheduler struct {
type inboundRPCGlobalReservation struct {
scheduler *inboundRPCScheduler
size int64
once sync.Once
released atomic.Bool
}
// inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次；
@ -92,6 +97,32 @@ type inboundRPCReservation struct {
once sync.Once
}
// inboundRPCSpec 是 container preflight 与 RPC scheduler 之间的有界 admission 描述。
// method 仅用于 metrics，size 是在 Copy 之前必须预留的 request body 字节数。
type inboundRPCSpec struct {
method string
size int
}
type inboundRPCBatchEntry struct {
global *inboundRPCGlobalReservation
method string
size int
}
// inboundRPCBatchReservation 把一个 container 内的 RPC 视为一个 admission 单元。
// reserve 一次预留整批的全局/单连接条数和字节预算；commit 一次 append
// 全部任务；abort 一次归还全部预算。这防止 container 只执行前半批。
type inboundRPCBatchReservation struct {
conn *Conn
ctx context.Context
entries []inboundRPCBatchEntry
totalSize int64
once sync.Once
}
var errInboundRPCBatchTaskCount = errors.New("inbound rpc batch task count mismatch")
func newInboundRPCScheduler(workers, maxTasks int, maxBytes int64) *inboundRPCScheduler {
if workers <= 0 {
workers = 1
@ -196,17 +227,89 @@ func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservat
return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil
}
// reserveGlobalBatch 在一次 budgetMu 临界区内检查并预留整批条数/字节。
// 返回的每个 reservation 仍由对应 task 单独归还，避免一个慢 RPC 持有
// 整个 container 已完成任务的预算。
func (s *inboundRPCScheduler) reserveGlobalBatch(sizes []int) ([]*inboundRPCGlobalReservation, string, error) {
reservations := make([]*inboundRPCGlobalReservation, len(sizes))
s.budgetMu.Lock()
defer s.budgetMu.Unlock()
select {
case <-s.stopCh:
return nil, "scheduler_closed", ErrConnClosed
default:
}
if len(sizes) > s.maxTasks-s.tasks {
return nil, "global_task_budget", ErrInboundRPCQueueFull
}
// 逐项从剩余预算中减，避免 total 和 s.bytes+total 溢出。
remaining := s.maxBytes - s.bytes
var total int64
for i, size := range sizes {
if size < 0 {
size = 0
}
size64 := int64(size)
if size64 > remaining-total {
return nil, "global_byte_budget", ErrInboundRPCQueueFull
}
total += size64
reservations[i] = &inboundRPCGlobalReservation{scheduler: s, size: size64}
}
s.tasks += len(sizes)
s.bytes += total
return reservations, "", nil
}
func (r *inboundRPCGlobalReservation) release() {
if r == nil || r.scheduler == nil {
return
}
r.once.Do(func() {
s := r.scheduler
s.budgetMu.Lock()
s.tasks--
s.bytes -= r.size
s.budgetMu.Unlock()
})
if !r.released.CompareAndSwap(false, true) {
return
}
s := r.scheduler
s.budgetMu.Lock()
s.tasks--
s.bytes -= r.size
s.budgetMu.Unlock()
}
// releaseInboundRPCGlobalBatch 使 batch abort/commit-failure 在一次全局锁内
// 归还它仍持有的全部预算。每张 ticket 的 CAS 保证与任何并发释放幂等。
func releaseInboundRPCGlobalBatch(reservations []*inboundRPCGlobalReservation) {
var (
scheduler *inboundRPCScheduler
tasks int
bytes int64
)
for _, reservation := range reservations {
if reservation == nil || reservation.scheduler == nil || !reservation.released.CompareAndSwap(false, true) {
continue
}
if scheduler == nil {
scheduler = reservation.scheduler
}
// 一个 batch 只能由同一 scheduler 创建。若未来出现混合调用，
// 仍通过单张 release 正确归还，而不会扣错 scheduler。
if reservation.scheduler != scheduler {
reservation.scheduler.budgetMu.Lock()
reservation.scheduler.tasks--
reservation.scheduler.bytes -= reservation.size
reservation.scheduler.budgetMu.Unlock()
continue
}
tasks++
bytes += reservation.size
}
if scheduler == nil || tasks == 0 {
return
}
scheduler.budgetMu.Lock()
scheduler.tasks -= tasks
scheduler.bytes -= bytes
scheduler.budgetMu.Unlock()
}
func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) {
@ -363,6 +466,10 @@ func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (
return nil, ctx.Err()
default:
}
if c.terminal.Load() {
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
}
if c.rpcScheduler == nil {
c.metrics.InboundRPCDropped(method, "scheduler_closed")
return nil, ErrConnClosed
@ -392,7 +499,7 @@ func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (
c.metrics.InboundRPCDropped(method, "context_done")
return nil, err
}
if c.rpcClosed {
if c.rpcClosed || c.terminal.Load() {
c.rpcMu.Unlock()
global.release()
c.metrics.InboundRPCDropped(method, "scheduler_closed")
@ -427,6 +534,102 @@ func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (
}, nil
}
// reserveInboundRPCBatch 必须在 container 内任何 request body Copy 前调用。
// 全局预算只锁一次，单连接预算也只锁一次；任一限制不满足时
// 整批失败，不会留下部分 task/字节 reservation。
func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpec) (*inboundRPCBatchReservation, error) {
if ctx == nil {
ctx = context.Background()
}
select {
case <-ctx.Done():
c.dropInboundRPCSpecs(specs, "context_done")
return nil, ctx.Err()
default:
}
if c.terminal.Load() {
c.dropInboundRPCSpecs(specs, "scheduler_closed")
return nil, ErrConnClosed
}
if c.rpcScheduler == nil {
c.dropInboundRPCSpecs(specs, "scheduler_closed")
return nil, ErrConnClosed
}
normalized := make([]inboundRPCSpec, len(specs))
sizes := make([]int, len(specs))
for i, spec := range specs {
if spec.size < 0 {
spec.size = 0
}
normalized[i] = spec
sizes[i] = spec.size
}
globals, reason, err := c.rpcScheduler.reserveGlobalBatch(sizes)
if err != nil {
c.dropInboundRPCSpecs(normalized, reason)
return nil, err
}
entries := make([]inboundRPCBatchEntry, len(normalized))
var totalSize int64
for i, spec := range normalized {
entries[i] = inboundRPCBatchEntry{
global: globals[i],
method: spec.method,
size: spec.size,
}
totalSize += int64(spec.size)
}
c.rpcMu.Lock()
if err := ctx.Err(); err != nil {
c.rpcMu.Unlock()
releaseInboundRPCGlobalBatch(globals)
c.dropInboundRPCSpecs(normalized, "context_done")
return nil, err
}
if c.rpcClosed || c.terminal.Load() {
c.rpcMu.Unlock()
releaseInboundRPCGlobalBatch(globals)
c.dropInboundRPCSpecs(normalized, "scheduler_closed")
return nil, ErrConnClosed
}
// 用减法比较，避免对抗性 batch 的 len 加法溢出。
availableSlots := c.rpcQueueSize - c.rpcReserved - len(c.rpcQueue)
if len(entries) > availableSlots {
c.rpcMu.Unlock()
releaseInboundRPCGlobalBatch(globals)
c.dropInboundRPCSpecs(normalized, "queue_full")
return nil, ErrInboundRPCQueueFull
}
if totalSize > maxInflightRPCBytes-c.inflightRPCBytes.Load() {
c.rpcMu.Unlock()
releaseInboundRPCGlobalBatch(globals)
c.dropInboundRPCSpecs(normalized, "byte_budget")
return nil, ErrInboundRPCQueueFull
}
c.rpcReserved += len(entries)
c.inflightRPCBytes.Add(totalSize)
// Add 与 close 的 Wait 由 rpcMu 排序：close 置 rpcClosed 后不会再发生 Add。
// 整批 reservation 只需一个 waiter；commit/abort 也只会完成一次。
c.rpcReservationWG.Add(1)
c.rpcMu.Unlock()
return &inboundRPCBatchReservation{
conn: c,
ctx: ctx,
entries: entries,
totalSize: totalSize,
}, nil
}
func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) {
for _, spec := range specs {
c.metrics.InboundRPCDropped(spec.method, reason)
}
}
// enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用
// reserveInboundRPC -> Copy -> commit，保证真正的 Copy 前预算。
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
@ -450,7 +653,7 @@ func (r *inboundRPCReservation) commit(task inboundRPC) error {
c := r.conn
c.rpcMu.Lock()
c.rpcReserved--
if c.rpcClosed {
if c.rpcClosed || c.terminal.Load() {
c.inflightRPCBytes.Add(-int64(r.size))
} else {
// The request deadline starts when admission succeeds, not when a worker
@ -525,6 +728,143 @@ func (r *inboundRPCReservation) abort() {
})
}
// commit 在一次 rpcMu 临界区内把整批 task append 到队列。
// deferSchedule=false 保持旧的立即调度语义；true 则返回一个幂等 activate
// 函数，让调用方先完成 new_session_created 等协议 barrier 再启动 worker。
//
// 延迟调度只能延迟本次 commit 新产生的 ready token；调用方应在连接的
// 首个 admission batch 使用它，不得把它当作已有 worker 的全局暂停锁。
func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bool) (activate func(), result error) {
if r == nil {
return nil, ErrConnClosed
}
result = ErrConnClosed
var (
committed bool
reschedule bool
firstQueueLen int
queueCap int
globals []*inboundRPCGlobalReservation
)
r.once.Do(func() {
c := r.conn
// A container reservation may deliberately span protocol-critical barriers
// (session ownership claim and new_session_created's physical write). Queue
// and execution latency starts only when the fully admitted batch becomes
// runnable, not while it is waiting behind those independent barriers.
enqueuedAt := time.Now()
deadline := time.Time{}
if c.rpcTimeout > 0 {
deadline = enqueuedAt.Add(c.rpcTimeout)
}
if ctxDeadline, ok := r.ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
deadline = ctxDeadline
}
globals = make([]*inboundRPCGlobalReservation, len(r.entries))
for i := range r.entries {
globals[i] = r.entries[i].global
}
c.rpcMu.Lock()
c.rpcReserved -= len(r.entries)
if len(tasks) != len(r.entries) {
c.inflightRPCBytes.Add(-r.totalSize)
result = errInboundRPCBatchTaskCount
} else if c.rpcClosed || c.terminal.Load() {
c.inflightRPCBytes.Add(-r.totalSize)
} else {
prepared := make([]inboundRPC, len(tasks))
// The request deadline starts when admission succeeds, not when a worker
// eventually dequeues the request. This bounds total queue + execution
// latency and lets a queued request emit its explicit timeout on time.
for i := range tasks {
task := tasks[i]
entry := r.entries[i]
if deadline.IsZero() {
task.ctx, task.cancel = context.WithCancel(r.ctx)
} else {
task.ctx, task.cancel = context.WithDeadline(r.ctx, deadline)
}
task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel)
task.method = entry.method
task.enqueuedAt = enqueuedAt
task.deadline = deadline
task.size = entry.size
task.budget = entry.global
ticket := &inboundRPCTicket{}
if task.onTimeout != nil {
onTimeout := task.onTimeout
var timeoutOnce sync.Once
ticket.onTimeout = func() {
timeoutOnce.Do(onTimeout)
}
task.onTimeout = ticket.onTimeout
}
task.ticket = ticket
if task.onTimeout != nil && !task.deadline.IsZero() {
taskCtx := task.ctx
task.stopTimeout = context.AfterFunc(taskCtx, func() {
if errors.Is(taskCtx.Err(), context.DeadlineExceeded) {
c.expireInboundRPCTicket(ticket)
}
})
}
prepared[i] = task
}
firstQueueLen = len(c.rpcQueue) + 1
c.rpcQueue = append(c.rpcQueue, prepared...)
queueCap = c.rpcQueueSize
if len(prepared) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
c.rpcReady = true
reschedule = true
}
committed = true
result = nil
}
c.rpcMu.Unlock()
c.rpcReservationWG.Done()
if !committed {
releaseInboundRPCGlobalBatch(globals)
}
})
if committed {
for i, entry := range r.entries {
r.conn.metrics.InboundRPCQueued(entry.method, firstQueueLen+i, queueCap)
}
if reschedule {
var once sync.Once
activate = func() {
once.Do(func() {
r.conn.rpcScheduler.schedule(r.conn)
})
}
if !deferSchedule {
activate()
activate = nil
}
}
}
return activate, result
}
func (r *inboundRPCBatchReservation) abort() {
if r == nil {
return
}
r.once.Do(func() {
c := r.conn
c.rpcMu.Lock()
c.rpcReserved -= len(r.entries)
c.inflightRPCBytes.Add(-r.totalSize)
c.rpcMu.Unlock()
c.rpcReservationWG.Done()
globals := make([]*inboundRPCGlobalReservation, len(r.entries))
for i := range r.entries {
globals[i] = r.entries[i].global
}
releaseInboundRPCGlobalBatch(globals)
})
}
func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) {
c.rpcMu.Lock()
defer c.rpcMu.Unlock()
@ -583,7 +923,7 @@ func (c *Conn) finishInboundRPC(task inboundRPC) {
if task.ticket != nil {
task.ticket.state.Store(inboundRPCTicketDone)
}
stopInboundRPCTask(task)
timeoutHandoff := stopInboundRPCTask(task)
var reschedule bool
c.rpcMu.Lock()
c.rpcRunning--
@ -594,11 +934,21 @@ func (c *Conn) finishInboundRPC(task inboundRPC) {
}
c.rpcMu.Unlock()
reservation := task.budget
release := task.release
// The scheduler budget may be reused immediately after release. Clear request-owned
// closures/context references first so slow metrics/rescheduling cannot overlap the old body
// with a newly admitted body under the same byte accounting.
task = inboundRPC{}
reservation.release()
if timeoutHandoff != nil {
// stopTimeout(false) means the deadline callback may already have read
// Running but not yet entered ticket.onTimeout. Calling the sync.Once wrapper
// here either performs or joins that response before owner release/Abort.
timeoutHandoff()
}
if release != nil {
release()
}
c.rpcWG.Done()
if reschedule {
c.rpcScheduler.schedule(c)
@ -648,7 +998,8 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
if found {
method := task.method
reservation := task.budget
stopInboundRPCTask(task)
release := task.release
_ = stopInboundRPCTask(task)
// Drop the run/context closures before returning the byte reservation. Otherwise an
// onTimeout callback that blocks or performs a slow write can keep the copied request body
// reachable after the global scheduler has advertised those bytes as available again.
@ -658,6 +1009,9 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
if ticket.onTimeout != nil {
ticket.onTimeout()
}
if release != nil {
release()
}
return
}
if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil {
@ -667,11 +1021,17 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
// stopInboundRPCTask disarms callbacks before canceling the context so a normal
// completion or connection close cannot manufacture an RPC_TIMEOUT response.
// A deadline callback already in flight is harmless because enqueueRPC's response
// gate makes timeout and normal rpc_result mutually exclusive.
func stopInboundRPCTask(task inboundRPC) {
// If the runtime already started a deadline callback, the returned sync.Once
// wrapper is a mandatory handoff: callers invoke it before owner release so a
// callback paused between ticket-state inspection and response-gate claim cannot
// publish into a later flight generation.
func stopInboundRPCTask(task inboundRPC) (timeoutHandoff func()) {
if task.stopTimeout != nil {
task.stopTimeout()
stopped := task.stopTimeout()
if !stopped && task.ctx != nil && errors.Is(task.ctx.Err(), context.DeadlineExceeded) &&
task.ticket != nil && task.ticket.onTimeout != nil {
timeoutHandoff = task.ticket.onTimeout
}
}
if task.stopRoot != nil {
task.stopRoot()
@ -679,6 +1039,7 @@ func stopInboundRPCTask(task inboundRPC) {
if task.cancel != nil {
task.cancel()
}
return timeoutHandoff
}
func (c *Conn) closeInboundRPCScheduler() {
@ -722,9 +1083,16 @@ func (c *Conn) beginCloseInboundRPCScheduler() {
}
method := task.method
reservation := task.budget
stopInboundRPCTask(task)
release := task.release
timeoutHandoff := stopInboundRPCTask(task)
task = inboundRPC{}
reservation.release()
if timeoutHandoff != nil {
timeoutHandoff()
}
if release != nil {
release()
}
c.metrics.InboundRPCDropped(method, "connection_closed")
}
})

343
internal/mtprotoedge/inbound_rpc_batch_test.go Normal file
View file

@ -0,0 +1,343 @@
package mtprotoedge
import (
"context"
"errors"
"testing"
"time"
)
func TestInboundRPCBatchReservationRejectsGloballyWithoutPartialBudget(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 2, 10)
c := newInboundTestConn(scheduler, 1, 8, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
_, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 3},
{method: "three", size: 3},
})
if !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("reserve over global task budget err = %v, want ErrInboundRPCQueueFull", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("global budget after atomic rejection = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after global rejection = %d, want zero", got)
}
c.rpcMu.Lock()
reserved := c.rpcReserved
queued := len(c.rpcQueue)
c.rpcMu.Unlock()
if reserved != 0 || queued != 0 {
t.Fatalf("connection state after global rejection = reserved %d queued %d, want zero", reserved, queued)
}
}
func TestInboundRPCBatchReservationRejectsConnectionWithoutLeakingGlobalBudget(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
c := newInboundTestConn(scheduler, 1, 2, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
_, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
{method: "three", size: 7},
})
if !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("reserve over connection queue budget err = %v, want ErrInboundRPCQueueFull", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("global budget after connection rejection = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after connection rejection = %d, want zero", got)
}
c.rpcMu.Lock()
reserved := c.rpcReserved
queued := len(c.rpcQueue)
c.rpcMu.Unlock()
if reserved != 0 || queued != 0 {
t.Fatalf("connection state after connection rejection = reserved %d queued %d, want zero", reserved, queued)
}
}
func TestInboundRPCBatchReservationRejectsAggregateConnectionBytes(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, int64(maxInflightRPCBytes)*2)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
halfPlusOne := maxInflightRPCBytes/2 + 1
_, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: halfPlusOne},
{method: "two", size: halfPlusOne},
})
if !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("reserve over aggregate connection byte budget err = %v, want ErrInboundRPCQueueFull", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("global budget after aggregate byte rejection = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after aggregate byte rejection = %d, want zero", got)
}
}
func TestInboundRPCBatchAbortReturnsEveryReservationExactlyOnce(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
{method: "three", size: 7},
})
if err != nil {
t.Fatalf("reserve batch: %v", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 3 || bytes != 15 {
t.Fatalf("budget after reserve = (%d, %d), want (3, 15)", tasks, bytes)
}
reservation.abort()
reservation.abort()
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after idempotent abort = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after abort = %d, want zero", got)
}
c.rpcMu.Lock()
reserved := c.rpcReserved
queued := len(c.rpcQueue)
c.rpcMu.Unlock()
if reserved != 0 || queued != 0 {
t.Fatalf("connection state after abort = reserved %d queued %d, want zero", reserved, queued)
}
}
func TestInboundRPCBatchCommitAppendsAllAndDefersSchedule(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
scheduler.start()
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
specs := []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
{method: "three", size: 7},
}
reservation, err := c.reserveInboundRPCBatch(context.Background(), specs)
if err != nil {
t.Fatalf("reserve batch: %v", err)
}
defer reservation.abort()
runs := make(chan string, len(specs))
tasks := make([]inboundRPC, len(specs))
for i, spec := range specs {
method := spec.method
tasks[i].run = func(context.Context) error {
runs <- method
return nil
}
}
activate, err := reservation.commit(tasks, true)
if err != nil {
t.Fatalf("commit batch: %v", err)
}
if activate == nil {
t.Fatal("deferred commit did not return an activation function")
}
c.rpcMu.Lock()
queued := len(c.rpcQueue)
ready := c.rpcReady
c.rpcMu.Unlock()
if queued != len(specs) || !ready {
t.Fatalf("atomic queue state after commit = queued %d ready %v, want %d/true", queued, ready, len(specs))
}
if got := scheduler.readyLen(); got != 0 {
t.Fatalf("scheduler ready tokens before activation = %d, want zero", got)
}
select {
case method := <-runs:
t.Fatalf("RPC %q ran before deferred activation", method)
default:
}
activate()
activate() // activation is idempotent.
for _, want := range []string{"one", "two", "three"} {
select {
case got := <-runs:
if got != want {
t.Fatalf("execution order = %q, want %q", got, want)
}
case <-time.After(time.Second):
t.Fatalf("timed out waiting for %q", want)
}
}
waitInboundRPCBatchBudget(t, scheduler, 0, 0)
}
func TestInboundRPCBatchCommitMismatchReleasesAllWithoutEnqueue(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
})
if err != nil {
t.Fatalf("reserve batch: %v", err)
}
if _, err := reservation.commit([]inboundRPC{{}}, false); !errors.Is(err, errInboundRPCBatchTaskCount) {
t.Fatalf("commit task mismatch err = %v, want %v", err, errInboundRPCBatchTaskCount)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after mismatched commit = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after mismatched commit = %d, want zero", got)
}
c.rpcMu.Lock()
reserved := c.rpcReserved
queued := len(c.rpcQueue)
c.rpcMu.Unlock()
if reserved != 0 || queued != 0 {
t.Fatalf("connection state after mismatched commit = reserved %d queued %d, want zero", reserved, queued)
}
}
func TestInboundRPCBatchCommitRacingCloseNeverPartiallyEnqueues(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer scheduler.stop(time.Second)
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
{method: "three", size: 7},
})
if err != nil {
t.Fatalf("reserve batch: %v", err)
}
closed := make(chan struct{})
go func() {
c.closeInboundRPCScheduler()
close(closed)
}()
waitInboundRPCBatchConnClosed(t, c)
if _, err := reservation.commit(make([]inboundRPC, 3), false); !errors.Is(err, ErrConnClosed) {
t.Fatalf("commit after close err = %v, want ErrConnClosed", err)
}
select {
case <-closed:
case <-time.After(time.Second):
t.Fatal("close did not finish after batch commit returned its reservation")
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after close/commit race = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after close/commit race = %d, want zero", got)
}
c.rpcMu.Lock()
queued := len(c.rpcQueue)
c.rpcMu.Unlock()
if queued != 0 {
t.Fatalf("queue after close/commit race = %d, want zero", queued)
}
}
func TestInboundRPCBatchCommitAfterTerminalFenceRejectsAll(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
})
if err != nil {
t.Fatalf("reserve batch: %v", err)
}
// Session replacement and revocation publish terminal before the slower
// physical-close path. A reservation held across that fence must not be able
// to append even one stale task.
c.terminal.Store(true)
if _, err := reservation.commit(make([]inboundRPC, 2), false); !errors.Is(err, ErrConnClosed) {
t.Fatalf("commit after terminal fence err = %v, want ErrConnClosed", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("budget after terminal commit rejection = (%d, %d), want zero", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection bytes after terminal commit rejection = %d, want zero", got)
}
c.rpcMu.Lock()
reserved := c.rpcReserved
queued := len(c.rpcQueue)
c.rpcMu.Unlock()
if reserved != 0 || queued != 0 {
t.Fatalf("connection state after terminal commit rejection = reserved %d queued %d, want zero", reserved, queued)
}
}
func waitInboundRPCBatchBudget(t *testing.T, scheduler *inboundRPCScheduler, wantTasks int, wantBytes int64) {
t.Helper()
deadline := time.Now().Add(time.Second)
for {
tasks, bytes := scheduler.budgetSnapshot()
if tasks == wantTasks && bytes == wantBytes {
return
}
if time.Now().After(deadline) {
t.Fatalf("budget = (%d, %d), want (%d, %d)", tasks, bytes, wantTasks, wantBytes)
}
time.Sleep(time.Millisecond)
}
}
func waitInboundRPCBatchConnClosed(t *testing.T, c *Conn) {
t.Helper()
deadline := time.Now().Add(time.Second)
for {
c.rpcMu.Lock()
closed := c.rpcClosed
c.rpcMu.Unlock()
if closed {
return
}
if time.Now().After(deadline) {
t.Fatal("connection scheduler was not marked closed")
}
time.Sleep(time.Millisecond)
}
}

47
internal/mtprotoedge/inbound_rpc_test.go
View file

@ -3,11 +3,58 @@ package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestStopInboundRPCTaskJoinsStartedDeadlineCallback(t *testing.T) {
ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second))
defer cancel()
started := make(chan struct{})
release := make(chan struct{})
var once sync.Once
ticket := &inboundRPCTicket{}
ticket.onTimeout = func() {
once.Do(func() {
close(started)
<-release
})
}
task := inboundRPC{
ctx: ctx,
ticket: ticket,
stopTimeout: func() bool { return false }, // runtime callback already scheduled
}
handoff := stopInboundRPCTask(task)
if handoff == nil {
t.Fatal("started deadline callback did not produce handoff")
}
go ticket.onTimeout()
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("deadline callback did not start")
}
joined := make(chan struct{})
go func() {
handoff()
close(joined)
}()
select {
case <-joined:
t.Fatal("handoff returned before in-flight callback completed")
case <-time.After(20 * time.Millisecond):
}
close(release)
select {
case <-joined:
case <-time.After(time.Second):
t.Fatal("handoff did not join completed callback")
}
}
func newInboundTestConn(s *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) *Conn {
c := &Conn{metrics: NopMetrics{}}
c.startInboundRPCScheduler(s, maxInflight, queueSize, timeout)

4
internal/mtprotoedge/metrics_test.go
View file

@ -61,8 +61,8 @@ func TestMetricsHooks(t *testing.T) {
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 验证发送计数，不约束同步时序。
// new_session_created 已经跨过 required-control 物理写屏障；msgs_ack 仍是异步控制帧，
// 因而这里只轮询等待最终发送计数。具体 boundary→RPC 顺序由 session boundary 测试锁定。
deadline := time.Now().Add(2 * time.Second)
for m.outbound.Load() < 3 && time.Now().Before(deadline) {
time.Sleep(5 * time.Millisecond)

47
internal/mtprotoedge/new_session_uid_test.go
View file

@ -43,3 +43,50 @@ func TestNewSessionCreatedUniqueIDPerSession(t *testing.T) {
t.Fatalf("new_session_created.unique_id reused across sessions: %d", created1.UniqueID)
}
}
// TestNewSessionCreatedMovesFloorForLaterLowerMessage locks the MTProto rule
// that a server-side session must publish a fresh boundary notification when it
// later accepts a smaller logical client msg_id. Official clients use this
// boundary to decide which requests must be regenerated and resent.
func TestNewSessionCreatedMovesFloorForLaterLowerMessage(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
freshIDs := proto.NewMessageIDGen(time.Now)
firstMsgID := freshIDs.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, firstMsgID, 0, &mt.HTTPWaitRequest{
MaxDelay: 0,
WaitAfter: 0,
MaxWait: 25_000,
})
first := collectReplies(t, conn, cipher, auth.AuthKey, mt.NewSessionCreatedTypeID)
var created1 mt.NewSessionCreated
if err := created1.Decode(mustHave(t, first, mt.NewSessionCreatedTypeID, "first new_session_created")); err != nil {
t.Fatalf("decode first new_session_created: %v", err)
}
if created1.FirstMsgID != firstMsgID {
t.Fatalf("first boundary = %d, want %d", created1.FirstMsgID, firstMsgID)
}
oldIDs := proto.NewMessageIDGen(func() time.Time { return time.Now().Add(-10 * time.Minute) })
lowerMsgID := oldIDs.New(proto.MessageFromClient)
outerMsgID := freshIDs.New(proto.MessageFromClient)
pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 99})
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{
Messages: []proto.Message{{
ID: lowerMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody,
}},
})
second := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
var created2 mt.NewSessionCreated
if err := created2.Decode(mustHave(t, second, mt.NewSessionCreatedTypeID, "lower-boundary new_session_created")); err != nil {
t.Fatalf("decode lower-boundary new_session_created: %v", err)
}
if created2.FirstMsgID != lowerMsgID {
t.Fatalf("lower boundary = %d, want %d", created2.FirstMsgID, lowerMsgID)
}
if created2.UniqueID == created1.UniqueID {
t.Fatalf("lower-boundary notification reused unique_id %d", created2.UniqueID)
}
}

119
internal/mtprotoedge/outbound.go
View file

@ -40,6 +40,11 @@ const (
defaultOutboundControlQueueSize = 32
defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server
defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server
// requiredControlMaxWait bounds protocol barriers such as new_session_created from
// the beginning of encoding through the completed physical write. These frames gate
// subsequent session state transitions, so timing out must close the connection instead
// of degrading to the best-effort control path.
requiredControlMaxWait = 5 * time.Second
maxTrackedServerMsgIDs = 4096
maxTrackedAckedMsgIDs = 1024
@ -332,6 +337,9 @@ func (c *Conn) Close() {
// transport.Close call, so a timed-out batch close cannot keep accepting memory/work.
func (c *Conn) beginTerminalShutdown() {
c.terminal.Store(true)
// Retirement is irreversible. SessionManager activation only uses CAS from
// provisional/claiming, so a stale goroutine cannot publish this Conn again.
c.lifecycle.Store(uint32(connLifecycleRetired))
c.signalOutboundStop()
c.beginCloseInboundRPCScheduler()
}
@ -342,6 +350,23 @@ func (c *Conn) waitOutboundShutdown() {
}
}
func (c *Conn) waitOutboundShutdownUntil(timeout time.Duration) bool {
if c == nil || c.outboundDone == nil {
return true
}
if timeout <= 0 {
return false
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-c.outboundDone:
return true
case <-timer.C:
return false
}
}
// ForceClose 停止连接并关闭底层 transport。
// 仅用于授权撤销 / destroy_auth_key 这类“必须让对端立即断线”的路径；普通生命周期仍由
// serveConn 统一关闭 transport，避免正常 push/索引清理把长连接误伤成硬断。
@ -355,6 +380,13 @@ func (c *Conn) ForceClose() {
// closeTransport 只关闭物理 transport，不等待 outbound actor。写失败路径运行在
// actor 自身 goroutine 中，若在这里调用 Close 会等待 outboundDone 而自锁。
func (c *Conn) closeTransport() {
if c == nil {
return
}
if c.transportLease != nil {
_ = c.transportLease.Close()
return
}
c.transportClose.Do(func() {
if c.transport != nil {
_ = c.transport.Close()
@ -372,6 +404,38 @@ func (c *Conn) failTransport() {
c.closeTransport()
}
// fenceUndeliveredRPCResult is the no-reentry terminal path used from a task's
// release callback. That callback may itself run while rpcClose.Do is draining
// queued tasks, so calling beginCloseInboundRPCScheduler again would deadlock on
// sync.Once. Closing the socket wakes serveConn, whose ordinary defer completes
// scheduler/index cleanup; when shutdown already owns the callback, that cleanup
// is already in progress.
func (c *Conn) fenceUndeliveredRPCResult() {
if c == nil {
return
}
// A replacement/shutdown that already published terminal owns physical
// lifecycle cleanup (and may intentionally transfer the lease). Only the
// resultless task that wins false->true is allowed to close this generation.
if !c.terminal.CompareAndSwap(false, true) {
return
}
c.lifecycle.Store(uint32(connLifecycleRetired))
c.signalOutboundStop()
if c.transportLease != nil {
c.transportLease.startCloseAlreadyFenced()
return
}
// Legacy construction-only Conns have no owner callback graph, so their
// exact transport close cannot re-enter logical lifecycle cleanup. Keep a
// pathological Close outside the shared RPC worker just like the lease path.
go c.transportClose.Do(func() {
if c.transport != nil {
_ = c.transport.Close()
}
})
}
// dropSlowConsumer 把出站队列持续拥塞的连接降级为离线连接。它不能等待 outbound
// actor：调用方位于 fan-out 热路径，等待单个慢 socket 会把同一用户的健康设备和
// transactional outbox lane 一起拖住。关闭 transport 会打断可能阻塞的写；serveConn
@ -402,6 +466,41 @@ func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.En
return c.send(ctx, t, msg, true)
}
// SendRequiredControl writes a protocol-critical control message before the caller commits
// the state transition guarded by that message. One absolute deadline covers encode admission,
// body-budget reservation, control-queue admission and the physical transport write. A failure
// is terminal: continuing on the same connection could expose state whose required notification
// never reached the client.
//
// Success only confirms the physical write; it does not wait for the client's msgs_ack.
func (c *Conn) SendRequiredControl(ctx context.Context, t proto.MessageType, msg bin.Encoder) error {
if ctx == nil {
ctx = context.Background()
}
now := time.Now()
deadline := now.Add(requiredControlMaxWait)
if c.writeTimeout > 0 {
if writeDeadline := now.Add(c.writeTimeout); writeDeadline.Before(deadline) {
deadline = writeDeadline
}
}
if parentDeadline, ok := ctx.Deadline(); ok && parentDeadline.Before(deadline) {
deadline = parentDeadline
}
requiredCtx, cancel := context.WithDeadline(ctx, deadline)
defer cancel()
if err := requiredCtx.Err(); err != nil {
c.failTransport()
return err
}
err := c.sendOutbound(requiredCtx, t, msg, nil, true)
if err != nil {
c.failTransport()
}
return err
}
// 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 {
@ -519,14 +618,27 @@ func (c *Conn) sendOutbound(ctx context.Context, t proto.MessageType, msg bin.En
case res := <-op.done:
return res.err
case <-ctx.Done():
// A physical write can complete at the same instant as the caller's
// deadline. Prefer the actor's terminal result when it is already
// available so required-control callers do not poison a healthy Conn.
select {
case res := <-op.done:
return res.err
default:
}
return ctx.Err()
case <-c.outboundStop:
select {
case res := <-op.done:
return res.err
default:
}
return ErrConnClosed
}
}
// SendAsync 入队一条 server 消息但不等待发送结果（fire-and-forget），用于读循环里的控制消息
// （ack/pong/new_session_created/bad_msg/future_salts/state_info）：避免读循环被 outbound 写
// （ack/pong/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 {
@ -1038,7 +1150,10 @@ func (c *Conn) failOutboundBudget(err error) {
if c.metrics != nil {
c.metrics.OutboundDropped("tracked_global_byte_budget")
}
c.failTransport()
// No socket bytes exist yet. If an intentional session handoff already won
// the terminal CAS, it owns close/transfer and this old producer must not close
// the still-current lease. A live connection still gets fenced and closed.
c.fenceUndeliveredRPCResult()
}
func (c *Conn) ensureOutboundTrackedBudget() *outboundTrackedBudget {

262
internal/mtprotoedge/outbound_required_control_test.go Normal file
View file

@ -0,0 +1,262 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
)
type gatedRequiredControlTransport struct {
started chan struct{}
release chan struct{}
sendErr error
startOnce sync.Once
closeOnce sync.Once
sends atomic.Int32
closes atomic.Int32
}
func newGatedRequiredControlTransport(sendErr error) *gatedRequiredControlTransport {
return &gatedRequiredControlTransport{
started: make(chan struct{}),
release: make(chan struct{}),
sendErr: sendErr,
}
}
func (t *gatedRequiredControlTransport) Send(context.Context, *bin.Buffer) error {
t.sends.Add(1)
t.startOnce.Do(func() { close(t.started) })
<-t.release
return t.sendErr
}
func (t *gatedRequiredControlTransport) Recv(context.Context, *bin.Buffer) error {
return io.EOF
}
func (t *gatedRequiredControlTransport) Close() error {
t.closes.Add(1)
t.closeOnce.Do(func() { close(t.release) })
return nil
}
func (t *gatedRequiredControlTransport) unblock() {
t.closeOnce.Do(func() { close(t.release) })
}
func TestSendRequiredControlWaitsForPhysicalWriteAndReturnsBudget(t *testing.T) {
tr := newGatedRequiredControlTransport(nil)
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
done := make(chan error, 1)
go func() {
done <- c.SendRequiredControl(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("required control did not reach the physical writer")
}
select {
case err := <-done:
t.Fatalf("SendRequiredControl returned before physical write completed: %v", err)
case <-time.After(20 * time.Millisecond):
}
tr.unblock()
select {
case err := <-done:
if err != nil {
t.Fatalf("SendRequiredControl: %v", err)
}
case <-time.After(time.Second):
t.Fatal("SendRequiredControl did not return after physical write")
}
if c.terminal.Load() {
t.Fatal("successful required control terminally closed the connection")
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after non-pending physical write = %d, want 0", got)
}
if got := tr.sends.Load(); got != 1 {
t.Fatalf("physical sends = %d, want 1", got)
}
}
func TestSendRequiredControlReturnsAfterWriteWithoutWaitingForAck(t *testing.T) {
tr := &failAfterTransport{}
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
created := &mt.NewSessionCreated{FirstMsgID: 1, UniqueID: 2, ServerSalt: 3}
encoded, err := encodeOutboundMessageWithoutSlot(created)
if err != nil {
t.Fatalf("encode new_session_created: %v", err)
}
if err := c.SendRequiredControl(context.Background(), proto.MessageFromServer, created); err != nil {
t.Fatalf("SendRequiredControl: %v", err)
}
if got := tr.stored.Load(); got != 1 {
t.Fatalf("completed physical sends = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != int64(len(encoded.body)) {
t.Fatalf("pending control budget = %d, want %d until client ACK", got, len(encoded.body))
}
frame, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt new_session_created: %v", err)
}
c.AckServerMessages([]int64{frame.MessageID})
deadline := time.Now().Add(time.Second)
for controlBudget.snapshot() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after ACK = %d, want 0", got)
}
}
func TestSendRequiredControlQueueDeadlineTerminatesAndReturnsBudget(t *testing.T) {
tr := &failAfterTransport{}
controlBudget := newOutboundTrackedBudget(1 << 20)
c := &Conn{
transport: tr,
writer: tr,
metrics: NopMetrics{},
writeTimeout: time.Second,
outboundTrackedBudget: newOutboundTrackedBudget(1 << 20),
outboundControlTrackedBudget: controlBudget,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
// No actor is running and the bounded control queue is full, so the parent
// deadline must cover queue admission and make the failure terminal.
c.outboundControl <- outboundOp{kind: outboundAck}
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
defer cancel()
started := time.Now()
err := c.SendRequiredControl(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("full control queue error = %v, want context deadline", err)
}
if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
t.Fatalf("full control queue waited %v, want parent-deadline-bounded admission", elapsed)
}
if !c.terminal.Load() {
t.Fatal("required control queue failure did not terminally close the connection")
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("physical sends = %d, want 0", got)
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after queue timeout = %d, want 0", got)
}
}
func TestSendRequiredControlBlockedWriteUsesWholeOperationDeadline(t *testing.T) {
tr := newGatedRequiredControlTransport(io.ErrClosedPipe)
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
c.writeTimeout = 25 * time.Millisecond
started := time.Now()
err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("blocked required control error = %v, want context deadline", err)
}
if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
t.Fatalf("blocked required control waited %v, want write-timeout-bounded operation", elapsed)
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control write timeout")
}
if !c.terminal.Load() {
t.Fatal("blocked required control did not terminally close the connection")
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after blocked write = %d, want 0", got)
}
}
func TestSendRequiredControlWriteFailureTerminatesAndReturnsBudget(t *testing.T) {
tr := &failAfterTransport{}
tr.failAt.Store(1)
controlBudget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if err == nil {
t.Fatal("write-failed required control unexpectedly succeeded")
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control write failure")
}
if !c.terminal.Load() {
t.Fatal("write-failed required control did not terminally close the connection")
}
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after write failure = %d, want 0", got)
}
}
func TestSendRequiredControlBudgetFailureIsTerminal(t *testing.T) {
tr := &failAfterTransport{}
controlBudget := newOutboundTrackedBudget(1)
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
c.outboundControlTrackedBudget = controlBudget
err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2})
if !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("required control over budget = %v, want ErrOutboundTrackedBudget", err)
}
select {
case <-c.outboundDone:
case <-time.After(time.Second):
t.Fatal("outbound actor did not stop after required-control budget failure")
}
if !c.terminal.Load() {
t.Fatal("required-control budget failure did not terminally close the connection")
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("budget-rejected required control wrote %d frames, want 0", got)
}
if got := controlBudget.snapshot(); got != 0 {
t.Fatalf("control budget after reservation failure = %d, want 0", got)
}
}

202
internal/mtprotoedge/provisional_revocation_test.go Normal file
View file

@ -0,0 +1,202 @@
package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/tg"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
type activationGatedAuthKeyStore struct {
store.AuthKeyStore
gets atomic.Int32
finalStarted chan struct{}
finalRelease chan struct{}
startOnce sync.Once
}
func (s *activationGatedAuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
if s.gets.Add(1) == 2 {
s.startOnce.Do(func() { close(s.finalStarted) })
select {
case <-s.finalRelease:
case <-ctx.Done():
return store.AuthKeyData{}, false, ctx.Err()
}
}
return s.AuthKeyStore.Get(ctx, id)
}
func waitForManagedSessionAbsent(t *testing.T, manager *SessionManager, key sessionKey) {
t.Helper()
deadline := time.Now().Add(2 * time.Second)
for {
manager.mu.RLock()
claim, active := manager.claims[key], manager.bySession[key]
manager.mu.RUnlock()
if claim == nil && active == nil {
return
}
if time.Now().After(deadline) {
t.Fatalf("managed session survived terminal rejection: claim=%p active=%p", claim, active)
}
time.Sleep(time.Millisecond)
}
}
func TestBadSaltStormRevalidatesStoreOnlyAtActivationBoundary(t *testing.T) {
const dc = 2
keys := &countingAuthKeyStore{AuthKeyStore: memory.NewAuthKeyStore()}
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, AuthKeys: keys, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
firstID := ids.New(proto.MessageFromClient)
const wrongFrames = 16
for i := 0; i < wrongFrames; i++ {
msgID := firstID
if i != 0 {
msgID = ids.New(proto.MessageFromClient)
}
sendEncryptedWithSalt(t, conn, cipher, auth, auth.ServerSalt+1, msgID, &tg.HelpGetConfigRequest{})
_, typeID, _ := readServerMessage(t, conn, cipher, auth.AuthKey)
if typeID != mt.BadServerSaltTypeID {
t.Fatalf("correction %d type = %#x", i, typeID)
}
}
if got := keys.gets.Load(); got != 1 {
t.Fatalf("AuthKeyStore.Get during bad-salt storm = %d, want initial lookup only", got)
}
sendEncrypted(t, conn, cipher, auth, firstID, &tg.HelpGetConfigRequest{})
collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 1})
if got := keys.gets.Load(); got != 2 {
t.Fatalf("AuthKeyStore.Get after activation boundary = %d, want 2", got)
}
waitForAtomicCalls(t, &handler.calls, 1)
}
func TestActivationFinalAuthKeyCheckRunsAfterClaim(t *testing.T) {
const dc = 2
base := memory.NewAuthKeyStore()
keys := &activationGatedAuthKeyStore{
AuthKeyStore: base,
finalStarted: make(chan struct{}),
finalRelease: make(chan struct{}),
}
defer func() {
select {
case <-keys.finalRelease:
default:
close(keys.finalRelease)
}
}()
handler := &admissionCountingRPC{}
addr, pub, srv := startTestServer(t, Options{DC: dc, AuthKeys: keys, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
// The first Get is serveConn's decrypt lookup. The second is deliberately
// blocked: it must start only after BeginActivation indexed the claim.
sendEncrypted(t, conn, cipher, auth, msgID, &tg.HelpGetConfigRequest{})
select {
case <-keys.finalStarted:
case <-time.After(2 * time.Second):
t.Fatal("activation final auth-key check did not start")
}
key := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID}
srv.conns.mu.RLock()
claim, active := srv.conns.claims[key], srv.conns.bySession[key]
srv.conns.mu.RUnlock()
if claim == nil || active != nil {
t.Fatalf("final auth-key check not protected by claim: claim=%p active=%p", claim, active)
}
// Model Delete committing after the initial decrypt lookup but before the
// activation check returns. The claimant must emit terminal -404, never publish
// or dispatch the request, even before revocation fan-out gets the manager lock.
if err := base.Delete(context.Background(), auth.AuthKey.ID); err != nil {
t.Fatalf("delete auth key during activation: %v", err)
}
close(keys.finalRelease)
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
var response bin.Buffer
err := conn.Recv(ctx, &response)
var protocolErr *codec.ProtocolErr
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
t.Fatalf("revoked activation recv = %T %v, want protocol -404", err, err)
}
if got := handler.calls.Load(); got != 0 {
t.Fatalf("revoked activation executed %d RPCs", got)
}
waitForManagedSessionAbsent(t, srv.conns, key)
}
func TestBadSaltProvisionalCannotReactivateDeletedAuthKey(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, srv := startTestServer(t, Options{DC: dc, RPC: handler})
provisional, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
reqMsgID := ids.New(proto.MessageFromClient)
// Socket A is retained as a bad-salt provisional and is intentionally absent
// from SessionManager active/claim indexes.
sendEncryptedWithSalt(t, provisional, cipher, auth, auth.ServerSalt+1, reqMsgID, &tg.HelpGetConfigRequest{})
_, typeID, _ := readServerMessage(t, provisional, cipher, auth.AuthKey)
if typeID != mt.BadServerSaltTypeID {
t.Fatalf("provisional correction type = %#x, want bad_server_salt", typeID)
}
key := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID}
srv.conns.mu.RLock()
activeBefore, claimBefore := srv.conns.bySession[key], srv.conns.claims[key]
srv.conns.mu.RUnlock()
if activeBefore != nil || claimBefore != nil {
t.Fatalf("bad-salt provisional leaked into manager: active=%p claim=%p", activeBefore, claimBefore)
}
// Socket B uses the same auth key with another session and deletes it. The
// provisional is not manager-indexed, so correctness depends on its next-frame
// AuthKeyStore recheck rather than fan-out close alone.
destroyer := dialTransportOnly(t, addr)
destroySessionID := auth.SessionID ^ 1
destroyBody := encodeClientMessageBodyForTest(t, &destroyAuthKeyRequest{})
sendEncryptedWithSessionSaltAndSeq(
t, destroyer, cipher, auth, destroySessionID, auth.ServerSalt,
ids.New(proto.MessageFromClient), 1, destroyBody,
)
destroyReplies := collectReplies(t, destroyer, cipher, auth.AuthKey, destroyAuthKeyOkTypeID)
mustHave(t, destroyReplies, destroyAuthKeyOkTypeID, "destroy_auth_key_ok")
if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || found {
t.Fatalf("auth key after destroy: found=%v err=%v", found, err)
}
// A corrected resend must now receive terminal -404; it must not activate or
// execute the previously rejected business request with its cached key.
sendEncrypted(t, provisional, cipher, auth, reqMsgID, &tg.HelpGetConfigRequest{})
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
var response bin.Buffer
err := provisional.Recv(ctx, &response)
var protocolErr *codec.ProtocolErr
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
t.Fatalf("corrected revoked provisional recv = %T %v, want protocol -404", err, err)
}
if got := handler.calls.Load(); got != 0 {
t.Fatalf("revoked provisional executed %d RPCs", got)
}
waitForManagedSessionAbsent(t, srv.conns, key)
}

207
internal/mtprotoedge/rpc_flight.go Normal file
View file

@ -0,0 +1,207 @@
package mtprotoedge
import (
"context"
"errors"
"sync/atomic"
)
const rpcResultFlightDefaultMaxPending = 8192
var (
// ErrRPCResultFlightCapacity is returned before installing a new owner when
// the process-wide in-flight claim table has reached its hard bound.
ErrRPCResultFlightCapacity = errors.New("mtproto rpc result in-flight capacity exhausted")
ErrRPCResultFlightInvalid = errors.New("mtproto rpc result in-flight claim is invalid")
)
type rpcResultAcquireState uint8
const (
rpcResultAcquireCompleted rpcResultAcquireState = iota + 1
rpcResultAcquirePending
rpcResultAcquireOwner
)
// rpcResultAcquire is the atomic outcome for one
// (auth_key_id, session_id, req_msg_id) claim.
//
// Exactly one state-specific field is non-nil:
// - completed: encoded contains the immutable completed rpc_result;
// - pending: waiter joins the already-running owner;
// - owner: owner must eventually complete through rpcResultCache.Put or Abort.
type rpcResultAcquire struct {
state rpcResultAcquireState
encoded *encodedOutboundMessage
waiter *rpcResultWaiter
owner *rpcResultOwnerLease
}
// rpcResultFlight is not part of the completed cache LRU/TTL lifecycle. Its
// done channel is closed exactly once while holding the owning cache shard lock;
// channel close publishes encoded/ok to all waiters without a waiter goroutine.
type rpcResultFlight struct {
done chan struct{}
encoded *encodedOutboundMessage
ok bool
}
type rpcResultWaiter struct {
flight *rpcResultFlight
}
// Wait blocks until the owner publishes through Put, aborts, or ctx expires.
// ok=false with err=nil means the owner aborted without a result.
func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMessage, ok bool, err error) {
if w == nil || w.flight == nil || ctx == nil {
return nil, false, ErrRPCResultFlightInvalid
}
// Prefer an already-published result over a concurrently canceled context.
select {
case <-w.flight.done:
return w.flight.encoded, w.flight.ok, nil
default:
}
select {
case <-w.flight.done:
return w.flight.encoded, w.flight.ok, nil
case <-ctx.Done():
// If completion raced with cancellation, prefer the terminal flight state.
select {
case <-w.flight.done:
return w.flight.encoded, w.flight.ok, nil
default:
return nil, false, ctx.Err()
}
}
}
type rpcResultOwnerLease struct {
cache *rpcResultCache
key rpcResultCacheKey
flight *rpcResultFlight
}
// Abort releases an unfinished owner claim and wakes every waiter with no
// result. Pointer identity prevents an old lease from deleting a later owner
// that reacquired the same key. It returns true only for the winning abort.
func (l *rpcResultOwnerLease) Abort() bool {
if l == nil || l.cache == nil || l.flight == nil {
return false
}
s := l.cache.shard(l.key)
s.mu.Lock()
defer s.mu.Unlock()
flight, ok := s.pending[l.key]
if !ok || flight != l.flight {
return false
}
delete(s.pending, l.key)
l.cache.flightLimit.release()
close(flight.done)
return true
}
type rpcResultFlightLimit struct {
max int64
used atomic.Int64
}
func (l *rpcResultFlightLimit) reserve() bool {
if l == nil || l.max <= 0 {
return false
}
for {
used := l.used.Load()
if used >= l.max {
return false
}
if l.used.CompareAndSwap(used, used+1) {
return true
}
}
}
func (l *rpcResultFlightLimit) release() {
if l == nil {
return
}
if remaining := l.used.Add(-1); remaining < 0 {
// Put/Abort use map removal and lease identity to make double release
// impossible. Fail fast instead of masking a capacity-accounting bug that
// could otherwise admit more owners than the configured hard limit.
panic("mtproto rpc result in-flight counter underflow")
}
}
func (l *rpcResultFlightLimit) snapshot() int64 {
if l == nil {
return 0
}
return l.used.Load()
}
// Acquire atomically returns a completed result, joins the existing in-flight
// owner, or installs the unique owner lease. Pending entries have a separate
// lifecycle from completed cache trim/TTL and consume one process-wide slot.
func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultAcquire, error) {
if c == nil || reqMsgID == 0 {
return rpcResultAcquire{}, ErrRPCResultFlightInvalid
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
if elem, ok := s.byKey[key]; ok {
entry := elem.Value.(*rpcResultCacheEntry)
if entry.expiresAt.After(now) {
return rpcResultAcquire{state: rpcResultAcquireCompleted, encoded: entry.encoded}, nil
}
s.removeElement(elem)
}
if flight, ok := s.pending[key]; ok {
return rpcResultAcquire{
state: rpcResultAcquirePending,
waiter: &rpcResultWaiter{flight: flight},
}, nil
}
if !c.flightLimit.reserve() {
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
flight := &rpcResultFlight{done: make(chan struct{})}
if s.pending == nil {
s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
}
s.pending[key] = flight
return rpcResultAcquire{
state: rpcResultAcquireOwner,
owner: &rpcResultOwnerLease{cache: c, key: key, flight: flight},
}, nil
}
// completeRPCResultFlightLocked publishes encoded to the current owner claim.
// The caller must hold s.mu and must publish the completed cache entry first.
func (c *rpcResultCache) completeRPCResultFlightLocked(
s *rpcResultCacheShard,
key rpcResultCacheKey,
encoded *encodedOutboundMessage,
) {
if c == nil || s == nil || encoded == nil {
return
}
flight, ok := s.pending[key]
if !ok {
return
}
delete(s.pending, key)
flight.encoded = encoded
flight.ok = true
c.flightLimit.release()
close(flight.done)
}

396
internal/mtprotoedge/rpc_flight_test.go Normal file
View file

@ -0,0 +1,396 @@
package mtprotoedge
import (
"context"
"errors"
"sync"
"testing"
"time"
)
func rpcFlightTestAuthID(seed byte) [8]byte {
return [8]byte{seed, seed + 1, seed + 2, seed + 3}
}
func TestRPCResultFlightConcurrentAcquireHasUniqueOwner(t *testing.T) {
const callers = 64
cache := newRPCResultCacheWithFlightLimit(time.Now, callers)
authKeyID := rpcFlightTestAuthID(1)
start := make(chan struct{})
results := make(chan rpcResultAcquire, callers)
errs := make(chan error, callers)
var wg sync.WaitGroup
wg.Add(callers)
for i := 0; i < callers; i++ {
go func() {
defer wg.Done()
<-start
claim, err := cache.Acquire(authKeyID, 10, 100)
if err != nil {
errs <- err
return
}
results <- claim
}()
}
close(start)
wg.Wait()
close(results)
close(errs)
for err := range errs {
t.Fatalf("Acquire: %v", err)
}
owners := 0
waiters := 0
var owner *rpcResultOwnerLease
for claim := range results {
switch claim.state {
case rpcResultAcquireOwner:
owners++
owner = claim.owner
case rpcResultAcquirePending:
waiters++
if claim.waiter == nil {
t.Fatal("pending claim has nil waiter")
}
default:
t.Fatalf("unexpected claim state %d", claim.state)
}
}
if owners != 1 || waiters != callers-1 {
t.Fatalf("claims = owners:%d waiters:%d, want 1/%d", owners, waiters, callers-1)
}
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("pending count = %d, want 1", got)
}
if owner == nil || !owner.Abort() {
t.Fatal("unique owner did not abort its claim")
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("pending count after abort = %d, want 0", got)
}
}
func TestRPCResultFlightPutPublishesAndWakesAllWaiters(t *testing.T) {
const waiters = 24
cache := newRPCResultCacheWithFlightLimit(time.Now, 32)
authKeyID := rpcFlightTestAuthID(10)
owner, err := cache.Acquire(authKeyID, 20, 200)
if err != nil || owner.state != rpcResultAcquireOwner || owner.owner == nil {
t.Fatalf("owner Acquire = state:%d err:%v", owner.state, err)
}
waiterClaims := make([]*rpcResultWaiter, 0, waiters)
for i := 0; i < waiters; i++ {
claim, acquireErr := cache.Acquire(authKeyID, 20, 200)
if acquireErr != nil || claim.state != rpcResultAcquirePending || claim.waiter == nil {
t.Fatalf("waiter %d Acquire = state:%d err:%v", i, claim.state, acquireErr)
}
waiterClaims = append(waiterClaims, claim.waiter)
}
type waiterResult struct {
encoded *encodedOutboundMessage
cached *encodedOutboundMessage
ok bool
err error
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
results := make(chan waiterResult, waiters)
for _, waiter := range waiterClaims {
go func(w *rpcResultWaiter) {
encoded, ok, waitErr := w.Wait(ctx)
cached, _ := cache.Get(authKeyID, 20, 200)
results <- waiterResult{encoded: encoded, cached: cached, ok: ok, err: waitErr}
}(waiter)
}
want := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, typeID: 42, reqMsgID: 200}
cache.Put(authKeyID, 20, 200, want)
for i := 0; i < waiters; i++ {
got := <-results
if got.err != nil || !got.ok {
t.Fatalf("waiter %d result = ok:%v err:%v", i, got.ok, got.err)
}
if got.encoded != want || got.cached != want {
t.Fatalf("waiter %d observed direct/cache pointers %p/%p, want %p", i, got.encoded, got.cached, want)
}
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("pending count after Put = %d, want 0", got)
}
completed, err := cache.Acquire(authKeyID, 20, 200)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want {
t.Fatalf("completed Acquire = state:%d encoded:%p err:%v", completed.state, completed.encoded, err)
}
if owner.owner.Abort() {
t.Fatal("completed owner's stale lease aborted a resolved claim")
}
}
func TestRPCResultFlightAbortWakesAndAllowsReclaim(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
authKeyID := rpcFlightTestAuthID(20)
first, err := cache.Acquire(authKeyID, 30, 300)
if err != nil || first.state != rpcResultAcquireOwner {
t.Fatalf("first Acquire = state:%d err:%v", first.state, err)
}
waiting, err := cache.Acquire(authKeyID, 30, 300)
if err != nil || waiting.state != rpcResultAcquirePending {
t.Fatalf("waiting Acquire = state:%d err:%v", waiting.state, err)
}
if !first.owner.Abort() {
t.Fatal("first Abort lost")
}
if encoded, ok, waitErr := waiting.waiter.Wait(context.Background()); waitErr != nil || ok || encoded != nil {
t.Fatalf("aborted Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr)
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("pending count after Abort = %d, want 0", got)
}
second, err := cache.Acquire(authKeyID, 30, 300)
if err != nil || second.state != rpcResultAcquireOwner || second.owner == nil {
t.Fatalf("reclaim = state:%d err:%v", second.state, err)
}
if first.owner.Abort() {
t.Fatal("stale first lease aborted the replacement owner")
}
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("pending count after reclaim = %d, want 1", got)
}
if !second.owner.Abort() || cache.flightLimit.snapshot() != 0 {
t.Fatal("replacement owner did not release its claim")
}
}
func TestRPCResultFlightCompletedCachePressureDoesNotEvictPending(t *testing.T) {
now := time.Unix(1_000, 0)
cache := newRPCResultCacheWithFlightLimit(func() time.Time { return now }, 4)
authKeyID := rpcFlightTestAuthID(30)
pending, err := cache.Acquire(authKeyID, 40, 400)
if err != nil || pending.state != rpcResultAcquireOwner {
t.Fatalf("pending Acquire = state:%d err:%v", pending.state, err)
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: 40, reqMsgID: 400}
shard := cache.shard(key)
shard.mu.Lock()
shard.maxEntries = 2
shard.maxBytes = 2
shard.mu.Unlock()
for i := int64(0); i < 16; i++ {
cache.Put(authKeyID, 40, 500+i, &encodedOutboundMessage{body: []byte{byte(i)}})
}
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("completed trim changed pending count to %d", got)
}
joined, err := cache.Acquire(authKeyID, 40, 400)
if err != nil || joined.state != rpcResultAcquirePending {
t.Fatalf("Acquire after completed trim = state:%d err:%v", joined.state, err)
}
// Expire the independent completed cache and prove the pending owner remains.
now = now.Add(rpcResultCacheTTL + time.Second)
_, _ = cache.Get(authKeyID, 40, 515)
joinedAfterTTL, err := cache.Acquire(authKeyID, 40, 400)
if err != nil || joinedAfterTTL.state != rpcResultAcquirePending {
t.Fatalf("Acquire after completed TTL = state:%d err:%v", joinedAfterTTL.state, err)
}
if !pending.owner.Abort() || cache.flightLimit.snapshot() != 0 {
t.Fatal("pending claim did not survive pressure through explicit Abort")
}
}
func TestRPCResultFlightCapacityAndCountReturn(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
authKeyID := rpcFlightTestAuthID(40)
first, err := cache.Acquire(authKeyID, 50, 501)
if err != nil || first.state != rpcResultAcquireOwner {
t.Fatalf("first Acquire = state:%d err:%v", first.state, err)
}
second, err := cache.Acquire(authKeyID, 50, 502)
if err != nil || second.state != rpcResultAcquireOwner {
t.Fatalf("second Acquire = state:%d err:%v", second.state, err)
}
joined, err := cache.Acquire(authKeyID, 50, 501)
if err != nil || joined.state != rpcResultAcquirePending {
t.Fatalf("join at capacity = state:%d err:%v", joined.state, err)
}
if _, err := cache.Acquire(authKeyID, 50, 503); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("over-capacity Acquire err = %v, want %v", err, ErrRPCResultFlightCapacity)
}
if got := cache.flightLimit.snapshot(); got != 2 {
t.Fatalf("pending count at capacity = %d, want 2", got)
}
want := &encodedOutboundMessage{body: []byte{9}, reqMsgID: 501}
cache.Put(authKeyID, 50, 501, want)
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("pending count after Put = %d, want 1", got)
}
third, err := cache.Acquire(authKeyID, 50, 503)
if err != nil || third.state != rpcResultAcquireOwner {
t.Fatalf("Acquire after returned slot = state:%d err:%v", third.state, err)
}
completed, err := cache.Acquire(authKeyID, 50, 501)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want {
t.Fatalf("completed Acquire at capacity = state:%d err:%v", completed.state, err)
}
if encoded, ok, waitErr := joined.waiter.Wait(context.Background()); waitErr != nil || !ok || encoded != want {
t.Fatalf("joined Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr)
}
if !second.owner.Abort() || !third.owner.Abort() {
t.Fatal("owners failed to return remaining capacity")
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("final pending count = %d, want 0", got)
}
}
func TestRPCResultFlightOversizedPutStillResolvesWaiters(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
authKeyID := rpcFlightTestAuthID(50)
owner, err := cache.Acquire(authKeyID, 60, 600)
if err != nil || owner.state != rpcResultAcquireOwner {
t.Fatalf("owner Acquire = state:%d err:%v", owner.state, err)
}
joined, err := cache.Acquire(authKeyID, 60, 600)
if err != nil || joined.state != rpcResultAcquirePending {
t.Fatalf("joined Acquire = state:%d err:%v", joined.state, err)
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: 60, reqMsgID: 600}
shard := cache.shard(key)
shard.mu.Lock()
shard.maxBytes = 1
shard.mu.Unlock()
want := &encodedOutboundMessage{body: []byte{1, 2}, reqMsgID: 600}
cache.Put(authKeyID, 60, 600, want)
if encoded, ok, waitErr := joined.waiter.Wait(context.Background()); waitErr != nil || !ok || encoded != want {
t.Fatalf("oversized Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr)
}
if _, ok := cache.Get(authKeyID, 60, 600); ok {
t.Fatal("oversized result changed completed-cache compatibility")
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("oversized Put leaked pending count %d", got)
}
if owner.owner.Abort() {
t.Fatal("oversized Put left its old owner abortable")
}
retry, err := cache.Acquire(authKeyID, 60, 600)
if err != nil || retry.state != rpcResultAcquireOwner {
t.Fatalf("retry after uncacheable completion = state:%d err:%v", retry.state, err)
}
if !retry.owner.Abort() {
t.Fatal("retry owner failed to abort")
}
}
func TestRPCResultFlightWaitContextDoesNotReleaseOwner(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
authKeyID := rpcFlightTestAuthID(60)
owner, err := cache.Acquire(authKeyID, 70, 700)
if err != nil {
t.Fatalf("owner Acquire: %v", err)
}
joined, err := cache.Acquire(authKeyID, 70, 700)
if err != nil {
t.Fatalf("joined Acquire: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
if encoded, ok, waitErr := joined.waiter.Wait(ctx); !errors.Is(waitErr, context.Canceled) || ok || encoded != nil {
t.Fatalf("canceled Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr)
}
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("waiter cancellation released owner count to %d", got)
}
if !owner.owner.Abort() || cache.flightLimit.snapshot() != 0 {
t.Fatal("owner did not retain and release claim after waiter cancellation")
}
}
func TestRPCResultFlightConcurrentCapacityReturnsAllSlots(t *testing.T) {
const (
limit = 32
callers = 512
)
cache := newRPCResultCacheWithFlightLimit(time.Now, limit)
authKeyID := rpcFlightTestAuthID(70)
start := make(chan struct{})
errs := make(chan error, callers)
var wg sync.WaitGroup
wg.Add(callers)
for i := 0; i < callers; i++ {
i := i
go func() {
defer wg.Done()
<-start
claim, err := cache.Acquire(authKeyID, 80+int64(i%4), 1_000+int64(i))
if errors.Is(err, ErrRPCResultFlightCapacity) {
return
}
if err != nil {
errs <- err
return
}
if claim.state != rpcResultAcquireOwner || claim.owner == nil {
errs <- errors.New("unique-key claim did not become owner")
return
}
if i%2 == 0 {
cache.Put(authKeyID, 80+int64(i%4), 1_000+int64(i), &encodedOutboundMessage{body: []byte{1}})
} else if !claim.owner.Abort() {
errs <- errors.New("owner Abort lost")
}
}()
}
close(start)
wg.Wait()
close(errs)
for err := range errs {
t.Errorf("concurrent claim: %v", err)
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("concurrent completion leaked %d pending slots", got)
}
}
func TestQueuedRPCConnectionCloseAbortsOwnerClaim(t *testing.T) {
s := New(Options{})
c := newInboundTestConn(s.rpcScheduler, 1, 4, time.Second)
claim, err := s.rpcResults.Acquire([8]byte{9}, 90, 900)
if err != nil || claim.state != rpcResultAcquireOwner || claim.owner == nil {
t.Fatalf("Acquire owner = state:%d err:%v", claim.state, err)
}
reservation, err := c.reserveInboundRPC(context.Background(), "test.queuedFlight", 4)
if err != nil {
t.Fatalf("reserve queued RPC: %v", err)
}
task := s.newInboundRPCTask(c, 900, "test.queuedFlight", []byte{1, 2, 3, 4}, claim.owner)
if err := reservation.commit(task); err != nil {
t.Fatalf("commit queued RPC: %v", err)
}
// The scheduler is intentionally not started, so close must drain the queued
// task and invoke its independent flight-release callback.
c.closeInboundRPCScheduler()
if got := s.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("connection close leaked %d queued owner claims", got)
}
if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("connection close leaked scheduler budget %d/%d", tasks, bytes)
}
retry, err := s.rpcResults.Acquire([8]byte{9}, 90, 900)
if err != nil || retry.state != rpcResultAcquireOwner {
t.Fatalf("reclaim after queued close = state:%d err:%v", retry.state, err)
}
if !retry.owner.Abort() {
t.Fatal("replacement owner did not abort")
}
}

60
internal/mtprotoedge/rpc_result_cache.go
View file

@ -30,11 +30,14 @@ type rpcResultCacheEntry struct {
expiresAt time.Time
}
// rpcResultCache 缓存已回发的 rpc_result（按 auth_key+session+req_msg_id），用于
// 跨连接重放重复请求。encodedOutboundMessage 构造后不可变（push fan-out 与 pending
// resend 均依赖该契约），因此 Get/Put 直接共享指针，不做防御性拷贝。
// rpcResultCache 缓存已有交付证明的 rpc_result（按 auth_key+session+req_msg_id），
// 用于跨连接重放重复请求。Put 的调用方必须先证明结果已物理写出，或原 logical Conn
// 已不可逆 fenced；绝不能发布“Conn 仍 current/open 但结果尚未上 wire”的完成态。
// encodedOutboundMessage 构造后不可变（push fan-out 与 pending resend 均依赖该契约），
// 因此 Get/Put 直接共享指针，不做防御性拷贝。
type rpcResultCache struct {
shards [rpcResultCacheShards]rpcResultCacheShard
shards [rpcResultCacheShards]rpcResultCacheShard
flightLimit rpcResultFlightLimit
}
type rpcResultCacheShard struct {
@ -46,13 +49,25 @@ type rpcResultCacheShard struct {
bytes int
order *list.List
byKey map[rpcResultCacheKey]*list.Element
// pending is deliberately independent from the completed-result order/byKey
// cache. In-flight owners and waiters must not disappear when completed
// results expire or are trimmed under entry/byte pressure.
pending map[rpcResultCacheKey]*rpcResultFlight
}
func newRPCResultCache(now func() time.Time) *rpcResultCache {
return newRPCResultCacheWithFlightLimit(now, rpcResultFlightDefaultMaxPending)
}
func newRPCResultCacheWithFlightLimit(now func() time.Time, maxPending int) *rpcResultCache {
if now == nil {
now = time.Now
}
if maxPending <= 0 {
maxPending = rpcResultFlightDefaultMaxPending
}
c := &rpcResultCache{}
c.flightLimit.max = int64(maxPending)
for i := range c.shards {
s := &c.shards[i]
s.now = now
@ -61,6 +76,7 @@ func newRPCResultCache(now func() time.Time) *rpcResultCache {
s.maxBytes = rpcResultCacheMaxBytes / rpcResultCacheShards
s.order = list.New()
s.byKey = make(map[rpcResultCacheKey]*list.Element)
s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
}
return c
}
@ -102,28 +118,32 @@ func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encod
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
size := len(encoded.body)
if s.maxBytes > 0 && size > s.maxBytes {
return
}
cacheable := s.maxBytes <= 0 || size <= s.maxBytes
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
s.expireLocked(now)
if elem, ok := s.byKey[key]; ok {
s.removeElement(elem)
if cacheable {
s.expireLocked(now)
if elem, ok := s.byKey[key]; ok {
s.removeElement(elem)
}
entry := &rpcResultCacheEntry{
key: key,
encoded: encoded,
size: size,
expiresAt: now.Add(s.ttl),
}
elem := s.order.PushBack(entry)
s.byKey[key] = elem
s.bytes += size
s.trimLocked()
}
entry := &rpcResultCacheEntry{
key: key,
encoded: encoded,
size: size,
expiresAt: now.Add(s.ttl),
}
elem := s.order.PushBack(entry)
s.byKey[key] = elem
s.bytes += size
s.trimLocked()
// Resolve the independent in-flight entry only after the completed cache has
// been published. Waiters awakened by this close can therefore immediately
// observe either the shared encoded result or the completed Get entry.
c.completeRPCResultFlightLocked(s, key, encoded)
}
func (s *rpcResultCacheShard) expireLocked(now time.Time) {

199
internal/mtprotoedge/rpc_result_delivery_test.go Normal file
View file

@ -0,0 +1,199 @@
package mtprotoedge
import (
"context"
"errors"
"io"
"sync"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
type blockingCloseRPCResultTransport struct {
started chan struct{}
release chan struct{}
once sync.Once
}
func newBlockingCloseRPCResultTransport() *blockingCloseRPCResultTransport {
return &blockingCloseRPCResultTransport{started: make(chan struct{}), release: make(chan struct{})}
}
func (*blockingCloseRPCResultTransport) Send(context.Context, *bin.Buffer) error { return nil }
func (*blockingCloseRPCResultTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
func (t *blockingCloseRPCResultTransport) Close() error {
t.once.Do(func() { close(t.started) })
<-t.release
return nil
}
func TestRPCResultCachePublishesOnlyAfterPhysicalWrite(t *testing.T) {
tr := newGatedRequiredControlTransport(nil)
s := New(Options{WriteTimeout: time.Second})
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74001, 1)
t.Cleanup(c.ForceClose)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
owner, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || owner.state != rpcResultAcquireOwner {
t.Fatalf("initial flight owner = %+v err=%v", owner, err)
}
done := make(chan error, 1)
go func() {
done <- s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2})
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("rpc_result did not reach physical writer")
}
if _, ok := s.rpcResults.Get(key.ID, c.sessionID, reqMsgID); ok {
t.Fatal("rpc_result became completed before physical write")
}
pending, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || pending.state != rpcResultAcquirePending || pending.waiter == nil {
t.Fatalf("flight while write blocked = %+v err=%v", pending, err)
}
tr.unblock()
select {
case err := <-done:
if err != nil {
t.Fatalf("sendResult after write: %v", err)
}
case <-time.After(time.Second):
t.Fatal("sendResult did not finish")
}
if encoded, ok, err := pending.waiter.Wait(context.Background()); err != nil || !ok || encoded == nil {
t.Fatalf("pending waiter after write = encoded:%p ok:%v err:%v", encoded, ok, err)
}
completed, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {
t.Fatalf("flight after write = %+v err=%v", completed, err)
}
}
func TestRPCResultPrewriteFailureFencesConnBeforeCachePublication(t *testing.T) {
tr := &collectingSessionTransport{}
s := New(Options{WriteTimeout: 20 * time.Millisecond})
// No rpc_result can reserve its conservative 3x wire scratch from one byte.
// The actor therefore fails before touching the socket, which used to leave a
// live Conn with a prematurely completed cache entry.
s.outboundScratchPool = newOutboundScratchPool(1)
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74002, 1)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
owner, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || owner.state != rpcResultAcquireOwner {
t.Fatalf("initial flight owner = %+v err=%v", owner, err)
}
err = s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2})
if err == nil || (!errors.Is(err, context.DeadlineExceeded) &&
!errors.Is(err, ErrConnClosed) && !errors.Is(err, ErrOutboundTrackedBudget)) {
t.Fatalf("prewrite sendResult error = %v", err)
}
closeDeadline := time.Now().Add(time.Second)
for !tr.closed.Load() && time.Now().Before(closeDeadline) {
time.Sleep(time.Millisecond)
}
if !c.terminal.Load() || !tr.closed.Load() || c.isPhysicalTransportCurrentOpen() {
t.Fatalf("failed delivery did not fence Conn: terminal=%v closed=%v current_open=%v", c.terminal.Load(), tr.closed.Load(), c.isPhysicalTransportCurrentOpen())
}
completed, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {
t.Fatalf("terminal replay cache = %+v err=%v", completed, acquireErr)
}
if got := s.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("failed delivery leaked flight slot: %d", got)
}
c.Close()
}
func TestRPCResultFailureAfterIntentionalTerminalDoesNotCloseTransferLease(t *testing.T) {
tr := &collectingSessionTransport{}
s := New(Options{WriteTimeout: time.Second})
s.outboundTrackedBudget = newOutboundTrackedBudget(1)
key := newTestAuthKey(t)
oldConn := s.newConn(tr, key, 74003, 1)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
owner, err := s.rpcResults.Acquire(key.ID, oldConn.sessionID, reqMsgID)
if err != nil || owner.state != rpcResultAcquireOwner {
t.Fatalf("initial flight owner = %+v err=%v", owner, err)
}
// Session replacement publishes terminal before transferring the physical
// lease. A late old-generation result sees ErrConnClosed at producer admission;
// it may publish cache-only, but must not upgrade that intentional fence into
// a physical close that makes Transfer fail.
oldConn.beginTerminalShutdown()
err = s.sendResult(context.Background(), oldConn, reqMsgID, &tg.Config{ThisDC: 2})
if !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("late result error = %v, want ErrOutboundTrackedBudget", err)
}
if tr.closed.Load() {
t.Fatal("late result closed intentionally transferable transport")
}
if !oldConn.waitOutboundShutdownUntil(time.Second) {
t.Fatal("old outbound actor did not stop")
}
nextLease, ok := oldConn.transferTransportOwnership()
if !ok || nextLease == nil {
t.Fatal("late result prevented physical transfer")
}
newConn := s.newConnWithLease(nextLease, key, 74004, 1)
oldConn.ForceClose()
if tr.closed.Load() || !newConn.isPhysicalTransportCurrentOpen() {
t.Fatalf("stale close after transfer: raw_closed=%v current_open=%v", tr.closed.Load(), newConn.isPhysicalTransportCurrentOpen())
}
completed, acquireErr := s.rpcResults.Acquire(key.ID, oldConn.sessionID, reqMsgID)
if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {
t.Fatalf("late result cache = %+v err=%v", completed, acquireErr)
}
newConn.ForceClose()
}
func TestRPCResultPublishesBeforePathologicalPhysicalCloseReturns(t *testing.T) {
tr := newBlockingCloseRPCResultTransport()
s := New(Options{WriteTimeout: time.Second})
s.outboundTrackedBudget = newOutboundTrackedBudget(1)
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74005, 1)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
owner, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || owner.state != rpcResultAcquireOwner {
t.Fatalf("initial flight owner = %+v err=%v", owner, err)
}
done := make(chan error, 1)
go func() { done <- s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2}) }()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("physical Close did not start")
}
select {
case err := <-done:
if !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("sendResult error = %v, want budget error", err)
}
case <-time.After(time.Second):
t.Fatal("pathological physical Close blocked result publication")
}
completed, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {
t.Fatalf("completed result before raw Close return = %+v err=%v", completed, acquireErr)
}
close(tr.release)
if c.transportLease != nil {
if err := c.transportLease.owner.waitClosed(); err != nil {
t.Fatalf("physical Close: %v", err)
}
}
c.Close()
}

14
internal/mtprotoedge/rpc_test.go
View file

@ -257,7 +257,8 @@ func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) {
func TestCanceledRPCErrorIsNotCachedAcrossReconnect(t *testing.T) {
const dc = 2
handler := &canceledInternalRPC{
firstDone: make(chan struct{}),
firstStarted: make(chan struct{}),
firstDone: make(chan struct{}),
}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
@ -266,6 +267,11 @@ func TestCanceledRPCErrorIsNotCachedAcrossReconnect(t *testing.T) {
reqMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, reqMsgID, &tg.HelpGetConfigRequest{})
select {
case <-handler.firstStarted:
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for first rpc to start after session barrier")
}
_ = conn.Close()
select {
case <-handler.firstDone:
@ -355,12 +361,14 @@ func (h *runningDeadlineRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _
func (h *runningDeadlineRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
type canceledInternalRPC struct {
calls atomic.Int32
firstDone chan struct{}
calls atomic.Int32
firstStarted chan struct{}
firstDone chan struct{}
}
func (h *canceledInternalRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
if h.calls.Add(1) == 1 {
close(h.firstStarted)
<-ctx.Done()
close(h.firstDone)
return nil, tgerr.New(500, "INTERNAL_SERVER_ERROR")

67
internal/mtprotoedge/server.go
View file

@ -172,7 +172,7 @@ func (o *Options) setDefaults() {
o.RPCGlobalWorkers = 256
}
if o.RPCGlobalMaxTasks <= 0 {
o.RPCGlobalMaxTasks = 8192
o.RPCGlobalMaxTasks = rpcResultFlightDefaultMaxPending
}
if o.RPCGlobalMaxBytes <= 0 {
o.RPCGlobalMaxBytes = 512 << 20
@ -295,7 +295,7 @@ func New(opts Options) *Server {
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
rpcResults: newRPCResultCache(opts.Clock.Now),
rpcResults: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks),
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
}
}
@ -307,8 +307,33 @@ func (s *Server) Conns() *SessionManager {
// newConn 基于一次解密结果创建一个可发送的连接对象。
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
if lease, ok := tc.(*physicalTransportLease); ok {
return s.newConnWithLease(lease, key, sessionID, salt)
}
if tc != nil {
_, lease := newPhysicalTransportOwner(tc)
return s.newConnWithLease(lease, key, sessionID, salt)
}
// Preserve the nil transport used by construction-only tests.
return s.buildConn(nil, nil, key, sessionID, salt)
}
// newConnWithLease attaches a logical Conn to an explicitly owned physical
// transport generation. Production session replacement must Transfer the old
// lease first; it must never wrap the same raw transport in a second owner.
func (s *Server) newConnWithLease(lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn {
if lease == nil {
panic("mtprotoedge: nil physical transport lease")
}
c := s.buildConn(lease, lease, key, sessionID, salt)
lease.bindLogicalConn(c)
return c
}
func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn {
c := &Conn{
transport: tc,
transportLease: lease,
writer: tc,
cipher: s.cipher,
msgID: proto.NewMessageIDGen(s.clock.Now),
@ -558,7 +583,8 @@ func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, err
// - auth_key_id 未注册：回 AuthKeyNotFound，促使客户端重新握手。
//
// 连接建立 session 后注册到 SessionManager，结束时注销。
func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) {
func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error) {
transportOwner, conn := newPhysicalTransportOwner(raw)
s.metrics.ConnOpened()
s.log.Debug("Connection accepted")
@ -568,9 +594,13 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
// this stack has stopped using b/plain; transport.Close may have raced us earlier and must
// not return that memory budget prematurely.
releaseInboundFrameOwnership(conn)
// 先同步关闭物理 socket，解除可能阻塞在 writer.Send 的 outbound actor；
// 再停止 logical Conn，避免 Close 等 actor 时反过来等到 write deadline。
_ = conn.Close()
// Publish the terminal/RPC-cancel gates before index removal or lifecycle
// observers. Physical close then releases a writer already inside Send; the
// final Close only waits for the now-fenced actors to converge.
if current != nil {
current.beginTerminalShutdown()
}
_ = transportOwner.CloseAny()
if current != nil {
s.conns.Unregister(current)
current.Close()
@ -584,7 +614,7 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
defer cancel()
go func() {
<-ctx.Done()
_ = conn.Close()
_ = transportOwner.CloseAny()
}()
cs := newConnState()
@ -601,7 +631,7 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
// 建立 session 前（current==nil，握手 + 首个加密消息之前）用较短的 handshakeTimeout
// 快速回收静默的半开 / 异常连接；建立 session 后用 readTimeout（客户端有 ping 心跳）。
timeout := s.readTimeout
if current == nil {
if current == nil || !current.isActive() {
timeout = s.handshakeTimeout
}
if err := s.recv(ctx, conn, &b, timeout); err != nil {
@ -618,6 +648,12 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
}
if authKeyID == emptyAuthKeyID {
// A physical socket may perform key exchange only before it owns an
// encrypted logical session. Mixing the direct exchange writer with an
// active outbound actor would bypass generation/write serialization.
if current != nil {
return errors.New("unencrypted exchange on established encrypted connection")
}
releaseHandshake, admitted := s.admission.tryAcquireHandshake()
if !admitted {
if err := s.sendProtoError(ctx, conn, codec.CodeTransportFlood); err != nil {
@ -647,7 +683,9 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
// 已建立连接复用缓存密钥走快路径（fetchedKey=nil）：避开每帧回查 AuthKeyStore——
// 这是 mtprotoedge 层最热的库访问点。密钥材料创建后不可变；销毁(destroy_auth_key)/
// 撤销由 SessionManager 主动 Close 连接保证失效，不依赖此被动回查。仅 destroy_auth_key
// 的发起连接置 keyDestroyed，使其下一帧回落到 Get→AuthKeyNotFound，维持原契约。
// 的发起连接置 keyDestroyed，使其下一帧回落到 Get→AuthKeyNotFound。尚未进入
// SessionManager 的 bad-salt provisional 会在 handleEncrypted 建立 activation claim
// 后精确复查一次，既把撤销与激活线性化，也不把 salt storm 放大成 PG 写风暴。
var fetchedKey *store.AuthKeyData
if current == nil || current.authKeyID != authKeyID || current.keyDestroyed.Load() {
d, found, err := s.authKeys.Get(ctx, authKeyID)
@ -655,7 +693,11 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
return fmt.Errorf("lookup auth key: %w", err)
}
if !found {
if err := s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound); err != nil {
writer := transport.Conn(conn)
if current != nil {
writer = current.transport
}
if err := s.sendProtoError(ctx, writer, codec.CodeAuthKeyNotFound); err != nil {
return err
}
// -404 对 TDesktop 是 terminal key failure；继续保留 socket 只会允许
@ -666,6 +708,11 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error)
}
current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b, &plain)
if errors.Is(err, errActivationAuthKeyRejected) {
// handleEncrypted writes -404 while its activation claim still owns the
// physical writer, then its deferred abort removes/closes the claim.
return nil
}
if err != nil {
return err
}

253
internal/mtprotoedge/session_activation_test.go Normal file
View file

@ -0,0 +1,253 @@
package mtprotoedge
import (
"context"
"errors"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
)
func newSessionActivationTestConn(t *testing.T, authKeyID [8]byte, sessionID int64) *Conn {
t.Helper()
c := &Conn{
authKeyID: authKeyID,
sessionID: sessionID,
metrics: NopMetrics{},
}
c.startOutbound()
t.Cleanup(c.Close)
return c
}
func TestSessionActivationGatesReplacementBeforePublishing(t *testing.T) {
manager := NewSessionManager(zaptest.NewLogger(t))
key := [8]byte{1, 2, 3, 4}
oldConn := newSessionActivationTestConn(t, key, 7001)
newConn := newSessionActivationTestConn(t, key, 7001)
if err := manager.Register(oldConn); err != nil {
t.Fatalf("register initial: %v", err)
}
if !oldConn.isActive() {
t.Fatal("initial connection was not activated")
}
if newConn.lifecycleState() != connLifecycleProvisional {
t.Fatal("provisional replacement was active before registration")
}
if err := manager.Register(newConn); err != nil {
t.Fatalf("register replacement: %v", err)
}
if !newConn.isActive() {
t.Fatal("replacement was not activated")
}
if oldConn.lifecycleState() != connLifecycleRetired || !oldConn.terminal.Load() {
t.Fatalf("old connection gates = lifecycle:%v terminal:%v", oldConn.lifecycleState(), oldConn.terminal.Load())
}
if err := oldConn.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("stale old connection send error = %v, want ErrConnClosed", err)
}
manager.mu.RLock()
got := manager.bySession[sessionKey{authKeyID: key, sessionID: 7001}]
manager.mu.RUnlock()
if got != newConn {
t.Fatalf("published session = %p, want replacement %p", got, newConn)
}
}
func TestSessionActivationClaimPreemptionCannotReversePublish(t *testing.T) {
manager := NewSessionManager(zaptest.NewLogger(t))
key := [8]byte{9, 8, 7, 6}
first := newSessionActivationTestConn(t, key, 8001)
second := newSessionActivationTestConn(t, key, 8001)
if err := manager.BeginActivation(first); err != nil {
t.Fatalf("begin first activation: %v", err)
}
if first.lifecycleState() != connLifecycleClaiming {
t.Fatalf("first lifecycle = %v, want claiming", first.lifecycleState())
}
if got := manager.Online(); got != 0 {
t.Fatalf("online during claim = %d, want 0", got)
}
if err := manager.BeginActivation(second); err != nil {
t.Fatalf("begin superseding activation: %v", err)
}
if first.lifecycleState() != connLifecycleRetired || !first.terminal.Load() {
t.Fatalf("superseded first lifecycle=%v terminal=%v", first.lifecycleState(), first.terminal.Load())
}
if err := manager.PublishActivation(first); !errors.Is(err, ErrSessionActivationSuperseded) {
t.Fatalf("stale publish error = %v, want superseded", err)
}
if err := manager.Register(first); !errors.Is(err, ErrSessionActivationSuperseded) {
t.Fatalf("stale register error = %v, want superseded", err)
}
if err := manager.PublishActivation(second); err != nil {
t.Fatalf("publish second activation: %v", err)
}
manager.mu.RLock()
got := manager.bySession[sessionKey{authKeyID: key, sessionID: 8001}]
claim := manager.claims[sessionKey{authKeyID: key, sessionID: 8001}]
manager.mu.RUnlock()
if got != second || claim != nil || !second.isActive() {
t.Fatalf("activation owner=%p claim=%p second_active=%v", got, claim, second.isActive())
}
}
func TestBeginActivationWaitsForPreviousPhysicalWriterFence(t *testing.T) {
manager := NewSessionManager(zaptest.NewLogger(t))
key := [8]byte{3, 3, 3, 3}
releaseClose := make(chan struct{})
transport := newSlowCloseTransport(0, releaseClose)
oldConn := &Conn{
authKeyID: key,
sessionID: 8501,
metrics: NopMetrics{},
transport: transport,
writer: transport,
}
oldConn.startOutbound()
t.Cleanup(oldConn.Close)
if err := manager.Register(oldConn); err != nil {
t.Fatalf("register old: %v", err)
}
newConn := newSessionActivationTestConn(t, key, oldConn.sessionID)
beginDone := make(chan error, 1)
go func() { beginDone <- manager.BeginActivation(newConn) }()
deadline := time.Now().Add(time.Second)
for transport.closes.Load() == 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if transport.closes.Load() == 0 {
t.Fatal("replacement did not start closing previous transport")
}
select {
case err := <-beginDone:
t.Fatalf("BeginActivation returned before old physical close: %v", err)
default:
}
close(releaseClose)
select {
case err := <-beginDone:
if err != nil {
t.Fatalf("BeginActivation after old close: %v", err)
}
case <-time.After(time.Second):
t.Fatal("BeginActivation did not converge after old close")
}
if err := manager.PublishActivation(newConn); err != nil {
t.Fatalf("publish replacement: %v", err)
}
}
func TestSessionActivationClaimIndexesCleanOnPublishAndAbort(t *testing.T) {
manager := NewSessionManager(zaptest.NewLogger(t))
key := [8]byte{4, 4, 4, 4}
published := newSessionActivationTestConn(t, key, 9001)
if err := manager.BeginActivation(published); err != nil {
t.Fatalf("begin publish claim: %v", err)
}
manager.mu.RLock()
if manager.claims[connSessionKey(published)] != published || manager.claimsByAuth[key][published.sessionID] != published {
manager.mu.RUnlock()
t.Fatal("claim indexes missing after BeginActivation")
}
manager.mu.RUnlock()
if err := manager.PublishActivation(published); err != nil {
t.Fatalf("publish claim: %v", err)
}
manager.mu.RLock()
globalClaims, authClaims := len(manager.claims), len(manager.claimsByAuth[key])
manager.mu.RUnlock()
if globalClaims != 0 || authClaims != 0 {
t.Fatalf("claim indexes after publish = %d/%d, want 0/0", globalClaims, authClaims)
}
aborted := newSessionActivationTestConn(t, key, 9002)
if err := manager.BeginActivation(aborted); err != nil {
t.Fatalf("begin abort claim: %v", err)
}
manager.AbortActivation(aborted)
manager.mu.RLock()
globalClaims, authClaims = len(manager.claims), len(manager.claimsByAuth[key])
manager.mu.RUnlock()
if globalClaims != 0 || authClaims != 0 || aborted.lifecycleState() != connLifecycleRetired {
t.Fatalf("claim indexes/lifecycle after abort = %d/%d/%v", globalClaims, authClaims, aborted.lifecycleState())
}
}
func TestRawAuthKeyCloseExactConnDoesNotExcludeSameSessionReplacement(t *testing.T) {
manager := NewSessionManager(zaptest.NewLogger(t))
key := [8]byte{4, 3, 2, 1}
const sessionID = 9050
// The destroy request may finish on a retired Conn after another physical
// connection has become the owner of the same logical session.
destroyer := newSessionActivationTestConn(t, key, sessionID)
destroyer.beginTerminalShutdown()
replacement := newSessionActivationTestConn(t, key, sessionID)
if err := manager.Register(replacement); err != nil {
t.Fatalf("register replacement: %v", err)
}
if got := manager.CloseSessionsForRawAuthKeyExceptConn(key, destroyer); got != 1 {
t.Fatalf("closed sessions = %d, want replacement only", got)
}
manager.mu.RLock()
active, claim := manager.bySession[sessionKey{authKeyID: key, sessionID: sessionID}], manager.claims[sessionKey{authKeyID: key, sessionID: sessionID}]
manager.mu.RUnlock()
if active != nil || claim != nil || !replacement.terminal.Load() {
t.Fatalf("same-session replacement escaped exact exclusion: active=%p claim=%p terminal=%v", active, claim, replacement.terminal.Load())
}
}
func TestDestroyFencesOutboundAndReservedRPCBeforeRemoval(t *testing.T) {
manager := NewSessionManager(zaptest.NewLogger(t))
key := [8]byte{5, 5, 5, 5}
c := newSessionActivationTestConn(t, key, 9101)
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c.startInboundRPCScheduler(scheduler, 1, 8, time.Second)
if err := manager.Register(c); err != nil {
t.Fatalf("register: %v", err)
}
reservation, err := c.reserveInboundRPC(context.Background(), "test.destroyFence", 8)
if err != nil {
t.Fatalf("reserve inbound RPC: %v", err)
}
destroyed := make(chan bool, 1)
go func() {
destroyed <- manager.DestroySessionForAuthKey(key, c.sessionID)
}()
select {
case <-c.rpcRootCtx.Done():
case <-time.After(time.Second):
t.Fatal("destroy did not synchronously cancel RPC admission")
}
if err := c.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("send after destroy fence = %v, want ErrConnClosed", err)
}
if err := reservation.commit(inboundRPC{}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("reserved RPC commit after destroy = %v, want ErrConnClosed", err)
}
select {
case ok := <-destroyed:
if !ok {
t.Fatal("destroy returned false")
}
case <-time.After(time.Second):
t.Fatal("destroy did not converge after reservation commit")
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("scheduler budget after destroy = %d/%d, want 0/0", tasks, bytes)
}
}

630
internal/mtprotoedge/session_barrier_integration_test.go Normal file
View file

@ -0,0 +1,630 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"errors"
"io"
"sync"
"sync/atomic"
"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"
"telesrv/internal/store"
)
func encryptedRPCFrameForBarrierTest(
t *testing.T,
key crypto.AuthKey,
salt, sessionID, msgID int64,
) (*bin.Buffer, store.AuthKeyData) {
return encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(t, key, salt, salt, sessionID, msgID, 1)
}
func encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t *testing.T,
key crypto.AuthKey,
frameSalt, authoritativeSalt, sessionID, msgID int64,
seqNo int32,
) (*bin.Buffer, store.AuthKeyData) {
t.Helper()
body := encodeClientMessageBodyForTest(t, &tg.HelpGetConfigRequest{})
var frame bin.Buffer
if err := crypto.NewClientCipher(rand.Reader).Encrypt(key, crypto.EncryptedMessageData{
Salt: frameSalt,
SessionID: sessionID,
MessageID: msgID,
SeqNo: seqNo,
MessageDataLen: int32(len(body)),
MessageDataWithPadding: body,
}, &frame); err != nil {
t.Fatalf("encrypt client RPC: %v", err)
}
return &frame, store.AuthKeyData{
ID: key.ID,
Value: [256]byte(key.Value),
ServerSalt: authoritativeSalt,
}
}
func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) {
handler := &admissionCountingRPC{}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
tr := &collectingSessionTransport{}
key := newTestAuthKey(t)
const (
serverSalt = int64(0x1020304050)
wrongSalt = serverSalt + 1
sessionID = int64(71001)
)
ids := proto.NewMessageIDGen(time.Now)
firstID := ids.New(proto.MessageFromClient)
firstWrong, stored := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, wrongSalt, serverSalt, sessionID, firstID, 1,
)
if err := s.authKeys.Save(context.Background(), stored); err != nil {
t.Fatalf("save auth key: %v", err)
}
cs := newConnState()
var plain bin.Buffer
firstConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, firstWrong, &plain)
if err != nil {
t.Fatalf("first bad salt: %v", err)
}
if firstConn == nil || firstConn.lifecycleState() != connLifecycleProvisional || firstConn.terminal.Load() {
t.Fatalf("first correction lifecycle conn=%p state=%v terminal=%v", firstConn, firstConn.lifecycleState(), firstConn.terminal.Load())
}
if cs.createdFloor != 0 || len(cs.seen) != 0 || handler.calls.Load() != 0 {
t.Fatalf("bad salt admitted state: floor=%d seen=%d calls=%d", cs.createdFloor, len(cs.seen), handler.calls.Load())
}
secondID := ids.New(proto.MessageFromClient)
secondWrong, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, wrongSalt, serverSalt, sessionID, secondID, 3,
)
secondConn, err := s.handleEncrypted(context.Background(), tr, cs, firstConn, nil, secondWrong, &plain)
if err != nil {
t.Fatalf("second bad salt: %v", err)
}
if secondConn != firstConn {
t.Fatalf("bad-salt retry replaced provisional Conn: first=%p second=%p", firstConn, secondConn)
}
if got := len(tr.snapshot()); got != 2 {
t.Fatalf("distinct bad msg ids produced %d corrections, want 2", got)
}
for i, wire := range tr.snapshot()[:2] {
data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire})
if decryptErr != nil {
t.Fatalf("decrypt correction %d: %v", i, decryptErr)
}
if data.Salt != serverSalt {
t.Fatalf("correction %d envelope salt = %#x, want %#x", i, data.Salt, serverSalt)
}
var bad mt.BadServerSalt
body := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)}
if decodeErr := bad.Decode(body); decodeErr != nil {
t.Fatalf("decode correction %d: %v", i, decodeErr)
}
if bad.NewServerSalt != serverSalt {
t.Fatalf("correction %d payload salt = %#x, want %#x", i, bad.NewServerSalt, serverSalt)
}
}
corrected, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, serverSalt, serverSalt, sessionID, firstID, 1,
)
activeConn, err := s.handleEncrypted(context.Background(), tr, cs, secondConn, nil, corrected, &plain)
if err != nil {
t.Fatalf("corrected retry: %v", err)
}
if activeConn != firstConn || !activeConn.isActive() {
t.Fatalf("corrected retry connection=%p first=%p active=%v", activeConn, firstConn, activeConn != nil && activeConn.isActive())
}
if cs.createdFloor != firstID {
t.Fatalf("corrected created floor = %d, want %d", cs.createdFloor, firstID)
}
waitForAtomicCalls(t, &handler.calls, 1)
flightDeadline := time.Now().Add(2 * time.Second)
for s.rpcResults.flightLimit.snapshot() != 0 && time.Now().Before(flightDeadline) {
time.Sleep(time.Millisecond)
}
if got := s.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("corrected retry leaked flight slots: %d", got)
}
activeConn.ForceClose()
}
func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) {
// The first handler remains pending until session replacement cancels it. Its
// owner Abort must observe the already-published terminal gate and must not
// close the physical lease that is about to transfer to the new session.
handler := &cancelThenRetryRPC{started: make(chan struct{})}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
tr := &collectingSessionTransport{}
key := newTestAuthKey(t)
const (
serverSalt = int64(0x66778899)
wrongSalt = serverSalt + 7
firstSID = int64(72001)
secondSID = int64(72002)
)
ids := proto.NewMessageIDGen(time.Now)
firstID := ids.New(proto.MessageFromClient)
firstFrame, stored := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, serverSalt, serverSalt, firstSID, firstID, 1,
)
if err := s.authKeys.Save(context.Background(), stored); err != nil {
t.Fatalf("save auth key: %v", err)
}
cs := newConnState()
var plain bin.Buffer
oldConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, firstFrame, &plain)
if err != nil {
t.Fatalf("activate first session: %v", err)
}
select {
case <-handler.started:
case <-time.After(2 * time.Second):
t.Fatal("first session RPC did not start")
}
secondID := ids.New(proto.MessageFromClient)
wrongFrame, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, wrongSalt, serverSalt, secondSID, secondID, 1,
)
newConn, err := s.handleEncrypted(context.Background(), tr, cs, oldConn, nil, wrongFrame, &plain)
if err != nil {
t.Fatalf("new session bad salt: %v", err)
}
if newConn == nil || newConn == oldConn || newConn.lifecycleState() != connLifecycleProvisional {
t.Fatalf("replacement conn old=%p new=%p state=%v", oldConn, newConn, newConn.lifecycleState())
}
if !oldConn.terminal.Load() || tr.closed.Load() {
t.Fatalf("transfer state old_terminal=%v raw_closed=%v", oldConn.terminal.Load(), tr.closed.Load())
}
// A delayed stale close must not tear down the generation already transferred
// to the new logical session.
oldConn.ForceClose()
if tr.closed.Load() || !newConn.isPhysicalTransportCurrentOpen() {
t.Fatalf("stale ForceClose closed replacement: raw_closed=%v current_open=%v", tr.closed.Load(), newConn.isPhysicalTransportCurrentOpen())
}
corrected, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, serverSalt, serverSalt, secondSID, secondID, 1,
)
activated, err := s.handleEncrypted(context.Background(), tr, cs, newConn, nil, corrected, &plain)
if err != nil {
t.Fatalf("activate transferred session: %v", err)
}
if activated != newConn || !activated.isActive() {
t.Fatalf("transferred session active=%v conn=%p want=%p", activated != nil && activated.isActive(), activated, newConn)
}
waitForAtomicCalls(t, &handler.calls, 2)
activated.ForceClose()
}
type collectingSessionTransport struct {
mu sync.Mutex
frames [][]byte
closed atomic.Bool
}
func (t *collectingSessionTransport) Send(_ context.Context, b *bin.Buffer) error {
if t.closed.Load() {
return io.ErrClosedPipe
}
t.mu.Lock()
t.frames = append(t.frames, append([]byte(nil), b.Raw()...))
t.mu.Unlock()
return nil
}
func (*collectingSessionTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
func (t *collectingSessionTransport) Close() error {
t.closed.Store(true)
return nil
}
func (t *collectingSessionTransport) snapshot() [][]byte {
t.mu.Lock()
defer t.mu.Unlock()
out := make([][]byte, len(t.frames))
for i := range t.frames {
out[i] = append([]byte(nil), t.frames[i]...)
}
return out
}
type reconnectFlightRPC struct {
calls atomic.Int32
started chan struct{}
release chan struct{}
once sync.Once
}
func (h *reconnectFlightRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) {
h.calls.Add(1)
h.once.Do(func() { close(h.started) })
<-h.release // deliberately ignore replacement cancellation after committing work
return &tg.Config{ThisDC: 2}, nil
}
func (*reconnectFlightRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
type cancelThenRetryRPC struct {
calls atomic.Int32
active atomic.Int32
max atomic.Int32
started chan struct{}
once sync.Once
}
func (h *cancelThenRetryRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
call := h.calls.Add(1)
active := h.active.Add(1)
defer h.active.Add(-1)
for {
max := h.max.Load()
if active <= max || h.max.CompareAndSwap(max, active) {
break
}
}
if call == 1 {
h.once.Do(func() { close(h.started) })
<-ctx.Done()
return nil, ctx.Err()
}
return &tg.Config{ThisDC: 2}, nil
}
func (*cancelThenRetryRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestHandleEncryptedRequiredSessionBarrierPrecedesStateRegistrationAndRPC(t *testing.T) {
handler := &admissionCountingRPC{}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
tr := newGatedRequiredControlTransport(nil)
key := newTestAuthKey(t)
const (
salt = int64(0x10203040)
sessionID = int64(70001)
)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
frame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID)
if err := s.authKeys.Save(context.Background(), stored); err != nil {
t.Fatalf("save auth key: %v", err)
}
cs := newConnState()
var plain bin.Buffer
type result struct {
conn *Conn
err error
}
done := make(chan result, 1)
go func() {
conn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, frame, &plain)
done <- result{conn: conn, err: err}
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("new_session_created did not reach the physical-write barrier")
}
if got := handler.calls.Load(); got != 0 {
t.Fatalf("handler calls while session barrier blocked = %d, want 0", got)
}
if cs.createdFloor != 0 || len(cs.seen) != 0 {
t.Fatalf("connState committed before session barrier: floor=%d seen=%d", cs.createdFloor, len(cs.seen))
}
s.conns.mu.RLock()
registered := s.conns.bySession[sessionKey{authKeyID: key.ID, sessionID: sessionID}]
claim := s.conns.claims[sessionKey{authKeyID: key.ID, sessionID: sessionID}]
s.conns.mu.RUnlock()
if registered != nil || claim == nil {
t.Fatalf("blocked barrier visibility = registered:%p claim:%p, want nil/non-nil", registered, claim)
}
select {
case got := <-done:
t.Fatalf("handleEncrypted returned before physical barrier: %v", got.err)
default:
}
tr.unblock()
var got result
select {
case got = <-done:
case <-time.After(2 * time.Second):
t.Fatal("handleEncrypted did not finish after session barrier write")
}
if got.err != nil {
t.Fatalf("handleEncrypted: %v", got.err)
}
if got.conn == nil || !got.conn.isActive() {
t.Fatalf("connection after successful barrier = %p lifecycle=%v", got.conn, got.conn.lifecycleState())
}
if cs.createdFloor != msgID {
t.Fatalf("created floor after barrier = %d, want %d", cs.createdFloor, msgID)
}
waitForAtomicCalls(t, &handler.calls, 1)
s.conns.mu.RLock()
registered = s.conns.bySession[sessionKey{authKeyID: key.ID, sessionID: sessionID}]
s.conns.mu.RUnlock()
if registered != got.conn {
t.Fatalf("registered connection = %p, want %p", registered, got.conn)
}
got.conn.ForceClose()
}
func TestHandleEncryptedRequiredSessionBarrierFailureIsAtomic(t *testing.T) {
handler := &admissionCountingRPC{}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
tr := newGatedRequiredControlTransport(io.ErrClosedPipe)
key := newTestAuthKey(t)
const (
salt = int64(0x50607080)
sessionID = int64(70002)
)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
frame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID)
if err := s.authKeys.Save(context.Background(), stored); err != nil {
t.Fatalf("save auth key: %v", err)
}
cs := newConnState()
var plain bin.Buffer
done := make(chan error, 1)
go func() {
_, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, frame, &plain)
done <- err
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("failed new_session_created did not reach physical writer")
}
if handler.calls.Load() != 0 || cs.createdFloor != 0 || len(cs.seen) != 0 {
t.Fatalf("state changed while failing barrier blocked: calls=%d floor=%d seen=%d", handler.calls.Load(), cs.createdFloor, len(cs.seen))
}
tr.unblock()
select {
case err := <-done:
if err == nil || !(errors.Is(err, io.ErrClosedPipe) || errors.Is(err, ErrConnClosed)) {
t.Fatalf("failed session barrier error = %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("failed session barrier did not return")
}
if got := handler.calls.Load(); got != 0 {
t.Fatalf("handler ran after failed session barrier: %d", got)
}
if cs.createdFloor != 0 || len(cs.seen) != 0 || len(cs.order) != 0 {
t.Fatalf("failed barrier committed connState: floor=%d seen=%d order=%d", cs.createdFloor, len(cs.seen), len(cs.order))
}
s.conns.mu.RLock()
registered := s.conns.bySession[sessionKey{authKeyID: key.ID, sessionID: sessionID}]
claim := s.conns.claims[sessionKey{authKeyID: key.ID, sessionID: sessionID}]
s.conns.mu.RUnlock()
if registered != nil || claim != nil {
t.Fatalf("failed barrier leaked manager ownership: registered=%p claim=%p", registered, claim)
}
if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("failed barrier leaked RPC budget: tasks=%d bytes=%d", tasks, bytes)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("failed barrier leaked frame budget: %d", got)
}
if got := s.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("failed barrier leaked RPC claim slots: %d", got)
}
}
func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testing.T) {
handler := &reconnectFlightRPC{started: make(chan struct{}), release: make(chan struct{})}
s := New(Options{RPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
key := newTestAuthKey(t)
const (
salt = int64(0x11223344)
sessionID = int64(70003)
)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
firstFrame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID)
if err := s.authKeys.Save(context.Background(), stored); err != nil {
t.Fatalf("save auth key: %v", err)
}
firstTransport := &collectingSessionTransport{}
firstState := newConnState()
var firstPlain bin.Buffer
firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, &stored, firstFrame, &firstPlain)
if err != nil {
t.Fatalf("first handleEncrypted: %v", err)
}
select {
case <-handler.started:
case <-time.After(2 * time.Second):
t.Fatal("first RPC owner did not start")
}
secondFrame, _ := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID)
secondTransport := &collectingSessionTransport{}
secondState := newConnState()
var secondPlain bin.Buffer
type handleResult struct {
conn *Conn
err error
}
secondDone := make(chan handleResult, 1)
go func() {
conn, handleErr := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain)
secondDone <- handleResult{conn: conn, err: handleErr}
}()
deadline := time.Now().Add(2 * time.Second)
for !firstConn.terminal.Load() && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if !firstConn.terminal.Load() {
t.Fatal("replacement did not fence the first physical connection")
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("overlapping reconnect executed %d business handlers, want 1", got)
}
select {
case got := <-secondDone:
t.Fatalf("duplicate reconnect returned before owner completion: %v", got.err)
default:
}
close(handler.release)
var second handleResult
select {
case second = <-secondDone:
case <-time.After(3 * time.Second):
t.Fatal("duplicate reconnect did not receive owner result")
}
if second.err != nil {
t.Fatalf("second handleEncrypted: %v", second.err)
}
if second.conn == nil || !second.conn.isActive() {
t.Fatalf("second connection was not active after replay: %p", second.conn)
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("cross-connection duplicate business calls = %d, want 1", got)
}
resultCount := 0
for _, wire := range secondTransport.snapshot() {
data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire})
if decryptErr != nil {
t.Fatalf("decrypt second connection reply: %v", decryptErr)
}
plain := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)}
typeID, peekErr := plain.PeekID()
if peekErr != nil {
t.Fatalf("peek second connection reply: %v", peekErr)
}
if typeID != proto.ResultTypeID {
continue
}
var result proto.Result
if decodeErr := result.Decode(plain); decodeErr != nil {
t.Fatalf("decode replayed rpc_result: %v", decodeErr)
}
if result.RequestMessageID != msgID {
t.Fatalf("replayed rpc_result request id = %d, want %d", result.RequestMessageID, msgID)
}
resultCount++
}
if resultCount != 1 {
t.Fatalf("replayed rpc_result count on replacement = %d, want 1", resultCount)
}
waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0)
if got := s.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("completed reconnect leaked flight claims: %d", got)
}
second.conn.ForceClose()
}
func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T) {
handler := &cancelThenRetryRPC{started: make(chan struct{})}
s := New(Options{RPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
key := newTestAuthKey(t)
const (
salt = int64(0x55667788)
sessionID = int64(70004)
)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
firstFrame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID)
if err := s.authKeys.Save(context.Background(), stored); err != nil {
t.Fatalf("save auth key: %v", err)
}
firstTransport := &collectingSessionTransport{}
firstState := newConnState()
var firstPlain bin.Buffer
firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, &stored, firstFrame, &firstPlain)
if err != nil {
t.Fatalf("first handleEncrypted: %v", err)
}
select {
case <-handler.started:
case <-time.After(2 * time.Second):
t.Fatal("first cancellation-aware owner did not start")
}
secondFrame, _ := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID)
secondTransport := &collectingSessionTransport{}
secondState := newConnState()
var secondPlain bin.Buffer
secondConn, err := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain)
if err != nil {
t.Fatalf("second handleEncrypted: %v", err)
}
waitForAtomicCalls(t, &handler.calls, 2)
waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0)
if got := handler.max.Load(); got != 1 {
t.Fatalf("old and retry handlers overlapped: max active=%d, want 1", got)
}
if got := s.rpcResults.flightLimit.snapshot(); got != 0 {
t.Fatalf("sequential retry leaked flight claims: %d", got)
}
resultCount := 0
for _, wire := range secondTransport.snapshot() {
data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire})
if decryptErr != nil {
t.Fatalf("decrypt sequential-retry reply: %v", decryptErr)
}
plain := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)}
typeID, peekErr := plain.PeekID()
if peekErr != nil {
t.Fatalf("peek sequential-retry reply: %v", peekErr)
}
if typeID == proto.ResultTypeID {
resultCount++
}
}
if resultCount != 1 {
t.Fatalf("sequential retry result count = %d, want 1", resultCount)
}
if !firstConn.terminal.Load() || secondConn == nil || !secondConn.isActive() {
t.Fatalf("replacement lifecycle = old terminal:%v new:%p active:%v", firstConn.terminal.Load(), secondConn, secondConn != nil && secondConn.isActive())
}
secondConn.ForceClose()
}
func waitForAtomicCalls(t *testing.T, calls interface{ Load() int32 }, want int32) {
t.Helper()
deadline := time.Now().Add(2 * time.Second)
for calls.Load() != want && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := calls.Load(); got != want {
t.Fatalf("handler calls = %d, want %d", got, want)
}
}

383
internal/mtprotoedge/session_boundary_test.go Normal file
View file

@ -0,0 +1,383 @@
package mtprotoedge
import (
"context"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
)
// TestFirstContainerBoundaryKeepsAndroidRequestMap models DrKLO's
// new_session_created handling: it drops every running request whose msg_id is
// lower than first_msg_id. Every inner request accepted from the first
// container must therefore remain addressable when its pong arrives.
func TestFirstContainerBoundaryKeepsAndroidRequestMap(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
msgIDs := proto.NewMessageIDGen(time.Now)
running := make(map[int64]int64, 3)
messages := make([]proto.Message, 0, 3)
for i := 0; i < 3; i++ {
msgID := msgIDs.New(proto.MessageFromClient)
pingID := int64(10_001 + i)
body := mustEncodeTL(t, &mt.PingRequest{PingID: pingID})
messages = append(messages, proto.Message{
ID: msgID,
SeqNo: 1 + i*2,
Bytes: len(body),
Body: body,
})
running[msgID] = pingID
}
outerMsgID := msgIDs.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages})
boundaryFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.NewSessionCreatedTypeID: 1,
})
var created mt.NewSessionCreated
for _, frame := range boundaryFrames {
if frame.TypeID == mt.PongTypeID {
t.Fatalf("pong arrived before new_session_created boundary")
}
if frame.TypeID != mt.NewSessionCreatedTypeID {
continue
}
if err := created.Decode(frame.Plain); err != nil {
t.Fatalf("decode new_session_created: %v", err)
}
}
// DrKLO ConnectionsManager.cpp clears running requests with
// request.messageId < first_msg_id when it receives this notification.
for msgID := range running {
if msgID < created.FirstMsgID {
delete(running, msgID)
}
}
for _, accepted := range messages {
if _, ok := running[accepted.ID]; !ok {
t.Fatalf(
"accepted inner msg_id %d was evicted by Android boundary %d (outer=%d)",
accepted.ID,
created.FirstMsgID,
outerMsgID,
)
}
}
pongFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.PongTypeID: len(messages),
})
seen := make(map[int64]struct{}, len(messages))
for _, frame := range pongFrames {
if frame.TypeID != mt.PongTypeID {
continue
}
var pong mt.Pong
if err := pong.Decode(frame.Plain); err != nil {
t.Fatalf("decode pong: %v", err)
}
wantPingID, ok := running[pong.MsgID]
if !ok {
t.Fatalf("orphan pong for msg_id %d after Android boundary cleanup", pong.MsgID)
}
if pong.PingID != wantPingID {
t.Fatalf("pong ping_id = %d for msg_id %d, want %d", pong.PingID, pong.MsgID, wantPingID)
}
if _, duplicate := seen[pong.MsgID]; duplicate {
t.Fatalf("duplicate pong for msg_id %d", pong.MsgID)
}
seen[pong.MsgID] = struct{}{}
delete(running, pong.MsgID)
}
if len(running) != 0 {
t.Fatalf("requests without correlated pong: %+v", running)
}
}
type admissionCountingRPC struct {
calls atomic.Int32
}
func (h *admissionCountingRPC) Dispatch(_ context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
h.calls.Add(1)
return &tg.Config{ThisDC: 2}, nil
}
func (*admissionCountingRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestContainerRPCAdmissionFailureIsAtomic(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCGlobalWorkers: 1,
})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
msgIDs := proto.NewMessageIDGen(time.Now)
messages := make([]proto.Message, 0, 2)
requestIDs := make(map[int64]struct{}, 2)
for i := 0; i < 2; i++ {
msgID := msgIDs.New(proto.MessageFromClient)
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
messages = append(messages, proto.Message{ID: msgID, SeqNo: 1 + i*2, Bytes: len(body), Body: body})
requestIDs[msgID] = struct{}{}
}
outerMsgID := msgIDs.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 2})
results := 0
for _, frame := range frames {
if frame.TypeID != proto.ResultTypeID {
continue
}
var result proto.Result
if err := result.Decode(frame.Plain); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if _, ok := requestIDs[result.RequestMessageID]; !ok {
t.Fatalf("unexpected or duplicate capacity rpc_result req_msg_id %d", result.RequestMessageID)
}
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode capacity rpc_error: %v", err)
}
if rpcErr.ErrorCode != 420 || rpcErr.ErrorMessage != "FLOOD_WAIT_1" {
t.Fatalf("capacity rpc_error = %+v", rpcErr)
}
delete(requestIDs, result.RequestMessageID)
results++
}
if results != 2 {
t.Fatalf("capacity rpc_results = %d, want 2", results)
}
if len(requestIDs) != 0 {
t.Fatalf("capacity requests without exactly one result: %+v", requestIDs)
}
if got := handler.calls.Load(); got != 0 {
t.Fatalf("partially executed handler calls = %d, want 0", got)
}
}
func TestGZIPWrappedRPCUsesLogicalEnvelopeBoundary(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
requestBody := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &proto.GZIP{Data: requestBody})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range frames {
if frame.TypeID == mt.NewSessionCreatedTypeID {
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode gzip RPC boundary: %v", err)
}
}
}
if boundary.FirstMsgID != msgID {
t.Fatalf("gzip RPC boundary = %d, want envelope %d", boundary.FirstMsgID, msgID)
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("gzip RPC handler calls = %d, want 1", got)
}
}
func TestGZIPWrappedContainerUsesLowestInnerBoundary(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
innerMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 7001})
containerBody := mustEncodeTL(t, &proto.MessageContainer{Messages: []proto.Message{{
ID: innerMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody,
}}})
sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, 2, &proto.GZIP{Data: containerBody})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.PongTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range frames {
if frame.TypeID == mt.NewSessionCreatedTypeID {
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode gzip container boundary: %v", err)
}
}
}
if boundary.FirstMsgID != innerMsgID {
t.Fatalf("gzip container boundary = %d, want inner %d (outer %d)", boundary.FirstMsgID, innerMsgID, outerMsgID)
}
}
func TestEmptyContainerUsesOuterBoundary(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
outerMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, 0, &proto.MessageContainer{})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.NewSessionCreatedTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range frames {
if frame.TypeID != mt.NewSessionCreatedTypeID {
continue
}
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode empty container boundary: %v", err)
}
}
if boundary.FirstMsgID != outerMsgID {
t.Fatalf("empty container boundary = %d, want outer %d", boundary.FirstMsgID, outerMsgID)
}
}
func TestDuplicateContainerAcksWithoutBusinessReexecution(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
messages := make([]proto.Message, 0, 2)
for i := 0; i < 2; i++ {
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
messages = append(messages, proto.Message{
ID: ids.New(proto.MessageFromClient), SeqNo: 1 + i*2, Bytes: len(body), Body: body,
})
}
outerMsgID := ids.New(proto.MessageFromClient)
container := &proto.MessageContainer{Messages: messages}
sendEncrypted(t, conn, cipher, auth, outerMsgID, container)
initialFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 2})
answerIDs := make([]int64, 0, 2)
for _, frame := range initialFrames {
if frame.TypeID == proto.ResultTypeID {
answerIDs = append(answerIDs, frame.Message.MessageID)
}
}
if len(answerIDs) != 2 {
t.Fatalf("initial rpc_result answer ids = %d, want 2", len(answerIDs))
}
// ACK both original server results before retransmitting the container. This is
// a wire barrier: any later rpc_result is necessarily a duplicate replay rather
// than an unconsumed response from the initial batch.
ackMsgID := ids.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, ackMsgID, 4, &mt.MsgsAck{MsgIDs: answerIDs})
if got := handler.calls.Load(); got != 2 {
t.Fatalf("initial handler calls = %d, want 2", got)
}
sendEncrypted(t, conn, cipher, auth, outerMsgID, container)
duplicateFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.MsgsAckTypeID: 1})
for _, frame := range duplicateFrames {
if frame.TypeID == proto.ResultTypeID {
t.Fatal("same-connection duplicate container replayed an extra rpc_result")
}
}
if got := handler.calls.Load(); got != 2 {
t.Fatalf("duplicate container reexecuted handlers: calls=%d, want 2", got)
}
}
type largeStartupBurstRPC struct {
calls atomic.Int32
body string
}
func (h *largeStartupBurstRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
call := h.calls.Add(1)
// Deliberately perturb completion order while remaining cancellation-aware.
timer := time.NewTimer(time.Duration(call%4) * time.Millisecond)
defer timer.Stop()
select {
case <-timer.C:
case <-ctx.Done():
return nil, ctx.Err()
}
return &tg.DataJSON{Data: h.body}, nil
}
func (*largeStartupBurstRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestAndroidStartupBurstKeepsAllLargeRPCResultsAddressable(t *testing.T) {
const (
dc = 2
requests = 30
)
handler := &largeStartupBurstRPC{body: strings.Repeat("x", 192<<10)}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
running := make(map[int64]struct{}, requests)
messages := make([]proto.Message, 0, requests)
for i := 0; i < requests; i++ {
msgID := ids.New(proto.MessageFromClient)
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
messages = append(messages, proto.Message{
ID: msgID, SeqNo: 1 + i*2, Bytes: len(body), Body: body,
})
running[msgID] = struct{}{}
}
outerMsgID := ids.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages})
boundaryFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.NewSessionCreatedTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range boundaryFrames {
if frame.TypeID == proto.ResultTypeID {
t.Fatal("large rpc_result arrived before session boundary")
}
if frame.TypeID == mt.NewSessionCreatedTypeID {
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode startup boundary: %v", err)
}
}
}
for msgID := range running {
if msgID < boundary.FirstMsgID {
delete(running, msgID)
}
}
if len(running) != requests {
t.Fatalf("Android boundary removed accepted startup requests: kept=%d want=%d floor=%d outer=%d", len(running), requests, boundary.FirstMsgID, outerMsgID)
}
resultFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: requests})
for _, frame := range resultFrames {
if frame.TypeID != proto.ResultTypeID {
continue
}
var result proto.Result
if err := result.Decode(frame.Plain); err != nil {
t.Fatalf("decode startup rpc_result: %v", err)
}
if _, ok := running[result.RequestMessageID]; !ok {
t.Fatalf("orphan large rpc_result for request %d", result.RequestMessageID)
}
delete(running, result.RequestMessageID)
}
if len(running) != 0 {
t.Fatalf("startup RPCs without result: %d", len(running))
}
if got := handler.calls.Load(); got != requests {
t.Fatalf("startup handler calls = %d, want %d", got, requests)
}
}

330
internal/mtprotoedge/session_manager.go
View file

@ -22,6 +22,11 @@ var ErrSessionNotFound = errors.New("session not found")
// ErrSessionAmbiguous 表示仅用 session_id 无法唯一定位连接。
var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys")
var (
ErrSessionActivationSuperseded = errors.New("session activation superseded")
ErrSessionActivationFence = errors.New("session activation could not fence previous writer")
)
const (
maxPendingPushesPerSession = 32
// maxFlushAttempts / flushRetryBackoff：排空暂存推送时 c.Send 失败（出站拥塞 5s 超时
@ -117,8 +122,13 @@ type SessionLifecycleObserver interface {
// 它管理运行态的在线连接，与持久化的 store.SessionStore 互补：后者记录 session 数据，
// 前者持有可发送的活跃连接。所有方法并发安全。
type SessionManager struct {
mu sync.RWMutex
bySession map[sessionKey]*Conn
mu sync.RWMutex
bySession map[sessionKey]*Conn
// claims owns the provisional -> active gap. A claimant is intentionally
// absent from every push/online index until its required session control frame
// is on the wire and PublishActivation validates the same owner.
claims map[sessionKey]*Conn
claimsByAuth map[[8]byte]map[int64]*Conn // raw authKeyID -> sessionID -> provisional claim
bySessionID map[int64]map[[8]byte]*Conn // sessionID → raw authKeyID → Conn，用于兼容旧 API 的唯一性检查
byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn
byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn
@ -142,6 +152,8 @@ func NewSessionManager(log *zap.Logger) *SessionManager {
}
return &SessionManager{
bySession: make(map[sessionKey]*Conn),
claims: make(map[sessionKey]*Conn),
claimsByAuth: make(map[[8]byte]map[int64]*Conn),
bySessionID: make(map[int64]map[[8]byte]*Conn),
byAuthKey: make(map[[8]byte]map[int64]*Conn),
byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn),
@ -164,32 +176,94 @@ func (m *SessionManager) SetLifecycleObserver(observer SessionLifecycleObserver)
m.mu.Unlock()
}
// Register 注册一个活跃连接。若同 raw auth_key_id + session_id 已存在（重连），旧连接被替换并移除索引。
func (m *SessionManager) Register(c *Conn) {
m.mu.Lock()
// BeginActivation atomically claims auth_key_id + session_id without publishing the
// new Conn. Under the manager lock it irreversibly fences every previous owner,
// removes active indexes and closes producer/RPC admission gates. Physical close and
// outbound-actor convergence happen outside the lock; the caller may send the
// required new_session_created frame only after this method returns nil.
func (m *SessionManager) BeginActivation(c *Conn) error {
if c == nil || !c.beginActivationClaim() {
return ErrSessionActivationSuperseded
}
key := connSessionKey(c)
var replaced *Conn
var evicted *Conn
if old, ok := m.bySession[key]; ok && old != c {
replaced = old
m.removeLocked(old, false)
} else if existing := m.byAuthKey[c.authKeyID]; len(existing) >= maxSessionsPerAuthKey {
// 同 raw auth_key 的 session 数达上限且本次是新 session：驱逐建连最早的 session 让位，
// 防对抗客户端用海量 session_id 撑爆索引。驱逐对象与新连接同属一个设备凭据，
// 触顶基本是该凭据自身异常；选最旧而非 map 随机，避免误杀刚建立的活跃下载/主连接。
// 被驱逐连接的 serveConn 会在下一帧因 actor 已关而退出。O(cap) 扫描仅在触顶时发生。
for _, ec := range existing {
if evicted == nil || ec.createdAt.Before(evicted.createdAt) {
evicted = ec
}
retired := make([]*Conn, 0, 2)
m.mu.Lock()
if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() || c.lifecycleState() != connLifecycleClaiming {
c.beginTerminalShutdown()
m.mu.Unlock()
return ErrConnClosed
}
if oldClaim := m.claims[key]; oldClaim != nil && oldClaim != c {
m.retireClaimLocked(key, oldClaim, false)
retired = append(retired, oldClaim)
}
if old := m.bySession[key]; old != nil && old != c {
m.retireConnLocked(old, false)
retired = append(retired, old)
}
// Claims reserve a cap slot just like published sessions. Otherwise many
// concurrent handshakes could all pass the old byAuthKey-only check and publish
// beyond maxSessionsPerAuthKey.
for len(m.byAuthKey[c.authKeyID])+m.claimCountForAuthLocked(c.authKeyID) >= maxSessionsPerAuthKey {
victimKey, victim, isClaim := m.oldestAuthOwnerLocked(c.authKeyID, c)
if victim == nil {
break
}
m.removeLocked(evicted, true)
m.log.Debug("Evicted oldest session for auth key at cap",
if isClaim {
m.retireClaimLocked(victimKey, victim, true)
} else {
m.retireConnLocked(victim, true)
}
retired = append(retired, victim)
m.log.Debug("Evicted oldest session activation for auth key at cap",
zap.String("auth_key_id", sessionKeyLog(c.authKeyID)),
zap.Int("cap", maxSessionsPerAuthKey),
)
}
m.addClaimLocked(key, c)
m.mu.Unlock()
// Do not wait for old business handlers: their root context/admission gate is
// already canceled. We only need the old physical writer and outbound actor to
// converge before the new Conn is allowed to write the session barrier.
if !closeConnBatch(retired, forceCloseBatchTimeout, false) {
m.AbortActivation(c)
return ErrSessionActivationFence
}
return nil
}
// PublishActivation makes the current claim visible to push/online lookups. It is
// deliberately a separate operation from BeginActivation so required protocol
// control can be written while the Conn remains provisional and unindexed.
func (m *SessionManager) PublishActivation(c *Conn) error {
if c == nil {
return ErrSessionActivationSuperseded
}
key := connSessionKey(c)
m.mu.Lock()
defer m.mu.Unlock()
if m.claims[key] != c {
return ErrSessionActivationSuperseded
}
if c.terminal.Load() || c.lifecycleState() != connLifecycleClaiming {
m.removeClaimLocked(key, c)
return ErrConnClosed
}
if old := m.bySession[key]; old != nil && old != c {
// An active owner without this claim can only be a stale/unsafe publisher.
// Never reverse-replace it from here; the caller must reconnect and claim again.
m.removeClaimLocked(key, c)
c.beginTerminalShutdown()
return ErrSessionActivationSuperseded
}
if !c.publishActivation() {
m.removeClaimLocked(key, c)
return ErrConnClosed
}
m.removeClaimLocked(key, c)
m.bySession[key] = c
addSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID, c)
addConnIndex(m.byAuthKey, c.authKeyID, c.sessionID, c)
@ -200,34 +274,88 @@ func (m *SessionManager) Register(c *Conn) {
c.userIDResolved.Store(true)
addUserIndex(m.byUser, uid, key, c)
}
m.log.Debug("Session registered",
m.log.Debug("Session activated",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
zap.Int64("session_id", c.sessionID),
zap.Int("online", len(m.bySession)),
)
m.mu.Unlock()
return nil
}
// 同 identity 的新物理连接已经原子接管索引；立即关闭旧 transport，不能只停
// actor 后让旧 FD/read goroutine 滞留到 read timeout。replacement 与 cap eviction
// 共用一个并发关闭批次，不能把每条 Conn 的 RPC 等待上界串行相加。
if replaced != nil || evicted != nil {
if !forceCloseConnBatch([]*Conn{replaced, evicted}, forceCloseBatchTimeout) {
m.log.Warn("Session replacement/eviction close exceeded shared deadline")
}
// AbortActivation removes a claim only when c still owns it, then terminally
// retires the provisional Conn. A superseded caller cannot delete the newer claim.
func (m *SessionManager) AbortActivation(c *Conn) {
if c == nil {
return
}
key := connSessionKey(c)
owned := false
m.mu.Lock()
if m.claims[key] == c {
c.beginTerminalShutdown()
m.removeClaimLocked(key, c)
if m.bySession[key] == nil {
m.deletePendingLocked(key)
delete(m.flushing, key)
}
owned = true
}
m.mu.Unlock()
if owned {
_ = closeConnBatch([]*Conn{c}, forceCloseBatchTimeout, false)
}
}
// Register is retained for tests and embedders that do not have a wire-level
// required-control barrier. Production encrypted traffic uses the explicit
// BeginActivation -> SendRequiredControl -> PublishActivation sequence.
func (m *SessionManager) Register(c *Conn) error {
if c == nil {
return ErrSessionActivationSuperseded
}
if c.isActive() {
m.mu.RLock()
current := m.bySession[connSessionKey(c)]
m.mu.RUnlock()
if current == c {
return nil
}
return ErrSessionActivationSuperseded
}
if err := m.BeginActivation(c); err != nil {
return err
}
if err := m.PublishActivation(c); err != nil {
m.AbortActivation(c)
return err
}
return nil
}
// Unregister 注销一个连接（仅当它仍是当前注册的同一对象，避免误删重连后的新连接）。
// 观察者对未登录连接（userID=0）也回调：业务层据此清理按 session 维度的缓存条目，
// 否则未登录连接的元数据只能等容量上限驱逐。
func (m *SessionManager) Unregister(c *Conn) {
if c == nil {
return
}
// Close admission/outbound producer gates before removing indexes or invoking
// a lifecycle observer. An observer may block, but no old RPC/push may continue
// to write or enqueue work during that interval.
c.beginTerminalShutdown()
m.mu.Lock()
var (
observer SessionLifecycleObserver
offlineUser int64
lastForUser bool
)
if cur, ok := m.bySession[connSessionKey(c)]; ok && cur == c {
key := connSessionKey(c)
if m.claims[key] == c {
m.removeClaimLocked(key, c)
m.deletePendingLocked(key)
delete(m.flushing, key)
}
if cur, ok := m.bySession[key]; ok && cur == c {
offlineUser = m.removeLocked(c, true)
if offlineUser != 0 {
lastForUser = len(m.byUser[offlineUser]) == 0
@ -256,7 +384,7 @@ func (m *SessionManager) DestroySession(sessionID int64) bool {
m.mu.Unlock()
return false
}
offlineUser := m.removeLocked(c, true)
offlineUser := m.retireConnLocked(c, true)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
observer := m.lifecycle
m.log.Debug("Session destroyed",
@ -283,11 +411,22 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
if claim := m.claims[key]; claim != nil {
m.retireClaimLocked(key, claim, true)
m.mu.Unlock()
if !forceCloseConnBatch([]*Conn{claim}, forceCloseBatchTimeout) {
m.log.Warn("Claimed session close exceeded shared deadline",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int64("session_id", sessionID),
)
}
return true
}
m.deletePendingLocked(key)
m.mu.Unlock()
return false
}
offlineUser := m.removeLocked(c, true)
offlineUser := m.retireConnLocked(c, true)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
observer := m.lifecycle
m.log.Debug("Session destroyed",
@ -505,10 +644,17 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int
if !connUsesBusinessAuthKey(c, authKeyID) {
continue
}
uid := m.removeLocked(c, true)
uid := m.retireConnLocked(c, true)
conns = append(conns, c)
events = append(events, offlineEvent{key: key, userID: uid, last: uid != 0 && len(m.byUser[uid]) == 0})
}
for key, c := range m.claims {
if !connUsesBusinessAuthKey(c, authKeyID) {
continue
}
m.retireClaimLocked(key, c, true)
conns = append(conns, c)
}
observer := m.lifecycle
if len(conns) > 0 {
m.log.Debug("Force close sessions for revoked auth key",
@ -531,10 +677,26 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int
return len(conns)
}
// CloseSessionsForRawAuthKeyExcept 强制断开指定 raw auth_key 的活跃连接，可排除
// 一个 session（destroy_auth_key 的发起连接：响应要先送达，它的密钥已删，下一帧
// 自然失效）。出站推送不回查密钥库，必须主动断开底层 transport 才能让销毁立即生效。
// CloseSessionsForRawAuthKeyExcept 强制断开指定 raw auth_key 的活跃连接，可按
// session ID 排除一个 session。该接口供业务层授权撤销使用；wire-level
// destroy_auth_key 必须改用精确 Conn 排除，避免同 session replacement 被误放过。
func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exceptSessionID int64) int {
return m.closeSessionsForRawAuthKey(authKeyID, func(sessionID int64, _ *Conn) bool {
return sessionID == exceptSessionID
})
}
// CloseSessionsForRawAuthKeyExceptConn closes every active/claiming owner for a raw
// auth key except the exact Conn executing destroy_auth_key. A session ID is not an
// identity: a concurrent replacement may already own the same logical session while
// the retired request handler finishes deletion.
func (m *SessionManager) CloseSessionsForRawAuthKeyExceptConn(authKeyID [8]byte, except *Conn) int {
return m.closeSessionsForRawAuthKey(authKeyID, func(_ int64, c *Conn) bool {
return c == except
})
}
func (m *SessionManager) closeSessionsForRawAuthKey(authKeyID [8]byte, skip func(int64, *Conn) bool) int {
type offlineEvent struct {
key sessionKey
userID int64
@ -544,14 +706,22 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc
var conns []*Conn
var events []offlineEvent
for sessionID, c := range m.byAuthKey[authKeyID] {
if sessionID == exceptSessionID {
if skip != nil && skip(sessionID, c) {
continue
}
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
uid := m.removeLocked(c, true)
uid := m.retireConnLocked(c, true)
conns = append(conns, c)
events = append(events, offlineEvent{key: key, userID: uid, last: uid != 0 && len(m.byUser[uid]) == 0})
}
for sessionID, c := range m.claimsByAuth[authKeyID] {
if skip != nil && skip(sessionID, c) {
continue
}
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
m.retireClaimLocked(key, c, true)
conns = append(conns, c)
}
observer := m.lifecycle
m.mu.Unlock()
if !forceCloseConnBatch(conns, forceCloseBatchTimeout) {
@ -574,6 +744,14 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc
// deadline, but no timed-out Conn can enqueue more work in that interval. Nil/duplicate entries are
// removed so Register's replacement/eviction slots cannot close the same Conn twice.
func forceCloseConnBatch(conns []*Conn, timeout time.Duration) bool {
return closeConnBatch(conns, timeout, true)
}
// closeConnBatch always converges physical writers/outbound actors. waitInbound
// is reserved for destructive control-plane operations; activation takeover sets
// it false so a canceled business handler that ignores its context cannot stall a
// healthy replacement. Its admission and response writer are already terminal.
func closeConnBatch(conns []*Conn, timeout time.Duration, waitInbound bool) bool {
if len(conns) == 0 {
return true
}
@ -640,7 +818,7 @@ func forceCloseConnBatch(conns []*Conn, timeout time.Duration) bool {
if remaining <= 0 {
return false
}
if c.rpcScheduler != nil && !c.waitInboundShutdown(remaining) {
if waitInbound && c.rpcScheduler != nil && !c.waitInboundShutdown(remaining) {
return false
}
if c.outboundDone == nil {
@ -1631,6 +1809,9 @@ func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64
func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
key := connSessionKey(c)
if m.bySession[key] != c {
return 0
}
delete(m.bySession, key)
removeSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID)
removeConnIndex(m.byAuthKey, c.authKeyID, c.sessionID)
@ -1650,6 +1831,75 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
return uid
}
// retireConnLocked closes every admission/producer gate before the Conn leaves
// manager indexes. Callers may close the physical transport and wait outside m.mu,
// but no pointer collected by an earlier fan-out can enqueue after this returns.
func (m *SessionManager) retireConnLocked(c *Conn, dropPending bool) int64 {
if c == nil {
return 0
}
c.beginTerminalShutdown()
return m.removeLocked(c, dropPending)
}
func (m *SessionManager) retireClaimLocked(key sessionKey, c *Conn, dropPending bool) {
if c == nil || m.claims[key] != c {
return
}
c.beginTerminalShutdown()
m.removeClaimLocked(key, c)
if dropPending && m.bySession[key] == nil {
m.deletePendingLocked(key)
}
delete(m.flushing, key)
}
func (m *SessionManager) claimCountForAuthLocked(authKeyID [8]byte) int {
return len(m.claimsByAuth[authKeyID])
}
func (m *SessionManager) oldestAuthOwnerLocked(authKeyID [8]byte, exclude *Conn) (sessionKey, *Conn, bool) {
var (
oldestKey sessionKey
oldest *Conn
isClaim bool
)
for sessionID, candidate := range m.byAuthKey[authKeyID] {
if candidate == nil || candidate == exclude {
continue
}
if oldest == nil || candidate.createdAt.Before(oldest.createdAt) {
oldestKey = sessionKey{authKeyID: authKeyID, sessionID: sessionID}
oldest = candidate
isClaim = false
}
}
for sessionID, candidate := range m.claimsByAuth[authKeyID] {
if candidate == nil || candidate == exclude {
continue
}
if oldest == nil || candidate.createdAt.Before(oldest.createdAt) {
oldestKey = sessionKey{authKeyID: authKeyID, sessionID: sessionID}
oldest = candidate
isClaim = true
}
}
return oldestKey, oldest, isClaim
}
func (m *SessionManager) addClaimLocked(key sessionKey, c *Conn) {
m.claims[key] = c
addConnIndex(m.claimsByAuth, key.authKeyID, key.sessionID, c)
}
func (m *SessionManager) removeClaimLocked(key sessionKey, c *Conn) {
if m.claims[key] != c {
return
}
delete(m.claims, key)
removeConnIndex(m.claimsByAuth, key.authKeyID, key.sessionID)
}
func (m *SessionManager) businessAuthKeyCandidatesLocked(authKeyID [8]byte) map[sessionKey]*Conn {
out := make(map[sessionKey]*Conn, len(m.byBusinessAuthKey[authKeyID])+len(m.byAuthKey[authKeyID]))
for key, c := range m.byBusinessAuthKey[authKeyID] {

462
internal/mtprotoedge/structural_limits_test.go
View file

@ -2,8 +2,12 @@ package mtprotoedge
import (
"context"
"encoding/binary"
"errors"
"io"
"strings"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
@ -39,6 +43,101 @@ func TestContainerMessageCountAndServiceVectorCaps(t *testing.T) {
}
}
func TestServiceLongVectorViewsAreBoundedExactAndZeroCopy(t *testing.T) {
tests := []struct {
name string
typeID uint32
value bin.Encoder
}{
{name: "msgs_ack", typeID: mt.MsgsAckTypeID, value: &mt.MsgsAck{MsgIDs: []int64{11, 22}}},
{name: "msgs_state_req", typeID: mt.MsgsStateReqTypeID, value: &mt.MsgsStateReq{MsgIDs: []int64{11, 22}}},
{name: "msg_resend_req", typeID: mt.MsgResendReqTypeID, value: &mt.MsgResendReq{MsgIDs: []int64{11, 22}}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var encoded bin.Buffer
if err := tt.value.Encode(&encoded); err != nil {
t.Fatalf("encode service vector: %v", err)
}
view, err := decodeInt64VectorView(&encoded, tt.typeID, maxServiceMessageIDs)
if err != nil {
t.Fatalf("decode service vector view: %v", err)
}
if view.count != 2 {
t.Fatalf("view count = %d, want 2", view.count)
}
// The retained plan payload is a view into the already-budgeted frame,
// not a generated-decoder []int64 copy.
binary.LittleEndian.PutUint64(encoded.Buf[12:20], uint64(99))
ids := view.materialize()
if len(ids) != 2 || ids[0] != 99 || ids[1] != 22 {
t.Fatalf("materialized ids = %v, want [99 22]", ids)
}
trailing := bin.Buffer{Buf: append(append([]byte(nil), encoded.Buf...), 0, 0, 0, 0)}
if _, err := decodeInt64VectorView(&trailing, tt.typeID, maxServiceMessageIDs); err == nil || !strings.Contains(err.Error(), "trailing") {
t.Fatalf("trailing vector err = %v, want exact-length rejection", err)
}
truncated := bin.Buffer{Buf: encoded.Buf[:len(encoded.Buf)-1]}
if _, err := decodeInt64VectorView(&truncated, tt.typeID, maxServiceMessageIDs); !errors.Is(err, io.ErrUnexpectedEOF) {
t.Fatalf("truncated vector err = %v, want unexpected EOF", err)
}
})
}
var oversized bin.Buffer
if err := (&mt.MsgsAck{MsgIDs: []int64{1}}).Encode(&oversized); err != nil {
t.Fatalf("encode oversized seed: %v", err)
}
binary.LittleEndian.PutUint32(oversized.Buf[8:12], uint32(maxServiceMessageIDs+1))
if _, err := decodeInt64VectorView(&oversized, mt.MsgsAckTypeID, maxServiceMessageIDs); err == nil || !strings.Contains(err.Error(), "exceeds") {
t.Fatalf("oversized vector err = %v, want capped rejection", err)
}
}
func TestInboundPlanContainerCapacityBoundary(t *testing.T) {
ids := proto.NewMessageIDGen(time.Now)
waitBody := encodeClientMessageBodyForTest(t, &mt.HTTPWaitRequest{MaxWait: 25_000})
messages := make([]proto.Message, maxContainerMessages)
for i := range messages {
messages[i] = proto.Message{
ID: ids.New(proto.MessageFromClient), SeqNo: 0, Bytes: len(waitBody), Body: waitBody,
}
}
outerMsgID := ids.New(proto.MessageFromClient)
body := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: messages})
s := New(Options{Logger: zaptest.NewLogger(t)})
before := s.frameBudget.usedBytes()
plan, err := s.preflightInbound(newConnState(), outerMsgID, 0, body)
if err != nil {
t.Fatalf("container at cap preflight: %v", err)
}
if len(plan.items) != maxContainerMessages || len(plan.staged) != maxContainerMessages+1 {
t.Fatalf("container plan sizes = items:%d staged:%d", len(plan.items), len(plan.staged))
}
if plan.logicalMin != messages[0].ID {
t.Fatalf("container plan floor = %d, want %d", plan.logicalMin, messages[0].ID)
}
plan.close()
if after := s.frameBudget.usedBytes(); after != before {
t.Fatalf("container plan leaked frame budget: before=%d after=%d", before, after)
}
over := append(append([]proto.Message(nil), messages...), proto.Message{
ID: outerMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody,
})
overBody := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: over})
if plan, err := s.preflightInbound(newConnState(), ids.New(proto.MessageFromClient), 0, overBody); err == nil {
plan.close()
t.Fatal("container above cap unexpectedly accepted")
}
if after := s.frameBudget.usedBytes(); after != before {
t.Fatalf("rejected container leaked frame budget: before=%d after=%d", before, after)
}
}
func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) {
encoded := bin.Buffer{}
wantBody := []byte{0x11, 0x22, 0x33, 0x44}
@ -83,6 +182,319 @@ func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) {
}
}
func TestContainerAndGZIPViewsRejectTrailingBytes(t *testing.T) {
var container bin.Buffer
if err := (&proto.MessageContainer{}).Encode(&container); err != nil {
t.Fatalf("encode empty container: %v", err)
}
container.PutInt(0x11223344)
s := New(Options{Logger: zaptest.NewLogger(t)})
if _, release, err := s.decodeMessageContainerViews(&container, 0); err == nil {
release()
t.Fatal("container trailing bytes unexpectedly accepted")
}
var packed bin.Buffer
if err := (proto.GZIP{Data: []byte{1, 2, 3, 4}}).Encode(&packed); err != nil {
t.Fatalf("encode gzip wrapper: %v", err)
}
packed.PutInt(0x55667788)
if _, err := gzipPackedBytesView(&packed); err == nil {
t.Fatal("gzip trailing bytes unexpectedly accepted")
}
}
func TestContainerInvalidSequenceTailIsAtomic(t *testing.T) {
ids := proto.NewMessageIDGen(time.Now)
firstMsgID := ids.New(proto.MessageFromClient)
secondMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
pingOne := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 1})
pingTwo := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2})
container := proto.MessageContainer{Messages: []proto.Message{
{ID: firstMsgID, SeqNo: 3, Bytes: len(pingOne), Body: pingOne},
{ID: secondMsgID, SeqNo: 1, Bytes: len(pingTwo), Body: pingTwo},
}}
encoded := encodeClientMessageBodyForTest(t, &container)
s := New(Options{Logger: zaptest.NewLogger(t)})
cs := newConnState()
c := &Conn{
metrics: NopMetrics{},
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outboundStop: make(chan struct{}),
}
var acks []int64
if err := s.dispatch(context.Background(), cs, c, outerMsgID, 4, &bin.Buffer{Buf: encoded}, &acks); err != nil {
t.Fatalf("dispatch invalid-tail container: %v", err)
}
if len(cs.seen) != 0 || len(cs.order) != 0 || cs.maxContentMsgID != 0 || cs.maxContentSeqNo != 0 {
t.Fatalf("invalid container partially committed connState: %+v", cs)
}
if len(acks) != 0 {
t.Fatalf("invalid container produced ACKs: %v", acks)
}
if got := len(c.outboundControl); got != 1 {
t.Fatalf("control frames = %d, want only bad_msg_notification", got)
}
op := <-c.outboundControl
if op.encoded == nil || op.encoded.typeID != mt.BadMsgNotificationTypeID {
t.Fatalf("control frame type = %#v, want bad_msg_notification", op.encoded)
}
var bad mt.BadMsgNotification
if err := bad.Decode(&bin.Buffer{Buf: op.encoded.body}); err != nil {
t.Fatalf("decode bad_msg_notification: %v", err)
}
if bad.BadMsgID != secondMsgID || bad.BadMsgSeqno != 1 || bad.ErrorCode != badMsgSeqTooLow {
t.Fatalf("bad_msg_notification = %+v", bad)
}
op.releaseReservation(c.outboundTrackedBudget)
}
func TestContainerMalformedGZIPTailHasNoStateEffects(t *testing.T) {
ids := proto.NewMessageIDGen(time.Now)
waitMsgID := ids.New(proto.MessageFromClient)
badMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
waitBody := encodeClientMessageBodyForTest(t, &mt.HTTPWaitRequest{MaxWait: 25_000})
var malformedGZIP bin.Buffer
malformedGZIP.PutID(proto.GZIPTypeID)
container := proto.MessageContainer{Messages: []proto.Message{
{ID: waitMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody},
{ID: badMsgID, SeqNo: 1, Bytes: malformedGZIP.Len(), Body: malformedGZIP.Buf},
}}
encoded := encodeClientMessageBodyForTest(t, &container)
s := New(Options{Logger: zaptest.NewLogger(t)})
cs := newConnState()
var acks []int64
if err := s.dispatch(context.Background(), cs, nil, outerMsgID, 2, &bin.Buffer{Buf: encoded}, &acks); err == nil {
t.Fatal("malformed gzip tail unexpectedly accepted")
}
if len(cs.seen) != 0 || len(cs.order) != 0 || cs.maxContentMsgID != 0 || cs.maxContentSeqNo != 0 {
t.Fatalf("malformed gzip tail partially committed connState: %+v", cs)
}
if len(acks) != 0 {
t.Fatalf("malformed gzip tail produced ACKs: %v", acks)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("malformed gzip tail leaked frame budget: %d", got)
}
}
func TestInvalidMessageIDRejectsBeforeGZIPExpansion(t *testing.T) {
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 1})
wrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
current := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
tests := []struct {
name string
msgID int64
badCode int
}{
{name: "stale", msgID: proto.NewMessageIDGen(func() time.Time { return time.Now().Add(-10 * time.Minute) }).New(proto.MessageFromClient), badCode: badMsgIDTooLow},
{name: "future", msgID: proto.NewMessageIDGen(func() time.Time { return time.Now().Add(time.Minute) }).New(proto.MessageFromClient), badCode: badMsgIDTooHigh},
{name: "invalid bits", msgID: current + 1, badCode: badMsgIDInvalidBits},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
s := New(Options{Logger: zaptest.NewLogger(t)})
// Any attempted gzip expansion would fail this deliberately tiny
// budget and mask the required bad_msg_notification.
s.frameBudget = newInboundFrameBudget(1)
plan, err := s.preflightInbound(newConnState(), tt.msgID, 1, wrapped)
if plan != nil {
plan.close()
t.Fatal("invalid gzip envelope unexpectedly produced a plan")
}
var bad *dispatchBadMsgError
if !errors.As(err, &bad) || bad.code != tt.badCode {
t.Fatalf("invalid gzip envelope err = %v, want bad_msg code %d", err, tt.badCode)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("invalid gzip envelope consumed expansion budget: %d", got)
}
})
}
}
func TestInvalidContainerInnerIDRejectsBeforeGZIPExpansion(t *testing.T) {
ids := proto.NewMessageIDGen(time.Now)
validInnerID := ids.New(proto.MessageFromClient)
invalidInnerID := validInnerID + 1
if proto.MessageID(invalidInnerID).Type() == proto.MessageFromClient {
t.Fatalf("test inner msg_id %d unexpectedly has client bits", invalidInnerID)
}
outerMsgID := ids.New(proto.MessageFromClient)
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2})
gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{
ID: invalidInnerID, SeqNo: 1, Bytes: len(gzipped), Body: gzipped,
}}})
s := New(Options{Logger: zaptest.NewLogger(t)})
// Enough for the one descriptor, intentionally not enough for a gzip
// expansion. Invalid inner bits must win and every descriptor charge unwind.
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
plan, err := s.preflightInbound(newConnState(), outerMsgID, 2, container)
if plan != nil {
plan.close()
t.Fatal("invalid inner msg_id unexpectedly produced a plan")
}
var bad *dispatchBadMsgError
if !errors.As(err, &bad) || bad.code != badMsgContainer || bad.msgID != invalidInnerID {
t.Fatalf("invalid inner gzip err = %v, want container rejection for %d", err, invalidInnerID)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("invalid inner gzip leaked descriptor/expansion budget: %d", got)
}
}
func TestSeenGZIPInnerShortCircuitsBeforeExpansion(t *testing.T) {
tests := []struct {
name string
seqNo int32
content bool
leaf bin.Encoder
wantACK bool
}{
{name: "content", seqNo: 1, content: true, leaf: &mt.PingRequest{PingID: 3}, wantACK: true},
{name: "non_content", seqNo: 0, content: false, leaf: &mt.HTTPWaitRequest{MaxWait: 25_000}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ids := proto.NewMessageIDGen(time.Now)
innerMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
leaf := encodeClientMessageBodyForTest(t, tt.leaf)
gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{
ID: innerMsgID, SeqNo: int(tt.seqNo), Bytes: len(gzipped), Body: gzipped,
}}})
cs := newConnState()
cs.track(innerMsgID, tt.seqNo, tt.content, msgStateReceived)
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
plan, err := s.preflightInbound(cs, outerMsgID, 2, container)
if err != nil {
t.Fatalf("preflight seen gzip inner: %v", err)
}
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || plan.items[0].content != tt.content {
plan.close()
t.Fatalf("seen gzip plan items = %+v, want one duplicate", plan.items)
}
if gotACK := len(plan.ackIDs) == 1 && plan.ackIDs[0] == innerMsgID; gotACK != tt.wantACK {
plan.close()
t.Fatalf("seen gzip ACKs = %v, wantACK=%v", plan.ackIDs, tt.wantACK)
}
if got := s.frameBudget.usedBytes(); got != containerDescriptorBudgetBytes {
plan.close()
t.Fatalf("seen gzip held budget = %d, want descriptor-only %d", got, containerDescriptorBudgetBytes)
}
plan.close()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("seen gzip close leaked budget: %d", got)
}
})
}
}
func TestSeenTopLevelContentGZIPShortCircuitsBeforeExpansion(t *testing.T) {
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 31})
wrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
cs := newConnState()
cs.track(msgID, 1, true, msgStateReceived)
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(1)
plan, err := s.preflightInbound(cs, msgID, 1, wrapped)
if err != nil {
t.Fatalf("preflight seen top-level gzip: %v", err)
}
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || len(plan.ackIDs) != 1 || plan.ackIDs[0] != msgID {
plan.close()
t.Fatalf("seen top-level gzip plan = items:%+v acks:%v", plan.items, plan.ackIDs)
}
plan.close()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("seen top-level gzip consumed expansion budget: %d", got)
}
}
func TestDuplicateOuterContainerOnlyInspectsInnerDescriptors(t *testing.T) {
ids := proto.NewMessageIDGen(time.Now)
innerMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 4})
gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{
ID: innerMsgID, SeqNo: 1, Bytes: len(gzipped), Body: gzipped,
}}})
wrappedContainer := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: container})
t.Run("seen inner replays without expansion", func(t *testing.T) {
cs := newConnState()
cs.track(outerMsgID, 2, false, msgStateReceived)
cs.track(innerMsgID, 1, true, msgStateReceived)
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
plan, err := s.preflightInbound(cs, outerMsgID, 2, container)
if err != nil {
t.Fatalf("preflight duplicate outer: %v", err)
}
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate {
plan.close()
t.Fatalf("duplicate outer items = %+v", plan.items)
}
plan.close()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("duplicate outer leaked descriptor budget: %d", got)
}
})
t.Run("top-level gzip container still locates inner", func(t *testing.T) {
cs := newConnState()
cs.track(outerMsgID, 2, false, msgStateReceived)
cs.track(innerMsgID, 1, true, msgStateReceived)
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes + containerDescriptorBudgetBytes)
plan, err := s.preflightInbound(cs, outerMsgID, 2, wrappedContainer)
if err != nil {
t.Fatalf("preflight gzip duplicate outer: %v", err)
}
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || plan.items[0].msgID != innerMsgID {
plan.close()
t.Fatalf("gzip duplicate outer items = %+v, want inner duplicate", plan.items)
}
plan.close()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("gzip duplicate outer leaked expansion/descriptor budget: %d", got)
}
})
t.Run("unseen inner rejects without expansion", func(t *testing.T) {
cs := newConnState()
cs.track(outerMsgID, 2, false, msgStateReceived)
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
plan, err := s.preflightInbound(cs, outerMsgID, 2, container)
if plan != nil {
plan.close()
t.Fatal("duplicate outer with unseen inner unexpectedly produced a plan")
}
var bad *dispatchBadMsgError
if !errors.As(err, &bad) || bad.code != badMsgContainer {
t.Fatalf("duplicate outer unseen-inner err = %v, want container rejection", err)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("unseen duplicate inner leaked descriptor/expansion budget: %d", got)
}
})
}
func TestServiceInfoViewsRejectOversizedBytesWithoutDecodeCopy(t *testing.T) {
state := mt.MsgsStateInfo{ReqMsgID: 7, Info: make([]byte, maxServiceMessageIDs)}
var encodedState bin.Buffer
@ -142,10 +554,58 @@ func TestDispatchRejectsExcessiveWrapperDepthBeforeRPC(t *testing.T) {
s := New(Options{Logger: zaptest.NewLogger(t)})
var acks []int64
err := s.dispatch(context.Background(), newConnState(), nil, 4, 0, &bin.Buffer{Buf: encoded}, &acks)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
err := s.dispatch(context.Background(), newConnState(), nil, msgID, 0, &bin.Buffer{Buf: encoded}, &acks)
if err == nil || !strings.Contains(err.Error(), "wrapper depth") {
t.Fatalf("deep wrapper err = %v, want wrapper depth rejection", err)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("deep wrapper rejection leaked frame budget: %d", got)
}
}
func TestNestedWrapperTailFailuresReleaseEveryBudget(t *testing.T) {
ping := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 5})
var gzipWithTail bin.Buffer
if err := (&proto.GZIP{Data: ping}).Encode(&gzipWithTail); err != nil {
t.Fatalf("encode inner gzip: %v", err)
}
gzipWithTail.PutInt(0x11223344)
gzipTailWrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: gzipWithTail.Copy()})
var containerWithTail bin.Buffer
if err := (&proto.MessageContainer{}).Encode(&containerWithTail); err != nil {
t.Fatalf("encode inner container: %v", err)
}
containerWithTail.PutInt(0x55667788)
containerTailWrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: containerWithTail.Copy()})
tests := []struct {
name string
body []byte
seqNo int32
}{
{name: "gzip tail after outer expansion", body: gzipTailWrapped, seqNo: 1},
{name: "container tail after outer expansion", body: containerTailWrapped, seqNo: 0},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
s := New(Options{Logger: zaptest.NewLogger(t)})
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
plan, err := s.preflightInbound(newConnState(), msgID, tt.seqNo, tt.body)
if plan != nil {
plan.close()
t.Fatal("invalid nested wrapper unexpectedly produced a plan")
}
if err == nil || !strings.Contains(err.Error(), "trailing") {
t.Fatalf("nested wrapper err = %v, want trailing-byte rejection", err)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("nested wrapper failure leaked frame budget: %d", got)
}
})
}
}
func TestOversizedConnectionBuffersAreReleasedAfterFrame(t *testing.T) {

339
internal/mtprotoedge/transport_ownership.go Normal file
View file

@ -0,0 +1,339 @@
package mtprotoedge
import (
"context"
"sync"
"sync/atomic"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/transport"
)
const (
physicalTransportClosedBit = uint64(1) << 63
physicalTransportGenerationMax = physicalTransportClosedBit - 1
)
// physicalTransportOwner separates the lifetime of one physical socket from
// the logical Conn generations that successively use it. state is either an
// open generation in [1, physicalTransportGenerationMax], or that generation
// with physicalTransportClosedBit set.
//
// Transfer and owner-close race through one atomic state transition. This is
// the critical property: when transfer wins, a stale generation can no longer
// close the socket; when close wins, no later generation can be published.
type physicalTransportOwner struct {
raw transport.Conn
// writeMu orders a completed send before generation transfer. CloseAny and
// owner-close deliberately do not wait for it: raw.Close must be able to
// interrupt a transport implementation blocked inside Send.
writeMu sync.Mutex
state atomic.Uint64
// binding identifies the logical Conn currently attached to the open
// generation. Physical close fences that exact Conn before touching raw;
// generation matching prevents a stale close from retiring a replacement.
bindingMu sync.Mutex
boundGeneration uint64
boundLogicalConn *Conn
closeDone chan struct{}
closeErr error
}
// physicalTransportLease is the capability owned by one logical Conn
// generation. It implements transport.Conn so exchange/protocol-error paths
// can use the initial lease from the moment serveConn starts.
type physicalTransportLease struct {
owner *physicalTransportOwner
generation uint64
}
var _ transport.Conn = (*physicalTransportLease)(nil)
func newPhysicalTransportOwner(raw transport.Conn) (*physicalTransportOwner, *physicalTransportLease) {
owner := &physicalTransportOwner{
raw: raw,
closeDone: make(chan struct{}),
}
owner.state.Store(1)
return owner, &physicalTransportLease{owner: owner, generation: 1}
}
// Transfer atomically hands the physical transport to the next logical Conn
// generation. Holding writeMu ensures every send admitted by the old lease has
// returned before the generation changes. CloseAny may still interrupt such a
// send and make the CAS fail, which is the required shutdown ordering.
func (l *physicalTransportLease) Transfer() (*physicalTransportLease, bool) {
if l == nil || l.owner == nil || l.generation == 0 {
return nil, false
}
owner := l.owner
owner.writeMu.Lock()
defer owner.writeMu.Unlock()
if l.generation >= physicalTransportGenerationMax {
return nil, false
}
if !owner.state.CompareAndSwap(l.generation, l.generation+1) {
return nil, false
}
owner.bindingMu.Lock()
if owner.boundGeneration == l.generation {
owner.boundGeneration = l.generation + 1
owner.boundLogicalConn = nil
}
owner.bindingMu.Unlock()
return &physicalTransportLease{owner: owner, generation: l.generation + 1}, true
}
// IsCurrentOpen reports whether this lease is still the unique open owner.
// Callers use it at protocol publication barriers after a required write.
func (l *physicalTransportLease) IsCurrentOpen() bool {
return l != nil && l.owner != nil && l.generation != 0 &&
l.owner.state.Load() == l.generation
}
// bindLogicalConn publishes the logical Conn for this exact generation. If
// physical close already won, the new Conn is terminally fenced immediately
// and can never pass SessionManager publication checks.
func (l *physicalTransportLease) bindLogicalConn(c *Conn) bool {
if l == nil || l.owner == nil || c == nil {
if c != nil {
c.beginTerminalShutdown()
}
return false
}
owner := l.owner
owner.bindingMu.Lock()
open := owner.state.Load() == l.generation
if open {
owner.boundGeneration = l.generation
owner.boundLogicalConn = c
}
owner.bindingMu.Unlock()
if !open {
c.beginTerminalShutdown()
}
return open
}
// Send admits a write only while this lease is the current open generation.
// The lock also serializes quick ACK/protocol writes with the outbound actor.
func (l *physicalTransportLease) Send(ctx context.Context, b *bin.Buffer) error {
return l.withCurrentWriter(func(raw transport.Conn) error {
return raw.Send(ctx, b)
})
}
// SendDeadline preserves the allocation-free fast path of compatTransportConn.
func (l *physicalTransportLease) SendDeadline(deadline time.Time, b *bin.Buffer) error {
return l.withCurrentWriter(func(raw transport.Conn) error {
if writer, ok := raw.(deadlineOutboundWriter); ok {
return writer.SendDeadline(deadline, b)
}
ctx := context.Background()
cancel := func() {}
if !deadline.IsZero() {
ctx, cancel = context.WithDeadline(ctx, deadline)
}
defer cancel()
return raw.Send(ctx, b)
})
}
func (l *physicalTransportLease) withCurrentWriter(send func(transport.Conn) error) error {
if l == nil || l.owner == nil || l.owner.raw == nil {
return ErrConnClosed
}
owner := l.owner
owner.writeMu.Lock()
defer owner.writeMu.Unlock()
if owner.state.Load() != l.generation {
return ErrConnClosed
}
return send(owner.raw)
}
// Recv is owned by serveConn rather than a logical Conn generation. It is a
// direct proxy so the read loop remains valid after ownership transfers.
func (l *physicalTransportLease) Recv(ctx context.Context, b *bin.Buffer) error {
if l == nil || l.owner == nil || l.owner.raw == nil {
return ErrConnClosed
}
return l.owner.raw.Recv(ctx, b)
}
// RecvDeadline preserves serveConn's direct-deadline fast path.
func (l *physicalTransportLease) RecvDeadline(deadline time.Time, b *bin.Buffer) error {
if l == nil || l.owner == nil || l.owner.raw == nil {
return ErrConnClosed
}
if receiver, ok := l.owner.raw.(deadlineReceiver); ok {
return receiver.RecvDeadline(deadline, b)
}
ctx := context.Background()
cancel := func() {}
if !deadline.IsZero() {
ctx, cancel = context.WithDeadline(ctx, deadline)
}
defer cancel()
return l.owner.raw.Recv(ctx, b)
}
// Close closes the physical transport only if this lease still owns the
// current generation. A stale logical Conn therefore cannot close a socket
// already transferred to its replacement.
func (l *physicalTransportLease) Close() error {
if l == nil || l.owner == nil || l.generation == 0 {
return nil
}
owner := l.owner
for {
state := owner.state.Load()
if state&physicalTransportClosedBit != 0 {
return owner.waitClosed()
}
if state != l.generation {
return nil
}
if owner.state.CompareAndSwap(state, state|physicalTransportClosedBit) {
owner.fenceBoundGeneration(l.generation)
return owner.closeRaw()
}
}
}
// startCloseAlreadyFenced is the non-reentrant close capability for the logical
// Conn that has already published terminal/lifecycle gates itself. It
// synchronously publishes the physical closed bit (so Transfer cannot win), then
// runs the potentially pathological raw.Close outside the RPC worker/flight
// handoff. It neither calls fenceBoundGeneration nor waits for a CloseAny that
// already won, avoiding both lifecycle reentry and close cycles.
func (l *physicalTransportLease) startCloseAlreadyFenced() {
if l == nil || l.owner == nil || l.generation == 0 {
return
}
owner := l.owner
for {
state := owner.state.Load()
if state&physicalTransportClosedBit != 0 || state != l.generation {
return
}
if owner.state.CompareAndSwap(state, state|physicalTransportClosedBit) {
owner.bindingMu.Lock()
if owner.boundGeneration == l.generation {
owner.boundLogicalConn = nil
}
owner.bindingMu.Unlock()
go func() { _ = owner.closeRaw() }()
return
}
}
}
// CloseAny is the unconditional physical-socket capability retained by
// serveConn. It marks the owner closed before calling raw.Close and never waits
// for writeMu, so it can break a blocked Send or Recv.
func (o *physicalTransportOwner) CloseAny() error {
if o == nil {
return nil
}
for {
state := o.state.Load()
if state&physicalTransportClosedBit != 0 {
return o.waitClosed()
}
if o.state.CompareAndSwap(state, state|physicalTransportClosedBit) {
o.fenceBoundGeneration(state)
return o.closeRaw()
}
}
}
func (o *physicalTransportOwner) fenceBoundGeneration(generation uint64) {
o.bindingMu.Lock()
var c *Conn
if o.boundGeneration == generation {
c = o.boundLogicalConn
o.boundLogicalConn = nil
}
o.bindingMu.Unlock()
if c != nil {
c.beginTerminalShutdown()
}
}
func (o *physicalTransportOwner) closeRaw() error {
if o.raw != nil {
o.closeErr = o.raw.Close()
}
close(o.closeDone)
return o.closeErr
}
func (o *physicalTransportOwner) waitClosed() error {
<-o.closeDone
return o.closeErr
}
// Forward the optional compat-transport capabilities hidden by the lease.
// These keep frame-budget ownership and quick-ack semantics unchanged while
// serveConn operates on the initial lease instead of the raw transport.
func (l *physicalTransportLease) releaseInboundFrame() {
if l == nil || l.owner == nil {
return
}
if releaser, ok := l.owner.raw.(inboundFrameOwnershipReleaser); ok {
releaser.releaseInboundFrame()
}
}
func (l *physicalTransportLease) retainInboundFrameBytes(n int64) bool {
if l == nil || l.owner == nil {
return true
}
if retainer, ok := l.owner.raw.(inboundFrameBackingRetainer); ok {
return retainer.retainInboundFrameBytes(n)
}
return true
}
func (l *physicalTransportLease) ConsumeQuickAckRequested() bool {
if l == nil || l.owner == nil {
return false
}
if quick, ok := l.owner.raw.(quickAckTransport); ok {
return quick.ConsumeQuickAckRequested()
}
return false
}
func (l *physicalTransportLease) SendQuickAck(ctx context.Context, token uint32) error {
return l.withCurrentWriter(func(raw transport.Conn) error {
if quick, ok := raw.(quickAckTransport); ok {
return quick.SendQuickAck(ctx, token)
}
return nil
})
}
func (l *physicalTransportLease) SendQuickAckDeadline(deadline time.Time, token uint32) error {
return l.withCurrentWriter(func(raw transport.Conn) error {
if quick, ok := raw.(deadlineQuickAckTransport); ok {
return quick.SendQuickAckDeadline(deadline, token)
}
if quick, ok := raw.(quickAckTransport); ok {
ctx := context.Background()
cancel := func() {}
if !deadline.IsZero() {
ctx, cancel = context.WithDeadline(ctx, deadline)
}
defer cancel()
return quick.SendQuickAck(ctx, token)
}
return nil
})
}

198
internal/mtprotoedge/transport_ownership_test.go Normal file
View file

@ -0,0 +1,198 @@
package mtprotoedge
import (
"context"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
)
type ownershipTestTransport struct {
closeCalls atomic.Int32
sendCalls atomic.Int32
closed chan struct{}
closeOnce sync.Once
blockSend bool
started chan struct{}
startOnce sync.Once
}
func newOwnershipTestTransport(blockSend bool) *ownershipTestTransport {
return &ownershipTestTransport{
closed: make(chan struct{}),
blockSend: blockSend,
started: make(chan struct{}),
}
}
func (t *ownershipTestTransport) Send(context.Context, *bin.Buffer) error {
t.sendCalls.Add(1)
t.startOnce.Do(func() { close(t.started) })
if !t.blockSend {
select {
case <-t.closed:
return io.ErrClosedPipe
default:
return nil
}
}
<-t.closed
return io.ErrClosedPipe
}
func (*ownershipTestTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
func (t *ownershipTestTransport) Close() error {
t.closeOnce.Do(func() {
t.closeCalls.Add(1)
close(t.closed)
})
return nil
}
func TestPhysicalTransportTransferFencesStaleLease(t *testing.T) {
raw := newOwnershipTestTransport(false)
owner, oldLease := newPhysicalTransportOwner(raw)
newLease, ok := oldLease.Transfer()
if !ok {
t.Fatal("transfer failed")
}
if err := oldLease.Send(context.Background(), &bin.Buffer{}); !errors.Is(err, ErrConnClosed) {
t.Fatalf("stale Send error = %v, want ErrConnClosed", err)
}
if err := oldLease.Close(); err != nil {
t.Fatalf("stale Close: %v", err)
}
if got := raw.closeCalls.Load(); got != 0 {
t.Fatalf("stale Close closed raw %d times", got)
}
if err := newLease.Send(context.Background(), &bin.Buffer{}); err != nil {
t.Fatalf("current Send: %v", err)
}
if err := owner.CloseAny(); err != nil {
t.Fatalf("owner CloseAny: %v", err)
}
if got := raw.closeCalls.Load(); got != 1 {
t.Fatalf("raw close calls = %d, want 1", got)
}
}
func TestPhysicalTransportCloseWinningPreventsTransfer(t *testing.T) {
raw := newOwnershipTestTransport(false)
owner, lease := newPhysicalTransportOwner(raw)
if err := owner.CloseAny(); err != nil {
t.Fatalf("CloseAny: %v", err)
}
if next, ok := lease.Transfer(); ok || next != nil {
t.Fatalf("transfer after close = (%p,%v), want nil,false", next, ok)
}
}
func TestPhysicalTransportCloseAnyInterruptsWriteAndDefeatsTransfer(t *testing.T) {
raw := newOwnershipTestTransport(true)
owner, lease := newPhysicalTransportOwner(raw)
sendDone := make(chan error, 1)
go func() { sendDone <- lease.Send(context.Background(), &bin.Buffer{}) }()
select {
case <-raw.started:
case <-time.After(time.Second):
t.Fatal("write did not block")
}
transferDone := make(chan bool, 1)
go func() {
_, ok := lease.Transfer()
transferDone <- ok
}()
select {
case ok := <-transferDone:
t.Fatalf("transfer returned before blocked write ended: %v", ok)
case <-time.After(20 * time.Millisecond):
}
if err := owner.CloseAny(); err != nil {
t.Fatalf("CloseAny: %v", err)
}
select {
case err := <-sendDone:
if !errors.Is(err, io.ErrClosedPipe) {
t.Fatalf("blocked Send error = %v", err)
}
case <-time.After(time.Second):
t.Fatal("CloseAny did not interrupt write")
}
select {
case ok := <-transferDone:
if ok {
t.Fatal("transfer won after physical close")
}
case <-time.After(time.Second):
t.Fatal("transfer did not finish")
}
}
func TestPhysicalCloseFencesActivationPublication(t *testing.T) {
raw := newOwnershipTestTransport(false)
owner, lease := newPhysicalTransportOwner(raw)
s := New(Options{})
c := s.newConnWithLease(lease, newTestAuthKey(t), 73001, 1)
if !c.beginActivationClaim() {
t.Fatal("activation claim failed")
}
if err := owner.CloseAny(); err != nil {
t.Fatalf("CloseAny: %v", err)
}
if c.publishActivation() {
t.Fatal("closed physical transport published active Conn")
}
if !c.terminal.Load() || c.lifecycleState() != connLifecycleRetired {
t.Fatalf("closed Conn terminal=%v lifecycle=%v", c.terminal.Load(), c.lifecycleState())
}
c.Close()
}
func TestPhysicalCloseBitPreventsActivationClaimBeforeLogicalFence(t *testing.T) {
raw := newOwnershipTestTransport(false)
owner, lease := newPhysicalTransportOwner(raw)
s := New(Options{})
c := s.newConnWithLease(lease, newTestAuthKey(t), 73002, 1)
// Hold the binding lock so CloseAny can linearize the physical closed bit but
// cannot yet publish c.terminal. beginActivationClaim must inspect the lease
// itself and refuse this otherwise-dangerous window.
owner.bindingMu.Lock()
closeDone := make(chan error, 1)
go func() { closeDone <- owner.CloseAny() }()
deadline := time.Now().Add(time.Second)
for owner.state.Load()&physicalTransportClosedBit == 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if owner.state.Load()&physicalTransportClosedBit == 0 {
owner.bindingMu.Unlock()
t.Fatal("CloseAny did not publish closed bit")
}
if c.terminal.Load() {
owner.bindingMu.Unlock()
t.Fatal("logical fence escaped held binding lock")
}
if c.beginActivationClaim() {
owner.bindingMu.Unlock()
t.Fatal("closed physical generation entered activation claim")
}
owner.bindingMu.Unlock()
select {
case err := <-closeDone:
if err != nil {
t.Fatalf("CloseAny: %v", err)
}
case <-time.After(time.Second):
t.Fatal("CloseAny did not finish")
}
if !c.terminal.Load() || c.lifecycleState() != connLifecycleRetired {
t.Fatalf("logical fence terminal=%v lifecycle=%v", c.terminal.Load(), c.lifecycleState())
}
c.Close()
}