Skip to content

Commit cafb660

Browse files
fallintoplacefallintoplace
authored
fix(parquet): validate DataPageV2 decoded length (#911)
1 parent 34dcc82 commit cafb660

2 files changed

Lines changed: 163 additions & 24 deletions

File tree

‎parquet/file/file_reader_test.go‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,108 @@ func (p *PageSerdeSuite) TestDataPageV2() {
215215
p.alloc.AssertSize(p.T(), 0)
216216
}
217217

218+
func (p *PageSerdeSuite) TestDataPageV2CorruptDecodedLength() {
219+
rawData := bytes.Repeat([]byte{0x42}, 128)
220+
codec, err := compress.GetCodec(compress.Codecs.Snappy)
221+
p.NoError(err)
222+
223+
buffer := make([]byte, int(codec.CompressBound(int64(len(rawData)))))
224+
compressed := codec.Encode(buffer, rawData)
225+
226+
p.dataPageHdrV2 = *format.NewDataPageHeaderV2()
227+
p.dataPageHdrV2.NumValues = 1
228+
p.dataPageHdrV2.NumRows = 1
229+
p.dataPageHdrV2.Encoding = format.Encoding_PLAIN
230+
231+
p.WriteDataPageHeaderV2(1024, int32(len(rawData)+1), int32(len(compressed)))
232+
_, err = p.sink.Write(compressed)
233+
p.NoError(err)
234+
235+
p.InitSerializedPageReader(1, compress.Codecs.Snappy)
236+
p.False(p.pageReader.Next())
237+
p.ErrorContains(p.pageReader.Err(), "metadata said 129 bytes uncompressed data page, got 128 bytes")
238+
p.NoError(p.pageReader.Close())
239+
p.alloc.AssertSize(p.T(), 0)
240+
}
241+
242+
func (p *PageSerdeSuite) TestDataPageV2Compressed() {
243+
rawData := bytes.Repeat([]byte{0x42}, 128)
244+
codec, err := compress.GetCodec(compress.Codecs.Snappy)
245+
p.NoError(err)
246+
247+
buffer := make([]byte, int(codec.CompressBound(int64(len(rawData)))))
248+
compressed := codec.Encode(buffer, rawData)
249+
250+
p.dataPageHdrV2 = *format.NewDataPageHeaderV2()
251+
p.dataPageHdrV2.NumValues = 1
252+
p.dataPageHdrV2.NumRows = 1
253+
p.dataPageHdrV2.Encoding = format.Encoding_PLAIN
254+
255+
p.WriteDataPageHeaderV2(1024, int32(len(rawData)), int32(len(compressed)))
256+
_, err = p.sink.Write(compressed)
257+
p.NoError(err)
258+
259+
p.InitSerializedPageReader(1, compress.Codecs.Snappy)
260+
p.True(p.pageReader.Next())
261+
p.Equal(rawData, p.pageReader.Page().Data())
262+
p.NoError(p.pageReader.Close())
263+
p.alloc.AssertSize(p.T(), 0)
264+
}
265+
266+
type mockDecryptor struct {
267+
decrypt func([]byte) []byte
268+
269+
aad []byte
270+
mem memory.Allocator
271+
}
272+
273+
func (m *mockDecryptor) FileAad() string { return "" }
274+
275+
func (m *mockDecryptor) Allocator() memory.Allocator {
276+
if m.mem == nil {
277+
return memory.DefaultAllocator
278+
}
279+
280+
return m.mem
281+
}
282+
283+
func (m *mockDecryptor) CiphertextSizeDelta() int { return 0 }
284+
285+
func (m *mockDecryptor) Decrypt(src []byte) []byte { return m.decrypt(src) }
286+
287+
func (m *mockDecryptor) DecryptFrom(r io.Reader) []byte {
288+
decrypted, _ := io.ReadAll(r)
289+
return m.Decrypt(decrypted)
290+
}
291+
292+
func (m *mockDecryptor) UpdateAad(aad string) { m.aad = []byte(aad) }
293+
294+
func (p *PageSerdeSuite) TestDataPageV2EncryptedUncompressedLengthMismatch() {
295+
p.dataPageHdrV2 = *format.NewDataPageHeaderV2()
296+
p.dataPageHdrV2.NumValues = 1
297+
p.dataPageHdrV2.NumRows = 1
298+
p.dataPageHdrV2.Encoding = format.Encoding_PLAIN
299+
p.dataPageHdrV2.IsCompressed = false
300+
301+
payload := bytes.Repeat([]byte{0xaa}, 4)
302+
decryptedPayload := []byte{0xaa}
303+
304+
p.WriteDataPageHeaderV2(1024, int32(len(payload)), int32(len(payload)))
305+
_, err := p.sink.Write(payload)
306+
p.NoError(err)
307+
308+
p.EndStream()
309+
p.alloc = memory.NewCheckedAllocator(memory.NewGoAllocator())
310+
p.pageReader, _ = file.NewPageReader(utils.NewByteReader(p.buffer.Bytes()), 1, compress.Codecs.Uncompressed, p.alloc, &file.CryptoContext{
311+
DataDecryptor: &mockDecryptor{decrypt: func([]byte) []byte { return decryptedPayload }},
312+
})
313+
314+
p.False(p.pageReader.Next())
315+
p.ErrorContains(p.pageReader.Err(), "metadata said 4 bytes uncompressed data page, got 1 bytes")
316+
p.NoError(p.pageReader.Close())
317+
p.alloc.AssertSize(p.T(), 0)
318+
}
319+
218320
func (p *PageSerdeSuite) TestLargePageHeaders() {
219321
const (
220322
statsSize = 256 * 1024 // 256KB

‎parquet/file/page_reader.go‎

Lines changed: 61 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -518,53 +518,85 @@ func (p *serializedPageReader) decompress(rd io.Reader, lenCompressed int, buf [
518518
return p.codec.Decode(buf, data), nil
519519
}
520520

521-
func (p *serializedPageReader) readV2Encrypted(rd io.Reader, lenCompressed int, levelsBytelen int, compressed bool, buf []byte) error {
521+
func (p *serializedPageReader) readV2Encrypted(rd io.Reader, lenCompressed int, lenUncompressed int, levelsBytelen int, compressed bool, buf []byte) (int, error) {
522522
// if encrypted, we need to decrypt before decompressing
523523
p.decompressBuffer.ResizeNoShrink(lenCompressed)
524524

525525
n, err := io.ReadFull(rd, p.decompressBuffer.Bytes()[:lenCompressed])
526526
if err != nil {
527-
return err
527+
return n, err
528528
}
529529
if n != lenCompressed {
530-
return fmt.Errorf("parquet: expected to read %d compressed bytes, got %d", lenCompressed, n)
530+
return n, fmt.Errorf("parquet: expected to read %d compressed bytes, got %d", lenCompressed, n)
531531
}
532532

533533
data := p.cryptoCtx.DataDecryptor.Decrypt(p.decompressBuffer.Bytes()[:lenCompressed])
534+
534535
// encrypted + uncompressed -> just copy the decrypted data to output buffer
535536
if !compressed {
536-
copy(buf, data)
537-
return nil
537+
if len(data) != lenUncompressed {
538+
return n, fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, len(data))
539+
}
540+
541+
copied := copy(buf, data)
542+
if copied != lenUncompressed {
543+
return copied, fmt.Errorf("parquet: expected to read %d uncompressed encrypted bytes, got %d", lenUncompressed, copied)
544+
}
545+
return copied, nil
538546
}
539547

540548
// definition + repetition levels are always uncompressed
541549
if levelsBytelen > 0 {
542-
copy(buf, data[:levelsBytelen])
550+
if len(data) < levelsBytelen {
551+
return n, fmt.Errorf("parquet: expected to read %d bytes for levels, got %d", levelsBytelen, len(data))
552+
}
553+
copied := copy(buf, data[:levelsBytelen])
554+
if copied != levelsBytelen {
555+
return copied, fmt.Errorf("parquet: expected to read %d bytes for levels, got %d", levelsBytelen, copied)
556+
}
543557
data = data[levelsBytelen:]
544558
}
545-
p.codec.Decode(buf[levelsBytelen:], data)
546-
return nil
559+
decoded := p.codec.Decode(buf[levelsBytelen:], data)
560+
if len(decoded) != lenUncompressed-levelsBytelen {
561+
return levelsBytelen + len(decoded), fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, levelsBytelen+len(decoded))
562+
}
563+
return levelsBytelen + len(decoded), nil
547564
}
548565

549-
func (p *serializedPageReader) readV2Unencrypted(rd io.Reader, lenCompressed int, levelsBytelen int, compressed bool, buf []byte) error {
566+
func (p *serializedPageReader) readV2Unencrypted(rd io.Reader, lenCompressed int, lenUncompressed int, levelsBytelen int, compressed bool, buf []byte) (int, error) {
550567
if !compressed {
551568
// uncompressed, just read into the buffer
552-
if _, err := io.ReadFull(rd, buf); err != nil {
553-
return err
569+
n, err := io.ReadFull(rd, buf)
570+
if err != nil {
571+
return n, err
572+
}
573+
if n != lenUncompressed {
574+
return n, fmt.Errorf("parquet: expected to read %d uncompressed bytes, got %d", lenUncompressed, n)
554575
}
555-
return nil
576+
return n, nil
556577
}
557578

579+
decodedLen := 0
558580
// definition + repetition levels are always uncompressed
559581
if levelsBytelen > 0 {
560-
if _, err := io.ReadFull(rd, buf[:levelsBytelen]); err != nil {
561-
return err
582+
n, err := io.ReadFull(rd, buf[:levelsBytelen])
583+
if err != nil {
584+
return n, err
585+
}
586+
decodedLen += n
587+
if n != levelsBytelen {
588+
return n, fmt.Errorf("parquet: expected to read %d bytes for levels, got %d", levelsBytelen, n)
562589
}
563590
}
564-
if _, err := p.decompress(p.r, lenCompressed-levelsBytelen, buf[levelsBytelen:]); err != nil {
565-
return err
591+
values, err := p.decompress(p.r, lenCompressed-levelsBytelen, buf[levelsBytelen:])
592+
if err != nil {
593+
return decodedLen, err
566594
}
567-
return nil
595+
decodedLen += len(values)
596+
if decodedLen != lenUncompressed {
597+
return decodedLen, fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, decodedLen)
598+
}
599+
return decodedLen, nil
568600
}
569601

570602
type dataheader interface {
@@ -864,20 +896,25 @@ func (p *serializedPageReader) Next() bool {
864896
}
865897

866898
if p.cryptoCtx.DataDecryptor != nil {
867-
if err := p.readV2Encrypted(p.r, lenCompressed, levelsBytelen, compressed, buf.Bytes()); err != nil {
899+
got, err := p.readV2Encrypted(p.r, lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes())
900+
if err != nil {
868901
p.err = err
869902
return false
870903
}
904+
if got != lenUncompressed {
905+
p.err = fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, got)
906+
return false
907+
}
871908
} else {
872-
if err := p.readV2Unencrypted(p.r, lenCompressed, levelsBytelen, compressed, buf.Bytes()); err != nil {
909+
got, err := p.readV2Unencrypted(p.r, lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes())
910+
if err != nil {
873911
p.err = err
874912
return false
875913
}
876-
}
877-
878-
if buf.Len() != lenUncompressed {
879-
p.err = fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, buf.Len())
880-
return false
914+
if got != lenUncompressed {
915+
p.err = fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, got)
916+
return false
917+
}
881918
}
882919

883920
// make datapage v2

0 commit comments

Comments
 (0)