fix(parquet): validate DataPageV2 decoded length

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

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