diff --git a/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnDescriptor.java b/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnDescriptor.java index 5d4dad3897d..a8b08fe15a6 100644 --- a/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnDescriptor.java +++ b/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnDescriptor.java @@ -69,7 +69,11 @@ public void ensureBlocksRead() throws IOException { // read block descriptors InputBuffer in = new InputBuffer(file, start); - int blockCount = in.readFixed32(); + // Each block descriptor occupies at least one byte on the wire, so a block + // count larger than the bytes remaining cannot be satisfied. Validate before + // allocating to avoid an oversized allocation from a malformed, corrupted, or + // truncated file. + int blockCount = in.checkLength(in.readFixed32(), 1); BlockDescriptor[] blocks = new BlockDescriptor[blockCount]; if (metaData.hasIndexValues()) firstValues = (T[]) new Comparable[blockCount]; diff --git a/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnFileReader.java b/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnFileReader.java index 1ae0f73232e..ad6e15c6611 100644 --- a/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnFileReader.java +++ b/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnFileReader.java @@ -100,7 +100,11 @@ private void readHeader() throws IOException { InputBuffer in = new InputBuffer(file, 0); readMagic(in); this.rowCount = in.readFixed64(); - this.columnCount = in.readFixed32(); + // Each column contributes at least one byte of metadata and column-start + // data that follows, so a column count larger than the bytes remaining + // cannot be satisfied. Validate before allocating to avoid an oversized + // allocation from a malformed, corrupted, or truncated file. + this.columnCount = in.checkLength(in.readFixed32(), 1); this.metaData = ColumnFileMetaData.read(in); this.columnsByName = new HashMap<>(columnCount); diff --git a/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnValues.java b/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnValues.java index be26783135f..a5488b9260f 100644 --- a/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnValues.java +++ b/lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnValues.java @@ -90,8 +90,24 @@ private void startBlock(int block) throws IOException { this.row = column.firstRows[block]; in.seek(column.blockStarts[block]); + // The block on disk is the compressed payload followed by the checksum + // bytes. Validate the combined length against the bytes remaining before + // allocating, computing in long to avoid integer overflow, so a malformed, + // corrupted, or truncated file fails fast with an IOException rather than an + // oversized/negative allocation or an unchecked ArithmeticException. + int checksumSize = checksum.size(); int end = column.blocks[block].compressedSize; - byte[] raw = new byte[end + checksum.size()]; + if (end < 0) + throw new IOException("Invalid negative block size: " + end); + if (end > Integer.MAX_VALUE - checksumSize) + throw new IOException( + "Block size " + end + " plus checksum size " + checksumSize + " exceeds the maximum " + "array size"); + int rawLength = end + checksumSize; + long remaining = in.remaining(); + if (remaining >= 0 && rawLength > remaining) + throw new IOException("Block size " + end + " plus checksum size " + checksumSize + " exceeds the " + remaining + + " bytes remaining in the input. The file is likely corrupted or truncated."); + byte[] raw = new byte[rawLength]; in.readFully(raw); ByteBuffer data = codec.decompress(ByteBuffer.wrap(raw, 0, end)); if (!checksum.compute(data).equals(ByteBuffer.wrap(raw, end, checksum.size()))) diff --git a/lang/java/trevni/core/src/main/java/org/apache/trevni/InputBuffer.java b/lang/java/trevni/core/src/main/java/org/apache/trevni/InputBuffer.java index 526bb46bc24..1a21ff01b8e 100644 --- a/lang/java/trevni/core/src/main/java/org/apache/trevni/InputBuffer.java +++ b/lang/java/trevni/core/src/main/java/org/apache/trevni/InputBuffer.java @@ -80,6 +80,37 @@ public long length() { return inLength; } + /** The number of bytes remaining to be read from the underlying input. */ + public long remaining() { + return inLength - tell(); + } + + /** + * Validate a length or item count read from the input before it is used to size + * an allocation. Rejects a negative value, and - when the number of bytes + * remaining in the input is known - a value that could not possibly be backed + * by the data that follows, assuming each counted element occupies at least + * {@code minBytesPerElement} bytes on the wire. This guards against a + * malformed, corrupted, or truncated file driving an oversized (or negative) + * allocation. + * + * @param count the length or item count read from the input + * @param minBytesPerElement the minimum number of input bytes each counted + * element occupies (use 1 for a raw byte length) + * @return {@code count}, if it is valid + * @throws IOException if {@code count} is negative or larger than the input can + * support + */ + public int checkLength(int count, long minBytesPerElement) throws IOException { + if (count < 0) + throw new IOException("Invalid negative length: " + count); + long remaining = remaining(); + if (remaining >= 0 && minBytesPerElement > 0 && count > remaining / minBytesPerElement) + throw new IOException("Length " + count + " exceeds the " + remaining + + " bytes remaining in the input. The file is likely corrupted or truncated."); + return count; + } + public T readValue(ValueType type) throws IOException { switch (type) { case NULL: @@ -306,7 +337,7 @@ public long readFixed64() throws IOException { } public String readString() throws IOException { - int length = readInt(); + int length = checkLength(readInt(), 1); if (length <= (limit - pos)) { // in buffer String result = utf8.decode(ByteBuffer.wrap(buf, pos, length)).toString(); pos += length; @@ -318,13 +349,13 @@ public String readString() throws IOException { } public byte[] readBytes() throws IOException { - byte[] result = new byte[readInt()]; + byte[] result = new byte[checkLength(readInt(), 1)]; readFully(result); return result; } public ByteBuffer readBytes(ByteBuffer old) throws IOException { - int length = readInt(); + int length = checkLength(readInt(), 1); ByteBuffer result; if (old != null && length <= old.capacity()) { result = old; diff --git a/lang/java/trevni/core/src/test/java/org/apache/trevni/TestColumnFile.java b/lang/java/trevni/core/src/test/java/org/apache/trevni/TestColumnFile.java index 781476abfc5..cd8dfa51e75 100644 --- a/lang/java/trevni/core/src/test/java/org/apache/trevni/TestColumnFile.java +++ b/lang/java/trevni/core/src/test/java/org/apache/trevni/TestColumnFile.java @@ -18,6 +18,8 @@ package org.apache.trevni; import java.io.File; +import java.io.IOException; +import java.io.RandomAccessFile; import java.util.Random; import java.util.Arrays; import java.util.Iterator; @@ -26,6 +28,7 @@ import java.util.stream.Stream; import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; @@ -58,6 +61,33 @@ void emptyFile(ColumnFileMetaData fileMeta) throws Exception { in.close(); } + /** Byte offset of the little-endian 4-byte columnCount field in the header. */ + private static final int COLUMN_COUNT_OFFSET = ColumnFileWriter.MAGIC.length + Long.BYTES; // MAGIC + rowCount fixed64 + + /** + * A header column count larger than the data present (from a malformed, + * corrupted, or truncated file) must be rejected before allocating, rather than + * attempting an oversized allocation or failing later. + */ + @Test + void oversizedColumnCountIsRejected() throws Exception { + FILE.delete(); + // A valid, minimal file (no columns) written by Trevni's own writer. + new ColumnFileWriter(new ColumnFileMetaData()).writeTo(FILE); + + // Overwrite the columnCount field with Integer.MAX_VALUE. + try (RandomAccessFile raf = new RandomAccessFile(FILE, "rw")) { + raf.seek(COLUMN_COUNT_OFFSET); + raf.write(0xFF); + raf.write(0xFF); + raf.write(0xFF); + raf.write(0x7F); + } + + IOException e = Assertions.assertThrows(IOException.class, () -> new ColumnFileReader(FILE).close()); + Assertions.assertNotNull(e.getMessage()); + } + @ParameterizedTest @MethodSource("codecs") void emptyColumn(ColumnFileMetaData fileMeta) throws Exception {