Add support for importing/exporting common limits to YAML environment configs.

PiperOrigin-RevId: 874676207

Author: jnthntatum

bundle/src/main/java/dev/cel/bundle/CelEnvironment.java

@@ -53,6 +53,7 @@
 import dev.cel.runtime.CelRuntimeLibrary;
 import java.util.Arrays;
 import java.util.Optional;
+import java.util.function.ObjIntConsumer;
 
 /**
  * CelEnvironment is a native representation of a CEL environment for compiler and runtime. This
@@ -74,6 +75,24 @@ public abstract class CelEnvironment {
           "strings", CanonicalCelExtension.STRINGS,
           "comprehensions", CanonicalCelExtension.COMPREHENSIONS);
 
+  private static final ImmutableMap<String, ObjIntConsumer<CelOptions.Builder>> LIMIT_HANDLERS =
+      ImmutableMap.of(
+          "cel.limit.expression_code_points",
+          (options, value) -> options.maxExpressionCodePointSize(value),
+          "cel.limit.parse_error_recovery",
+          (options, value) -> options.maxParseErrorRecoveryLimit(value),
+          "cel.limit.parse_recursion_depth",
+          (options, value) -> options.maxParseRecursionDepth(value));
+
+  private static final ImmutableMap<String, BooleanOptionConsumer> FEATURE_HANDLERS =
+      ImmutableMap.of(
+          "cel.feature.macro_call_tracking",
+          (options, enabled) -> options.populateMacroCalls(enabled),
+          "cel.feature.backtick_escape_syntax",
+          (options, enabled) -> options.enableQuotedIdentifierSyntax(enabled),
+          "cel.feature.cross_type_numeric_comparisons",
+          (options, enabled) -> options.enableHeterogeneousNumericComparisons(enabled));
+
   /** Environment source in textual format (ex: textproto, YAML). */
   public abstract Optional<Source> source();
 
