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

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

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