Fix FileDescriptorSetConverter to always reference WellKnownTypes descriptors from generated ones

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,18 @@
 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.ExtensionRegistry;
 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 +66,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 +96,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 +202,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 +2128,58 @@ 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(),
+ ExtensionRegistry.getEmptyRegistry());
+ // Setup CEL environment with the same descriptors obtained from FDS
+ Cel cel =
+ standardCelBuilderWithMacros()
+ .addMessageTypes(descriptors)
+ .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/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 {