Skip to content

Commit 422ba92

Browse files
Divjot Aroradmhilly
Divjot Arora
authored and
dmhilly
committed
Disallow partial intermediate chunks for GridFS.
GODRIVER-753 Change-Id: I07b2b6a05f07fa7ed86ffc881e9634d562df8a1f
1 parent 5521c38 commit 422ba92

File tree

4 files changed

+161
-12
lines changed

4 files changed

+161
-12
lines changed

‎data/gridfs/download.json

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,97 @@
355355
"error": "ChunkIsMissing"
356356
}
357357
},
358+
{
359+
"description": "Download when an intermediate chunk is the wrong size",
360+
"arrange": {
361+
"data": [
362+
{
363+
"update": "fs.chunks",
364+
"updates": [
365+
{
366+
"q": {
367+
"files_id": {
368+
"$oid": "000000000000000000000005"
369+
},
370+
"n": 1
371+
},
372+
"u": {
373+
"$set": {
374+
"data": {
375+
"$hex": "556677"
376+
}
377+
}
378+
}
379+
},
380+
{
381+
"q": {
382+
"files_id": {
383+
"$oid": "000000000000000000000005"
384+
},
385+
"n": 2
386+
},
387+
"u": {
388+
"$set": {
389+
"data": {
390+
"$hex": "8899aa"
391+
}
392+
}
393+
}
394+
}
395+
]
396+
}
397+
]
398+
},
399+
"act": {
400+
"operation": "download",
401+
"arguments": {
402+
"id": {
403+
"$oid": "000000000000000000000005"
404+
}
405+
}
406+
},
407+
"assert": {
408+
"error": "ChunkIsWrongSize"
409+
}
410+
},
411+
{
412+
"description": "Download when final chunk is the wrong size",
413+
"arrange": {
414+
"data": [
415+
{
416+
"update": "fs.chunks",
417+
"updates": [
418+
{
419+
"q": {
420+
"files_id": {
421+
"$oid": "000000000000000000000005"
422+
},
423+
"n": 2
424+
},
425+
"u": {
426+
"$set": {
427+
"data": {
428+
"$hex": "99"
429+
}
430+
}
431+
}
432+
}
433+
]
434+
}
435+
]
436+
},
437+
"act": {
438+
"operation": "download",
439+
"arguments": {
440+
"id": {
441+
"$oid": "000000000000000000000005"
442+
}
443+
}
444+
},
445+
"assert": {
446+
"error": "ChunkIsWrongSize"
447+
}
448+
},
358449
{
359450
"description": "Download legacy file with no name",
360451
"act": {

‎data/gridfs/download.yml

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,35 @@ tests:
152152
id: { "$oid" : "000000000000000000000005" }
153153
assert:
154154
error: "ChunkIsMissing"
155+
-
156+
description: "Download when an intermediate chunk is the wrong size"
157+
arrange:
158+
data:
159+
-
160+
{ update : "fs.chunks", updates : [
161+
{ q : { files_id : { "$oid" : "000000000000000000000005" }, n : 1 }, u : { $set : { data : { $hex : "556677" } } } },
162+
{ q : { files_id : { "$oid" : "000000000000000000000005" }, n : 2 }, u : { $set : { data : { $hex : "8899aa" } } } }
163+
] }
164+
act:
165+
operation: download
166+
arguments:
167+
id: { "$oid" : "000000000000000000000005" }
168+
assert:
169+
error: "ChunkIsWrongSize"
170+
-
171+
description: "Download when final chunk is the wrong size"
172+
arrange:
173+
data:
174+
-
175+
{ update : "fs.chunks", updates : [
176+
{ q : { files_id : { "$oid" : "000000000000000000000005" }, n : 2 }, u : { $set : { data : { $hex : "99" } } } }
177+
] }
178+
act:
179+
operation: download
180+
arguments:
181+
id: { "$oid" : "000000000000000000000005" }
182+
assert:
183+
error: "ChunkIsWrongSize"
155184
-
156185
description: "Download legacy file with no name"
157186
act:

‎mongo/gridfs/download_stream.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,20 @@ func (ds *DownloadStream) fillBuffer(ctx context.Context) error {
190190
_, dataBytes := data.Binary()
191191
copied := copy(ds.buffer, dataBytes)
192192

193+
bytesLen := int32(len(dataBytes))
194+
if ds.expectedChunk == ds.numChunks {
195+
// final chunk can be fewer than ds.chunkSize bytes
196+
bytesDownloaded := ds.chunkSize * (ds.expectedChunk - 1)
197+
bytesRemaining := ds.fileLen - int64(bytesDownloaded)
198+
199+
if int64(bytesLen) != bytesRemaining {
200+
return ErrWrongSize
201+
}
202+
} else if bytesLen != ds.chunkSize {
203+
// all intermediate chunks must have size ds.chunkSize
204+
return ErrWrongSize
205+
}
206+
193207
ds.bufferStart = 0
194208
ds.bufferEnd = copied
195209

‎mongo/gridfs/upload_stream.go

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,13 @@ package gridfs
99
import (
1010
"errors"
1111

12-
"math"
13-
1412
"context"
1513
"time"
1614

1715
"github.com/mongodb/mongo-go-driver/bson/primitive"
1816
"github.com/mongodb/mongo-go-driver/mongo"
1917
"github.com/mongodb/mongo-go-driver/x/bsonx"
18+
"math"
2019
)
2120

2221
// UploadBufferSize is the size in bytes of one stream batch. Chunks will be written to the db after the sum of chunk
@@ -67,7 +66,7 @@ func (us *UploadStream) Close() error {
6766
}
6867

6968
if us.bufferIndex != 0 {
70-
if err := us.uploadChunks(ctx); err != nil {
69+
if err := us.uploadChunks(ctx, true); err != nil {
7170
return err
7271
}
7372
}
@@ -115,11 +114,10 @@ func (us *UploadStream) Write(p []byte) (int, error) {
115114
us.bufferIndex += n
116115

117116
if us.bufferIndex == UploadBufferSize {
118-
err := us.uploadChunks(ctx)
117+
err := us.uploadChunks(ctx, false)
119118
if err != nil {
120119
return 0, err
121120
}
122-
us.bufferIndex = 0
123121
}
124122
}
125123
return origLen, nil
@@ -145,25 +143,36 @@ func (us *UploadStream) Abort() error {
145143
return nil
146144
}
147145

148-
func (us *UploadStream) uploadChunks(ctx context.Context) error {
149-
numChunks := math.Ceil(float64(us.bufferIndex) / float64(us.chunkSize))
146+
// uploadChunks uploads the current buffer as a series of chunks to the bucket
147+
// if uploadPartial is true, any data at the end of the buffer that is smaller than a chunk will be uploaded as a partial
148+
// chunk. if it is false, the data will be moved to the front of the buffer.
149+
// uploadChunks sets us.bufferIndex to the next available index in the buffer after uploading
150+
func (us *UploadStream) uploadChunks(ctx context.Context, uploadPartial bool) error {
151+
chunks := float64(us.bufferIndex) / float64(us.chunkSize)
152+
numChunks := int(math.Ceil(chunks))
153+
if !uploadPartial {
154+
numChunks = int(math.Floor(chunks))
155+
}
150156

151157
docs := make([]interface{}, int(numChunks))
158+
152159
begChunkIndex := us.chunkIndex
153160
for i := 0; i < us.bufferIndex; i += int(us.chunkSize) {
154-
var chunkData []byte
161+
endIndex := i + int(us.chunkSize)
155162
if us.bufferIndex-i < int(us.chunkSize) {
156-
chunkData = us.buffer[i:us.bufferIndex]
157-
} else {
158-
chunkData = us.buffer[i : i+int(us.chunkSize)]
163+
// partial chunk
164+
if !uploadPartial {
165+
break
166+
}
167+
endIndex = us.bufferIndex
159168
}
169+
chunkData := us.buffer[i:endIndex]
160170
docs[us.chunkIndex-begChunkIndex] = bsonx.Doc{
161171
{"_id", bsonx.ObjectID(primitive.NewObjectID())},
162172
{"files_id", bsonx.ObjectID(us.FileID)},
163173
{"n", bsonx.Int32(int32(us.chunkIndex))},
164174
{"data", bsonx.Binary(0x00, chunkData)},
165175
}
166-
167176
us.chunkIndex++
168177
us.fileLen += int64(len(chunkData))
169178
}
@@ -173,6 +182,12 @@ func (us *UploadStream) uploadChunks(ctx context.Context) error {
173182
return err
174183
}
175184

185+
// copy any remaining bytes to beginning of buffer and set buffer index
186+
bytesUploaded := numChunks * int(us.chunkSize)
187+
if bytesUploaded != UploadBufferSize && !uploadPartial {
188+
copy(us.buffer[0:], us.buffer[bytesUploaded:us.bufferIndex])
189+
}
190+
us.bufferIndex = UploadBufferSize - bytesUploaded
176191
return nil
177192
}
178193

0 commit comments

Comments
 (0)