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

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