Skip to content

Commit d05331e

Browse files
jnthntatumcopybara-github
jnthntatum
authored and
copybara-github
committed
Avoid copying complex target in optMap/optFlatMap
cross ref: cel-expr/cel-go#1387 PiperOrigin-RevId: 954914524
1 parent f502672 commit d05331e

10 files changed

Lines changed: 257 additions & 18 deletions

File tree

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,9 @@ public abstract class CelEnvironment {
8585
"cel.limit.parse_error_recovery",
8686
CelOptions.Builder::maxParseErrorRecoveryLimit,
8787
"cel.limit.parse_recursion_depth",
88-
CelOptions.Builder::maxParseRecursionDepth);
88+
CelOptions.Builder::maxParseRecursionDepth,
89+
"cel.limit.expression_node_count",
90+
CelOptions.Builder::maxParseExpressionNodeCount);
8991

9092
private static final ImmutableMap<String, BooleanOptionConsumer> FEATURE_HANDLERS =
9193
ImmutableMap.of(

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,11 @@ private void addOptions(CelEnvironment.Builder envBuilder, CelOptions options) {
237237
CelEnvironment.Limit.create(
238238
"cel.limit.parse_recursion_depth", options.maxParseRecursionDepth()));
239239
}
240+
if (options.maxParseExpressionNodeCount() != CelOptions.DEFAULT.maxParseExpressionNodeCount()) {
241+
limits.add(
242+
CelEnvironment.Limit.create(
243+
"cel.limit.expression_node_count", options.maxParseExpressionNodeCount()));
244+
}
240245
envBuilder.setLimits(limits.build());
241246
}
242247

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,7 @@ public void options() {
348348
.maxExpressionCodePointSize(100)
349349
.maxParseErrorRecoveryLimit(10)
350350
.maxParseRecursionDepth(10)
351+
.maxParseExpressionNodeCount(500)
351352
.enableQuotedIdentifierSyntax(true)
352353
.enableHeterogeneousNumericComparisons(true)
353354
.populateMacroCalls(true)
@@ -365,6 +366,7 @@ public void options() {
365366
.containsExactly(
366367
CelEnvironment.Limit.create("cel.limit.expression_code_points", 100),
367368
CelEnvironment.Limit.create("cel.limit.parse_error_recovery", 10),
368-
CelEnvironment.Limit.create("cel.limit.parse_recursion_depth", 10));
369+
CelEnvironment.Limit.create("cel.limit.parse_recursion_depth", 10),
370+
CelEnvironment.Limit.create("cel.limit.expression_node_count", 500));
369371
}
370372
}

‎bundle/src/test/java/dev/cel/bundle/CelEnvironmentTest.java‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,8 @@ public void extend_allLimits() throws Exception {
134134
.setLimits(
135135
CelEnvironment.Limit.create("cel.limit.expression_code_points", 20),
136136
CelEnvironment.Limit.create("cel.limit.parse_error_recovery", 10),
137-
CelEnvironment.Limit.create("cel.limit.parse_recursion_depth", 10))
137+
CelEnvironment.Limit.create("cel.limit.parse_recursion_depth", 10),
138+
CelEnvironment.Limit.create("cel.limit.expression_node_count", 500))
138139
.build();
139140

