Wrap classes before validating result

jasmith-hs · jasmith-hs · commit 3be951e64650 · 2026-02-13T13:56:53.000-05:00
diff --git a/src/main/java/com/hubspot/jinjava/el/ext/MethodValidator.java b/src/main/java/com/hubspot/jinjava/el/ext/MethodValidator.java
@@ -1,6 +1,7 @@
 package com.hubspot.jinjava.el.ext;
 
 import com.google.common.collect.ImmutableSet;
+import com.hubspot.jinjava.interpret.JinjavaInterpreter;
 import java.lang.reflect.Method;
 
 public final class MethodValidator {
@@ -39,32 +40,39 @@ public Method validateMethod(Method m) {
     if (m == null) {
       return null;
     }
-    String canonicalDeclaringClassName = m.getDeclaringClass().getCanonicalName();
-    return (
-        allowedMethods.contains(m) ||
-        allowedDeclaredMethodsFromCanonicalClassNames.contains(
-          canonicalDeclaringClassName
-        ) ||
-        allowedDeclaredMethodsFromCanonicalClassPrefixes
-          .stream()
-          .anyMatch(canonicalDeclaringClassName::startsWith)
-      )
-      ? m
-      : null;
+    Class<?> clazz = m.getDeclaringClass();
+    String canonicalClassName = clazz.getCanonicalName();
+    if (
+      allowedMethods.contains(m) ||
+      allowedDeclaredMethodsFromCanonicalClassNames.contains(canonicalClassName) ||
+      allowedDeclaredMethodsFromCanonicalClassPrefixes
+        .stream()
+        .anyMatch(canonicalClassName::startsWith)
+    ) {
+      return m;
+    }
+    return null;
   }
 
   public Object validateResult(Object o) {
-    if (o == null) {
+    Object wrapped = JinjavaInterpreter
+      .getCurrentMaybe()
+      .map(jinjavaInterpreter -> jinjavaInterpreter.wrap(o))
+      .orElse(o);
+
+    if (wrapped == null) {
       return null;
     }
-    String canonicalClassName = o.getClass().getCanonicalName();
-    return (
-        allowedResultCanonicalClassNames.contains(canonicalClassName) ||
-        allowedResultCanonicalClassPrefixes
-          .stream()
-          .anyMatch(canonicalClassName::startsWith)
-      )
-      ? o
-      : null;
+    Class<?> clazz = wrapped.getClass();
+    String canonicalClassName = clazz.getCanonicalName();
+    if (
+      allowedResultCanonicalClassNames.contains(canonicalClassName) ||
+      allowedResultCanonicalClassPrefixes
+        .stream()
+        .anyMatch(canonicalClassName::startsWith)
+    ) {
+      return wrapped;
+    }
+    return null;
   }
 }
diff --git a/src/main/java/com/hubspot/jinjava/el/ext/MethodValidatorConfig.java b/src/main/java/com/hubspot/jinjava/el/ext/MethodValidatorConfig.java
@@ -1,16 +1,28 @@
 package com.hubspot.jinjava.el.ext;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ForwardingCollection;
+import com.google.common.collect.ForwardingList;
+import com.google.common.collect.ForwardingMap;
 import com.google.common.collect.ImmutableSet;
 import com.hubspot.jinjava.JinjavaImmutableStyle;
 import com.hubspot.jinjava.lib.exptest.ExpTest;
 import com.hubspot.jinjava.lib.filter.Filter;
+import com.hubspot.jinjava.objects.DummyObject;
+import com.hubspot.jinjava.objects.Namespace;
+import com.hubspot.jinjava.objects.SafeString;
 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.collections.SnakeCaseAccessibleMap;
+import com.hubspot.jinjava.objects.date.FormattedDate;
+import com.hubspot.jinjava.objects.date.PyishDate;
 import com.hubspot.jinjava.util.ForLoop;
 import java.lang.reflect.Method;
+import java.math.BigDecimal;
+import java.util.AbstractCollection;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.function.Function;
 import java.util.stream.Stream;
@@ -145,6 +157,7 @@ public enum AllowlistGroup {
           float.class,
           boolean.class,
           short.class,
+          BigDecimal.class,
         };
 
         @Override
@@ -160,12 +173,36 @@ String[] allowedDeclaredMethodsFromCanonicalClassPrefixes() {
           return ARRAY;
         }
       },
