add size limited pymaps and pylists

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

@@ -46,15 +46,17 @@ public class JinjavaConfig {
   private final boolean enableRecursiveMacroCalls;
   private final int maxMacroRecursionDepth;
 
-  private Map<Context.Library, Set<String>> disabled;
+  private final Map<Context.Library, Set<String>> disabled;
   private final boolean failOnUnknownTokens;
   private final boolean nestedInterpretationEnabled;
   private final RandomNumberGeneratorStrategy randomNumberGenerator;
   private final boolean validationMode;
   private final long maxStringLength;
-  private InterpreterFactory interpreterFactory;
+  private final int maxListSize;
+  private final int maxMapSize;
+  private final InterpreterFactory interpreterFactory;
   private TokenScannerSymbols tokenScannerSymbols;
-  private ELResolver elResolver;
+  private final ELResolver elResolver;
   private final boolean iterateOverMapKeys;
   private final boolean preserveForFinalPass;
 
@@ -101,6 +103,8 @@ private JinjavaConfig(Builder builder) {
     randomNumberGenerator = builder.randomNumberGeneratorStrategy;
     validationMode = builder.validationMode;
     maxStringLength = builder.maxStringLength;
+    maxListSize = builder.maxListSize;
+    maxMapSize = builder.maxMapSize;
     interpreterFactory = builder.interpreterFactory;
     tokenScannerSymbols = builder.tokenScannerSymbols;
     elResolver = builder.elResolver;
@@ -128,6 +132,14 @@ public long getMaxOutputSize() {
     return maxOutputSize;
   }
 
+  public int getMaxListSize() {
+    return maxListSize;
+  }
+
+  public int getMaxMapSize() {
+    return maxMapSize;
+  }
+
   public RandomNumberGeneratorStrategy getRandomNumberGeneratorStrategy() {
     return randomNumberGenerator;
   }
@@ -216,6 +228,8 @@ public static class Builder {
     private ELResolver elResolver = JinjavaInterpreterResolver.DEFAULT_RESOLVER_READ_ONLY;
     private boolean iterateOverMapKeys;
     private boolean preserveForFinalPass;
+    private int maxListSize = Integer.MAX_VALUE;
+    private int maxMapSize = Integer.MAX_VALUE;
 
     private Builder() {}
 
@@ -309,6 +323,16 @@ public Builder withMaxStringLength(long maxStringLength) {
       return this;
     }
 
+    public Builder withMaxListSize(int maxListSize) {
+      this.maxListSize = maxListSize;
+      return this;
+    }
+
+    public Builder withMaxMapSize(int maxMapSize) {
+      this.maxMapSize = maxMapSize;
+      return this;
+    }
+
     public Builder withInterperterFactory(InterpreterFactory interperterFactory) {
       this.interpreterFactory = interperterFactory;
       return this;

src/main/java/com/hubspot/jinjava/el/JinjavaInterpreterResolver.java

@@ -19,8 +19,8 @@
 import com.hubspot.jinjava.interpret.TemplateError.ErrorType;
 import com.hubspot.jinjava.interpret.errorcategory.BasicTemplateErrorCategory;
 import com.hubspot.jinjava.objects.PyWrapper;
-import com.hubspot.jinjava.objects.collections.PyList;
-import com.hubspot.jinjava.objects.collections.PyMap;
+import com.hubspot.jinjava.objects.collections.SizeLimitingPyList;
+import com.hubspot.jinjava.objects.collections.SizeLimitingPyMap;
 import com.hubspot.jinjava.objects.date.FormattedDate;
 import com.hubspot.jinjava.objects.date.PyishDate;
 import com.hubspot.jinjava.objects.date.StrftimeFormatter;
@@ -295,11 +295,17 @@ Object wrap(Object value) {
     }
 
     if (List.class.isAssignableFrom(value.getClass())) {
-      return new PyList((List<Object>) value);
+      return new SizeLimitingPyList(
+        (List<Object>) value,
+        interpreter.getConfig().getMaxListSize()
+      );
     }
     if (Map.class.isAssignableFrom(value.getClass())) {
       // FIXME: ensure keys are actually strings, if not, convert them
-      return new PyMap((Map<String, Object>) value);
+      return new SizeLimitingPyMap(
+        (Map<String, Object>) value,
+        interpreter.getConfig().getMaxMapSize()
+      );
     }
 
     if (Date.class.isAssignableFrom(value.getClass())) {

src/main/java/com/hubspot/jinjava/el/ext/AstDict.java

@@ -1,7 +1,8 @@
 package com.hubspot.jinjava.el.ext;
 
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.interpret.TemplateStateException;
-import com.hubspot.jinjava.objects.collections.PyMap;
+import com.hubspot.jinjava.objects.collections.SizeLimitingPyMap;
 import de.odysseus.el.tree.Bindings;
 import de.odysseus.el.tree.impl.ast.AstIdentifier;
 import de.odysseus.el.tree.impl.ast.AstLiteral;
@@ -39,7 +40,11 @@ public Object eval(Bindings bindings, ELContext context) {
       resolved.put(key, entry.getValue().eval(bindings, context));
     }
 
-    return new PyMap(resolved);
+    JinjavaInterpreter interpreter = (JinjavaInterpreter) context
+      .getELResolver()
+      .getValue(context, null, ExtendedParser.INTERPRETER);
+
+    return new SizeLimitingPyMap(resolved, interpreter.getConfig().getMaxMapSize());
   }
 
   @Override
@@ -54,10 +59,7 @@ public String toString() {
     StringBuilder s = new StringBuilder("{");
 
     for (Map.Entry<AstNode, AstNode> entry : dict.entrySet()) {
-      s
-        .append(Objects.toString(entry.getKey()))
-        .append(":")
-        .append(Objects.toString(entry.getValue()));
+      s.append(entry.getKey()).append(":").append(entry.getValue());
     }
 
     return s.append("}").toString();

src/main/java/com/hubspot/jinjava/el/ext/AstList.java

@@ -1,6 +1,7 @@
 package com.hubspot.jinjava.el.ext;
 
-import com.hubspot.jinjava.objects.collections.PyList;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
+import com.hubspot.jinjava.objects.collections.SizeLimitingPyList;
 import de.odysseus.el.tree.Bindings;
 import de.odysseus.el.tree.impl.ast.AstLiteral;
 import de.odysseus.el.tree.impl.ast.AstParameters;
@@ -24,7 +25,11 @@ public Object eval(Bindings bindings, ELContext context) {
       list.add(elements.getChild(i).eval(bindings, context));
     }
 
-    return new PyList(list);
+    JinjavaInterpreter interpreter = (JinjavaInterpreter) context
+      .getELResolver()
+      .getValue(context, null, ExtendedParser.INTERPRETER);
+
+    return new SizeLimitingPyList(list, interpreter.getConfig().getMaxListSize());
   }
 
   @Override

src/main/java/com/hubspot/jinjava/el/ext/AstRangeBracket.java

@@ -1,7 +1,9 @@
 package com.hubspot.jinjava.el.ext;
 
 import com.google.common.collect.Iterables;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import com.hubspot.jinjava.objects.collections.PyList;
+import com.hubspot.jinjava.objects.collections.SizeLimitingPyList;
 import de.odysseus.el.misc.LocalMessages;
 import de.odysseus.el.tree.Bindings;
 import de.odysseus.el.tree.impl.ast.AstBracket;
@@ -76,7 +78,14 @@ public Object eval(Bindings bindings, ELContext context) {
     int startNum = ((Number) start).intValue();
     int endNum = ((Number) end).intValue();
 
-    PyList result = new PyList(new ArrayList<>());
+    JinjavaInterpreter interpreter = (JinjavaInterpreter) context
+      .getELResolver()
+      .getValue(context, null, ExtendedParser.INTERPRETER);
+
+    PyList result = new SizeLimitingPyList(
+      new ArrayList<>(),
+      interpreter.getConfig().getMaxListSize()
+    );
     int index = 0;
 
     // Handle negative indices.

src/main/java/com/hubspot/jinjava/el/ext/ExtendedParser.java

@@ -308,7 +308,6 @@ protected AstNode literal() throws ScanException, ParseException {
     switch (getToken().getSymbol()) {
       case LBRACK:
         v = new AstList(params(LBRACK, RBRACK));
-
         break;
       case LPAREN:
         v = new AstTuple(params());

src/main/java/com/hubspot/jinjava/objects/collections/PyMap.java

@@ -6,7 +6,7 @@
 import java.util.Set;
 
 public class PyMap extends ForwardingMap<String, Object> implements PyWrapper {
-  private Map<String, Object> map;
+  private final Map<String, Object> map;
 
   public PyMap(Map<String, Object> map) {
     this.map = map;

src/main/java/com/hubspot/jinjava/objects/collections/SizeLimitingPyList.java

@@ -0,0 +1,90 @@
+package com.hubspot.jinjava.objects.collections;
+
+import com.hubspot.jinjava.interpret.IndexOutOfRangeException;
+import com.hubspot.jinjava.objects.PyWrapper;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+public class SizeLimitingPyList extends PyList implements PyWrapper {
+  private int maxSize;
+
+  private SizeLimitingPyList(List<Object> list) {
+    super(list);
+  }
+
+  public SizeLimitingPyList(List<Object> list, int maxSize) {
+    super(list);
+    if (maxSize <= 0) {
+      throw new IllegalArgumentException("maxSize must be >= 1");
+    }
+
+    this.maxSize = maxSize;
+    if (list.size() >= maxSize) {
+      throw createOutOfRangeException(list.size());
+    }
+  }
+
+  @Override
+  public boolean append(Object e) {
+    if (size() >= maxSize) {
+      throw createOutOfRangeException(size() + 1);
+    }
+    return super.append(e);
+  }
+
+  @Override
+  public void insert(int i, Object e) {
+    if (size() >= maxSize) {
+      throw createOutOfRangeException(size() + 1);
+    }
+    super.insert(i, e);
+  }
+
+  @Override
+  public boolean add(Object element) {
+    if (size() >= maxSize) {
+      throw createOutOfRangeException(size() + 1);
+    }
+    return super.add(element);
+  }
+
+  @Override
+  public void add(int index, Object element) {
+    if (size() >= maxSize) {
+      throw createOutOfRangeException(size() + 1);
+    }
+    super.add(index, element);
+  }
+
+  @Override
+  public boolean addAll(int index, Collection<?> elements) {
+    if (size() + elements.size() >= maxSize) {
+      throw createOutOfRangeException(size() + elements.size());
+    }
+    return super.addAll(index, elements);
+  }
+
+  @Override
+  public boolean addAll(Collection<?> collection) {
+    if (size() + collection.size() >= maxSize) {
+      throw createOutOfRangeException(size() + collection.size());
+    }
+    return super.addAll(collection);
+  }
+
+  @Override
+  public PyList copy() {
+    return new SizeLimitingPyList(new ArrayList<>(delegate()));
+  }
+
+  IndexOutOfRangeException createOutOfRangeException(int index) {
+    return new IndexOutOfRangeException(
+      String.format(
+        "Index %d is out of range for list of maximum size %d",
+        index,
+        maxSize
+      )
+    );
+  }
+}

src/main/java/com/hubspot/jinjava/objects/collections/SizeLimitingPyMap.java

@@ -0,0 +1,53 @@
+package com.hubspot.jinjava.objects.collections;
+
+import com.hubspot.jinjava.interpret.IndexOutOfRangeException;
+import com.hubspot.jinjava.objects.PyWrapper;
+import java.util.HashSet;
+import java.util.Map;
+
+public class SizeLimitingPyMap extends PyMap implements PyWrapper {
+  private int maxSize;
+
+  private SizeLimitingPyMap(Map<String, Object> map) {
+    super(map);
+  }
+
+  public SizeLimitingPyMap(Map<String, Object> map, int maxSize) {
+    super(map);
+    if (maxSize <= 0) {
+      throw new IllegalArgumentException("maxSize must be >= 1");
+    }
+
+    this.maxSize = maxSize;
+    if (map.size() > maxSize) {
+      throw createOutOfRangeException(map.size());
+    }
+  }
+
+  @Override
+  public Object put(String s, Object o) {
+    if (delegate().size() >= maxSize && !delegate().containsKey(s)) {
+      throw createOutOfRangeException(delegate().size() + 1);
+    }
+
+    return super.put(s, o);
+  }
+
+  @Override
+  public void putAll(Map<? extends String, ?> m) {
+    HashSet<String> keys = new HashSet<>(delegate().keySet());
+    int newKeys = (int) m.keySet().stream().filter(k -> !keys.contains(k)).count();
+
+    if (newKeys + delegate().size() > maxSize) {
+      throw createOutOfRangeException(newKeys + delegate().size());
+    }
+
+    super.putAll(m);
+  }
+
+  IndexOutOfRangeException createOutOfRangeException(int index) {
+    return new IndexOutOfRangeException(
+      String.format("%d is out of range for map of maximum size %d", index, maxSize)
+    );
+  }
+}

src/test/java/com/hubspot/jinjava/el/ext/CollectionMembershipOperatorTest.java

@@ -7,7 +7,6 @@
 import org.junit.Before;
 import org.junit.Test;
 
-@SuppressWarnings("unchecked")
 public class CollectionMembershipOperatorTest {
   private JinjavaInterpreter interpreter;