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

@@ -53,6 +53,8 @@ public abstract class CelPolicy {
 
   public abstract ImmutableList<Import> imports();
 
+  public abstract ImmutableList<Invariant> invariants();
+
   /** Creates a new builder to construct a {@link CelPolicy} instance. */
   public static Builder newBuilder() {
     return new AutoValue_CelPolicy.Builder()
@@ -90,6 +92,14 @@ public List<Import> imports() {
       return Collections.unmodifiableList(importList);
     }
 
+    private final ArrayList<Invariant> invariantList = new ArrayList<>();
+
+    abstract Builder setInvariants(ImmutableList<Invariant> value);
+
+    public List<Invariant> invariants() {
+      return Collections.unmodifiableList(invariantList);
+    }
+
     public Map<String, Object> metadata() {
       return Collections.unmodifiableMap(metadata);
     }
@@ -106,6 +116,18 @@ public Builder addImports(Collection<Import> values) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder addInvariant(Invariant value) {
+      invariantList.add(value);
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder addInvariants(Collection<Invariant> values) {
+      invariantList.addAll(values);
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder putMetadata(String key, Object value) {
       metadata.put(key, value);
@@ -122,6 +144,7 @@ public Builder putMetadata(Map<String, Object> map) {
 
     public CelPolicy build() {
       setImports(ImmutableList.copyOf(importList));
+      setInvariants(ImmutableList.copyOf(invariantList));
       setMetadata(ImmutableMap.copyOf(metadata));
       return autoBuild();
     }
@@ -328,4 +351,52 @@ 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 Optional<ValueString> assume();
+
+    public abstract ValueString assertClause();
+
+    /** Builder for {@link Invariant}. */
+    @AutoValue.Builder
+    public abstract static class Builder implements RequiredFieldsChecker {
+      public abstract Builder setId(long value);
+
+      abstract Optional<ValueString> invariantId();
+
+      abstract Optional<ValueString> assertClause();
+
+      public abstract Builder setInvariantId(ValueString value);
+
+      public abstract Builder setDescription(ValueString value);
+
+      public abstract Builder setAssume(ValueString value);
+
+      public abstract Builder setAssertClause(ValueString value);
+
+      @Override
+      public ImmutableList<RequiredField> requiredFields() {
+        return ImmutableList.of(
+            RequiredField.of("id", this::invariantId),
+            RequiredField.of("assert", this::assertClause));
+      }
+
+      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

@@ -28,6 +28,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 +48,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,6 +140,9 @@ 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;
@@ -148,6 +154,36 @@ public CelPolicy parsePolicy(PolicyParserContext<Node> ctx, Node node) {
           .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();
+        if (fieldName.equals("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));
+          }
+        } else {
+          ctx.reportError(keyId, "Unexpected key in verification block: " + fieldName);
+        }
+      }
+    }
+
     private void parseImports(
         CelPolicy.Builder policyBuilder, PolicyParserContext<Node> ctx, Node node) {
       long id = ctx.collectMetadata(node);
@@ -409,6 +445,50 @@ 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":
+            builder.setAssume(ctx.newSourceString(valueNode));
+            break;
+          case "assert":
+            builder.setAssertClause(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);
 }

policy/src/test/java/dev/cel/policy/BUILD.bazel

@@ -8,9 +8,7 @@ package(
 java_library(
     name = "tests",
     testonly = True,
-    srcs = glob(
-        ["*.java"],
-    ),
+    srcs = glob(["*.java"]),
     data = [
         "@cel_policy//conformance:testdata",
     ],

policy/src/test/java/dev/cel/policy/CelPolicyYamlParserTest.java

@@ -17,11 +17,13 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import dev.cel.common.formats.ValueString;
 import dev.cel.policy.CelPolicy.Import;
+import dev.cel.policy.CelPolicy.Invariant;
 import dev.cel.policy.PolicyTestHelper.TestYamlPolicy;
 import dev.cel.policy.testing.K8sTagHandler;
 import org.junit.Test;
@@ -194,6 +196,43 @@ public void parseYamlPolicy_errors(@TestParameter PolicyParseErrorTestCase testC
     assertThat(e).hasMessageThat().isEqualTo(testCase.expectedErrorMessage);
   }
 
+  @Test
+  public void policyBuilder_addInvariants() {
+    Invariant invariant =
+        Invariant.newBuilder(1L)
+            .setInvariantId(ValueString.newBuilder().setValue("id").build())
+            .setAssertClause(ValueString.newBuilder().setValue("true").build())
+            .build();
+    CelPolicy policy =
+        CelPolicy.newBuilder()
+            .setName(ValueString.of(0, "test"))
+            .setPolicySource(CelPolicySource.newBuilder("").build())
+            .addInvariants(ImmutableList.of(invariant))
+            .build();
+    assertThat(policy.invariants()).containsExactly(invariant);
+  }
+
+  @Test
+  public void parseYamlPolicy_invariants_success() throws Exception {
+    String policySource =
+        "name: 'policy_with_invariants'\n"
+            + "verification:\n"
+            + "  invariants:\n"
+            + "    - id: 'inv_1'\n"
+            + "      description: 'invariant description'\n"
+            + "      assume: 'true'\n"
+            + "      assert: 'invariants.result == true'";
+
+    CelPolicy policy = POLICY_PARSER.parse(policySource);
+
+    assertThat(policy.invariants()).hasSize(1);
+    Invariant invariant = Iterables.getOnlyElement(policy.invariants());
+    assertThat(invariant.invariantId().value()).isEqualTo("inv_1");
+    assertThat(invariant.description().get().value()).isEqualTo("invariant description");
+    assertThat(invariant.assume().get().value()).isEqualTo("true");
+    assertThat(invariant.assertClause().value()).isEqualTo("invariants.result == true");
+  }
+
   private enum PolicyParseErrorTestCase {
     MALFORMED_YAML_DOCUMENT(
         "a:\na",
@@ -400,7 +439,70 @@ private enum PolicyParseErrorTestCase {
             + "- foo: bar",
         "ERROR: <input>:2:3: Invalid import key: foo, expected 'name'\n"
             + " | - foo: bar\n"
-            + " | ..^");
+            + " | ..^"),
+    UNSUPPORTED_VERIFICATION_TAG(
+        "verification:\n" //
+            + "  bad_key: true",
+        "ERROR: <input>:2:3: Unexpected key in verification block: bad_key\n"
+            + " |   bad_key: true\n"
+            + " | ..^"),
+    UNSUPPORTED_INVARIANT_TAG(
+        "verification:\n" //
+            + "  invariants:\n" //
+            + "    - id: foo\n" //
+            + "      bad_inv_key: true\n" //
+            + "      assert: 'true'",
+        "ERROR: <input>:4:7: Unexpected key in invariant block: bad_inv_key\n"
+            + " |       bad_inv_key: true\n"
+            + " | ......^"),
+    MISSING_INVARIANT_ID(
+        "verification:\n" //
+            + "  invariants:\n" //
+            + "    - assert: 'true'",
+        "ERROR: <input>:3:7: Missing required attribute(s): id\n"
+            + " |     - assert: 'true'\n"
+            + " | ......^"),
+    MISSING_INVARIANT_ASSERT(
+        "verification:\n" //
+            + "  invariants:\n" //
+            + "    - id: foo",
+        "ERROR: <input>:3:7: Missing required attribute(s): assert\n"
+            + " |     - id: foo\n"
+            + " | ......^"),
+    ILLEGAL_YAML_TYPE_ON_VERIFICATION_VALUE(
+        "verification: illegal\n",
+        "ERROR: <input>:1:15: Got yaml node type tag:yaml.org,2002:str, wanted type(s)"
+            + " [tag:yaml.org,2002:map]\n"
+            + " | verification: illegal\n"
+            + " | ..............^"),
+    ILLEGAL_YAML_TYPE_ON_VERIFICATION_MAP_KEY(
+        "verification:\n" + "  1: foo",
+        "ERROR: <input>:2:3: Got yaml node type tag:yaml.org,2002:int, wanted type(s)"
+            + " [tag:yaml.org,2002:str !txt]\n"
+            + " |   1: foo\n"
+            + " | ..^"),
+    ILLEGAL_YAML_TYPE_ON_INVARIANTS_VALUE(
+        "verification:\n" + "  invariants: illegal\n",
+        "ERROR: <input>:2:15: Got yaml node type tag:yaml.org,2002:str, wanted type(s)"
+            + " [tag:yaml.org,2002:seq]\n"
+            + " |   invariants: illegal\n"
+            + " | ..............^"),
+    ILLEGAL_YAML_TYPE_ON_INVARIANTS_LIST(
+        "verification:\n" + "  invariants:\n" + "    - illegal",
+        "ERROR: <input>:3:7: Got yaml node type tag:yaml.org,2002:str, wanted type(s)"
+            + " [tag:yaml.org,2002:map]\n"
+            + " |     - illegal\n"
+            + " | ......^"),
+    ILLEGAL_YAML_TYPE_ON_INVARIANT_MAP_KEY(
+        "verification:\n"
+            + "  invariants:\n"
+            + "    - 1: foo\n"
+            + "      id: 'hi'\n"
+            + "      assert: 'true'",
+        "ERROR: <input>:3:7: Got yaml node type tag:yaml.org,2002:int, wanted type(s)"
+            + " [tag:yaml.org,2002:str !txt]\n"
+            + " |     - 1: foo\n"
+            + " | ......^");
 
     private final String yamlPolicy;
     private final String expectedErrorMessage;

testing/src/test/resources/policy/verification/flawed_policy.yaml

@@ -0,0 +1,24 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: flawed_policy
+rule:
+  match:
+    - condition: port == 80
+      output: 'true'
+    - output: 'false'
+verification:
+  invariants:
+    - id: always_secure
+      assert: invariants.result == false