140141
Cel cel =
@@ -147,6 +148,7 @@ public void extend_allLimits() throws Exception {
147148
assertThat(checkerOptions.maxExpressionCodePointSize()).isEqualTo(20);
148149
assertThat(checkerOptions.maxParseErrorRecoveryLimit()).isEqualTo(10);
149150
assertThat(checkerOptions.maxParseRecursionDepth()).isEqualTo(10);
151+
assertThat(checkerOptions.maxParseExpressionNodeCount()).isEqualTo(500);
150152

151153
CelAbstractSyntaxTree ast = cel.compile("1 + 2 + 3 + 4 + 5").getAst();
152154
Long result = (Long) cel.createProgram(ast).eval();
@@ -158,6 +160,27 @@ public void extend_allLimits() throws Exception {
158160
.contains("expression code point size exceeds limit: size: 21, limit 20");
159161
}
160162

163+
@Test
164+
public void extend_expressionNodeCountLimit() throws Exception {
165+
CelEnvironment environment =
166+
CelEnvironment.newBuilder()
167+
.setLimits(CelEnvironment.Limit.create("cel.limit.expression_node_count", 2))
168+
.build();
169+
170+
Cel cel =
171+
environment.extend(
172+
CelFactory.legacyCelBuilder()
173+
.setStandardMacros(CelStandardMacro.STANDARD_MACROS)
174+
.build(),
175+
CelOptions.DEFAULT);
176+
CelOptions checkerOptions = cel.toCheckerBuilder().options();
177+
assertThat(checkerOptions.maxParseExpressionNodeCount()).isEqualTo(2);
178+
179+
CelValidationResult validationResult = cel.compile("1 + 2 + 3");
180+
assertThat(validationResult.hasError()).isTrue();
181+
assertThat(validationResult.getErrorString()).contains("expression node limit exceeded");
182+
}
183+
161184
@Test
162185
public void extend_unsupportedFeatureFlag_throws() throws Exception {
163186
CelEnvironment environment =

‎common/src/main/java/dev/cel/common/CelOptions.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ public enum ProtoUnsetFieldOptions {
6060

6161
public abstract int maxParseRecursionDepth();
6262

63+
public abstract int maxParseExpressionNodeCount();
64+
6365
public abstract boolean populateMacroCalls();
6466

6567
public abstract boolean retainRepeatedUnaryOperators();
@@ -134,6 +136,7 @@ public static Builder newBuilder() {
134136
.maxExpressionCodePointSize(100_000)
135137
.maxParseErrorRecoveryLimit(30)
136138
.maxParseRecursionDepth(250)
139+
.maxParseExpressionNodeCount(100_000)
137140
.populateMacroCalls(false)
138141
.retainRepeatedUnaryOperators(false)
139142
.retainUnbalancedLogicalExpressions(false)
@@ -223,6 +226,14 @@ public abstract static class Builder {
223226
/** Limit the amount of recursion within parse expressions. */
224227
public abstract Builder maxParseRecursionDepth(int value);
225228

229+
/**
230+
* Set a limit on the number of expression nodes in the abstract syntax tree for the expression.
231+
* This prevents cases where macro expansion results in an AST that is larger than expected from
232+
* the source expression. Once exceeded, the parser will record an error and stop expanding
233+
* macros but continue parsing to report other errors.
234+
*/
235+
public abstract Builder maxParseExpressionNodeCount(int value);
236+
226237
/** Populate macro_calls map in source_info with macro calls parsed from the expression. */
227238
public abstract Builder populateMacroCalls(boolean value);
228239

‎extensions/src/main/java/dev/cel/extensions/CelOptionalLibrary.java‎

Lines changed: 67 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,7 @@ static CelExtensionLibrary<CelOptionalLibrary> library() {
297297
public static final CelOptionalLibrary INSTANCE = CelOptionalLibrary.library().latest();
298298

299299
private static final String UNUSED_ITER_VAR = "#unused";
300+
private static final String OPTIONAL_MAP_VAR = "@target";
300301

301302
private final int version;
302303
private final ImmutableSet<CelFunctionDecl> functions;
@@ -524,21 +525,51 @@ private static Optional<CelExpr> expandOptMap(
524525
CelExpr mapExpr = checkNotNull(arguments.get(1));
525526
String varName = varIdent.ident().name();
526527

527-
return Optional.of(
528+
if (target.exprKind().getKind() == CelExpr.ExprKind.Kind.IDENT) {
529+
return Optional.of(
530+
exprFactory.newGlobalCall(
531+
Operator.CONDITIONAL.getFunction(),
532+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
533+
exprFactory.newGlobalCall(
534+
OPTIONAL_OF.getFunction(),
535+
exprFactory.fold(
536+
UNUSED_ITER_VAR,
537+
exprFactory.newList(),
538+
varName,
539+
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
540+
exprFactory.newBoolLiteral(true),
541+
exprFactory.newIdentifier(varName),
542+
mapExpr)),
543+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
544+
}
545+
546+
CelExpr localVar = exprFactory.newIdentifier(OPTIONAL_MAP_VAR);
547+
CelExpr localVarCopy = exprFactory.copy(localVar);
548+
CelExpr conditionalExpr =
528549
exprFactory.newGlobalCall(
529550
Operator.CONDITIONAL.getFunction(),
530-
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
551+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), localVar),
531552
exprFactory.newGlobalCall(
532553
OPTIONAL_OF.getFunction(),
533554
exprFactory.fold(
534555
UNUSED_ITER_VAR,
535556
exprFactory.newList(),
536557
varName,
537-
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
558+
exprFactory.newReceiverCall(VALUE.getFunction(), localVarCopy),
538559
exprFactory.newBoolLiteral(true),
539560
exprFactory.newIdentifier(varName),
540561
mapExpr)),
541-
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
562+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction()));
563+
564+
return Optional.of(
565+
exprFactory.fold(
566+
UNUSED_ITER_VAR,
567+
exprFactory.newList(),
568+
OPTIONAL_MAP_VAR,
569+
localVar,
570+
exprFactory.newBoolLiteral(false),
571+
exprFactory.newIdentifier(OPTIONAL_MAP_VAR),
572+
conditionalExpr));
542573
}
543574

544575
private static Optional<CelExpr> expandOptFlatMap(
@@ -558,19 +589,47 @@ private static Optional<CelExpr> expandOptFlatMap(
558589
CelExpr mapExpr = checkNotNull(arguments.get(1));
559590
String varName = varIdent.ident().name();
560591

561-
return Optional.of(
592+
if (target.exprKind().getKind() == CelExpr.ExprKind.Kind.IDENT) {
593+
return Optional.of(
594+
exprFactory.newGlobalCall(
595+
Operator.CONDITIONAL.getFunction(),
596+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
597+
exprFactory.fold(
598+
UNUSED_ITER_VAR,
599+
exprFactory.newList(),
600+
varName,
601+
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
602+
exprFactory.newBoolLiteral(true),
603+
exprFactory.newIdentifier(varName),
604+
mapExpr),
605+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
606+
}
607+
608+
CelExpr localVar = exprFactory.newIdentifier(OPTIONAL_MAP_VAR);
609+
CelExpr localVarCopy = exprFactory.copy(localVar);
610+
CelExpr conditionalExpr =
562611
exprFactory.newGlobalCall(
563612
Operator.CONDITIONAL.getFunction(),
564-
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
613+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), localVar),
565614
exprFactory.fold(
566615
UNUSED_ITER_VAR,
567616
exprFactory.newList(),
568617
varName,
569-
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
618+
exprFactory.newReceiverCall(VALUE.getFunction(), localVarCopy),
570619
exprFactory.newBoolLiteral(true),
571620
exprFactory.newIdentifier(varName),
572621
mapExpr),
573-
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
622+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction()));
623+
624+
return Optional.of(
625+
exprFactory.fold(
626+
UNUSED_ITER_VAR,
627+
exprFactory.newList(),
628+
OPTIONAL_MAP_VAR,
629+
target,
630+
exprFactory.newBoolLiteral(false),
631+
exprFactory.newIdentifier(OPTIONAL_MAP_VAR),
632+
conditionalExpr));
574633
}
575634

576635
private static Object indexOptionalMap(

‎extensions/src/test/java/dev/cel/extensions/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ java_library(
1717
"//common:compiler_common",
1818
"//common:container",
1919
"//common:options",
20+
"//common/ast",
2021
"//common/exceptions:attribute_not_found",
2122
"//common/exceptions:divide_by_zero",
2223
"//common/exceptions:index_out_of_bounds",

‎extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import dev.cel.common.CelOverloadDecl;
3535
import dev.cel.common.CelValidationException;
3636
import dev.cel.common.CelVarDecl;
37+
import dev.cel.common.ast.CelExpr;
3738
import dev.cel.common.types.CelType;
3839
import dev.cel.common.types.ListType;
3940
import dev.cel.common.types.MapType;
@@ -1571,6 +1572,68 @@ public void optionalFlatMapMacro_receiverHasValue_returnsOptionalValue() throws
15711572
assertThat(result).hasValue(43L);
15721573
}
15731574

1575+
@Test
1576+
public void optionalMapMacro_simpleTarget_notWrappedInComprehension() throws Exception {
1577+
Cel cel =
1578+
newCelBuilder()
1579+
.addVar("x", OptionalType.create(SimpleType.INT))
1580+
.setResultType(OptionalType.create(SimpleType.INT))
1581+
.build();
1582+
CelAbstractSyntaxTree ast = compile(cel, "x.optMap(y, y + 1)");
1583+
1584+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.CALL);
1585+
}
1586+
1587+
@Test
1588+
public void optionalMapMacro_complexTarget_astWrappedInComprehension() throws Exception {
1589+
Cel cel =
1590+
newCelBuilder()
1591+
.setResultType(OptionalType.create(SimpleType.INT))
1592+
.addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
1593+
.build();
1594+
CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32.optMap(y, y + 1)");
1595+
1596+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.COMPREHENSION);
1597+
assertThat(ast.getExpr().comprehension().accuVar()).isEqualTo("@target");
1598+
1599+
Optional<Long> result =
1600+
(Optional<Long>)
1601+
cel.createProgram(ast)
1602+
.eval(ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()));
1603+
assertThat(result).hasValue(43L);
1604+
}
1605+
1606+
@Test
1607+
public void optionalFlatMapMacro_simpleTarget_notWrappedInComprehension() throws Exception {
1608+
Cel cel =
1609+
newCelBuilder()
1610+
.addVar("x", OptionalType.create(SimpleType.INT))
1611+
.setResultType(OptionalType.create(SimpleType.INT))
1612+
.build();
1613+
CelAbstractSyntaxTree ast = compile(cel, "x.optFlatMap(y, optional.of(y + 1))");
1614+
1615+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.CALL);
1616+
}
1617+
1618+
@Test
1619+
public void optionalFlatMapMacro_complexTarget_astWrappedInComprehension() throws Exception {
1620+
Cel cel =
1621+
newCelBuilder()
1622+
.setResultType(OptionalType.create(SimpleType.INT))
1623+
.addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
1624+
.build();
1625+
CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32.optFlatMap(y, optional.of(y + 1))");
1626+
1627+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.COMPREHENSION);
1628+
assertThat(ast.getExpr().comprehension().accuVar()).isEqualTo("@target");
1629+
1630+
Optional<Long> result =
1631+
(Optional<Long>)
1632+
cel.createProgram(ast)
1633+
.eval(ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()));
1634+
assertThat(result).hasValue(43L);
1635+
}
1636+
15741637
@Test
15751638
public void optionalFlatMapMacro_withOptionalOfNonZeroValue_optionalEmptyWhenValueIsZero()
15761639
throws Exception {

0 commit comments

Comments
 (0)