@@ -112,6 +131,9 @@ public abstract class CelEnvironment {
   /** Feature flags to enable in the environment. */
   public abstract ImmutableSet<FeatureFlag> features();
 
+  /** Limits to set in the environment. */
+  public abstract ImmutableSet<Limit> limits();
+
   /** Builder for {@link CelEnvironment}. */
   @AutoValue.Builder
   public abstract static class Builder {
@@ -168,7 +190,14 @@ public Builder setFeatures(FeatureFlag... featureFlags) {
       return setFeatures(ImmutableSet.copyOf(featureFlags));
     }
 
-    public abstract Builder setFeatures(ImmutableSet<FeatureFlag> macros);
+    public abstract Builder setFeatures(ImmutableSet<FeatureFlag> featureFlags);
+
+    @CanIgnoreReturnValue
+    public Builder setLimits(Limit... limits) {
+      return setLimits(ImmutableSet.copyOf(limits));
+    }
+
+    public abstract Builder setLimits(ImmutableSet<Limit> limits);
 
     abstract CelEnvironment autoBuild();
 
@@ -200,13 +229,14 @@ public static Builder newBuilder() {
         .setContainer(CelContainer.ofName(""))
         .setVariables(ImmutableSet.of())
         .setFunctions(ImmutableSet.of())
-        .setFeatures(ImmutableSet.of());
+        .setFeatures(ImmutableSet.of())
+        .setLimits(ImmutableSet.of());
   }
 
   /** Extends the provided {@link CelCompiler} environment with this configuration. */
   public CelCompiler extend(CelCompiler celCompiler, CelOptions celOptions)
       throws CelEnvironmentException {
-    celOptions = applyFeatureFlags(celOptions);
+    celOptions = applyEnvironmentOptions(celOptions);
     try {
       CelTypeProvider celTypeProvider = celCompiler.getTypeProvider();
       CelCompilerBuilder compilerBuilder =
@@ -236,7 +266,7 @@ public CelCompiler extend(CelCompiler celCompiler, CelOptions celOptions)
 
   /** Extends the provided {@link Cel} environment with this configuration. */
   public Cel extend(Cel cel, CelOptions celOptions) throws CelEnvironmentException {
-    celOptions = applyFeatureFlags(celOptions);
+    celOptions = applyEnvironmentOptions(celOptions);
     try {
       // Casting is necessary to only extend the compiler here
       CelCompiler celCompiler = extend((CelCompiler) cel, celOptions);
@@ -249,18 +279,22 @@ public Cel extend(Cel cel, CelOptions celOptions) throws CelEnvironmentException
     }
   }
 
-  private CelOptions applyFeatureFlags(CelOptions celOptions) {
+  private CelOptions applyEnvironmentOptions(CelOptions celOptions) {
     CelOptions.Builder optionsBuilder = celOptions.toBuilder();
     for (FeatureFlag featureFlag : features()) {
-      if (featureFlag.name().equals("cel.feature.macro_call_tracking")) {
-        optionsBuilder.populateMacroCalls(featureFlag.enabled());
-      } else if (featureFlag.name().equals("cel.feature.backtick_escape_syntax")) {
-        optionsBuilder.enableQuotedIdentifierSyntax(featureFlag.enabled());
-      } else if (featureFlag.name().equals("cel.feature.cross_type_numeric_comparisons")) {
-        optionsBuilder.enableHeterogeneousNumericComparisons(featureFlag.enabled());
-      } else {
+      BooleanOptionConsumer consumer = FEATURE_HANDLERS.get(featureFlag.name());
+      if (consumer == null) {
         throw new IllegalArgumentException("Unknown feature flag: " + featureFlag.name());
       }
+      consumer.accept(optionsBuilder, featureFlag.enabled());
+    }
+    for (Limit limit : limits()) {
+      int value = limit.value() < 0 ? -1 : limit.value();
+      ObjIntConsumer<CelOptions.Builder> consumer = LIMIT_HANDLERS.get(limit.name());
+      if (consumer == null) {
+        throw new IllegalArgumentException("Unknown limit: " + limit.name());
+      }
+      consumer.accept(optionsBuilder, value);
     }
     return optionsBuilder.build();
   }
@@ -672,6 +706,25 @@ public static FeatureFlag create(String name, boolean enabled) {
     }
   }
 
+  /**
+   * Represents a configurable limit in the environment.
+   *
+   * <p>A negative value indicates no limit. If not specified, the limit should be set to the
+   * library default.
+   */
+  @AutoValue
+  public abstract static class Limit {
+    /** Normalized name of the limit (e.g. cel.limit.expression_code_points */
+    public abstract String name();
+
+    /** The value of the limit, -1 means no limit. */
+    public abstract int value();
+
+    public static Limit create(String name, int value) {
+      return new AutoValue_CelEnvironment_Limit(name, value);
+    }
+  }
+
   /**
    * Represents a configuration for a canonical CEL extension that can be enabled in the
    * environment.
@@ -995,4 +1048,9 @@ public static OverloadSelector.Builder newBuilder() {
       }
     }
   }
+
+  @FunctionalInterface
+  private static interface BooleanOptionConsumer {
+    void accept(CelOptions.Builder options, boolean value);
+  }
 }

bundle/src/main/java/dev/cel/bundle/CelEnvironmentExporter.java

@@ -221,6 +221,23 @@ private void addOptions(CelEnvironment.Builder envBuilder, CelOptions options) {
       featureFlags.add(CelEnvironment.FeatureFlag.create("cel.feature.macro_call_tracking", true));
     }
     envBuilder.setFeatures(featureFlags.build());
+    ImmutableSet.Builder<CelEnvironment.Limit> limits = ImmutableSet.builder();
+    if (options.maxExpressionCodePointSize() != CelOptions.DEFAULT.maxExpressionCodePointSize()) {
+      limits.add(
+          CelEnvironment.Limit.create(
+              "cel.limit.expression_code_points", options.maxExpressionCodePointSize()));
+    }
+    if (options.maxParseErrorRecoveryLimit() != CelOptions.DEFAULT.maxParseErrorRecoveryLimit()) {
+      limits.add(
+          CelEnvironment.Limit.create(
+              "cel.limit.parse_error_recovery", options.maxParseErrorRecoveryLimit()));
+    }
+    if (options.maxParseRecursionDepth() != CelOptions.DEFAULT.maxParseRecursionDepth()) {
+      limits.add(
+          CelEnvironment.Limit.create(
+              "cel.limit.parse_recursion_depth", options.maxParseRecursionDepth()));
+    }
+    envBuilder.setLimits(limits.build());
   }
 
   /**

bundle/src/main/java/dev/cel/bundle/CelEnvironmentYamlParser.java

@@ -43,6 +43,7 @@
 import dev.cel.common.formats.YamlHelper.YamlNodeType;
 import dev.cel.common.formats.YamlParserContextImpl;
 import dev.cel.common.internal.CelCodePointArray;
+import java.util.Optional;
 import org.jspecify.annotations.Nullable;
 import org.yaml.snakeyaml.DumperOptions.FlowStyle;
 import org.yaml.snakeyaml.nodes.MappingNode;
@@ -188,6 +189,71 @@ private ImmutableSet<CelEnvironment.FeatureFlag> parseFeatures(
     return featureFlags.build();
   }
 
+  private ImmutableSet<CelEnvironment.Limit> parseLimits(ParserContext<Node> ctx, Node node) {
+    long valueId = ctx.collectMetadata(node);
+    if (!validateYamlType(node, YamlNodeType.LIST, YamlNodeType.TEXT)) {
+      ctx.reportError(valueId, "Unsupported limits format");
+    }
+
+    ImmutableSet.Builder<CelEnvironment.Limit> limits = ImmutableSet.builder();
+
+    SequenceNode featureListNode = (SequenceNode) node;
+    for (Node featureMapNode : featureListNode.getValue()) {
+      long featureMapId = ctx.collectMetadata(featureMapNode);
+      if (!assertYamlType(ctx, featureMapId, featureMapNode, YamlNodeType.MAP)) {
+        continue;
+      }
+
+      MappingNode featureMap = (MappingNode) featureMapNode;
+      String name = "";
+      Optional<Integer> value = Optional.empty();
+      // Shorthand syntax for limit: "cel.limit.foo: 1"
+      if (featureMap.getValue().size() == 1) {
+        NodeTuple nodeTuple = featureMap.getValue().get(0);
+        Node keyNode = nodeTuple.getKeyNode();
+        long keyId = ctx.collectMetadata(keyNode);
+        Node valueNode = nodeTuple.getValueNode();
+        String keyName = ((ScalarNode) keyNode).getValue();
+        if (!keyName.equals("name") && !keyName.equals("value")) {
+          limits.add(CelEnvironment.Limit.create(keyName, newInteger(ctx, valueNode)));
+          continue;
+        }
+        // Fall through to check against the long syntax.
+      }
+      // Long syntax for limit:
+      // limits:
+      // - name: cel.limit.foo
+      //   value: 1
+      for (NodeTuple nodeTuple : featureMap.getValue()) {
+        Node keyNode = nodeTuple.getKeyNode();
+        long keyId = ctx.collectMetadata(keyNode);
+        Node valueNode = nodeTuple.getValueNode();
+        String keyName = ((ScalarNode) keyNode).getValue();
+        switch (keyName) {
+          case "name":
+            name = newString(ctx, valueNode);
+            break;
+          case "value":
+            value = Optional.of(newInteger(ctx, valueNode));
+            break;
+          default:
+            ctx.reportError(keyId, String.format("Unsupported limits tag: %s", keyName));
+            break;
+        }
+      }
+      if (name.isEmpty()) {
+        ctx.reportError(featureMapId, "Missing required attribute(s): name");
+        continue;
+      }
+      if (!value.isPresent()) {
+        ctx.reportError(featureMapId, "Missing required attribute(s): value");
+        continue;
+      }
+      limits.add(CelEnvironment.Limit.create(name, value.get()));
+    }
+    return limits.build();
+  }
+
   private ImmutableSet<Alias> parseAliases(ParserContext<Node> ctx, Node node) {
     ImmutableSet.Builder<Alias> aliasSetBuilder = ImmutableSet.builder();
     long valueId = ctx.collectMetadata(node);
@@ -804,6 +870,9 @@ private CelEnvironment.Builder parseConfig(ParserContext<Node> ctx, Node node) {
           case "features":
             builder.setFeatures(parseFeatures(ctx, valueNode));
             break;
+          case "limits":
+            builder.setLimits(parseLimits(ctx, valueNode));
+            break;
           default:
             ctx.reportError(id, "Unknown config tag: " + fieldName);
             // continue handling the rest of the nodes

bundle/src/main/java/dev/cel/bundle/CelEnvironmentYamlSerializer.java

@@ -61,6 +61,7 @@ private CelEnvironmentYamlSerializer() {
     this.multiRepresenters.put(CelEnvironment.Alias.class, new RepresentAlias());
     this.multiRepresenters.put(CelContainer.class, new RepresentContainer());
     this.multiRepresenters.put(CelEnvironment.FeatureFlag.class, new RepresentFeatureFlag());
+    this.multiRepresenters.put(CelEnvironment.Limit.class, new RepresentLimit());
   }
 
   public static String toYaml(CelEnvironment environment) {
@@ -98,6 +99,9 @@ public Node representData(Object data) {
       if (!environment.features().isEmpty()) {
         configMap.put("features", environment.features().asList());
       }
+      if (!environment.limits().isEmpty()) {
+        configMap.put("limits", environment.limits().asList());
+      }
       return represent(configMap.buildOrThrow());
     }
   }
@@ -275,4 +279,17 @@ public Node representData(Object data) {
               .buildOrThrow());
     }
   }
+
+  private final class RepresentLimit implements Represent {
+
+    @Override
+    public Node representData(Object data) {
+      CelEnvironment.Limit limit = (CelEnvironment.Limit) data;
+      return represent(
+          ImmutableMap.builder()
+              .put("name", limit.name())
+              .put("value", limit.value() < 0 ? -1 : limit.value())
+              .buildOrThrow());
+    }
+  }
 }

bundle/src/test/java/dev/cel/bundle/CelEnvironmentExporterTest.java

@@ -260,5 +260,34 @@ public void container() {
     assertThat(container.abbreviations()).containsExactly("foo.Bar", "baz.Qux").inOrder();
     assertThat(container.aliases()).containsAtLeast("nm", "user.name", "id", "user.id").inOrder();
   }
+
+  @Test
+  public void options() {
+    Cel cel =
+        CelFactory.standardCelBuilder()
+            .setOptions(
+                CelOptions.current()
+                    .maxExpressionCodePointSize(100)
+                    .maxParseErrorRecoveryLimit(10)
+                    .maxParseRecursionDepth(10)
+                    .enableQuotedIdentifierSyntax(true)
+                    .enableHeterogeneousNumericComparisons(true)
+                    .populateMacroCalls(true)
+                    .build())
+            .build();
+
+    CelEnvironmentExporter exporter = CelEnvironmentExporter.newBuilder().build();
+    CelEnvironment celEnvironment = exporter.export(cel);
+    assertThat(celEnvironment.features())
+        .containsExactly(
+            CelEnvironment.FeatureFlag.create("cel.feature.backtick_escape_syntax", true),
+            CelEnvironment.FeatureFlag.create("cel.feature.cross_type_numeric_comparisons", true),
+            CelEnvironment.FeatureFlag.create("cel.feature.macro_call_tracking", true));
+    assertThat(celEnvironment.limits())
+        .containsExactly(
+            CelEnvironment.Limit.create("cel.limit.expression_code_points", 100),
+            CelEnvironment.Limit.create("cel.limit.parse_error_recovery", 10),
+            CelEnvironment.Limit.create("cel.limit.parse_recursion_depth", 10));
+  }
 }