Limit length of strings while building

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

@@ -22,6 +22,7 @@
 import com.hubspot.jinjava.doc.annotations.JinjavaParam;
 import com.hubspot.jinjava.doc.annotations.JinjavaSnippet;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 @JinjavaDoc(
     value = "Escapes strings so that they can be safely inserted into a JavaScript variable declaration",
@@ -38,7 +39,7 @@ public class EscapeJsFilter implements Filter {
   @Override
   public Object filter(Object objectToFilter, JinjavaInterpreter jinjavaInterpreter, String... strings) {
     String input = Objects.toString(objectToFilter, "");
-    StringBuilder builder = new StringBuilder();
+    LengthLimitingStringBuilder builder = new LengthLimitingStringBuilder(jinjavaInterpreter.getConfig().getMaxOutputSize());
 
     for (int i = 0; i < input.length(); i++) {
       char ch = input.charAt(i);

src/main/java/com/hubspot/jinjava/lib/fn/Functions.java

@@ -24,6 +24,7 @@
 import com.hubspot.jinjava.objects.date.PyishDate;
 import com.hubspot.jinjava.objects.date.StrftimeFormatter;
 import com.hubspot.jinjava.tree.Node;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 public class Functions {
 
@@ -38,7 +39,7 @@ public class Functions {
   })
   public static String renderSuperBlock() {
     JinjavaInterpreter interpreter = JinjavaInterpreter.getCurrent();
-    StringBuilder result = new StringBuilder();
+    LengthLimitingStringBuilder result = new LengthLimitingStringBuilder(interpreter.getConfig().getMaxOutputSize());
 
     List<? extends Node> superBlock = interpreter.getContext().getSuperBlock();
     if (superBlock != null) {

src/main/java/com/hubspot/jinjava/lib/fn/MacroFunction.java

@@ -9,6 +9,7 @@
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.interpret.JinjavaInterpreter.InterpreterScopeClosable;
 import com.hubspot.jinjava.tree.Node;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 /**
  * Function definition parsed from a jinjava template, stored in global macros registry in interpreter context.
@@ -63,7 +64,7 @@ public Object doEvaluate(Map<String, Object> argMap, Map<String, Object> kwargMa
       // varargs list
       interpreter.getContext().put("varargs", varArgs);
 
-      StringBuilder result = new StringBuilder();
+      LengthLimitingStringBuilder result = new LengthLimitingStringBuilder(interpreter.getConfig().getMaxOutputSize());
 
       for (Node node : content) {
         result.append(node.render(interpreter));

src/main/java/com/hubspot/jinjava/lib/tag/AutoEscapeTag.java

@@ -9,6 +9,7 @@
 import com.hubspot.jinjava.interpret.JinjavaInterpreter.InterpreterScopeClosable;
 import com.hubspot.jinjava.tree.Node;
 import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 @JinjavaDoc(
     value = "Autoescape the tag's contents",
@@ -39,7 +40,7 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
       boolean escapeFlag = BooleanUtils.toBoolean(StringUtils.isNotBlank(boolFlagStr) ? boolFlagStr : "true");
       interpreter.getContext().setAutoEscape(escapeFlag);
 
-      StringBuilder result = new StringBuilder();
+      LengthLimitingStringBuilder result = new LengthLimitingStringBuilder(interpreter.getConfig().getMaxOutputSize());
 
       for (Node child : tagNode.getChildren()) {
         result.append(child.render(interpreter));

src/main/java/com/hubspot/jinjava/lib/tag/ForTag.java

@@ -36,6 +36,7 @@
 import com.hubspot.jinjava.util.ForLoop;
 import com.hubspot.jinjava.util.HelperStringTokenizer;
 import com.hubspot.jinjava.util.ObjectIterator;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 /**
  * {% for a in b|f1:d,c %}
@@ -125,7 +126,7 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
     try (InterpreterScopeClosable c = interpreter.enterScope()) {
       interpreter.getContext().put(LOOP, loop);
 
-      StringBuilder buff = new StringBuilder();
+      LengthLimitingStringBuilder buff = new LengthLimitingStringBuilder(interpreter.getConfig().getMaxOutputSize());
       while (loop.hasNext()) {
         Object val = loop.next();
 

src/main/java/com/hubspot/jinjava/lib/tag/IfTag.java

@@ -25,6 +25,7 @@
 import com.hubspot.jinjava.interpret.TemplateSyntaxException;
 import com.hubspot.jinjava.tree.Node;
 import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 import com.hubspot.jinjava.util.ObjectTruthValue;
 
 @JinjavaDoc(
@@ -64,7 +65,7 @@ public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
       nextIfElseTagNode = findNextIfElseTagNode(nodeIterator);
     }
 
-    StringBuilder sb = new StringBuilder();
+    LengthLimitingStringBuilder sb = new LengthLimitingStringBuilder(interpreter.getConfig().getMaxOutputSize());
     if (nextIfElseTagNode != null) {
       while (nodeIterator.hasNext()) {
         Node n = nodeIterator.next();

src/main/java/com/hubspot/jinjava/lib/tag/RawTag.java

@@ -7,6 +7,7 @@
 import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.tree.Node;
 import com.hubspot.jinjava.tree.TagNode;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 @JinjavaDoc(
     value = "Process all inner HubL as plain text",
@@ -32,7 +33,7 @@ public String getEndTagName() {
 
   @Override
   public String interpret(TagNode tagNode, JinjavaInterpreter interpreter) {
-    StringBuilder result = new StringBuilder();
+    LengthLimitingStringBuilder result = new LengthLimitingStringBuilder(interpreter.getConfig().getMaxOutputSize());
 
     for (Node n : tagNode.getChildren()) {
       result.append(renderNodeRaw(n));

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

@@ -4,6 +4,7 @@
 import java.util.List;
 
 import com.hubspot.jinjava.interpret.OutputTooBigException;
+import com.hubspot.jinjava.util.LengthLimitingStringBuilder;
 
 public class OutputList {
 
@@ -42,15 +43,9 @@ public List<BlockPlaceholderOutputNode> getBlocks() {
   }
 
   public String getValue() {
-    StringBuilder val = new StringBuilder();
-
-    long valueSize = 0;
+    LengthLimitingStringBuilder val = new LengthLimitingStringBuilder(maxOutputSize);
 
     for (OutputNode node : nodes) {
-      if (maxOutputSize > 0 && valueSize + node.getSize() > maxOutputSize) {
-        throw new OutputTooBigException(maxOutputSize, valueSize + node.getSize());
-      }
-      valueSize += node.getSize();
       val.append(node.getValue());
     }
 

src/main/java/com/hubspot/jinjava/tree/parse/ExpressionToken.java

@@ -32,8 +32,7 @@ public ExpressionToken(String image, int lineNumber) {
 
   @Override
   public String toString() {
-    StringBuilder s = new StringBuilder("{{ ").append(getExpr()).append("}}");
-    return s.toString();
+    return "{{ " + getExpr() + "}}";
   }
 
   @Override

src/main/java/com/hubspot/jinjava/util/LengthLimitingStringBuilder.java

@@ -0,0 +1,62 @@
+package com.hubspot.jinjava.util;
+
+import java.io.Serializable;
+import java.util.stream.IntStream;
+
+import com.hubspot.jinjava.interpret.OutputTooBigException;
+
+public class LengthLimitingStringBuilder implements Serializable, CharSequence {
+
+  private static final long serialVersionUID = -1891922886257965755L;
+
+  private final StringBuilder builder;
+  private long length = 0;
+  private final long maxLength;
+
+  public LengthLimitingStringBuilder(long maxLength) {
+    builder = new StringBuilder();
+    this.maxLength = maxLength;
+  }
+
+  @Override
+  public int length() {
+    return builder.length();
+  }
+
+  @Override
+  public char charAt(int index) {
+    return builder.charAt(index);
+  }
+
+  @Override
+  public CharSequence subSequence(int start, int end) {
+    return builder.subSequence(start, end);
+  }
+
+  @Override
+  public String toString() {
+    return builder.toString();
+  }
+
+  @Override
+  public IntStream chars() {
+    return builder.chars();
+  }
+
+  @Override
+  public IntStream codePoints() {
+    return builder.codePoints();
+  }
+
+  public void append(Object obj) {
+    append(String.valueOf(obj));
+  }
+
+  public void append(String str) {
+    length += str.length();
+    if (maxLength > 0 && length > maxLength) {
+      throw new OutputTooBigException(maxLength, length);
+    }
+    builder.append(str);
+  }
+}

src/test/java/com/hubspot/jinjava/util/LengthLimitingStringBuilderTest.java

@@ -0,0 +1,24 @@
+package com.hubspot.jinjava.util;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import org.junit.Test;
+
+import com.hubspot.jinjava.interpret.OutputTooBigException;
+
+public class LengthLimitingStringBuilderTest {
+
+  @Test
+  public void itLimitsStringLength() throws Exception {
+    LengthLimitingStringBuilder sb = new LengthLimitingStringBuilder(10);
+    sb.append("0123456789");
+    assertThatThrownBy(() -> sb.append("1")).isInstanceOf(OutputTooBigException.class);
+  }
+
+  @Test
+  public void itDoesNotLimitWithZeroLength() throws Exception {
+    LengthLimitingStringBuilder sb = new LengthLimitingStringBuilder(0);
+    sb.append("0123456789");
+  }
+
+}