feat: Add configurable stream read constraints for JSON input codec

- Introduced a configurable option for setting the maximum event length in the JSON input codec.
- Updated JsonDecoder to accept max_event_length parameter and apply it to StreamReadConstraints.
- Added validation to ensure the maximum event length is within acceptable limits.
Fixes #5466

Signed-off-by: Pallempati Saketh <[email protected]>

Author: saketh-pallempati

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java

@@ -8,6 +8,7 @@
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.core.StreamReadConstraints;
 import com.fasterxml.jackson.databind.ObjectMapper;
 
 import org.opensearch.dataprepper.model.event.Event;
@@ -30,10 +31,13 @@ public class JsonDecoder implements ByteDecoder {
     private Collection<String> includeKeys;
     private Collection<String> includeKeysMetadata;
 
-    public JsonDecoder(String keyName, Collection<String> includeKeys, Collection<String> includeKeysMetadata) {
+    public JsonDecoder(String keyName, Collection<String> includeKeys, Collection<String> includeKeysMetadata, int maxEventLength) {
         this.keyName = keyName;
         this.includeKeys = includeKeys;
         this.includeKeysMetadata = includeKeysMetadata;
+        jsonFactory.setStreamReadConstraints(StreamReadConstraints.builder()
+                .maxStringLength(maxEventLength)
+                .build());
     }
 
     public JsonDecoder() {

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java

@@ -22,9 +22,13 @@
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+
 
 import org.junit.jupiter.api.BeforeEach;
 
@@ -48,32 +52,68 @@ void test_basicJsonDecoder() {
         String stringValue = UUID.randomUUID().toString();
         Random r = new Random();
         int intValue = r.nextInt();
-        String inputString = "[{\"key1\":\""+stringValue+"\", \"key2\":"+intValue+"}]";
+        String inputString = "[{\"key1\":\"" + stringValue + "\", \"key2\":" + intValue + "}]";
         try {
             jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), null, (record) -> {
                 receivedRecord = record;
             });
-        } catch (Exception e){}
-        
+        } catch (Exception e) {
+        }
+
         assertNotEquals(receivedRecord, null);
         Map<String, Object> map = receivedRecord.getData().toMap();
         assertThat(map.get("key1"), equalTo(stringValue));
         assertThat(map.get("key2"), equalTo(intValue));
     }
 
+    @Test
+    void test_basicJsonDecoder_exceedingMaxEventLength_throwsException() {
+        String largeString = "x".repeat(200);
+        String inputString = "[{\"key1\":\"" + largeString + "\"}]";
+
+        jsonDecoder = new JsonDecoder(null, null, null, 100);
+
+        Exception exception = assertThrows(Exception.class, () -> {
+            jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), null, (record) -> {
+                receivedRecord = record;
+            });
+        });
+
+        assertEquals("String value length (200) exceeds the maximum allowed (100, from `StreamReadConstraints.getMaxStringLength()`)", exception.getMessage());
+    }
+
+    @Test
+    void test_basicJsonDecoder_withMaxEventLength() {
+        String validString = "Short string";
+        String inputString = "[{\"key1\":\"" + validString + "\"}]";
+
+        jsonDecoder = new JsonDecoder(null, null, null, 100);
+
+        assertDoesNotThrow(() -> {
+            jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), null, (record) -> {
+                receivedRecord = record;
+            });
+        });
+
+        assertNotNull(receivedRecord);
+        Map<String, Object> map = receivedRecord.getData().toMap();
+        assertThat(map.get("key1"), equalTo(validString));
+    }
+
     @Test
     void test_basicJsonDecoder_withTimeReceived() {
         String stringValue = UUID.randomUUID().toString();
         Random r = new Random();
         int intValue = r.nextInt();
-        String inputString = "[{\"key1\":\""+stringValue+"\", \"key2\":"+intValue+"}]";
+        String inputString = "[{\"key1\":\"" + stringValue + "\", \"key2\":" + intValue + "}]";
         final Instant now = Instant.now();
         try {
             jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), now, (record) -> {
                 receivedRecord = record;
                 receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
             });
-        } catch (Exception e){}
+        } catch (Exception e) {
+        }
 
         assertNotEquals(receivedRecord, null);
         Map<String, Object> map = receivedRecord.getData().toMap();
