astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

sync: # file upload improve

Browse source
This commit is contained in:
iamxvbaba 2026-08-03 17:33:49 +08:00
parent 9041abd3d3
commit e70bb0f167
6 changed files with 489 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

101
internal/mtprotoedge/inbound_layer_rpc.go
View file

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

298
internal/mtprotoedge/layer_admission_budget_test.go
View file

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

11
internal/mtprotoedge/server.go
View file

@ -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

2
internal/rpc/layer_dispatch_test.go
View file

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

52
internal/rpc/request_preflight.go
View file

@ -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")

56
internal/rpc/request_preflight_test.go
View file

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