diff --git a/internal/mtprotoedge/inbound_layer_rpc.go b/internal/mtprotoedge/inbound_layer_rpc.go index 92c8010d..8ab2e468 100644 --- a/internal/mtprotoedge/inbound_layer_rpc.go +++ b/internal/mtprotoedge/inbound_layer_rpc.go @@ -42,6 +42,7 @@ const ( // 32-MiB per-connection budget and therefore remains admissible by default. layerRPCAdmissionStaticObjectBytes = 512 layerRPCAdmissionWireFactor = 60 + layerRPCAdmissionFlatBytesFactor = 2 layerRPCAdmissionGraphSlack = layerRPCAdmissionStaticObjectBytes * 32 ) @@ -164,33 +165,91 @@ func (s *Server) initialLayerRPCAdmissionCursor(ctx context.Context, c *Conn) (l // Saturation deliberately turns hostile integer-sized inputs into ordinary // capacity rejection; it must never wrap into a small accepted reservation. func layerRPCAdmissionReservationSize(wireBytes int) int { + return saturatingLayerRPCAdmissionCharge( + layerRPCAdmissionGraphSlack, + layerRPCAdmissionWireCharge(wireBytes), + ) +} + +func layerRPCAdmissionWireCharge(wireBytes int) int { if wireBytes < 0 { wireBytes = 0 } maxInt := int(^uint(0) >> 1) - if wireBytes > (maxInt-layerRPCAdmissionGraphSlack)/layerRPCAdmissionWireFactor { + if wireBytes > maxInt/layerRPCAdmissionWireFactor { return maxInt } - return wireBytes*layerRPCAdmissionWireFactor + layerRPCAdmissionGraphSlack + return wireBytes * layerRPCAdmissionWireFactor +} + +// layerRPCFlatBytesWireCharge keeps the generic factor on fixed TL wire and +// applies a tight copy factor only to a payload which the production Router has +// already proven to be one bounded flat bytes field. The temporary expanded +// buffer remains independently owned by inboundFrameBudget while generated +// admission materializes the request; 2x covers the retained bytes copy plus +// allocator rounding without treating every payload byte as a possible nested +// object/vector node. The per-request graph slack is owned by the base +// reservation and must not be repeated for each nested gzip expansion. +func layerRPCFlatBytesWireCharge(wireBytes, payloadBytes int) int { + if wireBytes < 0 || payloadBytes < 0 || payloadBytes > wireBytes { + return layerRPCAdmissionWireCharge(wireBytes) + } + fixedBytes := wireBytes - payloadBytes + maxInt := int(^uint(0) >> 1) + if fixedBytes > maxInt/layerRPCAdmissionWireFactor { + return maxInt + } + charge := fixedBytes * layerRPCAdmissionWireFactor + if payloadBytes > (maxInt-charge)/layerRPCAdmissionFlatBytesFactor { + return maxInt + } + return charge + payloadBytes*layerRPCAdmissionFlatBytesFactor +} + +func saturatingLayerRPCAdmissionCharge(left, right int) int { + if left < 0 { + left = 0 + } + if right < 0 { + right = 0 + } + maxInt := int(^uint(0) >> 1) + if right > maxInt-left { + return maxInt + } + return left + right +} + +func (s *Server) layerRPCExpandedWireCharge(wire []byte) int { + if s != nil { + if sizer, ok := s.layerRPC.(LayerRPCFlatBytesPayloadSizer); ok { + if payloadBytes, flat := sizer.LayerRPCFlatBytesPayloadSize(wire); flat && payloadBytes >= 0 && payloadBytes <= len(wire) { + return layerRPCFlatBytesWireCharge(len(wire), payloadBytes) + } + } + } + return layerRPCAdmissionWireCharge(len(wire)) } // 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. +// to the durable scheduler charge of one exact typed request. The original +// compressed wire retains the generic graph charge. Every successful expansion +// adds its own charge; only a handler-proven flat bytes terminal can use the +// tighter payload factor. An authoritative-profile re-decode reuses the largest +// already-held charge instead of double-counting sequential attempts. type layerRPCGZIPExpansionBudget struct { - server *Server - plan *inboundPlan - reservation *inboundRPCBatchReservation - entry int - baseSourceBytes int - attemptExpandedBytes int - chargedSourceBytes int + server *Server + plan *inboundPlan + reservation *inboundRPCBatchReservation + entry int + baseCharge int + attemptExpandedCharge int + chargedSize int } func (b *layerRPCGZIPExpansionBudget) beginAttempt() { if b != nil { - b.attemptExpandedBytes = 0 + b.attemptExpandedCharge = 0 } } @@ -221,17 +280,19 @@ func (b *layerRPCGZIPExpansionBudget) expand(wire []byte, admissionLimit int) ([ 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 { + b.attemptExpandedCharge = saturatingLayerRPCAdmissionCharge( + b.attemptExpandedCharge, + b.server.layerRPCExpandedWireCharge(data), + ) + targetCharge := saturatingLayerRPCAdmissionCharge(b.baseCharge, b.attemptExpandedCharge) + if targetCharge > b.chargedSize { 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 + b.chargedSize = targetCharge } return data, release, nil } @@ -293,7 +354,9 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan } expansionBudget := &layerRPCGZIPExpansionBudget{ server: s, plan: plan, reservation: reservation, - entry: reservationIndex, baseSourceBytes: len(item.body), chargedSourceBytes: len(item.body), + entry: reservationIndex, + baseCharge: provisionalSpecs[reservationIndex].size, + chargedSize: provisionalSpecs[reservationIndex].size, } expansionBudgets[reservationIndex] = expansionBudget options := tlprofile.AdmissionOptions{ diff --git a/internal/mtprotoedge/layer_admission_budget_test.go b/internal/mtprotoedge/layer_admission_budget_test.go index 35b9e040..20489819 100644 --- a/internal/mtprotoedge/layer_admission_budget_test.go +++ b/internal/mtprotoedge/layer_admission_budget_test.go @@ -12,6 +12,7 @@ import ( "github.com/iamxvbaba/td/clock" "github.com/iamxvbaba/td/proto" "github.com/iamxvbaba/td/tg" + "github.com/iamxvbaba/td/tgerr" "go.uber.org/zap/zaptest" "github.com/iamxvbaba/td/tlprofile" @@ -85,6 +86,18 @@ func TestLayerRPCAdmissionMaterializationConstants(t *testing.T) { if got := layerRPCAdmissionReservationSize(-1); got != layerRPCAdmissionGraphSlack { t.Fatalf("negative wire charge = %d, want fixed slack %d", got, layerRPCAdmissionGraphSlack) } + if got := layerRPCFlatBytesWireCharge(100, 80); got != 20*layerRPCAdmissionWireFactor+80*layerRPCAdmissionFlatBytesFactor { + t.Fatalf("flat bytes wire charge = %d", got) + } + if got := layerRPCFlatBytesWireCharge(10, 11); got != layerRPCAdmissionWireCharge(10) { + t.Fatalf("invalid flat bytes hint charge = %d, want generic %d", got, layerRPCAdmissionWireCharge(10)) + } + if got := layerRPCFlatBytesWireCharge(maxInt, maxInt); got != maxInt { + t.Fatalf("saturating flat bytes wire charge = %d, want max int %d", got, maxInt) + } + if got := saturatingLayerRPCAdmissionCharge(maxInt, 1); got != maxInt { + t.Fatalf("saturating addition = %d, want max int %d", got, maxInt) + } } type countingLayerRPCAdmission struct { @@ -196,6 +209,274 @@ func TestLayerRPCAdmissionExpandsTDLibNestedGZIPUnderTransferredBudget(t *testin } } +func TestLayerRPCAdmissionAdmitsTDLibUploadParts(t *testing.T) { + tests := []struct { + name string + method string + body bin.Object + withoutGZIP bool + bare bool + }{ + { + name: "small_file_part", + method: "upload.saveFilePart", + body: &tg.UploadSaveFilePartRequest{ + FileID: 91, + FilePart: 0, + Bytes: make([]byte, 1071), + }, + }, + { + name: "pixel_9a_big_file_part_gzip", + method: "upload.saveBigFilePart", + body: &tg.UploadSaveBigFilePartRequest{ + FileID: 92, + FilePart: 0, + FileTotalParts: 364, + Bytes: make([]byte, 64<<10), + }, + }, + { + name: "pixel_9a_big_file_part_plain", + method: "upload.saveBigFilePart", + withoutGZIP: true, + body: &tg.UploadSaveBigFilePartRequest{ + FileID: 93, + FilePart: 7, + FileTotalParts: 364, + Bytes: make([]byte, 64<<10), + }, + }, + { + name: "pixel_9a_big_file_part_bare_upload_session", + method: "upload.saveBigFilePart", + bare: true, + body: &tg.UploadSaveBigFilePartRequest{ + FileID: 94, + FilePart: 15, + FileTotalParts: 364, + Bytes: make([]byte, 64<<10), + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(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, 31}, sessionID: 831, metrics: NopMetrics{}} + c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + s.rpcScheduler.stop(time.Second) + }() + + var ( + body []byte + expandedBytes int + ) + if tc.bare { + body = exactOutboundLayerRPCBody(t, tlprofile.Profile228, tc.body) + } else if tc.withoutGZIP { + body = tdlibWrappedBody(t, tlprofile.Profile228, tc.body) + } else { + body, expandedBytes = tdlibNestedGZIPBody(t, tlprofile.Profile228, tc.body) + } + 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 upload plan = kind:%d reservation:%v tasks:%d payload:%v", + plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks), plan.items[0].payload) + } + if got := plan.gzipExpandedBytes; got != expandedBytes { + t.Fatalf("nested gzip upload cumulative expansion = %d, want %d", got, expandedBytes) + } + + // Admission alone is insufficient: the prepared wrapper chain must remain + // executable after the temporary gzip expansion buffer has been released. + // With no Files dependency configured, reaching the upload handler has the + // stable terminal NOT_IMPLEMENTED; INPUT_REQUEST_INVALID means the wrapper or + // prepared-call boundary corrupted the otherwise valid request. + requestBody := &bin.Buffer{Buf: append([]byte(nil), body...)} + admitted, err := router.AdmitLayerWithOptions(tlprofile.Profile228, requestBody, tlprofile.AdmissionOptions{ + Limits: inboundLayerDecodeLimits, + ExpandGZIP: func(wire []byte, limit int) ([]byte, func(), error) { + return s.decodeGZIPWithGlobalBudgetLimit(&bin.Buffer{Buf: wire}, limit) + }, + }) + if err != nil { + t.Fatal(err) + } + _, method, err := router.DispatchAdmitted( + rpc.WithUserID(context.Background(), 42), + [8]byte{8, 31}, + 831, + 100, + 1, + admitted, + ) + if method != tc.method || !tgerr.Is(err, "NOT_IMPLEMENTED") { + t.Fatalf("nested gzip upload dispatch = method:%q err:%v, want %s/NOT_IMPLEMENTED", method, err, tc.method) + } + }) + } +} + +func TestLayerRPCAdmissionAdmitsTDLibFirstUploadContainer(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, 32}, sessionID: 832, metrics: NopMetrics{}} + c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second) + defer func() { + c.closeInboundRPCScheduler() + s.rpcScheduler.stop(time.Second) + }() + + // A newly opened TDLib upload Session applies its invokeWithLayer / + // initConnection header to every query in the first MTProto container. The + // Pixel 9a trace contains eight gzip-packed 64 KiB saveBigFilePart requests + // in that first container, all with the same known total and distinct + // parts/message IDs. The fixture is an ELF and its first chunks compress to + // roughly 11-14 KiB each, yielding a 129 KiB encrypted write. + plan, legacyGenericCharge := tdlibFirstUploadPlan(t) + defer plan.close() + if legacyGenericCharge <= maxInflightRPCBytes { + t.Fatalf("test fixture generic charge = %d, must exceed old connection ceiling %d", legacyGenericCharge, maxInflightRPCBytes) + } + + if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil { + t.Fatal(err) + } + if plan.rpcReservation == nil || len(plan.rpcTasks) != len(plan.items) { + t.Fatalf("first upload container reservation/tasks = %v/%d, want retained/%d", + plan.rpcReservation != nil, len(plan.rpcTasks), len(plan.items)) + } + if got := plan.rpcReservation.totalSize; got <= 0 || got >= legacyGenericCharge || got > maxInflightRPCBytes { + t.Fatalf("first upload container retained charge = %d, want 0 < charge < legacy %d and <= %d", + got, legacyGenericCharge, maxInflightRPCBytes) + } + for index := range plan.items { + item := &plan.items[index] + if item.kind != inboundItemRPC { + t.Fatalf("first upload container item %d kind = %d, payload = %v", index, item.kind, item.payload) + } + if method := plan.rpcTasks[index].method; method != "upload.saveBigFilePart" { + t.Fatalf("first upload container task %d method = %q, want upload.saveBigFilePart", index, method) + } + } +} + +func TestLayerRPCAdmissionTDLibFirstUploadContainerRequiresFlatBytesCapability(t *testing.T) { + router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System) + // The wrapper deliberately exposes exact admission but not the optional + // flat-bytes sizing capability. The edge must therefore retain the generic + // worst-case charge and reject the whole oversized batch atomically. + handler := &countingLayerRPCAdmission{LayerRPCHandler: router} + s := New(Options{DC: 2, LayerRPC: handler, Logger: zaptest.NewLogger(t)}) + c := &Conn{authKeyID: [8]byte{8, 33}, sessionID: 833, metrics: NopMetrics{}} + c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second) + defer func() { + c.closeInboundRPCScheduler() + s.rpcScheduler.stop(time.Second) + }() + + plan, _ := tdlibFirstUploadPlan(t) + 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("generic fallback retained reservation/tasks = %v/%d", plan.rpcReservation != nil, len(plan.rpcTasks)) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("generic fallback leaked connection charge %d", got) + } + if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("generic fallback leaked global budget %d/%d", tasks, bytes) + } +} + +func tdlibFirstUploadPlan(t *testing.T) (*inboundPlan, int64) { + t.Helper() + plan := &inboundPlan{items: make([]inboundItem, 8)} + var legacyGenericCharge int64 + for index := range plan.items { + payload := make([]byte, 64<<10) + state := uint32(index + 1) + for byteIndex := 0; byteIndex < 13<<10; byteIndex++ { + state ^= state << 13 + state ^= state >> 17 + state ^= state << 5 + payload[byteIndex] = byte(state) + } + body, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.UploadSaveBigFilePartRequest{ + FileID: 95, + FilePart: index, + FileTotalParts: 364, + Bytes: payload, + }) + plan.items[index] = inboundItem{ + kind: inboundItemRPC, + msgID: int64(100 + index*4), + body: body, + } + legacyGenericCharge += int64(layerRPCAdmissionReservationSize(len(body) + expandedBytes)) + } + return plan, legacyGenericCharge +} + +func TestLayerRPCAdmissionAdmitsTDLibWrappedBindTempAuthKey(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, 41}, sessionID: -7940676790771565328, metrics: NopMetrics{}} + c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + s.rpcScheduler.stop(time.Second) + }() + + request := &tg.AuthBindTempAuthKeyRequest{ + PermAuthKeyID: 9179421154451858694, + Nonce: 5318578202586482454, + ExpiresAt: 1785817822, + EncryptedMessage: make([]byte, 104), + } + body := tdlibWrappedBody(t, tlprofile.Profile228, request) + 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 TDLib bind plan = kind:%d reservation:%v tasks:%d payload:%v", + plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks), plan.items[0].payload) + } + requestBody := &bin.Buffer{Buf: append([]byte(nil), body...)} + admitted, err := router.AdmitUnprofiled(requestBody, inboundLayerDecodeLimits) + if err != nil { + t.Fatal(err) + } + result, method, err := router.DispatchAdmitted( + context.Background(), + c.authKeyID, + c.sessionID, + 100, + 1, + admitted, + ) + if err != nil || method != "auth.bindTempAuthKey" || result == nil || !result.WireInvariant() { + t.Fatalf("TDLib wrapped bind dispatch = method:%q result:%T invariant:%v err:%v", + method, result, result != nil && result.WireInvariant(), err) + } +} + 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)}) @@ -313,7 +594,7 @@ func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testi plan225 := &inboundPlan{} budget225 := &layerRPCGZIPExpansionBudget{ server: s, plan: plan225, reservation: reservation225, - baseSourceBytes: len(body), chargedSourceBytes: len(body), + baseCharge: initialCharge, chargedSize: initialCharge, } options225 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget225.expand} item225 := inboundItem{msgID: 100, body: body} @@ -327,7 +608,7 @@ func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testi plan227 := &inboundPlan{} budget227 := &layerRPCGZIPExpansionBudget{ server: s, plan: plan227, reservation: reservation227, - baseSourceBytes: len(body), chargedSourceBytes: len(body), + baseCharge: initialCharge, chargedSize: initialCharge, } options227 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget227.expand} item227 := inboundItem{msgID: 100, body: body} @@ -413,6 +694,11 @@ 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) + return tdlibWrappedBody(t, profile, &proto.GZIP{Data: terminalWire}), len(terminalWire) +} + +func tdlibWrappedBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) []byte { + t.Helper() request := &tg.InvokeWithLayerRequest{ Layer: int(profile), Query: &tg.InitConnectionRequest{ @@ -423,10 +709,14 @@ func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.O SystemLangCode: "en", LangPack: "", LangCode: "en", - Query: &proto.GZIP{Data: terminalWire}, + Params: &tg.JSONObject{Value: []tg.JSONObjectValue{{ + Key: "tz_offset", + Value: &tg.JSONNumber{Value: 8 * 60 * 60}, + }}}, + Query: terminal, }, } - return exactLayerRPCBody(t, request), len(terminalWire) + return exactLayerRPCBody(t, request) } func TestLayerRPCAdmissionPendingReplayReleasesProvisionalEntry(t *testing.T) { diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index 89ba8a76..a34b85c8 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -93,6 +93,17 @@ type LayerRPCDefaultProfileOptionsAdmitter interface { AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) } +// LayerRPCFlatBytesPayloadSizer is an optional, allocation-free admission +// capability for exact terminal requests whose generated object graph contains +// one already-bounded flat bytes payload. The handler may return ok only after +// proving the complete terminal wire shape and every semantic field cap it +// relies on. The edge still owns the multiplier, graph slack and all process / +// connection budgets; an absent or invalid hint falls back to the conservative +// generic graph charge. +type LayerRPCFlatBytesPayloadSizer interface { + LayerRPCFlatBytesPayloadSize(wire []byte) (payloadBytes int, ok bool) +} + // 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/rpc/layer_dispatch_test.go b/internal/rpc/layer_dispatch_test.go index 986defe6..f3822539 100644 --- a/internal/rpc/layer_dispatch_test.go +++ b/internal/rpc/layer_dispatch_test.go @@ -586,7 +586,7 @@ func TestLayerAdmissionFieldPoliciesCoverEveryRoutableProfile(t *testing.T) { r := New(Config{DC: 2, IP: "127.0.0.1", Port: 2398}, Deps{}, zaptest.NewLogger(t), clock.System) limits := tlprofile.Limits{MaxVectorElements: 8 << 10} - for profile := tlprofile.Profile225; profile <= tlprofile.Profile227; profile++ { + for profile := tlprofile.Profile225; profile <= tlprofile.Profile228; profile++ { for _, tc := range cases { tc := tc if _, available := tlprofile.WireID(profile, tc.method); !available { diff --git a/internal/rpc/request_preflight.go b/internal/rpc/request_preflight.go index 43c6993e..10a039d0 100644 --- a/internal/rpc/request_preflight.go +++ b/internal/rpc/request_preflight.go @@ -178,6 +178,37 @@ func preflightRPCWire(id uint32, wire rpcPreflightWire) error { } } +// LayerRPCFlatBytesPayloadSize exposes the only two byte-heavy, flat request +// graphs to mtprotoedge's materialization budget. Returning ok is deliberately +// stricter than recognizing the constructor: the complete TL bytes tail, +// per-part byte ceiling and big-file total-parts policy must all pass the same +// allocation-free preflight used by generated exact admission. All other RPCs +// retain the generic worst-case graph charge. +func (r *Router) LayerRPCFlatBytesPayloadSize(wire []byte) (int, bool) { + if r == nil || len(wire) < 4 { + return 0, false + } + raw := rawRPCPreflightWire(wire) + id := binary.LittleEndian.Uint32(wire[:4]) + bytesOffset := 0 + switch id { + case tg.UploadSaveFilePartRequestTypeID: + bytesOffset = 16 + case tg.UploadSaveBigFilePartRequestTypeID: + bytesOffset = 20 + default: + return 0, false + } + if err := preflightRPCWire(id, raw); err != nil { + return 0, false + } + payloadBytes, encodedBytes, err := tlBytesSizeAt(raw, bytesOffset) + if err != nil || encodedBytes != len(wire)-bytesOffset { + return 0, false + } + return payloadBytes, true +} + func preflightFixedVector(wire rpcPreflightWire, policy requestVectorPolicy) error { if policy.vectorOffset < 4 || policy.minElemBytes <= 0 || wire.WireSize() < policy.vectorOffset+8 { return inputRequestInvalidErr() @@ -210,10 +241,18 @@ func preflightFixedVector(wire rpcPreflightWire, policy requestVectorPolicy) err } func preflightUploadPart(wire rpcPreflightWire, bytesOffset int, big bool) error { + if wire.WireSize() < bytesOffset { + return inputRequestInvalidErr() + } + rawPart, err := wire.Uint32At(12) + if err != nil { + return inputRequestInvalidErr() + } + part := int32(rawPart) + if part < 0 || part >= int32(appfiles.MaxUploadParts) { + return filePartInvalidErr() + } if big { - if wire.WireSize() < 20 { - return inputRequestInvalidErr() - } rawTotalParts, err := wire.Uint32At(16) if err != nil { return inputRequestInvalidErr() @@ -230,14 +269,17 @@ func preflightUploadPart(wire rpcPreflightWire, bytesOffset int, big bool) error if encoded != wire.WireSize()-bytesOffset { return inputRequestInvalidErr() } + if n == 0 { + return filePartInvalidErr() + } if n > appfiles.MaxUploadPartBytes { return filePartTooBigErr() } return nil } -// tlBytesSizeAt parses a TL bytes prefix without copying the payload. encoded includes prefix, -// payload and 4-byte padding. +// tlBytesSizeAt parses a TL bytes prefix without copying the payload. encoded +// includes prefix, payload and 4-byte padding. func tlBytesSizeAt(wire rpcPreflightWire, offset int) (n, encoded int, err error) { if offset < 0 || offset >= wire.WireSize() { return 0, 0, fmt.Errorf("bytes prefix out of range") diff --git a/internal/rpc/request_preflight_test.go b/internal/rpc/request_preflight_test.go index 7d3768db..71c1ed6d 100644 --- a/internal/rpc/request_preflight_test.go +++ b/internal/rpc/request_preflight_test.go @@ -70,6 +70,7 @@ func TestUploadPartPreflightBeforeBytesDecode(t *testing.T) { id uint32 offset int big bool + part int parts int size int want string @@ -78,11 +79,14 @@ func TestUploadPartPreflightBeforeBytesDecode(t *testing.T) { {name: "small_at_cap", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: appfiles.MaxUploadPartBytes}, {name: "small_over_cap", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: appfiles.MaxUploadPartBytes + 1, want: "FILE_PART_TOO_BIG"}, {name: "big_at_cap", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, big: true, parts: appfiles.MaxUploadParts, size: appfiles.MaxUploadPartBytes}, + {name: "part_negative", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, part: -1, size: 1, want: "FILE_PART_INVALID"}, + {name: "part_at_limit", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, part: appfiles.MaxUploadParts, parts: appfiles.MaxUploadParts, size: 1, want: "FILE_PART_INVALID"}, {name: "big_parts_over_cap", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, big: true, parts: appfiles.MaxUploadParts + 1, size: 1, want: "FILE_PART_INVALID"}, + {name: "empty", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 0, want: "FILE_PART_INVALID"}, {name: "truncated", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 1024, truncateBy: 1, want: "INPUT_REQUEST_INVALID"}, } { t.Run(tc.name, func(t *testing.T) { - raw := uploadPartRequest(tc.id, tc.offset, tc.parts, tc.size) + raw := uploadPartRequest(tc.id, tc.offset, tc.part, tc.parts, tc.size) if tc.truncateBy > 0 { raw = raw[:len(raw)-tc.truncateBy] } @@ -100,6 +104,51 @@ func TestUploadPartPreflightBeforeBytesDecode(t *testing.T) { } } +func TestLayerRPCFlatBytesPayloadSizeRequiresCompleteLegalUpload(t *testing.T) { + router := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + for _, tc := range []struct { + name string + id uint32 + offset int + part int + parts int + size int + mutate func([]byte) []byte + wantOK bool + }{ + {name: "small", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 64 << 10, wantOK: true}, + {name: "big", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, parts: 364, size: 64 << 10, wantOK: true}, + {name: "part_negative", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, part: -1, size: 1}, + {name: "part_at_limit", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, part: appfiles.MaxUploadParts, parts: 364, size: 1}, + {name: "big_total_invalid", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, parts: 0, size: 64 << 10}, + {name: "empty_payload", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 0}, + {name: "payload_over_cap", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: appfiles.MaxUploadPartBytes + 1}, + { + name: "truncated", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 64 << 10, + mutate: func(wire []byte) []byte { return wire[:len(wire)-1] }, + }, + { + name: "trailing", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 64 << 10, + mutate: func(wire []byte) []byte { return append(wire, 0, 0, 0, 0) }, + }, + {name: "other_rpc", id: tg.HelpGetConfigRequestTypeID, offset: 4}, + } { + t.Run(tc.name, func(t *testing.T) { + wire := uploadPartRequest(tc.id, tc.offset, tc.part, tc.parts, tc.size) + if tc.mutate != nil { + wire = tc.mutate(wire) + } + payloadBytes, ok := router.LayerRPCFlatBytesPayloadSize(wire) + if ok != tc.wantOK { + t.Fatalf("flat payload = %d/%v, want ok=%v", payloadBytes, ok, tc.wantOK) + } + if ok && payloadBytes != tc.size { + t.Fatalf("flat payload size = %d, want %d", payloadBytes, tc.size) + } + }) + } +} + func fixedVectorRequest(id uint32, policy requestVectorPolicy, count int) []byte { raw := make([]byte, policy.vectorOffset+8+count*policy.minElemBytes) binary.LittleEndian.PutUint32(raw[0:4], id) @@ -108,9 +157,12 @@ func fixedVectorRequest(id uint32, policy requestVectorPolicy, count int) []byte return raw } -func uploadPartRequest(id uint32, offset, parts, size int) []byte { +func uploadPartRequest(id uint32, offset, part, parts, size int) []byte { raw := make([]byte, offset) binary.LittleEndian.PutUint32(raw[:4], id) + if offset >= 16 { + binary.LittleEndian.PutUint32(raw[12:16], uint32(part)) + } if offset == 20 { binary.LittleEndian.PutUint32(raw[16:20], uint32(parts)) }