@@ -91,21 +131,23 @@ class JsonDecoderWithInputConfig {
         private static final int numKeyPerRecord = 3;
         private Map<String, Object> jsonObject;
         private final String key_name = "logEvents";
+        private final Integer maxEventLength = 20000000;
 
         @BeforeEach
         void setup() {
             objectMapper = new ObjectMapper();
-            for (int i=0; i<10; i++) {
+            for (int i = 0; i < 10; i++) {
                 includeKeys.add(UUID.randomUUID().toString());
                 includeMetadataKeys.add(UUID.randomUUID().toString());
             }
             jsonObject = generateJsonWithSpecificKeys(includeKeys, includeMetadataKeys, key_name, numKeyRecords, numKeyPerRecord);
         }
+
         @Test
         void test_basicJsonDecoder_withInputConfig() throws IOException {
             final Instant now = Instant.now();
             List<Record<Event>> records = new ArrayList<>();
-            jsonDecoder = new JsonDecoder(key_name, includeKeys, includeMetadataKeys);
+            jsonDecoder = new JsonDecoder(key_name, includeKeys, includeMetadataKeys, maxEventLength);
             jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
                 records.add(record);
                 receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
@@ -118,10 +160,10 @@ void test_basicJsonDecoder_withInputConfig() throws IOException {
                 Map<String, Object> dataMap = record.getData().toMap();
                 Map<String, Object> metadataMap = record.getData().getMetadata().getAttributes();
 
-                for (String includeKey: includeKeys) {
+                for (String includeKey : includeKeys) {
                     assertThat(dataMap.get(includeKey), equalTo(jsonObject.get(includeKey)));
                 }
-                for (String includeMetadataKey: includeMetadataKeys) {
+                for (String includeMetadataKey : includeMetadataKeys) {
                     assertThat(metadataMap.get(includeMetadataKey), equalTo(jsonObject.get(includeMetadataKey)));
                 }
             });
@@ -133,7 +175,7 @@ void test_basicJsonDecoder_withInputConfig() throws IOException {
         void test_basicJsonDecoder_withInputConfig_withoutEvents_empty_metadata_keys() throws IOException {
             final Instant now = Instant.now();
             List<Record<Event>> records = new ArrayList<>();
-            jsonDecoder = new JsonDecoder("", includeKeys, Collections.emptyList());
+            jsonDecoder = new JsonDecoder("", includeKeys, Collections.emptyList(), maxEventLength);
             jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
                 records.add(record);
                 receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
@@ -145,7 +187,7 @@ void test_basicJsonDecoder_withInputConfig_withoutEvents_empty_metadata_keys() t
         void test_basicJsonDecoder_withInputConfig_withoutEvents_null_include_metadata_keys() throws IOException {
             final Instant now = Instant.now();
             List<Record<Event>> records = new ArrayList<>();
-            jsonDecoder = new JsonDecoder("", includeKeys, null);
+            jsonDecoder = new JsonDecoder("", includeKeys, null, maxEventLength);
             jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
                 records.add(record);
                 receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
@@ -158,7 +200,7 @@ void test_basicJsonDecoder_withInputConfig_withoutEvents_null_include_metadata_k
         void test_basicJsonDecoder_withInputConfig_withoutEvents_empty_include_keys() throws IOException {
             final Instant now = Instant.now();
             List<Record<Event>> records = new ArrayList<>();
-            jsonDecoder = new JsonDecoder("", Collections.emptyList(), includeMetadataKeys);
+            jsonDecoder = new JsonDecoder("", Collections.emptyList(), includeMetadataKeys, maxEventLength);
             jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
                 records.add(record);
                 receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
@@ -170,7 +212,7 @@ void test_basicJsonDecoder_withInputConfig_withoutEvents_empty_include_keys() th
         void test_basicJsonDecoder_withInputConfig_withoutEvents_null_include_keys() throws IOException {
             final Instant now = Instant.now();
             List<Record<Event>> records = new ArrayList<>();
-            jsonDecoder = new JsonDecoder("", null, includeMetadataKeys);
+            jsonDecoder = new JsonDecoder("", null, includeMetadataKeys, maxEventLength);
             jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
                 records.add(record);
                 receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
@@ -187,17 +229,17 @@ private Map<String, Object> generateJsonWithSpecificKeys(final List<String> incl
             final Map<String, Object> jsonObject = new LinkedHashMap<>();
             final List<Map<String, Object>> innerObjects = new ArrayList<>();
 
-            for (String includeKey: includeKeys) {
+            for (String includeKey : includeKeys) {
                 jsonObject.put(includeKey, UUID.randomUUID().toString());
             }
 
-            for (String includeMetadataKey: includeMetadataKeys) {
+            for (String includeMetadataKey : includeMetadataKeys) {
                 jsonObject.put(includeMetadataKey, UUID.randomUUID().toString());
             }
 
-            for (int i=0; i<numKeyRecords; i++) {
+            for (int i = 0; i < numKeyRecords; i++) {
                 final Map<String, Object> innerJsonMap = new LinkedHashMap<>();
-                for (int j=0; j<numKeyPerRecord; j++) {
+                for (int j = 0; j < numKeyPerRecord; j++) {
                     innerJsonMap.put(UUID.randomUUID().toString(), UUID.randomUUID().toString());
                 }
                 innerObjects.add(innerJsonMap);

data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodec.java

@@ -25,11 +25,10 @@ public class JsonInputCodec extends JsonDecoder implements InputCodec {
 
     @DataPrepperPluginConstructor
     public JsonInputCodec(final JsonInputCodecConfig config) {
-        super(Objects.requireNonNull(config).getKeyName(), config.getIncludeKeys(), config.getIncludeKeysMetadata());
+        super(Objects.requireNonNull(config).getKeyName(), config.getIncludeKeys(), config.getIncludeKeysMetadata(), config.getMaxEventLength());
     }
 
     public void parse(InputStream inputStream, Consumer<Record<Event>> eventConsumer) throws IOException {
         parse(inputStream, null, eventConsumer);
     }
-
 }

data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodecConfig.java

@@ -11,11 +11,13 @@
 package org.opensearch.dataprepper.plugins.codec.json;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.Min;
 import jakarta.validation.constraints.Size;
 
 import java.util.List;
 
 public class JsonInputCodecConfig {
+    static final Integer DEFAULT_MAX_EVENT_LENGTH = 20000000;
 
     @JsonProperty("key_name")
     @Size(min = 1, max = 2048)
@@ -38,4 +40,12 @@ public List<String> getIncludeKeys() {
     public List<String> getIncludeKeysMetadata() {
         return includeKeysMetadata;
     }
+
+    @JsonProperty("max_event_length")
+    @Min(1)
+    private Integer maxEventLength = DEFAULT_MAX_EVENT_LENGTH;
+
+    public Integer getMaxEventLength(){
+        return maxEventLength;
+    }
 }

data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/codec/json/JsonCodecsIT.java

@@ -37,6 +37,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.dataprepper.plugins.codec.json.JsonInputCodecConfig.DEFAULT_MAX_EVENT_LENGTH;
 
 public class JsonCodecsIT {
 
@@ -51,6 +52,7 @@ void setUp() {
         when(jsonInputCodecConfig.getIncludeKeysMetadata()).thenReturn(Collections.emptyList());
         when(jsonInputCodecConfig.getIncludeKeys()).thenReturn(Collections.emptyList());
         when(jsonInputCodecConfig.getKeyName()).thenReturn(null);
+        when(jsonInputCodecConfig.getMaxEventLength()).thenReturn(DEFAULT_MAX_EVENT_LENGTH);
         eventConsumer = mock(Consumer.class);
     }
 

data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodecConfigTest.java

@@ -7,7 +7,10 @@
 
 import org.junit.jupiter.api.Test;
 
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.opensearch.dataprepper.plugins.codec.json.JsonInputCodecConfig.DEFAULT_MAX_EVENT_LENGTH;
 
 public class JsonInputCodecConfigTest {
 
@@ -21,5 +24,6 @@ public void testJsonInputCodecConfig() {
         assertNull(jsonInputCodecConfig.getKeyName());
         assertNull(jsonInputCodecConfig.getIncludeKeys());
         assertNull(jsonInputCodecConfig.getIncludeKeysMetadata());
+        assertThat(jsonInputCodecConfig.getMaxEventLength(), equalTo(DEFAULT_MAX_EVENT_LENGTH));
     }
 }

data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/codec/json/JsonInputCodecTest.java

@@ -54,6 +54,7 @@
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
+import static org.opensearch.dataprepper.plugins.codec.json.JsonInputCodecConfig.DEFAULT_MAX_EVENT_LENGTH;
 
 class JsonInputCodecTest {
 
@@ -68,6 +69,7 @@ void setUp() {
         when(jsonInputCodecConfig.getIncludeKeysMetadata()).thenReturn(null);
         when(jsonInputCodecConfig.getIncludeKeys()).thenReturn(null);
         when(jsonInputCodecConfig.getKeyName()).thenReturn(null);
+        when(jsonInputCodecConfig.getMaxEventLength()).thenReturn(DEFAULT_MAX_EVENT_LENGTH);
         eventConsumer = mock(Consumer.class);
     }