Skip to content

Commit 4e61adc

Browse files
sfc-gh-hkaraucursoragentholdenk
committed
[MINOR][SQL] Bounds-check RLE/bit-packed group counts in the vectorized Parquet reader
Group counts read from a Parquet page body were trusted blindly in VectorizedRleValuesReader.readNextGroup(). A corrupt page could make `numGroups * 8` overflow int to a negative currentCount (sending the consumers' leftInPage accounting backwards) or size a multi-GB int[] allocation the page has no bytes for, and a truncated varint header looped forever because read() keeps returning -1 at end of input, whose continuation bit is always set. Remember the page's declared value count from initFromPage and validate each group header against it (allowing 7 values of bit-packing padding) and against the remaining byte budget before the count sizes an allocation or drives loop control. Corrupt pages now fail fast with ParquetDecodingException. Co-authored-by: Cursor <[email protected]> Co-Authored-By: Holden Karau <[email protected]>
1 parent bb76812 commit 4e61adc

2 files changed

Lines changed: 151 additions & 1 deletion

File tree

‎sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,10 @@ private enum MODE {
5656
// Encoded data.
5757
private ByteBufferInputStream in;
5858

59+
// Total number of values in the current page, as declared by the page header. Group
60+
// counts read from the page body are validated against this bound.
61+
private int pageValueCount;
62+
5963
// bit/byte width of decoded data and utility to batch unpack them.
6064
private int bitWidth;
6165
private int bytesWidth;
@@ -95,6 +99,7 @@ public VectorizedRleValuesReader(int bitWidth, boolean readLength) {
9599
@Override
96100
public void initFromPage(int valueCount, ByteBufferInputStream in) throws IOException {
97101
this.in = in;
102+
this.pageValueCount = valueCount;
98103
if (fixedWidth) {
99104
// Initialize for repetition and definition levels
100105
if (readLength) {
@@ -860,6 +865,11 @@ private int readUnsignedVarInt() throws IOException {
860865
int b;
861866
do {
862867
b = in.read();
868+
if (b < 0) {
869+
// A truncated varint would otherwise spin forever: read() keeps returning -1 at
870+
// end of input, whose continuation bit is always set.
871+
throw new ParquetDecodingException("Reached end of page while reading a varint header");
872+
}
863873
value |= (b & 0x7F) << shift;
864874
shift += 7;
865875
} while ((b & 0x80) != 0);
@@ -919,11 +929,30 @@ private boolean readNextGroup() {
919929
switch (mode) {
920930
case RLE:
921931
this.currentCount = header >>> 1;
932+
if (this.currentCount > pageValueCount) {
933+
throw new ParquetDecodingException("Invalid RLE run of " + this.currentCount +
934+
" values in a page declaring " + pageValueCount);
935+
}
922936
this.currentValue = readIntLittleEndianPaddedOnBitWidth();
923937
break;
924938
case PACKED:
925939
int numGroups = header >>> 1;
926-
this.currentCount = numGroups * 8;
940+
// The group count comes from the page body and may be garbage if the file is
941+
// corrupt: `numGroups * 8` can overflow int, and a huge count would size a
942+
// multi-GB allocation the page has no bytes for. The last bit-packed group may
943+
// be zero-padded to a multiple of 8, so allow count up to pageValueCount + 7.
944+
long count = (long) numGroups * 8;
945+
long bytesNeeded = (long) numGroups * bitWidth;
946+
int bytesAvailable = in.available();
947+
// The count check alone still admits count = 2^31 when pageValueCount is within
948+
// 7 of Integer.MAX_VALUE, which would overflow the int cast below.
949+
if (count > (long) pageValueCount + 7 || count > Integer.MAX_VALUE ||
950+
bytesNeeded > bytesAvailable) {
951+
throw new ParquetDecodingException("Invalid bit-packed run: " + numGroups +
952+
" groups in a page declaring " + pageValueCount + " values with " +
953+
bytesAvailable + " bytes remaining");
954+
}
955+
this.currentCount = (int) count;
927956

928957
if (this.currentBuffer.length < this.currentCount) {
929958
this.currentBuffer = new int[this.currentCount];

‎sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.spark.sql.execution.datasources.parquet
19+
20+
import java.io.ByteArrayOutputStream
21+
import java.nio.ByteBuffer
22+
23+
import org.apache.parquet.bytes.ByteBufferInputStream
24+
import org.apache.parquet.io.ParquetDecodingException
25+
26+
import org.apache.spark.SparkFunSuite
27+
28+
/**
29+
* Group headers in a Parquet page body are attacker-controlled if the file is, so
30+
* `readNextGroup` validates them against the page's declared value count and remaining
31+
* byte budget before a count can size an allocation or drive loop control.
32+
*/
33+
class VectorizedRleValuesReaderSuite extends SparkFunSuite {
34+
35+
import VectorizedRleValuesReaderSuite._
36+
37+
test("invalid PACKED header: overflowing value count is rejected") {
38+
// numGroups = 2^28 overflows `numGroups * 8` to a negative Int currentCount, which
39+
// would send readBatch's leftInPage accounting backwards. Assert on the message:
40+
// without the validation this also throws, but as a wrapped end-of-stream failure.
41+
val e = intercept[ParquetDecodingException] {
42+
packedPageReader(numGroups = 1 << 28, pageValueCount = 1024).readInteger()
43+
}
44+
assert(e.getMessage.contains("Invalid bit-packed run"))
45+
}
46+
47+
test("invalid PACKED header: count beyond the page byte budget is rejected pre-allocation") {
48+
// numGroups = 2^27 would allocate a ~4GB int buffer for a page with only a handful
49+
// of bytes left.
50+
val e = intercept[ParquetDecodingException] {
51+
packedPageReader(numGroups = 1 << 27, pageValueCount = 1024).readInteger()
52+
}
53+
assert(e.getMessage.contains("Invalid bit-packed run"))
54+
}
55+
56+
test("invalid RLE header: run longer than the page's value count is rejected") {
57+
val out = new ByteArrayOutputStream()
58+
writeUnsignedVarInt(out, 1L << 21) // RLE header (LSB 0): run of 2^20 values
59+
out.write(1) // the repeated value (bitWidth 1 -> one byte)
60+
val reader = new VectorizedRleValuesReader(1, false)
61+
reader.initFromPage(1024, ByteBufferInputStream.wrap(ByteBuffer.wrap(out.toByteArray)))
62+
val e = intercept[ParquetDecodingException] {
63+
reader.readInteger()
64+
}
65+
assert(e.getMessage.contains("Invalid RLE run"))
66+
}
67+
68+
test("truncated varint header is rejected") {
69+
// A lone continuation byte: read() returns -1 at end of input, whose set 0x80 bit
70+
// would otherwise keep the varint continuation loop alive.
71+
val out = new ByteArrayOutputStream()
72+
out.write(0x80)
73+
val reader = new VectorizedRleValuesReader(1, false)
74+
reader.initFromPage(1024, ByteBufferInputStream.wrap(ByteBuffer.wrap(out.toByteArray)))
75+
intercept[ParquetDecodingException] {
76+
reader.readInteger()
77+
}
78+
}
79+
80+
test("PACKED: value count not a multiple of 8 is accepted (final group is zero-padded)") {
81+
// The last bit-packed group is padded to a multiple of 8 values, so a page of 1023
82+
// values legitimately carries a run counting 1024 -- the validation's +7 allowance.
83+
// Guards against the bound being too tight to read valid files.
84+
val numGroups = 128
85+
val out = new ByteArrayOutputStream()
86+
writeUnsignedVarInt(out, (numGroups.toLong << 1) | 1L)
87+
out.write(Array.fill[Byte](numGroups)(0x55)) // bitWidth 1 -> one byte per group
88+
val reader = new VectorizedRleValuesReader(1, false)
89+
reader.initFromPage(1023, ByteBufferInputStream.wrap(ByteBuffer.wrap(out.toByteArray)))
90+
reader.readInteger()
91+
}
92+
}
93+
94+
private object VectorizedRleValuesReaderSuite {
95+
96+
/**
97+
* Builds a reader over a page whose body starts with a PACKED header (LSB 1) declaring
98+
* `numGroups` bit-packed groups of 8 values each, followed by far fewer bytes than the
99+
* numGroups * bitWidth the run requires.
100+
*/
101+
private def packedPageReader(
102+
numGroups: Int,
103+
pageValueCount: Int): VectorizedRleValuesReader = {
104+
val out = new ByteArrayOutputStream()
105+
writeUnsignedVarInt(out, (numGroups.toLong << 1) | 1L)
106+
out.write(Array[Byte](0, 0, 0, 0))
107+
val reader = new VectorizedRleValuesReader(1, false)
108+
reader.initFromPage(
109+
pageValueCount, ByteBufferInputStream.wrap(ByteBuffer.wrap(out.toByteArray)))
110+
reader
111+
}
112+
113+
private def writeUnsignedVarInt(out: ByteArrayOutputStream, value: Long): Unit = {
114+
var v = value
115+
while ((v & ~0x7fL) != 0) {
116+
out.write(((v & 0x7f) | 0x80).toInt)
117+
v >>>= 7
118+
}
119+
out.write(v.toInt)
120+
}
121+
}

0 commit comments

Comments
 (0)