fix(mtprotoedge): sync admit nested gzip RPC envelopes
This commit is contained in:
parent
9cb76b7c5a
commit
e756b0f9f9
12 changed files with 783 additions and 26 deletions
|
|
@ -349,6 +349,35 @@ func (e *dispatchBadMsgError) Error() string {
|
|||
return fmt.Sprintf("bad client message %d/%d: code %d", e.msgID, e.seqNo, e.code)
|
||||
}
|
||||
|
||||
var errGZIPExpansionLimit = errors.New("gzip expansion limit exceeded")
|
||||
|
||||
type gzipExpansionWorkError struct {
|
||||
expanded int
|
||||
cause error
|
||||
}
|
||||
|
||||
func (e *gzipExpansionWorkError) Error() string {
|
||||
if e == nil || e.cause == nil {
|
||||
return "gzip expansion failed"
|
||||
}
|
||||
return e.cause.Error()
|
||||
}
|
||||
|
||||
func (e *gzipExpansionWorkError) Unwrap() error {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return e.cause
|
||||
}
|
||||
|
||||
func gzipExpansionWork(err error) int {
|
||||
var work *gzipExpansionWorkError
|
||||
if errors.As(err, &work) && work.expanded > 0 {
|
||||
return work.expanded
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// decodeGZIPWithGlobalBudget reserves the maximum single-wrapper output before
|
||||
// decompression starts. Once the actual size is known the excess reservation is
|
||||
// returned, while the actual output remains charged until the inbound plan is
|
||||
|
|
@ -356,6 +385,16 @@ func (e *dispatchBadMsgError) Error() string {
|
|||
// This closes the gap where every connection read goroutine could otherwise hold
|
||||
// an unaccounted 10 MiB expansion before the shared RPC scheduler saw the body.
|
||||
func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), error) {
|
||||
return s.decodeGZIPWithGlobalBudgetLimit(b, maxSingleGZIPExpandedBytes)
|
||||
}
|
||||
|
||||
// decodeGZIPWithGlobalBudgetLimit is the caller-bounded form used by exact
|
||||
// Layer admission. limit is also capped by the protocol's single-wrapper
|
||||
// ceiling; the returned bytes remain charged until release is called.
|
||||
func (s *Server) decodeGZIPWithGlobalBudgetLimit(b *bin.Buffer, limit int) ([]byte, func(), error) {
|
||||
if limit <= 0 || limit > maxSingleGZIPExpandedBytes {
|
||||
return nil, func() {}, fmt.Errorf("invalid gzip expansion limit %d", limit)
|
||||
}
|
||||
compressed, err := gzipPackedBytesView(b)
|
||||
if err != nil {
|
||||
return nil, func() {}, err
|
||||
|
|
@ -368,7 +407,7 @@ func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), erro
|
|||
}
|
||||
}
|
||||
if s.frameBudget != nil {
|
||||
reserved, err = s.frameBudget.reserve(maxSingleGZIPExpandedBytes, 0)
|
||||
reserved, err = s.frameBudget.reserve(int64(limit), 0)
|
||||
if err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
|
|
@ -379,19 +418,22 @@ func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), erro
|
|||
release()
|
||||
return nil, func() {}, err
|
||||
}
|
||||
data, readErr := io.ReadAll(io.LimitReader(r, maxSingleGZIPExpandedBytes+1))
|
||||
data, readErr := io.ReadAll(io.LimitReader(r, int64(limit)+1))
|
||||
closeErr := r.Close()
|
||||
if readErr != nil {
|
||||
release()
|
||||
return nil, func() {}, readErr
|
||||
return nil, func() {}, &gzipExpansionWorkError{expanded: len(data), cause: readErr}
|
||||
}
|
||||
if closeErr != nil {
|
||||
release()
|
||||
return nil, func() {}, closeErr
|
||||
return nil, func() {}, &gzipExpansionWorkError{expanded: len(data), cause: closeErr}
|
||||
}
|
||||
if len(data) > maxSingleGZIPExpandedBytes {
|
||||
if len(data) > limit {
|
||||
release()
|
||||
return nil, func() {}, fmt.Errorf("gzip expansion %d exceeds %d", len(data), maxSingleGZIPExpandedBytes)
|
||||
return nil, func() {}, &gzipExpansionWorkError{
|
||||
expanded: len(data),
|
||||
cause: fmt.Errorf("%w: expansion %d exceeds %d", errGZIPExpansionLimit, len(data), limit),
|
||||
}
|
||||
}
|
||||
if reserved > int64(len(data)) {
|
||||
s.frameBudget.release(reserved - int64(len(data)))
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ var inboundLayerDecodeLimits = tlprofile.Limits{
|
|||
}
|
||||
|
||||
var errDefaultLayerAdmission = errors.New("selected layer profile rejected naked RPC")
|
||||
var errLayerRPCGZIPCapacity = errors.New("exact layer gzip admission capacity exhausted")
|
||||
var errLayerRPCAdmissionCapability = errors.New("exact layer RPC admission capability unavailable")
|
||||
|
||||
const (
|
||||
maxLayerRPCDependencyIDs = 128
|
||||
|
|
@ -172,6 +174,68 @@ func layerRPCAdmissionReservationSize(wireBytes int) int {
|
|||
return wireBytes*layerRPCAdmissionWireFactor + layerRPCAdmissionGraphSlack
|
||||
}
|
||||
|
||||
// layerRPCGZIPExpansionBudget bridges transient process-wide expansion memory
|
||||
// to the durable scheduler charge of one exact typed request. sourceBytes is
|
||||
// conservative: it retains the original compressed wire plus every successful
|
||||
// expansion seen across nested envelopes or an authoritative-profile re-decode.
|
||||
type layerRPCGZIPExpansionBudget struct {
|
||||
server *Server
|
||||
plan *inboundPlan
|
||||
reservation *inboundRPCBatchReservation
|
||||
entry int
|
||||
baseSourceBytes int
|
||||
attemptExpandedBytes int
|
||||
chargedSourceBytes int
|
||||
}
|
||||
|
||||
func (b *layerRPCGZIPExpansionBudget) beginAttempt() {
|
||||
if b != nil {
|
||||
b.attemptExpandedBytes = 0
|
||||
}
|
||||
}
|
||||
|
||||
func (b *layerRPCGZIPExpansionBudget) expand(wire []byte, admissionLimit int) ([]byte, func(), error) {
|
||||
noop := func() {}
|
||||
if b == nil || b.server == nil || b.plan == nil || b.reservation == nil {
|
||||
return nil, noop, errors.New("invalid exact layer gzip expansion budget")
|
||||
}
|
||||
frameRemaining := maxDispatchExpandedBytes - b.plan.gzipExpandedBytes
|
||||
limit := min(admissionLimit, frameRemaining)
|
||||
if limit <= 0 {
|
||||
return nil, noop, errors.Join(
|
||||
errLayerRPCGZIPCapacity,
|
||||
fmt.Errorf("cumulative gzip expansion reached %d bytes", maxDispatchExpandedBytes),
|
||||
)
|
||||
}
|
||||
data, release, err := b.server.decodeGZIPWithGlobalBudgetLimit(&bin.Buffer{Buf: wire}, limit)
|
||||
if err != nil {
|
||||
if work := gzipExpansionWork(err); work > 0 {
|
||||
b.plan.gzipExpandedBytes += work
|
||||
}
|
||||
if errors.Is(err, ErrInboundFrameBudgetExceeded) {
|
||||
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
|
||||
}
|
||||
if limit < admissionLimit && errors.Is(err, errGZIPExpansionLimit) {
|
||||
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
|
||||
}
|
||||
return nil, release, err
|
||||
}
|
||||
b.plan.gzipExpandedBytes += len(data)
|
||||
b.attemptExpandedBytes += len(data)
|
||||
targetSourceBytes := b.baseSourceBytes + b.attemptExpandedBytes
|
||||
targetCharge := layerRPCAdmissionReservationSize(targetSourceBytes)
|
||||
if targetSourceBytes > b.chargedSourceBytes {
|
||||
if err := b.reservation.growEntry(b.entry, targetCharge); err != nil {
|
||||
if errors.Is(err, ErrInboundRPCQueueFull) {
|
||||
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
|
||||
}
|
||||
return nil, release, err
|
||||
}
|
||||
b.chargedSourceBytes = targetSourceBytes
|
||||
}
|
||||
return data, release, nil
|
||||
}
|
||||
|
||||
// prepareInboundLayerRPCBatch is the production API path. The whole container
|
||||
// reserves conservative task/materialization capacity before the first exact
|
||||
// decoder callback. Admission then classifies every request; fresh owners keep
|
||||
|
|
@ -217,15 +281,38 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
|
|||
return err
|
||||
}
|
||||
evidence := make([]layerRPCProfileEvidence, len(plan.items))
|
||||
for _, index := range candidateItems {
|
||||
decodeOptions := make([]tlprofile.AdmissionOptions, len(candidateItems))
|
||||
expansionBudgets := make([]*layerRPCGZIPExpansionBudget, len(candidateItems))
|
||||
materializationCapacity := false
|
||||
for reservationIndex, index := range candidateItems {
|
||||
item := &plan.items[index]
|
||||
itemState := admissionCursor.state
|
||||
existingProfile, existing := s.rpcResults.ExactAdmissionProfile(c.authKeyID, c.sessionID, item.msgID)
|
||||
if existing {
|
||||
itemState = LayerProfileSnapshot{Profile: existingProfile, Origin: LayerProfileExplicit}
|
||||
}
|
||||
admitted, method, err := s.decodeInboundLayerRPC(itemState, item.body)
|
||||
expansionBudget := &layerRPCGZIPExpansionBudget{
|
||||
server: s, plan: plan, reservation: reservation,
|
||||
entry: reservationIndex, baseSourceBytes: len(item.body), chargedSourceBytes: len(item.body),
|
||||
}
|
||||
expansionBudgets[reservationIndex] = expansionBudget
|
||||
options := tlprofile.AdmissionOptions{
|
||||
Limits: inboundLayerDecodeLimits,
|
||||
ExpandGZIP: expansionBudget.expand,
|
||||
}
|
||||
decodeOptions[reservationIndex] = options
|
||||
admitted, method, err := s.decodeInboundLayerRPCWithOptions(itemState, item.body, options)
|
||||
if err != nil {
|
||||
if errors.Is(err, errLayerRPCGZIPCapacity) {
|
||||
materializationCapacity = true
|
||||
break
|
||||
}
|
||||
if errors.Is(err, ErrConnClosed) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return err
|
||||
}
|
||||
if errors.Is(err, errLayerRPCAdmissionCapability) {
|
||||
return err
|
||||
}
|
||||
if terminal, recognized := wrappedDestroyAuthKeyTerminal(err); recognized {
|
||||
if terminal.WireSize != bin.Word || !validWrappedDestroyAuthKeyChain(terminal) {
|
||||
s.log.Debug("Wrapped destroy_auth_key terminal rejected",
|
||||
|
|
@ -278,6 +365,19 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
|
|||
}
|
||||
}
|
||||
}
|
||||
if materializationCapacity {
|
||||
for _, index := range candidateItems {
|
||||
item := &plan.items[index]
|
||||
item.kind = inboundItemCapacityError
|
||||
item.admitted = tlprofile.Admission{}
|
||||
c.metrics.InboundRPCDropped(s.typeName(item.typeID), "materialization_capacity")
|
||||
}
|
||||
if err := reservation.retain(nil, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
plan.rpcReservation = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
var indices []int
|
||||
var reservationIndices []int
|
||||
|
|
@ -304,7 +404,16 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
|
|||
item.admitted.Prepared().SemanticIdentity(),
|
||||
item.admitted.Call().Identity(),
|
||||
); candidate != nil {
|
||||
claim, err := s.acquireAdmittedLayerRPC(c, item, &evidence[index])
|
||||
claim, err := s.acquireAdmittedLayerRPC(
|
||||
c, item, &evidence[index], decodeOptions[reservationIndex], expansionBudgets[reservationIndex],
|
||||
)
|
||||
if errors.Is(err, errLayerRPCGZIPCapacity) {
|
||||
s.rpcRewrap.release(candidate)
|
||||
c.metrics.InboundRPCDropped(candidate.method, "materialization_capacity")
|
||||
flightCapacity = true
|
||||
item.kind = inboundItemCapacityError
|
||||
continue
|
||||
}
|
||||
if errors.Is(err, ErrRPCResultFlightCapacity) {
|
||||
s.rpcRewrap.release(candidate)
|
||||
c.metrics.InboundRPCDropped(candidate.method, "flight_capacity")
|
||||
|
|
@ -387,7 +496,15 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
|
|||
clearedPostInitCandidates = true
|
||||
}
|
||||
|
||||
claim, err := s.acquireAdmittedLayerRPC(c, item, &evidence[index])
|
||||
claim, err := s.acquireAdmittedLayerRPC(
|
||||
c, item, &evidence[index], decodeOptions[reservationIndex], expansionBudgets[reservationIndex],
|
||||
)
|
||||
if errors.Is(err, errLayerRPCGZIPCapacity) {
|
||||
c.metrics.InboundRPCDropped(method, "materialization_capacity")
|
||||
flightCapacity = true
|
||||
item.kind = inboundItemCapacityError
|
||||
continue
|
||||
}
|
||||
if errors.Is(err, ErrRPCResultFlightCapacity) {
|
||||
c.metrics.InboundRPCDropped(method, "flight_capacity")
|
||||
flightCapacity = true
|
||||
|
|
@ -624,6 +741,8 @@ func (s *Server) acquireAdmittedLayerRPC(
|
|||
c *Conn,
|
||||
item *inboundItem,
|
||||
evidence *layerRPCProfileEvidence,
|
||||
options tlprofile.AdmissionOptions,
|
||||
expansionBudget *layerRPCGZIPExpansionBudget,
|
||||
) (rpcResultAcquire, error) {
|
||||
if s == nil || c == nil || item == nil {
|
||||
return rpcResultAcquire{}, ErrRPCResultFlightInvalid
|
||||
|
|
@ -652,9 +771,17 @@ func (s *Server) acquireAdmittedLayerRPC(
|
|||
// mutation, not an inherited-default race.
|
||||
return rpcResultAcquire{}, err
|
||||
}
|
||||
redecoded, method, decodeErr := s.decodeInboundLayerRPC(
|
||||
// Drop the losing typed graph before materializing its authoritative-profile
|
||||
// replacement. The grown reservation therefore needs to cover the larger
|
||||
// graph, not two simultaneous copies.
|
||||
item.admitted = tlprofile.Admission{}
|
||||
if expansionBudget != nil {
|
||||
expansionBudget.beginAttempt()
|
||||
}
|
||||
redecoded, method, decodeErr := s.decodeInboundLayerRPCWithOptions(
|
||||
LayerProfileSnapshot{Profile: winnerProfile, Origin: LayerProfileExplicit},
|
||||
item.body,
|
||||
options,
|
||||
)
|
||||
if decodeErr != nil {
|
||||
return rpcResultAcquire{}, decodeErr
|
||||
|
|
@ -786,6 +913,14 @@ func layerRPCAdmissionHasExplicitSelector(body []byte, admissionErr error) bool
|
|||
// evidence only after the full request identity has acquired an owner (or a
|
||||
// genuine new-msg_id rewrap alias).
|
||||
func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte) (tlprofile.Admission, string, error) {
|
||||
return s.decodeInboundLayerRPCWithLimits(state, body, inboundLayerDecodeLimits)
|
||||
}
|
||||
|
||||
func (s *Server) decodeInboundLayerRPCWithLimits(state LayerProfileSnapshot, body []byte, limits tlprofile.Limits) (tlprofile.Admission, string, error) {
|
||||
return s.decodeInboundLayerRPCWithOptions(state, body, tlprofile.AdmissionOptions{Limits: limits})
|
||||
}
|
||||
|
||||
func (s *Server) decodeInboundLayerRPCWithOptions(state LayerProfileSnapshot, body []byte, options tlprofile.AdmissionOptions) (tlprofile.Admission, string, error) {
|
||||
if s == nil || s.layerRPC == nil || len(body) < bin.Word {
|
||||
return tlprofile.Admission{}, "unknown", fmt.Errorf("invalid exact RPC admission input")
|
||||
}
|
||||
|
|
@ -795,15 +930,30 @@ func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte)
|
|||
err error
|
||||
)
|
||||
if state.Origin != LayerProfileUnknown {
|
||||
if admitter, ok := s.layerRPC.(LayerRPCDefaultProfileAdmitter); ok {
|
||||
request, err = admitter.AdmitDefaultLayer(state.Profile, b, inboundLayerDecodeLimits)
|
||||
if admitter, ok := s.layerRPC.(LayerRPCDefaultProfileOptionsAdmitter); ok {
|
||||
request, err = admitter.AdmitDefaultLayerWithOptions(state.Profile, b, options)
|
||||
} else if admitter, ok := s.layerRPC.(LayerRPCDefaultProfileAdmitter); ok {
|
||||
// Preserve the stable pre-options capability first: unlike strict
|
||||
// exact admission, default admission lets an explicit invokeWithLayer
|
||||
// correct inherited or restored profile evidence.
|
||||
request, err = admitter.AdmitDefaultLayer(state.Profile, b, options.Limits)
|
||||
} else if admitter, ok := s.layerRPC.(LayerRPCOptionsAdmitter); ok {
|
||||
request, err = admitter.AdmitLayerWithOptions(state.Profile, b, options)
|
||||
} else {
|
||||
// Compatibility fallback for old package tests/mocks. Production Router
|
||||
// implements default admission so explicit invokeWithLayer can correct.
|
||||
request, err = s.layerRPC.AdmitLayer(state.Profile, b, inboundLayerDecodeLimits)
|
||||
request, err = s.layerRPC.AdmitLayer(state.Profile, b, options.Limits)
|
||||
}
|
||||
} else if admitter, ok := s.layerRPC.(LayerRPCOptionsAdmitter); ok {
|
||||
request, err = admitter.AdmitUnprofiledWithOptions(b, options)
|
||||
} else {
|
||||
request, err = s.layerRPC.AdmitUnprofiled(b, inboundLayerDecodeLimits)
|
||||
request, err = s.layerRPC.AdmitUnprofiled(b, options.Limits)
|
||||
}
|
||||
if err != nil && options.ExpandGZIP != nil && errors.Is(err, tlprofile.ErrGZIPExpanderMissing) {
|
||||
// A handler compiled against the stable Limits-only boundary remains
|
||||
// valid for plain requests. Encountering gzip_packed without the optional
|
||||
// capability is instead a server wiring/programming error: returning the
|
||||
// generated error as INPUT_REQUEST_INVALID would silently blame a valid
|
||||
// client envelope and make recovery impossible.
|
||||
err = fmt.Errorf("%w: handler cannot use caller-owned bounded gzip expansion: %w", errLayerRPCAdmissionCapability, err)
|
||||
}
|
||||
method := "unknown"
|
||||
if err == nil {
|
||||
|
|
@ -847,6 +997,10 @@ func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte)
|
|||
method = s.typeName(codecErr.WireID)
|
||||
}
|
||||
}
|
||||
var unknownTerminal *tlprofile.UnknownTerminalError
|
||||
if errors.As(err, &unknownTerminal) && unknownTerminal.WireID != 0 {
|
||||
method = s.typeName(unknownTerminal.WireID)
|
||||
}
|
||||
if errors.Is(err, tlprofile.ErrUnknownRPCMethod) && s.log != nil {
|
||||
if terminal, recognized := wrappedDestroyAuthKeyTerminal(err); recognized {
|
||||
method = "destroy_auth_key"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
|
@ -46,6 +47,14 @@ type admissionOnlyLayerRPC struct {
|
|||
published []publishedLayerEvidence
|
||||
}
|
||||
|
||||
// legacyAdmissionOnlyLayerRPC intentionally exposes only the original
|
||||
// Limits-based admission interfaces. It guards source and runtime compatibility
|
||||
// for implementations compiled before caller-owned AdmissionOptions existed.
|
||||
type legacyAdmissionOnlyLayerRPC struct {
|
||||
dispatcher *tlprofile.Dispatcher
|
||||
lastRemaining int
|
||||
}
|
||||
|
||||
type orderedAdmissionOnlyLayerRPC struct {
|
||||
*admissionOnlyLayerRPC
|
||||
exactMu sync.Mutex
|
||||
|
|
@ -221,6 +230,10 @@ func newAdmissionOnlyLayerRPC() *admissionOnlyLayerRPC {
|
|||
return &admissionOnlyLayerRPC{dispatcher: tlprofile.NewDispatcher()}
|
||||
}
|
||||
|
||||
func newLegacyAdmissionOnlyLayerRPC() *legacyAdmissionOnlyLayerRPC {
|
||||
return &legacyAdmissionOnlyLayerRPC{dispatcher: tlprofile.NewDispatcher()}
|
||||
}
|
||||
|
||||
func newOrderedAdmissionOnlyLayerRPC() *orderedAdmissionOnlyLayerRPC {
|
||||
return &orderedAdmissionOnlyLayerRPC{
|
||||
admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC(),
|
||||
|
|
@ -265,14 +278,48 @@ func (h *admissionOnlyLayerRPC) AdmitLayer(profile tlprofile.Profile, b *bin.Buf
|
|||
return h.dispatcher.Admit(profile, b, limits)
|
||||
}
|
||||
|
||||
func (h *admissionOnlyLayerRPC) AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
return h.dispatcher.AdmitWithOptions(profile, b, options)
|
||||
}
|
||||
|
||||
func (h *admissionOnlyLayerRPC) AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
return h.dispatcher.AdmitDefault(profile, b, limits)
|
||||
}
|
||||
|
||||
func (h *admissionOnlyLayerRPC) AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
return h.dispatcher.AdmitDefaultWithOptions(profile, b, options)
|
||||
}
|
||||
|
||||
func (h *admissionOnlyLayerRPC) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
return h.dispatcher.AdmitUnprofiled(b, limits)
|
||||
}
|
||||
|
||||
func (h *admissionOnlyLayerRPC) AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
return h.dispatcher.AdmitUnprofiledWithOptions(b, options)
|
||||
}
|
||||
|
||||
func (h *legacyAdmissionOnlyLayerRPC) AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
request, err := h.dispatcher.Admit(profile, b, limits)
|
||||
h.lastRemaining = b.Len()
|
||||
return request, err
|
||||
}
|
||||
|
||||
func (h *legacyAdmissionOnlyLayerRPC) AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
request, err := h.dispatcher.AdmitDefault(profile, b, limits)
|
||||
h.lastRemaining = b.Len()
|
||||
return request, err
|
||||
}
|
||||
|
||||
func (h *legacyAdmissionOnlyLayerRPC) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
request, err := h.dispatcher.AdmitUnprofiled(b, limits)
|
||||
h.lastRemaining = b.Len()
|
||||
return request, err
|
||||
}
|
||||
|
||||
func (*legacyAdmissionOnlyLayerRPC) DispatchAdmitted(context.Context, [8]byte, int64, int64, uint64, tlprofile.Admission) (tlprofile.Result, string, error) {
|
||||
return nil, "", fmt.Errorf("admission-only handler")
|
||||
}
|
||||
|
||||
func (*admissionOnlyLayerRPC) DispatchAdmitted(context.Context, [8]byte, int64, int64, uint64, tlprofile.Admission) (tlprofile.Result, string, error) {
|
||||
return nil, "", fmt.Errorf("admission-only handler")
|
||||
}
|
||||
|
|
@ -301,6 +348,99 @@ func (h *admissionOnlyLayerRPC) publications() []publishedLayerEvidence {
|
|||
return append([]publishedLayerEvidence(nil), h.published...)
|
||||
}
|
||||
|
||||
func TestLegacyLayerRPCAdmissionInterfacesRemainCompatible(t *testing.T) {
|
||||
t.Run("inherited default accepts explicit correction", func(t *testing.T) {
|
||||
handler := newLegacyAdmissionOnlyLayerRPC()
|
||||
s := New(Options{DC: 2, LayerRPC: handler})
|
||||
body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{
|
||||
Layer: 225,
|
||||
Query: &tg.HelpGetConfigRequest{},
|
||||
})
|
||||
|
||||
request, method, err := s.decodeInboundLayerRPCWithOptions(
|
||||
LayerProfileSnapshot{Profile: tlprofile.Profile228, Origin: LayerProfileInherited},
|
||||
body,
|
||||
tlprofile.AdmissionOptions{
|
||||
Limits: inboundLayerDecodeLimits,
|
||||
ExpandGZIP: func([]byte, int) ([]byte, func(), error) {
|
||||
t.Fatal("plain legacy admission unexpectedly requested gzip expansion")
|
||||
return nil, nil, nil
|
||||
},
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("legacy default admission: %v", err)
|
||||
}
|
||||
if method != "help.getConfig" {
|
||||
t.Fatalf("method = %q, want help.getConfig", method)
|
||||
}
|
||||
if request.Call().Profile() != tlprofile.Profile225 {
|
||||
t.Fatalf("call profile = %d, want 225", request.Call().Profile())
|
||||
}
|
||||
if profile, ok := request.ProfileEvidence(); !ok || profile != tlprofile.Profile225 {
|
||||
t.Fatalf("profile evidence = %d/%v, want 225/true", profile, ok)
|
||||
}
|
||||
if handler.lastRemaining != 0 {
|
||||
t.Fatalf("successful legacy admission left %d bytes", handler.lastRemaining)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("unprofiled selector remains supported", func(t *testing.T) {
|
||||
handler := newLegacyAdmissionOnlyLayerRPC()
|
||||
s := New(Options{DC: 2, LayerRPC: handler})
|
||||
body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{
|
||||
Layer: 225,
|
||||
Query: &tg.HelpGetNearestDCRequest{},
|
||||
})
|
||||
|
||||
request, method, err := s.decodeInboundLayerRPCWithOptions(
|
||||
LayerProfileSnapshot{},
|
||||
body,
|
||||
tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("legacy unprofiled admission: %v", err)
|
||||
}
|
||||
if method != "help.getNearestDc" {
|
||||
t.Fatalf("method = %q, want help.getNearestDc", method)
|
||||
}
|
||||
if request.Call().Profile() != tlprofile.Profile225 {
|
||||
t.Fatalf("call profile = %d, want 225", request.Call().Profile())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("nested gzip requires explicit optional capability", func(t *testing.T) {
|
||||
handler := newLegacyAdmissionOnlyLayerRPC()
|
||||
s := New(Options{DC: 2, LayerRPC: handler})
|
||||
body, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
|
||||
original := append([]byte(nil), body...)
|
||||
|
||||
_, _, err := s.decodeInboundLayerRPCWithOptions(
|
||||
LayerProfileSnapshot{},
|
||||
body,
|
||||
tlprofile.AdmissionOptions{
|
||||
Limits: inboundLayerDecodeLimits,
|
||||
ExpandGZIP: func([]byte, int) ([]byte, func(), error) {
|
||||
t.Fatal("legacy handler unexpectedly received options-only gzip expander")
|
||||
return nil, nil, nil
|
||||
},
|
||||
},
|
||||
)
|
||||
if !errors.Is(err, errLayerRPCAdmissionCapability) {
|
||||
t.Fatalf("nested gzip error = %v, want admission capability error", err)
|
||||
}
|
||||
if !errors.Is(err, tlprofile.ErrGZIPExpanderMissing) {
|
||||
t.Fatalf("nested gzip error = %v, want generated missing-expander cause", err)
|
||||
}
|
||||
if handler.lastRemaining != len(body) {
|
||||
t.Fatalf("failed legacy admission retained %d/%d input bytes", handler.lastRemaining, len(body))
|
||||
}
|
||||
if !bytes.Equal(body, original) {
|
||||
t.Fatal("failed legacy admission mutated caller wire bytes")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestNestedExplicitLayerAdmissionErrorsAreNotDefaultFailures(t *testing.T) {
|
||||
handler := newAdmissionOnlyLayerRPC()
|
||||
s := New(Options{DC: 2, LayerRPC: handler})
|
||||
|
|
@ -833,11 +973,12 @@ func TestSameMsgIDNakedReplayUsesWinnerAdmissionProfile(t *testing.T) {
|
|||
}
|
||||
c220 := &Conn{authKeyID: authKeyID, sessionID: sessionID}
|
||||
c227 := &Conn{authKeyID: authKeyID, sessionID: sessionID}
|
||||
winner, err := s.acquireAdmittedLayerRPC(c220, &item220, nil)
|
||||
options := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits}
|
||||
winner, err := s.acquireAdmittedLayerRPC(c220, &item220, nil, options, nil)
|
||||
if err != nil || winner.state != rpcResultAcquireOwner || winner.owner == nil {
|
||||
t.Fatalf("winner = state:%d err:%v", winner.state, err)
|
||||
}
|
||||
loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil)
|
||||
loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil, options, nil)
|
||||
if err != nil || loser.state != rpcResultAcquirePending || loser.admissionSeq != winner.admissionSeq {
|
||||
t.Fatalf("loser join = state:%d seq:%d err:%v, winner seq:%d", loser.state, loser.admissionSeq, err, winner.admissionSeq)
|
||||
}
|
||||
|
|
@ -852,7 +993,7 @@ func TestSameMsgIDNakedReplayUsesWinnerAdmissionProfile(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.acquireAdmittedLayerRPC(c220, &changed, nil); !errors.Is(err, ErrRPCResultIdentityMismatch) {
|
||||
if _, err := s.acquireAdmittedLayerRPC(c220, &changed, nil, options, nil); !errors.Is(err, ErrRPCResultIdentityMismatch) {
|
||||
t.Fatalf("same-msg_id changed body err=%v, want identity mismatch", err)
|
||||
}
|
||||
winner.owner.Abort()
|
||||
|
|
|
|||
|
|
@ -92,6 +92,11 @@ type inboundPlan struct {
|
|||
ackIDs []int64
|
||||
logicalMin int64
|
||||
releases []func()
|
||||
// gzipExpandedBytes is the non-refundable per-frame decompression work
|
||||
// already performed by outer and exact-layer nested gzip envelopes. Memory
|
||||
// reservations are released when their buffers die, but this cumulative
|
||||
// counter prevents sibling RPCs from recycling the same CPU budget.
|
||||
gzipExpandedBytes int
|
||||
|
||||
rpcReservation *inboundRPCBatchReservation
|
||||
rpcTasks []inboundRPC
|
||||
|
|
@ -314,6 +319,7 @@ func (s *Server) preflightInbound(cs *connState, msgID int64, seqNo int32, body
|
|||
plan.close()
|
||||
return nil, err
|
||||
}
|
||||
plan.gzipExpandedBytes = budget.expanded
|
||||
plan.staged = overlay.staged
|
||||
if plan.logicalMin == 0 {
|
||||
plan.close()
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"container/list"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
|
@ -544,6 +545,60 @@ func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) {
|
|||
}
|
||||
}
|
||||
|
||||
// growEntry raises one provisional task's materialization charge without a
|
||||
// release/reacquire window. Exact gzip admission calls this after the expanded
|
||||
// size is known but before the generated typed decoder can allocate the request
|
||||
// graph. A failure leaves every existing reservation unchanged so the caller
|
||||
// can reject and abort the whole container atomically.
|
||||
func (r *inboundRPCBatchReservation) growEntry(index, targetSize int) error {
|
||||
if r == nil || index < 0 || index >= len(r.entries) || targetSize < 0 {
|
||||
return errInboundRPCBatchSelection
|
||||
}
|
||||
entry := &r.entries[index]
|
||||
if targetSize <= entry.size {
|
||||
return nil
|
||||
}
|
||||
if entry.global == nil || entry.global.scheduler == nil || entry.global.released.Load() {
|
||||
return ErrConnClosed
|
||||
}
|
||||
delta := int64(targetSize) - int64(entry.size)
|
||||
scheduler := entry.global.scheduler
|
||||
conn := r.conn
|
||||
|
||||
// Match initial admission's lock order: global scheduler budget, then the
|
||||
// connection queue budget. Release paths never hold rpcMu while acquiring
|
||||
// budgetMu, so this cannot invert task completion or close.
|
||||
scheduler.budgetMu.Lock()
|
||||
defer scheduler.budgetMu.Unlock()
|
||||
select {
|
||||
case <-scheduler.stopCh:
|
||||
return ErrConnClosed
|
||||
default:
|
||||
}
|
||||
if delta > scheduler.maxBytes-scheduler.bytes {
|
||||
return fmt.Errorf("%w: grow exact admission global byte budget by %d", ErrInboundRPCQueueFull, delta)
|
||||
}
|
||||
|
||||
conn.rpcMu.Lock()
|
||||
defer conn.rpcMu.Unlock()
|
||||
if err := r.ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if conn.rpcClosed || conn.isRetired() {
|
||||
return ErrConnClosed
|
||||
}
|
||||
if delta > int64(maxInflightRPCBytes)-conn.inflightRPCBytes.Load() {
|
||||
return fmt.Errorf("%w: grow exact admission connection byte budget by %d", ErrInboundRPCQueueFull, delta)
|
||||
}
|
||||
|
||||
scheduler.bytes += delta
|
||||
entry.global.size += delta
|
||||
entry.size = targetSize
|
||||
r.totalSize += delta
|
||||
conn.inflightRPCBytes.Add(delta)
|
||||
return nil
|
||||
}
|
||||
|
||||
// retain keeps a subset of a provisional batch on the original connection and
|
||||
// global reservations. Exact-layer admission uses this after typed decode has
|
||||
// classified completed replays, pending joins, admission errors, and fresh
|
||||
|
|
|
|||
|
|
@ -130,6 +130,80 @@ func TestInboundRPCBatchAbortReturnsEveryReservationExactlyOnce(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCBatchReservationGrowTransfersBudgetAtomically(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
||||
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
|
||||
{method: "one", size: 3},
|
||||
{method: "two", size: 5},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := reservation.growEntry(0, 11); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 2 || bytes != 16 {
|
||||
t.Fatalf("grown global budget = %d/%d, want 2/16", tasks, bytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 16 {
|
||||
t.Fatalf("grown connection budget = %d, want 16", got)
|
||||
}
|
||||
if reservation.entries[0].size != 11 || reservation.totalSize != 16 {
|
||||
t.Fatalf("grown reservation entry/total = %d/%d", reservation.entries[0].size, reservation.totalSize)
|
||||
}
|
||||
reservation.abort()
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("grown reservation abort leaked global budget %d/%d", tasks, bytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("grown reservation abort leaked connection budget %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCBatchReservationGrowFailureKeepsOriginalBudget(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
globalMax int64
|
||||
targetSize int
|
||||
}{
|
||||
{name: "global", globalMax: 5, targetSize: 6},
|
||||
{name: "connection", globalMax: int64(maxInflightRPCBytes) * 2, targetSize: maxInflightRPCBytes + 1},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, test.globalMax)
|
||||
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "one", size: 3}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := reservation.growEntry(0, test.targetSize); !errors.Is(err, ErrInboundRPCQueueFull) {
|
||||
t.Fatalf("grow error = %v, want ErrInboundRPCQueueFull", err)
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 3 {
|
||||
t.Fatalf("failed grow changed global budget %d/%d", tasks, bytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 3 {
|
||||
t.Fatalf("failed grow changed connection budget %d", got)
|
||||
}
|
||||
if reservation.entries[0].size != 3 || reservation.totalSize != 3 {
|
||||
t.Fatalf("failed grow changed reservation entry/total = %d/%d", reservation.entries[0].size, reservation.totalSize)
|
||||
}
|
||||
reservation.abort()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCBatchCommitAppendsAllAndSchedulesAtomically(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
||||
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import (
|
|||
|
||||
"github.com/iamxvbaba/td/bin"
|
||||
"github.com/iamxvbaba/td/clock"
|
||||
"github.com/iamxvbaba/td/proto"
|
||||
"github.com/iamxvbaba/td/tg"
|
||||
"go.uber.org/zap/zaptest"
|
||||
|
||||
|
|
@ -113,11 +114,27 @@ func (h *countingLayerRPCAdmission) AdmitLayer(profile tlprofile.Profile, b *bin
|
|||
return h.LayerRPCHandler.AdmitLayer(profile, b, limits)
|
||||
}
|
||||
|
||||
func (h *countingLayerRPCAdmission) AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
h.decodeCalls.Add(1)
|
||||
if admitter, ok := h.LayerRPCHandler.(LayerRPCOptionsAdmitter); ok {
|
||||
return admitter.AdmitLayerWithOptions(profile, b, options)
|
||||
}
|
||||
return h.LayerRPCHandler.AdmitLayer(profile, b, options.Limits)
|
||||
}
|
||||
|
||||
func (h *countingLayerRPCAdmission) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
h.decodeCalls.Add(1)
|
||||
return h.LayerRPCHandler.AdmitUnprofiled(b, limits)
|
||||
}
|
||||
|
||||
func (h *countingLayerRPCAdmission) AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
h.decodeCalls.Add(1)
|
||||
if admitter, ok := h.LayerRPCHandler.(LayerRPCOptionsAdmitter); ok {
|
||||
return admitter.AdmitUnprofiledWithOptions(b, options)
|
||||
}
|
||||
return h.LayerRPCHandler.AdmitUnprofiled(b, options.Limits)
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionCapacityRejectsBeforeDecoder(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
|
|
@ -162,6 +179,209 @@ func TestLayerRPCAdmissionCapacityRejectsBeforeDecoder(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionExpandsTDLibNestedGZIPUnderTransferredBudget(t *testing.T) {
|
||||
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
|
||||
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
|
||||
c := &Conn{authKeyID: [8]byte{8, 21}, sessionID: 821, metrics: NopMetrics{}}
|
||||
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
s.rpcScheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
body, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
|
||||
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
|
||||
defer plan.close()
|
||||
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.items[0].kind != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
|
||||
t.Fatalf("admitted nested gzip plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
|
||||
}
|
||||
wantCharge := int64(layerRPCAdmissionReservationSize(len(body) + expandedBytes))
|
||||
if got := c.inflightRPCBytes.Load(); got != wantCharge {
|
||||
t.Fatalf("nested gzip connection charge = %d, want %d", got, wantCharge)
|
||||
}
|
||||
if got := plan.rpcReservation.entries[0].size; int64(got) != wantCharge {
|
||||
t.Fatalf("nested gzip reservation charge = %d, want %d", got, wantCharge)
|
||||
}
|
||||
if got := plan.gzipExpandedBytes; got != expandedBytes {
|
||||
t.Fatalf("nested gzip cumulative expansion = %d, want %d", got, expandedBytes)
|
||||
}
|
||||
if got := s.frameBudget.usedBytes(); got != 0 {
|
||||
t.Fatalf("nested gzip temporary frame budget retained after materialization: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionNestedGZIPGrowFailureRejectsWholeBatch(t *testing.T) {
|
||||
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
|
||||
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
|
||||
first, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
|
||||
second, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetNearestDCRequest{})
|
||||
initialCharge := int64(layerRPCAdmissionReservationSize(len(first)) + layerRPCAdmissionReservationSize(len(second)))
|
||||
s.rpcScheduler = newInboundRPCScheduler(1, 4, initialCharge)
|
||||
c := &Conn{authKeyID: [8]byte{8, 22}, sessionID: 822, metrics: NopMetrics{}}
|
||||
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
s.rpcScheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
plan := &inboundPlan{items: []inboundItem{
|
||||
{kind: inboundItemRPC, msgID: 100, body: first},
|
||||
{kind: inboundItemRPC, msgID: 104, body: second},
|
||||
}}
|
||||
defer plan.close()
|
||||
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for index := range plan.items {
|
||||
if plan.items[index].kind != inboundItemCapacityError {
|
||||
t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
|
||||
}
|
||||
}
|
||||
if plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
|
||||
t.Fatalf("capacity plan retained reservation/tasks = %v/%d", plan.rpcReservation != nil, len(plan.rpcTasks))
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("grow failure leaked connection charge %d", got)
|
||||
}
|
||||
if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("grow failure leaked global budget %d/%d", tasks, bytes)
|
||||
}
|
||||
if got := s.frameBudget.usedBytes(); got != 0 {
|
||||
t.Fatalf("grow failure leaked temporary frame budget %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionNestedGZIPSiblingsShareFrameExpansionLimit(t *testing.T) {
|
||||
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
|
||||
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
|
||||
c := &Conn{authKeyID: [8]byte{8, 23}, sessionID: 823, metrics: NopMetrics{}}
|
||||
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
s.rpcScheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
first, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
|
||||
second, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetNearestDCRequest{})
|
||||
plan := &inboundPlan{
|
||||
gzipExpandedBytes: maxDispatchExpandedBytes - expandedBytes,
|
||||
items: []inboundItem{
|
||||
{kind: inboundItemRPC, msgID: 100, body: first},
|
||||
{kind: inboundItemRPC, msgID: 104, body: second},
|
||||
},
|
||||
}
|
||||
defer plan.close()
|
||||
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for index := range plan.items {
|
||||
if plan.items[index].kind != inboundItemCapacityError {
|
||||
t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
|
||||
}
|
||||
}
|
||||
if got := plan.gzipExpandedBytes; got != maxDispatchExpandedBytes {
|
||||
t.Fatalf("shared cumulative expansion = %d, want %d", got, maxDispatchExpandedBytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("shared-limit rejection leaked connection charge %d", got)
|
||||
}
|
||||
if got := s.frameBudget.usedBytes(); got != 0 {
|
||||
t.Fatalf("shared-limit rejection leaked temporary frame budget %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testing.T) {
|
||||
handler := newAdmissionOnlyLayerRPC()
|
||||
s := New(Options{DC: 2, LayerRPC: handler, Logger: zaptest.NewLogger(t)})
|
||||
s.rpcResults = newRPCResultCacheWithFlightLimit(time.Now, 8)
|
||||
scheduler := newInboundRPCScheduler(1, 4, 1<<30)
|
||||
s.rpcScheduler = scheduler
|
||||
authKeyID := [8]byte{8, 24}
|
||||
const sessionID = int64(824)
|
||||
c225 := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
|
||||
c227 := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
|
||||
c225.startInboundRPCScheduler(scheduler, 1, 2, time.Second)
|
||||
c227.startInboundRPCScheduler(scheduler, 1, 2, time.Second)
|
||||
defer func() {
|
||||
c225.closeInboundRPCScheduler()
|
||||
c227.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
terminal := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.MessagesGetHistoryRequest{
|
||||
Peer: &tg.InputPeerSelf{}, Limit: 1,
|
||||
})
|
||||
body := exactLayerRPCBody(t, &tg.InvokeWithoutUpdatesRequest{Query: &proto.GZIP{Data: terminal}})
|
||||
initialCharge := layerRPCAdmissionReservationSize(len(body))
|
||||
reservation225, err := c225.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "messages.getHistory", size: initialCharge}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer reservation225.abort()
|
||||
reservation227, err := c227.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "messages.getHistory", size: initialCharge}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer reservation227.abort()
|
||||
|
||||
plan225 := &inboundPlan{}
|
||||
budget225 := &layerRPCGZIPExpansionBudget{
|
||||
server: s, plan: plan225, reservation: reservation225,
|
||||
baseSourceBytes: len(body), chargedSourceBytes: len(body),
|
||||
}
|
||||
options225 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget225.expand}
|
||||
item225 := inboundItem{msgID: 100, body: body}
|
||||
item225.admitted, item225.method, err = s.decodeInboundLayerRPCWithOptions(
|
||||
LayerProfileSnapshot{Profile: tlprofile.Profile225, Origin: LayerProfileInherited}, body, options225,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
plan227 := &inboundPlan{}
|
||||
budget227 := &layerRPCGZIPExpansionBudget{
|
||||
server: s, plan: plan227, reservation: reservation227,
|
||||
baseSourceBytes: len(body), chargedSourceBytes: len(body),
|
||||
}
|
||||
options227 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget227.expand}
|
||||
item227 := inboundItem{msgID: 100, body: body}
|
||||
item227.admitted, item227.method, err = s.decodeInboundLayerRPCWithOptions(
|
||||
LayerProfileSnapshot{Profile: tlprofile.Profile227, Origin: LayerProfileInherited}, body, options227,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if item225.admitted.Prepared().Identity() == item227.admitted.Prepared().Identity() {
|
||||
t.Fatal("test request identity is invariant; need authoritative-profile re-decode")
|
||||
}
|
||||
|
||||
winner, err := s.acquireAdmittedLayerRPC(c225, &item225, nil, options225, budget225)
|
||||
if err != nil || winner.state != rpcResultAcquireOwner || winner.owner == nil {
|
||||
t.Fatalf("winner = state:%d err:%v", winner.state, err)
|
||||
}
|
||||
defer winner.owner.Abort()
|
||||
loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil, options227, budget227)
|
||||
if err != nil || loser.state != rpcResultAcquirePending {
|
||||
t.Fatalf("loser = state:%d err:%v", loser.state, err)
|
||||
}
|
||||
if got := item227.admitted.Call().Profile(); got != tlprofile.Profile225 {
|
||||
t.Fatalf("loser re-admitted profile = %d, want 225", got)
|
||||
}
|
||||
wantCharge := layerRPCAdmissionReservationSize(len(body) + len(terminal))
|
||||
if got := reservation227.entries[0].size; got != wantCharge {
|
||||
t.Fatalf("re-decode reservation charge = %d, want single-graph maximum %d", got, wantCharge)
|
||||
}
|
||||
if got := plan227.gzipExpandedBytes; got != 2*len(terminal) {
|
||||
t.Fatalf("re-decode cumulative work = %d, want %d", got, 2*len(terminal))
|
||||
}
|
||||
if got := s.frameBudget.usedBytes(); got != 0 {
|
||||
t.Fatalf("re-decode leaked temporary frame budget %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T) {
|
||||
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
|
||||
s := New(Options{DC: 2, LayerRPC: router})
|
||||
|
|
@ -207,6 +427,25 @@ func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T)
|
|||
}
|
||||
}
|
||||
|
||||
func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) ([]byte, int) {
|
||||
t.Helper()
|
||||
terminalWire := exactOutboundLayerRPCBody(t, profile, terminal)
|
||||
request := &tg.InvokeWithLayerRequest{
|
||||
Layer: int(profile),
|
||||
Query: &tg.InitConnectionRequest{
|
||||
APIID: 1,
|
||||
DeviceModel: "android",
|
||||
SystemVersion: "test",
|
||||
AppVersion: "1.0",
|
||||
SystemLangCode: "en",
|
||||
LangPack: "",
|
||||
LangCode: "en",
|
||||
Query: &proto.GZIP{Data: terminalWire},
|
||||
},
|
||||
}
|
||||
return exactLayerRPCBody(t, request), len(terminalWire)
|
||||
}
|
||||
|
||||
func TestLayerRPCAdmissionPendingReplayReleasesProvisionalEntry(t *testing.T) {
|
||||
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
|
||||
s := New(Options{DC: 2, LayerRPC: router})
|
||||
|
|
|
|||
|
|
@ -70,6 +70,14 @@ type LayerRPCHandler interface {
|
|||
) (tlprofile.Result, string, error)
|
||||
}
|
||||
|
||||
// LayerRPCOptionsAdmitter extends the stable handler boundary with
|
||||
// caller-owned admission capabilities. Implementations that do not expose it
|
||||
// remain usable for requests that need only Limits.
|
||||
type LayerRPCOptionsAdmitter interface {
|
||||
AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
|
||||
AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
|
||||
}
|
||||
|
||||
// LayerRPCDefaultProfileAdmitter decodes with a recoverable inherited/default
|
||||
// profile. Production handlers should implement it with the same sparse
|
||||
// tlprofile dispatcher and semantic adapter registry used by AdmitLayer. The split keeps old
|
||||
|
|
@ -79,6 +87,12 @@ type LayerRPCDefaultProfileAdmitter interface {
|
|||
AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error)
|
||||
}
|
||||
|
||||
// LayerRPCDefaultProfileOptionsAdmitter is the capability-aware form used when
|
||||
// exact admission needs caller-owned resources such as bounded gzip expansion.
|
||||
type LayerRPCDefaultProfileOptionsAdmitter interface {
|
||||
AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
|
||||
}
|
||||
|
||||
// LayerRPCSessionProfileResolver may restore an exact profile only when it was
|
||||
// previously proven for this same (auth_key_id, session_id). Auth-key-wide
|
||||
// device metadata is intentionally ineligible: a client upgrade can reuse its
|
||||
|
|
|
|||
|
|
@ -664,4 +664,24 @@ func TestGZIPExpansionUsesProcessBudgetBeforeDecode(t *testing.T) {
|
|||
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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -57,13 +57,17 @@ const layerRPCReplayRestoreTimeout = 5 * time.Second
|
|||
// touching auth/session stores. The MTProto edge must call it before acquiring
|
||||
// an RPC flight/cache slot or scheduling business work.
|
||||
func (r *Router) AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
return r.AdmitLayerWithOptions(profile, b, tlprofile.AdmissionOptions{Limits: limits})
|
||||
}
|
||||
|
||||
func (r *Router) AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
if r == nil || r.dispatcher == nil {
|
||||
return tlprofile.Admission{}, internalErr()
|
||||
}
|
||||
if b == nil {
|
||||
return tlprofile.Admission{}, inputRequestInvalidErr()
|
||||
}
|
||||
return r.dispatcher.Admit(profile, b, limits)
|
||||
return r.dispatcher.AdmitWithOptions(profile, b, options)
|
||||
}
|
||||
|
||||
// AdmitDefaultLayer admits a request using an inherited auth-key profile as
|
||||
|
|
@ -71,13 +75,17 @@ func (r *Router) AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlp
|
|||
// same wrapper chain to correct that default. Generated admission preserves
|
||||
// the distinction through EffectiveProfile and ProfileEvidence.
|
||||
func (r *Router) AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
return r.AdmitDefaultLayerWithOptions(profile, b, tlprofile.AdmissionOptions{Limits: limits})
|
||||
}
|
||||
|
||||
func (r *Router) AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
if r == nil || r.dispatcher == nil {
|
||||
return tlprofile.Admission{}, internalErr()
|
||||
}
|
||||
if b == nil {
|
||||
return tlprofile.Admission{}, inputRequestInvalidErr()
|
||||
}
|
||||
return r.dispatcher.AdmitDefault(profile, b, limits)
|
||||
return r.dispatcher.AdmitDefaultWithOptions(profile, b, options)
|
||||
}
|
||||
|
||||
// registerAndroidLayerRPCAdapter installs the only client-private schema seam.
|
||||
|
|
@ -191,10 +199,14 @@ func (r *Router) PrepareAdmittedReplay(
|
|||
// a closed terminal whose complete request and result wire graphs were proven
|
||||
// invariant across every generated profile. The latter never freezes a layer.
|
||||
func (r *Router) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
|
||||
return r.AdmitUnprofiledWithOptions(b, tlprofile.AdmissionOptions{Limits: limits})
|
||||
}
|
||||
|
||||
func (r *Router) AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
|
||||
if r == nil || r.dispatcher == nil {
|
||||
return tlprofile.Admission{}, internalErr()
|
||||
}
|
||||
return r.dispatcher.AdmitUnprofiled(b, limits)
|
||||
return r.dispatcher.AdmitUnprofiledWithOptions(b, options)
|
||||
}
|
||||
|
||||
// DispatchAdmitted executes one generated admission lease. invokeAfterMsg(s)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue