808 lines
31 KiB
Go
808 lines
31 KiB
Go
package mtprotoedge
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/flate"
|
|
"compress/zlib"
|
|
"context"
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/iamxvbaba/td/bin"
|
|
"github.com/iamxvbaba/td/mt"
|
|
"github.com/iamxvbaba/td/proto"
|
|
"go.uber.org/zap/zaptest"
|
|
)
|
|
|
|
func TestContainerMessageCountAndServiceVectorCaps(t *testing.T) {
|
|
var container bin.Buffer
|
|
container.PutID(proto.MessageContainerTypeID)
|
|
container.PutInt(maxContainerMessages)
|
|
if got, err := containerMessageCount(&container); err != nil || got != maxContainerMessages {
|
|
t.Fatalf("container count = %d/%v, want %d/nil", got, err, maxContainerMessages)
|
|
}
|
|
container.Buf[4]++
|
|
if got, err := containerMessageCount(&container); err != nil || got != maxContainerMessages+1 {
|
|
t.Fatalf("oversized container preflight = %d/%v, want %d/nil", got, err, maxContainerMessages+1)
|
|
}
|
|
|
|
ack := mt.MsgsAck{MsgIDs: make([]int64, maxServiceMessageIDs)}
|
|
var encoded bin.Buffer
|
|
if err := ack.Encode(&encoded); err != nil {
|
|
t.Fatalf("encode msgs_ack: %v", err)
|
|
}
|
|
if err := validateFirstVectorCount(&encoded, maxServiceMessageIDs); err != nil {
|
|
t.Fatalf("service vector at cap: %v", err)
|
|
}
|
|
// Count lives after constructor + vector constructor. We only mutate the declared count: the
|
|
// preflight must reject before generated Decode attempts a long loop/allocation.
|
|
encoded.Buf[8]++
|
|
if err := validateFirstVectorCount(&encoded, maxServiceMessageIDs); err == nil {
|
|
t.Fatal("service vector above cap unexpectedly accepted")
|
|
}
|
|
}
|
|
|
|
func TestServiceLongVectorViewsAreBoundedExactAndZeroCopy(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
typeID uint32
|
|
value bin.Encoder
|
|
}{
|
|
{name: "msgs_ack", typeID: mt.MsgsAckTypeID, value: &mt.MsgsAck{MsgIDs: []int64{11, 22}}},
|
|
{name: "msgs_state_req", typeID: mt.MsgsStateReqTypeID, value: &mt.MsgsStateReq{MsgIDs: []int64{11, 22}}},
|
|
{name: "msg_resend_req", typeID: mt.MsgResendReqTypeID, value: &mt.MsgResendReq{MsgIDs: []int64{11, 22}}},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var encoded bin.Buffer
|
|
if err := tt.value.Encode(&encoded); err != nil {
|
|
t.Fatalf("encode service vector: %v", err)
|
|
}
|
|
view, err := decodeInt64VectorView(&encoded, tt.typeID, maxServiceMessageIDs)
|
|
if err != nil {
|
|
t.Fatalf("decode service vector view: %v", err)
|
|
}
|
|
if view.count != 2 {
|
|
t.Fatalf("view count = %d, want 2", view.count)
|
|
}
|
|
|
|
// The retained plan payload is a view into the already-budgeted frame,
|
|
// not a generated-decoder []int64 copy.
|
|
binary.LittleEndian.PutUint64(encoded.Buf[12:20], uint64(99))
|
|
ids := view.materialize()
|
|
if len(ids) != 2 || ids[0] != 99 || ids[1] != 22 {
|
|
t.Fatalf("materialized ids = %v, want [99 22]", ids)
|
|
}
|
|
|
|
trailing := bin.Buffer{Buf: append(append([]byte(nil), encoded.Buf...), 0, 0, 0, 0)}
|
|
if _, err := decodeInt64VectorView(&trailing, tt.typeID, maxServiceMessageIDs); err == nil || !strings.Contains(err.Error(), "trailing") {
|
|
t.Fatalf("trailing vector err = %v, want exact-length rejection", err)
|
|
}
|
|
truncated := bin.Buffer{Buf: encoded.Buf[:len(encoded.Buf)-1]}
|
|
if _, err := decodeInt64VectorView(&truncated, tt.typeID, maxServiceMessageIDs); !errors.Is(err, io.ErrUnexpectedEOF) {
|
|
t.Fatalf("truncated vector err = %v, want unexpected EOF", err)
|
|
}
|
|
})
|
|
}
|
|
|
|
var oversized bin.Buffer
|
|
if err := (&mt.MsgsAck{MsgIDs: []int64{1}}).Encode(&oversized); err != nil {
|
|
t.Fatalf("encode oversized seed: %v", err)
|
|
}
|
|
binary.LittleEndian.PutUint32(oversized.Buf[8:12], uint32(maxServiceMessageIDs+1))
|
|
if _, err := decodeInt64VectorView(&oversized, mt.MsgsAckTypeID, maxServiceMessageIDs); err == nil || !strings.Contains(err.Error(), "exceeds") {
|
|
t.Fatalf("oversized vector err = %v, want capped rejection", err)
|
|
}
|
|
}
|
|
|
|
func TestInboundPlanContainerCapacityBoundary(t *testing.T) {
|
|
ids := proto.NewMessageIDGen(time.Now)
|
|
waitBody := encodeClientMessageBodyForTest(t, &mt.HTTPWaitRequest{MaxWait: 25_000})
|
|
messages := make([]proto.Message, maxContainerMessages)
|
|
for i := range messages {
|
|
messages[i] = proto.Message{
|
|
ID: ids.New(proto.MessageFromClient), SeqNo: 0, Bytes: len(waitBody), Body: waitBody,
|
|
}
|
|
}
|
|
outerMsgID := ids.New(proto.MessageFromClient)
|
|
body := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: messages})
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
before := s.frameBudget.usedBytes()
|
|
plan, err := s.preflightInbound(newConnState(), outerMsgID, 0, body)
|
|
if err != nil {
|
|
t.Fatalf("container at cap preflight: %v", err)
|
|
}
|
|
if len(plan.items) != maxContainerMessages || len(plan.staged) != maxContainerMessages+1 {
|
|
t.Fatalf("container plan sizes = items:%d staged:%d", len(plan.items), len(plan.staged))
|
|
}
|
|
if plan.logicalMin != messages[0].ID {
|
|
t.Fatalf("container plan floor = %d, want %d", plan.logicalMin, messages[0].ID)
|
|
}
|
|
plan.close()
|
|
if after := s.frameBudget.usedBytes(); after != before {
|
|
t.Fatalf("container plan leaked frame budget: before=%d after=%d", before, after)
|
|
}
|
|
|
|
over := append(append([]proto.Message(nil), messages...), proto.Message{
|
|
ID: outerMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody,
|
|
})
|
|
overBody := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: over})
|
|
if plan, err := s.preflightInbound(newConnState(), ids.New(proto.MessageFromClient), 0, overBody); err == nil {
|
|
plan.close()
|
|
t.Fatal("container above cap unexpectedly accepted")
|
|
}
|
|
if after := s.frameBudget.usedBytes(); after != before {
|
|
t.Fatalf("rejected container leaked frame budget: before=%d after=%d", before, after)
|
|
}
|
|
}
|
|
|
|
func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) {
|
|
encoded := bin.Buffer{}
|
|
wantBody := []byte{0x11, 0x22, 0x33, 0x44}
|
|
message := proto.Message{ID: 1, SeqNo: 1, Bytes: len(wantBody), Body: wantBody}
|
|
if err := (&proto.MessageContainer{Messages: []proto.Message{message}}).Encode(&encoded); err != nil {
|
|
t.Fatalf("encode container: %v", err)
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes - 1)
|
|
if _, release, err := s.decodeMessageContainerViews(&encoded, 1); err == nil {
|
|
release()
|
|
t.Fatal("descriptor allocation unexpectedly bypassed process budget")
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("failed descriptor reservation leaked %d bytes", got)
|
|
}
|
|
|
|
s.frameBudget = newInboundFrameBudget(2 * containerDescriptorBudgetBytes)
|
|
container, release, err := s.decodeMessageContainerViews(&encoded, 1)
|
|
if err != nil {
|
|
t.Fatalf("decode budgeted container: %v", err)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != containerDescriptorBudgetBytes {
|
|
t.Fatalf("descriptor budget = %d, want %d", got, containerDescriptorBudgetBytes)
|
|
}
|
|
container.Messages[0].Body[0] = 0x99
|
|
if encoded.Buf[8+16] != 0x99 {
|
|
t.Fatal("container body was copied instead of viewing the charged input frame")
|
|
}
|
|
release()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("released descriptor budget = %d, want zero", got)
|
|
}
|
|
|
|
truncated := bin.Buffer{Buf: encoded.Buf[:len(encoded.Buf)-1]}
|
|
if _, _, err := s.decodeMessageContainerViews(&truncated, 1); err == nil {
|
|
t.Fatal("truncated container unexpectedly decoded")
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("failed container decode leaked %d bytes", got)
|
|
}
|
|
}
|
|
|
|
func TestContainerAndGZIPViewsRejectTrailingBytes(t *testing.T) {
|
|
var container bin.Buffer
|
|
if err := (&proto.MessageContainer{}).Encode(&container); err != nil {
|
|
t.Fatalf("encode empty container: %v", err)
|
|
}
|
|
container.PutInt(0x11223344)
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
if _, release, err := s.decodeMessageContainerViews(&container, 0); err == nil {
|
|
release()
|
|
t.Fatal("container trailing bytes unexpectedly accepted")
|
|
}
|
|
|
|
var packed bin.Buffer
|
|
if err := (proto.GZIP{Data: []byte{1, 2, 3, 4}}).Encode(&packed); err != nil {
|
|
t.Fatalf("encode gzip wrapper: %v", err)
|
|
}
|
|
packed.PutInt(0x55667788)
|
|
if _, err := gzipPackedBytesView(&packed); err == nil {
|
|
t.Fatal("gzip trailing bytes unexpectedly accepted")
|
|
}
|
|
}
|
|
|
|
func TestContainerInvalidSequenceTailIsAtomic(t *testing.T) {
|
|
ids := proto.NewMessageIDGen(time.Now)
|
|
firstMsgID := ids.New(proto.MessageFromClient)
|
|
secondMsgID := ids.New(proto.MessageFromClient)
|
|
outerMsgID := ids.New(proto.MessageFromClient)
|
|
pingOne := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 1})
|
|
pingTwo := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2})
|
|
container := proto.MessageContainer{Messages: []proto.Message{
|
|
{ID: firstMsgID, SeqNo: 3, Bytes: len(pingOne), Body: pingOne},
|
|
{ID: secondMsgID, SeqNo: 1, Bytes: len(pingTwo), Body: pingTwo},
|
|
}}
|
|
encoded := encodeClientMessageBodyForTest(t, &container)
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
cs := newConnState()
|
|
c := &Conn{
|
|
metrics: NopMetrics{},
|
|
outbound: make(chan *outboundOp, 4),
|
|
outboundControl: make(chan *outboundOp, 4),
|
|
outboundStop: make(chan struct{}),
|
|
}
|
|
var acks []int64
|
|
if err := s.dispatch(context.Background(), cs, c, outerMsgID, 4, &bin.Buffer{Buf: encoded}, &acks); err != nil {
|
|
t.Fatalf("dispatch invalid-tail container: %v", err)
|
|
}
|
|
|
|
if len(cs.seen) != 0 || len(cs.order) != 0 || cs.maxContentMsgID != 0 || cs.maxContentSeqNo != 0 {
|
|
t.Fatalf("invalid container partially committed connState: %+v", cs)
|
|
}
|
|
if len(acks) != 0 {
|
|
t.Fatalf("invalid container produced ACKs: %v", acks)
|
|
}
|
|
if got := len(c.outboundControl); got != 1 {
|
|
t.Fatalf("control frames = %d, want only bad_msg_notification", got)
|
|
}
|
|
op := <-c.outboundControl
|
|
if op.encoded == nil || op.encoded.typeID != mt.BadMsgNotificationTypeID {
|
|
t.Fatalf("control frame type = %#v, want bad_msg_notification", op.encoded)
|
|
}
|
|
var bad mt.BadMsgNotification
|
|
if err := bad.Decode(&bin.Buffer{Buf: op.encoded.body}); err != nil {
|
|
t.Fatalf("decode bad_msg_notification: %v", err)
|
|
}
|
|
if bad.BadMsgID != secondMsgID || bad.BadMsgSeqno != 1 || bad.ErrorCode != badMsgSeqTooLow {
|
|
t.Fatalf("bad_msg_notification = %+v", bad)
|
|
}
|
|
op.releaseReservation(c.outboundTrackedBudget)
|
|
}
|
|
|
|
func TestContainerMalformedGZIPTailHasNoStateEffects(t *testing.T) {
|
|
ids := proto.NewMessageIDGen(time.Now)
|
|
waitMsgID := ids.New(proto.MessageFromClient)
|
|
badMsgID := ids.New(proto.MessageFromClient)
|
|
outerMsgID := ids.New(proto.MessageFromClient)
|
|
waitBody := encodeClientMessageBodyForTest(t, &mt.HTTPWaitRequest{MaxWait: 25_000})
|
|
var malformedGZIP bin.Buffer
|
|
malformedGZIP.PutID(proto.GZIPTypeID)
|
|
container := proto.MessageContainer{Messages: []proto.Message{
|
|
{ID: waitMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody},
|
|
{ID: badMsgID, SeqNo: 1, Bytes: malformedGZIP.Len(), Body: malformedGZIP.Buf},
|
|
}}
|
|
encoded := encodeClientMessageBodyForTest(t, &container)
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
cs := newConnState()
|
|
var acks []int64
|
|
if err := s.dispatch(context.Background(), cs, nil, outerMsgID, 2, &bin.Buffer{Buf: encoded}, &acks); err == nil {
|
|
t.Fatal("malformed gzip tail unexpectedly accepted")
|
|
}
|
|
if len(cs.seen) != 0 || len(cs.order) != 0 || cs.maxContentMsgID != 0 || cs.maxContentSeqNo != 0 {
|
|
t.Fatalf("malformed gzip tail partially committed connState: %+v", cs)
|
|
}
|
|
if len(acks) != 0 {
|
|
t.Fatalf("malformed gzip tail produced ACKs: %v", acks)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("malformed gzip tail leaked frame budget: %d", got)
|
|
}
|
|
}
|
|
|
|
func TestInvalidMessageIDRejectsBeforeGZIPExpansion(t *testing.T) {
|
|
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 1})
|
|
wrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
|
|
current := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
|
|
tests := []struct {
|
|
name string
|
|
msgID int64
|
|
badCode int
|
|
}{
|
|
{name: "stale", msgID: proto.NewMessageIDGen(func() time.Time { return time.Now().Add(-10 * time.Minute) }).New(proto.MessageFromClient), badCode: badMsgIDTooLow},
|
|
{name: "future", msgID: proto.NewMessageIDGen(func() time.Time { return time.Now().Add(time.Minute) }).New(proto.MessageFromClient), badCode: badMsgIDTooHigh},
|
|
{name: "invalid bits", msgID: current + 1, badCode: badMsgIDInvalidBits},
|
|
{name: "empty fractional bits", msgID: time.Now().Unix() << 32, badCode: badMsgIDInvalidBits},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
// Any attempted gzip expansion would fail this deliberately tiny
|
|
// budget and mask the required bad_msg_notification.
|
|
s.frameBudget = newInboundFrameBudget(1)
|
|
plan, err := s.preflightInbound(newConnState(), tt.msgID, 1, wrapped)
|
|
if plan != nil {
|
|
plan.close()
|
|
t.Fatal("invalid gzip envelope unexpectedly produced a plan")
|
|
}
|
|
var bad *dispatchBadMsgError
|
|
if !errors.As(err, &bad) || bad.code != tt.badCode {
|
|
t.Fatalf("invalid gzip envelope err = %v, want bad_msg code %d", err, tt.badCode)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("invalid gzip envelope consumed expansion budget: %d", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestInvalidContainerInnerIDRejectsBeforeGZIPExpansion(t *testing.T) {
|
|
ids := proto.NewMessageIDGen(time.Now)
|
|
validInnerID := ids.New(proto.MessageFromClient)
|
|
invalidInnerID := validInnerID + 1
|
|
if proto.MessageID(invalidInnerID).Type() == proto.MessageFromClient {
|
|
t.Fatalf("test inner msg_id %d unexpectedly has client bits", invalidInnerID)
|
|
}
|
|
outerMsgID := ids.New(proto.MessageFromClient)
|
|
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2})
|
|
gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
|
|
container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{
|
|
ID: invalidInnerID, SeqNo: 1, Bytes: len(gzipped), Body: gzipped,
|
|
}}})
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
// Enough for the one descriptor, intentionally not enough for a gzip
|
|
// expansion. Invalid inner bits must win and every descriptor charge unwind.
|
|
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
|
|
plan, err := s.preflightInbound(newConnState(), outerMsgID, 2, container)
|
|
if plan != nil {
|
|
plan.close()
|
|
t.Fatal("invalid inner msg_id unexpectedly produced a plan")
|
|
}
|
|
var bad *dispatchBadMsgError
|
|
if !errors.As(err, &bad) || bad.code != badMsgContainer || bad.msgID != invalidInnerID {
|
|
t.Fatalf("invalid inner gzip err = %v, want container rejection for %d", err, invalidInnerID)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("invalid inner gzip leaked descriptor/expansion budget: %d", got)
|
|
}
|
|
}
|
|
|
|
func TestSeenGZIPInnerShortCircuitsBeforeExpansion(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
seqNo int32
|
|
content bool
|
|
leaf bin.Encoder
|
|
wantACK bool
|
|
}{
|
|
{name: "content", seqNo: 1, content: true, leaf: &mt.PingRequest{PingID: 3}, wantACK: true},
|
|
{name: "non_content", seqNo: 0, content: false, leaf: &mt.HTTPWaitRequest{MaxWait: 25_000}},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
ids := proto.NewMessageIDGen(time.Now)
|
|
innerMsgID := ids.New(proto.MessageFromClient)
|
|
outerMsgID := ids.New(proto.MessageFromClient)
|
|
leaf := encodeClientMessageBodyForTest(t, tt.leaf)
|
|
gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
|
|
container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{
|
|
ID: innerMsgID, SeqNo: int(tt.seqNo), Bytes: len(gzipped), Body: gzipped,
|
|
}}})
|
|
|
|
cs := newConnState()
|
|
cs.track(innerMsgID, tt.seqNo, tt.content, msgStateReceived)
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
|
|
plan, err := s.preflightInbound(cs, outerMsgID, 2, container)
|
|
if err != nil {
|
|
t.Fatalf("preflight seen gzip inner: %v", err)
|
|
}
|
|
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || plan.items[0].content != tt.content {
|
|
plan.close()
|
|
t.Fatalf("seen gzip plan items = %+v, want one duplicate", plan.items)
|
|
}
|
|
if gotACK := len(plan.ackIDs) == 1 && plan.ackIDs[0] == innerMsgID; gotACK != tt.wantACK {
|
|
plan.close()
|
|
t.Fatalf("seen gzip ACKs = %v, wantACK=%v", plan.ackIDs, tt.wantACK)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != containerDescriptorBudgetBytes {
|
|
plan.close()
|
|
t.Fatalf("seen gzip held budget = %d, want descriptor-only %d", got, containerDescriptorBudgetBytes)
|
|
}
|
|
plan.close()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("seen gzip close leaked budget: %d", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSeenTopLevelContentGZIPShortCircuitsBeforeExpansion(t *testing.T) {
|
|
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
|
|
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 31})
|
|
wrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
|
|
cs := newConnState()
|
|
cs.track(msgID, 1, true, msgStateReceived)
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(1)
|
|
plan, err := s.preflightInbound(cs, msgID, 1, wrapped)
|
|
if err != nil {
|
|
t.Fatalf("preflight seen top-level gzip: %v", err)
|
|
}
|
|
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || len(plan.ackIDs) != 1 || plan.ackIDs[0] != msgID {
|
|
plan.close()
|
|
t.Fatalf("seen top-level gzip plan = items:%+v acks:%v", plan.items, plan.ackIDs)
|
|
}
|
|
plan.close()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("seen top-level gzip consumed expansion budget: %d", got)
|
|
}
|
|
}
|
|
|
|
func TestDuplicateOuterContainerOnlyInspectsInnerDescriptors(t *testing.T) {
|
|
ids := proto.NewMessageIDGen(time.Now)
|
|
innerMsgID := ids.New(proto.MessageFromClient)
|
|
outerMsgID := ids.New(proto.MessageFromClient)
|
|
leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 4})
|
|
gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf})
|
|
container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{
|
|
ID: innerMsgID, SeqNo: 1, Bytes: len(gzipped), Body: gzipped,
|
|
}}})
|
|
wrappedContainer := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: container})
|
|
|
|
t.Run("seen inner replays without expansion", func(t *testing.T) {
|
|
cs := newConnState()
|
|
cs.track(outerMsgID, 2, false, msgStateReceived)
|
|
cs.track(innerMsgID, 1, true, msgStateReceived)
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
|
|
plan, err := s.preflightInbound(cs, outerMsgID, 2, container)
|
|
if err != nil {
|
|
t.Fatalf("preflight duplicate outer: %v", err)
|
|
}
|
|
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate {
|
|
plan.close()
|
|
t.Fatalf("duplicate outer items = %+v", plan.items)
|
|
}
|
|
plan.close()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("duplicate outer leaked descriptor budget: %d", got)
|
|
}
|
|
})
|
|
|
|
t.Run("top-level gzip container still locates inner", func(t *testing.T) {
|
|
cs := newConnState()
|
|
cs.track(outerMsgID, 2, false, msgStateReceived)
|
|
cs.track(innerMsgID, 1, true, msgStateReceived)
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes + containerDescriptorBudgetBytes)
|
|
plan, err := s.preflightInbound(cs, outerMsgID, 2, wrappedContainer)
|
|
if err != nil {
|
|
t.Fatalf("preflight gzip duplicate outer: %v", err)
|
|
}
|
|
if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || plan.items[0].msgID != innerMsgID {
|
|
plan.close()
|
|
t.Fatalf("gzip duplicate outer items = %+v, want inner duplicate", plan.items)
|
|
}
|
|
plan.close()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("gzip duplicate outer leaked expansion/descriptor budget: %d", got)
|
|
}
|
|
})
|
|
|
|
t.Run("unseen inner rejects without expansion", func(t *testing.T) {
|
|
cs := newConnState()
|
|
cs.track(outerMsgID, 2, false, msgStateReceived)
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes)
|
|
plan, err := s.preflightInbound(cs, outerMsgID, 2, container)
|
|
if plan != nil {
|
|
plan.close()
|
|
t.Fatal("duplicate outer with unseen inner unexpectedly produced a plan")
|
|
}
|
|
var bad *dispatchBadMsgError
|
|
if !errors.As(err, &bad) || bad.code != badMsgContainer {
|
|
t.Fatalf("duplicate outer unseen-inner err = %v, want container rejection", err)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("unseen duplicate inner leaked descriptor/expansion budget: %d", got)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestServiceInfoViewsRejectOversizedBytesWithoutDecodeCopy(t *testing.T) {
|
|
state := mt.MsgsStateInfo{ReqMsgID: 7, Info: make([]byte, maxServiceMessageIDs)}
|
|
var encodedState bin.Buffer
|
|
if err := state.Encode(&encodedState); err != nil {
|
|
t.Fatalf("encode msgs_state_info: %v", err)
|
|
}
|
|
reqMsgID, info, err := msgsStateInfoView(&encodedState)
|
|
if err != nil || reqMsgID != state.ReqMsgID || len(info) != maxServiceMessageIDs {
|
|
t.Fatalf("state info view = id %d len %d err %v", reqMsgID, len(info), err)
|
|
}
|
|
info[0] = 0x7f
|
|
if encodedState.Buf[16] != 0x7f {
|
|
t.Fatal("msgs_state_info view unexpectedly copied info")
|
|
}
|
|
|
|
state.Info = make([]byte, maxServiceMessageIDs+1)
|
|
encodedState.Reset()
|
|
if err := state.Encode(&encodedState); err != nil {
|
|
t.Fatalf("encode oversized msgs_state_info: %v", err)
|
|
}
|
|
if _, _, err := msgsStateInfoView(&encodedState); err == nil || !strings.Contains(err.Error(), "exceeds") {
|
|
t.Fatalf("oversized msgs_state_info err = %v, want capped rejection", err)
|
|
}
|
|
|
|
all := mt.MsgsAllInfo{MsgIDs: []int64{1, 2}, Info: []byte{4, 4}}
|
|
var encodedAll bin.Buffer
|
|
if err := all.Encode(&encodedAll); err != nil {
|
|
t.Fatalf("encode msgs_all_info: %v", err)
|
|
}
|
|
count, allInfo, err := msgsAllInfoView(&encodedAll)
|
|
if err != nil || count != 2 || len(allInfo) != 2 {
|
|
t.Fatalf("all info view = count %d len %d err %v", count, len(allInfo), err)
|
|
}
|
|
all.Info = make([]byte, maxServiceMessageIDs+1)
|
|
encodedAll.Reset()
|
|
if err := all.Encode(&encodedAll); err != nil {
|
|
t.Fatalf("encode oversized msgs_all_info: %v", err)
|
|
}
|
|
if _, _, err := msgsAllInfoView(&encodedAll); err == nil || !strings.Contains(err.Error(), "exceeds") {
|
|
t.Fatalf("oversized msgs_all_info err = %v, want capped rejection", err)
|
|
}
|
|
}
|
|
|
|
func TestDispatchRejectsExcessiveWrapperDepthBeforeRPC(t *testing.T) {
|
|
var body bin.Buffer
|
|
if err := (&mt.MsgsStateInfo{ReqMsgID: 1, Info: []byte{4}}).Encode(&body); err != nil {
|
|
t.Fatalf("encode leaf: %v", err)
|
|
}
|
|
encoded := body.Copy()
|
|
for i := 0; i < maxDispatchDepth+1; i++ {
|
|
var wrapped bin.Buffer
|
|
if err := (proto.GZIP{Data: encoded}).Encode(&wrapped); err != nil {
|
|
t.Fatalf("encode gzip depth %d: %v", i+1, err)
|
|
}
|
|
encoded = wrapped.Copy()
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
var acks []int64
|
|
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
|
|
err := s.dispatch(context.Background(), newConnState(), nil, msgID, 0, &bin.Buffer{Buf: encoded}, &acks)
|
|
if err == nil || !strings.Contains(err.Error(), "wrapper depth") {
|
|
t.Fatalf("deep wrapper err = %v, want wrapper depth rejection", err)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("deep wrapper rejection leaked frame budget: %d", got)
|
|
}
|
|
}
|
|
|
|
func TestNestedWrapperTailFailuresReleaseEveryBudget(t *testing.T) {
|
|
ping := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 5})
|
|
|
|
var gzipWithTail bin.Buffer
|
|
if err := (&proto.GZIP{Data: ping}).Encode(&gzipWithTail); err != nil {
|
|
t.Fatalf("encode inner gzip: %v", err)
|
|
}
|
|
gzipWithTail.PutInt(0x11223344)
|
|
gzipTailWrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: gzipWithTail.Copy()})
|
|
|
|
var containerWithTail bin.Buffer
|
|
if err := (&proto.MessageContainer{}).Encode(&containerWithTail); err != nil {
|
|
t.Fatalf("encode inner container: %v", err)
|
|
}
|
|
containerWithTail.PutInt(0x55667788)
|
|
containerTailWrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: containerWithTail.Copy()})
|
|
|
|
tests := []struct {
|
|
name string
|
|
body []byte
|
|
seqNo int32
|
|
}{
|
|
{name: "gzip tail after outer expansion", body: gzipTailWrapped, seqNo: 1},
|
|
{name: "container tail after outer expansion", body: containerTailWrapped, seqNo: 0},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
|
|
plan, err := s.preflightInbound(newConnState(), msgID, tt.seqNo, tt.body)
|
|
if plan != nil {
|
|
plan.close()
|
|
t.Fatal("invalid nested wrapper unexpectedly produced a plan")
|
|
}
|
|
if err == nil || !strings.Contains(err.Error(), "trailing") {
|
|
t.Fatalf("nested wrapper err = %v, want trailing-byte rejection", err)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("nested wrapper failure leaked frame budget: %d", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOversizedConnectionBuffersAreReleasedAfterFrame(t *testing.T) {
|
|
inbound := &bin.Buffer{Buf: make([]byte, 1, maxRetainedConnBuffer+1)}
|
|
trimOversizedInboundBuffer(inbound)
|
|
if inbound.Buf != nil {
|
|
t.Fatalf("oversized inbound buffer cap=%d, want released", cap(inbound.Buf))
|
|
}
|
|
regular := &bin.Buffer{Buf: make([]byte, 1, maxRetainedConnBuffer)}
|
|
trimOversizedInboundBuffer(regular)
|
|
if cap(regular.Buf) != maxRetainedConnBuffer {
|
|
t.Fatalf("regular inbound buffer cap=%d, want retained", cap(regular.Buf))
|
|
}
|
|
|
|
pool := newOutboundScratchPool(16 << 20)
|
|
scratch, err := pool.acquire(context.Background(), nil, maxRetainedConnBuffer+1)
|
|
if err != nil {
|
|
t.Fatalf("acquire oversized outbound scratch: %v", err)
|
|
}
|
|
pool.release(scratch)
|
|
if got := pool.snapshot(); got != 0 {
|
|
t.Fatalf("oversized outbound scratch retained %d bytes, want 0", got)
|
|
}
|
|
}
|
|
|
|
func TestGZIPExpansionUsesProcessBudgetBeforeDecode(t *testing.T) {
|
|
payload := make([]byte, 1<<20)
|
|
var wrapped bin.Buffer
|
|
if err := (proto.GZIP{Data: payload}).Encode(&wrapped); err != nil {
|
|
t.Fatalf("encode gzip: %v", err)
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes - 1)
|
|
if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
|
|
release()
|
|
t.Fatal("gzip decode unexpectedly bypassed saturated process budget")
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("failed gzip reservation leaked %d bytes", got)
|
|
}
|
|
|
|
s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
|
|
decoded, release, err := s.decodeGZIPWithGlobalBudget(&wrapped)
|
|
if err != nil {
|
|
t.Fatalf("budgeted gzip decode: %v", err)
|
|
}
|
|
if len(decoded) != len(payload) {
|
|
t.Fatalf("decoded bytes = %d, want %d", len(decoded), len(payload))
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != int64(len(payload)) {
|
|
t.Fatalf("held expansion budget = %d, want %d", got, len(payload))
|
|
}
|
|
release()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("released expansion budget = %d, want zero", got)
|
|
}
|
|
|
|
s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
|
|
if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)-1); err == nil {
|
|
release()
|
|
t.Fatal("caller-bounded gzip decode accepted an oversized expansion")
|
|
} else if got := gzipExpansionWork(err); got != len(payload) {
|
|
t.Fatalf("caller-bounded rejected expansion work = %d, want %d", got, len(payload))
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("caller-bounded rejection leaked %d bytes", got)
|
|
}
|
|
|
|
s.frameBudget = newInboundFrameBudget(int64(len(payload) - 1))
|
|
if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
|
|
release()
|
|
t.Fatalf("caller-bounded process budget error = %v, want ErrInboundFrameBudgetExceeded", err)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("caller-bounded reservation failure leaked %d bytes", got)
|
|
}
|
|
}
|
|
|
|
func TestGZIPPackedAcceptsOfficialClientWrappers(t *testing.T) {
|
|
payload := []byte("official Telegram gzip_packed payload")
|
|
tests := []struct {
|
|
name string
|
|
packed bin.Encoder
|
|
}{
|
|
{name: "tdlib_zlib", packed: zlibPackedObjectForTest(t, payload)},
|
|
{name: "drklo_gotd_gzip", packed: &proto.GZIP{Data: payload}},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var wrapped bin.Buffer
|
|
if err := tt.packed.Encode(&wrapped); err != nil {
|
|
t.Fatalf("encode gzip_packed: %v", err)
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
|
|
decoded, release, err := s.decodeGZIPWithGlobalBudget(&wrapped)
|
|
if err != nil {
|
|
t.Fatalf("decode gzip_packed: %v", err)
|
|
}
|
|
if !bytes.Equal(decoded, payload) {
|
|
t.Fatalf("decoded payload = %q, want %q", decoded, payload)
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != int64(len(payload)) {
|
|
t.Fatalf("held expansion budget = %d, want %d", got, len(payload))
|
|
}
|
|
release()
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("released expansion budget = %d, want zero", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGZIPPackedRejectsCorruptZLIBChecksumWithoutBudgetLeak(t *testing.T) {
|
|
packed := zlibPackedObjectForTest(t, []byte("checksum-protected payload"))
|
|
packed.PackedData[len(packed.PackedData)-1] ^= 0xff
|
|
var wrapped bin.Buffer
|
|
if err := packed.Encode(&wrapped); err != nil {
|
|
t.Fatalf("encode gzip_packed: %v", err)
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
|
|
if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
|
|
release()
|
|
t.Fatal("corrupt zlib checksum unexpectedly accepted")
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("corrupt zlib checksum leaked %d budget bytes", got)
|
|
}
|
|
}
|
|
|
|
func TestGZIPPackedZLIBHonorsExpansionLimit(t *testing.T) {
|
|
payload := make([]byte, 1<<20)
|
|
packed := zlibPackedObjectForTest(t, payload)
|
|
var wrapped bin.Buffer
|
|
if err := packed.Encode(&wrapped); err != nil {
|
|
t.Fatalf("encode gzip_packed: %v", err)
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
|
|
if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)-1); err == nil {
|
|
release()
|
|
t.Fatal("caller-bounded zlib decode accepted an oversized expansion")
|
|
} else if got := gzipExpansionWork(err); got != len(payload) {
|
|
t.Fatalf("caller-bounded zlib expansion work = %d, want %d", got, len(payload))
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("caller-bounded zlib rejection leaked %d bytes", got)
|
|
}
|
|
}
|
|
|
|
func TestGZIPPackedRejectsRawDEFLATEWithoutBudgetLeak(t *testing.T) {
|
|
var compressed bytes.Buffer
|
|
w, err := flate.NewWriter(&compressed, flate.DefaultCompression)
|
|
if err != nil {
|
|
t.Fatalf("create raw deflate writer: %v", err)
|
|
}
|
|
if _, err := w.Write([]byte("raw deflate is not a supported Telegram wrapper")); err != nil {
|
|
t.Fatalf("write raw deflate payload: %v", err)
|
|
}
|
|
if err := w.Close(); err != nil {
|
|
t.Fatalf("close raw deflate payload: %v", err)
|
|
}
|
|
packed := &mt.GzipPacked{PackedData: compressed.Bytes()}
|
|
var wrapped bin.Buffer
|
|
if err := packed.Encode(&wrapped); err != nil {
|
|
t.Fatalf("encode gzip_packed: %v", err)
|
|
}
|
|
|
|
s := New(Options{Logger: zaptest.NewLogger(t)})
|
|
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
|
|
if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
|
|
release()
|
|
t.Fatal("raw deflate unexpectedly accepted")
|
|
}
|
|
if got := s.frameBudget.usedBytes(); got != 0 {
|
|
t.Fatalf("raw deflate rejection leaked %d budget bytes", got)
|
|
}
|
|
}
|
|
|
|
func zlibPackedObjectForTest(t testing.TB, payload []byte) *mt.GzipPacked {
|
|
t.Helper()
|
|
var compressed bytes.Buffer
|
|
w := zlib.NewWriter(&compressed)
|
|
if _, err := w.Write(payload); err != nil {
|
|
t.Fatalf("write zlib payload: %v", err)
|
|
}
|
|
if err := w.Close(); err != nil {
|
|
t.Fatalf("close zlib payload: %v", err)
|
|
}
|
|
return &mt.GzipPacked{PackedData: compressed.Bytes()}
|
|
}
|