Implement custom policy invariants verification

Enables policy authors to declare custom logical invariants (`assume` preconditions and `assert` clauses) on `CelPolicy` definitions, mathematically verifying that properties hold across all possible input states.

PiperOrigin-RevId: 915170572

Author: l46kok

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

@@ -15,6 +15,7 @@
 package dev.cel.policy;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static java.util.stream.Collectors.joining;
 
 import com.google.auto.value.AutoOneOf;
 import com.google.auto.value.AutoValue;
@@ -53,6 +54,10 @@ public abstract class CelPolicy {
 
   public abstract ImmutableList<Import> imports();
 
+  public abstract ImmutableList<Invariant> invariants();
+
+  public abstract ImmutableList<Variable> verificationVariables();
+
   /** Creates a new builder to construct a {@link CelPolicy} instance. */
   public static Builder newBuilder() {
     return new AutoValue_CelPolicy.Builder()
@@ -74,6 +79,8 @@ public abstract static class Builder {
 
     public abstract Builder setDisplayName(ValueString displayName);
 
+    public abstract Rule rule();
+
     public abstract Builder setRule(Rule rule);
 
     public abstract Builder setPolicySource(CelPolicySource policySource);
@@ -90,6 +97,14 @@ public List<Import> imports() {
       return Collections.unmodifiableList(importList);
     }
 
+    abstract ImmutableList<Invariant> invariants();
+
+    abstract ImmutableList.Builder<Invariant> invariantsBuilder();
+
+    abstract ImmutableList<Variable> verificationVariables();
+
+    abstract ImmutableList.Builder<Variable> verificationVariablesBuilder();
+
     public Map<String, Object> metadata() {
       return Collections.unmodifiableMap(metadata);
     }
@@ -106,6 +121,24 @@ public Builder addImports(Collection<Import> values) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder addInvariant(Invariant value) {
+      invariantsBuilder().add(value);
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder addVerificationVariable(Variable value) {
+      verificationVariablesBuilder().add(value);
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder addVerificationVariables(Collection<Variable> values) {
+      verificationVariablesBuilder().addAll(values);
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder putMetadata(String key, Object value) {
       metadata.put(key, value);
@@ -328,4 +361,115 @@ public static Import create(long id, ValueString name) {
       return new AutoValue_CelPolicy_Import(id, name);
     }
   }
+
+  /**
+   * Invariant declares a required logical property that must hold true under specified
+   * preconditions.
+   */
+  @AutoValue
+  public abstract static class Invariant {
+    public abstract long id();
+
+    public abstract ValueString invariantId();
+
+    public abstract Optional<ValueString> description();
+
+    public abstract ImmutableList<ValueString> assume();
+
+    public abstract ImmutableList<ValueString> assertClause();
+
+    public String assumeSourceString() {
+      if (assume().isEmpty()) {
+        return "true";
+      }
+      if (assume().size() == 1) {
+        return assume().get(0).value();
+      }
+      return assume().stream().map(v -> "(" + v.value() + ")").collect(joining(" && "));
+    }
+
+    public String assertSourceString() {
+      if (assertClause().isEmpty()) {
+        return "true";
+      }
+      if (assertClause().size() == 1) {
+        return assertClause().get(0).value();
+      }
+      return assertClause().stream().map(v -> "(" + v.value() + ")").collect(joining(" && "));
+    }
+
+    /** Builder for {@link Invariant}. */
+    @AutoValue.Builder
+    public abstract static class Builder implements RequiredFieldsChecker {
+      public abstract Builder setId(long value);
+
+      abstract Optional<ValueString> invariantId();
+
+      abstract ImmutableList<ValueString> assume();
+
+      abstract ImmutableList.Builder<ValueString> assumeBuilder();
+
+      abstract ImmutableList<ValueString> assertClause();
+
+      abstract ImmutableList.Builder<ValueString> assertClauseBuilder();
+
+      public abstract Builder setInvariantId(ValueString value);
+
+      public abstract Builder setDescription(ValueString value);
+
+      public Builder setAssume(ValueString value) {
+        return setAssume(ImmutableList.of(value));
+      }
+
+      abstract Builder setAssume(ImmutableList<ValueString> values);
+
+      @CanIgnoreReturnValue
+      public Builder addAssume(ValueString value) {
+        assumeBuilder().add(value);
+        return this;
+      }
+
+      @CanIgnoreReturnValue
+      public Builder addAssume(Iterable<ValueString> values) {
+        assumeBuilder().addAll(values);
+        return this;
+      }
+
+      public Builder setAssertClause(ValueString value) {
+        return setAssertClause(ImmutableList.of(value));
+      }
+
+      abstract Builder setAssertClause(ImmutableList<ValueString> values);
+
+      @CanIgnoreReturnValue
+      public Builder addAssertClause(ValueString value) {
+        assertClauseBuilder().add(value);
+        return this;
+      }
+
+      @CanIgnoreReturnValue
+      public Builder addAssertClause(Iterable<ValueString> values) {
+        assertClauseBuilder().addAll(values);
+        return this;
+      }
+
+      @Override
+      public ImmutableList<RequiredField> requiredFields() {
+        return ImmutableList.of(
+            RequiredField.of("id", this::invariantId),
+            RequiredField.of(
+                "assert",
+                () ->
+                    assertClause().isEmpty()
+                        ? Optional.empty()
+                        : Optional.of(assertClause().get(0))));
+      }
+
+      public abstract Invariant build();
+    }
+
+    public static Builder newBuilder(long id) {
+      return new AutoValue_CelPolicy_Invariant.Builder().setId(id);
+    }
+  }
 }

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

@@ -15,6 +15,7 @@
 package dev.cel.policy;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static dev.cel.common.formats.YamlHelper.ERROR;
 import static dev.cel.common.formats.YamlHelper.assertRequiredFields;
 import static dev.cel.common.formats.YamlHelper.assertYamlType;
@@ -28,6 +29,7 @@
 import dev.cel.common.formats.YamlParserContextImpl;
 import dev.cel.common.internal.CelCodePointArray;
 import dev.cel.policy.CelPolicy.Import;
+import dev.cel.policy.CelPolicy.Invariant;
 import dev.cel.policy.CelPolicy.Match;
 import dev.cel.policy.CelPolicy.Match.Result;
 import dev.cel.policy.CelPolicy.Variable;
@@ -47,6 +49,8 @@ final class CelPolicyYamlParser implements CelPolicyParser {
       Match.newBuilder(0).setCondition(ERROR_VALUE).setResult(Result.ofOutput(ERROR_VALUE)).build();
   private static final Variable ERROR_VARIABLE =
       Variable.newBuilder().setExpression(ERROR_VALUE).setName(ERROR_VALUE).build();
+  private static final Invariant ERROR_INVARIANT =
+      Invariant.newBuilder(0).setInvariantId(ERROR_VALUE).setAssertClause(ERROR_VALUE).build();
 
   private final TagVisitor<Node> tagVisitor;
   private final boolean enableSimpleVariables;
@@ -137,17 +141,82 @@ public CelPolicy parsePolicy(PolicyParserContext<Node> ctx, Node node) {
           case "rule":
             policyBuilder.setRule(parseRule(ctx, policyBuilder, valueNode));
             break;
+          case "verification":
+            parseVerification(policyBuilder, ctx, valueNode);
+            break;
           default:
             tagVisitor.visitPolicyTag(ctx, keyId, fieldName, valueNode, policyBuilder);
             break;
         }
       }
 
+      ImmutableSet<String> ruleVarNames =
+          policyBuilder.rule().variables().stream()
+              .map(CelPolicy.Variable::name)
+              .filter(name -> !name.equals(ERROR_VALUE))
+              .map(ValueString::value)
+              .collect(toImmutableSet());
+      for (Variable verVar : policyBuilder.verificationVariables()) {
+        if (!verVar.name().equals(ERROR_VALUE)
+            && ruleVarNames.contains(verVar.name().value())) {
+          ctx.reportError(
+              verVar.name().id(),
+              "Duplicate variable name '"
+                  + verVar.name().value()
+                  + "' in verification.variables; already defined in rule.variables");
+        }
+      }
+
       return policyBuilder
           .setPolicySource(policySource.toBuilder().setPositionsMap(ctx.getIdToOffsetMap()).build())
           .build();
     }
 
+    private void parseVerification(
+        CelPolicy.Builder policyBuilder, PolicyParserContext<Node> ctx, Node node) {
+      long id = ctx.collectMetadata(node);
+      if (!assertYamlType(ctx, id, node, YamlNodeType.MAP)) {
+        return;
+      }
+      MappingNode mappingNode = (MappingNode) node;
+      for (NodeTuple nodeTuple : mappingNode.getValue()) {
+        Node key = nodeTuple.getKeyNode();
+        long keyId = ctx.collectMetadata(key);
+        if (!assertYamlType(ctx, keyId, key, YamlNodeType.STRING, YamlNodeType.TEXT)) {
+          continue;
+        }
+        String fieldName = ((ScalarNode) key).getValue();
+        Node valueNode = nodeTuple.getValueNode();
+        switch (fieldName) {
+          case "invariants": {
+            long valueId = ctx.collectMetadata(valueNode);
+            if (!assertYamlType(ctx, valueId, valueNode, YamlNodeType.LIST)) {
+              continue;
+            }
+            SequenceNode invariantListNode = (SequenceNode) valueNode;
+            for (Node invariantNode : invariantListNode.getValue()) {
+              policyBuilder.addInvariant(parseInvariant(ctx, policyBuilder, invariantNode));
+            }
+            break;
+          }
+          case "variables": {
+            long valueId = ctx.collectMetadata(valueNode);
+            if (!assertYamlType(ctx, valueId, valueNode, YamlNodeType.LIST)) {
+              continue;
+            }
+            SequenceNode variableListNode = (SequenceNode) valueNode;
+            for (Node varNode : variableListNode.getValue()) {
+              policyBuilder.addVerificationVariable(parseVariable(ctx, policyBuilder, varNode));
+            }
+            break;
+          }
+          default:
+            ctx.reportError(keyId, "Unexpected key in verification block: " + fieldName);
+            break;
+        }
+      }
+    }
+
     private void parseImports(
         CelPolicy.Builder policyBuilder, PolicyParserContext<Node> ctx, Node node) {
       long id = ctx.collectMetadata(node);
@@ -409,6 +478,82 @@ private Variable parseVariableObject(
       return builder.build();
     }
 
+    @Override
+    public CelPolicy.Invariant parseInvariant(
+        PolicyParserContext<Node> ctx, CelPolicy.Builder policyBuilder, Node node) {
+      long id = ctx.collectMetadata(node);
+      Invariant.Builder builder = Invariant.newBuilder(id);
+      if (!assertYamlType(ctx, id, node, YamlNodeType.MAP)) {
+        return ERROR_INVARIANT;
+      }
+
+      MappingNode invariantMap = (MappingNode) node;
+      for (NodeTuple nodeTuple : invariantMap.getValue()) {
+        Node keyNode = nodeTuple.getKeyNode();
+        long keyId = ctx.collectMetadata(keyNode);
+        if (!assertYamlType(ctx, keyId, keyNode, YamlNodeType.STRING, YamlNodeType.TEXT)) {
+          continue;
+        }
+        Node valueNode = nodeTuple.getValueNode();
+        String keyName = ((ScalarNode) keyNode).getValue();
+        switch (keyName) {
+          case "id":
+            builder.setInvariantId(ctx.newYamlString(valueNode));
+            break;
+          case "description":
+            builder.setDescription(ctx.newYamlString(valueNode));
+            break;
+          case "assume": {
+            if (!assertYamlType(
+                ctx,
+                ctx.collectMetadata(valueNode),
+                valueNode,
+                YamlNodeType.STRING,
+                YamlNodeType.TEXT,
+                YamlNodeType.LIST)) {
+              break;
+            }
+            if (valueNode instanceof SequenceNode) {
+              for (Node itemNode : ((SequenceNode) valueNode).getValue()) {
+                builder.addAssume(ctx.newSourceString(itemNode));
+              }
+            } else {
+              builder.addAssume(ctx.newSourceString(valueNode));
+            }
+            break;
+          }
+          case "assert": {
+            if (!assertYamlType(
+                ctx,
+                ctx.collectMetadata(valueNode),
+                valueNode,
+                YamlNodeType.STRING,
+                YamlNodeType.TEXT,
+                YamlNodeType.LIST)) {
+              break;
+            }
+            if (valueNode instanceof SequenceNode) {
+              for (Node itemNode : ((SequenceNode) valueNode).getValue()) {
+                builder.addAssertClause(ctx.newSourceString(itemNode));
+              }
+            } else {
+              builder.addAssertClause(ctx.newSourceString(valueNode));
+            }
+            break;
+          }
+          default:
+            ctx.reportError(keyId, "Unexpected key in invariant block: " + keyName);
+            break;
+        }
+      }
+
+      if (!assertRequiredFields(ctx, id, builder.getMissingRequiredFieldNames())) {
+        return ERROR_INVARIANT;
+      }
+
+      return builder.build();
+    }
+
     private ParserImpl(
         TagVisitor<Node> tagVisitor,
         boolean enableSimpleVariables,

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

@@ -16,6 +16,7 @@
 
 import com.google.auto.value.AutoValue;
 import dev.cel.common.formats.ParserContext;
+import dev.cel.policy.CelPolicy.Invariant;
 import dev.cel.policy.CelPolicy.Match;
 import dev.cel.policy.CelPolicy.Rule;
 import dev.cel.policy.CelPolicy.Variable;
@@ -51,4 +52,6 @@ static NewPolicyMetadata create(CelPolicySource source, long id) {
   Match parseMatch(PolicyParserContext<T> ctx, CelPolicy.Builder policyBuilder, T node);
 
   Variable parseVariable(PolicyParserContext<T> ctx, CelPolicy.Builder policyBuilder, T node);
+
+  Invariant parseInvariant(PolicyParserContext<T> ctx, CelPolicy.Builder policyBuilder, T node);
 }