internal/mtprotoedge/frame_budget_test.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"bytes"
"context"
"encoding/binary"
"errors"
"io"
"net"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
type frameBudgetTestConn struct {
reader bytes.Reader
read int
closed bool
}
func newFrameBudgetTestConn(packet []byte) *frameBudgetTestConn {
c := &frameBudgetTestConn{}
c.reader.Reset(packet)
return c
}
func (c *frameBudgetTestConn) Read(p []byte) (int, error) {
n, err := c.reader.Read(p)
c.read += n
return n, err
}
func (*frameBudgetTestConn) Write(p []byte) (int, error) { return len(p), nil }
func (c *frameBudgetTestConn) Close() error {
c.closed = true
return nil
}
func (*frameBudgetTestConn) LocalAddr() net.Addr { return frameBudgetTestAddr("local") }
func (*frameBudgetTestConn) RemoteAddr() net.Addr { return frameBudgetTestAddr("remote") }
func (*frameBudgetTestConn) SetDeadline(time.Time) error { return nil }
func (*frameBudgetTestConn) SetReadDeadline(time.Time) error { return nil }
func (*frameBudgetTestConn) SetWriteDeadline(time.Time) error { return nil }
type frameBudgetTestAddr string
func (a frameBudgetTestAddr) Network() string { return "frame-budget-test" }
func (a frameBudgetTestAddr) String() string { return string(a) }
func newFrameBudgetTestTransport(packet []byte, c transport.Codec, budget *inboundFrameBudget) (*compatTransportConn, *frameBudgetTestConn) {
raw := newFrameBudgetTestConn(packet)
return &compatTransportConn{
conn: raw,
codec: c,
codecKind: classifyInboundFrameCodec(c),
budgetedCodec: unwrapInboundFrameBudgetedCodec(c),
budget: budget,
}, raw
}
func TestInboundFrameBudgetSupportsBuiltInCodecs(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
abridged := append([]byte{byte(len(payload) / bin.Word)}, payload...)
intermediate := make([]byte, bin.Word+len(payload))
binary.LittleEndian.PutUint32(intermediate, uint32(len(payload)))
copy(intermediate[bin.Word:], payload)
padded := make([]byte, bin.Word+len(payload)+1)
binary.LittleEndian.PutUint32(padded, uint32(len(payload)+1))
copy(padded[bin.Word:], payload)
padded[len(padded)-1] = 0xa5
var full bytes.Buffer
fullCodec := &codec.Full{}
fullPayload := &bin.Buffer{Buf: append([]byte(nil), payload...)}
if err := fullCodec.Write(&full, fullPayload); err != nil {
t.Fatalf("encode full frame: %v", err)
}
tests := []struct {
name string
packet []byte
codec transport.Codec
reservation int64
}{
{name: "abridged", packet: abridged, codec: &quickAckAbridgedCodec{}, reservation: 2 * int64(len(payload))},
{name: "intermediate", packet: intermediate, codec: &quickAckIntermediateCodec{}, reservation: 2 * int64(len(payload))},
{name: "padded_intermediate", packet: padded, codec: &quickAckPaddedIntermediateCodec{}, reservation: 2 * int64(len(payload)+1)},
{name: "full", packet: full.Bytes(), codec: &codec.Full{}, reservation: int64(full.Len() + len(payload))},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
budget := newInboundFrameBudget(tt.reservation)
conn, _ := newFrameBudgetTestTransport(tt.packet, tt.codec, budget)
var got bin.Buffer
if err := conn.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv: %v", err)
}
if !bytes.Equal(got.Raw(), payload) {
t.Fatalf("payload = %x, want %x", got.Raw(), payload)
}
if used := budget.usedBytes(); used != tt.reservation {
t.Fatalf("held budget = %d, want %d", used, tt.reservation)
}
if err := conn.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
if used := budget.usedBytes(); used != tt.reservation {
t.Fatalf("budget after concurrent Close = %d, want delivered ownership %d", used, tt.reservation)
}
conn.releaseInboundFrame()
if used := budget.usedBytes(); used != 0 {
t.Fatalf("budget after ownership release = %d, want 0", used)
}
})
}
}
func TestInboundFrameBudgetRejectsBeforePayloadAllocation(t *testing.T) {
const payloadBytes = 1 << 20
var header [bin.Word]byte
binary.LittleEndian.PutUint32(header[:], payloadBytes)
budget := newInboundFrameBudget(2*payloadBytes - 1)
conn, raw := newFrameBudgetTestTransport(header[:], &quickAckIntermediateCodec{}, budget)
var got bin.Buffer
err := conn.Recv(context.Background(), &got)
if !errors.Is(err, ErrInboundFrameBudgetExceeded) {
t.Fatalf("Recv error = %v, want ErrInboundFrameBudgetExceeded", err)
}
if raw.read != bin.Word {
t.Fatalf("wire bytes read = %d, want only %d-byte length prefix", raw.read, bin.Word)
}
if cap(got.Buf) != 0 {
t.Fatalf("payload buffer capacity = %d, want 0 before admission", cap(got.Buf))
}
if used := budget.usedBytes(); used != 0 {
t.Fatalf("budget after rejected preflight = %d, want 0", used)
}
}
func TestInboundFrameBudgetAbridgedPreflightMatchesCodecSemantics(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
quickPacket := append([]byte{0x80 | byte(len(payload)/bin.Word)}, payload...)
quickBudget := newInboundFrameBudget(int64(2 * len(payload)))
quick, _ := newFrameBudgetTestTransport(quickPacket, &quickAckAbridgedCodec{}, quickBudget)
var got bin.Buffer
if err := quick.Recv(context.Background(), &got); err != nil {
t.Fatalf("quick-ack abridged Recv: %v", err)
}
requested := quick.ConsumeQuickAckRequested()
if !bytes.Equal(got.Raw(), payload) || !requested {
t.Fatalf("quick-ack frame = %x requested=%v", got.Raw(), requested)
}
quick.releaseInboundFrame()
_ = quick.Close()
// gotd's plain codec treats every first byte >= 0x7f as the extended form (it does not
// implement the quick-ack high bit). The preflight parser must mirror that behavior; treating
// 0x82 as a short two-word frame would let the codec allocate from the following three bytes.
malicious := []byte{0x82, 0xff, 0xff, 0xff}
plainBudget := newInboundFrameBudget(defaultInboundFrameGlobalMaxBytes)
plain, raw := newFrameBudgetTestTransport(malicious, codec.Abridged{}, plainBudget)
got.Reset()
err := plain.Recv(context.Background(), &got)
if err == nil {
t.Fatal("plain abridged accepted oversized extended length")
}
if raw.read != 4 || cap(got.Buf) > 2*bin.Word {
t.Fatalf("plain abridged read=%d buffer_cap=%d, want prefix-only allocation", raw.read, cap(got.Buf))
}
_ = plain.Close()
}
func TestInboundFrameBudgetReleasedAtNextRecvAndReusable(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
frame := make([]byte, bin.Word+len(payload))
binary.LittleEndian.PutUint32(frame, uint32(len(payload)))
copy(frame[bin.Word:], payload)
packet := append(append([]byte(nil), frame...), frame...)
reservation := int64(2 * len(payload))
budget := newInboundFrameBudget(reservation)
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
for i := 0; i < 2; i++ {
var got bin.Buffer
if err := conn.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv %d: %v", i+1, err)
}
if used := budget.usedBytes(); used != reservation {
t.Fatalf("held budget after frame %d = %d, want %d", i+1, used, reservation)
}
}
conn.releaseInboundFrame()
_ = conn.Close()
}
func TestInboundFrameRetainedBackingStaysChargedAcrossSmallFrame(t *testing.T) {
const largeBytes = 1 << 20
large := make([]byte, bin.Word+largeBytes)
binary.LittleEndian.PutUint32(large, largeBytes)
smallPayload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
small := make([]byte, bin.Word+len(smallPayload))
binary.LittleEndian.PutUint32(small, uint32(len(smallPayload)))
copy(small[bin.Word:], smallPayload)
packet := append(large, small...)
budget := newInboundFrameBudget(2 * largeBytes)
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
var wire bin.Buffer
if err := conn.Recv(context.Background(), &wire); err != nil {
t.Fatalf("large Recv: %v", err)
}
// Model decryptClientFrame's exact-size plaintext reuse buffer.
plain := bin.Buffer{Buf: make([]byte, largeBytes)}
retainInboundFrameBackings(conn, &wire, &plain)
retained := int64(cap(wire.Buf) + cap(plain.Buf))
if got := budget.usedBytes(); got != retained {
t.Fatalf("retained budget after large frame = %d, want capacities %d", got, retained)
}
wire.Reset()
if err := conn.Recv(context.Background(), &wire); err != nil {
t.Fatalf("small Recv: %v", err)
}
// The small announcement must not release the large backing's charge. This was the
// warm-many-connections bypass: each socket retained MiBs while the global budget saw bytes.
if got := budget.usedBytes(); got != retained {
t.Fatalf("budget after small frame = %d, want retained high-water %d", got, retained)
}
wire.Buf = nil
plain.Buf = nil
retainInboundFrameBackings(conn, &wire, &plain)
if got := budget.usedBytes(); got != 0 {
t.Fatalf("budget after dropping reusable backings = %d, want 0", got)
}
_ = conn.Close()
}
func TestInboundFrameBudgetClosePreservesDeliveredOwnershipUntilRelease(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
frame := make([]byte, bin.Word+len(payload))
binary.LittleEndian.PutUint32(frame, uint32(len(payload)))
copy(frame[bin.Word:], payload)
budget := newInboundFrameBudget(int64(2 * len(payload)))
first, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
var got bin.Buffer
if err := first.Recv(context.Background(), &got); err != nil {
t.Fatalf("first Recv: %v", err)
}
blocked, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
var blockedPayload bin.Buffer
if err := blocked.Recv(context.Background(), &blockedPayload); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
t.Fatalf("concurrent Recv error = %v, want global budget rejection", err)
}
if cap(blockedPayload.Buf) != 0 {
t.Fatalf("blocked connection allocated payload capacity %d", cap(blockedPayload.Buf))
}
_ = blocked.Close()
if err := first.Close(); err != nil {
t.Fatalf("first Close: %v", err)
}
if used := budget.usedBytes(); used != int64(2*len(payload)) {
t.Fatalf("budget after concurrent Close = %d, want delivered frame still charged", used)
}
second, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
got.Reset()
if err := second.Recv(context.Background(), &got); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
t.Fatalf("second Recv before ownership release = %v, want budget rejection", err)
}
_ = second.Close()
first.releaseInboundFrame()
third, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
got.Reset()
if err := third.Recv(context.Background(), &got); err != nil {
t.Fatalf("third Recv after ownership release: %v", err)
}
third.releaseInboundFrame()
_ = third.Close()
}
type unsafeFrameBudgetCodec struct {
readCalled bool
}
func (*unsafeFrameBudgetCodec) WriteHeader(io.Writer) error { return nil }
func (*unsafeFrameBudgetCodec) ReadHeader(io.Reader) error { return nil }
func (*unsafeFrameBudgetCodec) Write(io.Writer, *bin.Buffer) error { return nil }
func (c *unsafeFrameBudgetCodec) Read(io.Reader, *bin.Buffer) error { c.readCalled = true; return nil }
func TestCustomCodecWithoutPreflightFailsClosed(t *testing.T) {
raw := newFrameBudgetTestConn([]byte{1, 2, 3, 4})
listener := newSingleConnListener(raw)
custom := &unsafeFrameBudgetCodec{}
budgeted := newCompatTransportListener(func() transport.Codec { return custom }, listener, newInboundFrameBudget(1024))
conn, err := budgeted.Accept()
if !errors.Is(err, errInboundFrameCodecUnsupported) {
t.Fatalf("Accept error = %v, want unsupported preflight codec", err)
}
if conn != nil {
t.Fatal("unsupported custom codec unexpectedly accepted")
}
if custom.readCalled || raw.read != 0 {
t.Fatalf("custom codec touched frame before rejection: read_called=%v wire_read=%d", custom.readCalled, raw.read)
}
}
func TestExplicitBuiltInCodecUsesBudgetedTransport(t *testing.T) {
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
packet := append([]byte(nil), codec.IntermediateClientStart[:]...)
var header [bin.Word]byte
binary.LittleEndian.PutUint32(header[:], uint32(len(payload)))
packet = append(packet, header[:]...)
packet = append(packet, payload...)
raw := newFrameBudgetTestConn(packet)
budget := newInboundFrameBudget(int64(2 * len(payload)))
listener := newCompatTransportListener(
func() transport.Codec { return codec.Intermediate{} },
newSingleConnListener(raw),
budget,
)
conn, err := listener.Accept()
if err != nil {
t.Fatalf("Accept: %v", err)
}
var got bin.Buffer
if err := conn.Recv(context.Background(), &got); err != nil {
t.Fatalf("Recv: %v", err)
}
if !bytes.Equal(got.Raw(), payload) || budget.usedBytes() != int64(2*len(payload)) {
t.Fatalf("payload=%x budget=%d", got.Raw(), budget.usedBytes())
}
conn.(*compatTransportConn).releaseInboundFrame()
_ = conn.Close()
}