Java protobuf verified lazy message extension field prototype

PiperOrigin-RevId: 823639590

Author: protobuf-github-bot

java/core/BUILD.bazel

@@ -1,6 +1,5 @@
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
 load("@rules_java//java:java_library.bzl", "java_library")
-load("@rules_pkg//pkg:mappings.bzl", "pkg_files", "strip_prefix")
 load("//:protobuf.bzl", "internal_gen_well_known_protos_java")
 load("//:protobuf_version.bzl", "PROTOBUF_JAVA_VERSION")
 load("//bazel:cc_proto_library.bzl", "cc_proto_library")
@@ -563,6 +562,7 @@ LITE_TEST_EXCLUSIONS = [
     "src/test/java/com/google/protobuf/LegacyUnredactedTextFormatTest.java",
     "src/test/java/com/google/protobuf/MapForProto2Test.java",
     "src/test/java/com/google/protobuf/MapTest.java",
+    "src/test/java/com/google/protobuf/LazyExtensionFieldsTest.java",
     "src/test/java/com/google/protobuf/MessageTest.java",
     "src/test/java/com/google/protobuf/NestedBuildersTest.java",
     "src/test/java/com/google/protobuf/PackedFieldTest.java",
@@ -654,6 +654,7 @@ junit_tests(
             "src/test/java/com/google/protobuf/TestUtil.java",
             "src/test/java/com/google/protobuf/TestUtilLite.java",
             "src/test/java/com/google/protobuf/RuntimeVersionTest.java",
+            "src/test/java/com/google/protobuf/LazyExtensionFieldsTest.java",
         ],
     ),
     test_prefix = "v25SrcJar",
@@ -708,6 +709,7 @@ junit_tests(
             "src/test/java/com/google/protobuf/TestUtil.java",
             "src/test/java/com/google/protobuf/TestUtilLite.java",
             "src/test/java/com/google/protobuf/RuntimeVersionTest.java",
+            "src/test/java/com/google/protobuf/LazyExtensionFieldsTest.java",
         ],
     ),
     test_prefix = "v25Jar",

java/core/src/main/java/com/google/protobuf/CodedInputStream.java

@@ -189,6 +189,22 @@ private CodedInputStream() {}
    */
   public abstract boolean skipField(final int tag) throws IOException;
 
