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

perf: sync protocol and core hardening updates

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

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

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

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

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

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

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

50
internal/mtprotoedge/conn.go
View file

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

428
internal/mtprotoedge/encrypted.go
View file

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

64
internal/mtprotoedge/exchange.go
View file

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

75
internal/mtprotoedge/exchange_compat.go
View file

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

285
internal/mtprotoedge/exchange_test.go
View file

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

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

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

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

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

813
internal/mtprotoedge/inbound_rpc.go
View file

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

422
internal/mtprotoedge/inbound_rpc_test.go
View file

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

9
internal/mtprotoedge/login_email_e2e_test.go
View file

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

1000
internal/mtprotoedge/outbound.go
View file

File diff suppressed because it is too large Load diff

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

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

805
internal/mtprotoedge/outbound_test.go
View file

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

11
internal/mtprotoedge/passkey_e2e_test.go
View file

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

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

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

146
internal/mtprotoedge/rpc_test.go
View file

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

126
internal/mtprotoedge/same_port_mux.go
View file

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

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

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

315
internal/mtprotoedge/server.go
View file

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

2
internal/mtprotoedge/server_test.go
View file

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

556
internal/mtprotoedge/session_manager.go
View file

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

519
internal/mtprotoedge/session_manager_test.go
View file

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

60
internal/mtprotoedge/session_membership_gen_test.go
View file

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

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

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

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

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

257
internal/mtprotoedge/transport_compat.go
View file

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

14
internal/mtprotoedge/transport_compat_test.go
View file

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