internal/mtprotoedge/session_boundary_test.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"context"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestFirstContainerBoundaryKeepsAndroidRequestMap models DrKLO's
// new_session_created handling: it drops every running request whose msg_id is
// lower than first_msg_id. Every inner request accepted from the first
// container must therefore remain addressable when its pong arrives.
func TestFirstContainerBoundaryKeepsAndroidRequestMap(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
msgIDs := proto.NewMessageIDGen(time.Now)
running := make(map[int64]int64, 3)
messages := make([]proto.Message, 0, 3)
for i := 0; i < 3; i++ {
msgID := msgIDs.New(proto.MessageFromClient)
pingID := int64(10_001 + i)
body := mustEncodeTL(t, &mt.PingRequest{PingID: pingID})
messages = append(messages, proto.Message{
ID: msgID,
SeqNo: 1 + i*2,
Bytes: len(body),
Body: body,
})
running[msgID] = pingID
}
outerMsgID := msgIDs.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages})
boundaryFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.NewSessionCreatedTypeID: 1,
})
var created mt.NewSessionCreated
for _, frame := range boundaryFrames {
if frame.TypeID == mt.PongTypeID {
t.Fatalf("pong arrived before new_session_created boundary")
}
if frame.TypeID != mt.NewSessionCreatedTypeID {
continue
}
if err := created.Decode(frame.Plain); err != nil {
t.Fatalf("decode new_session_created: %v", err)
}
}
// DrKLO ConnectionsManager.cpp clears running requests with
// request.messageId < first_msg_id when it receives this notification.
for msgID := range running {
if msgID < created.FirstMsgID {
delete(running, msgID)
}
}
for _, accepted := range messages {
if _, ok := running[accepted.ID]; !ok {
t.Fatalf(
"accepted inner msg_id %d was evicted by Android boundary %d (outer=%d)",
accepted.ID,
created.FirstMsgID,
outerMsgID,
)
}
}
pongFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.PongTypeID: len(messages),
})
seen := make(map[int64]struct{}, len(messages))
for _, frame := range pongFrames {
if frame.TypeID != mt.PongTypeID {
continue
}
var pong mt.Pong
if err := pong.Decode(frame.Plain); err != nil {
t.Fatalf("decode pong: %v", err)
}
wantPingID, ok := running[pong.MsgID]
if !ok {
t.Fatalf("orphan pong for msg_id %d after Android boundary cleanup", pong.MsgID)
}
if pong.PingID != wantPingID {
t.Fatalf("pong ping_id = %d for msg_id %d, want %d", pong.PingID, pong.MsgID, wantPingID)
}
if _, duplicate := seen[pong.MsgID]; duplicate {
t.Fatalf("duplicate pong for msg_id %d", pong.MsgID)
}
seen[pong.MsgID] = struct{}{}
delete(running, pong.MsgID)
}
if len(running) != 0 {
t.Fatalf("requests without correlated pong: %+v", running)
}
}
type admissionCountingRPC struct {
calls atomic.Int32
}
func (h *admissionCountingRPC) Dispatch(_ context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
h.calls.Add(1)
return &tg.Config{ThisDC: 2}, nil
}
func (*admissionCountingRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestContainerRPCAdmissionFailureIsAtomic(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
legacyRPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCGlobalWorkers: 1,
})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
msgIDs := proto.NewMessageIDGen(time.Now)
messages := make([]proto.Message, 0, 2)
requestIDs := make(map[int64]struct{}, 2)
for i := 0; i < 2; i++ {
msgID := msgIDs.New(proto.MessageFromClient)
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
messages = append(messages, proto.Message{ID: msgID, SeqNo: 1 + i*2, Bytes: len(body), Body: body})
requestIDs[msgID] = struct{}{}
}
outerMsgID := msgIDs.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 2})
results := 0
for _, frame := range frames {
if frame.TypeID != proto.ResultTypeID {
continue
}
var result proto.Result
if err := result.Decode(frame.Plain); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if _, ok := requestIDs[result.RequestMessageID]; !ok {
t.Fatalf("unexpected or duplicate capacity rpc_result req_msg_id %d", result.RequestMessageID)
}
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode capacity rpc_error: %v", err)
}
if rpcErr.ErrorCode != rpcWorkerBusyErrorCode || rpcErr.ErrorMessage != rpcWorkerBusyErrorMessage {
t.Fatalf("capacity rpc_error = %+v", rpcErr)
}
delete(requestIDs, result.RequestMessageID)
results++
}
if results != 2 {
t.Fatalf("capacity rpc_results = %d, want 2", results)
}
if len(requestIDs) != 0 {
t.Fatalf("capacity requests without exactly one result: %+v", requestIDs)
}
if got := handler.calls.Load(); got != 0 {
t.Fatalf("partially executed handler calls = %d, want 0", got)
}
}
func TestGZIPWrappedRPCUsesLogicalEnvelopeBoundary(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
requestBody := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &proto.GZIP{Data: requestBody})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range frames {
if frame.TypeID == mt.NewSessionCreatedTypeID {
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode gzip RPC boundary: %v", err)
}
}
}
if boundary.FirstMsgID != msgID {
t.Fatalf("gzip RPC boundary = %d, want envelope %d", boundary.FirstMsgID, msgID)
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("gzip RPC handler calls = %d, want 1", got)
}
}
func TestGZIPWrappedContainerUsesLowestInnerBoundary(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
innerMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 7001})
containerBody := mustEncodeTL(t, &proto.MessageContainer{Messages: []proto.Message{{
ID: innerMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody,
}}})
sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, 2, &proto.GZIP{Data: containerBody})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.PongTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range frames {
if frame.TypeID == mt.NewSessionCreatedTypeID {
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode gzip container boundary: %v", err)
}
}
}
if boundary.FirstMsgID != innerMsgID {
t.Fatalf("gzip container boundary = %d, want inner %d (outer %d)", boundary.FirstMsgID, innerMsgID, outerMsgID)
}
}
func TestEmptyContainerUsesOuterBoundary(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
outerMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, 0, &proto.MessageContainer{})
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.NewSessionCreatedTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range frames {
if frame.TypeID != mt.NewSessionCreatedTypeID {
continue
}
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode empty container boundary: %v", err)
}
}
if boundary.FirstMsgID != outerMsgID {
t.Fatalf("empty container boundary = %d, want outer %d", boundary.FirstMsgID, outerMsgID)
}
}
func TestDuplicateContainerAcksWithoutBusinessReexecution(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
messages := make([]proto.Message, 0, 2)
for i := 0; i < 2; i++ {
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
messages = append(messages, proto.Message{
ID: ids.New(proto.MessageFromClient), SeqNo: 1 + i*2, Bytes: len(body), Body: body,
})
}
outerMsgID := ids.New(proto.MessageFromClient)
container := &proto.MessageContainer{Messages: messages}
sendEncrypted(t, conn, cipher, auth, outerMsgID, container)
initialFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 2})
answerIDs := make([]int64, 0, 2)
for _, frame := range initialFrames {
if frame.TypeID == proto.ResultTypeID {
answerIDs = append(answerIDs, frame.Message.MessageID)
}
}
if len(answerIDs) != 2 {
t.Fatalf("initial rpc_result answer ids = %d, want 2", len(answerIDs))
}
// ACK both original server results before retransmitting the container. This is
// a wire barrier: any later rpc_result is necessarily a duplicate replay rather
// than an unconsumed response from the initial batch.
ackMsgID := ids.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, ackMsgID, 4, &mt.MsgsAck{MsgIDs: answerIDs})
if got := handler.calls.Load(); got != 2 {
t.Fatalf("initial handler calls = %d, want 2", got)
}
sendEncrypted(t, conn, cipher, auth, outerMsgID, container)
duplicateFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.MsgsAckTypeID: 1})
for _, frame := range duplicateFrames {
if frame.TypeID == proto.ResultTypeID {
t.Fatal("same-connection duplicate container replayed an extra rpc_result")
}
}
if got := handler.calls.Load(); got != 2 {
t.Fatalf("duplicate container reexecuted handlers: calls=%d, want 2", got)
}
}
type largeStartupBurstRPC struct {
calls atomic.Int32
body string
}
func (h *largeStartupBurstRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) {
call := h.calls.Add(1)
// Deliberately perturb completion order while remaining cancellation-aware.
timer := time.NewTimer(time.Duration(call%4) * time.Millisecond)
defer timer.Stop()
select {
case <-timer.C:
case <-ctx.Done():
return nil, ctx.Err()
}
return &tg.DataJSON{Data: h.body}, nil
}
func (*largeStartupBurstRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestAndroidStartupBurstKeepsAllLargeRPCResultsAddressable(t *testing.T) {
const (
dc = 2
requests = 30
)
handler := &largeStartupBurstRPC{body: strings.Repeat("x", 192<<10)}
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
running := make(map[int64]struct{}, requests)
messages := make([]proto.Message, 0, requests)
for i := 0; i < requests; i++ {
msgID := ids.New(proto.MessageFromClient)
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
messages = append(messages, proto.Message{
ID: msgID, SeqNo: 1 + i*2, Bytes: len(body), Body: body,
})
running[msgID] = struct{}{}
}
outerMsgID := ids.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages})
boundaryFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.NewSessionCreatedTypeID: 1})
var boundary mt.NewSessionCreated
for _, frame := range boundaryFrames {
if frame.TypeID == proto.ResultTypeID {
t.Fatal("large rpc_result arrived before session boundary")
}
if frame.TypeID == mt.NewSessionCreatedTypeID {
if err := boundary.Decode(frame.Plain); err != nil {
t.Fatalf("decode startup boundary: %v", err)
}
}
}
for msgID := range running {
if msgID < boundary.FirstMsgID {
delete(running, msgID)
}
}
if len(running) != requests {
t.Fatalf("Android boundary removed accepted startup requests: kept=%d want=%d floor=%d outer=%d", len(running), requests, boundary.FirstMsgID, outerMsgID)
}
resultFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: requests})
for _, frame := range resultFrames {
if frame.TypeID != proto.ResultTypeID {
continue
}
var result proto.Result
if err := result.Decode(frame.Plain); err != nil {
t.Fatalf("decode startup rpc_result: %v", err)
}
if _, ok := running[result.RequestMessageID]; !ok {
t.Fatalf("orphan large rpc_result for request %d", result.RequestMessageID)
}
delete(running, result.RequestMessageID)
}
if len(running) != 0 {
t.Fatalf("startup RPCs without result: %d", len(running))
}
if got := handler.calls.Load(); got != requests {
t.Fatalf("startup handler calls = %d, want %d", got, requests)
}
}