Java protobuf unverified lazy message extension field prototype

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

@@ -52,6 +52,7 @@ public class ExtensionRegistryLite {
   // parsing feature for MessageSet, which may be too crude for some
   // applications. Need to support this feature on smaller granularity.
   private static volatile boolean eagerlyParseMessageSets = false;
+  private static volatile boolean eagerlyParseExtensionFields = false;
 
   // Visible for testing.
   static final String EXTENSION_CLASS_NAME = "com.google.protobuf.Extension";
@@ -77,6 +78,14 @@ public static void setEagerlyParseMessageSets(boolean isEagerlyParse) {
     eagerlyParseMessageSets = isEagerlyParse;
   }
 
+  public static boolean isEagerlyParseExtensionFields() {
+    return eagerlyParseExtensionFields;
+  }
+
+  public static void setEagerlyParseExtensionFields(boolean isEagerlyParse) {
+    eagerlyParseExtensionFields = isEagerlyParse;
+  }
+
   /**
    * Construct a new, empty instance.
    *

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

@@ -272,6 +272,14 @@ public Object getField(final T descriptor) {
     return o;
   }
 
+  public LazyField getLazyField(final T descriptor) {
+    Object o = fields.get(descriptor);
+    if (o instanceof LazyField lazyField) {
+      return lazyField;
+    }
+    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);
@@ -1136,6 +1144,14 @@ public Object getField(final T descriptor) {
       return replaceBuilders(descriptor, value, true);
     }
 
+    public LazyField getLazyField(final T descriptor) {
+      Object value = fields.get(descriptor);
+      if (value instanceof LazyField lazyField) {
+        return lazyField;
+      }
+      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

@@ -1218,7 +1218,14 @@ public void writeUntil(final int end, final CodedOutputStream output) throws IOE
             // 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();

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

@@ -326,9 +326,19 @@ void mergeMessage(
         CodedInputStream input,
         ExtensionRegistryLite extensionRegistry,
         FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean lazy)
         throws IOException;
 
+    default void mergeMessage(
+        CodedInputStream input,
+        ExtensionRegistryLite extensionRegistry,
+        FieldDescriptor field,
+        Message defaultInstance)
+        throws IOException {
+      mergeMessage(input, extensionRegistry, field, defaultInstance, /* lazy= */ false);
+    }
+
     /** Returns the UTF8 validation level for the field. */
     WireFormat.Utf8Validation getUtf8Validation(Descriptors.FieldDescriptor descriptor);
 
@@ -564,7 +574,8 @@ public void mergeMessage(
         CodedInputStream input,
         ExtensionRegistryLite extensionRegistry,
         Descriptors.FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean lazy)
         throws IOException {
       if (!field.isRepeated()) {
         Message.Builder subBuilder;
@@ -796,15 +807,27 @@ public void mergeMessage(
         CodedInputStream input,
         ExtensionRegistryLite extensionRegistry,
         FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean lazy)
         throws IOException {
       if (!field.isRepeated()) {
         if (hasField(field)) {
+          // If the field is present and lazy, merge the bytes into the lazy field.
+          LazyField lazyField = extensions.getLazyField(field);
+          if (lazy && lazyField != null) {
+            lazyField.mergeFrom(input, extensionRegistry);
+            return;
+          }
           MessageLite.Builder current = ((MessageLite) getField(field)).toBuilder();
           input.readMessage(current, extensionRegistry);
           Object unused = setField(field, current.buildPartial());
           return;
         }
+        if (lazy) {
+          extensions.setField(
+              field, new LazyField(defaultInstance, extensionRegistry, input.readBytes()));
+          return;
+        }
         Message.Builder subBuilder = defaultInstance.newBuilderForType();
         input.readMessage(subBuilder, extensionRegistry);
         Object unused = setField(field, subBuilder.buildPartial());
@@ -1017,10 +1040,17 @@ public void mergeMessage(
         CodedInputStream input,
         ExtensionRegistryLite extensionRegistry,
         FieldDescriptor field,
-        Message defaultInstance)
+        Message defaultInstance,
+        boolean lazy)
         throws IOException {
       if (!field.isRepeated()) {
         if (hasField(field)) {
+          // If the field is present and lazy, merge the bytes into the lazy field.
+          LazyField lazyField = extensions.getLazyField(field);
+          if (lazy && lazyField != null) {
+            lazyField.mergeFrom(input, extensionRegistry);
+            return;
+          }
           Object fieldOrBuilder = extensions.getFieldAllowBuilders(field);
           MessageLite.Builder subBuilder;
           if (fieldOrBuilder instanceof MessageLite.Builder) {
@@ -1032,6 +1062,12 @@ public void mergeMessage(
           input.readMessage(subBuilder, extensionRegistry);
           return;
         }
+
+        if (lazy) {
+          extensions.setField(
+              field, new LazyField(defaultInstance, extensionRegistry, input.readBytes()));
+          return;
+        }
         Message.Builder subBuilder = defaultInstance.newBuilderForType();
         input.readMessage(subBuilder, extensionRegistry);
         Object unused = setField(field, subBuilder);
@@ -1210,7 +1246,13 @@ static boolean mergeFieldFrom(
           }
         case MESSAGE:
           {
-            target.mergeMessage(input, extensionRegistry, field, defaultInstance);
+            target.mergeMessage(
+                input,
+                extensionRegistry,
+                field,
+                defaultInstance,
+                !ExtensionRegistryLite.isEagerlyParseExtensionFields()
+                    && type.isExtensionNumber(fieldNumber));
             return true;
           }
         case ENUM:

java/core/src/test/java/com/google/protobuf/MessageTest.java

@@ -10,9 +10,12 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import proto2_unittest.UnittestProto;
 import proto2_unittest.UnittestProto.ForeignMessage;
 import proto2_unittest.UnittestProto.TestAllExtensions;
 import proto2_unittest.UnittestProto.TestAllTypes;
+import proto2_unittest.UnittestProto.TestLazyExtensions;
+import proto2_unittest.UnittestProto.TestLazyMessage;
 import proto2_unittest.UnittestProto.TestRequired;
 import proto2_unittest.UnittestProto.TestRequiredForeign;
 import java.util.List;
@@ -416,6 +419,23 @@ public void testUnpairedSurrogatesReplacedByQuestionMark() throws InvalidProtoco
         .isTrue();
   }
 
+  @Test
+  public void dummy() throws Exception {
+    ExtensionRegistry.setEagerlyParseExtensionFields(/* isEagerlyParse= */ false);
+    TestAllTypes.Builder builder = TestUtil.getAllSetBuilder();
+    TestLazyExtensions message =
+        TestLazyExtensions.newBuilder()
+            .setExtension(
+                UnittestProto.lazyExtensionsLazyMessage,
+                TestLazyMessage.newBuilder().setSubMessage(builder).build())
+            .build();
+    byte[] data = message.toByteArray();
+    TestLazyExtensions proto =
+        TestLazyExtensions.parseFrom(data, ExtensionRegistry.getGeneratedRegistry());
+    byte[] temp = proto.toByteArray();
+    System.out.println("temp: " + temp.length);
+  }
+
   private static String encodeHex(ByteString bytes) {
     String hexDigits = "0123456789abcdef";
     StringBuilder stringBuilder = new StringBuilder(bytes.size() * 2);