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
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))
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue