perf: sync dispatch hot path optimizations
This commit is contained in:
A
parent
03b785ebf4
commit
7e64d9c30e
13 changed files with 824 additions and 215 deletions
|
|
@ -1,7 +1,9 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/hex"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
|
@ -32,6 +34,9 @@ type Conn struct {
|
|||
metrics Metrics
|
||||
|
||||
authKeyID [8]byte
|
||||
// authKeyHex 是 authKeyID 的 hex 缓存：每条 RPC 的结构化日志都会带它，
|
||||
// 建连时算一次，避免热路径反复 hex 编码分配。
|
||||
authKeyHex string
|
||||
sessionID int64
|
||||
salt int64
|
||||
key crypto.AuthKey
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
93
internal/mtprotoedge/decrypt.go
Normal file
93
internal/mtprotoedge/decrypt.go
Normal file
|
|
@ -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
|
||||
}
|
||||
181
internal/mtprotoedge/decrypt_test.go
Normal file
181
internal/mtprotoedge/decrypt_test.go
Normal file
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
if timeout == 0 {
|
||||
// 快路径：非阻塞入队。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 {
|
||||
c.metrics.OutboundDropped("push_queue_full")
|
||||
return ErrOutboundQueueFull
|
||||
}
|
||||
c.metrics.OutboundQueueWait(len(c.outbound), cap(c.outbound))
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
enqueueCtx := ctx
|
||||
if enqueueCtx == nil {
|
||||
enqueueCtx = 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 {
|
||||
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 err
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
68
internal/mtprotoedge/rpc_result_cache_test.go
Normal file
68
internal/mtprotoedge/rpc_result_cache_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
74
internal/mtprotoedge/seq_validate_test.go
Normal file
74
internal/mtprotoedge/seq_validate_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 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,14 +165,10 @@ 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 {
|
||||
deadline, _ := ctx.Deadline()
|
||||
if err := c.conn.SetWriteDeadline(deadline); err != nil {
|
||||
return errors.Wrap(err, "set write deadline")
|
||||
}
|
||||
}
|
||||
|
||||
raw := q.quickAckResponse(token)
|
||||
if err := writeAll(c.conn, raw[:]); err != nil {
|
||||
|
|
@ -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 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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue