755 lines
29 KiB
Go
755 lines
29 KiB
Go
package mtprotoedge
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rsa"
|
||
"encoding/hex"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net"
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"go.uber.org/zap"
|
||
|
||
"github.com/gotd/td/bin"
|
||
"github.com/gotd/td/clock"
|
||
"github.com/gotd/td/crypto"
|
||
"github.com/gotd/td/exchange"
|
||
"github.com/gotd/td/mt"
|
||
"github.com/gotd/td/proto"
|
||
"github.com/gotd/td/proto/codec"
|
||
"github.com/gotd/td/tg"
|
||
"github.com/gotd/td/tmap"
|
||
"github.com/gotd/td/transport"
|
||
|
||
"telesrv/internal/store"
|
||
"telesrv/internal/store/memory"
|
||
)
|
||
|
||
// RPCHandler 把解密后的 RPC 请求体路由到响应。由 internal/rpc 实现。
|
||
//
|
||
// b 是明文 RPC 请求（已剥离 MTProto 外壳）；返回的 bin.Encoder 会被包成 rpc_result。
|
||
// 返回 *tgerr.Error 时连接层将其转为 rpc_error 回发；其他 error 视为连接级故障。
|
||
type RPCHandler interface {
|
||
Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error)
|
||
// NegotiatedLayer returns the TL layer the session negotiated via
|
||
// invokeWithLayer and whether one was ever observed. Used to downgrade
|
||
// outbound objects for clients compiled on an older layer. ok=false means
|
||
// unknown (cold/evicted) — the caller must keep the connection's last-known
|
||
// layer rather than overwrite it.
|
||
NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool)
|
||
}
|
||
|
||
// Options 配置 Server。
|
||
type Options struct {
|
||
// Logger 日志器。默认 zap.NewNop()。
|
||
Logger *zap.Logger
|
||
// 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 字节前缀流程。
|
||
ObfuscatedTCP bool
|
||
// WebSocket 在同一个 listener 上接受 MTProto over WebSocket(/apiws*)。
|
||
// 开启后仅在连接建立时读取前 4 字节做 HTTP/TCP 分流；MTProto TCP
|
||
// 后续仍走原 ObfuscatedTCP + codec 热路径。
|
||
WebSocket bool
|
||
// WebSocketAllowedOrigins 是允许浏览器发起 WebSocket upgrade 的页面 origin。
|
||
// 空列表表示只接受无 Origin 的非浏览器客户端；"*" 表示允许所有来源（仅调试）。
|
||
WebSocketAllowedOrigins []string
|
||
// ReadTimeout 单次读取超时。默认 5m。
|
||
ReadTimeout time.Duration
|
||
// HandshakeIdleTimeout 是连接「建立 session 前」（握手 + 首个加密消息之前）的读超时，
|
||
// 比 ReadTimeout 短，用于快速回收握手后静默的半开 / 异常连接。默认 60s。
|
||
HandshakeIdleTimeout time.Duration
|
||
// HandshakeMaxDuration 是单次密钥交换（serverExchange）的总时长上界。HandshakeIdleTimeout
|
||
// 只约束「单次读 idle」，对一个持续发包的客户端无效——若客户端陷入「收到 ResPQ→nonce 失步
|
||
// →重发 req_pq」的握手重启死循环（见 docs/client-compat-notes.md），无界的 serverExchange 会
|
||
// 对每个 req_pq 盲回 ResPQ、永不收敛地空转刷日志/占 CPU。本上界给整个握手设总预算，超时即
|
||
// 放弃并断开，客户端重连发起全新握手（无残留相位差）即恢复。正常握手 <1s。默认 20s。
|
||
HandshakeMaxDuration time.Duration
|
||
// WriteTimeout 单次写入超时。默认 30s。
|
||
WriteTimeout time.Duration
|
||
// 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 队列长度。默认 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
|
||
// RSAKey 是 server RSA 私钥，用于密钥交换。nil 时无法完成握手。
|
||
RSAKey *rsa.PrivateKey
|
||
// AuthKeys 持久化 auth key。默认内存实现。
|
||
AuthKeys store.AuthKeyStore
|
||
// ActiveSessions 管理活跃连接。默认新建；传入时可让 RPC 层共享同一注册表。
|
||
ActiveSessions *SessionManager
|
||
// RPC 是 typed RPC 路由。nil 时加密 RPC 被丢弃并记录。
|
||
RPC RPCHandler
|
||
// Metrics 接收连接层指标。默认 NopMetrics。
|
||
Metrics Metrics
|
||
// Clock 用于消息 ID 与时间戳。默认 clock.System。
|
||
Clock clock.Clock
|
||
// Rand 随机源。默认 crypto.DefaultRand()。
|
||
Rand io.Reader
|
||
}
|
||
|
||
func (o *Options) setDefaults() {
|
||
if o.Logger == nil {
|
||
o.Logger = zap.NewNop()
|
||
}
|
||
if o.ReadTimeout == 0 {
|
||
o.ReadTimeout = 5 * time.Minute
|
||
}
|
||
if o.HandshakeIdleTimeout == 0 {
|
||
o.HandshakeIdleTimeout = 60 * time.Second
|
||
}
|
||
if o.HandshakeMaxDuration == 0 {
|
||
o.HandshakeMaxDuration = 20 * time.Second
|
||
}
|
||
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 = 64
|
||
}
|
||
if o.RPCTimeout == 0 {
|
||
o.RPCTimeout = 30 * time.Second
|
||
}
|
||
if o.RPCGlobalWorkers <= 0 {
|
||
o.RPCGlobalWorkers = 256
|
||
}
|
||
if o.RPCGlobalMaxTasks <= 0 {
|
||
o.RPCGlobalMaxTasks = rpcResultFlightDefaultMaxPending
|
||
}
|
||
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
|
||
}
|
||
if o.AuthKeys == nil {
|
||
o.AuthKeys = memory.NewAuthKeyStore()
|
||
}
|
||
if o.Metrics == nil {
|
||
o.Metrics = NopMetrics{}
|
||
}
|
||
if o.Clock == nil {
|
||
o.Clock = clock.System
|
||
}
|
||
if o.Rand == nil {
|
||
o.Rand = crypto.DefaultRand()
|
||
}
|
||
}
|
||
|
||
// Server 是 MTProto 连接层（mtprotoedge）。
|
||
//
|
||
// 职责见 doc.go。它把原始 TCP 字节流转换为「已解密、已识别 session 的 RPC 请求」：
|
||
// 接受连接、协商 codec、完成密钥交换、解密并分发加密消息到 RPC 路由，处理服务消息，
|
||
// 并把活跃连接注册到 SessionManager 以支持主动推送（updates 等）。不含业务逻辑。
|
||
type Server struct {
|
||
log *zap.Logger
|
||
codec func() transport.Codec
|
||
obfuscated bool
|
||
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
|
||
conns *SessionManager
|
||
rpc RPCHandler
|
||
metrics Metrics
|
||
cipher crypto.Cipher
|
||
clock clock.Clock
|
||
rand io.Reader
|
||
types *tmap.Map
|
||
admission *admissionController
|
||
|
||
rpcResults *rpcResultCache
|
||
|
||
// onFrame 是测试钩子：收到一帧时回调其字节数；生产为 nil。
|
||
onFrame func(n int)
|
||
}
|
||
|
||
// New 创建 Server。
|
||
func New(opts Options) *Server {
|
||
opts.setDefaults()
|
||
conns := opts.ActiveSessions
|
||
if conns == nil {
|
||
conns = NewSessionManager(opts.Logger.Named("sessions"))
|
||
}
|
||
return &Server{
|
||
log: opts.Logger,
|
||
codec: opts.Codec,
|
||
obfuscated: opts.ObfuscatedTCP,
|
||
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,
|
||
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: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks),
|
||
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
|
||
}
|
||
}
|
||
|
||
// newConn 基于一次解密结果创建一个可发送的连接对象。
|
||
func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn {
|
||
if lease, ok := tc.(*physicalTransportLease); ok {
|
||
return s.newConnWithLease(lease, key, sessionID, salt)
|
||
}
|
||
if tc != nil {
|
||
_, lease := newPhysicalTransportOwner(tc)
|
||
return s.newConnWithLease(lease, key, sessionID, salt)
|
||
}
|
||
// Preserve the nil transport used by construction-only tests.
|
||
return s.buildConn(nil, nil, key, sessionID, salt)
|
||
}
|
||
|
||
// newConnWithLease attaches a logical Conn to an explicitly owned physical
|
||
// transport generation. Production session replacement must Transfer the old
|
||
// lease first; it must never wrap the same raw transport in a second owner.
|
||
func (s *Server) newConnWithLease(lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn {
|
||
if lease == nil {
|
||
panic("mtprotoedge: nil physical transport lease")
|
||
}
|
||
c := s.buildConn(lease, lease, key, sessionID, salt)
|
||
lease.bindLogicalConn(c)
|
||
return c
|
||
}
|
||
|
||
func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn {
|
||
c := &Conn{
|
||
transport: tc,
|
||
transportLease: lease,
|
||
writer: tc,
|
||
cipher: s.cipher,
|
||
msgID: proto.NewMessageIDGen(s.clock.Now),
|
||
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.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)
|
||
}
|
||
return s.serveTCP(ctx, ln)
|
||
}
|
||
|
||
func (s *Server) serveTCP(ctx context.Context, ln net.Listener) error {
|
||
ctx, cancel := context.WithCancel(ctx)
|
||
defer cancel()
|
||
|
||
s.log.Info("Serving", zap.String("addr", ln.Addr().String()), zap.Int("dc", s.dc), zap.Bool("obfuscated_tcp", s.obfuscated))
|
||
defer s.log.Info("Stopped")
|
||
|
||
return s.acceptLoop(ctx, ln, s.obfuscated)
|
||
}
|
||
|
||
func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
|
||
ctx, cancel := context.WithCancel(ctx)
|
||
defer cancel()
|
||
|
||
// 嗅探(读首 4 字节做 HTTP/TCP 分流)的读超时必须对齐「建立 session 前」的读超时
|
||
// s.handshakeTimeout(默认 60s)——与 serveDetectedConn/serveConn 的 pre-session 读
|
||
// deadline 完全一致。一条尚未发出首帧的连接正处于「pre-session idle」状态：合法 MTProto
|
||
// 客户端(如 DrKLO)会预开「暖」连接、在有请求前并不立即发送 obfuscated2 init。此前用
|
||
// minDuration(5s,...) 把嗅探压到 5s，比非 mux 路径激进 12 倍，会把这些暖连接在 5s 误杀，
|
||
// 触发客户端 6s 重连风暴并误判「后端不健康」回退到外部 DNS。per-conn goroutine 模型已消解
|
||
// slow-loris 接入饥饿，故嗅探用满 handshakeTimeout 是安全的。
|
||
mux := newSamePortMux(ln, s.handshakeTimeout)
|
||
wsRawLn, wsHandler := transport.WebsocketListener(ln.Addr())
|
||
wsLn := newTransportPacketMessageListener(wsRawLn)
|
||
|
||
httpServer := &http.Server{
|
||
Handler: websocketRouteHandler(wsHandler, s.websocketOrigins),
|
||
ReadHeaderTimeout: minDuration(10*time.Second, s.handshakeTimeout),
|
||
BaseContext: func(net.Listener) context.Context {
|
||
return ctx
|
||
},
|
||
}
|
||
|
||
s.log.Info("Serving",
|
||
zap.String("addr", ln.Addr().String()),
|
||
zap.Int("dc", s.dc),
|
||
zap.Bool("obfuscated_tcp", s.obfuscated),
|
||
zap.Bool("websocket", true),
|
||
zap.Strings("websocket_origins", s.websocketOrigins),
|
||
)
|
||
defer s.log.Info("Stopped")
|
||
|
||
stopAll := func() {
|
||
cancel()
|
||
_ = mux.Close()
|
||
_ = httpServer.Close()
|
||
_ = wsLn.Close()
|
||
}
|
||
go func() {
|
||
<-ctx.Done()
|
||
stopAll()
|
||
}()
|
||
|
||
errCh := make(chan error, 4)
|
||
var wg sync.WaitGroup
|
||
wg.Add(4)
|
||
// 分流器：窥探前 4 字节把 HTTP(WebSocket 升级) 与裸 MTProto TCP 拆开。
|
||
go func() {
|
||
defer wg.Done()
|
||
errCh <- mux.Serve(ctx)
|
||
}()
|
||
// 裸 MTProto TCP：每条连接在自己的 goroutine 里完成去混淆 + codec 探测。
|
||
go func() {
|
||
defer wg.Done()
|
||
errCh <- s.acceptLoop(ctx, mux.TCP(), s.obfuscated)
|
||
}()
|
||
// WebSocket：gotd 升级处理器已剥离 obfuscated2 并补回 codec tag，这里只需探测 codec。
|
||
go func() {
|
||
defer wg.Done()
|
||
errCh <- s.acceptLoop(ctx, wsLn, false)
|
||
}()
|
||
go func() {
|
||
defer wg.Done()
|
||
if err := httpServer.Serve(mux.HTTP()); err != nil {
|
||
if errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
|
||
errCh <- nil
|
||
return
|
||
}
|
||
errCh <- fmt.Errorf("websocket http serve: %w", err)
|
||
return
|
||
}
|
||
errCh <- nil
|
||
}()
|
||
|
||
// 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
|
||
if err := <-errCh; err != nil {
|
||
firstErr = err
|
||
}
|
||
stopAll()
|
||
for i := 1; i < 4; i++ {
|
||
if err := <-errCh; err != nil && firstErr == nil {
|
||
firstErr = err
|
||
}
|
||
}
|
||
wg.Wait()
|
||
return firstErr
|
||
}
|
||
|
||
// acceptLoop 接受裸连接，并为每条连接单独起 goroutine 完成「去混淆 + codec 探测 +
|
||
// serveConn」。探测在 accept 循环之外、带握手超时进行——慢/半开/坏 init 的客户端只占用
|
||
// 自己的 goroutine，绝不阻塞其他连接的接入；单条连接的握手失败也只关闭该连接，不会拖垮
|
||
// 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆（裸 MTProto TCP）；WebSocket
|
||
// 连接传 false（gotd 升级处理器已完成去混淆）。
|
||
func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error {
|
||
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 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() {
|
||
defer wg.Done()
|
||
s.serveDetectedConn(ctx, raw, obfuscated)
|
||
}()
|
||
}
|
||
}
|
||
|
||
// serveDetectedConn 把一条裸连接提升为 transport.Conn（去混淆 + codec 探测）后运行 MTProto
|
||
// 连接循环。提升过程的读取放在本 goroutine、且受握手读超时约束，而非塞在 accept 循环里，
|
||
// 这样慢连接不会阻塞其他连接接入，去混淆/codec 握手本身也有时间上界。
|
||
func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated bool) {
|
||
// 握手读超时只覆盖去混淆 + codec 探测这一小段；用真实墙钟时间（SetReadDeadline 语义），
|
||
// 不走可能被测试注入的逻辑 clock。
|
||
if err := raw.SetReadDeadline(time.Now().Add(s.handshakeTimeout)); err != nil {
|
||
_ = raw.Close()
|
||
return
|
||
}
|
||
|
||
// 探测阶段若 ctx 取消，主动关闭 raw 解除阻塞读取——否则去混淆读会一直挂到握手超时，
|
||
// 把半开连接拖进优雅退出的等待里。探测结束即停掉该 watcher，连接服务期由 serveConn
|
||
// 自己的 ctx watcher 接管。
|
||
promoted := make(chan struct{})
|
||
go func() {
|
||
select {
|
||
case <-ctx.Done():
|
||
_ = raw.Close()
|
||
case <-promoted:
|
||
}
|
||
}()
|
||
|
||
conn, err := s.promoteConn(raw, obfuscated)
|
||
close(promoted)
|
||
if err != nil {
|
||
// 去混淆/codec 探测失败（读超时、客户端中途断开、坏 init 等）只影响这一条连接，
|
||
// 记 debug 即可。
|
||
if !isClientDisconnect(err) {
|
||
s.log.Debug("Transport handshake failed", zap.Error(err))
|
||
}
|
||
_ = raw.Close()
|
||
return
|
||
}
|
||
// 探测完成，撤掉握手读超时；后续每帧读写由 serveConn / 传输层各自管理超时。
|
||
if err := raw.SetReadDeadline(time.Time{}); err != nil {
|
||
_ = conn.Close()
|
||
return
|
||
}
|
||
if err := s.serveConn(ctx, conn); err != nil && !isClientDisconnect(err) {
|
||
s.log.Info("Connection closed with error", zap.Error(err))
|
||
}
|
||
}
|
||
|
||
// promoteConn 复用与 listener 组合完全一致的「obfuscated2 去混淆 + codec 探测」管线，但针对
|
||
// 单条连接，使其可在 accept 循环之外执行。obfuscated 对 WebSocket 连接必须为 false（gotd
|
||
// 升级处理器已剥离 obfuscated2 并补回 codec tag）。
|
||
func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, error) {
|
||
var ln net.Listener = newSingleConnListener(raw)
|
||
if obfuscated {
|
||
ln = transport.ObfuscatedListener(ln)
|
||
}
|
||
return newCompatTransportListener(s.codec, ln, s.frameBudget).Accept()
|
||
}
|
||
|
||
// serveConn 处理单个传输连接：读帧并按 auth_key_id 分流。
|
||
//
|
||
// - auth_key_id == 0：未加密的密钥交换起始消息，执行握手并落地 auth key。
|
||
// - auth_key_id 已注册：加密消息，解密、注册连接并分发到 RPC 路由。
|
||
// - auth_key_id 未注册：回 AuthKeyNotFound，促使客户端重新握手。
|
||
//
|
||
// 连接建立 session 后注册到 SessionManager，结束时注销。
|
||
func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error) {
|
||
transportOwner, conn := newPhysicalTransportOwner(raw)
|
||
s.metrics.ConnOpened()
|
||
s.log.Debug("Connection accepted")
|
||
|
||
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)
|
||
// Publish the terminal/RPC-cancel gates before index removal or lifecycle
|
||
// observers. Physical close then releases a writer already inside Send; the
|
||
// final Close only waits for the now-fenced actors to converge.
|
||
if current != nil {
|
||
current.beginTerminalShutdown()
|
||
}
|
||
_ = transportOwner.CloseAny()
|
||
if current != nil {
|
||
s.conns.Unregister(current)
|
||
current.Close()
|
||
}
|
||
s.metrics.ConnClosed()
|
||
s.log.Debug("Connection closed", zap.Error(err))
|
||
}()
|
||
|
||
// ctx 取消或处理结束时关闭连接，解除 Recv 阻塞。
|
||
ctx, cancel := context.WithCancel(ctx)
|
||
defer cancel()
|
||
go func() {
|
||
<-ctx.Done()
|
||
_ = transportOwner.CloseAny()
|
||
}()
|
||
|
||
cs := newConnState()
|
||
var b bin.Buffer
|
||
// plain 是本连接的复用明文缓冲：decryptClientFrame 把每帧解密进它，免去
|
||
// per-frame 整帧明文分配；帧内 slice 在下一帧读取前有效（RPC body 已在 dispatch 拷贝）。
|
||
var plain bin.Buffer
|
||
var replay *bin.Buffer
|
||
for {
|
||
if replay != nil {
|
||
b.ResetTo(replay.Buf)
|
||
replay = nil
|
||
} else {
|
||
// 建立 session 前（current==nil，握手 + 首个加密消息之前）用较短的 handshakeTimeout
|
||
// 快速回收静默的半开 / 异常连接；建立 session 后用 readTimeout（客户端有 ping 心跳）。
|
||
timeout := s.readTimeout
|
||
if current == nil || !current.isActive() {
|
||
timeout = s.handshakeTimeout
|
||
}
|
||
if err := s.recv(ctx, conn, &b, timeout); err != nil {
|
||
return err
|
||
}
|
||
if s.onFrame != nil {
|
||
s.onFrame(b.Len())
|
||
}
|
||
}
|
||
|
||
authKeyID, err := peekAuthKeyID(&b)
|
||
if err != nil {
|
||
return fmt.Errorf("peek auth key id: %w", err)
|
||
}
|
||
|
||
if authKeyID == emptyAuthKeyID {
|
||
// A physical socket may perform key exchange only before it owns an
|
||
// encrypted logical session. Mixing the direct exchange writer with an
|
||
// active outbound actor would bypass generation/write serialization.
|
||
if current != nil {
|
||
return errors.New("unencrypted exchange on established encrypted connection")
|
||
}
|
||
releaseHandshake, admitted := s.admission.tryAcquireHandshake()
|
||
if !admitted {
|
||
if err := s.sendProtoError(ctx, conn, codec.CodeTransportFlood); err != nil {
|
||
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
|
||
}
|
||
|
||
// 已建立连接复用缓存密钥走快路径（fetchedKey=nil）：避开每帧回查 AuthKeyStore——
|
||
// 这是 mtprotoedge 层最热的库访问点。密钥材料创建后不可变；销毁(destroy_auth_key)/
|
||
// 撤销由 SessionManager 主动 Close 连接保证失效，不依赖被动的“下一帧 -404”。
|
||
// 尚未进入 SessionManager 的 bad-salt provisional 会在 handleEncrypted 建立 activation claim
|
||
// 后精确复查一次，既把撤销与激活线性化，也不把 salt storm 放大成 PG 写风暴。
|
||
var fetchedKey *store.AuthKeyData
|
||
if current == nil || current.authKeyID != authKeyID {
|
||
d, found, err := s.authKeys.Get(ctx, authKeyID)
|
||
if err != nil {
|
||
return fmt.Errorf("lookup auth key: %w", err)
|
||
}
|
||
if !found {
|
||
writer := transport.Conn(conn)
|
||
if current != nil {
|
||
writer = current.transport
|
||
}
|
||
if err := s.sendProtoError(ctx, writer, codec.CodeAuthKeyNotFound); err != nil {
|
||
return err
|
||
}
|
||
// -404 对 TDesktop 是 terminal key failure；继续保留 socket 只会允许
|
||
// 同一客户端反复触发 AuthKeyStore 查询。回包一次后立即断开。
|
||
return nil
|
||
}
|
||
fetchedKey = &d
|
||
}
|
||
|
||
current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b, &plain)
|
||
if errors.Is(err, errActivationAuthKeyRejected) {
|
||
// handleEncrypted writes -404 while its activation claim still owns the
|
||
// physical writer, then its deferred abort removes/closes the claim.
|
||
return nil
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
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
|
||
}
|
||
}
|
||
|
||
// deadlineReceiver 是可选的直管读超时接口：telesrv-owned compat transport 实现它，
|
||
// 让每帧读只做一次 SetReadDeadline，不再分配 per-frame context timer。ctx 取消仍由
|
||
// serveConn 的 watcher 关闭底层连接来解除阻塞读（与 ctx deadline 路径行为一致）。
|
||
type deadlineReceiver interface {
|
||
RecvDeadline(deadline time.Time, b *bin.Buffer) error
|
||
}
|
||
|
||
func (s *Server) recv(ctx context.Context, conn transport.Conn, b *bin.Buffer, timeout time.Duration) error {
|
||
b.Reset()
|
||
if dr, ok := conn.(deadlineReceiver); ok {
|
||
return dr.RecvDeadline(time.Now().Add(timeout), b)
|
||
}
|
||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||
defer cancel()
|
||
return conn.Recv(ctx, b)
|
||
}
|
||
|
||
// isClientDisconnect 判断错误是否为正常的客户端断开/服务关闭，不应作为异常记录。
|
||
func isClientDisconnect(err error) bool {
|
||
switch {
|
||
case errors.Is(err, io.EOF),
|
||
errors.Is(err, net.ErrClosed),
|
||
errors.Is(err, context.Canceled),
|
||
errors.Is(err, context.DeadlineExceeded):
|
||
return true
|
||
}
|
||
var nerr *net.OpError
|
||
if errors.As(err, &nerr) && (nerr.Op == "read" || nerr.Op == "write") {
|
||
return true
|
||
}
|
||
return false
|
||
}
|