internal/mtprotoedge/frame_budget.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"encoding/binary"
"errors"
"fmt"
"io"
"sync/atomic"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
const defaultInboundFrameGlobalMaxBytes int64 = 512 << 20
var (
// ErrInboundFrameBudgetExceeded means the process-wide wire+plaintext reservation for a
// newly announced transport frame could not be acquired. The length prefix has been read,
// but the payload buffer has not been allocated and the connection must be closed.
ErrInboundFrameBudgetExceeded = errors.New("inbound frame global byte budget exceeded")
errInboundFrameCodecUnsupported = errors.New("transport codec cannot preflight inbound frame length")
errInboundFrameNotReserved = errors.New("transport codec returned a frame without reserving inbound bytes")
)
// InboundFrameBudgetedCodec is the fail-safe extension point for a custom Options.Codec.
// Implementations must parse and validate the frame length, call reserve exactly once before
// allocating or growing the payload buffer, and keep the reservation valid until Read returns.
// Built-in abridged/intermediate/padded-intermediate/full codecs are recognized directly.
type InboundFrameBudgetedCodec interface {
transport.Codec
ReadWithInboundFrameBudget(r io.Reader, b *bin.Buffer, reserve func(wireBytes, plaintextBytes int64) error) error
}
// inboundFrameBudget accounts the two per-frame buffers that can coexist while an encrypted
// request is handled: transport/wire bytes and decrypted plaintext. It deliberately charges the
// maximum plaintext size announced by framing even for an unencrypted handshake frame; that
// conservative rule makes admission independent of auth state and prevents allocation before
// auth_key_id can be inspected.
type inboundFrameBudget struct {
max int64
used atomic.Int64
}
func newInboundFrameBudget(max int64) *inboundFrameBudget {
if max <= 0 {
max = defaultInboundFrameGlobalMaxBytes
}
return &inboundFrameBudget{max: max}
}
func (b *inboundFrameBudget) reserve(wireBytes, plaintextBytes int64) (int64, error) {
return b.growReservation(0, wireBytes, plaintextBytes)
}
// growReservation atomically raises one connection's existing retained/frame reservation to
// cover a newly announced frame. Keeping the old charge until this transition is what makes a
// reused transport/plaintext backing remain accounted between frames; a small next frame cannot
// release a previously large allocation while still retaining its capacity.
func (b *inboundFrameBudget) growReservation(current, wireBytes, plaintextBytes int64) (int64, error) {
if current < 0 || wireBytes <= 0 || plaintextBytes < 0 || wireBytes > b.max || plaintextBytes > b.max-wireBytes {
return 0, fmt.Errorf("%w: wire=%d plaintext=%d limit=%d", ErrInboundFrameBudgetExceeded, wireBytes, plaintextBytes, b.max)
}
target := wireBytes + plaintextBytes
if target <= current {
return current, nil
}
n := target - current
for {
used := b.used.Load()
if n > b.max-used {
return 0, fmt.Errorf("%w: requested=%d used=%d limit=%d", ErrInboundFrameBudgetExceeded, n, used, b.max)
}
if b.used.CompareAndSwap(used, used+n) {
return target, nil
}
}
}
func (b *inboundFrameBudget) release(n int64) {
if n == 0 {
return
}
used := b.used.Add(-n)
if used < 0 {
// This is an internal ownership invariant, not recoverable input. A negative value would
// silently disable admission for subsequent frames, so fail loudly during development.
panic("mtprotoedge: inbound frame budget released more than reserved")
}
}
func (b *inboundFrameBudget) usedBytes() int64 {
return b.used.Load()
}
type inboundFrameCodecKind uint8
const (
inboundFrameCodecUnknown inboundFrameCodecKind = iota
inboundFrameCodecQuickAckAbridged
inboundFrameCodecAbridged
inboundFrameCodecIntermediate
inboundFrameCodecPaddedIntermediate
inboundFrameCodecFull
inboundFrameCodecCustom
)
func classifyInboundFrameCodec(c transport.Codec) inboundFrameCodecKind {
switch v := c.(type) {
case *quickAckAbridgedCodec:
return inboundFrameCodecQuickAckAbridged
case codec.Abridged, *codec.Abridged:
return inboundFrameCodecAbridged
case *quickAckIntermediateCodec, codec.Intermediate, *codec.Intermediate:
return inboundFrameCodecIntermediate
case *quickAckPaddedIntermediateCodec, codec.PaddedIntermediate, *codec.PaddedIntermediate:
return inboundFrameCodecPaddedIntermediate
case *codec.Full:
return inboundFrameCodecFull
case codec.NoHeader:
return classifyInboundFrameCodec(v.Codec)
case *codec.NoHeader:
if v == nil {
return inboundFrameCodecUnknown
}
return classifyInboundFrameCodec(v.Codec)
case InboundFrameBudgetedCodec:
return inboundFrameCodecCustom
default:
return inboundFrameCodecUnknown
}
}
func unwrapInboundFrameBudgetedCodec(c transport.Codec) InboundFrameBudgetedCodec {
switch v := c.(type) {
case InboundFrameBudgetedCodec:
return v
case codec.NoHeader:
return unwrapInboundFrameBudgetedCodec(v.Codec)
case *codec.NoHeader:
if v != nil {
return unwrapInboundFrameBudgetedCodec(v.Codec)
}
}
return nil
}
// inboundFramePreflightReader consumes only the framing length prefix, reserves the announced
// wire+plaintext bytes, and only then exposes the final prefix bytes to the codec. Consequently a
// budget error is observed by codec.Read before it can ResetN/Expand the payload buffer.
type inboundFramePreflightReader struct {
r io.Reader
kind inboundFrameCodecKind
reserve func(wireBytes, plaintextBytes int64) error
abridgedFirstDelivered bool
done bool
}
func (r *inboundFramePreflightReader) Read(p []byte) (int, error) {
if len(p) == 0 {
return 0, nil
}
if r.done {
return r.r.Read(p)
}
switch r.kind {
case inboundFrameCodecQuickAckAbridged:
return r.readAbridgedPrefix(p, true)
case inboundFrameCodecAbridged:
return r.readAbridgedPrefix(p, false)
case inboundFrameCodecIntermediate, inboundFrameCodecPaddedIntermediate:
return r.readWordPrefix(p, false)
case inboundFrameCodecFull:
return r.readWordPrefix(p, true)
default:
return 0, errInboundFrameCodecUnsupported
}
}
func (r *inboundFramePreflightReader) readAbridgedPrefix(p []byte, quickAck bool) (int, error) {
if !r.abridgedFirstDelivered {
var first [1]byte
if _, err := io.ReadFull(r.r, first[:]); err != nil {
return 0, err
}
lengthByte := first[0]
extended := lengthByte >= 0x7f
if quickAck {
lengthByte &= 0x7f
extended = lengthByte == 0x7f
}
if !extended {
n := int64(lengthByte) * bin.Word
if err := reserveCompatFrame(r.reserve, n, n); err != nil {
return 0, err
}
r.done = true
}
r.abridgedFirstDelivered = true
p[0] = first[0]
return 1, nil
}
var tail [3]byte
if _, err := io.ReadFull(r.r, tail[:]); err != nil {
return 0, err
}
words := uint32(tail[0]) | uint32(tail[1])<<8 | uint32(tail[2])<<16
n := int64(words) * bin.Word
if err := reserveCompatFrame(r.reserve, n, n); err != nil {
return 0, err
}
r.done = true
return copy(p, tail[:]), nil
}
func (r *inboundFramePreflightReader) readWordPrefix(p []byte, full bool) (int, error) {
var header [bin.Word]byte
if _, err := io.ReadFull(r.r, header[:]); err != nil {
return 0, err
}
raw := int64(binary.LittleEndian.Uint32(header[:]))
var wireBytes, plaintextBytes int64
if full {
// Full transport length includes length + sequence + payload + CRC.
if raw < 3*bin.Word || raw > maxTransportMessageSize {
return 0, fmt.Errorf("invalid full transport message length %d", raw)
}
wireBytes = raw
plaintextBytes = raw - 3*bin.Word
} else {
wireBytes = raw &^ int64(quickAckResponseFlag)
plaintextBytes = wireBytes
}
if err := reserveCompatFrame(r.reserve, wireBytes, plaintextBytes); err != nil {
return 0, err
}
r.done = true
return copy(p, header[:]), nil
}
func reserveCompatFrame(reserve func(wireBytes, plaintextBytes int64) error, wireBytes, plaintextBytes int64) error {
if wireBytes <= 0 || wireBytes > maxTransportMessageSize {
return fmt.Errorf("invalid transport message length %d", wireBytes)
}
return reserve(wireBytes, plaintextBytes)
}