Add option for using protobuf WKTs as context decl in type checking.

PiperOrigin-RevId: 761662037

Author: jnthntatum

checker/BUILD

@@ -118,6 +118,9 @@ cc_test(
     name = "type_checker_builder_factory_test",
     srcs = ["type_checker_builder_factory_test.cc"],
     deps = [
+        ":checker_options",
+        ":standard_library",
+        ":type_checker",
         ":type_checker_builder",
         ":type_checker_builder_factory",
         ":validation_result",

checker/checker_options.h

@@ -43,6 +43,19 @@ struct CheckerOptions {
   // as parsed.
   bool update_struct_type_names = true;
 
+  // Well-known types defined by protobuf are treated specially in CEL, and
+  // generally don't behave like other messages as runtime values. When used as
+  // context declarations, this introduces some ambiguity about the intended
+  // types of the field declarations, so it is disallowed by default.
+  //
+  // When enabled, the well-known types are treated like a normal message type
+  // for the purposes for declaring context bindings (i.e no unpacking or
+  // adapting), and use the Descriptor that is assumed by CEL.
+  //
+  // E.g. for google.protobuf.Any, the type checker will add a context binding
+  // with `type_url: string` and `value: bytes` as top level variables.
+  bool allow_well_known_type_context_declarations = false;
+
   // Maximum number (inclusive) of expression nodes to check for an input
   // expression.
   //

checker/internal/BUILD

@@ -197,6 +197,8 @@ cc_test(
     deps = [
         ":test_ast_helpers",
         ":type_checker_impl",
+        "//checker:checker_options",
+        "//checker:standard_library",
         "//checker:type_checker",
         "//checker:validation_result",
         "//common:decl",

checker/internal/type_checker_builder_impl.cc

@@ -287,7 +287,8 @@ absl::Status TypeCheckerBuilderImpl::AddContextDeclaration(
         absl::StrCat("context declaration '", type, "' not found"));
   }
 
-  if (IsWellKnownMessageType(desc)) {
+  if (IsWellKnownMessageType(desc) &&
+      !options_.allow_well_known_type_context_declarations) {
     return absl::InvalidArgumentError(
         absl::StrCat("context declaration '", type, "' is not a struct"));
   }

checker/type_checker_builder_factory_test.cc

@@ -21,7 +21,10 @@
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
+#include "checker/checker_options.h"
 #include "checker/internal/test_ast_helpers.h"
+#include "checker/standard_library.h"
+#include "checker/type_checker.h"
 #include "checker/type_checker_builder.h"
 #include "checker/validation_result.h"
 #include "common/decl.h"
@@ -391,6 +394,108 @@ TEST(TypeCheckerBuilderTest, AddContextDeclaration) {
   EXPECT_TRUE(result.IsValid());
 }
 
+TEST(TypeCheckerBuilderTest, WellKnownTypeContextDeclarationError) {
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TypeCheckerBuilder> builder,
+      CreateTypeCheckerBuilder(GetSharedTestingDescriptorPool()));
+
+  ASSERT_THAT(builder->AddContextDeclaration("google.protobuf.Any"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("'google.protobuf.Any' is not a struct")));
+}
+
+TEST(TypeCheckerBuilderTest, AllowWellKnownTypeContextDeclaration) {
+  CheckerOptions options;
+  options.allow_well_known_type_context_declarations = true;
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TypeCheckerBuilder> builder,
+      CreateTypeCheckerBuilder(GetSharedTestingDescriptorPool(), options));
+
+  ASSERT_THAT(builder->AddContextDeclaration("google.protobuf.Any"), IsOk());
+  ASSERT_THAT(builder->AddLibrary(StandardCheckerLibrary()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TypeChecker> type_checker,
+                       builder->Build());
+  ASSERT_OK_AND_ASSIGN(
+      auto ast,
+      MakeTestParsedAst(
+          R"cel(value == b'' && type_url == 'type.googleapis.com/google.protobuf.Duration')cel"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result,
+                       type_checker->Check(std::move(ast)));
+
+  ASSERT_TRUE(result.IsValid());
+}
+
+TEST(TypeCheckerBuilderTest, AllowWellKnownTypeContextDeclarationStruct) {
+  CheckerOptions options;
+  options.allow_well_known_type_context_declarations = true;
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TypeCheckerBuilder> builder,
+      CreateTypeCheckerBuilder(GetSharedTestingDescriptorPool(), options));
+
+  ASSERT_THAT(builder->AddContextDeclaration("google.protobuf.Struct"), IsOk());
+  ASSERT_THAT(builder->AddLibrary(StandardCheckerLibrary()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TypeChecker> type_checker,
+                       builder->Build());
+  ASSERT_OK_AND_ASSIGN(
+      auto ast,
+      MakeTestParsedAst(R"cel(fields.foo.bar_list.exists(x, x == 1))cel"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result,
+                       type_checker->Check(std::move(ast)));
+
+  ASSERT_TRUE(result.IsValid());
+}
+
+TEST(TypeCheckerBuilderTest, AllowWellKnownTypeContextDeclarationValue) {
+  CheckerOptions options;
+  options.allow_well_known_type_context_declarations = true;
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TypeCheckerBuilder> builder,
+      CreateTypeCheckerBuilder(GetSharedTestingDescriptorPool(), options));
+
+  ASSERT_THAT(builder->AddContextDeclaration("google.protobuf.Value"), IsOk());
+  ASSERT_THAT(builder->AddLibrary(StandardCheckerLibrary()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TypeChecker> type_checker,
+                       builder->Build());
+  ASSERT_OK_AND_ASSIGN(
+      auto ast, MakeTestParsedAst(
+                    // Note: one of fields are all added with safe traversal, so
+                    // we lose the union discriminator information.
+                    R"cel(
+            null_value == null &&
+            number_value == 0.0 &&
+            string_value == '' &&
+            list_value == [] &&
+            struct_value == {} &&
+            bool_value == false)cel"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result,
+                       type_checker->Check(std::move(ast)));
+
+  ASSERT_TRUE(result.IsValid());
+}
+
+TEST(TypeCheckerBuilderTest, AllowWellKnownTypeContextDeclarationInt64Value) {
+  CheckerOptions options;
+  options.allow_well_known_type_context_declarations = true;
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TypeCheckerBuilder> builder,
+      CreateTypeCheckerBuilder(GetSharedTestingDescriptorPool(), options));
+
+  ASSERT_THAT(builder->AddContextDeclaration("google.protobuf.Int64Value"),
+              IsOk());
+  ASSERT_THAT(builder->AddLibrary(StandardCheckerLibrary()), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TypeChecker> type_checker,
+                       builder->Build());
+  ASSERT_OK_AND_ASSIGN(auto ast, MakeTestParsedAst(R"cel(value == 0)cel"));
+  ASSERT_OK_AND_ASSIGN(ValidationResult result,
+                       type_checker->Check(std::move(ast)));
+
+  ASSERT_TRUE(result.IsValid());
+}
+
 TEST(TypeCheckerBuilderTest, AddLibraryRedeclaredError) {
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<TypeCheckerBuilder> builder,