(kafka) branch trunk updated: MINOR: improve ProcessorRecordContext (#23158)

[mj...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mjsax pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git


The following commit(s) were added to refs/heads/trunk by this push:
     new a85ec49904b MINOR: improve ProcessorRecordContext (#23158)
a85ec49904b is described below

commit a85ec49904b5dc8a9f8c8c29b8598fe89b9433ae
Author: Lucy Liu <lu...@confluent.io>
AuthorDate: Thu Aug 20 12:10:02 2026 -0500

    MINOR: improve ProcessorRecordContext (#23158)
    
    Reviewers: Matthias J. Sax <ma...@confluent.io>
---
 .../internals/ProcessorRecordContext.java          |  4 +++
 .../internals/ProcessorRecordContextTest.java      | 39 ++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java
index 8198645eb65..72ffcdd3a16 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContext.java
@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
+import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeader;
@@ -195,6 +196,9 @@ public class ProcessorRecordContext implements RecordContext, RecordMetadata {
         if (headerCount == -1) { // keep for backward compatibility
             headers = new RecordHeaders();
         } else {
+            if (headerCount > buffer.remaining() / (2 * Integer.BYTES)) {
+                throw new SerializationException();
+            }
             final Header[] headerArr = new Header[headerCount];
             for (int i = 0; i < headerCount; i++) {
                 final byte[] keyBytes = requireNonNull(getNullableSizePrefixedArray(buffer));
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContextTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContextTest.java
index 1324e706268..62aeac5dc1b 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContextTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorRecordContextTest.java
@@ -16,11 +16,15 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
+import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 
 import org.junit.jupiter.api.Test;
 
+import java.nio.ByteBuffer;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
@@ -110,4 +114,39 @@ public class ProcessorRecordContextTest {
 
         assertEquals(MIN_SIZE + 10L, context.residentMemorySizeEstimate());
     }
+
+    @Test
+    public void shouldRejectHeaderCountLargerThanRemainingBuffer() {
+        final Headers headers = new RecordHeaders();
+        headers.add("header-key", "header-value".getBytes(UTF_8));
+        final ProcessorRecordContext context = new ProcessorRecordContext(
+                42L, 73L, 0, "topic", headers);
+
+        final byte[] serialized = context.serialize();
+        final ByteBuffer buffer = ByteBuffer.wrap(serialized);
+
+        // Locate headerCount's position within the real serialized bytes
+        // timestamp(8) + offset(8) + topicLen(4) + "topic"(5) + partition(4)
+        final int headerCountOffset = 8 + 8 + 4 + "topic".getBytes(UTF_8).length + 4;
+        final int bytesAfterHeaderCount = serialized.length - (headerCountOffset + 4);
+        final int maxPlausibleHeaderCount = bytesAfterHeaderCount / (2 * Integer.BYTES);
+
+        // Overwrite the real headerCount with a little more that it could support
+        buffer.putInt(headerCountOffset, maxPlausibleHeaderCount + 1);
+
+        assertThrows(SerializationException.class, () -> ProcessorRecordContext.deserialize(buffer));
+    }
+
+    @Test
+    public void shouldDeserializeValidHeaderCountWithoutRejecting() {
+        final Headers headers = new RecordHeaders();
+        headers.add("header-key", "header-value".getBytes(UTF_8));
+        final ProcessorRecordContext context = new ProcessorRecordContext(
+                42L, 73L, 0, "topic", headers);
+
+        final ProcessorRecordContext roundTripped =
+                ProcessorRecordContext.deserialize(ByteBuffer.wrap(context.serialize()));
+
+        assertEquals(context, roundTripped);
+    }
 }