Internal Changes

PiperOrigin-RevId: 574480963

Author: l46kok

bundle/src/main/java/dev/cel/bundle/CelBuilder.java

@@ -20,7 +20,6 @@
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FileDescriptor;
-import com.google.protobuf.ExtensionRegistry;
 import com.google.protobuf.Message;
 import dev.cel.checker.ProtoTypeMask;
 import dev.cel.checker.TypeProvider;
@@ -283,13 +282,6 @@ public interface CelBuilder {
   @CanIgnoreReturnValue
   CelBuilder addRuntimeLibraries(Iterable<CelRuntimeLibrary> libraries);
 
-  /**
-   * Sets a proto ExtensionRegistry to assist with unpacking Any messages containing a proto2
-   extension field.
-   */
-  @CanIgnoreReturnValue
-  CelBuilder setExtensionRegistry(ExtensionRegistry extensionRegistry);
-
   /** Construct a new {@code Cel} instance from the provided configuration. */
   Cel build();
 }

bundle/src/main/java/dev/cel/bundle/CelImpl.java

@@ -25,7 +25,6 @@
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FileDescriptor;
-import com.google.protobuf.ExtensionRegistry;
 import com.google.protobuf.Message;
 import dev.cel.checker.CelCheckerBuilder;
 import dev.cel.checker.ProtoTypeMask;
@@ -340,13 +339,6 @@ public Builder addRuntimeLibraries(Iterable<CelRuntimeLibrary> libraries) {
       return this;
     }
 
