diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index 6eced5e6..53e4ce6d 100644 --- a/internal/mtprotoedge/conn.go +++ b/internal/mtprotoedge/conn.go @@ -1,7 +1,9 @@ package mtprotoedge import ( + "bufio" "context" + "encoding/hex" "sync" "sync/atomic" "time" @@ -32,9 +34,12 @@ type Conn struct { metrics Metrics authKeyID [8]byte - sessionID int64 - salt int64 - key crypto.AuthKey + // authKeyHex 是 authKeyID 的 hex 缓存:每条 RPC 的结构化日志都会带它, + // 建连时算一次,避免热路径反复 hex 编码分配。 + authKeyHex string + sessionID int64 + salt int64 + key crypto.AuthKey outbound chan outboundOp outboundControl chan outboundOp @@ -62,9 +67,13 @@ type Conn struct { // outboundPlain/outboundWire 只由 outbound actor 访问,用于复用出站加密缓冲。 outboundPlain bin.Buffer outboundWire bin.Buffer + // outboundRand 只由 outbound actor 访问:对 cipher 随机源的缓冲预读, + // 把每帧 padding 的 getrandom syscall 摊薄成 ~1KiB 一次。 + outboundRand *bufio.Reader identityMu sync.RWMutex businessAuthKeyID [8]byte + businessAuthKeyHex string businessAuthKeyResolved bool userID atomic.Int64 userIDResolved atomic.Bool @@ -74,6 +83,13 @@ type Conn struct { // 同步失败时保持 false,让置位短路放行、下一条 RPC 重试同步,避免 // 「已置位但 channel 路由缺失」的 session 静默漏收超级群推送。 membershipsSynced atomic.Bool + // membershipGen 是本连接 channel membership 索引的修订号:任何增量修订 + // (join/leave/kick 的 Add/Remove、身份切换/下线的整体清除)都递增。全量同步方 + // 在读取持久成员列表前采样、落地时带回比对,检测「读取窗口内发生增量修订」的 + // 丢失更新竞态(SetSessionChannelMemberships 改走合并路径并保持未就绪重试)。 + membershipGen atomic.Int64 + // createdAt 是连接建立时刻,供同 auth_key session 数触顶时驱逐真正最旧的连接。 + createdAt time.Time // keyDestroyed 标记本连接的 auth_key 已被 destroy_auth_key 删除。serveConn 对已建立 // 连接复用缓存密钥跳过每帧 AuthKeyStore 回查;置位后强制回落到 Get→AuthKeyNotFound, // 维持「destroy_auth_key 发起连接下一帧自然失效」契约。只由 destroy_auth_key 处理器置位。 @@ -112,10 +128,20 @@ func (c *Conn) BusinessAuthKeyID() ([8]byte, bool) { return c.businessAuthKeyID, c.businessAuthKeyResolved } +// BusinessAuthKeyHex 返回业务视角 auth_key_id 的 hex 缓存(每 RPC 日志用,免重复编码)。 +func (c *Conn) BusinessAuthKeyHex() (string, bool) { + c.identityMu.RLock() + defer c.identityMu.RUnlock() + return c.businessAuthKeyHex, c.businessAuthKeyResolved +} + // SetBusinessAuthKeyID 缓存业务视角 auth_key_id。 func (c *Conn) SetBusinessAuthKeyID(id [8]byte) { c.identityMu.Lock() changed := !c.businessAuthKeyResolved || c.businessAuthKeyID != id + if changed || c.businessAuthKeyHex == "" { + c.businessAuthKeyHex = hex.EncodeToString(id[:]) + } c.businessAuthKeyID = id c.businessAuthKeyResolved = true c.identityMu.Unlock() diff --git a/internal/mtprotoedge/decrypt.go b/internal/mtprotoedge/decrypt.go new file mode 100644 index 00000000..0fe93e9d --- /dev/null +++ b/internal/mtprotoedge/decrypt.go @@ -0,0 +1,93 @@ +package mtprotoedge + +import ( + "crypto/aes" + "encoding/binary" + "errors" + "fmt" + + "github.com/gotd/ige" + + "github.com/gotd/td/bin" + "github.com/gotd/td/crypto" +) + +// clientFrame 是解密后的单帧客户端消息视图。data/plaintext 引用调用方持有的复用明文 +// 缓冲,仅在解密下一帧前有效;需要跨帧保留的字节必须拷贝(dispatch 对 RPC body 已 +// b.Copy(),container/gzip 均在当前帧内同步消费)。 +type clientFrame struct { + salt int64 + sessionID int64 + messageID int64 + seqNo int32 + // data 是去掉 32 字节头与 padding 后的消息体。 + data []byte + // plaintext 是完整明文(头 + 数据 + padding),quick ack token 直接对它做 SHA256, + // 免去 gotd 路径上「为算 token 把解密结果整帧重编码一遍」的拷贝。 + plaintext []byte +} + +// encryptedFrameHeaderLen 是 MTProto 2.0 明文头长度:salt(8)+session_id(8)+msg_id(8)+seq_no(4)+len(4)。 +const encryptedFrameHeaderLen = 32 + +// decryptClientFrame 把一帧客户端加密消息解密进 plain 复用缓冲(telesrv-owned, +// 与出站 encryptOutboundFrame 对称)。校验集合与 gotd Cipher.Decrypt 逐条一致: +// auth_key_id 匹配、密文 16 字节块对齐、msg_key(SHA256,client side x=0)、明文头 +// 长度、message_data_len 非负 / 4 字节对齐 / 不越界、padding ≤ 1024。区别只在 +// 明文缓冲复用与免结构体堆分配——gotd DecryptFromBuffer 每帧 make 整帧明文 +// (上传帧可达 512KiB+),是入站 Dispatch 链路最大的稳态分配点。 +func decryptClientFrame(key crypto.AuthKey, b *bin.Buffer, plain *bin.Buffer) (clientFrame, error) { + buf := b.Buf + if len(buf) < 24 { + return clientFrame{}, errors.New("encrypted message is too short") + } + var authKeyID [8]byte + copy(authKeyID[:], buf[:8]) + if authKeyID != key.ID { + return clientFrame{}, errors.New("unknown auth key id") + } + var msgKey bin.Int128 + copy(msgKey[:], buf[8:24]) + encrypted := buf[24:] + if len(encrypted) == 0 || len(encrypted)%16 != 0 { + return clientFrame{}, errors.New("invalid encrypted data padding") + } + + aesKey, iv := crypto.Keys(key.Value, msgKey, crypto.Client) + aesBlock, err := aes.NewCipher(aesKey[:]) + if err != nil { + return clientFrame{}, err + } + ensureBinBufferLen(plain, len(encrypted)) + ige.DecryptBlocks(aesBlock, iv[:], plain.Buf, encrypted) + plaintext := plain.Buf + + if crypto.MessageKey(key.Value, plaintext, crypto.Client) != msgKey { + return clientFrame{}, errors.New("msg_key is invalid") + } + if len(plaintext) < encryptedFrameHeaderLen { + return clientFrame{}, errors.New("message data is too short") + } + + dataLen := int(int32(binary.LittleEndian.Uint32(plaintext[28:32]))) + withPadding := plaintext[encryptedFrameHeaderLen:] + switch { + case dataLen < 0: + return clientFrame{}, fmt.Errorf("message length is invalid: %d less than zero", dataLen) + case dataLen%4 != 0: + return clientFrame{}, fmt.Errorf("message length is invalid: %d is not divisible by 4", dataLen) + case dataLen > len(withPadding): + return clientFrame{}, fmt.Errorf("message length %d is bigger than data length %d", dataLen, len(withPadding)) + case len(withPadding)-dataLen > 1024: + return clientFrame{}, fmt.Errorf("padding %d of message is too big", len(withPadding)-dataLen) + } + + return clientFrame{ + salt: int64(binary.LittleEndian.Uint64(plaintext[0:8])), + sessionID: int64(binary.LittleEndian.Uint64(plaintext[8:16])), + messageID: int64(binary.LittleEndian.Uint64(plaintext[16:24])), + seqNo: int32(binary.LittleEndian.Uint32(plaintext[24:28])), + data: withPadding[:dataLen], + plaintext: plaintext, + }, nil +} diff --git a/internal/mtprotoedge/decrypt_test.go b/internal/mtprotoedge/decrypt_test.go new file mode 100644 index 00000000..3eabfab6 --- /dev/null +++ b/internal/mtprotoedge/decrypt_test.go @@ -0,0 +1,181 @@ +package mtprotoedge + +import ( + "bytes" + "crypto/aes" + "crypto/rand" + "encoding/binary" + "testing" + + "github.com/gotd/ige" + + "github.com/gotd/td/bin" + "github.com/gotd/td/crypto" +) + +func newTestAuthKey(t *testing.T) crypto.AuthKey { + t.Helper() + var key crypto.Key + if _, err := rand.Read(key[:]); err != nil { + t.Fatalf("rand: %v", err) + } + return key.WithID() +} + +// encryptRawClientPlaintext 用 client side(x=0)把一段已含 32 字节头与 padding 的 +// 原始明文加密成完整入站帧,用于构造 gotd Cipher.Encrypt 不允许生成的畸形明文。 +func encryptRawClientPlaintext(t *testing.T, key crypto.AuthKey, plaintext []byte) *bin.Buffer { + t.Helper() + if len(plaintext)%16 != 0 { + t.Fatalf("plaintext must be 16-aligned, got %d", len(plaintext)) + } + msgKey := crypto.MessageKey(key.Value, plaintext, crypto.Client) + aesKey, iv := crypto.Keys(key.Value, msgKey, crypto.Client) + blk, err := aes.NewCipher(aesKey[:]) + if err != nil { + t.Fatalf("aes: %v", err) + } + encrypted := make([]byte, len(plaintext)) + ige.EncryptBlocks(blk, iv[:], encrypted, plaintext) + + var b bin.Buffer + b.Put(key.ID[:]) + b.Put(msgKey[:]) + b.Put(encrypted) + return &b +} + +// buildRawPlaintext 构造 salt/session/msg_id/seq_no + dataLen 头与指定 data/padding 的明文。 +// dataLen 允许与真实 data 长度不一致,用于打边界。 +func buildRawPlaintext(salt, sessionID, msgID int64, seqNo, dataLen int32, data, padding []byte) []byte { + out := make([]byte, 0, 32+len(data)+len(padding)) + var hdr [32]byte + binary.LittleEndian.PutUint64(hdr[0:8], uint64(salt)) + binary.LittleEndian.PutUint64(hdr[8:16], uint64(sessionID)) + binary.LittleEndian.PutUint64(hdr[16:24], uint64(msgID)) + binary.LittleEndian.PutUint32(hdr[24:28], uint32(seqNo)) + binary.LittleEndian.PutUint32(hdr[28:32], uint32(dataLen)) + out = append(out, hdr[:]...) + out = append(out, data...) + return append(out, padding...) +} + +// TestDecryptClientFrameParityWithGotd 逐字节对照 telesrv 自建解密与 gotd server cipher: +// 同一帧要么双方都接受且字段/数据一致,要么双方都拒绝。覆盖正常帧、畸形长度、 +// 超限 padding、篡改 msg_key/auth_key_id、截断帧。 +func TestDecryptClientFrameParityWithGotd(t *testing.T) { + key := newTestAuthKey(t) + serverCipher := crypto.NewServerCipher(rand.Reader) + + pad := func(n int) []byte { + p := make([]byte, n) + if _, err := rand.Read(p); err != nil { + t.Fatalf("rand: %v", err) + } + return p + } + data16 := pad(16) + + cases := []struct { + name string + frame *bin.Buffer + }{ + {"valid_small", encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7000, 1, 16, data16, pad(16)))}, + {"valid_large", encryptRawClientPlaintext(t, key, buildRawPlaintext(9, 8, 7002, 3, 4096, pad(4096), pad(16)))}, + {"zero_len_data", encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7004, 5, 0, nil, pad(16)))}, + {"data_len_negative", encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7006, 7, -4, data16, pad(16)))}, + {"data_len_unaligned", encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7008, 9, 6, data16, pad(16)))}, + {"data_len_overflow", encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7010, 11, 64, data16, pad(16)))}, + {"padding_too_big", encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7012, 13, 16, data16, pad(1040)))}, + {"plaintext_only_header_block", encryptRawClientPlaintext(t, key, pad(16))}, + } + + // 篡改 msg_key。 + tampered := encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7014, 15, 16, data16, pad(16))) + tampered.Buf[8] ^= 0xff + cases = append(cases, struct { + name string + frame *bin.Buffer + }{"tampered_msg_key", tampered}) + + // 错误 auth_key_id。 + wrongKey := encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7016, 17, 16, data16, pad(16))) + wrongKey.Buf[0] ^= 0xff + cases = append(cases, struct { + name string + frame *bin.Buffer + }{"wrong_auth_key_id", wrongKey}) + + // 截断帧。 + cases = append(cases, + struct { + name string + frame *bin.Buffer + }{"truncated_header", &bin.Buffer{Buf: pad(16)}}, + struct { + name string + frame *bin.Buffer + }{"unaligned_ciphertext", &bin.Buffer{Buf: pad(24 + 15)}}, + ) + + var plain bin.Buffer + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + gotdData, gotdErr := serverCipher.DecryptFromBuffer(key, &bin.Buffer{Buf: append([]byte(nil), tc.frame.Buf...)}) + frame, ourErr := decryptClientFrame(key, &bin.Buffer{Buf: append([]byte(nil), tc.frame.Buf...)}, &plain) + + if (gotdErr == nil) != (ourErr == nil) { + t.Fatalf("accept/reject mismatch: gotd err=%v, ours err=%v", gotdErr, ourErr) + } + if gotdErr != nil { + return + } + if frame.salt != gotdData.Salt || frame.sessionID != gotdData.SessionID || + frame.messageID != gotdData.MessageID || frame.seqNo != gotdData.SeqNo { + t.Fatalf("header mismatch: ours=%+v gotd salt=%d session=%d msg=%d seq=%d", + frame, gotdData.Salt, gotdData.SessionID, gotdData.MessageID, gotdData.SeqNo) + } + if !bytes.Equal(frame.data, gotdData.Data()) { + t.Fatalf("data mismatch: ours %d bytes, gotd %d bytes", len(frame.data), len(gotdData.Data())) + } + }) + } +} + +// TestDecryptClientFrameReusesPlainBuffer 验证同一 plain 缓冲跨帧复用:先大帧后小帧, +// 解密结果仍正确且不受前一帧残留字节影响。 +func TestDecryptClientFrameReusesPlainBuffer(t *testing.T) { + key := newTestAuthKey(t) + + big := make([]byte, 2048) + for i := range big { + big[i] = byte(i) + } + small := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16} + + pad := make([]byte, 16) + frame1 := encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7100, 1, int32(len(big)), big, pad)) + frame2 := encryptRawClientPlaintext(t, key, buildRawPlaintext(1, 2, 7102, 3, int32(len(small)), small, pad)) + + var plain bin.Buffer + f1, err := decryptClientFrame(key, frame1, &plain) + if err != nil { + t.Fatalf("decrypt big frame: %v", err) + } + if !bytes.Equal(f1.data, big) { + t.Fatal("big frame data mismatch") + } + f2, err := decryptClientFrame(key, frame2, &plain) + if err != nil { + t.Fatalf("decrypt small frame: %v", err) + } + if !bytes.Equal(f2.data, small) { + t.Fatal("small frame data mismatch after buffer reuse") + } + if f2.messageID != 7102 || f2.seqNo != 3 || f2.salt != 1 || f2.sessionID != 2 { + t.Fatalf("small frame header mismatch: %+v", f2) + } + if len(f2.plaintext) != 32+len(small)+len(pad) { + t.Fatalf("plaintext length not shrunk on reuse: %d", len(f2.plaintext)) + } +} diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index 207f0d0c..e514fc49 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -4,7 +4,6 @@ import ( "context" "crypto/sha256" "encoding/binary" - "encoding/hex" "errors" "fmt" "io" @@ -33,6 +32,11 @@ type connState struct { order []int64 minSeen int64 maxSeen int64 + // maxContentMsgID/maxContentSeqNo 是已接受 content 消息的 msg_id / seq_no 高水位, + // 供 validateSeq 的 O(1) 快路径使用(客户端正常发送严格递增)。二者只增不减、 + // 不随 seen 淘汰回退——快路径只接受「全扫描也必然接受」的子集,其余回落全扫描。 + maxContentMsgID int64 + maxContentSeqNo int32 } type clientMsgRecord struct { @@ -76,7 +80,8 @@ const ( // fetchedKey 非 nil 表示本帧的 auth key 是刚从 AuthKeyStore 查出的(首帧/换 auth key/被销毁 // 后回落);为 nil 表示走快路径——serveConn 判定 current 仍持同一未销毁的 auth key,直接复用 // current.key/current.salt 解密,既不回查 AuthKeyStore 也不重建 store.AuthKeyData。 -func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b *bin.Buffer) (*Conn, error) { +// plain 是 serveConn 持有的复用明文缓冲,frame 的 slice 仅在下一帧解密前有效。 +func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b, plain *bin.Buffer) (*Conn, error) { var key crypto.AuthKey var serverSalt int64 if fetchedKey != nil { @@ -88,19 +93,19 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con serverSalt = current.salt } - data, err := s.cipher.DecryptFromBuffer(key, b) + frame, err := decryptClientFrame(key, b, plain) if err != nil { return current, fmt.Errorf("decrypt: %w", err) } - if data.Salt != serverSalt { + if frame.salt != serverSalt { c := current temp := false - if c == nil || c.sessionID != data.SessionID { - c = s.newConn(tc, key, data.SessionID, serverSalt) + if c == nil || c.sessionID != frame.sessionID { + c = s.newConn(tc, key, frame.sessionID, serverSalt) temp = true } - err := s.sendBadServerSalt(ctx, c, data.MessageID, data.SeqNo, serverSalt) + err := s.sendBadServerSalt(ctx, c, frame.messageID, frame.seqNo, serverSalt) if temp { c.Close() } @@ -108,7 +113,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con } // 首个加密消息或 session 变化时(重新)注册连接到 SessionManager。 - if current == nil || current.sessionID != data.SessionID { + if current == nil || current.sessionID != frame.sessionID { if current != nil { cs.reset() } @@ -116,62 +121,71 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con s.conns.Unregister(current) current.Close() } - current = s.newConn(tc, key, data.SessionID, serverSalt) + current = s.newConn(tc, key, frame.sessionID, serverSalt) + // 注册即播种协商 layer:新 Conn 的 clientLayer 为 0(=canonical 227),若等到 + // 首条 RPC 的 Dispatch 返回后才刷新,重连老客户端在首条 RPC handler 执行期间 + // 收到的 pending flush / 并发 push 会漏降级。进程内重连时 rpc 层留有 + // (auth_key, session) / auth_key 两级协商记录,这里一次查询即可闭合该空窗。 + if s.rpc != nil { + if layer, ok := s.rpc.NegotiatedLayer(current.authKeyID, current.sessionID); ok { + current.SetClientLayer(layer) + } + } s.conns.Register(current) } - s.maybePersistSession(ctx, current, data.SessionID, key.ID, serverSalt) + s.maybePersistSession(ctx, current, frame.sessionID, key.ID, serverSalt) - body := data.Data() + body := frame.data typeID, err := (&bin.Buffer{Buf: body}).PeekID() if err != nil { return current, fmt.Errorf("peek encrypted payload type id: %w", err) } - if code := validateClientEnvelope(s.clock.Now(), data.MessageID, data.SeqNo, typeID); code != 0 { + if code := validateClientEnvelope(s.clock.Now(), frame.messageID, frame.seqNo, typeID); code != 0 { s.log.Debug("Sending bad_msg_notification", - zap.Int64("msg_id", data.MessageID), - zap.Int32("seq_no", data.SeqNo), + zap.Int64("msg_id", frame.messageID), + zap.Int32("seq_no", frame.seqNo), zap.Uint32("type_id", typeID), zap.Int("code", code), ) - return current, s.sendBadMsg(ctx, current, data.MessageID, data.SeqNo, code) + return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code) } - if err := sendQuickAckIfRequested(ctx, tc, key, data); err != nil { + if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext); err != nil { return current, err } content := clientMessageNeedsAck(typeID) - if record, ok := cs.seenRecord(data.MessageID); ok { - s.log.Debug("Duplicate msg_id; replay cached result if available", zap.Int64("msg_id", data.MessageID)) - if err := s.replayRPCResultByRequest(ctx, current, data.MessageID); err != nil { + if record, ok := cs.seenRecord(frame.messageID); ok { + s.log.Debug("Duplicate msg_id; replay cached result if available", zap.Int64("msg_id", frame.messageID)) + if err := s.replayRPCResultByRequest(ctx, current, frame.messageID); err != nil { return current, err } if !record.content { return current, nil } - return current, s.sendAck(ctx, current, data.MessageID) + return current, s.sendAck(ctx, current, frame.messageID) } - if code := cs.validateSeq(data.MessageID, data.SeqNo, content); code != 0 { + if code := cs.validateSeq(frame.messageID, frame.seqNo, content); code != 0 { s.log.Debug("Sending bad_msg_notification", - zap.Int64("msg_id", data.MessageID), - zap.Int32("seq_no", data.SeqNo), + zap.Int64("msg_id", frame.messageID), + zap.Int32("seq_no", frame.seqNo), zap.Uint32("type_id", typeID), zap.Int("code", code), ) - return current, s.sendBadMsg(ctx, current, data.MessageID, data.SeqNo, code) + return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code) } - cs.track(data.MessageID, data.SeqNo, content, msgStateReceived) + cs.track(frame.messageID, frame.seqNo, content, msgStateReceived) if !cs.sentCreated { cs.sentCreated = true - s.log.Debug("Sending new_session_created", zap.Int64("msg_id", data.MessageID), zap.Int32("seq_no", data.SeqNo)) - if err := s.sendNewSessionCreated(ctx, current, data.MessageID); err != nil { + s.log.Debug("Sending new_session_created", zap.Int64("msg_id", frame.messageID), zap.Int32("seq_no", frame.seqNo)) + if err := s.sendNewSessionCreated(ctx, current, frame.messageID); err != nil { return current, err } } var acks []int64 - if err := s.dispatch(ctx, cs, current, data.MessageID, data.SeqNo, &bin.Buffer{Buf: body}, &acks); err != nil { + if err := s.dispatch(ctx, cs, current, frame.messageID, frame.seqNo, &bin.Buffer{Buf: body}, &acks); err != nil { return current, err } if len(acks) > 0 { @@ -210,28 +224,23 @@ func (s *Server) maybePersistSession(ctx context.Context, c *Conn, sessionID int } } -func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, data *crypto.EncryptedMessageData) error { +func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, plaintext []byte) error { q, ok := tc.(quickAckTransport) if !ok || !q.ConsumeQuickAckRequested() { return nil } - token, err := clientQuickAckToken(key, data) - if err != nil { - return err - } - return q.SendQuickAck(ctx, token) + return q.SendQuickAck(ctx, clientQuickAckToken(key, plaintext)) } -func clientQuickAckToken(key crypto.AuthKey, data *crypto.EncryptedMessageData) (uint32, error) { - var plain bin.Buffer - if err := data.Encode(&plain); err != nil { - return 0, err - } +// clientQuickAckToken 按 Android MTProto v2 公式计算 quick ack:SHA256(auth_key[88:120] + +// 完整明文)[:4]。plaintext 直接来自解密复用缓冲(decryptClientFrame.plaintext), +// 与旧实现「把解密结果重编码一遍再哈希」字节一致但零拷贝。 +func clientQuickAckToken(key crypto.AuthKey, plaintext []byte) uint32 { h := sha256.New() _, _ = h.Write(key.Value[88:120]) - _, _ = h.Write(plain.Raw()) + _, _ = h.Write(plaintext) sum := h.Sum(nil) - return binary.LittleEndian.Uint32(sum[:4]) &^ quickAckResponseFlag, nil + return binary.LittleEndian.Uint32(sum[:4]) &^ quickAckResponseFlag } // dispatch 处理一条明文消息:解包 container/gzip,处理服务消息,其余转 RPC 路由。 @@ -397,14 +406,14 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int return err } ackContent() - s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:]))) + s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", c.authKeyHex)) // 真正销毁:删密钥库记录(每帧回查,删除后该 key 的入站帧立即失效)并主动 // 断开同 key 的其他连接——出站推送用连接持有的密钥副本加密、不回查密钥库, // 不断开的话被销毁 key 的空闲连接仍能持续收到推送。发起连接除外:响应要 // 先送达,它的下一帧会因密钥缺失自然断开。授权(authorizations)不在此清理, // destroy_auth_key 是 PFS 密钥轮换的清理动作,不等于登出。 if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil { - s.log.Warn("Delete auth key failed", zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), zap.Error(err)) + s.log.Warn("Delete auth key failed", zap.String("auth_key_id", c.authKeyHex), zap.Error(err)) return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{}) } // 标记密钥已销毁:发起连接被 CloseSessionsForRawAuthKeyExcept 排除(响应需先送达), @@ -416,7 +425,7 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int default: ackContent() body := b.Copy() - return s.enqueueRPC(ctx, c, msgID, body) + return s.enqueueRPC(ctx, c, msgID, id, body) } } @@ -437,14 +446,15 @@ func mergeStateInfo(primary, fallback []byte) []byte { return info } -func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, body []byte) error { - id, _ := (&bin.Buffer{Buf: body}).PeekID() - method := s.typeName(id) +// 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 { + method := s.typeName(typeID) if cached, ok := s.cachedRPCResult(c, msgID); ok { s.log.Info("RPC duplicate replay from session cache", zap.String("method", method), zap.Int64("msg_id", msgID), - zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), ) return c.SendEncoded(ctx, proto.MessageServerResponse, cached) @@ -455,10 +465,10 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, body []by run: func(taskCtx context.Context) error { // body 已是 enqueueRPC 入参的独立副本(dispatch 里 b.Copy()),且每个任务只 run 一次, // 无需再 append 拷贝;直接复用,省掉一份 inbound 在途内存。 - if err := s.handleRPC(taskCtx, c, msgID, &bin.Buffer{Buf: body}); err != nil { + if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}); err != nil { fields := []zap.Field{ zap.Int64("msg_id", msgID), - zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), zap.Error(err), } @@ -476,7 +486,7 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, body []by s.log.Debug("Inbound RPC queue full", zap.String("method", method), zap.Int64("msg_id", msgID), - zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), ) return s.sendResult(ctx, c, msgID, &mt.RPCError{ @@ -488,9 +498,7 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, body []by } // handleRPC 把明文 RPC 请求交给 RPC 路由,并将结果或错误包成 rpc_result 回发。 -func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buffer) error { - id, _ := b.PeekID() - method := s.typeName(id) +func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer) error { if s.rpc == nil { s.log.Warn("No RPC handler configured; dropping request", zap.String("method", method)) return nil @@ -510,15 +518,16 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, b *bin.Buf c.SetClientLayer(layer) } - fields := []zap.Field{ + fields := make([]zap.Field, 0, 12) + fields = append(fields, zap.String("method", method), - zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), zap.Int64("msg_id", msgID), zap.Duration("dur", dur), - } - if businessAuthKeyID, ok := c.BusinessAuthKeyID(); ok { - fields = append(fields, zap.String("business_auth_key_id", hex.EncodeToString(businessAuthKeyID[:]))) + ) + if businessAuthKeyHex, ok := c.BusinessAuthKeyHex(); ok { + fields = append(fields, zap.String("business_auth_key_id", businessAuthKeyHex)) } if userID := c.UserID(); userID != 0 { fields = append(fields, zap.Int64("user_id", userID)) @@ -566,31 +575,37 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result return c.SendEncoded(ctx, proto.MessageServerResponse, encoded) } -// encodeRPCResult 编码 rpc_result。proto.Result.Result 是裸 boxed 对象字节,故在包入 -// rpc_result 之前对其按连接协商 layer 降级(layer==227 直通,零开销)。降级失败 fail-safe: -// 记日志并发送 canonical 字节——宁可老客户端对个别长尾对象渲染异常,也不让连接/流崩。 +// encodeRPCResult 编码 rpc_result。内层对象与 rpc_result 头(type_id + req_msg_id) +// 一次性编码进同一 buffer——旧实现先编码内层、再经 proto.Result.Encode 整体拷贝一遍, +// 每条响应多一份全量 body 拷贝。内层按连接协商 layer 降级(layer==227 直通,零开销), +// 降级改写字节时才重建整条消息。降级失败 fail-safe:记日志并发送 canonical 字节—— +// 宁可老客户端对个别长尾对象渲染异常,也不让连接/流崩。 func (s *Server) encodeRPCResult(c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) { + const headerLen = 4 + 8 // rpc_result#f35c6d01 type_id + req_msg_id var buf bin.Buffer + buf.PutID(proto.ResultTypeID) + buf.PutLong(reqMsgID) if err := result.Encode(&buf); err != nil { return nil, fmt.Errorf("encode rpc result: %w", err) } - inner := buf.Raw() if layer := c.ClientLayer(); layer < layerwire.CanonicalLayer { + inner := buf.Buf[headerLen:] if down, err := layerwire.Transcode(inner, layer); err != nil { s.log.Warn("layerwire downgrade failed; sending canonical rpc_result", zap.Int("layer", layer), zap.Int64("req_msg_id", reqMsgID), zap.Error(err)) - } else { - inner = down + } else if !sameBacking(down, inner) { + var rebuilt bin.Buffer + rebuilt.PutID(proto.ResultTypeID) + rebuilt.PutLong(reqMsgID) + rebuilt.Put(down) + buf = rebuilt } } - encoded, err := encodeOutboundMessage(&proto.Result{ - RequestMessageID: reqMsgID, - Result: inner, - }) - if err != nil { - return nil, err - } - return encoded, nil + return &encodedOutboundMessage{ + typeID: proto.ResultTypeID, + body: buf.Raw(), + reqMsgID: reqMsgID, + }, nil } func (s *Server) cachedRPCResult(c *Conn, reqMsgID int64) (*encodedOutboundMessage, bool) { @@ -694,7 +709,7 @@ func (s *Server) sendDestroySession(ctx context.Context, c *Conn, sessionID int6 removed = s.conns.DestroySessionForAuthKey(c.authKeyID, sessionID) if err := s.sessions.Delete(ctx, sessionID); err != nil { s.log.Debug("Delete session record failed", - zap.String("auth_key_id", hex.EncodeToString(c.authKeyID[:])), + zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", sessionID), zap.Error(err), ) @@ -834,6 +849,11 @@ func (cs *connState) validateSeq(msgID int64, seqNo int32, content bool) int { if !content { return 0 } + // 快路径:msg_id 与 seq_no 都严格高于已接受 content 高水位时,任何已见记录都不可能 + // 与本条构成 too_low/too_high 反转,免去 O(len(seen)) 全扫描(正常客户端恒命中)。 + if msgID > cs.maxContentMsgID && seqNo > cs.maxContentSeqNo { + return 0 + } for seenMsgID, record := range cs.seen { if !record.content { continue @@ -854,6 +874,14 @@ func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) { seqNo: seqNo, content: content, } + if content { + if msgID > cs.maxContentMsgID { + cs.maxContentMsgID = msgID + } + if seqNo > cs.maxContentSeqNo { + cs.maxContentSeqNo = seqNo + } + } cs.order = append(cs.order, msgID) if msgID < cs.minSeen { cs.minSeen = msgID diff --git a/internal/mtprotoedge/inbound_rpc.go b/internal/mtprotoedge/inbound_rpc.go index b2c6bd9b..15d05ece 100644 --- a/internal/mtprotoedge/inbound_rpc.go +++ b/internal/mtprotoedge/inbound_rpc.go @@ -136,15 +136,17 @@ func (c *Conn) runInboundRPC(rootCtx context.Context, task inboundRPC) { if ctx == nil { ctx = context.Background() } - ctx, cancel := context.WithCancel(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) + } defer cancel() stopRoot := context.AfterFunc(rootCtx, cancel) defer stopRoot() - if c.rpcTimeout > 0 { - var timeoutCancel context.CancelFunc - ctx, timeoutCancel = context.WithTimeout(ctx, c.rpcTimeout) - defer timeoutCancel() - } _ = task.run(ctx) } diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index a748cbfe..53359c24 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -1,6 +1,7 @@ package mtprotoedge import ( + "bufio" "context" "crypto/aes" "errors" @@ -165,34 +166,40 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. encoded: encoded, enqueuedAt: time.Now(), } + // 快路径:非阻塞入队。fan-out 每 (conn × push) 都走这里,队列有空位时不为 + // 本次推送分配任何 timer(此前 timeout>0 无条件 WithTimeout,稳态白建 timer)。 + select { + case c.outbound <- op: + return nil + case <-c.outboundStop: + return ErrConnClosed + default: + } if timeout == 0 { - select { - case c.outbound <- op: - return nil - case <-c.outboundStop: - return ErrConnClosed - default: - c.metrics.OutboundDropped("push_queue_full") - return ErrOutboundQueueFull - } + c.metrics.OutboundDropped("push_queue_full") + return ErrOutboundQueueFull } - enqueueCtx := ctx - if enqueueCtx == nil { - enqueueCtx = context.Background() + c.metrics.OutboundQueueWait(len(c.outbound), cap(c.outbound)) + if ctx == nil { + ctx = context.Background() } - var cancel context.CancelFunc + var timeoutC <-chan time.Time if timeout > 0 { - enqueueCtx, cancel = context.WithTimeout(enqueueCtx, timeout) - defer cancel() + timer := time.NewTimer(timeout) + defer timer.Stop() + timeoutC = timer.C } - if err := c.enqueueOutbound(enqueueCtx, op); err != nil { - if errors.Is(err, context.DeadlineExceeded) && timeout > 0 { - c.metrics.OutboundDropped("push_queue_timeout") - return ErrOutboundQueueFull - } - return err + select { + case c.outbound <- op: + return nil + case <-timeoutC: + c.metrics.OutboundDropped("push_queue_timeout") + return ErrOutboundQueueFull + case <-ctx.Done(): + return ctx.Err() + case <-c.outboundStop: + return ErrConnClosed } - return nil } func (c *Conn) send(ctx context.Context, t proto.MessageType, msg bin.Encoder, control bool) error { @@ -600,6 +607,13 @@ func (c *Conn) commitContentSeqNo() { c.sentContentMessages++ } +// deadlineOutboundWriter 是可选的直管写超时接口:telesrv-owned compat transport 实现它, +// 让 outbound actor 每帧只做一次 SetWriteDeadline,不再为写超时分配 context timer +// (gotd transport.Conn 的 Send 本身也只消费 ctx.Deadline,不监听 ctx.Done,语义等价)。 +type deadlineOutboundWriter interface { + SendDeadline(deadline time.Time, b *bin.Buffer) error +} + func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { if ctx == nil { ctx = context.Background() @@ -609,17 +623,30 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { return fmt.Errorf("encrypt: %w", err) } - sendCtx := ctx - cancel := func() {} - if c.writeTimeout > 0 { - sendCtx, cancel = context.WithTimeout(ctx, c.writeTimeout) - } - defer cancel() writer := c.writer if writer == nil { writer = c.transport } - if err := writer.Send(sendCtx, out); err != nil { + var deadline time.Time + if c.writeTimeout > 0 { + deadline = time.Now().Add(c.writeTimeout) + } + if d, ok := ctx.Deadline(); ok && (deadline.IsZero() || d.Before(deadline)) { + deadline = d + } + if dw, ok := writer.(deadlineOutboundWriter); ok { + err = dw.SendDeadline(deadline, out) + } else { + // 回落路径:gotd full codec / 测试注入 codec 仍走 ctx deadline。 + sendCtx := ctx + cancel := func() {} + if !deadline.IsZero() { + sendCtx, cancel = context.WithDeadline(ctx, deadline) + } + err = writer.Send(sendCtx, out) + cancel() + } + if err != nil { return fmt.Errorf("send: %w", err) } if frame.sentAt.IsZero() { @@ -642,7 +669,13 @@ func (c *Conn) encryptOutboundFrame(frame *outboundFrame) (*bin.Buffer, error) { paddingOffset := plain.Len() paddingLen := encryptedPaddingLen(paddingOffset) growBinBufferLen(plain, paddingOffset+paddingLen) - if _, err := io.ReadFull(c.cipher.Rand(), plain.Buf[paddingOffset:]); err != nil { + // padding 随机数走 per-Conn 缓冲读:每帧 12..1024 字节直读 crypto/rand 是一次 + // getrandom syscall,缓冲后按 ~1KiB 批量取。只由 outbound actor 单 goroutine 访问, + // 随机源本身不变(仍是 cipher 的 CSPRNG),只是预读。 + if c.outboundRand == nil { + c.outboundRand = bufio.NewReaderSize(c.cipher.Rand(), 1024) + } + if _, err := io.ReadFull(c.outboundRand, plain.Buf[paddingOffset:]); err != nil { return nil, err } diff --git a/internal/mtprotoedge/outbound_test.go b/internal/mtprotoedge/outbound_test.go index 38466b93..95b6c0b6 100644 --- a/internal/mtprotoedge/outbound_test.go +++ b/internal/mtprotoedge/outbound_test.go @@ -104,6 +104,42 @@ func TestOutboundActorSerializesConcurrentSends(t *testing.T) { } } +func TestSendBestEffortQueueFullBehavior(t *testing.T) { + c := &Conn{metrics: NopMetrics{}} + c.outbound = make(chan outboundOp, 1) + c.outboundControl = make(chan outboundOp, 1) + c.outboundStop = make(chan struct{}) + // 占满普通队列,模拟出站拥塞。 + c.outbound <- outboundOp{} + + if err := c.SendBestEffort(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}, 0); err != ErrOutboundQueueFull { + t.Fatalf("timeout=0 on full queue: err = %v, want ErrOutboundQueueFull", err) + } + + start := time.Now() + if err := c.SendBestEffort(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}, 30*time.Millisecond); err != ErrOutboundQueueFull { + t.Fatalf("timeout=30ms on full queue: err = %v, want ErrOutboundQueueFull", err) + } + if waited := time.Since(start); waited < 30*time.Millisecond { + t.Fatalf("timeout wait = %v, want >= 30ms", waited) + } + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + if err := c.SendBestEffort(canceled, proto.MessageFromServer, &mt.MsgsAck{}, time.Second); err != context.Canceled { + t.Fatalf("canceled ctx on full queue: err = %v, want context.Canceled", err) + } + + // 腾出队列后快路径应直接入队成功。 + <-c.outbound + if err := c.SendBestEffort(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}, 0); err != nil { + t.Fatalf("enqueue after drain: %v", err) + } + if got := len(c.outbound); got != 1 { + t.Fatalf("queued ops = %d, want 1", got) + } +} + func TestFrameNeedsAckServiceExceptions(t *testing.T) { cases := []struct { name string diff --git a/internal/mtprotoedge/rpc_result_cache.go b/internal/mtprotoedge/rpc_result_cache.go index 6e8ed756..28b46754 100644 --- a/internal/mtprotoedge/rpc_result_cache.go +++ b/internal/mtprotoedge/rpc_result_cache.go @@ -2,6 +2,7 @@ package mtprotoedge import ( "container/list" + "encoding/binary" "sync" "time" ) @@ -10,6 +11,10 @@ const ( rpcResultCacheTTL = 3 * time.Minute rpcResultCacheMaxEntries = 4096 rpcResultCacheMaxBytes = 64 << 20 + // rpcResultCacheShards 把缓存按 (auth_key_id, session_id) 分片:每条 RPC 都要 + // Get(重复检测)+ Put(结果缓存),单把全局锁会让所有连接的 RPC 热路径在 + // 一个 mutex 上汇聚(同 P0-5 的 SessionManager 教训)。分片数为 2 的幂。 + rpcResultCacheShards = 16 ) type rpcResultCacheKey struct { @@ -25,7 +30,14 @@ type rpcResultCacheEntry struct { expiresAt time.Time } +// rpcResultCache 缓存已回发的 rpc_result(按 auth_key+session+req_msg_id),用于 +// 跨连接重放重复请求。encodedOutboundMessage 构造后不可变(push fan-out 与 pending +// resend 均依赖该契约),因此 Get/Put 直接共享指针,不做防御性拷贝。 type rpcResultCache struct { + shards [rpcResultCacheShards]rpcResultCacheShard +} + +type rpcResultCacheShard struct { mu sync.Mutex now func() time.Time ttl time.Duration @@ -40,14 +52,24 @@ func newRPCResultCache(now func() time.Time) *rpcResultCache { if now == nil { now = time.Now } - return &rpcResultCache{ - now: now, - ttl: rpcResultCacheTTL, - maxEntries: rpcResultCacheMaxEntries, - maxBytes: rpcResultCacheMaxBytes, - order: list.New(), - byKey: make(map[rpcResultCacheKey]*list.Element), + c := &rpcResultCache{} + for i := range c.shards { + s := &c.shards[i] + s.now = now + s.ttl = rpcResultCacheTTL + s.maxEntries = rpcResultCacheMaxEntries / rpcResultCacheShards + s.maxBytes = rpcResultCacheMaxBytes / rpcResultCacheShards + s.order = list.New() + s.byKey = make(map[rpcResultCacheKey]*list.Element) } + return c +} + +func (c *rpcResultCache) shard(key rpcResultCacheKey) *rpcResultCacheShard { + // auth_key_id 与 session_id 都是均匀随机的 64-bit 值,异或折叠后取低位即可。 + h := binary.LittleEndian.Uint64(key.authKeyID[:]) ^ uint64(key.sessionID) + h ^= h >> 32 + return &c.shards[h&(rpcResultCacheShards-1)] } func (c *rpcResultCache) Get(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool) { @@ -55,102 +77,87 @@ func (c *rpcResultCache) Get(authKeyID [8]byte, sessionID, reqMsgID int64) (*enc return nil, false } key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID} - now := c.now() + s := c.shard(key) + now := s.now() - c.mu.Lock() - defer c.mu.Unlock() + s.mu.Lock() + defer s.mu.Unlock() - elem, ok := c.byKey[key] + elem, ok := s.byKey[key] if !ok { return nil, false } entry := elem.Value.(*rpcResultCacheEntry) if !entry.expiresAt.After(now) { - c.removeElement(elem) + s.removeElement(elem) return nil, false } - return cloneEncodedOutboundMessage(entry.encoded), true + return entry.encoded, true } func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encoded *encodedOutboundMessage) { if c == nil || reqMsgID == 0 || encoded == nil { return } - copied := cloneEncodedOutboundMessage(encoded) - if copied == nil { - return - } - size := len(copied.body) - if c.maxBytes > 0 && size > c.maxBytes { - return - } - key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID} - now := c.now() + s := c.shard(key) + size := len(encoded.body) + if s.maxBytes > 0 && size > s.maxBytes { + return + } + now := s.now() - c.mu.Lock() - defer c.mu.Unlock() + s.mu.Lock() + defer s.mu.Unlock() - c.expireLocked(now) - if elem, ok := c.byKey[key]; ok { - c.removeElement(elem) + s.expireLocked(now) + if elem, ok := s.byKey[key]; ok { + s.removeElement(elem) } entry := &rpcResultCacheEntry{ key: key, - encoded: copied, + encoded: encoded, size: size, - expiresAt: now.Add(c.ttl), + expiresAt: now.Add(s.ttl), } - elem := c.order.PushBack(entry) - c.byKey[key] = elem - c.bytes += size - c.trimLocked() + elem := s.order.PushBack(entry) + s.byKey[key] = elem + s.bytes += size + s.trimLocked() } -func (c *rpcResultCache) expireLocked(now time.Time) { - for elem := c.order.Front(); elem != nil; { +func (s *rpcResultCacheShard) expireLocked(now time.Time) { + for elem := s.order.Front(); elem != nil; { next := elem.Next() entry := elem.Value.(*rpcResultCacheEntry) if entry.expiresAt.After(now) { return } - c.removeElement(elem) + s.removeElement(elem) elem = next } } -func (c *rpcResultCache) trimLocked() { - for c.order.Len() > 0 { - tooManyEntries := c.maxEntries > 0 && c.order.Len() > c.maxEntries - tooManyBytes := c.maxBytes > 0 && c.bytes > c.maxBytes +func (s *rpcResultCacheShard) trimLocked() { + for s.order.Len() > 0 { + tooManyEntries := s.maxEntries > 0 && s.order.Len() > s.maxEntries + tooManyBytes := s.maxBytes > 0 && s.bytes > s.maxBytes if !tooManyEntries && !tooManyBytes { return } - c.removeElement(c.order.Front()) + s.removeElement(s.order.Front()) } } -func (c *rpcResultCache) removeElement(elem *list.Element) { +func (s *rpcResultCacheShard) removeElement(elem *list.Element) { if elem == nil { return } entry := elem.Value.(*rpcResultCacheEntry) - delete(c.byKey, entry.key) - c.bytes -= entry.size - if c.bytes < 0 { - c.bytes = 0 - } - c.order.Remove(elem) -} - -func cloneEncodedOutboundMessage(src *encodedOutboundMessage) *encodedOutboundMessage { - if src == nil { - return nil - } - body := append([]byte(nil), src.body...) - return &encodedOutboundMessage{ - body: body, - typeID: src.typeID, - reqMsgID: src.reqMsgID, + delete(s.byKey, entry.key) + s.bytes -= entry.size + if s.bytes < 0 { + s.bytes = 0 } + s.order.Remove(elem) } diff --git a/internal/mtprotoedge/rpc_result_cache_test.go b/internal/mtprotoedge/rpc_result_cache_test.go new file mode 100644 index 00000000..9d6ac353 --- /dev/null +++ b/internal/mtprotoedge/rpc_result_cache_test.go @@ -0,0 +1,68 @@ +package mtprotoedge + +import ( + "testing" + "time" +) + +func TestRPCResultCacheRoundTripAndTTL(t *testing.T) { + now := time.Unix(1000, 0) + cache := newRPCResultCache(func() time.Time { return now }) + + var keyID [8]byte + keyID[0] = 0xab + encoded := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, typeID: 42, reqMsgID: 7} + + if _, ok := cache.Get(keyID, 5, 7); ok { + t.Fatal("unexpected hit on empty cache") + } + cache.Put(keyID, 5, 7, encoded) + + got, ok := cache.Get(keyID, 5, 7) + if !ok { + t.Fatal("expected hit") + } + // encodedOutboundMessage 不可变契约下 Get/Put 共享指针,不做防御性拷贝。 + if got != encoded { + t.Fatal("expected shared pointer, got clone") + } + + // 不同 session / msg_id 不串。 + if _, ok := cache.Get(keyID, 6, 7); ok { + t.Fatal("hit with wrong session id") + } + if _, ok := cache.Get(keyID, 5, 8); ok { + t.Fatal("hit with wrong msg id") + } + + // TTL 过期。 + now = now.Add(rpcResultCacheTTL + time.Second) + if _, ok := cache.Get(keyID, 5, 7); ok { + t.Fatal("expected expiry after TTL") + } +} + +func TestRPCResultCacheShardTrim(t *testing.T) { + now := time.Unix(1000, 0) + cache := newRPCResultCache(func() time.Time { return now }) + + var keyID [8]byte + // 同一 (auth_key, session) 固定落在同一 shard;塞超过单 shard 条数上限,最旧的被逐出。 + perShard := rpcResultCacheMaxEntries / rpcResultCacheShards + for i := 0; i < perShard+1; i++ { + cache.Put(keyID, 1, int64(100+i), &encodedOutboundMessage{body: []byte{byte(i)}}) + } + if _, ok := cache.Get(keyID, 1, 100); ok { + t.Fatal("oldest entry should have been evicted by per-shard entry limit") + } + if _, ok := cache.Get(keyID, 1, int64(100+perShard)); !ok { + t.Fatal("newest entry should survive") + } + + // 单条超过单 shard 字节预算的结果不入缓存。 + huge := &encodedOutboundMessage{body: make([]byte, rpcResultCacheMaxBytes/rpcResultCacheShards+1)} + cache.Put(keyID, 2, 999, huge) + if _, ok := cache.Get(keyID, 2, 999); ok { + t.Fatal("oversized entry should be rejected") + } +} diff --git a/internal/mtprotoedge/seq_validate_test.go b/internal/mtprotoedge/seq_validate_test.go new file mode 100644 index 00000000..0ffb3393 --- /dev/null +++ b/internal/mtprotoedge/seq_validate_test.go @@ -0,0 +1,74 @@ +package mtprotoedge + +import ( + "math/rand" + "testing" +) + +// referenceValidateSeq 是 validateSeq 的旧全扫描语义(无高水位快路径), +// 作为随机对拍的行为基准:快路径只允许接受「全扫描也必然接受」的子集。 +func referenceValidateSeq(cs *connState, msgID int64, seqNo int32, content bool) int { + if !content { + return 0 + } + for seenMsgID, record := range cs.seen { + if !record.content { + continue + } + if seenMsgID < msgID && record.seqNo >= seqNo { + return badMsgSeqTooLow + } + if seenMsgID > msgID && record.seqNo <= seqNo { + return badMsgSeqTooHigh + } + } + return 0 +} + +func TestValidateSeqFastPathMatchesFullScan(t *testing.T) { + rng := rand.New(rand.NewSource(20260705)) + for round := 0; round < 32; round++ { + fast := newConnState() + ref := newConnState() + for i := 0; i < 2000; i++ { + // 小值域制造乱序、重复 seq 与 too_low/too_high 反转;大 i 也覆盖淘汰窗口。 + msgID := int64(rng.Intn(3000) + 1) + seqNo := int32(rng.Intn(600)) + content := rng.Intn(4) != 0 + if _, ok := fast.seen[msgID]; ok { + continue // 真实调用链在 seenRecord 命中时不会走 validateSeq + } + got := fast.validateSeq(msgID, seqNo, content) + want := referenceValidateSeq(ref, msgID, seqNo, content) + if got != want { + t.Fatalf("round %d step %d: validateSeq(msg_id=%d seq=%d content=%v) = %d, want %d", + round, i, msgID, seqNo, content, got, want) + } + if got == 0 { + fast.track(msgID, seqNo, content, msgStateReceived) + ref.track(msgID, seqNo, content, msgStateReceived) + } + } + } +} + +func TestValidateSeqOrderedFastPath(t *testing.T) { + cs := newConnState() + // 正常客户端:msg_id 与 content seq_no 严格递增,应全部通过。 + for i := 0; i < 1000; i++ { + msgID := int64(1000 + i*4) + seqNo := int32(i*2 + 1) + if code := cs.validateSeq(msgID, seqNo, true); code != 0 { + t.Fatalf("ordered message %d rejected with code %d", i, code) + } + cs.track(msgID, seqNo, true, msgStateReceived) + } + // seq 回退必须仍被拒绝(快路径不放行)。 + if code := cs.validateSeq(1000+1000*4, 3, true); code != badMsgSeqTooLow { + t.Fatalf("seq regression code = %d, want badMsgSeqTooLow", code) + } + // 旧 msg_id 配新 seq 也必须仍被拒绝。 + if code := cs.validateSeq(500, 5000, true); code != badMsgSeqTooHigh { + t.Fatalf("old msg_id with high seq code = %d, want badMsgSeqTooHigh", code) + } +} diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index 1c27020b..c1c22d9a 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -3,6 +3,7 @@ package mtprotoedge import ( "context" "crypto/rsa" + "encoding/hex" "errors" "fmt" "io" @@ -230,9 +231,11 @@ func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt writeTimeout: s.writeTimeout, metrics: s.metrics, authKeyID: key.ID, + authKeyHex: hex.EncodeToString(key.ID[:]), sessionID: sessionID, salt: salt, key: key, + createdAt: s.clock.Now(), } c.startOutbound() c.startInboundRPCScheduler(s.rpcInflight, s.rpcQueueSize, s.rpcTimeout) @@ -459,6 +462,9 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) 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 { @@ -512,15 +518,25 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) fetchedKey = &d } - current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b) + current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b, &plain) if err != nil { return err } } } +// 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) diff --git a/internal/mtprotoedge/transport_compat.go b/internal/mtprotoedge/transport_compat.go index 2d518873..7353b78e 100644 --- a/internal/mtprotoedge/transport_compat.go +++ b/internal/mtprotoedge/transport_compat.go @@ -1,6 +1,7 @@ package mtprotoedge import ( + "bufio" "bytes" "context" "encoding/binary" @@ -128,18 +129,19 @@ type compatTransportConn struct { } func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error { + deadline, _ := ctx.Deadline() + return c.SendDeadline(deadline, b) +} + +// SendDeadline 按显式写超时发送一帧(deadline 为零值表示不设超时)。 +// 出站热路径(Conn.writeFrame)走这里,免去 per-frame context timer 分配。 +func (c *compatTransportConn) SendDeadline(deadline time.Time, b *bin.Buffer) error { c.writeMux.Lock() defer c.writeMux.Unlock() - if err := c.conn.SetWriteDeadline(time.Time{}); err != nil { - return errors.Wrap(err, "reset write deadline") + if err := c.conn.SetWriteDeadline(deadline); err != nil { + return errors.Wrap(err, "set write deadline") } - if deadline, ok := ctx.Deadline(); ok { - if err := c.conn.SetWriteDeadline(deadline); err != nil { - return errors.Wrap(err, "set write deadline") - } - } - if err := c.codec.Write(c.conn, b); err != nil { return errors.Wrap(err, "write") } @@ -163,13 +165,9 @@ func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) er c.writeMux.Lock() defer c.writeMux.Unlock() - if err := c.conn.SetWriteDeadline(time.Time{}); err != nil { - return errors.Wrap(err, "reset write deadline") - } - if deadline, ok := ctx.Deadline(); ok { - if err := c.conn.SetWriteDeadline(deadline); err != nil { - return errors.Wrap(err, "set write deadline") - } + deadline, _ := ctx.Deadline() + if err := c.conn.SetWriteDeadline(deadline); err != nil { + return errors.Wrap(err, "set write deadline") } raw := q.quickAckResponse(token) @@ -180,18 +178,20 @@ func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) er } func (c *compatTransportConn) Recv(ctx context.Context, b *bin.Buffer) error { + deadline, _ := ctx.Deadline() + return c.RecvDeadline(deadline, b) +} + +// RecvDeadline 按显式读超时收一帧(deadline 为零值表示不设超时)。 +// serveConn 的每帧读走这里,免去 per-frame context timer 分配;连接取消仍由 +// serveConn 的 ctx watcher 主动 Close 底层连接来解除阻塞(与旧行为一致)。 +func (c *compatTransportConn) RecvDeadline(deadline time.Time, b *bin.Buffer) error { c.readMux.Lock() defer c.readMux.Unlock() - if err := c.conn.SetReadDeadline(time.Time{}); err != nil { - return errors.Wrap(err, "reset read deadline") + if err := c.conn.SetReadDeadline(deadline); err != nil { + return errors.Wrap(err, "set read deadline") } - if deadline, ok := ctx.Deadline(); ok { - if err := c.conn.SetReadDeadline(deadline); err != nil { - return errors.Wrap(err, "set read deadline") - } - } - if err := c.codec.Read(c.conn, b); err != nil { return errors.Wrap(err, "read") } @@ -233,6 +233,7 @@ type quickAckCodec interface { type quickAckAbridgedCodec struct { quickAckRequested bool + wbuf []byte } func (*quickAckAbridgedCodec) WriteHeader(w io.Writer) error { @@ -243,7 +244,7 @@ func (*quickAckAbridgedCodec) ReadHeader(r io.Reader) error { return (codec.Abridged{}).ReadHeader(r) } -func (*quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error { +func (q *quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error { if err := validateOutgoingCompatMessage(b); err != nil { return err } @@ -260,7 +261,7 @@ func (*quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error { header[3] = byte(words >> 16) headerLen = 4 } - return writeCompatPacket(w, header[:headerLen], b.Raw()) + return writeCompatPacket(w, &q.wbuf, header[:headerLen], b.Raw()) } func (q *quickAckAbridgedCodec) Read(r io.Reader, b *bin.Buffer) error { @@ -286,6 +287,7 @@ func (*quickAckAbridgedCodec) quickAckResponse(token uint32) [4]byte { type quickAckIntermediateCodec struct { quickAckRequested bool + wbuf []byte } func (*quickAckIntermediateCodec) WriteHeader(w io.Writer) error { @@ -296,13 +298,13 @@ func (*quickAckIntermediateCodec) ReadHeader(r io.Reader) error { return (codec.Intermediate{}).ReadHeader(r) } -func (*quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error { +func (q *quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error { if err := validateOutgoingCompatMessage(b); err != nil { return err } var header [4]byte binary.LittleEndian.PutUint32(header[:], uint32(b.Len())) - return writeCompatPacket(w, header[:], b.Raw()) + return writeCompatPacket(w, &q.wbuf, header[:], b.Raw()) } func (q *quickAckIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error { @@ -328,6 +330,8 @@ func (*quickAckIntermediateCodec) quickAckResponse(token uint32) [4]byte { type quickAckPaddedIntermediateCodec struct { quickAckRequested bool + wbuf []byte + rand *bufio.Reader } func (*quickAckPaddedIntermediateCodec) WriteHeader(w io.Writer) error { @@ -338,22 +342,27 @@ func (*quickAckPaddedIntermediateCodec) ReadHeader(r io.Reader) error { return (codec.PaddedIntermediate{}).ReadHeader(r) } -func (*quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error { +func (q *quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error { if err := validateOutgoingCompatMessage(b); err != nil { return err } + // padding 随机数走 per-codec 缓冲预读;codec 写入被 compatTransportConn.writeMux + // 串行化,单 goroutine 访问安全。 + if q.rand == nil { + q.rand = bufio.NewReaderSize(tdcrypto.DefaultRand(), 64) + } var padding [4]byte - if _, err := io.ReadFull(tdcrypto.DefaultRand(), padding[:]); err != nil { + if _, err := io.ReadFull(q.rand, padding[:]); err != nil { return err } n := int(padding[0] % 4) - payload := b.Raw() - if n > 0 { - payload = append(append([]byte(nil), payload...), padding[:n]...) - } - var header [4]byte - binary.LittleEndian.PutUint32(header[:], uint32(len(payload))) - return writeCompatPacket(w, header[:], payload) + // 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) } func (q *quickAckPaddedIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error { @@ -440,11 +449,13 @@ func validateOutgoingCompatMessage(b *bin.Buffer) error { return nil } -func writeCompatPacket(w io.Writer, header, payload []byte) error { - packet := make([]byte, 0, len(header)+len(payload)) - packet = append(packet, header...) - packet = append(packet, payload...) - return writeAll(w, packet) +// 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) } func writeAll(w io.Writer, p []byte) error { diff --git a/internal/mtprotoedge/transport_compat_test.go b/internal/mtprotoedge/transport_compat_test.go index e2ebf042..c863021c 100644 --- a/internal/mtprotoedge/transport_compat_test.go +++ b/internal/mtprotoedge/transport_compat_test.go @@ -63,6 +63,40 @@ func TestQuickAckReadFlags(t *testing.T) { } } +func TestCompatPaddedIntermediateWriteRoundTrip(t *testing.T) { + codec := &quickAckPaddedIntermediateCodec{} + // 连写多帧:验证复用写缓冲不串包,且 padding 后仍能被读端正确剥离。 + for i := 0; i < 8; i++ { + var payload bin.Buffer + payload.PutInt32(int32(0x11220000 + i)) + payload.PutInt32(int32(0x33440000 + i)) + + var out countWriteBuffer + 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) + } + total := binary.LittleEndian.Uint32(out.Bytes()[:4]) + if int(total) != len(out.Bytes())-4 { + t.Fatalf("write %d: header length = %d, body = %d", i, total, len(out.Bytes())-4) + } + + var got bin.Buffer + requested, err := readQuickAckIntermediate(bytes.NewReader(out.Bytes()), &got, true) + if err != nil { + t.Fatalf("read back %d: %v", i, err) + } + if requested { + t.Fatalf("read back %d: unexpected quick ack flag", i) + } + if !bytes.Equal(got.Raw(), payload.Raw()) { + t.Fatalf("read back %d: payload = %x, want %x", i, got.Raw(), payload.Raw()) + } + } +} + func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) { var payload bin.Buffer payload.PutInt32(0x01020304)