+      JinjavaObjects {
+        private static final Class<?>[] ARRAY = {
+          PyList.class,
+          PyMap.class,
+          SizeLimitingPyMap.class,
+          SizeLimitingPyList.class,
+          SnakeCaseAccessibleMap.class,
+          FormattedDate.class,
+          PyishDate.class,
+          DummyObject.class,
+          Namespace.class,
+          SafeString.class,
+        };
+
+        @Override
+        Class<?>[] allowedDeclaredMethodsFromClasses() {
+          return ARRAY;
+        }
+      },
       Collections {
         private static final Class<?>[] ARRAY = {
           PyList.class,
           PyMap.class,
           SizeLimitingPyMap.class,
           SizeLimitingPyList.class,
+          ArrayList.class,
+          ForwardingList.class,
+          ForwardingMap.class,
+          ForwardingCollection.class,
+          AbstractCollection.class,
         };
 
         @Override
diff --git a/src/test/java/com/hubspot/jinjava/el/ExpressionResolverTest.java b/src/test/java/com/hubspot/jinjava/el/ExpressionResolverTest.java
@@ -29,6 +29,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -44,6 +45,12 @@ public void setup() {
     jinjava = new Jinjava();
     interpreter = jinjava.newInterpreter();
     context = interpreter.getContext();
+    JinjavaInterpreter.pushCurrent(interpreter);
+  }
+
+  @After
+  public void teardown() {
+    JinjavaInterpreter.popCurrent();
   }
 
   @Test
@@ -467,6 +474,7 @@ public void itBlocksDisabledFunctions() {
     );
 
     String template = "hi {% for i in range(1, 3) %}{{i}} {% endfor %}";
+    JinjavaInterpreter.popCurrent();
 
     String rendered = jinjava.render(template, context);
     assertEquals("hi 1 2 ", rendered);
diff --git a/src/test/java/com/hubspot/jinjava/el/ext/AstDictTest.java b/src/test/java/com/hubspot/jinjava/el/ext/AstDictTest.java
@@ -8,63 +8,105 @@
 import com.hubspot.jinjava.interpret.TemplateError.ErrorType;
 import java.util.Map;
 import java.util.Set;
-import org.junit.Before;
 import org.junit.Test;
 
 public class AstDictTest {
 
-  private JinjavaInterpreter interpreter;
-
-  @Before
-  public void setup() {
-    interpreter = new Jinjava().newInterpreter();
-  }
-
   @Test
   public void itGetsDictValues() {
-    interpreter.getContext().put("foo", ImmutableMap.of("bar", "test"));
-    assertThat(interpreter.resolveELExpression("foo.bar", -1)).isEqualTo("test");
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter.getContext().put("foo", ImmutableMap.of("bar", "test"));
+      assertThat(interpreter.resolveELExpression("foo.bar", -1)).isEqualTo("test");
+    }
   }
 
   @Test
   public void itGetsDictValuesWithEnumKeys() {
-    interpreter.getContext().put("foo", ImmutableMap.of(ErrorType.FATAL, "test"));
-    assertThat(interpreter.resolveELExpression("foo.FATAL", -1)).isEqualTo("test");
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter.getContext().put("foo", ImmutableMap.of(ErrorType.FATAL, "test"));
+      assertThat(interpreter.resolveELExpression("foo.FATAL", -1)).isEqualTo("test");
+    }
   }
 
   @Test
   public void itGetsDictValuesWithEnumKeysUsingToString() {
-    interpreter.getContext().put("foo", ImmutableMap.of(TestEnum.BAR, "test"));
-    assertThat(interpreter.resolveELExpression("foo.barName", -1)).isEqualTo("test");
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter.getContext().put("foo", ImmutableMap.of(TestEnum.BAR, "test"));
+      assertThat(interpreter.resolveELExpression("foo.barName", -1)).isEqualTo("test");
+    }
   }
 
   @Test
   public void itDoesItemsMethodCall() {
-    interpreter.getContext().put("foo", ImmutableMap.of(TestEnum.BAR, "test"));
-    assertThat(interpreter.resolveELExpression("foo.items()", -1))
-      .isInstanceOf(Set.class);
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter.getContext().put("foo", ImmutableMap.of(TestEnum.BAR, "test"));
+      assertThat(interpreter.resolveELExpression("foo.items()", -1))
+        .isInstanceOf(Set.class);
+    }
   }
 
   @Test
   public void itDoesKeysMethodCall() {
-    interpreter.getContext().put("foo", ImmutableMap.of(TestEnum.BAR, "test"));
-    assertThat(interpreter.resolveELExpression("foo.keys()", -1)).isInstanceOf(Set.class);
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter.getContext().put("foo", ImmutableMap.of(TestEnum.BAR, "test"));
+      assertThat(interpreter.resolveELExpression("foo.keys()", -1))
+        .isInstanceOf(Set.class);
+    }
   }
 
   @Test
   public void itHandlesEmptyMaps() {
-    interpreter.getContext().put("foo", ImmutableMap.of());
-    assertThat(interpreter.resolveELExpression("foo.FATAL", -1)).isNull();
-    assertThat(interpreter.getErrors()).isEmpty();
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter.getContext().put("foo", ImmutableMap.of());
+      assertThat(interpreter.resolveELExpression("foo.FATAL", -1)).isNull();
+      assertThat(interpreter.getErrors()).isEmpty();
+    }
   }
 
   @Test
   public void itGetsDictValuesWithEnumKeysInObjects() {
-    interpreter
-      .getContext()
-      .put("test", new TestClass(ImmutableMap.of(ErrorType.FATAL, "test")));
-    assertThat(interpreter.resolveELExpression("test.my_map.FATAL", -1))
-      .isEqualTo("test");
+    try (
+      var a = JinjavaInterpreter
+        .closeablePushCurrent(new Jinjava().newInterpreter())
+        .get()
+    ) {
+      JinjavaInterpreter interpreter = a.value();
+      interpreter
+        .getContext()
+        .put("test", new TestClass(ImmutableMap.of(ErrorType.FATAL, "test")));
+      assertThat(interpreter.resolveELExpression("test.my_map.FATAL", -1))
+        .isEqualTo("test");
+    }
   }
 
   public class TestClass {
diff --git a/src/test/java/com/hubspot/jinjava/el/ext/eager/EagerAstDotTest.java b/src/test/java/com/hubspot/jinjava/el/ext/eager/EagerAstDotTest.java
@@ -6,6 +6,8 @@
 import com.hubspot.jinjava.BaseInterpretingTest;
 import com.hubspot.jinjava.JinjavaConfig;
 import com.hubspot.jinjava.LegacyOverrides;
+import com.hubspot.jinjava.el.ext.MethodValidator;
+import com.hubspot.jinjava.el.ext.MethodValidatorConfig;
 import com.hubspot.jinjava.interpret.Context;
 import com.hubspot.jinjava.interpret.DeferredValue;
 import com.hubspot.jinjava.interpret.DeferredValueException;
@@ -28,6 +30,17 @@ public void setup() {
       .withLegacyOverrides(
         LegacyOverrides.newBuilder().withUsePyishObjectMapper(true).build()
       )
+      .withMethodValidator(
+        MethodValidator.create(
+          MethodValidatorConfig
+            .builder()
+            .addDefaultAllowlistGroups()
+            .addAllowedDeclaredMethodsFromCanonicalClassPrefixes(
+              EagerAstDotTest.class.getCanonicalName()
+            )
+            .build()
+        )
+      )
       .withMaxMacroRecursionDepth(5)
       .withEnableRecursiveMacroCalls(true)
       .build();
diff --git a/src/test/java/com/hubspot/jinjava/el/ext/eager/EagerAstMethodTest.java b/src/test/java/com/hubspot/jinjava/el/ext/eager/EagerAstMethodTest.java
@@ -7,6 +7,8 @@
 import com.hubspot.jinjava.JinjavaConfig;
 import com.hubspot.jinjava.LegacyOverrides;
 import com.hubspot.jinjava.el.ext.DeferredParsingException;
+import com.hubspot.jinjava.el.ext.MethodValidator;
+import com.hubspot.jinjava.el.ext.MethodValidatorConfig;
 import com.hubspot.jinjava.el.ext.eager.EagerAstDotTest.Foo;
 import com.hubspot.jinjava.interpret.Context;
 import com.hubspot.jinjava.interpret.DeferredValue;
@@ -35,6 +37,17 @@ public void setup() {
       .withLegacyOverrides(
         LegacyOverrides.newBuilder().withUsePyishObjectMapper(true).build()
       )
+      .withMethodValidator(
+        MethodValidator.create(
+          MethodValidatorConfig
+            .builder()
+            .addDefaultAllowlistGroups()
+            .addAllowedDeclaredMethodsFromCanonicalClassPrefixes(
+              EagerAstDotTest.class.getCanonicalName()
+            )
+            .build()
+        )
+      )
       .withMaxMacroRecursionDepth(5)
       .withEnableRecursiveMacroCalls(true)
       .build();
diff --git a/src/test/java/com/hubspot/jinjava/lib/expression/EagerExpressionStrategyTest.java b/src/test/java/com/hubspot/jinjava/lib/expression/EagerExpressionStrategyTest.java
diff --git a/src/test/java/com/hubspot/jinjava/tree/ExpressionNodeTest.java b/src/test/java/com/hubspot/jinjava/tree/ExpressionNodeTest.java
