diff --git a/go.mod b/go.mod index 5a2a8fba..8ec2b2e3 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,7 @@ require ( github.com/golang-migrate/migrate/v4 v4.19.1 github.com/gotd/ige v0.3.0 github.com/gotd/log/logzap v0.1.1 - github.com/iamxvbaba/td v1.1.5 + github.com/iamxvbaba/td v1.2.0 github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa github.com/jackc/pgx/v5 v5.9.2 github.com/lestrrat-go/jwx/v3 v3.1.1 diff --git a/go.sum b/go.sum index 5b369364..f4775f77 100644 --- a/go.sum +++ b/go.sum @@ -81,8 +81,8 @@ github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= github.com/hashicorp/go-multierror v1.1.0 h1:B9UzwGQJehnUY1yNrnwREHc3fGbC2xefo8g4TbElacI= github.com/hashicorp/go-multierror v1.1.0/go.mod h1:spPvp8C1qA32ftKqdAHm4hHTbPw+vmowP0z+KUhOZdA= -github.com/iamxvbaba/td v1.1.5 h1:fy6wWQIMzKKOXBwBRcnDaLFhiHyQazJt9eMOrdhl22E= -github.com/iamxvbaba/td v1.1.5/go.mod h1:INkZJi18XbXtVOrldDnPtmQCJvIXhgDZdbo9MV/NY7M= +github.com/iamxvbaba/td v1.2.0 h1:3RqZir1Yk2uEkECh8JKhXW8cW+dYIGa/21+Q7FtsHnE= +github.com/iamxvbaba/td v1.2.0/go.mod h1:INkZJi18XbXtVOrldDnPtmQCJvIXhgDZdbo9MV/NY7M= github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa h1:s+4MhCQ6YrzisK6hFJUX53drDT4UsSW3DEhKn0ifuHw= github.com/jackc/pgerrcode v0.0.0-20220416144525-469b46aa5efa/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index 0338ca81..11947849 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -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))) diff --git a/internal/mtprotoedge/inbound_layer_rpc.go b/internal/mtprotoedge/inbound_layer_rpc.go index 2b23eee2..5c419372 100644 --- a/internal/mtprotoedge/inbound_layer_rpc.go +++ b/internal/mtprotoedge/inbound_layer_rpc.go @@ -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" diff --git a/internal/mtprotoedge/inbound_layer_rpc_test.go b/internal/mtprotoedge/inbound_layer_rpc_test.go index f0f3fe32..90905557 100644 --- a/internal/mtprotoedge/inbound_layer_rpc_test.go +++ b/internal/mtprotoedge/inbound_layer_rpc_test.go @@ -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() diff --git a/internal/mtprotoedge/inbound_preflight.go b/internal/mtprotoedge/inbound_preflight.go index 256fa120..1b8296bc 100644 --- a/internal/mtprotoedge/inbound_preflight.go +++ b/internal/mtprotoedge/inbound_preflight.go @@ -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() diff --git a/internal/mtprotoedge/inbound_rpc.go b/internal/mtprotoedge/inbound_rpc.go index f1e2abeb..7f0c9bb1 100644 --- a/internal/mtprotoedge/inbound_rpc.go +++ b/internal/mtprotoedge/inbound_rpc.go @@ -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 diff --git a/internal/mtprotoedge/inbound_rpc_batch_test.go b/internal/mtprotoedge/inbound_rpc_batch_test.go index 423b1ab2..5efc4fba 100644 --- a/internal/mtprotoedge/inbound_rpc_batch_test.go +++ b/internal/mtprotoedge/inbound_rpc_batch_test.go @@ -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) diff --git a/internal/mtprotoedge/layer_admission_budget_test.go b/internal/mtprotoedge/layer_admission_budget_test.go index 3202a7c9..0a4e9d3b 100644 --- a/internal/mtprotoedge/layer_admission_budget_test.go +++ b/internal/mtprotoedge/layer_admission_budget_test.go @@ -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}) diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index 4524b144..de0cc6a6 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -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 diff --git a/internal/mtprotoedge/structural_limits_test.go b/internal/mtprotoedge/structural_limits_test.go index 12ef3512..652019c6 100644 --- a/internal/mtprotoedge/structural_limits_test.go +++ b/internal/mtprotoedge/structural_limits_test.go @@ -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) + } } diff --git a/internal/rpc/layer_dispatch.go b/internal/rpc/layer_dispatch.go index 69aca08a..1180be82 100644 --- a/internal/rpc/layer_dispatch.go +++ b/internal/rpc/layer_dispatch.go @@ -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)