Bound decompressed record size (#23429)

checkstyle/suppressions.xml

@@ -362,9 +362,11 @@
               files="(LocalLog|LogLoader|LogValidator|RemoteLogManager|RemoteIndexCache|UnifiedLog).java"/>
     <suppress checks="ParameterNumber"
               files="(LogAppendInfo|LogLoader|RemoteLogManagerConfig|UnifiedLog).java"/>
+    <suppress checks="ParameterNumber"
+              files="[/\\]Cleaner\.java"/>
     <suppress checks="(ClassDataAbstractionCoupling|ClassFanOutComplexity)"
               files="(UnifiedLog|RemoteLogManager|RemoteLogManagerTest).java"/>
-    <suppress checks="MethodLength" files="(RemoteLogManager|RemoteLogManagerConfig).java"/>
+    <suppress checks="MethodLength" files="(RemoteLogManager|RemoteLogManagerConfig|UnifiedLog).java"/>
     <suppress checks="JavaNCSS" files="RemoteLogManagerTest.java"/>
 
     <!-- benchmarks -->

clients/src/main/java/org/apache/kafka/common/config/TopicConfig.java

@@ -110,6 +110,18 @@ public class TopicConfig {
     public static final String MAX_MESSAGE_BYTES_DOC =
         "The largest record batch size allowed by Kafka (after compression if compression is enabled).";
 
+    public static final String MAX_DECOMPRESSED_MESSAGE_BYTES_CONFIG = "max.decompressed.message.bytes";
+    public static final String MAX_DECOMPRESSED_MESSAGE_BYTES_DOC = "The maximum decompressed size, in bytes, of a " +
+        "single record the broker will accept; larger records are rejected with an " +
+        "<code>InvalidRecordException</code>. It is enforced where the broker decompresses records: produce " +
+        "validation of compressed batches, log compaction, and offset lookups by timestamp or max timestamp against " +
+        "the local log or remote storage. " +
+        "Unlike <code>max.message.bytes</code>, which bounds the compressed batch on the wire, this bounds the " +
+        "decompressed size of an individual record. Lowering it below the size of already-stored records makes " +
+        "compaction and timestamp offset lookups of affected partitions fail until it is raised again. The " +
+        "default imposes no limit beyond the largest record the broker could allocate anyway; set a smaller value " +
+        "to enforce a per-record cap.";
+
     public static final String INDEX_INTERVAL_BYTES_CONFIG = "index.interval.bytes";
     public static final String INDEX_INTERVAL_BYTES_DOC = "This setting controls how frequently Kafka " +
             "adds entries to its offset index and, conditionally, to its time index. " +

clients/src/main/java/org/apache/kafka/common/record/AbstractLegacyRecordBatch.java

@@ -232,8 +232,12 @@ public Iterator<Record> iterator() {
     }
 
     CloseableIterator<Record> iterator(BufferSupplier bufferSupplier) {
+        return iterator(bufferSupplier, Records.SOFT_MAX_ARRAY_LENGTH);
+    }
+
+    CloseableIterator<Record> iterator(BufferSupplier bufferSupplier, int maxRecordBodySize) {
         if (isCompressed())
-            return new DeepRecordsIterator(this, false, Integer.MAX_VALUE, bufferSupplier);
+            return new DeepRecordsIterator(this, false, Integer.MAX_VALUE, bufferSupplier, maxRecordBodySize);
 
         return new CloseableIterator<>() {
             private boolean hasNext = true;
@@ -267,6 +271,12 @@ public CloseableIterator<Record> streamingIterator(BufferSupplier bufferSupplier
         return iterator(bufferSupplier);
     }
 
+    @Override
+    public CloseableIterator<Record> streamingIterator(BufferSupplier bufferSupplier, int maxRecordBodySize) {
+        // the older message format versions do not support streaming, so we return the normal iterator
+        return iterator(bufferSupplier, maxRecordBodySize);
+    }
+
     static void writeHeader(ByteBuffer buffer, long offset, int size) {
         buffer.putLong(offset);
         buffer.putInt(size);
@@ -280,11 +290,13 @@ static void writeHeader(DataOutputStream out, long offset, int size) throws IOEx
     private static final class DataLogInputStream implements LogInputStream<AbstractLegacyRecordBatch> {
         private final InputStream stream;
         private final int maxMessageSize;
+        private final int maxRecordBodySize;
         private final ByteBuffer offsetAndSizeBuffer;
 
-        DataLogInputStream(InputStream stream, int maxMessageSize) {
+        DataLogInputStream(InputStream stream, int maxMessageSize, int maxRecordBodySize) {
             this.stream = stream;
             this.maxMessageSize = maxMessageSize;
+            this.maxRecordBodySize = maxRecordBodySize;
             this.offsetAndSizeBuffer = ByteBuffer.allocate(Records.LOG_OVERHEAD);
         }
 
@@ -300,6 +312,15 @@ public AbstractLegacyRecordBatch nextBatch() throws IOException {
                 throw new CorruptRecordException(String.format("Record size is less than the minimum record overhead (%d)", LegacyRecord.RECORD_OVERHEAD_V0));
             if (size > maxMessageSize)
                 throw new CorruptRecordException(String.format("Record size exceeds the largest allowable message size (%d).", maxMessageSize));
+            // The maxMessageSize check above does not bound the allocation on the compressed deep-decode
+            // path, where maxMessageSize can be up to Integer.MAX_VALUE. Reject, before allocating, any
+            // inner record whose declared size exceeds the configured per-record maximum. maxRecordBodySize
+            // never exceeds SOFT_MAX_ARRAY_LENGTH (the default here, and the config's upper bound), so this
+            // also stops an adversarial inner size from triggering an OutOfMemoryError (see the equivalent
+            // guard in DefaultRecord for the V2 format).
+            if (size > maxRecordBodySize)
+                throw new InvalidRecordException("Invalid record size: " + size +
+                    " exceeds the configured maximum record size of " + maxRecordBodySize + ".");
 
             ByteBuffer batchBuffer = ByteBuffer.allocate(size);
             Utils.readFully(stream, batchBuffer);
@@ -319,7 +340,8 @@ private static class DeepRecordsIterator extends AbstractIterator<Record> implem
         private DeepRecordsIterator(AbstractLegacyRecordBatch wrapperEntry,
                                     boolean ensureMatchingMagic,
                                     int maxMessageSize,
-                                    BufferSupplier bufferSupplier) {
+                                    BufferSupplier bufferSupplier,
+                                    int maxRecordBodySize) {
             LegacyRecord wrapperRecord = wrapperEntry.outerRecord();
             this.wrapperMagic = wrapperRecord.magic();
             if (wrapperMagic != RecordBatch.MAGIC_VALUE_V0 && wrapperMagic != RecordBatch.MAGIC_VALUE_V1)
@@ -334,7 +356,7 @@ private DeepRecordsIterator(AbstractLegacyRecordBatch wrapperEntry,
                         wrapperMagic + ")");
 
             InputStream stream = Compression.of(compressionType).build().wrapForInput(wrapperValue, wrapperRecord.magic(), bufferSupplier);
-            LogInputStream<AbstractLegacyRecordBatch> logStream = new DataLogInputStream(stream, maxMessageSize);
+            LogInputStream<AbstractLegacyRecordBatch> logStream = new DataLogInputStream(stream, maxMessageSize, maxRecordBodySize);
 
             long lastOffsetFromWrapper = wrapperEntry.lastOffset();
             long timestampFromWrapper = wrapperRecord.timestamp();
@@ -517,6 +539,12 @@ public CloseableIterator<Record> skipKeyValueIterator(BufferSupplier bufferSuppl
             return CloseableIterator.wrap(iterator(bufferSupplier));
         }
 
+        @Override
+        public CloseableIterator<Record> skipKeyValueIterator(BufferSupplier bufferSupplier, int maxRecordBodySize) {
+            // legacy batches cannot cheaply skip the record body, so this is a full decode
+            return CloseableIterator.wrap(iterator(bufferSupplier, maxRecordBodySize));
+        }
+
         @Override
         public void writeTo(ByteBufferOutputStream outputStream) {
             outputStream.write(buffer.duplicate());

clients/src/main/java/org/apache/kafka/common/record/DefaultRecord.java

@@ -276,12 +276,28 @@ public int hashCode() {
         return result;
     }
 
+    /**
+     * Decode a record from the (decompressed) stream, rejecting any record whose declared body size
+     * exceeds {@code maxRecordBodySize} before allocating the body buffer. Callers that do not
+     * enforce a configured limit pass {@link Records#SOFT_MAX_ARRAY_LENGTH}, effectively the
+     * array-length allocation limit.
+     */
     public static DefaultRecord readFrom(InputStream input,
                                          long baseOffset,
                                          long baseTimestamp,
                                          int baseSequence,
-                                         Long logAppendTime) throws IOException {
+                                         Long logAppendTime,
+                                         int maxRecordBodySize) throws IOException {
         int sizeOfBodyInBytes = ByteUtils.readVarint(input);
+        if (sizeOfBodyInBytes < 0)
+            throw new InvalidRecordException("Invalid record size: " + sizeOfBodyInBytes + " is negative.");
+        // Reject, before allocating, any record whose declared (decompressed) body exceeds the
+        // configured per-record maximum. maxRecordBodySize never exceeds SOFT_MAX_ARRAY_LENGTH (the
+        // default here, and the config's upper bound), so this also stops an adversarial size from
+        // reaching ByteBuffer.allocate and raising OutOfMemoryError.
+        if (sizeOfBodyInBytes > maxRecordBodySize)
+            throw new InvalidRecordException("Invalid record size: " + sizeOfBodyInBytes +
+                " exceeds the configured maximum record size of " + maxRecordBodySize + ".");
         ByteBuffer recordBuffer = ByteBuffer.allocate(sizeOfBodyInBytes);
         int bytesRead = Utils.readFully(input, recordBuffer);
         if (bytesRead != sizeOfBodyInBytes)
@@ -358,12 +374,26 @@ private static DefaultRecord readFrom(ByteBuffer buffer,
         }
     }
 
+    /**
+     * Skip-decode a record from the (decompressed) stream, rejecting any record whose declared body
+     * size exceeds {@code maxRecordBodySize}. See
+     * {@link #readFrom(InputStream, long, long, int, Long, int)}.
+     */
     public static PartialDefaultRecord readPartiallyFrom(InputStream input,
                                                          long baseOffset,
                                                          long baseTimestamp,
                                                          int baseSequence,
-                                                         Long logAppendTime) throws IOException {
+                                                         Long logAppendTime,
+                                                         int maxRecordBodySize) throws IOException {
         int sizeOfBodyInBytes = ByteUtils.readVarint(input);
+        if (sizeOfBodyInBytes < 0)
+            throw new InvalidRecordException("Invalid record size: " + sizeOfBodyInBytes + " is negative.");
+        // Reject records whose declared (decompressed) body exceeds the configured per-record
+        // maximum; as in readFrom, this doubles as the array-length allocation guard because
+        // maxRecordBodySize never exceeds SOFT_MAX_ARRAY_LENGTH.
+        if (sizeOfBodyInBytes > maxRecordBodySize)
+            throw new InvalidRecordException("Invalid record size: " + sizeOfBodyInBytes +
+                " exceeds the configured maximum record size of " + maxRecordBodySize + ".");
         int totalSizeInBytes = ByteUtils.sizeOfVarint(sizeOfBodyInBytes) + sizeOfBodyInBytes;
 
         return readPartiallyFrom(input, totalSizeInBytes, baseOffset, baseTimestamp,

clients/src/main/java/org/apache/kafka/common/record/DefaultRecordBatch.java

@@ -275,21 +275,21 @@ public InputStream recordInputStream(BufferSupplier bufferSupplier) {
         return Compression.of(compressionType()).build().wrapForInput(buffer, magic(), bufferSupplier);
     }
 
-    private CloseableIterator<Record> compressedIterator(BufferSupplier bufferSupplier, boolean skipKeyValue) {
+    private CloseableIterator<Record> compressedIterator(BufferSupplier bufferSupplier, boolean skipKeyValue, int maxRecordBodySize) {
         final InputStream inputStream = recordInputStream(bufferSupplier);
 
         if (skipKeyValue) {
             return new StreamRecordIterator(inputStream) {
                 @Override
                 protected Record doReadRecord(long baseOffset, long baseTimestamp, int baseSequence, Long logAppendTime) throws IOException {
-                    return DefaultRecord.readPartiallyFrom(inputStream, baseOffset, baseTimestamp, baseSequence, logAppendTime);
+                    return DefaultRecord.readPartiallyFrom(inputStream, baseOffset, baseTimestamp, baseSequence, logAppendTime, maxRecordBodySize);
                 }
             };
         } else {
             return new StreamRecordIterator(inputStream) {
                 @Override
                 protected Record doReadRecord(long baseOffset, long baseTimestamp, int baseSequence, Long logAppendTime) throws IOException {
-                    return DefaultRecord.readFrom(inputStream, baseOffset, baseTimestamp, baseSequence, logAppendTime);
+                    return DefaultRecord.readFrom(inputStream, baseOffset, baseTimestamp, baseSequence, logAppendTime, maxRecordBodySize);
                 }
             };
         }
@@ -327,7 +327,7 @@ public Iterator<Record> iterator() {
         // for a normal iterator, we cannot ensure that the underlying compression stream is closed,
         // so we decompress the full record set here. Use cases which call for a lower memory footprint
         // can use `streamingIterator` at the cost of additional complexity
-        try (CloseableIterator<Record> iterator = compressedIterator(BufferSupplier.NO_CACHING, false)) {
+        try (CloseableIterator<Record> iterator = compressedIterator(BufferSupplier.NO_CACHING, false, Records.SOFT_MAX_ARRAY_LENGTH)) {
             List<Record> records = new ArrayList<>(count());
             while (iterator.hasNext())
                 records.add(iterator.next());
@@ -337,6 +337,11 @@ public Iterator<Record> iterator() {
 
     @Override
     public CloseableIterator<Record> skipKeyValueIterator(BufferSupplier bufferSupplier) {
+        return skipKeyValueIterator(bufferSupplier, Records.SOFT_MAX_ARRAY_LENGTH);
+    }
+
+    @Override
+    public CloseableIterator<Record> skipKeyValueIterator(BufferSupplier bufferSupplier, int maxRecordBodySize) {
         if (count() == 0) {
             return CloseableIterator.wrap(Collections.emptyIterator());
         }
@@ -351,13 +356,18 @@ public CloseableIterator<Record> skipKeyValueIterator(BufferSupplier bufferSuppl
 
         // we define this to be a closable iterator so that caller (i.e. the log validator) needs to close it
         // while we can save memory footprint of not decompressing the full record set ahead of time
-        return compressedIterator(bufferSupplier, true);
+        return compressedIterator(bufferSupplier, true, maxRecordBodySize);
     }
 
     @Override
     public CloseableIterator<Record> streamingIterator(BufferSupplier bufferSupplier) {
+        return streamingIterator(bufferSupplier, Records.SOFT_MAX_ARRAY_LENGTH);
+    }
+
+    @Override
+    public CloseableIterator<Record> streamingIterator(BufferSupplier bufferSupplier, int maxRecordBodySize) {
         if (isCompressed())
-            return compressedIterator(bufferSupplier, false);
+            return compressedIterator(bufferSupplier, false, maxRecordBodySize);
         else
             return uncompressedIterator();
     }

clients/src/main/java/org/apache/kafka/common/record/FileLogInputStream.java

@@ -158,6 +158,11 @@ public CloseableIterator<Record> streamingIterator(BufferSupplier bufferSupplier
             return loadFullBatch().streamingIterator(bufferSupplier);
         }
 
+        @Override
+        public CloseableIterator<Record> streamingIterator(BufferSupplier bufferSupplier, int maxRecordBodySize) {
+            return loadFullBatch().streamingIterator(bufferSupplier, maxRecordBodySize);
+        }
+
         @Override
         public boolean isValid() {
             return loadFullBatch().isValid();

clients/src/main/java/org/apache/kafka/common/record/FileRecords.java

@@ -20,6 +20,8 @@
 import org.apache.kafka.common.network.TransferableChannel;
 import org.apache.kafka.common.record.FileLogInputStream.FileChannelRecordBatch;
 import org.apache.kafka.common.utils.AbstractIterator;
+import org.apache.kafka.common.utils.BufferSupplier;
+import org.apache.kafka.common.utils.CloseableIterator;
 import org.apache.kafka.common.utils.Utils;
 
 import java.io.Closeable;
@@ -351,17 +353,25 @@ public LogOffsetPosition searchForOffsetFromPosition(long targetOffset, int star
      * @param targetTimestamp The timestamp to search for.
      * @param startingPosition The starting position to search.
      * @param startingOffset The starting offset to search.
+     * @param maxRecordBodySize The maximum declared (decompressed) body size of a single record; a compressed
+     *                          record exceeding it is rejected with an InvalidRecordException before its body is
+     *                          allocated. Pass {@link Records#SOFT_MAX_ARRAY_LENGTH} for no limit beyond the
+     *                          array-length ceiling.
      * @return The timestamp and offset of the message found. Null if no message is found.
      */
-    public TimestampAndOffset searchForTimestamp(long targetTimestamp, int startingPosition, long startingOffset) {
+    public TimestampAndOffset searchForTimestamp(long targetTimestamp, int startingPosition, long startingOffset,
+                                                 int maxRecordBodySize) {
         for (RecordBatch batch : batchesFrom(startingPosition)) {
             if (batch.maxTimestamp() >= targetTimestamp) {
                 // We found a message
-                for (Record record : batch) {
-                    long timestamp = record.timestamp();
-                    if (timestamp >= targetTimestamp && record.offset() >= startingOffset)
-                        return new TimestampAndOffset(timestamp, record.offset(),
-                                maybeLeaderEpoch(batch.partitionLeaderEpoch()));
+                try (CloseableIterator<Record> iterator = batch.streamingIterator(BufferSupplier.NO_CACHING, maxRecordBodySize)) {
+                    while (iterator.hasNext()) {
+                        Record record = iterator.next();
+                        long timestamp = record.timestamp();
+                        if (timestamp >= targetTimestamp && record.offset() >= startingOffset)
+                            return new TimestampAndOffset(timestamp, record.offset(),
+                                    maybeLeaderEpoch(batch.partitionLeaderEpoch()));
+                    }
                 }
             }
         }