Merge pull request #1131 from HubSpot/length-limit-tojson-filter

Add length-limiting to ToJsonFilter to avoid serializing too large a value

Author: jasmith-hs

src/main/java/com/hubspot/jinjava/lib/filter/ToJsonFilter.java

@@ -1,13 +1,18 @@
 package com.hubspot.jinjava.lib.filter;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.hubspot.jinjava.doc.annotations.JinjavaDoc;
 import com.hubspot.jinjava.doc.annotations.JinjavaParam;
 import com.hubspot.jinjava.doc.annotations.JinjavaSnippet;
 import com.hubspot.jinjava.interpret.DeferredValueException;
 import com.hubspot.jinjava.interpret.InvalidInputException;
 import com.hubspot.jinjava.interpret.InvalidReason;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.objects.serialization.LengthLimitingWriter;
+import com.hubspot.jinjava.objects.serialization.PyishObjectMapper;
+import java.io.CharArrayWriter;
+import java.io.IOException;
+import java.io.Writer;
+import java.util.concurrent.atomic.AtomicInteger;
 
 @JinjavaDoc(
   value = "Writes object as a JSON string",
@@ -23,11 +28,21 @@ public class ToJsonFilter implements Filter {
   @Override
   public Object filter(Object var, JinjavaInterpreter interpreter, String... args) {
     try {
-      return interpreter.getConfig().getObjectMapper().writeValueAsString(var);
-    } catch (JsonProcessingException e) {
+      if (interpreter.getConfig().getMaxOutputSize() > 0) {
+        AtomicInteger remainingLength = new AtomicInteger(
+          (int) Math.min(Integer.MAX_VALUE, interpreter.getConfig().getMaxOutputSize())
+        );
+        Writer writer = new LengthLimitingWriter(new CharArrayWriter(), remainingLength);
+        interpreter.getConfig().getObjectMapper().writeValue(writer, var);
+        return writer.toString();
+      } else {
+        return interpreter.getConfig().getObjectMapper().writeValueAsString(var);
+      }
+    } catch (IOException e) {
       if (e.getCause() instanceof DeferredValueException) {
         throw (DeferredValueException) e.getCause();
       }
+      PyishObjectMapper.handleLengthLimitingException(e);
       throw new InvalidInputException(interpreter, this, InvalidReason.JSON_WRITE);
     }
   }

src/main/java/com/hubspot/jinjava/objects/serialization/PyishObjectMapper.java

@@ -67,22 +67,26 @@ private static String getAsPyishString(Object val, boolean forOutput) {
     try {
       return getAsPyishStringOrThrow(val, forOutput);
     } catch (IOException e) {
-      Throwable unwrapped = e;
-      if (e instanceof JsonMappingException) {
-        unwrapped = unwrapped.getCause();
-      }
-      if (unwrapped instanceof LengthLimitingJsonProcessingException) {
-        throw new OutputTooBigException(
-          ((LengthLimitingJsonProcessingException) unwrapped).getMaxSize(),
-          ((LengthLimitingJsonProcessingException) unwrapped).getAttemptedSize()
-        );
-      } else if (unwrapped instanceof OutputTooBigException) {
-        throw (OutputTooBigException) unwrapped;
-      }
+      handleLengthLimitingException(e);
       return Objects.toString(val, "");
     }
   }
 
+  public static void handleLengthLimitingException(IOException e) {
+    Throwable unwrapped = e;
+    if (e instanceof JsonMappingException) {
+      unwrapped = unwrapped.getCause();
+    }
+    if (unwrapped instanceof LengthLimitingJsonProcessingException) {
+      throw new OutputTooBigException(
+        ((LengthLimitingJsonProcessingException) unwrapped).getMaxSize(),
+        ((LengthLimitingJsonProcessingException) unwrapped).getAttemptedSize()
+      );
+    } else if (unwrapped instanceof OutputTooBigException) {
+      throw (OutputTooBigException) unwrapped;
+    }
+  }
+
   public static String getAsPyishStringOrThrow(Object val) throws IOException {
     return getAsPyishStringOrThrow(val, false);
   }

src/test/java/com/hubspot/jinjava/lib/filter/ToJsonFilterTest.java

@@ -3,7 +3,12 @@
 import static org.assertj.core.api.Java6Assertions.assertThat;
 
 import com.hubspot.jinjava.BaseInterpretingTest;
+import com.hubspot.jinjava.Jinjava;
+import com.hubspot.jinjava.JinjavaConfig;
+import com.hubspot.jinjava.interpret.OutputTooBigException;
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import org.junit.Before;
 import org.junit.Test;
@@ -27,4 +32,29 @@ public void itWritesObjectAsString() {
     assertThat(filter.filter(testMap, interpreter))
       .isEqualTo("{\"testArray\":[4,1,2],\"testString\":\"testString\"}");
   }
+
+  @Test
+  public void itLimitsLength() {
+    List<List<?>> original = new ArrayList<>();
+    List<List<?>> temp = original;
+    for (int i = 0; i < 100; i++) {
+      List<List<?>> nested = new ArrayList<>();
+      temp.add(nested);
+      temp = nested;
+    }
+    interpreter =
+      new Jinjava(JinjavaConfig.newBuilder().withMaxOutputSize(500).build())
+      .newInterpreter();
+    assertThat(filter.filter(original, interpreter)).asString().contains("[[]]]]");
+    for (int i = 0; i < 400; i++) {
+      List<List<?>> nested = new ArrayList<>();
+      temp.add(nested);
+      temp = nested;
+    }
+    try {
+      filter.filter(original, interpreter);
+    } catch (Exception e) {
+      assertThat(e).isInstanceOf(OutputTooBigException.class);
+    }
+  }
 }