diff --git a/lang/java/avro/src/main/java/org/apache/avro/file/DataFileStream.java b/lang/java/avro/src/main/java/org/apache/avro/file/DataFileStream.java index 9bb15183ec7..d2051827d93 100644 --- a/lang/java/avro/src/main/java/org/apache/avro/file/DataFileStream.java +++ b/lang/java/avro/src/main/java/org/apache/avro/file/DataFileStream.java @@ -320,6 +320,18 @@ boolean hasNextBlock() { if (blockSize > Integer.MAX_VALUE || blockSize < 0) { throw new IOException("Block size invalid or too large for this implementation: " + blockSize); } + // When the number of bytes remaining in the input is known (e.g. a + // byte-array- or known-length-stream-backed decoder), reject a declared + // block size that could not possibly be satisfied by the data available. + // This avoids eagerly allocating a large block buffer (see the DataBlock + // constructor) for a malformed, corrupted, or truncated file before any + // block byte has been read. A value of -1 means the remaining count is + // unknown, in which case the check is skipped. + int remaining = vin.remainingBytes(); + if (remaining >= 0 && blockSize > remaining) { + throw new IOException("Block size " + blockSize + " exceeds the number of bytes remaining in the input (" + + remaining + "). The file is likely corrupted or truncated."); + } blockCount = blockRemaining; availableBlock = true; return true; diff --git a/lang/java/avro/src/test/java/org/apache/avro/TestDataFileReader.java b/lang/java/avro/src/test/java/org/apache/avro/TestDataFileReader.java index 0e13c8388c0..85cf405bf78 100644 --- a/lang/java/avro/src/test/java/org/apache/avro/TestDataFileReader.java +++ b/lang/java/avro/src/test/java/org/apache/avro/TestDataFileReader.java @@ -319,4 +319,63 @@ private static byte[] buildVersion12ContainerWithoutSchema() throws IOException return output.toByteArray(); } + + /** + * A block header may declare a block size much larger than the data actually + * present in a corrupted or truncated file. When the remaining byte count is + * known, the reader must reject such a block up front rather than attempting to + * allocate a buffer of the declared size. + */ + @Test + void oversizedBlockSizeIsRejectedBeforeAllocation() throws IOException { + Schema schema = new Schema.Parser().parse("{\"type\":\"int\"}"); + + // A spec-correct header with zero records, so no real data block is written. + ByteArrayOutputStream fileBytes = new ByteArrayOutputStream(); + try (DataFileWriter w = new DataFileWriter<>(new GenericDatumWriter<>(schema))) { + w.create(schema, fileBytes); + } + + // Append a single block header that declares a huge block size but no data. + BinaryEncoder encoder = EncoderFactory.get().binaryEncoder(fileBytes, null); + encoder.writeLong(1L); // block entry count + encoder.writeLong(2_000_000_000L); // block size in bytes, far larger than what follows + encoder.flush(); + + byte[] malformed = fileBytes.toByteArray(); + + // hasNextBlock() surfaces block-header IOExceptions wrapped in an + // AvroRuntimeException, the same way the existing "block size too large" + // check does. + AvroRuntimeException exception = assertThrows(AvroRuntimeException.class, () -> { + DataFileStream reader = new DataFileStream<>(new ByteArrayInputStream(malformed), + new GenericDatumReader<>()); + while (reader.hasNext()) { + reader.next(); + } + }); + assertNotNull(exception.getMessage()); + assertTrue(exception.getMessage().contains("Block size"), "Unexpected message: " + exception.getMessage()); + } + + /** + * A valid single-record file must still read normally after the guard is added. + */ + @Test + void validFileWithSingleRecordStillReads() throws IOException { + Schema schema = new Schema.Parser().parse("{\"type\":\"int\"}"); + + ByteArrayOutputStream fileBytes = new ByteArrayOutputStream(); + try (DataFileWriter w = new DataFileWriter<>(new GenericDatumWriter<>(schema))) { + w.create(schema, fileBytes); + w.append(42); + } + + try (DataFileStream reader = new DataFileStream<>(new ByteArrayInputStream(fileBytes.toByteArray()), + new GenericDatumReader<>())) { + assertTrue(reader.hasNext()); + assertEquals(42, reader.next()); + assertFalse(reader.hasNext()); + } + } }