Limit output size

src/main/java/com/hubspot/jinjava/JinjavaConfig.java

@@ -33,6 +33,7 @@ public class JinjavaConfig {
   private final Locale locale;
   private final ZoneId timeZone;
   private final int maxRenderDepth;
+  private final long maxOutputSize;
 
   private final boolean trimBlocks;
   private final boolean lstripBlocks;
@@ -48,11 +49,11 @@ public static Builder newBuilder() {
   }
 
   public JinjavaConfig() {
-    this(StandardCharsets.UTF_8, Locale.ENGLISH, ZoneOffset.UTC, 10, new HashMap<>(), false, false, true, false, false);
+    this(StandardCharsets.UTF_8, Locale.ENGLISH, ZoneOffset.UTC, 10, new HashMap<>(), false, false, true, false, false, 0);
   }
 
   public JinjavaConfig(Charset charset, Locale locale, ZoneId timeZone, int maxRenderDepth) {
-    this(charset, locale, timeZone, maxRenderDepth, new HashMap<>(), false, false, true, false, false);
+    this(charset, locale, timeZone, maxRenderDepth, new HashMap<>(), false, false, true, false, false, 0);
   }
 
   private JinjavaConfig(Charset charset,
@@ -65,7 +66,8 @@ private JinjavaConfig(Charset charset,
                         boolean lstripBlocks,
                         boolean readOnlyResolver,
                         boolean enableRecursiveMacroCalls,
-                        boolean failOnUnknownTokens) {
+                        boolean failOnUnknownTokens,
+                        long maxOutputSize) {
     this.charset = charset;
     this.locale = locale;
     this.timeZone = timeZone;
@@ -76,6 +78,7 @@ private JinjavaConfig(Charset charset,
     this.readOnlyResolver = readOnlyResolver;
     this.enableRecursiveMacroCalls = enableRecursiveMacroCalls;
     this.failOnUnknownTokens = failOnUnknownTokens;
+    this.maxOutputSize = maxOutputSize;
   }
 
   public Charset getCharset() {
@@ -94,6 +97,10 @@ public int getMaxRenderDepth() {
     return maxRenderDepth;
   }
 
+  public long getMaxOutputSize() {
+    return maxOutputSize;
+  }
+
   public boolean isTrimBlocks() {
     return trimBlocks;
   }
@@ -123,6 +130,7 @@ public static class Builder {
     private Locale locale = Locale.ENGLISH;
     private ZoneId timeZone = ZoneOffset.UTC;
     private int maxRenderDepth = 10;
+    private long maxOutputSize = 0; // in bytes
     private Map<Context.Library, Set<String>> disabled = new HashMap<>();
 
     private boolean trimBlocks;
@@ -184,8 +192,13 @@ public Builder withFailOnUnknownTokens(boolean failOnUnknownTokens) {
       return this;
     }
 
+    public Builder withMaxOutputSize(long maxOutputSize) {
+      this.maxOutputSize = maxOutputSize;
+      return this;
+    }
+
     public JinjavaConfig build() {
-      return new JinjavaConfig(charset, locale, timeZone, maxRenderDepth, disabled, trimBlocks, lstripBlocks, readOnlyResolver, enableRecursiveMacroCalls, failOnUnknownTokens);
+      return new JinjavaConfig(charset, locale, timeZone, maxRenderDepth, disabled, trimBlocks, lstripBlocks, readOnlyResolver, enableRecursiveMacroCalls, failOnUnknownTokens, maxOutputSize);
     }
 
   }

src/main/java/com/hubspot/jinjava/interpret/JinjavaInterpreter.java

@@ -40,6 +40,7 @@
 import com.hubspot.jinjava.tree.TreeParser;
 import com.hubspot.jinjava.tree.output.BlockPlaceholderOutputNode;
 import com.hubspot.jinjava.tree.output.OutputList;
+import com.hubspot.jinjava.tree.output.OutputNode;
 import com.hubspot.jinjava.util.Variable;
 import com.hubspot.jinjava.util.WhitespaceUtils;
 
@@ -185,21 +186,23 @@ public String render(Node root) {
    * @return rendered result
    */
   public String render(Node root, boolean processExtendRoots) {
-    OutputList output = new OutputList();
+    OutputList output = new OutputList(config.getMaxOutputSize());
 
     for (Node node : root.getChildren()) {
       lineNumber = node.getLineNumber();
-      output.addNode(node.render(this));
+      OutputNode out = node.render(this);
+      output.addNode(out);
     }
 
     // render all extend parents, keeping the last as the root output
     if (processExtendRoots) {
       while (!extendParentRoots.isEmpty()) {
         Node parentRoot = extendParentRoots.removeFirst();
-        output = new OutputList();
+        output = new OutputList(config.getMaxOutputSize());
 
         for (Node node : parentRoot.getChildren()) {
-          output.addNode(node.render(this));
+          OutputNode out = node.render(this);
+          output.addNode(out);
         }
 
         context.getExtendPathStack().pop();
@@ -226,7 +229,7 @@ private void resolveBlockStubs(OutputList output, Stack<String> blockNames) {
           List<? extends Node> superBlock = Iterables.get(blockChain, 1, null);
           context.setSuperBlock(superBlock);
 
-          OutputList blockValueBuilder = new OutputList();
+          OutputList blockValueBuilder = new OutputList(config.getMaxOutputSize());
 
           for (Node child : block) {
             blockValueBuilder.addNode(child.render(this));

src/main/java/com/hubspot/jinjava/interpret/OutputTooBigException.java

@@ -0,0 +1,16 @@
+package com.hubspot.jinjava.interpret;
+
+public class OutputTooBigException extends RuntimeException {
+
+  private final long size;
+
+  public OutputTooBigException(long size) {
+    this.size = size;
+  }
+
+  @Override
+  public String getMessage() {
+    return String.format("%d byte output rendered", size);
+  }
+
+}

src/main/java/com/hubspot/jinjava/tree/output/BlockPlaceholderOutputNode.java

@@ -30,6 +30,11 @@ public String getValue() {
     return output;
   }
 
+  @Override
+  public long getSize() {
+    return output == null ? 0 : output.getBytes().length;
+  }
+
   @Override
   public String toString() {
     return getValue();

src/main/java/com/hubspot/jinjava/tree/output/OutputList.java

@@ -3,12 +3,25 @@
 import java.util.LinkedList;
 import java.util.List;
 
+import com.hubspot.jinjava.interpret.OutputTooBigException;
+
 public class OutputList {
 
   private final List<OutputNode> nodes = new LinkedList<>();
   private final List<BlockPlaceholderOutputNode> blocks = new LinkedList<>();
+  private long maxOutputSize;
+
+  public OutputList(long maxOutputSize) {
+
+    this.maxOutputSize = maxOutputSize;
+  }
 
   public void addNode(OutputNode node) {
+
+    if (maxOutputSize > 0 && node.getSize() > maxOutputSize) {
+      throw new OutputTooBigException(node.getSize());
+    }
+
     nodes.add(node);
 
     if (node instanceof BlockPlaceholderOutputNode) {

src/main/java/com/hubspot/jinjava/tree/output/OutputNode.java

@@ -4,4 +4,6 @@ public interface OutputNode {
 
   String getValue();
 
+  long getSize();
+
 }

src/main/java/com/hubspot/jinjava/tree/output/RenderedOutputNode.java

@@ -18,4 +18,8 @@ public String toString() {
     return getValue();
   }
 
+  @Override
+  public long getSize() {
+    return output == null ? 0 : output.length();
+  }
 }

src/test/java/com/hubspot/jinjava/interpret/JinjavaInterpreterTest.java

@@ -10,6 +10,7 @@
 
 import com.google.common.collect.Lists;
 import com.hubspot.jinjava.Jinjava;
+import com.hubspot.jinjava.JinjavaConfig;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter.InterpreterScopeClosable;
 import com.hubspot.jinjava.interpret.TemplateError.ErrorReason;
 import com.hubspot.jinjava.tree.TextNode;
@@ -101,7 +102,8 @@ public void multiWordNumberSnakeCase() {
 
   @Test
   public void triesBeanMethodFirst() {
-    assertThat(interpreter.resolveProperty(ZonedDateTime.parse("2013-09-19T12:12:12+00:00"), "year").toString()).isEqualTo("2013");
+    assertThat(interpreter.resolveProperty(ZonedDateTime.parse("2013-09-19T12:12:12+00:00"), "year")
+        .toString()).isEqualTo("2013");
   }
 
   @Test
@@ -152,4 +154,17 @@ public void parseWithSyntaxError() {
     assertThat(result.getErrors().get(0).getReason()).isEqualTo(ErrorReason.SYNTAX_ERROR);
   }
 
+  @Test
+  public void itLimitsOutputSize() throws Exception {
+
+    JinjavaConfig outputSizeLimitedConfig = JinjavaConfig.newBuilder().withMaxOutputSize(20).build();
+    String output = "123456789012345678901234567890";
+
+    RenderResult renderResult = new Jinjava().renderForResult(output, new HashMap<>());
+    assertThat(renderResult.getOutput()).isEqualTo(output);
+    assertThat(renderResult.hasErrors()).isFalse();
+
+    renderResult = new Jinjava(outputSizeLimitedConfig).renderForResult(output, new HashMap<>());
+    assertThat(renderResult.getErrors().get(0).getMessage()).contains("OutputTooBigException");
+  }
 }