Fix FileDescriptorSetConverter to always reference WellKnownTypes descriptors from generated ones

PiperOrigin-RevId: 815892539

Author: l46kok

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

@@ -25,6 +25,7 @@ java_library(
         "//checker:checker_legacy_environment",
         "//checker:proto_type_mask",
         "//common:cel_ast",
+        "//common:cel_descriptor_util",
         "//common:cel_source",
         "//common:compiler_common",
         "//common:container",

bundle/src/test/java/dev/cel/bundle/CelImplTest.java

@@ -43,14 +43,17 @@
 import com.google.protobuf.ByteString;
 import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
+import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FileDescriptor;
 import com.google.protobuf.Duration;
+import com.google.protobuf.DynamicMessage;
 import com.google.protobuf.Empty;
 import com.google.protobuf.FieldMask;
 import com.google.protobuf.Message;
 import com.google.protobuf.Struct;
 import com.google.protobuf.TextFormat;
 import com.google.protobuf.Timestamp;
+import com.google.protobuf.TypeRegistry;
 import com.google.protobuf.WrappersProto;
 import com.google.rpc.context.AttributeContext;
 import com.google.testing.junit.testparameterinjector.TestParameter;
@@ -62,6 +65,7 @@
 import dev.cel.checker.TypeProvider;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelContainer;
+import dev.cel.common.CelDescriptorUtil;
 import dev.cel.common.CelErrorCode;
 import dev.cel.common.CelIssue;
 import dev.cel.common.CelOptions;
@@ -91,6 +95,7 @@
 import dev.cel.compiler.CelCompilerFactory;
 import dev.cel.compiler.CelCompilerImpl;
 import dev.cel.expr.conformance.proto2.Proto2ExtensionScopedMessage;
+import dev.cel.expr.conformance.proto2.TestAllTypesExtensions;
 import dev.cel.expr.conformance.proto3.TestAllTypes;
 import dev.cel.parser.CelParserImpl;
 import dev.cel.parser.CelStandardMacro;
@@ -196,13 +201,11 @@ public void build_badFileDescriptorSet() {
             IllegalArgumentException.class,
             () ->
                 standardCelBuilderWithMacros()
-                    .setContainer(CelContainer.ofName("google.rpc.context.AttributeContext"))
+                    .setContainer(CelContainer.ofName("cel.expr.conformance.proto2"))
                     .addFileTypes(
                         FileDescriptorSet.newBuilder()
-                            .addFile(AttributeContext.getDescriptor().getFile().toProto())
+                            .addFile(TestAllTypesExtensions.getDescriptor().getFile().toProto())
                             .build())
-                    .setProtoResultType(
-                        CelProtoTypes.createMessage("google.rpc.context.AttributeContext.Resource"))
                     .build());
     assertThat(e).hasMessageThat().contains("file descriptor set with unresolved proto file");
   }
@@ -2124,6 +2127,56 @@ public void program_evaluateCanonicalTypesToNativeTypesDisabled_producesBytesPro
     assertThat(result).isEqualTo(ByteString.copyFromUtf8("abc"));
   }
 
