Make iteration limit configurable for rule composer

PiperOrigin-RevId: 651551461

Author: l46kok

policy/src/main/java/dev/cel/policy/CelPolicyCompilerBuilder.java

@@ -20,9 +20,17 @@
 /** Interface for building an instance of {@link CelPolicyCompiler} */
 public interface CelPolicyCompilerBuilder {
 
+  /** Sets the prefix for the policy variables. Default is `variables.`. */
   @CanIgnoreReturnValue
   CelPolicyCompilerBuilder setVariablesPrefix(String prefix);
 
+  /**
+   * Limit the number of iteration while composing rules into a single AST. An exception is thrown
+   * if the iteration count exceeds the set value.
+   */
+  @CanIgnoreReturnValue
+  CelPolicyCompilerBuilder setIterationLimit(int iterationLimit);
+
   @CheckReturnValue
   CelPolicyCompiler build();
 }

policy/src/main/java/dev/cel/policy/CelPolicyCompilerImpl.java

@@ -44,8 +44,10 @@
 /** Package-private implementation for policy compiler. */
 final class CelPolicyCompilerImpl implements CelPolicyCompiler {
   private static final String DEFAULT_VARIABLE_PREFIX = "variables.";
+  private static final int DEFAULT_ITERATION_LIMIT = 1000;
   private final Cel cel;
   private final String variablesPrefix;
+  private final int iterationLimit;
 
   @Override
   public CelAbstractSyntaxTree compile(CelPolicy policy) throws CelPolicyValidationException {
@@ -59,7 +61,10 @@ public CelAbstractSyntaxTree compile(CelPolicy policy) throws CelPolicyValidatio
         CelOptimizerFactory.standardCelOptimizerBuilder(compiledRule.cel())
             .addAstOptimizers(
                 RuleComposer.newInstance(
-                    compiledRule, compilerContext.newVariableDeclarations, variablesPrefix))
+                    compiledRule,
+                    compilerContext.newVariableDeclarations,
+                    variablesPrefix,
+                    iterationLimit))
             .build();
 
     CelAbstractSyntaxTree ast;
@@ -194,6 +199,7 @@ private CompilerContext(CelPolicySource celPolicySource) {
   static final class Builder implements CelPolicyCompilerBuilder {
     private final Cel cel;
     private String variablesPrefix;
+    private int iterationLimit;
 
     private Builder(Cel cel) {
       this.cel = cel;
@@ -206,18 +212,28 @@ public Builder setVariablesPrefix(String prefix) {
       return this;
     }
 
+    @Override
+    @CanIgnoreReturnValue
+    public Builder setIterationLimit(int iterationLimit) {
+      this.iterationLimit = iterationLimit;
+      return this;
+    }
+
     @Override
     public CelPolicyCompiler build() {
-      return new CelPolicyCompilerImpl(cel, this.variablesPrefix);
+      return new CelPolicyCompilerImpl(cel, this.variablesPrefix, this.iterationLimit);
     }
   }
 
   static Builder newBuilder(Cel cel) {
-    return new Builder(cel).setVariablesPrefix(DEFAULT_VARIABLE_PREFIX);
+    return new Builder(cel)
+        .setVariablesPrefix(DEFAULT_VARIABLE_PREFIX)
+        .setIterationLimit(DEFAULT_ITERATION_LIMIT);
   }
 
-  private CelPolicyCompilerImpl(Cel cel, String variablesPrefix) {
+  private CelPolicyCompilerImpl(Cel cel, String variablesPrefix, int iterationLimit) {
     this.cel = checkNotNull(cel);
     this.variablesPrefix = checkNotNull(variablesPrefix);
+    this.iterationLimit = iterationLimit;
   }
 }

policy/src/main/java/dev/cel/policy/RuleComposer.java

@@ -34,7 +34,6 @@
 
 /** Package-private class for composing various rules into a single expression using optimizer. */
 final class RuleComposer implements CelAstOptimizer {
-  private static final int AST_MUTATOR_ITERATION_LIMIT = 1000;
   private final CelCompiledRule compiledRule;
   private final ImmutableList<CelVarDecl> newVarDecls;
   private final String variablePrefix;
@@ -121,15 +120,21 @@ private RuleOptimizationResult optimizeRule(CelCompiledRule compiledRule) {
   }
 
   static RuleComposer newInstance(
-      CelCompiledRule compiledRule, List<CelVarDecl> newVarDecls, String variablePrefix) {
-    return new RuleComposer(compiledRule, newVarDecls, variablePrefix);
+      CelCompiledRule compiledRule,
+      List<CelVarDecl> newVarDecls,
+      String variablePrefix,
+      int iterationLimit) {
+    return new RuleComposer(compiledRule, newVarDecls, variablePrefix, iterationLimit);
   }
 
   private RuleComposer(
-      CelCompiledRule compiledRule, List<CelVarDecl> newVarDecls, String variablePrefix) {
+      CelCompiledRule compiledRule,
+      List<CelVarDecl> newVarDecls,
+      String variablePrefix,
+      int iterationLimit) {
     this.compiledRule = checkNotNull(compiledRule);
     this.newVarDecls = ImmutableList.copyOf(checkNotNull(newVarDecls));
     this.variablePrefix = variablePrefix;
-    this.astMutator = AstMutator.newInstance(AST_MUTATOR_ITERATION_LIMIT);
+    this.astMutator = AstMutator.newInstance(iterationLimit);
   }
 }