Fix map entry wrapping with limiting and adjust exceptions

jasmith-hs · jasmith-hs · commit d34bc172166a · 2023-01-25T10:51:59.000-05:00
diff --git a/src/main/java/com/hubspot/jinjava/objects/serialization/MapEntrySerializer.java b/src/main/java/com/hubspot/jinjava/objects/serialization/MapEntrySerializer.java
@@ -36,10 +36,7 @@ public void serialize(
         new CharArrayWriter(),
         remainingLength
       );
-      objectWriter.writeValue(
-        new SizeLimitingWriter(new CharArrayWriter(), remainingLength),
-        entry.getValue()
-      );
+      objectWriter.writeValue(sizeLimitingWriter, entry.getValue());
       value = sizeLimitingWriter.toString();
     } else {
       key = PyishObjectMapper.PYISH_OBJECT_WRITER.writeValueAsString(entry.getKey());
diff --git a/src/main/java/com/hubspot/jinjava/objects/serialization/PyishObjectMapper.java b/src/main/java/com/hubspot/jinjava/objects/serialization/PyishObjectMapper.java
@@ -9,6 +9,7 @@
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import com.hubspot.jinjava.interpret.DeferredValueException;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.interpret.OutputTooBigException;
 import com.hubspot.jinjava.util.WhitespaceUtils;
 import java.io.CharArrayWriter;
 import java.io.IOException;
@@ -48,7 +49,21 @@ public static String getAsPyishString(Object val) {
       return getAsPyishStringOrThrow(val);
     } catch (IOException e) {
       if (e instanceof SizeLimitingJsonProcessingException) {
-        throw new DeferredValueException(String.format("%s: %s", e.getMessage(), val));
+        if (
+          JinjavaInterpreter
+            .getCurrentMaybe()
+            .map(
+              interpreter -> interpreter.getConfig().getExecutionMode().useEagerParser()
+            )
+            .orElse(false)
+        ) {
+          throw new DeferredValueException(String.format("%s: %s", e.getMessage(), val));
+        } else {
+          throw new OutputTooBigException(
+            ((SizeLimitingJsonProcessingException) e).getMaxSize(),
+            ((SizeLimitingJsonProcessingException) e).getAttemptedSize()
+          );
+        }
       }
       return Objects.toString(val, "");
     }
diff --git a/src/main/java/com/hubspot/jinjava/objects/serialization/SizeLimitingJsonProcessingException.java b/src/main/java/com/hubspot/jinjava/objects/serialization/SizeLimitingJsonProcessingException.java
@@ -3,8 +3,26 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 
 public class SizeLimitingJsonProcessingException extends JsonProcessingException {
+  private final int maxSize;
+  private final int attemptedSize;
 
   protected SizeLimitingJsonProcessingException(int maxSize, int attemptedSize) {
-    super("Max length of {} chars reached when serializing. {} chars attempted.");
+    super(
+      String.format(
+        "Max length of %d chars reached when serializing. %d chars attempted.",
+        maxSize,
+        attemptedSize
+      )
+    );
+    this.maxSize = maxSize;
+    this.attemptedSize = attemptedSize;
+  }
+
+  public int getAttemptedSize() {
+    return attemptedSize;
+  }
+
+  public int getMaxSize() {
+    return maxSize;
   }
 }
diff --git a/src/main/java/com/hubspot/jinjava/objects/serialization/SizeLimitingWriter.java b/src/main/java/com/hubspot/jinjava/objects/serialization/SizeLimitingWriter.java
@@ -43,8 +43,8 @@ public void write(String str, int off, int len)
   private void checkMaxSize(int extra) throws SizeLimitingJsonProcessingException {
     if (remainingLength.addAndGet(extra * -1) < 0) {
       throw new SizeLimitingJsonProcessingException(
-        charArrayWriter.size() + extra,
-        startingLength
+        startingLength,
+        charArrayWriter.size() + extra
       );
     }
   }
diff --git a/src/test/java/com/hubspot/jinjava/lib/tag/eager/EagerTagDecoratorTest.java b/src/test/java/com/hubspot/jinjava/lib/tag/eager/EagerTagDecoratorTest.java
@@ -2,15 +2,16 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import com.hubspot.jinjava.BaseInterpretingTest;
 import com.hubspot.jinjava.JinjavaConfig;
 import com.hubspot.jinjava.interpret.DeferredValue;
 import com.hubspot.jinjava.interpret.DeferredValueException;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.interpret.OutputTooBigException;
-import com.hubspot.jinjava.interpret.TemplateError.ErrorReason;
 import com.hubspot.jinjava.lib.fn.ELFunctionDefinition;
 import com.hubspot.jinjava.lib.tag.Tag;
 import com.hubspot.jinjava.mode.EagerExecutionMode;
@@ -124,7 +125,7 @@ public void itLimitsTagLength() {
         () ->
           eagerTagDecorator.getEagerTagImage((TagToken) tagNode.getMaster(), interpreter)
       )
-      .isInstanceOf(OutputTooBigException.class);
+      .isInstanceOf(DeferredValueException.class);
   }
 
   @Test
diff --git a/src/test/java/com/hubspot/jinjava/objects/serialization/PyishObjectMapperTest.java b/src/test/java/com/hubspot/jinjava/objects/serialization/PyishObjectMapperTest.java
@@ -5,6 +5,9 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.google.common.collect.ImmutableMap;
+import com.hubspot.jinjava.Jinjava;
+import com.hubspot.jinjava.JinjavaConfig;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.objects.collections.SizeLimitingPyMap;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -33,6 +36,25 @@ public void itSerializesMapEntrySet() throws JsonProcessingException {
       .isEqualTo("[fn:map_entry('bar', {'foobar': []} ), fn:map_entry('foo', 'bar')]");
   }
 
+  @Test
+  public void itSerializesMapEntrySetWithLimit() throws JsonProcessingException {
+    SizeLimitingPyMap map = new SizeLimitingPyMap(new HashMap<>(), 10);
+    map.put("foo", "bar");
+    map.put("bar", ImmutableMap.of("foobar", new ArrayList<>()));
+
+    Jinjava jinjava = new Jinjava(
+      JinjavaConfig.newBuilder().withMaxOutputSize(10000).build()
+    );
+    try {
+      JinjavaInterpreter.pushCurrent(jinjava.newInterpreter());
+      String result = PyishObjectMapper.getAsPyishString(map.items());
+      assertThat(result)
+        .isEqualTo("[fn:map_entry('bar', {'foobar': []} ), fn:map_entry('foo', 'bar')]");
+    } finally {
+      JinjavaInterpreter.popCurrent();
+    }
+  }
+
   @Test
   public void itSerializesMapWithNullValues() {
     Map<String, Object> map = new SizeLimitingPyMap(new HashMap<>(), 10);

Original file line number	Diff line number	Diff line change
`@@ -43,8 +43,8 @@ public void write(String str, int off, int len)`
`43`	`43`	`private void checkMaxSize(int extra) throws SizeLimitingJsonProcessingException {`
`44`	`44`	`if (remainingLength.addAndGet(extra * -1) < 0) {`
`45`	`45`	`throw new SizeLimitingJsonProcessingException(`
`46`		`- charArrayWriter.size() + extra,`
`47`		`- startingLength`
	`46`	`+ startingLength,`
	`47`	`+ charArrayWriter.size() + extra`
`48`	`48`	`);`
`49`	`49`	`}`
`50`	`50`	`}`