+  /**
+   * Reads and discards a single varint field.
+   *
+   * @throws InvalidProtocolBufferException if the varint is malformed.
+   */
+  abstract void skipRawVarint() throws IOException;
+
+  /**
+   * Reads and discards a single UTF-8 string field.
+   *
+   * @throws InvalidProtocolBufferException if the string is not valid UTF-8.
+   */
+  void skipUtf8String(int length) throws IOException {
+    throw new UnsupportedOperationException("Utf8 string skipping is not implemented.");
+  }
+
   /**
    * Reads a single field and writes it to output in wire format, given its tag value.
    *
@@ -1061,14 +1077,31 @@ public int readRawVarint32() throws IOException {
       return (int) readRawVarint64SlowPath();
     }
 
-    private void skipRawVarint() throws IOException {
+    @Override
+    void skipRawVarint() throws IOException {
       if (limit - pos >= MAX_VARINT_SIZE) {
         skipRawVarintFastPath();
       } else {
         skipRawVarintSlowPath();
       }
     }
 
+    @Override
+    void skipUtf8String(int length) throws IOException {
+      if (length >= 0 && length <= (limit - pos)) {
+        if (!Utf8.isValidUtf8(buffer, pos, pos + length)) {
+          throw InvalidProtocolBufferException.invalidUtf8();
+        }
+        pos += length;
+        return;
+      }
+
+      if (length < 0) {
+        throw InvalidProtocolBufferException.negativeSize();
+      }
+      throw InvalidProtocolBufferException.truncatedMessage();
+    }
+
     private void skipRawVarintFastPath() throws IOException {
       for (int i = 0; i < MAX_VARINT_SIZE; i++) {
         if (buffer[pos++] >= 0) {
@@ -1834,7 +1867,8 @@ public int readRawVarint32() throws IOException {
       return (int) readRawVarint64SlowPath();
     }
 
-    private void skipRawVarint() throws IOException {
+    @Override
+    void skipRawVarint() throws IOException {
       if (bufferSize - pos >= MAX_VARINT_SIZE) {
         skipRawVarintFastPath();
       } else {

java/core/src/main/java/com/google/protobuf/ExtensionRegistryLite.java

@@ -53,6 +53,59 @@ public class ExtensionRegistryLite {
   // applications. Need to support this feature on smaller granularity.
   private static volatile boolean eagerlyParseMessageSets = false;
 
+  enum LazyExtensionFieldsExperimentMode {
+    EAGER("EAGER"),
+    // Throw exceptions on malformed data, but will defer parsing until first access.
+    VERIFIED_LAZY("VERIFIED_LAZY"),
+    // Caution: This mode is unsafe as it postpone parsing errors such as required fields missing
+    // until first access.
+    UNVERIFIED_LAZY("UNVERIFIED_LAZY");
+
+    private final String name;
+
+    LazyExtensionFieldsExperimentMode(String name) {
+      this.name = name;
+    }
+
+    String getName() {
+      return name;
+    }
+
+    @SuppressWarnings("StatementSwitchToExpressionSwitch")
+    static LazyExtensionFieldsExperimentMode fromString(String name) {
+      switch (name) {
+        case "EAGER":
+          return EAGER;
+        case "VERIFIED_LAZY":
+          return VERIFIED_LAZY;
+        case "UNVERIFIED_LAZY":
+          return UNVERIFIED_LAZY;
+        default:
+          throw new IllegalArgumentException("Unknown LazyExtensionFieldsExperimentMode: " + name);
+      }
+    }
+  }
+
+  private static volatile LazyExtensionFieldsExperimentMode lazyExtensionFieldsExperimentMode =
+      LazyExtensionFieldsExperimentMode.EAGER;
+
+  static void setLazyExtensionFieldsExperimentMode(LazyExtensionFieldsExperimentMode mode) {
+    lazyExtensionFieldsExperimentMode = mode;
+  }
+
+  static LazyExtensionFieldsExperimentMode getLazyExtensionFieldsExperimentMode() {
+    return lazyExtensionFieldsExperimentMode;
+  }
+
+  static boolean lazyExtensionFieldEnabled() {
+    return !lazyExtensionFieldsExperimentMode.equals(LazyExtensionFieldsExperimentMode.EAGER);
+  }
+
+  static boolean lazyExtensionFieldValidationEnabled() {
+    return lazyExtensionFieldsExperimentMode.equals(
+        LazyExtensionFieldsExperimentMode.VERIFIED_LAZY);
+  }
+
   // Visible for testing.
   static final String EXTENSION_CLASS_NAME = "com.google.protobuf.Extension";
 

java/core/src/main/java/com/google/protobuf/FieldSet.java

@@ -273,6 +273,15 @@ public Object getField(final T descriptor) {
     return o;
   }
 
+  @SuppressWarnings({"ReturnMissingNullable", "PatternMatchingInstanceof"})
+  LazyField getLazyField(final T descriptor) {
+    Object o = fields.get(descriptor);
+    if (o instanceof LazyField) {
+      return (LazyField) o;
+    }
+    return null;
+  }
+
   /** Returns true if the field is a lazy field and it is corrupted. */
   boolean lazyFieldCorrupted(final T descriptor) {
     Object o = fields.get(descriptor);
@@ -493,6 +502,9 @@ private static boolean isMessageFieldValueInitialized(Object value) {
       // this method is used by FieldSet.Builder.isInitialized.
       return ((MessageLiteOrBuilder) value).isInitialized();
     } else if (value instanceof LazyField) {
+      if (ExtensionRegistryLite.lazyExtensionFieldValidationEnabled()) {
+        return !((LazyField) value).containsMissingRequiredFields();
+      }
       return true;
     } else {
       throw new IllegalArgumentException(
@@ -1137,6 +1149,15 @@ public Object getField(final T descriptor) {
       return replaceBuilders(descriptor, value, true);
     }
 
+    @SuppressWarnings({"ReturnMissingNullable", "PatternMatchingInstanceof"})
+    public LazyField getLazyField(final T descriptor) {
+      Object value = fields.get(descriptor);
+      if (value instanceof LazyField) {
+        return (LazyField) value;
+      }
+      return null;
+    }
+
     /** Same as {@link #getField(F)}, but allow a {@link MessageLite.Builder} to be returned. */
     Object getFieldAllowBuilders(final T descriptor) {
       Object o = fields.get(descriptor);

java/core/src/main/java/com/google/protobuf/GeneratedMessage.java

@@ -1180,14 +1180,14 @@ public void writeUntil(final int end, final CodedOutputStream output) throws IOE
               output.writeMessageSetExtension(descriptor.getNumber(), (Message) next.getValue());
             }
           } else {
-            // TODO: Taken care of following code, it may cause
-            // problem when we use LazyField for normal fields/extensions.
-            // Due to the optional field can be duplicated at the end of
-            // serialized bytes, which will make the serialized size change
-            // after lazy field parsed. So when we use LazyField globally,
-            // we need to change the following write method to write cached
-            // bytes directly rather than write the parsed message.
-            FieldSet.writeField(descriptor, next.getValue(), output);
+            if (descriptor.getLiteJavaType() == WireFormat.JavaType.MESSAGE
+                && next instanceof LazyField.LazyEntry<?>) {
+              output.writeBytes(
+                  descriptor.getNumber(),
+                  ((LazyField.LazyEntry<?>) next).getField().toByteString());
+            } else {
+              FieldSet.writeField(descriptor, next.getValue(), output);
+            }
           }
           if (iter.hasNext()) {
             next = iter.next();
@@ -2014,13 +2014,21 @@ protected Object fromReflectionType(final Object value) {
      * single element.
      */
     @Override
+    @SuppressWarnings("PatternMatchingInstanceof")
     protected Object singularFromReflectionType(final Object value) {
       FieldDescriptor descriptor = getDescriptor();
       switch (descriptor.getJavaType()) {
         case MESSAGE:
           if (singularType.isInstance(value)) {
             return value;
           } else {
+            if (value instanceof LazyField) {
+              LazyField lazyField = (LazyField) value;
+              return messageDefaultInstance
+                  .newBuilderForType()
+                  .mergeFrom(lazyField.getValue())
+                  .build();
+            }
             return messageDefaultInstance.newBuilderForType().mergeFrom((Message) value).build();
           }
         case ENUM:

java/core/src/main/java/com/google/protobuf/LazyFieldLite.java

@@ -91,6 +91,14 @@ public class LazyFieldLite {
 
   private volatile boolean corrupted;
 
+  /**
+   * Required field presence is usually checked by calling `isInitialized` on the parsed message
+   * instead of at decoding phase (when LazyField is initialized), so we need to defer the check
+   * until it is actually needed. For example, `builder.mergeFrom` will not trigger the check, but
+   * the following `builder.build()` will.
+   */
+  private volatile boolean containsMissingRequiredFields;
+
   /**
    * Carry a message's default instance which is used by {@code hashCode()}, {@code equals()}, and
    * {@code toString()}. Can be null.
@@ -118,6 +126,14 @@ public LazyFieldLite() {
     defaultInstance = null;
   }
 
+  public boolean containsMissingRequiredFields() {
+    return containsMissingRequiredFields;
+  }
+
+  public void setContainsMissingRequiredFields(boolean containsMissingRequiredFields) {
+    this.containsMissingRequiredFields = containsMissingRequiredFields;
+  }
+
   /**
    * Constructs a LazyFieldLite instance with a value. The LazyFieldLite may not be able to parse
    * the extensions in the value as it has no ExtensionRegistry.
@@ -331,10 +347,9 @@ public void merge(LazyFieldLite other) {
    * <p>LazyField is not thread-safe for write access. Synchronizations are needed under read/write
    * situations.
    */
-  public void mergeFrom(CodedInputStream input, ExtensionRegistryLite extensionRegistry)
-      throws IOException {
+  void mergeFrom(ByteString bytes, ExtensionRegistryLite extensionRegistry) throws IOException {
     if (this.containsDefaultInstance()) {
-      setByteString(input.readBytes(), extensionRegistry);
+      setByteString(bytes, extensionRegistry);
       return;
     }
 
@@ -350,21 +365,29 @@ public void mergeFrom(CodedInputStream input, ExtensionRegistryLite extensionReg
     // to outway the benefits of combining the extension registries, which is not normally done for
     // lite protos anyways.
     if (this.delayedBytes != null) {
-      setByteString(this.delayedBytes.concat(input.readBytes()), this.extensionRegistry);
+      setByteString(this.delayedBytes.concat(bytes), this.extensionRegistry);
       return;
     }
 
     // We are parsed and both contain data. We won't drop any extensions here directly, but in the
     // case that the extension registries are not the same then we might in the future if we
     // need to serialize and parse a message again.
     try {
-      setValue(value.toBuilder().mergeFrom(input, extensionRegistry).build());
+      setValue(value.toBuilder().mergeFrom(bytes, extensionRegistry).build());
     } catch (InvalidProtocolBufferException e) {
       // Nothing is logged and no exceptions are thrown. Clients will be unaware that a proto
       // was invalid.
     }
   }
 
+  /**
+   * Merges another instance's contents from a stream.
+   */
+  public void mergeFrom(CodedInputStream input, ExtensionRegistryLite extensionRegistry)
+      throws IOException {
+    mergeFrom(input.readBytes(), extensionRegistry);
+  }
+
   private static MessageLite mergeValueAndBytes(
       MessageLite value, ByteString otherBytes, ExtensionRegistryLite extensionRegistry) {
     try {
@@ -477,8 +500,15 @@ protected void ensureInitialized(MessageLite defaultInstance) {
       try {
         if (delayedBytes != null) {
           // The extensionRegistry shouldn't be null here since we have delayedBytes.
+          // We don't enforce required field presence in lazy mode here, as `isInitialized()` at
+          // `build()` or `parseFrom()` has already checked missing required fields before this
+          // step, so exception has already been thrown or `parsePartialFrom()` is used.
           MessageLite parsedValue =
-              defaultInstance.getParserForType().parseFrom(delayedBytes, extensionRegistry);
+              ExtensionRegistryLite.lazyExtensionFieldValidationEnabled()
+                  ? defaultInstance
+                      .getParserForType()
+                      .parsePartialFrom(delayedBytes, extensionRegistry)
+                  : defaultInstance.getParserForType().parseFrom(delayedBytes, extensionRegistry);
           this.value = parsedValue;
           this.memoizedBytes = delayedBytes;
         } else {