AVRO-4325: [Trevni] Validate column-file header counts and lengths before allocating

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];

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);
 

lang/java/trevni/core/src/main/java/org/apache/trevni/ColumnValues.java

@@ -90,8 +90,12 @@ private void startBlock(int block) throws IOException {
     this.row = column.firstRows[block];
 
     in.seek(column.blockStarts[block]);
-    int end = column.blocks[block].compressedSize;
-    byte[] raw = new byte[end + checksum.size()];
+    // Validate the declared compressed size against the bytes remaining (after
+    // the seek) before allocating, and guard the addition of the checksum size
+    // against integer overflow, so a malformed, corrupted, or truncated file
+    // cannot drive an oversized or negative allocation.
+    int end = in.checkLength(column.blocks[block].compressedSize, 1);
+    byte[] raw = new byte[Math.addExact(end, checksum.size())];
     in.readFully(raw);
     ByteBuffer data = codec.decompress(ByteBuffer.wrap(raw, 0, end));
     if (!checksum.compute(data).equals(ByteBuffer.wrap(raw, end, checksum.size())))

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 extends Comparable> 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;

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 = 12; // MAGIC(4) + rowCount fixed64(8)
+
+  /**
+   * 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 {