From da3f8b6622c311b8135802e3f1d4c519ee82af39 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Wed, 27 Nov 2024 10:17:56 -0800
Subject: [PATCH] Add CelOptions to disable comprehension

Fixes https://github.com/google/cel-java/issues/484

PiperOrigin-RevId: 700739705
---
 .../src/test/java/dev/cel/bundle/CelImplTest.java  | 14 ++++++++++++++
 .../src/main/java/dev/cel/common/CelOptions.java   | 12 +++++++++++-
 .../java/dev/cel/runtime/DefaultInterpreter.java   |  5 +++--
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/bundle/src/test/java/dev/cel/bundle/CelImplTest.java b/bundle/src/test/java/dev/cel/bundle/CelImplTest.java
index f4647fe0..7b14d825 100644
--- a/bundle/src/test/java/dev/cel/bundle/CelImplTest.java
+++ b/bundle/src/test/java/dev/cel/bundle/CelImplTest.java
@@ -2015,6 +2015,20 @@ public void program_listConcatenationDisabled_throws() throws Exception {
     assertThat(e.getErrorCode()).isEqualTo(CelErrorCode.OVERLOAD_NOT_FOUND);
   }
 
+  @Test
+  public void program_comprehensionDisabled_throws() throws Exception {
+    Cel cel =
+        standardCelBuilderWithMacros()
+            .setOptions(CelOptions.current().enableComprehension(false).build())
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("['foo', 'bar'].map(x, x)").getAst();
+
+    CelEvaluationException e =
+        assertThrows(CelEvaluationException.class, () -> cel.createProgram(ast).eval());
+    assertThat(e).hasMessageThat().contains("Iteration budget exceeded: 0");
+    assertThat(e.getErrorCode()).isEqualTo(CelErrorCode.ITERATION_BUDGET_EXCEEDED);
+  }
+
   @Test
   public void toBuilder_isImmutable() {
     CelBuilder celBuilder = CelFactory.standardCelBuilder();
diff --git a/common/src/main/java/dev/cel/common/CelOptions.java b/common/src/main/java/dev/cel/common/CelOptions.java
index d68afcaf..a1222d14 100644
--- a/common/src/main/java/dev/cel/common/CelOptions.java
+++ b/common/src/main/java/dev/cel/common/CelOptions.java
@@ -115,6 +115,8 @@ public enum ProtoUnsetFieldOptions {
 
   public abstract boolean enableListConcatenation();
 
+  public abstract boolean enableComprehension();
+
   public abstract Builder toBuilder();
 
   public ImmutableSet<ExprFeatures> toExprFeatures() {
@@ -209,7 +211,8 @@ public static Builder newBuilder() {
         .fromProtoUnsetFieldOption(ProtoUnsetFieldOptions.BIND_DEFAULT)
         .enableStringConversion(true)
         .enableStringConcatenation(true)
-        .enableListConcatenation(true);
+        .enableListConcatenation(true)
+        .enableComprehension(true);
   }
 
   /**
@@ -531,6 +534,13 @@ public abstract static class Builder {
      */
     public abstract Builder enableListConcatenation(boolean value);
 
+    /**
+     * Enables comprehension (macros) for the runtime. Setting false has the same effect with
+     * assigning 0 for {@link #comprehensionMaxIterations()}. This option exists to maintain parity
+     * with cel-cpp interpreter options.
+     */
+    public abstract Builder enableComprehension(boolean value);
+
     public abstract CelOptions build();
   }
 }
diff --git a/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java b/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
index 11889e54..86f8fbbd 100644
--- a/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
+++ b/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
@@ -179,9 +179,10 @@ public Object evalTrackingUnknowns(
         Optional<? extends FunctionResolver> functionResolver,
         CelEvaluationListener listener)
         throws InterpreterException {
+      int comprehensionMaxIterations =
+          celOptions.enableComprehension() ? celOptions.comprehensionMaxIterations() : 0;
       ExecutionFrame frame =
-          new ExecutionFrame(
-              listener, resolver, functionResolver, celOptions.comprehensionMaxIterations());
+          new ExecutionFrame(listener, resolver, functionResolver, comprehensionMaxIterations);
       IntermediateResult internalResult = evalInternal(frame, ast.getExpr());
       return internalResult.value();
     }