-    @Override
-    public CelBuilder setExtensionRegistry(ExtensionRegistry extensionRegistry) {
-      checkNotNull(extensionRegistry);
-      runtimeBuilder.setExtensionRegistry(extensionRegistry);
-      return this;
-    }
-
     @Override
     public Cel build() {
       return new CelImpl(

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

@@ -91,9 +91,11 @@ java_library(
     ],
 )
 
+# keep sorted
 DYNAMIC_PROTO_SOURCES = [
     "DynamicProto.java",
     "ProtoAdapter.java",
+    "ProtoRegistryProvider.java",
 ]
 
 java_library(
@@ -117,12 +119,13 @@ java_library(
     ],
     deps = [
         ":converter",
-        ":proto_message_factory",
-        ":well_known_proto",
+        ":default_instance_message_factory",
         "//:auto_value",
+        "//common",
         "//common:error_codes",
         "//common:runtime_exception",
         "//common/annotations",
+        "//common/types:cel_types",
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@maven//:com_google_code_findbugs_annotations",
         "@maven//:com_google_errorprone_error_prone_annotations",

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

@@ -74,8 +74,7 @@ public ExtensionRegistry getExtensionRegistry() {
   private CombinedDescriptorPool(ImmutableList<CelDescriptorPool> descriptorPools) {
     this.descriptorPools = descriptorPools;
     // TODO: Combine the extension registry. This will become necessary once we accept
-    // ExtensionRegistry through runtime builder. Ideally, proto team should open source this
-    // implementation but we may have to create our own.
+    // ExtensionRegistry through runtime builder.
     this.extensionRegistry =
         descriptorPools.stream()
             .map(CelDescriptorPool::getExtensionRegistry)

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

@@ -46,10 +46,7 @@ public final class DefaultDescriptorPool implements CelDescriptorPool {
 
   /** A DefaultDescriptorPool instance with just well known types loaded. */
   public static final DefaultDescriptorPool INSTANCE =
-      new DefaultDescriptorPool(
-          WELL_KNOWN_TYPE_DESCRIPTORS,
-          ImmutableMultimap.of(),
-          ExtensionRegistry.getEmptyRegistry());
+      new DefaultDescriptorPool(WELL_KNOWN_TYPE_DESCRIPTORS, ImmutableMultimap.of());
 
   // K: Fully qualified message type name, V: Message descriptor
   private final ImmutableMap<String, Descriptor> descriptorMap;
@@ -58,15 +55,7 @@ public final class DefaultDescriptorPool implements CelDescriptorPool {
   // V: Field descriptor for the extension message
   private final ImmutableMultimap<String, FieldDescriptor> extensionDescriptorMap;
 
-  @SuppressWarnings("Immutable") // ExtensionRegistry is immutable, just not marked as such.
-  private final ExtensionRegistry extensionRegistry;
-
   public static DefaultDescriptorPool create(CelDescriptors celDescriptors) {
-    return create(celDescriptors, ExtensionRegistry.getEmptyRegistry());
-  }
-
-  public static DefaultDescriptorPool create(
-      CelDescriptors celDescriptors, ExtensionRegistry extensionRegistry) {
     Map<String, Descriptor> descriptorMap = new HashMap<>(); // Using a hashmap to allow deduping
     stream(WellKnownProto.values()).forEach(d -> descriptorMap.put(d.typeName(), d.descriptor()));
 
@@ -75,9 +64,7 @@ public static DefaultDescriptorPool create(
     }
 
     return new DefaultDescriptorPool(
-        ImmutableMap.copyOf(descriptorMap),
-        celDescriptors.extensionDescriptors(),
-        extensionRegistry);
+        ImmutableMap.copyOf(descriptorMap), celDescriptors.extensionDescriptors());
   }
 
   @Override
@@ -96,15 +83,14 @@ public Optional<FieldDescriptor> findExtensionDescriptor(
 
   @Override
   public ExtensionRegistry getExtensionRegistry() {
-    return extensionRegistry;
+    // TODO: Populate one from runtime builder.
+    return ExtensionRegistry.getEmptyRegistry();
   }
 
   private DefaultDescriptorPool(
       ImmutableMap<String, Descriptor> descriptorMap,
-      ImmutableMultimap<String, FieldDescriptor> extensionDescriptorMap,
-      ExtensionRegistry extensionRegistry) {
+      ImmutableMultimap<String, FieldDescriptor> extensionDescriptorMap) {
     this.descriptorMap = checkNotNull(descriptorMap);
     this.extensionDescriptorMap = checkNotNull(extensionDescriptorMap);
-    this.extensionRegistry = checkNotNull(extensionRegistry);
   }
 }

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

@@ -15,16 +15,29 @@
 package dev.cel.common.internal;
 
 import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static java.util.Arrays.stream;
 
+import com.google.auto.value.AutoBuilder;
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.errorprone.annotations.CheckReturnValue;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.Any;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Descriptors.Descriptor;
+import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.DynamicMessage;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.Message;
+import dev.cel.common.CelDescriptorUtil;
+import dev.cel.common.CelDescriptors;
 import dev.cel.common.annotations.Internal;
+import dev.cel.common.types.CelTypes;
+import java.util.Map.Entry;
 import java.util.Optional;
+import org.jspecify.nullness.Nullable;
 
 /**
  * The {@code DynamicProto} class supports the conversion of {@link Any} values to concrete {@code
@@ -36,30 +49,59 @@
 @CheckReturnValue
 @Internal
 public final class DynamicProto {
+
+  private static final ImmutableMap<String, Descriptor> WELL_KNOWN_DESCRIPTORS =
+      stream(ProtoAdapter.WellKnownProto.values())
+          .collect(toImmutableMap(d -> d.typeName(), d -> d.descriptor()));
+
+  private final ImmutableMap<String, Descriptor> dynamicDescriptors;
+  private final ImmutableMultimap<String, FieldDescriptor> dynamicExtensionDescriptors;
   private final ProtoMessageFactory protoMessageFactory;
 
-  public static DynamicProto create(ProtoMessageFactory protoMessageFactory) {
-    return new DynamicProto(protoMessageFactory);
+  /** {@code ProtoMessageFactory} provides a method to create a protobuf builder objects by name. */
+  @Immutable
+  @FunctionalInterface
+  public interface ProtoMessageFactory {
+    Message.@Nullable Builder newBuilder(String messageName);
   }
 
-  DynamicProto(ProtoMessageFactory protoMessageFactory) {
-    this.protoMessageFactory = checkNotNull(protoMessageFactory);
+  /** Builder for configuring the {@link DynamicProto}. */
+  @AutoBuilder(ofClass = DynamicProto.class)
+  public abstract static class Builder {
+
+    /** Sets {@link CelDescriptors} to unpack any message types. */
+    public abstract Builder setDynamicDescriptors(CelDescriptors celDescriptors);
+
+    /** Sets a custom type factory to unpack any message types. */
+    public abstract Builder setProtoMessageFactory(ProtoMessageFactory factory);
+
+    /** Builds a new instance of {@link DynamicProto} */
+    @CheckReturnValue
+    public abstract DynamicProto build();
   }
 
-  /** Attempts to unpack an Any message. */
-  public Optional<Message> maybeUnpackAny(Message msg) {
-    try {
-      Any any =
-          msg instanceof Any
-              ? (Any) msg
-              : Any.parseFrom(
-                  msg.toByteString(),
-                  protoMessageFactory.getDescriptorPool().getExtensionRegistry());
-
-      return Optional.of(unpack(any));
-    } catch (InvalidProtocolBufferException e) {
-      return Optional.empty();
-    }
+  public static Builder newBuilder() {
+    return new AutoBuilder_DynamicProto_Builder()
+        .setDynamicDescriptors(CelDescriptors.builder().build())
+        .setProtoMessageFactory((typeName) -> null);
+  }
+
+  DynamicProto(
+      CelDescriptors dynamicDescriptors,
+      ProtoMessageFactory protoMessageFactory) {
+    ImmutableMap<String, Descriptor> messageTypeDescriptorMap =
+        CelDescriptorUtil.descriptorCollectionToMap(dynamicDescriptors.messageTypeDescriptors());
+    ImmutableMap<String, Descriptor> filteredDescriptors =
+        messageTypeDescriptorMap.entrySet().stream()
+            .filter(e -> !WELL_KNOWN_DESCRIPTORS.containsKey(e.getKey()))
+            .collect(toImmutableMap(Entry::getKey, Entry::getValue));
+    this.dynamicDescriptors =
+        ImmutableMap.<String, Descriptor>builder()
+            .putAll(WELL_KNOWN_DESCRIPTORS)
+            .putAll(filteredDescriptors)
+            .buildOrThrow();
+    this.dynamicExtensionDescriptors = checkNotNull(dynamicDescriptors.extensionDescriptors());
+    this.protoMessageFactory = checkNotNull(protoMessageFactory);
   }
 
   /**
@@ -78,8 +120,7 @@ public Message unpack(Any any) throws InvalidProtocolBufferException {
                         String.format("malformed type URL: %s", any.getTypeUrl())));
 
     Message.Builder builder =
-        protoMessageFactory
-            .newBuilder(messageTypeName)
+        newMessageBuilder(messageTypeName)
             .orElseThrow(
                 () ->
                     new InvalidProtocolBufferException(
@@ -95,7 +136,7 @@ public Message unpack(Any any) throws InvalidProtocolBufferException {
    */
   public Message maybeAdaptDynamicMessage(DynamicMessage input) {
     Optional<Message.Builder> maybeBuilder =
-        protoMessageFactory.newBuilder(input.getDescriptorForType().getFullName());
+        newMessageBuilder(input.getDescriptorForType().getFullName());
     if (!maybeBuilder.isPresent() || maybeBuilder.get() instanceof DynamicMessage.Builder) {
       // Just return the same input if:
       // 1. We didn't get a builder back because there's no descriptor (nothing we can do)
@@ -107,6 +148,61 @@ public Message maybeAdaptDynamicMessage(DynamicMessage input) {
     return merge(maybeBuilder.get(), input.toByteString());
   }
 
+  /**
+   * This method instantiates a builder for the given {@code typeName} assuming one is configured
+   * within the descriptor set provided to the {@code DynamicProto} constructor.
+   *
+   * <p>When the {@code useLinkedTypes} flag is set, the {@code Message.Builder} returned will be
+   * the concrete builder instance linked into the binary if it is present; otherwise, the result
+   * will be a {@code DynamicMessageBuilder}.
+   */
+  public Optional<Message.Builder> newMessageBuilder(String typeName) {
+    if (!CelTypes.isWellKnownType(typeName)) {
+      // Check if the message factory can produce a concrete message via custom type factory
+      // first.
+      Message.Builder builder = protoMessageFactory.newBuilder(typeName);
+      if (builder != null) {
+        return Optional.of(builder);
+      }
+    }
+
+    Optional<Descriptor> descriptor = maybeGetDescriptor(typeName);
+    if (!descriptor.isPresent()) {
+      return Optional.empty();
+    }
+    // If the descriptor that's resolved does not match the descriptor instance in the message
+    // factory, the call to fetch the prototype will return null, and a dynamic proto message
+    // should be used as a fallback.
+    Optional<Message> message =
+        DefaultInstanceMessageFactory.getInstance().getPrototype(descriptor.get());
+    if (message.isPresent()) {
+      return Optional.of(message.get().toBuilder());
+    }
+
+    // Fallback to a dynamic proto instance.
+    return Optional.of(DynamicMessage.newBuilder(descriptor.get()));
+  }
+
+  private Optional<Descriptor> maybeGetDescriptor(String typeName) {
+
+    Descriptor descriptor = ProtoRegistryProvider.getTypeRegistry().find(typeName);
+    return Optional.ofNullable(descriptor != null ? descriptor : dynamicDescriptors.get(typeName));
+  }
+
+  /** Gets the corresponding field descriptor for an extension field on a message. */
+  public Optional<FieldDescriptor> maybeGetExtensionDescriptor(
+      Descriptor containingDescriptor, String fieldName) {
+
+    String typeName = containingDescriptor.getFullName();
+    ImmutableCollection<FieldDescriptor> fieldDescriptors =
+        dynamicExtensionDescriptors.get(typeName);
+    if (fieldDescriptors.isEmpty()) {
+      return Optional.empty();
+    }
+
+    return fieldDescriptors.stream().filter(d -> d.getFullName().equals(fieldName)).findFirst();
+  }
+
   /**
    * Merge takes in a Message builder and merges another message bytes into the builder. Some
    * example usages are:
@@ -118,9 +214,7 @@ public Message maybeAdaptDynamicMessage(DynamicMessage input) {
    */
   private Message merge(Message.Builder builder, ByteString inputBytes) {
     try {
-      return builder
-          .mergeFrom(inputBytes, protoMessageFactory.getDescriptorPool().getExtensionRegistry())
-          .build();
+      return builder.mergeFrom(inputBytes, ProtoRegistryProvider.getExtensionRegistry()).build();
     } catch (InvalidProtocolBufferException e) {
       throw new AssertionError("Failed to merge input message into the message builder", e);
     }

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

@@ -143,6 +143,53 @@ public final class ProtoAdapter {
       stream(WellKnownProto.values())
           .collect(toImmutableMap(WellKnownProto::typeName, Function.identity()));
 
+  /**
+   * WellKnownProto types used throughout CEL. These types are specially handled to ensure that
+   * bidirectional conversion between CEL native values and these well-known types is performed
+   * consistently across runtimes.
+   */
+  enum WellKnownProto {
+    JSON_VALUE(Value.getDescriptor()),
+    JSON_STRUCT_VALUE(Struct.getDescriptor()),
+    JSON_LIST_VALUE(ListValue.getDescriptor()),
+    ANY_VALUE(Any.getDescriptor()),
+    BOOL_VALUE(BoolValue.getDescriptor(), true),
+    BYTES_VALUE(BytesValue.getDescriptor(), true),
+    DOUBLE_VALUE(DoubleValue.getDescriptor(), true),
+    FLOAT_VALUE(FloatValue.getDescriptor(), true),
+    INT32_VALUE(Int32Value.getDescriptor(), true),
+    INT64_VALUE(Int64Value.getDescriptor(), true),
+    STRING_VALUE(StringValue.getDescriptor(), true),
+    UINT32_VALUE(UInt32Value.getDescriptor(), true),
+    UINT64_VALUE(UInt64Value.getDescriptor(), true),
+    DURATION_VALUE(Duration.getDescriptor()),
+    TIMESTAMP_VALUE(Timestamp.getDescriptor());
+
+    private final Descriptor descriptor;
+    private final boolean isWrapperType;
+
+    WellKnownProto(Descriptor descriptor) {
+      this(descriptor, /* isWrapperType= */ false);
+    }
+
+    WellKnownProto(Descriptor descriptor, boolean isWrapperType) {
+      this.descriptor = descriptor;
+      this.isWrapperType = isWrapperType;
+    }
+
+    Descriptor descriptor() {
+      return descriptor;
+    }
+
+    String typeName() {
+      return descriptor.getFullName();
+    }
+
+    boolean isWrapperType() {
+      return isWrapperType;
+    }
+  }
+
   private final DynamicProto dynamicProto;
   private final boolean enableUnsignedLongs;