sync: # file upload improve

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
+	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{

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) {

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

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 {

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 big {
-		if wire.WireSize() < 20 {
+	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 {
 		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")

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))
 	}