+  @Test
+  public void program_fdsContainsWktDependency_descriptorInstancesMatch() throws Exception {
+    // Force serialization of the descriptor to get a unique instance
+    FileDescriptorProto proto = TestAllTypes.getDescriptor().getFile().toProto();
+    FileDescriptorSet fds = FileDescriptorSet.newBuilder().addFile(proto).build();
+    ImmutableSet<FileDescriptor> fileDescriptors =
+        CelDescriptorUtil.getFileDescriptorsFromFileDescriptorSet(fds);
+    ImmutableSet<Descriptor> descriptors =
+        CelDescriptorUtil.getAllDescriptorsFromFileDescriptor(fileDescriptors)
+            .messageTypeDescriptors();
+    Descriptor testAllTypesDescriptor =
+        descriptors.stream()
+            .filter(x -> x.getFullName().equals(TestAllTypes.getDescriptor().getFullName()))
+            .findAny()
+            .get();
+
+    // Parse text proto using this fds
+    TypeRegistry typeRegistry = TypeRegistry.newBuilder().add(descriptors).build();
+    TestAllTypes.Builder testAllTypesBuilder = TestAllTypes.newBuilder();
+    TextFormat.Parser textFormatParser =
+        TextFormat.Parser.newBuilder().setTypeRegistry(typeRegistry).build();
+    String textProto =
+        "single_timestamp {\n" //
+            + "  seconds: 100\n" //
+            + "}";
+    textFormatParser.merge(textProto, testAllTypesBuilder);
+    TestAllTypes testAllTypesFromTextProto = testAllTypesBuilder.build();
+    DynamicMessage dynamicMessage =
+        DynamicMessage.parseFrom(testAllTypesDescriptor, testAllTypesFromTextProto.toByteArray());
+    // Setup CEL environment with the same descriptors obtained from FDS
+    Cel cel =
+        standardCelBuilderWithMacros()
+            .addMessageTypes(descriptors)
+            .setEvaluateLinkedMessageTypes(false)
+            .setOptions(
+                CelOptions.current()
+                    .evaluateCanonicalTypesToNativeValues(true)
+                    .enableTimestampEpoch(true)
+                    .build())
+            .setContainer(CelContainer.ofName("cel.expr.conformance.proto3"))
+            .build();
+    CelAbstractSyntaxTree ast =
+        cel.compile("TestAllTypes{single_timestamp: timestamp(100)}").getAst();
+
+    DynamicMessage evalResult = (DynamicMessage) cel.createProgram(ast).eval();
+
+    // This should strictly equal regardless of where the descriptors came from for WKTs
+    assertThat(evalResult).isEqualTo(dynamicMessage);
+  }
+
   @Test
   public void toBuilder_isImmutable() {
     CelBuilder celBuilder = CelFactory.standardCelBuilder();

checker/src/main/java/dev/cel/checker/BUILD.bazel

@@ -73,7 +73,7 @@ java_library(
         ":type_provider_legacy_impl",
         "//:auto_value",
         "//common:cel_ast",
-        "//common:cel_descriptors",
+        "//common:cel_descriptor_util",
         "//common:cel_source",
         "//common:compiler_common",
         "//common:container",

common/BUILD.bazel

@@ -106,5 +106,18 @@ java_library(
 
 java_library(
     name = "cel_descriptors",
+    visibility = ["//:internal"],
     exports = ["//common/src/main/java/dev/cel/common:cel_descriptors"],
 )
+
+java_library(
+    name = "cel_descriptor_util",
+    visibility = [
+        "//:internal",
+        # TODO: Remove references to the following clients
+        "//java/com/google/abuse/admin/notebook/compiler/checkedtypes:__pkg__",
+        "//java/com/google/paymentfraud/v2/util/featurereplay/common/risklogrecordio:__pkg__",
+        "//java/com/google/payments/consumer/growth/treatmentconfig/management/backend/service/config/utils:__pkg__",
+    ],
+    exports = ["//common/src/main/java/dev/cel/common:cel_descriptor_util"],
+)

common/src/main/java/dev/cel/common/BUILD.bazel

@@ -36,18 +36,31 @@ PROTO_V1ALPHA1_AST_SOURCE = [
 ]
 
 java_library(
-    name = "cel_descriptors",
+    name = "cel_descriptor_util",
     srcs = [
         "CelDescriptorUtil.java",
-        "CelDescriptors.java",
     ],
     tags = [
     ],
     deps = [
-        "//:auto_value",
+        ":cel_descriptors",
         "//common/annotations",
         "//common/internal:file_descriptor_converter",
         "//common/types:cel_types",
+        "@maven//:com_google_guava_guava",
+        "@maven//:com_google_protobuf_protobuf_java",
+    ],
+)
+
+java_library(
+    name = "cel_descriptors",
+    srcs = [
+        "CelDescriptors.java",
+    ],
+    tags = [
+    ],
+    deps = [
+        "//:auto_value",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",

common/src/main/java/dev/cel/common/CelDescriptorUtil.java

@@ -166,10 +166,12 @@ private static void collectMessageTypeDescriptors(
     if (visited.contains(messageName)) {
       return;
     }
+
     if (!descriptor.getOptions().getMapEntry()) {
       visited.add(messageName);
       celDescriptors.addMessageTypeDescriptors(descriptor);
     }
+
     if (CelTypes.getWellKnownCelType(messageName).isPresent()) {
       return;
     }
@@ -234,6 +236,7 @@ private static void copyToFileDescriptorSet(
     if (visited.contains(fd.getFullName())) {
       return;
     }
+
     visited.add(fd.getFullName());
     for (FileDescriptor dep : fd.getDependencies()) {
       copyToFileDescriptorSet(visited, dep, files);

common/src/main/java/dev/cel/common/internal/BUILD.bazel

@@ -232,7 +232,9 @@ java_library(
     tags = [
     ],
     deps = [
+        ":cel_descriptor_pools",
         "//:auto_value",
+        "//common/internal:well_known_proto",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",

common/src/main/java/dev/cel/common/internal/DefaultDescriptorPool.java

@@ -119,6 +119,10 @@ public static DefaultDescriptorPool create(
         extensionRegistry);
   }
 
+  public static Descriptor getWellKnownProtoDescriptor(WellKnownProto wellKnownProto) {
+    return WELL_KNOWN_PROTO_TO_DESCRIPTORS.get(wellKnownProto);
+  }
+
   @Override
   public Optional<Descriptor> findDescriptor(String name) {
     return Optional.ofNullable(descriptorMap.get(name));

common/src/main/java/dev/cel/common/internal/DefaultInstanceMessageFactory.java

@@ -59,6 +59,8 @@ public Optional<Message> getPrototype(Descriptor descriptor) {
     }
 
     Message fullMessage = (Message) defaultInstance;
+    Descriptor fullDescriptor = fullMessage.getDescriptorForType();
+    System.out.println(fullDescriptor);
 
     // Reference equality is intended. We want to make sure the descriptors are equal
     // to guarantee types to be hermetic if linked types is disabled.

common/src/main/java/dev/cel/common/internal/FileDescriptorSetConverter.java

@@ -15,6 +15,7 @@
 package dev.cel.common.internal;
 
 import com.google.common.base.VerifyException;
+import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
@@ -76,7 +77,15 @@ private static FileDescriptor readDescriptor(
     // Read dependencies first, they are needed to create the logical descriptor from the proto.
     List<FileDescriptor> deps = new ArrayList<>();
     for (String dep : fileProto.getDependencyList()) {
-      deps.add(readDescriptor(dep, descriptorProtos, descriptors));
+      ImmutableCollection<WellKnownProto> wktProtos = WellKnownProto.getByPathName(dep);
+      if (wktProtos.isEmpty()) {
+        deps.add(readDescriptor(dep, descriptorProtos, descriptors));
+      } else {
+        // Ensure the generated message's descriptor is used as a dependency for WKTs to avoid
+        // issues with descriptor instance mismatch.
+        WellKnownProto wellKnownProto = wktProtos.iterator().next();
+        deps.add(DefaultDescriptorPool.getWellKnownProtoDescriptor(wellKnownProto).getFile());
+      }
     }
     // Create the file descriptor, cache, and return.
     try {