add ConfigurableExtensionCustomizer (#671)

* add DisallowExtensionCustomizerBuilder * Code cleanup --------- Co-authored-by: Eric Bussieres <[email protected]>
PebbleTemplates · Mar 3, 2024 · 82ad7fc · 82ad7fc
1 parent 14e59e6
commit 82ad7fc
Show file tree

Hide file tree

Showing 3 changed files with 249 additions and 8 deletions.
diff --git a/docs/src/orchid/resources/wiki/guide/customize-defaults.md b/docs/src/orchid/resources/wiki/guide/customize-defaults.md
@@ -35,3 +35,22 @@ The `ExtensionCustomizer` will be used to wrap any Pebble-extension which is pro
 ```java
 PebbleEngine engine = new PebbleEngine.Builder().registerExtensionCustomizer(ExampleOptOuts::new).build();
 ```
+
+### Default implementation of ExtensionCustomizer
+
+The `DisallowExtensionCustomizerBuilder` class can be used to disallow some default functionality, make pebble more controllable.
+
+For example of use, see below:
+
+```java
+PebbleEngine engine = new PebbleEngine.Builder()
+        .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                .disallowedTokenParserTags(singletonList("flush"))
+                .disallowedFunctionKeys(singletonList("max"))
+                .disallowedFilterKeys(singletonList("upper"))
+                .disallowedTestKeys(singletonList("null"))
+                .disallowedBinaryOperatorSymbols(singletonList(">"))
+                .disallowedUnaryOperatorSymbols(singletonList("-"))
+                .build())
+        .build();
+```
diff --git a/...ain/java/io/pebbletemplates/pebble/extension/core/DisallowExtensionCustomizerBuilder.java b/...ain/java/io/pebbletemplates/pebble/extension/core/DisallowExtensionCustomizerBuilder.java
@@ -0,0 +1,136 @@
+package io.pebbletemplates.pebble.extension.core;
+
+import io.pebbletemplates.pebble.extension.Extension;
+import io.pebbletemplates.pebble.extension.ExtensionCustomizer;
+import io.pebbletemplates.pebble.extension.Filter;
+import io.pebbletemplates.pebble.extension.Test;
+import io.pebbletemplates.pebble.operator.BinaryOperator;
+import io.pebbletemplates.pebble.operator.UnaryOperator;
+import io.pebbletemplates.pebble.tokenParser.TokenParser;
+
+import java.util.*;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * @author freshchen
+ * @since 2024/3/2
+ */
+public class DisallowExtensionCustomizerBuilder {
+
+  private Collection<String> disallowedFilterKeys;
+
+  private Collection<String> disallowedTokenParserTags;
+
+  private Collection<String> disallowedFunctionKeys;
+
+  private Collection<String> disallowedBinaryOperatorSymbols;
+
+  private Collection<String> disallowedUnaryOperatorSymbols;
+
+  private Collection<String> disallowedTestKeys;
+
+  public DisallowExtensionCustomizerBuilder disallowedFunctionKeys(Collection<String> disallowedFunctionKeys) {
+    this.disallowedFunctionKeys = disallowedFunctionKeys;
+    return this;
+  }
+
+  public DisallowExtensionCustomizerBuilder disallowedTokenParserTags(Collection<String> disallowedTokenParserTags) {
+    this.disallowedTokenParserTags = disallowedTokenParserTags;
+    return this;
+  }
+
+  public DisallowExtensionCustomizerBuilder disallowedFilterKeys(Collection<String> disallowedFilterKeys) {
+    this.disallowedFilterKeys = disallowedFilterKeys;
+    return this;
+  }
+
+  public DisallowExtensionCustomizerBuilder disallowedUnaryOperatorSymbols(Collection<String> disallowedUnaryOperatorSymbols) {
+    this.disallowedUnaryOperatorSymbols = disallowedUnaryOperatorSymbols;
+    return this;
+  }
+
+  public DisallowExtensionCustomizerBuilder disallowedBinaryOperatorSymbols(Collection<String> disallowedBinaryOperatorSymbols) {
+    this.disallowedBinaryOperatorSymbols = disallowedBinaryOperatorSymbols;
+    return this;
+  }
+
+  public DisallowExtensionCustomizerBuilder disallowedTestKeys(Collection<String> disallowedTestKeys) {
+    this.disallowedTestKeys = disallowedTestKeys;
+    return this;
+  }
+
+  public Function<Extension, ExtensionCustomizer> build() {
+
+    return extension -> new ExtensionCustomizer(extension) {
+
+      @Override
+      public Map<String, Test> getTests() {
+        return this.disallow(super::getTests, DisallowExtensionCustomizerBuilder.this.disallowedTestKeys);
+      }
+
+      @Override
+      public List<UnaryOperator> getUnaryOperators() {
+        return this.disallow(super::getUnaryOperators, DisallowExtensionCustomizerBuilder.this.disallowedUnaryOperatorSymbols, UnaryOperator::getSymbol);
+      }
+
+      @Override
+      public List<BinaryOperator> getBinaryOperators() {
+        return this.disallow(super::getBinaryOperators, DisallowExtensionCustomizerBuilder.this.disallowedBinaryOperatorSymbols, BinaryOperator::getSymbol);
+      }
+
+      @Override
+      public Map<String, io.pebbletemplates.pebble.extension.Function> getFunctions() {
+        return this.disallow(super::getFunctions, DisallowExtensionCustomizerBuilder.this.disallowedFunctionKeys);
+      }
+
+      @Override
+      public Map<String, Filter> getFilters() {
+        return this.disallow(super::getFilters, DisallowExtensionCustomizerBuilder.this.disallowedFilterKeys);
+      }
+
+      @Override
+      public List<TokenParser> getTokenParsers() {
+        return this.disallow(super::getTokenParsers, DisallowExtensionCustomizerBuilder.this.disallowedTokenParserTags, TokenParser::getTag);
+      }
+
+      private <T> List<T> disallow(Supplier<List<T>> superGetter,
+                                   Collection<String> disallowedList,
+                                   Function<T, String> keyGetter) {
+        List<T> superList = superGetter.get();
+        if (disallowedList == null || disallowedList.isEmpty()) {
+          return superList;
+        }
+
+        List<T> result = Optional.ofNullable(superList).map(ArrayList::new)
+                .orElseGet(ArrayList::new);
+
+        disallowedList.stream()
+                .filter(Objects::nonNull)
+                .forEach(v -> result.removeIf(t -> v.equals(keyGetter.apply(t))));
+
+        return result;
+      }
+
+      private <T> Map<String, T> disallow(Supplier<Map<String, T>> superGetter,
+                                          Collection<String> disallowedList) {
+        Map<String, T> superMap = superGetter.get();
+        if (disallowedList == null || disallowedList.isEmpty()) {
+          return superMap;
+        }
+
+        Map<String, T> result = Optional.ofNullable(superMap).map(HashMap::new)
+                .orElseGet(HashMap::new);
+
+        disallowedList.stream()
+                .filter(Objects::nonNull)
+                .forEach(result::remove);
+
+        return result;
+      }
+
+    };
+  }
+
+
+}
diff --git a/pebble/src/test/java/io/pebbletemplates/pebble/extension/ExtensionCustomizerTest.java b/pebble/src/test/java/io/pebbletemplates/pebble/extension/ExtensionCustomizerTest.java
@@ -2,12 +2,13 @@
 
 import io.pebbletemplates.pebble.PebbleEngine;
 import io.pebbletemplates.pebble.error.PebbleException;
+import io.pebbletemplates.pebble.extension.core.DisallowExtensionCustomizerBuilder;
 import io.pebbletemplates.pebble.template.PebbleTemplate;
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
 import java.io.StringWriter;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
@@ -17,26 +18,111 @@
 
 class ExtensionCustomizerTest {
 
-  PebbleEngine pebble;
-
-  @BeforeEach
-  void setUp() {
-    pebble = new PebbleEngine.Builder()
+  @Test
+  void upperFilterCannotBeUsed() throws IOException {
+    PebbleEngine pebble = new PebbleEngine.Builder()
             .registerExtensionCustomizer(RemoveUpperCustomizer::new)
             .build();
+
+    Map<String, Object> obj = new HashMap<>();
+    obj.put("test", "abc");
+    PebbleTemplate template = pebble.getLiteralTemplate("{{ test | upper }}");
+
+    PebbleException exception = assertThrows(PebbleException.class, () -> template.evaluate(new StringWriter(), obj));
+    assertTrue(exception.getMessage().contains("upper"),
+            () -> "Expect upper-Filter to not exist, actual Problem: " + exception.getMessage());
+  }
+
+  @Test
+  void setDisallowedTokenParserTags() {
+    PebbleEngine pebbleEngine = new PebbleEngine.Builder()
+            .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                    .disallowedTokenParserTags(Collections.singletonList("flush"))
+                    .build())
+            .build();
+
+    PebbleException exception = assertThrows(PebbleException.class,
+            () -> pebbleEngine.getLiteralTemplate("{{ k1 }}\n" +
+                    "{% flush %}\n" +
+                    "{{ k2 }}"));
+    assertTrue(exception.getMessage().contains("Unexpected tag name \"flush\""));
   }
 
   @Test
-  void upperFilterCannotBeUsed() throws IOException {
+  void setDisallowedFilters() {
+    PebbleEngine pebbleEngine = new PebbleEngine.Builder()
+            .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                    .disallowedFilterKeys(Collections.singletonList("upper"))
+                    .build())
+            .build();
+
     Map<String, Object> obj = new HashMap<>();
     obj.put("test", "abc");
-    PebbleTemplate template = pebble.getLiteralTemplate("{{ test | upper }}");
+    PebbleTemplate template = pebbleEngine.getLiteralTemplate("{{ test | upper }}");
 
     PebbleException exception = assertThrows(PebbleException.class, () -> template.evaluate(new StringWriter(), obj));
     assertTrue(exception.getMessage().contains("upper"),
             () -> "Expect upper-Filter to not exist, actual Problem: " + exception.getMessage());
   }
 
+  @Test
+  void setDisallowedFunctions() {
+    PebbleEngine pebbleEngine = new PebbleEngine.Builder()
+            .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                    .disallowedFunctionKeys(Collections.singletonList("max"))
+                    .build())
+            .build();
+
+    Map<String, Object> obj = new HashMap<>();
+    obj.put("age", 30);
+    PebbleTemplate template = pebbleEngine.getLiteralTemplate("{{ max(age, 80) }}");
+
+    PebbleException exception = assertThrows(PebbleException.class, () -> template.evaluate(new StringWriter(), obj));
+    assertTrue(exception.getMessage().contains("Function or Macro [max] does not exist"));
+  }
+
+  @Test
+  void setDisallowedBinaryOperatorSymbols() {
+    PebbleEngine pebbleEngine = new PebbleEngine.Builder()
+            .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                    .disallowedBinaryOperatorSymbols(Collections.singletonList(">"))
+                    .build())
+            .build();
+
+    PebbleException exception = assertThrows(PebbleException.class, () -> pebbleEngine.getLiteralTemplate("{% if 10 > 9 %}\n" +
+            "{{ name }}" +
+            "{% endif %}"));
+    assertTrue(exception.getMessage().contains("Unexpected character [>]"));
+  }
+
+  @Test
+  void setDisallowedUnaryOperatorSymbols() throws IOException {
+    PebbleEngine pebbleEngine = new PebbleEngine.Builder()
+            .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                    .disallowedUnaryOperatorSymbols(Collections.singletonList("-"))
+                    .build())
+            .build();
+
+    PebbleException exception = assertThrows(PebbleException.class, () -> pebbleEngine.getLiteralTemplate("{{ -num }}"));
+    assertTrue(exception.getMessage().contains("Unexpected token \"OPERATOR\" of value \"-\""));
+  }
+
+  @Test
+  void setDisallowedTestKeys() throws IOException {
+    PebbleEngine pebbleEngine = new PebbleEngine.Builder()
+            .registerExtensionCustomizer(new DisallowExtensionCustomizerBuilder()
+                    .disallowedTestKeys(Collections.singletonList("null"))
+                    .build())
+            .build();
+
+    PebbleTemplate template = pebbleEngine.getLiteralTemplate("{% if 123 is null %}\n" +
+            "{{ name }}" +
+            "{% endif %}");
+
+    PebbleException exception = assertThrows(PebbleException.class, () -> template.evaluate(new StringWriter()));
+    assertTrue(exception.getMessage().contains("Wrong operand(s) type in conditional expression"));
+  }
+
   private static class RemoveUpperCustomizer extends ExtensionCustomizer {
 
     public RemoveUpperCustomizer(Extension core) {