Skip to content

Commit 21b84bf

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Add support for optional field pruning in verifier
PiperOrigin-RevId: 951556122
1 parent 077d9a9 commit 21b84bf

3 files changed

Lines changed: 83 additions & 18 deletions

File tree

‎verifier/src/main/java/dev/cel/verifier/CelAstAlphaHasher.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,10 @@ private static void hashAst(CelExpr expr, @Nullable Scope scope, HasherContext c
121121
for (CelExpr elem : expr.list().elements()) {
122122
hashAst(elem, scope, context);
123123
}
124+
context.hasher.putInt(expr.list().optionalIndices().size());
125+
for (int optIndex : expr.list().optionalIndices()) {
126+
context.hasher.putInt(optIndex);
127+
}
124128
break;
125129
case STRUCT:
126130
context.hasher.putString(expr.struct().messageName(), UTF_8);

‎verifier/src/main/java/dev/cel/verifier/CelAstToZ3Translator.java‎

Lines changed: 65 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -284,11 +284,20 @@ private TranslatedValue translateList(CelExpr celExpr, CelAbstractSyntaxTree ast
284284
// check to a trivial identity check (e.g., `list_ref_0 == list_ref_0`).
285285
if (listRef == null) {
286286
SeqExpr seq = ctx.mkEmptySeq(ctx.mkSeqSort(typeSystem.celValueSort()));
287-
for (CelExpr element : createList.elements()) {
287+
ImmutableList<Integer> optionalIndices = createList.optionalIndices();
288+
for (int i = 0; i < createList.elements().size(); i++) {
289+
CelExpr element = createList.elements().get(i);
288290
TranslatedValue elem = translateExpr(element, ast);
289291
elementsTv.add(elem);
290292

291-
seq = typeSystem.mkConcatSafe(seq, ctx.mkUnit(elem.z3Expr()));
293+
if (optionalIndices.contains(i)) {
294+
Expr<?> optRef = typeSystem.getOptionalRef(elem.z3Expr());
295+
BoolExpr hasVal = typeSystem.optHasValue(optRef);
296+
Expr<?> val = typeSystem.getOptionalValue(optRef);
297+
seq = (SeqExpr) ctx.mkITE(hasVal, typeSystem.mkConcatSafe(seq, ctx.mkUnit(val)), seq);
298+
} else {
299+
seq = typeSystem.mkConcatSafe(seq, ctx.mkUnit(elem.z3Expr()));
300+
}
292301
}
293302
listRef = typeSystem.mkListRefConst(LIST_REF_PREFIX);
294303
typeConstraints.add(ctx.mkEq(typeSystem.getSeq(listRef), seq));
@@ -297,7 +306,9 @@ private TranslatedValue translateList(CelExpr celExpr, CelAbstractSyntaxTree ast
297306
}
298307

299308
Expr<?> result = typeSystem.wrapList(listRef);
300-
return TranslatedValue.propagateStrict(ctx, typeSystem, result, celExpr, elementsTv);
309+
BoolExpr baseTaint = ctx.mkFalse();
310+
return TranslatedValue.propagateStrict(
311+
ctx, typeSystem, result, Optional.of(celExpr), baseTaint, elementsTv);
301312
}
302313

303314
private TranslatedValue translateMap(CelExpr celExpr, CelAbstractSyntaxTree ast) {
@@ -319,19 +330,37 @@ private TranslatedValue translateMap(CelExpr celExpr, CelAbstractSyntaxTree ast)
319330
elementsTv.add(valueTv);
320331

321332
BoolExpr keyAlreadyPresent = (BoolExpr) ctx.mkSelect(mapPresence, key);
322-
keysSeq =
323-
ctx.mkITE(keyAlreadyPresent, keysSeq, typeSystem.mkConcatSafe(keysSeq, ctx.mkUnit(key)));
324-
325-
mapValues = ctx.mkStore(mapValues, key, value);
326-
mapPresence = ctx.mkStore(mapPresence, key, ctx.mkTrue());
333+
if (entryAst.optionalEntry()) {
334+
Expr<?> optRef = typeSystem.getOptionalRef(value);
335+
BoolExpr hasVal = typeSystem.optHasValue(optRef);
336+
Expr<?> optVal = typeSystem.getOptionalValue(optRef);
337+
338+
keysSeq =
339+
ctx.mkITE(
340+
hasVal,
341+
ctx.mkITE(
342+
keyAlreadyPresent, keysSeq, typeSystem.mkConcatSafe(keysSeq, ctx.mkUnit(key))),
343+
keysSeq);
344+
mapValues = (ArrayExpr) ctx.mkITE(hasVal, ctx.mkStore(mapValues, key, optVal), mapValues);
345+
mapPresence =
346+
(ArrayExpr) ctx.mkITE(hasVal, ctx.mkStore(mapPresence, key, ctx.mkTrue()), mapPresence);
347+
} else {
348+
keysSeq =
349+
ctx.mkITE(
350+
keyAlreadyPresent, keysSeq, typeSystem.mkConcatSafe(keysSeq, ctx.mkUnit(key)));
351+
mapValues = ctx.mkStore(mapValues, key, value);
352+
mapPresence = ctx.mkStore(mapPresence, key, ctx.mkTrue());
353+
}
327354
}
328355

329356
typeConstraints.add(ctx.mkEq(typeSystem.getMapValues(mapRef), mapValues));
330357
typeConstraints.add(ctx.mkEq(typeSystem.getMapPresence(mapRef), mapPresence));
331358
typeConstraints.add(ctx.mkEq(typeSystem.getMapKeys(mapRef), keysSeq));
332359

333360
Expr<?> result = typeSystem.wrapMap(mapRef);
334-
return TranslatedValue.propagateStrict(ctx, typeSystem, result, celExpr, elementsTv);
361+
BoolExpr baseTaint = ctx.mkFalse();
362+
return TranslatedValue.propagateStrict(
363+
ctx, typeSystem, result, Optional.of(celExpr), baseTaint, elementsTv);
335364
}
336365

337366
private TranslatedValue translateStruct(CelExpr celExpr, CelAbstractSyntaxTree ast) {
@@ -379,15 +408,29 @@ private TranslatedValue translateStruct(CelExpr celExpr, CelAbstractSyntaxTree a
379408
// (`msg1 == msg2`) to work without using quantifiers (which avoids MBQI loops).
380409
// Because proto3 singular primitives do not have field presence, we also skip setting
381410
// `msgPresence`.
411+
Expr<?> effectiveValue;
412+
BoolExpr isPresent;
413+
if (entryAst.optionalEntry()) {
414+
Expr<?> optRef = typeSystem.getOptionalRef(value);
415+
isPresent = typeSystem.optHasValue(optRef);
416+
effectiveValue = typeSystem.getOptionalValue(optRef);
417+
} else {
418+
isPresent = ctx.mkTrue();
419+
effectiveValue = value;
420+
}
421+
382422
BoolExpr shouldBypass =
383-
fieldType.kind().isPrimitive() ? ctx.mkEq(value, defaultVal) : ctx.mkFalse();
423+
fieldType.kind().isPrimitive() ? ctx.mkEq(effectiveValue, defaultVal) : ctx.mkFalse();
424+
425+
BoolExpr shouldStore = ctx.mkAnd(isPresent, ctx.mkNot(shouldBypass));
384426

385427
msgValues =
386-
(ArrayExpr) ctx.mkITE(shouldBypass, msgValues, ctx.mkStore(msgValues, key, value));
428+
(ArrayExpr)
429+
ctx.mkITE(shouldStore, ctx.mkStore(msgValues, key, effectiveValue), msgValues);
387430

388431
msgPresence =
389432
(ArrayExpr)
390-
ctx.mkITE(shouldBypass, msgPresence, ctx.mkStore(msgPresence, key, ctx.mkTrue()));
433+
ctx.mkITE(shouldStore, ctx.mkStore(msgPresence, key, ctx.mkTrue()), msgPresence);
391434
}
392435

393436
typeConstraints.add(
@@ -396,7 +439,9 @@ private TranslatedValue translateStruct(CelExpr celExpr, CelAbstractSyntaxTree a
396439
typeConstraints.add(ctx.mkEq(typeSystem.getMsgPresence(msgRef), msgPresence));
397440

398441
Expr<?> result = typeSystem.wrapMessage(msgRef);
399-
return TranslatedValue.propagateStrict(ctx, typeSystem, result, celExpr, elementsTv);
442+
BoolExpr baseTaint = ctx.mkFalse();
443+
return TranslatedValue.propagateStrict(
444+
ctx, typeSystem, result, Optional.of(celExpr), baseTaint, elementsTv);
400445
}
401446

402447
private Expr<?> getDefaultValueForType(CelType type) {
@@ -644,6 +689,12 @@ private <T> T withScope(String varName, TranslatedValue value, Supplier<T> actio
644689
private TranslatedValue translateComprehension(CelExpr celExpr, CelAbstractSyntaxTree ast) {
645690
CelComprehension comp = celExpr.comprehension();
646691
CelExpr iterRangeExpr = comp.iterRange();
692+
if (iterRangeExpr.exprKind().getKind() == ExprKind.Kind.IDENT) {
693+
TranslatedValue boundTv = symbolTable.get(iterRangeExpr.ident().name());
694+
if (boundTv != null) {
695+
iterRangeExpr = boundTv.celExpr().orElse(iterRangeExpr);
696+
}
697+
}
647698
List<IterationElement> iterationElements = new ArrayList<>();
648699
List<BoolExpr> taints = new ArrayList<>();
649700
List<Expr<?>> allRangeElems = new ArrayList<>();
@@ -1243,6 +1294,7 @@ private Optional<Object> toCacheKey(CelExpr expr) {
12431294
}
12441295
builder.add(elemKey.get());
12451296
}
1297+
builder.add(expr.list().optionalIndices());
12461298
return Optional.of(builder.build());
12471299
default:
12481300
return Optional.empty();

‎verifier/src/test/java/dev/cel/verifier/CelVerifierZ3ImplTest.java‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
import dev.cel.common.ast.CelExpr.CelCall;
4343
import dev.cel.common.types.ListType;
4444
import dev.cel.common.types.MapType;
45+
import dev.cel.common.types.OptionalType;
4546
import dev.cel.common.types.ProtoMessageTypeProvider;
4647
import dev.cel.common.types.SimpleType;
4748
import dev.cel.common.types.StructTypeReference;
@@ -87,6 +88,7 @@ public final class CelVerifierZ3ImplTest {
8788
.addVar("y", SimpleType.INT)
8889
.addVar("a", SimpleType.BOOL)
8990
.addVar("b", SimpleType.BOOL)
91+
.addVar("opt_var", OptionalType.create(SimpleType.INT))
9092
.addVar("role", SimpleType.STRING)
9193
.addVar("country", SimpleType.STRING)
9294
.addVar("port", SimpleType.INT)
@@ -1250,8 +1252,8 @@ private enum EquivalenceInconclusiveTestCase {
12501252
"size(int_list) == 6 ? int_list.map(x, 2.0) : [1.0]"),
12511253
TRUNCATION_DIVERGENCE_DIFFERENT_BYTES(
12521254
"size(int_list) == 6 ? int_list.map(x, b'a') : [b'a']",
1253-
"size(int_list) == 6 ? int_list.map(x, b'b') : [b'a']");
1254-
1255+
"size(int_list) == 6 ? int_list.map(x, b'b') : [b'a']"),
1256+
;
12551257
final String exprA;
12561258
final String exprB;
12571259

@@ -1320,7 +1322,6 @@ private enum EquivalenceTestCase {
13201322
MACRO_STRUCTURAL_EQUIVALENCE_PRESERVED(
13211323
"request.auth.claims.groups.all(g, g == 'admin')",
13221324
"request.auth.claims.groups.all(x, x == 'admin')"),
1323-
13241325
CONSTANTS_BYTES("by == b\"abc\"", "b\"abc\" == by"),
13251326
STRING_CONCATENATION("\"abc\" + \"def\"", "\"abcdef\""),
13261327
COMPLEX_MESSAGE(
@@ -1334,6 +1335,12 @@ private enum EquivalenceTestCase {
13341335
OPTIONAL_VALUE_EQUIVALENCE("optional.of(x).value()", "x"),
13351336
OPTIONAL_HAS_VALUE_EQUIVALENCE("optional.of(x).hasValue()", "true"),
13361337
OPTIONAL_NONE_HAS_VALUE_EQUIVALENCE("optional.none().hasValue()", "false"),
1338+
OPTIONAL_PRUNE_LIST_LITERAL("[1, ?optional.of(3)]", "[1,3]"),
1339+
OPTIONAL_PRUNE_LIST_NONE("[?optional.none(), ?opt_var]", "[?opt_var]"),
1340+
OPTIONAL_PRUNE_MAP_NONE("{?1: optional.none()}", "{}"),
1341+
OPTIONAL_PRUNE_STRUCT_LIST(
1342+
"TestAllTypes{?repeated_int32: optional.of([1, 2])}",
1343+
"cel.expr.conformance.proto3.TestAllTypes{repeated_int32: [1, 2]}"),
13371344
FUNCTIONS("size(\"abc\") == size(role)", "size(role) == size(\"abc\")"),
13381345
NOT_EQUALS("x != y", "!(x == y)"),
13391346
LESS("x < y", "y > x"),
@@ -1420,7 +1427,8 @@ private enum EquivalenceTestCase {
14201427
"has(dyn({'a': 1}).a) && has(dyn(TestAllTypes{single_int32: 1}).single_int32)"),
14211428
DYNAMIC_INDEXING_TYPE_MISMATCH(
14221429
"type(request) == type(1) && request[1] == 1 && request[2] == 2",
1423-
"type(request) == type(1) && 1 / 0 == 1 && request[2] == 2");
1430+
"type(request) == type(1) && 1 / 0 == 1 && request[2] == 2"),
1431+
;
14241432

14251433
private final String exprA;
14261434
private final String exprB;
@@ -1463,7 +1471,8 @@ private enum EquivalenceViolationTestCase {
14631471
CROSS_TYPE_SYMBOLIC_INEQUALITY_UINT_INT("dyn(u) == dyn(x)", "false"),
14641472
OPTIONAL_OR_VALUE_VIOLATION("optional.of(x).orValue(y)", "y"),
14651473
OPTIONAL_VALUE_VIOLATION("optional.of(x).value()", "y"),
1466-
CROSS_NUMERIC_EQUALITY_INT_DYN_VIOLATION("1 == request", "false");
1474+
CROSS_NUMERIC_EQUALITY_INT_DYN_VIOLATION("1 == request", "false"),
1475+
OPTIONAL_INVALID_PRUNE_OPT_VAR("[1, ?opt_var]", "[1]");
14671476

14681477
final String exprA;
14691478
final String exprB;

0 commit comments

Comments
 (0)