diff --git a/parquet/file/file_reader_test.go b/parquet/file/file_reader_test.go index c764ff07..5ae9ae74 100644 --- a/parquet/file/file_reader_test.go +++ b/parquet/file/file_reader_test.go @@ -215,6 +215,108 @@ func (p *PageSerdeSuite) TestDataPageV2() { p.alloc.AssertSize(p.T(), 0) } +func (p *PageSerdeSuite) TestDataPageV2CorruptDecodedLength() { + rawData := bytes.Repeat([]byte{0x42}, 128) + codec, err := compress.GetCodec(compress.Codecs.Snappy) + p.NoError(err) + + buffer := make([]byte, int(codec.CompressBound(int64(len(rawData))))) + compressed := codec.Encode(buffer, rawData) + + p.dataPageHdrV2 = *format.NewDataPageHeaderV2() + p.dataPageHdrV2.NumValues = 1 + p.dataPageHdrV2.NumRows = 1 + p.dataPageHdrV2.Encoding = format.Encoding_PLAIN + + p.WriteDataPageHeaderV2(1024, int32(len(rawData)+1), int32(len(compressed))) + _, err = p.sink.Write(compressed) + p.NoError(err) + + p.InitSerializedPageReader(1, compress.Codecs.Snappy) + p.False(p.pageReader.Next()) + p.ErrorContains(p.pageReader.Err(), "metadata said 129 bytes uncompressed data page, got 128 bytes") + p.NoError(p.pageReader.Close()) + p.alloc.AssertSize(p.T(), 0) +} + +func (p *PageSerdeSuite) TestDataPageV2Compressed() { + rawData := bytes.Repeat([]byte{0x42}, 128) + codec, err := compress.GetCodec(compress.Codecs.Snappy) + p.NoError(err) + + buffer := make([]byte, int(codec.CompressBound(int64(len(rawData))))) + compressed := codec.Encode(buffer, rawData) + + p.dataPageHdrV2 = *format.NewDataPageHeaderV2() + p.dataPageHdrV2.NumValues = 1 + p.dataPageHdrV2.NumRows = 1 + p.dataPageHdrV2.Encoding = format.Encoding_PLAIN + + p.WriteDataPageHeaderV2(1024, int32(len(rawData)), int32(len(compressed))) + _, err = p.sink.Write(compressed) + p.NoError(err) + + p.InitSerializedPageReader(1, compress.Codecs.Snappy) + p.True(p.pageReader.Next()) + p.Equal(rawData, p.pageReader.Page().Data()) + p.NoError(p.pageReader.Close()) + p.alloc.AssertSize(p.T(), 0) +} + +type mockDecryptor struct { + decrypt func([]byte) []byte + + aad []byte + mem memory.Allocator +} + +func (m *mockDecryptor) FileAad() string { return "" } + +func (m *mockDecryptor) Allocator() memory.Allocator { + if m.mem == nil { + return memory.DefaultAllocator + } + + return m.mem +} + +func (m *mockDecryptor) CiphertextSizeDelta() int { return 0 } + +func (m *mockDecryptor) Decrypt(src []byte) []byte { return m.decrypt(src) } + +func (m *mockDecryptor) DecryptFrom(r io.Reader) []byte { + decrypted, _ := io.ReadAll(r) + return m.Decrypt(decrypted) +} + +func (m *mockDecryptor) UpdateAad(aad string) { m.aad = []byte(aad) } + +func (p *PageSerdeSuite) TestDataPageV2EncryptedUncompressedLengthMismatch() { + p.dataPageHdrV2 = *format.NewDataPageHeaderV2() + p.dataPageHdrV2.NumValues = 1 + p.dataPageHdrV2.NumRows = 1 + p.dataPageHdrV2.Encoding = format.Encoding_PLAIN + p.dataPageHdrV2.IsCompressed = false + + payload := bytes.Repeat([]byte{0xaa}, 4) + decryptedPayload := []byte{0xaa} + + p.WriteDataPageHeaderV2(1024, int32(len(payload)), int32(len(payload))) + _, err := p.sink.Write(payload) + p.NoError(err) + + p.EndStream() + p.alloc = memory.NewCheckedAllocator(memory.NewGoAllocator()) + p.pageReader, _ = file.NewPageReader(utils.NewByteReader(p.buffer.Bytes()), 1, compress.Codecs.Uncompressed, p.alloc, &file.CryptoContext{ + DataDecryptor: &mockDecryptor{decrypt: func([]byte) []byte { return decryptedPayload }}, + }) + + p.False(p.pageReader.Next()) + p.ErrorContains(p.pageReader.Err(), "metadata said 4 bytes uncompressed data page, got 1 bytes") + p.NoError(p.pageReader.Close()) + p.alloc.AssertSize(p.T(), 0) +} + func (p *PageSerdeSuite) TestLargePageHeaders() { const ( statsSize = 256 * 1024 // 256KB diff --git a/parquet/file/page_reader.go b/parquet/file/page_reader.go index 8c6c6ace..66b8473a 100644 --- a/parquet/file/page_reader.go +++ b/parquet/file/page_reader.go @@ -518,53 +518,85 @@ func (p *serializedPageReader) decompress(rd io.Reader, lenCompressed int, buf [ return p.codec.Decode(buf, data), nil } -func (p *serializedPageReader) readV2Encrypted(rd io.Reader, lenCompressed int, levelsBytelen int, compressed bool, buf []byte) error { +func (p *serializedPageReader) readV2Encrypted(rd io.Reader, lenCompressed int, lenUncompressed int, levelsBytelen int, compressed bool, buf []byte) (int, error) { // if encrypted, we need to decrypt before decompressing p.decompressBuffer.ResizeNoShrink(lenCompressed) n, err := io.ReadFull(rd, p.decompressBuffer.Bytes()[:lenCompressed]) if err != nil { - return err + return n, err } if n != lenCompressed { - return fmt.Errorf("parquet: expected to read %d compressed bytes, got %d", lenCompressed, n) + return n, fmt.Errorf("parquet: expected to read %d compressed bytes, got %d", lenCompressed, n) } data := p.cryptoCtx.DataDecryptor.Decrypt(p.decompressBuffer.Bytes()[:lenCompressed]) + // encrypted + uncompressed -> just copy the decrypted data to output buffer if !compressed { - copy(buf, data) - return nil + if len(data) != lenUncompressed { + return n, fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, len(data)) + } + + copied := copy(buf, data) + if copied != lenUncompressed { + return copied, fmt.Errorf("parquet: expected to read %d uncompressed encrypted bytes, got %d", lenUncompressed, copied) + } + return copied, nil } // definition + repetition levels are always uncompressed if levelsBytelen > 0 { - copy(buf, data[:levelsBytelen]) + if len(data) < levelsBytelen { + return n, fmt.Errorf("parquet: expected to read %d bytes for levels, got %d", levelsBytelen, len(data)) + } + copied := copy(buf, data[:levelsBytelen]) + if copied != levelsBytelen { + return copied, fmt.Errorf("parquet: expected to read %d bytes for levels, got %d", levelsBytelen, copied) + } data = data[levelsBytelen:] } - p.codec.Decode(buf[levelsBytelen:], data) - return nil + decoded := p.codec.Decode(buf[levelsBytelen:], data) + if len(decoded) != lenUncompressed-levelsBytelen { + return levelsBytelen + len(decoded), fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, levelsBytelen+len(decoded)) + } + return levelsBytelen + len(decoded), nil } -func (p *serializedPageReader) readV2Unencrypted(rd io.Reader, lenCompressed int, levelsBytelen int, compressed bool, buf []byte) error { +func (p *serializedPageReader) readV2Unencrypted(rd io.Reader, lenCompressed int, lenUncompressed int, levelsBytelen int, compressed bool, buf []byte) (int, error) { if !compressed { // uncompressed, just read into the buffer - if _, err := io.ReadFull(rd, buf); err != nil { - return err + n, err := io.ReadFull(rd, buf) + if err != nil { + return n, err + } + if n != lenUncompressed { + return n, fmt.Errorf("parquet: expected to read %d uncompressed bytes, got %d", lenUncompressed, n) } - return nil + return n, nil } + decodedLen := 0 // definition + repetition levels are always uncompressed if levelsBytelen > 0 { - if _, err := io.ReadFull(rd, buf[:levelsBytelen]); err != nil { - return err + n, err := io.ReadFull(rd, buf[:levelsBytelen]) + if err != nil { + return n, err + } + decodedLen += n + if n != levelsBytelen { + return n, fmt.Errorf("parquet: expected to read %d bytes for levels, got %d", levelsBytelen, n) } } - if _, err := p.decompress(p.r, lenCompressed-levelsBytelen, buf[levelsBytelen:]); err != nil { - return err + values, err := p.decompress(p.r, lenCompressed-levelsBytelen, buf[levelsBytelen:]) + if err != nil { + return decodedLen, err } - return nil + decodedLen += len(values) + if decodedLen != lenUncompressed { + return decodedLen, fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, decodedLen) + } + return decodedLen, nil } type dataheader interface { @@ -864,20 +896,25 @@ func (p *serializedPageReader) Next() bool { } if p.cryptoCtx.DataDecryptor != nil { - if err := p.readV2Encrypted(p.r, lenCompressed, levelsBytelen, compressed, buf.Bytes()); err != nil { + got, err := p.readV2Encrypted(p.r, lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes()) + if err != nil { p.err = err return false } + if got != lenUncompressed { + p.err = fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, got) + return false + } } else { - if err := p.readV2Unencrypted(p.r, lenCompressed, levelsBytelen, compressed, buf.Bytes()); err != nil { + got, err := p.readV2Unencrypted(p.r, lenCompressed, lenUncompressed, levelsBytelen, compressed, buf.Bytes()) + if err != nil { p.err = err return false } - } - - if buf.Len() != lenUncompressed { - p.err = fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, buf.Len()) - return false + if got != lenUncompressed { + p.err = fmt.Errorf("parquet: metadata said %d bytes uncompressed data page, got %d bytes", lenUncompressed, got) + return false + } } // make datapage v2