You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ma...@apache.org on 2026/09/18 08:49:19 UTC
(kafka) branch 4.4 updated: Bound array and tagged-field allocation in generated message readers
This is an automated email from the ASF dual-hosted git repository.
omkreddy pushed a commit to branch 4.4
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/4.4 by this push:
new 1a770734fe9 Bound array and tagged-field allocation in generated message readers
1a770734fe9 is described below
commit 1a770734fe936ea5c6f46a9e748bee190ceba8e0
Author: Manikumar Reddy <ma...@confluent.io>
AuthorDate: Wed Sep 16 16:30:47 2026 +0530
Bound array and tagged-field allocation in generated message readers
Generated request/response readers sized array and tagged-field
collections directly from the sender-declared count. The existing
check bounds a declared count against the remaining buffer bytes, but
that bound assumes roughly one byte per element; a count that passes
it can still translate into a large initial allocation once elements
are unpacked into Java objects, and tagged-field counts had no
remaining-bytes check at all.
Adds, mirroring the existing array-length guard:
- Initial pre-allocation for array/collection fields is capped at
MessageUtil.MAX_PREALLOCATED_ARRAY_CAPACITY (1000); a larger
declared count still parses correctly, growing the collection on
demand as elements are read.
- A declared array/collection length above
MessageUtil.MAX_ARRAY_LENGTH (1,000,000) is now rejected up front.
- The tagged-field count for flexible-versioned messages is now
checked the same way array lengths already are: rejected if
negative, if it exceeds the remaining bytes, or if it exceeds
MessageUtil.MAX_TAGGED_FIELD_COUNT (10,000).
New tests cover the pre-allocation cap and hard-length cap on both the
ArrayList and keyed-collection (ImplicitLinkedHashMultiCollection)
branches, at the generator level and in a runtime round-trip, the new
tagged-field guards, and end-to-end RequestContext/RequestHeader cases
using the actual ProduceRequestData (v3) and RequestHeaderData (v2)
production classes at large declared-count magnitudes backed by real
padding bytes, so the pre-existing remaining-bytes check alone would
not reject them and the new hard caps are what's being exercised.
Co-authored-by: Izzy Harker <10...@users.noreply.github.com>
Co-authored-by: Alyssa Huang <ah...@confluent.io>
Co-authored-by: Truc Nguyen <11...@users.noreply.github.com>
Co-authored-by: Shay Elkin <20...@users.noreply.github.com>
---
.../apache/kafka/common/protocol/MessageUtil.java | 6 ++
.../apache/kafka/common/message/MessageTest.java | 63 +++++++++++++++++
.../common/message/SimpleArraysMessageTest.java | 56 +++++++++++++++
.../kafka/common/requests/RequestContextTest.java | 72 ++++++++++++++++++-
.../kafka/common/requests/RequestHeaderTest.java | 40 +++++++++++
.../common/message/SimpleKeyedArraysMessage.json | 28 ++++++++
.../apache/kafka/message/MessageDataGenerator.java | 32 ++++++++-
.../kafka/message/MessageDataGeneratorTest.java | 80 ++++++++++++++++++++++
8 files changed, 373 insertions(+), 4 deletions(-)
diff --git a/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java b/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
index fe79bb94e65..f2c25325cf4 100644
--- a/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
+++ b/clients/src/main/java/org/apache/kafka/common/protocol/MessageUtil.java
@@ -36,6 +36,12 @@ public final class MessageUtil {
public static final int UNSIGNED_SHORT_MAX = 65535;
+ public static final int MAX_PREALLOCATED_ARRAY_CAPACITY = 1000;
+
+ public static final int MAX_ARRAY_LENGTH = 1_000_000;
+
+ public static final int MAX_TAGGED_FIELD_COUNT = 10_000;
+
/**
* Copy a byte buffer into an array. This will not affect the buffer's
* position or mark.
diff --git a/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java b/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
index c339db8f35e..870961c2617 100644
--- a/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/message/MessageTest.java
@@ -57,6 +57,7 @@ import org.apache.kafka.common.protocol.MessageUtil;
import org.apache.kafka.common.protocol.ObjectSerializationCache;
import org.apache.kafka.common.protocol.types.RawTaggedField;
import org.apache.kafka.common.utils.annotation.ApiKeyVersionsSource;
+import org.apache.kafka.common.utils.internals.ByteUtils;
import com.fasterxml.jackson.databind.JsonNode;
@@ -857,6 +858,68 @@ public final class MessageTest {
verifyWriteSucceeds((short) 6, createTopics);
}
+ private byte[] rawTaggedFieldsSection(int declaredCount, int... tagsAndSizes) {
+ ByteBuffer scratch = ByteBuffer.allocate(64);
+ ByteUtils.writeUnsignedVarint(declaredCount, scratch);
+ for (int i = 0; i < tagsAndSizes.length; i += 2) {
+ ByteUtils.writeUnsignedVarint(tagsAndSizes[i], scratch);
+ ByteUtils.writeUnsignedVarint(tagsAndSizes[i + 1], scratch);
+ }
+ scratch.flip();
+ byte[] bytes = new byte[scratch.remaining()];
+ scratch.get(bytes);
+ return bytes;
+ }
+
+ private ByteBuffer messageWithTaggedFieldsSection(short version, byte[] taggedFieldsSection) {
+ SimpleExampleMessageData message = new SimpleExampleMessageData();
+ ObjectSerializationCache cache = new ObjectSerializationCache();
+ int size = message.size(cache, version);
+ ByteBuffer prefix = ByteBuffer.allocate(size);
+ message.write(new ByteBufferAccessor(prefix), cache, version);
+ prefix.flip();
+ byte[] prefixBytes = new byte[prefix.remaining() - 1];
+ prefix.get(prefixBytes);
+ assertEquals((byte) 0, prefix.get(), "expected an empty tagged-fields section to replace");
+ ByteBuffer result = ByteBuffer.allocate(prefixBytes.length + taggedFieldsSection.length);
+ result.put(prefixBytes);
+ result.put(taggedFieldsSection);
+ result.flip();
+ return result;
+ }
+
+ @Test
+ public void testTaggedFieldCountRejectedWhenLargerThanRemainingBytes() {
+ short version = 1;
+ ByteBuffer buf = messageWithTaggedFieldsSection(version, rawTaggedFieldsSection(1_000_000));
+ SimpleExampleMessageData message = new SimpleExampleMessageData();
+ RuntimeException e = assertThrows(RuntimeException.class,
+ () -> message.read(new ByteBufferAccessor(buf), version));
+ assertTrue(e.getMessage().contains("tagged fields"),
+ "Expected a bounded-count rejection, but got: " + e.getMessage());
+ }
+
+ @Test
+ public void testTaggedFieldCountRejectedWhenExceedingHardCap() {
+ short version = 1;
+ int declaredCount = MessageUtil.MAX_TAGGED_FIELD_COUNT + 1;
+ ByteBuffer scratch = ByteBuffer.allocate(16);
+ ByteUtils.writeUnsignedVarint(declaredCount, scratch);
+ scratch.flip();
+ byte[] countBytes = new byte[scratch.remaining()];
+ scratch.get(countBytes);
+ byte[] padding = new byte[declaredCount];
+ ByteBuffer section = ByteBuffer.allocate(countBytes.length + padding.length);
+ section.put(countBytes);
+ section.put(padding);
+ ByteBuffer buf = messageWithTaggedFieldsSection(version, section.array());
+ SimpleExampleMessageData message = new SimpleExampleMessageData();
+ RuntimeException e = assertThrows(RuntimeException.class,
+ () -> message.read(new ByteBufferAccessor(buf), version));
+ assertTrue(e.getMessage().contains("exceeds the maximum allowed count"),
+ "Expected a hard-cap rejection, but got: " + e.getMessage());
+ }
+
@Test
public void testLongTaggedString() {
char[] chars = new char[1024];
diff --git a/clients/src/test/java/org/apache/kafka/common/message/SimpleArraysMessageTest.java b/clients/src/test/java/org/apache/kafka/common/message/SimpleArraysMessageTest.java
index 963553a6f2f..7601f0ccf11 100644
--- a/clients/src/test/java/org/apache/kafka/common/message/SimpleArraysMessageTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/message/SimpleArraysMessageTest.java
@@ -17,6 +17,7 @@
package org.apache.kafka.common.message;
import org.apache.kafka.common.protocol.ByteBufferAccessor;
+import org.apache.kafka.common.protocol.MessageUtil;
import org.junit.jupiter.api.Test;
@@ -52,4 +53,59 @@ public class SimpleArraysMessageTest {
assertEquals("Tried to allocate a collection of size 125, but there are only 6 bytes remaining.",
assertThrows(RuntimeException.class, () -> out.read(accessor, (short) 2)).getMessage());
}
+
+ @Test
+ public void testDeclaredLengthAboveInitialCapacityStillParses() {
+ final int count = 1010;
+ final ByteBuffer buf = ByteBuffer.allocate(4 + (count * 4));
+ buf.putInt(count);
+ for (int i = 0; i < count; i++) {
+ buf.putInt(i);
+ }
+ buf.flip();
+ final SimpleArraysMessageData out = new SimpleArraysMessageData();
+ out.read(new ByteBufferAccessor(buf), (short) 0);
+ assertEquals(count, out.sheep().size());
+ }
+
+ @Test
+ public void testKeyedDeclaredLengthAboveInitialCapacityStillParses() {
+ final int count = 1010;
+ final ByteBuffer buf = ByteBuffer.allocate(4 + (count * 8));
+ buf.putInt(count);
+ for (int i = 0; i < count; i++) {
+ buf.putInt(i);
+ buf.putInt(i * 2);
+ }
+ buf.flip();
+ final SimpleKeyedArraysMessageData out = new SimpleKeyedArraysMessageData();
+ out.read(new ByteBufferAccessor(buf), (short) 0);
+ assertEquals(count, out.keyedStructs().size());
+ }
+
+ @Test
+ public void testArrayLengthAboveMaxIsRejected() {
+ final int count = MessageUtil.MAX_ARRAY_LENGTH + 1;
+ final ByteBuffer buf = ByteBuffer.allocate(4 + count);
+ buf.putInt(count);
+ buf.rewind();
+ final SimpleArraysMessageData out = new SimpleArraysMessageData();
+ final ByteBufferAccessor accessor = new ByteBufferAccessor(buf);
+ assertEquals("Tried to read a collection of size " + count + ", which exceeds the maximum allowed size of "
+ + MessageUtil.MAX_ARRAY_LENGTH + ".",
+ assertThrows(RuntimeException.class, () -> out.read(accessor, (short) 0)).getMessage());
+ }
+
+ @Test
+ public void testKeyedArrayLengthAboveMaxIsRejected() {
+ final int count = MessageUtil.MAX_ARRAY_LENGTH + 1;
+ final ByteBuffer buf = ByteBuffer.allocate(4 + count);
+ buf.putInt(count);
+ buf.rewind();
+ final SimpleKeyedArraysMessageData out = new SimpleKeyedArraysMessageData();
+ final ByteBufferAccessor accessor = new ByteBufferAccessor(buf);
+ assertEquals("Tried to read a collection of size " + count + ", which exceeds the maximum allowed size of "
+ + MessageUtil.MAX_ARRAY_LENGTH + ".",
+ assertThrows(RuntimeException.class, () -> out.read(accessor, (short) 0)).getMessage());
+ }
}
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java b/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java
index ba53edcacf8..06c2370de7c 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/RequestContextTest.java
@@ -132,6 +132,55 @@ public class RequestContextTest {
assertEquals("Tried to allocate a collection of size 1073741823, but there are only 17 bytes remaining.", msg);
}
+ @Test
+ public void testInvalidRequestForImplicitHashCollectionWithMaxPreAllocation() throws UnknownHostException {
+ short version = (short) 7; // non-flexible, name-based and no topicId
+ int actualTopicCount = 1010; // (greater than the 1000 pre-alloc cap)
+ ProduceRequestData data = new ProduceRequestData().setAcks((short) -1).setTimeoutMs(1);
+ for (int i = 0; i < actualTopicCount; i++) {
+ data.topicData().add(new ProduceRequestData.TopicProduceData()
+ .setName("t" + i)
+ .setPartitionData(Collections.emptyList()));
+ }
+ ByteBuffer buffer = serialize(version, data);
+
+ // Inflate the declared topic_data array length.
+ buffer.putInt(8, 5000);
+
+ RequestHeader header = new RequestHeader(ApiKeys.PRODUCE, version, "console-producer", 3);
+ RequestContext context = new RequestContext(header, "0", InetAddress.getLocalHost(),
+ KafkaPrincipal.ANONYMOUS, new ListenerName("ssl"), SecurityProtocol.SASL_SSL,
+ ClientInformation.EMPTY, true);
+
+ assertThrows(InvalidRequestException.class, () -> context.parseRequest(buffer));
+ }
+
+ @Test
+ public void testProduceRequestV3WithHugeDeclaredTopicCountIsRejected() throws UnknownHostException {
+ short version = 3; // ProduceRequestData.LOWEST_SUPPORTED_VERSION
+ int declaredTopicCount = 89_999_978;
+
+ // Real ~90MB buffer, backed by declaredTopicCount padding bytes, so the pre-existing
+ // count-vs-remaining-bytes guard alone would let this through and the
+ // new hard cap is what has to catch it.
+ ByteBuffer buffer = ByteBuffer.allocate(12 + declaredTopicCount);
+ buffer.putShort((short) -1); // null transactionalId
+ buffer.putShort((short) -1); // acks
+ buffer.putInt(1); // timeoutMs
+ buffer.putInt(declaredTopicCount); // topic_data declared length
+ buffer.position(12 + declaredTopicCount);
+ buffer.flip();
+
+ RequestHeader header = new RequestHeader(ApiKeys.PRODUCE, version, "console-producer", 3);
+ RequestContext context = new RequestContext(header, "0", InetAddress.getLocalHost(),
+ KafkaPrincipal.ANONYMOUS, new ListenerName("plaintext"), SecurityProtocol.PLAINTEXT,
+ ClientInformation.EMPTY, false);
+
+ InvalidRequestException e = assertThrows(InvalidRequestException.class, () -> context.parseRequest(buffer));
+ assertTrue(e.getCause().getMessage().contains("exceeds the maximum allowed size"),
+ "Expected a hard-cap rejection, but got: " + e.getCause().getMessage());
+ }
+
@Test
public void testInvalidRequestForArrayList() throws UnknownHostException {
short version = (short) 7; // choose a version with fixed length encoding, for simplicity
@@ -150,6 +199,25 @@ public class RequestContextTest {
"Tried to allocate a collection of size 2147483647, but there are only 8 bytes remaining.", msg);
}
+ @Test
+ public void testKeyedCollectionAboveInitialCapacityStillParses() {
+ short version = (short) 7; // non-flexible, name-based and no topicId
+ int count = 1010;
+ ProduceRequestData data = new ProduceRequestData()
+ .setAcks((short) -1)
+ .setTimeoutMs(1);
+ for (int i = 0; i < count; i++) {
+ data.topicData().add(new ProduceRequestData.TopicProduceData()
+ .setName("t" + i)
+ .setPartitionData(Collections.emptyList()));
+ }
+
+ ByteBuffer buffer = serialize(version, data);
+ ProduceRequestData parsed = new ProduceRequestData(new ByteBufferAccessor(buffer), version);
+
+ assertEquals(count, parsed.topicData().size());
+ }
+
private ByteBuffer produceRequest(short version) {
ProduceRequestData data = new ProduceRequestData()
.setAcks((short) -1)
@@ -165,8 +233,8 @@ public class RequestContextTest {
private ByteBuffer serialize(short version, ApiMessage data) {
ObjectSerializationCache cache = new ObjectSerializationCache();
- data.size(cache, version);
- ByteBuffer buffer = ByteBuffer.allocate(1024);
+ int size = data.size(cache, version);
+ ByteBuffer buffer = ByteBuffer.allocate(size);
data.write(new ByteBufferAccessor(buffer), cache, version);
buffer.flip();
return buffer;
diff --git a/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java b/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java
index 738a06444ed..5c441c82d2a 100644
--- a/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java
+++ b/clients/src/test/java/org/apache/kafka/common/requests/RequestHeaderTest.java
@@ -16,16 +16,20 @@
*/
package org.apache.kafka.common.requests;
+import org.apache.kafka.common.errors.InvalidRequestException;
import org.apache.kafka.common.message.RequestHeaderData;
import org.apache.kafka.common.protocol.ApiKeys;
import org.apache.kafka.common.protocol.ByteBufferAccessor;
import org.apache.kafka.common.protocol.ObjectSerializationCache;
+import org.apache.kafka.common.utils.internals.ByteUtils;
import org.junit.jupiter.api.Test;
import java.nio.ByteBuffer;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
@@ -92,6 +96,42 @@ public class RequestHeaderTest {
assertEquals((short) 10, parsed.apiVersion());
}
+ @Test
+ public void testHugeDeclaredTaggedFieldCountIsRejected() {
+ RequestHeaderData headerData = new RequestHeaderData().
+ setClientId("client").
+ setCorrelationId(123).
+ setRequestApiKey(ApiKeys.FIND_COORDINATOR.id).
+ setRequestApiVersion((short) 10);
+ ObjectSerializationCache serializationCache = new ObjectSerializationCache();
+ ByteBuffer prefix = ByteBuffer.allocate(headerData.size(serializationCache, (short) 2));
+ headerData.write(new ByteBufferAccessor(prefix), serializationCache, (short) 2);
+ prefix.flip();
+ byte[] prefixBytes = new byte[prefix.remaining() - 1];
+ prefix.get(prefixBytes);
+ assertEquals((byte) 0, prefix.get(), "expected an empty tagged-fields section to replace");
+
+ int declaredCount = 20_000_000;
+ ByteBuffer countBuf = ByteBuffer.allocate(8);
+ ByteUtils.writeUnsignedVarint(declaredCount, countBuf);
+ countBuf.flip();
+ byte[] countBytes = new byte[countBuf.remaining()];
+ countBuf.get(countBytes);
+
+ // Real buffer backed by declaredCount padding bytes, so the pre-existing
+ // count-vs-remaining-bytes guard alone would let this through and the new hard cap is
+ // what has to catch it.
+ ByteBuffer buffer = ByteBuffer.allocate(prefixBytes.length + countBytes.length + declaredCount);
+ buffer.put(prefixBytes);
+ buffer.put(countBytes);
+ buffer.position(prefixBytes.length + countBytes.length + declaredCount);
+ buffer.flip();
+
+ InvalidRequestException e = assertThrows(InvalidRequestException.class, () -> RequestHeader.parse(buffer));
+ assertTrue(e.getCause().getMessage().contains("exceeds the maximum allowed count"),
+ "Expected a hard-cap rejection, but got: " + e.getCause().getMessage());
+ }
+
@Test
public void verifySizeMethodsReturnSameValue() {
// Create a dummy RequestHeaderData
diff --git a/clients/src/test/resources/common/message/SimpleKeyedArraysMessage.json b/clients/src/test/resources/common/message/SimpleKeyedArraysMessage.json
new file mode 100644
index 00000000000..637bcdd71d4
--- /dev/null
+++ b/clients/src/test/resources/common/message/SimpleKeyedArraysMessage.json
@@ -0,0 +1,28 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+{
+ "name": "SimpleKeyedArraysMessage",
+ "type": "header",
+ "validVersions": "0-2",
+ "flexibleVersions": "1+",
+ "fields": [
+ { "name": "KeyedStructs", "type": "[]KeyedStruct", "versions": "0+",
+ "fields": [
+ { "name": "Key", "type": "int32", "versions": "0+", "mapKey": true },
+ { "name": "Value", "type": "int32", "versions": "0+" }
+ ]
+ }
+ ]
+}
diff --git a/generator/src/main/java/org/apache/kafka/message/MessageDataGenerator.java b/generator/src/main/java/org/apache/kafka/message/MessageDataGenerator.java
index a09da80f3e3..66f9aa8d8bb 100644
--- a/generator/src/main/java/org/apache/kafka/message/MessageDataGenerator.java
+++ b/generator/src/main/java/org/apache/kafka/message/MessageDataGenerator.java
@@ -487,6 +487,26 @@ public final class MessageDataGenerator implements MessageClassGenerator {
VersionConditional.forVersions(messageFlexibleVersions, curVersions).
ifMember(curFlexibleVersions -> {
buffer.printf("int _numTaggedFields = _readable.readUnsignedVarint();%n");
+ buffer.printf("if (_numTaggedFields < 0) {%n");
+ buffer.incrementIndent();
+ buffer.printf("throw new RuntimeException(\"Invalid negative number of tagged fields \" + " +
+ "_numTaggedFields);%n");
+ buffer.decrementIndent();
+ buffer.printf("}%n");
+ buffer.printf("if (_numTaggedFields > _readable.remaining()) {%n");
+ buffer.incrementIndent();
+ buffer.printf("throw new RuntimeException(\"Tried to read \" + _numTaggedFields + \" tagged " +
+ "fields, but there are only \" + _readable.remaining() + \" bytes remaining.\");%n");
+ buffer.decrementIndent();
+ buffer.printf("}%n");
+ headerGenerator.addImport(MessageGenerator.MESSAGE_UTIL_CLASS);
+ buffer.printf("if (_numTaggedFields > MessageUtil.MAX_TAGGED_FIELD_COUNT) {%n");
+ buffer.incrementIndent();
+ buffer.printf("throw new RuntimeException(\"Tried to read \" + _numTaggedFields + \" tagged " +
+ "fields, which exceeds the maximum allowed count of \" + " +
+ "MessageUtil.MAX_TAGGED_FIELD_COUNT + \".\");%n");
+ buffer.decrementIndent();
+ buffer.printf("}%n");
buffer.printf("for (int _i = 0; _i < _numTaggedFields; _i++) {%n");
buffer.incrementIndent();
buffer.printf("int _tag = _readable.readUnsignedVarint();%n");
@@ -660,16 +680,24 @@ public final class MessageDataGenerator implements MessageClassGenerator {
"there are only \" + _readable.remaining() + \" bytes remaining.\");%n", lengthVar);
buffer.decrementIndent();
buffer.printf("}%n");
+ headerGenerator.addImport(MessageGenerator.MESSAGE_UTIL_CLASS);
+ buffer.printf("if (%s > MessageUtil.MAX_ARRAY_LENGTH) {%n", lengthVar);
+ buffer.incrementIndent();
+ buffer.printf("throw new RuntimeException(\"Tried to read a collection of size \" + %s + \", which " +
+ "exceeds the maximum allowed size of \" + MessageUtil.MAX_ARRAY_LENGTH + \".\");%n", lengthVar);
+ buffer.decrementIndent();
+ buffer.printf("}%n");
if (isStructArrayWithKeys) {
headerGenerator.addImport(MessageGenerator.IMPLICIT_LINKED_HASH_MULTI_COLLECTION_CLASS);
- buffer.printf("%s newCollection = new %s(%s);%n",
+ buffer.printf("%s newCollection = new %s(Math.min(%s, MessageUtil.MAX_PREALLOCATED_ARRAY_CAPACITY));%n",
FieldSpec.collectionType(arrayType.elementType().toString()),
FieldSpec.collectionType(arrayType.elementType().toString()), lengthVar);
} else {
headerGenerator.addImport(MessageGenerator.ARRAYLIST_CLASS);
String boxedArrayType =
arrayType.elementType().getBoxedJavaType(headerGenerator);
- buffer.printf("ArrayList<%s> newCollection = new ArrayList<>(%s);%n", boxedArrayType, lengthVar);
+ buffer.printf("ArrayList<%s> newCollection = new ArrayList<>(Math.min(%s, MessageUtil.MAX_PREALLOCATED_ARRAY_CAPACITY));%n",
+ boxedArrayType, lengthVar);
}
buffer.printf("for (int i = 0; i < %s; i++) {%n", lengthVar);
buffer.incrementIndent();
diff --git a/generator/src/test/java/org/apache/kafka/message/MessageDataGeneratorTest.java b/generator/src/test/java/org/apache/kafka/message/MessageDataGeneratorTest.java
index 716f9e25687..c0f9cb4ab03 100644
--- a/generator/src/test/java/org/apache/kafka/message/MessageDataGeneratorTest.java
+++ b/generator/src/test/java/org/apache/kafka/message/MessageDataGeneratorTest.java
@@ -20,8 +20,11 @@ package org.apache.kafka.message;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
+import java.io.BufferedWriter;
+import java.io.StringWriter;
import java.util.Arrays;
+import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
@@ -77,6 +80,83 @@ public class MessageDataGeneratorTest {
"Expected string to contain '" + substring + "', but it was " + value);
}
+ private String generateMessageSource(MessageSpec spec) throws Exception {
+ MessageDataGenerator generator = new MessageDataGenerator("org.apache.kafka.common.message");
+ generator.generate(spec);
+ StringWriter writer = new StringWriter();
+ try (BufferedWriter buffered = new BufferedWriter(writer)) {
+ generator.write(buffered);
+ }
+ return writer.toString();
+ }
+
+ private MessageSpec parseSpec(String... lines) throws Exception {
+ return MessageGenerator.JSON_SERDE.readValue(String.join("", Arrays.asList(lines)), MessageSpec.class);
+ }
+
+ @Test
+ public void testArrayPreallocationIsCapped() throws Exception {
+ MessageSpec spec = parseSpec(
+ "{",
+ " \"type\": \"request\", \"name\": \"CapTest\", \"validVersions\": \"0-2\", \"flexibleVersions\": \"none\",",
+ " \"fields\": [",
+ " { \"name\": \"Foo\", \"type\": \"[]int32\", \"versions\": \"0+\" },",
+ " { \"name\": \"Bar\", \"type\": \"[]Baz\", \"versions\": \"1+\", \"fields\": [",
+ " { \"name\": \"Key\", \"type\": \"int32\", \"versions\": \"1+\", \"mapKey\": true },",
+ " { \"name\": \"Value\", \"type\": \"[]Bam\", \"versions\": \"2+\", \"fields\": [",
+ " {\"name\": \"NestedKey\", \"type\": \"string\", \"versions\": \"2+\", \"mapKey\": true },",
+ " { \"name\": \"NestedValue\", \"type\": \"string\", \"versions\": \"2+\" }",
+ " ]}",
+ " ]}",
+ " ]",
+ "}");
+ String source = generateMessageSource(spec);
+ assertStringContains("new ArrayList<>(Math.min(arrayLength, MessageUtil.MAX_PREALLOCATED_ARRAY_CAPACITY))", source);
+ assertStringContains("BazCollection(Math.min(arrayLength, MessageUtil.MAX_PREALLOCATED_ARRAY_CAPACITY))", source);
+ assertStringContains("BamCollection(Math.min(arrayLength, MessageUtil.MAX_PREALLOCATED_ARRAY_CAPACITY))", source);
+ }
+
+ @Test
+ public void testArrayLengthIsCapped() throws Exception {
+ MessageSpec spec = parseSpec(
+ "{",
+ " \"type\": \"request\", \"name\": \"CapTest\", \"validVersions\": \"0-2\", \"flexibleVersions\": \"none\",",
+ " \"fields\": [",
+ " { \"name\": \"Foo\", \"type\": \"[]int32\", \"versions\": \"0+\" },",
+ " { \"name\": \"Bar\", \"type\": \"[]Baz\", \"versions\": \"1+\", \"fields\": [",
+ " { \"name\": \"Key\", \"type\": \"int32\", \"versions\": \"1+\", \"mapKey\": true },",
+ " { \"name\": \"Value\", \"type\": \"[]Bam\", \"versions\": \"2+\", \"fields\": [",
+ " {\"name\": \"NestedKey\", \"type\": \"string\", \"versions\": \"2+\", \"mapKey\": true },",
+ " { \"name\": \"NestedValue\", \"type\": \"string\", \"versions\": \"2+\" }",
+ " ]}",
+ " ]}",
+ " ]",
+ "}");
+ String source = generateMessageSource(spec);
+ int occurrences = source.split("if \\(arrayLength > MessageUtil.MAX_ARRAY_LENGTH\\) \\{", -1).length - 1;
+ assertEquals(3, occurrences,
+ "Expected the cap check for all 3 array fields (Foo, Bar, Value), but found " + occurrences);
+ }
+
+ @Test
+ public void testTaggedFieldCountIsBounded() throws Exception {
+ MessageSpec spec = parseSpec(
+ "{",
+ " \"type\": \"request\", \"name\": \"TagBoundTest\", \"validVersions\": \"0-1\", \"flexibleVersions\": \"1+\",",
+ " \"fields\": [",
+ " { \"name\": \"Foo\", \"type\": \"int32\", \"versions\": \"1+\", \"taggedVersions\": \"1+\", \"tag\": 0, \"default\": \"0\" }",
+ " ]",
+ "}");
+ String source = generateMessageSource(spec);
+ assertStringContains("if (_numTaggedFields < 0) {", source);
+ assertStringContains("if (_numTaggedFields > _readable.remaining()) {", source);
+ assertStringContains("Tried to read \" + _numTaggedFields + \" tagged fields, but there are only \" + " +
+ "_readable.remaining() + \" bytes remaining.", source);
+ assertStringContains("if (_numTaggedFields > MessageUtil.MAX_TAGGED_FIELD_COUNT) {", source);
+ assertStringContains("Tried to read \" + _numTaggedFields + \" tagged fields, which exceeds the maximum " +
+ "allowed count of \" + MessageUtil.MAX_TAGGED_FIELD_COUNT + \".\"", source);
+ }
+
@Test
public void testInvalidNullDefaultForInt() throws Exception {
MessageSpec testMessageSpec = MessageGenerator.JSON_SERDE.readValue(String.join("", Arrays.asList(