Force re-computing the serialized size after the lazy fields with corrupted data are replaced with empty bytes

PiperOrigin-RevId: 698781661

Author: protobuf-github-bot

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

@@ -269,6 +269,12 @@ public Object getField(final T descriptor) {
     return o;
   }
 
+  /** Returns true if the field is a lazy field and it is corrupted. */
+  boolean lazyFieldCorrupted(final T descriptor) {
+    Object o = fields.get(descriptor);
+    return o instanceof LazyField && ((LazyField) o).isCorrupted();
+  }
+
   /**
    * Useful for implementing {@link Message.Builder#setField(Descriptors.FieldDescriptor,Object)}.
    */

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

@@ -986,17 +986,26 @@ public final <T> T getExtension(final ExtensionLite<? extends MessageT, T> exten
       verifyExtensionContainingType(extension);
       FieldDescriptor descriptor = extension.getDescriptor();
       final Object value = extensions.getField(descriptor);
+      T result = null;
       if (value == null) {
         if (descriptor.isRepeated()) {
-          return (T) ProtobufArrayList.emptyList();
+          result = (T) ProtobufArrayList.emptyList();
         } else if (descriptor.getJavaType() == FieldDescriptor.JavaType.MESSAGE) {
-          return (T) extension.getMessageDefaultInstance();
+          result = (T) extension.getMessageDefaultInstance();
         } else {
-          return (T) extension.fromReflectionType(descriptor.getDefaultValue());
+          result = (T) extension.fromReflectionType(descriptor.getDefaultValue());
         }
       } else {
-        return (T) extension.fromReflectionType(value);
+        result = (T) extension.fromReflectionType(value);
       }
+
+      // If the lazy field is corrupted, we need to invalidate the memoized size in case the
+      // corrupted message data was replaced with an empty ByteString and yet a previous serialized
+      // size was memoized.
+      if (extensions.lazyFieldCorrupted(descriptor)) {
+        setMemoizedSerializedSize(-1);
+      }
+      return result;
     }
 
     /** Get one element of a repeated extension. */

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

@@ -89,6 +89,8 @@ public class LazyFieldLite {
    */
   private volatile ByteString memoizedBytes;
 
+  private volatile boolean corrupted;
+
   /** Constructs a LazyFieldLite with bytes that will be parsed lazily. */
   public LazyFieldLite(ExtensionRegistryLite extensionRegistry, ByteString bytes) {
     checkArguments(extensionRegistry, bytes);
@@ -400,6 +402,7 @@ protected void ensureInitialized(MessageLite defaultInstance) {
       } catch (InvalidProtocolBufferException e) {
         // Nothing is logged and no exceptions are thrown. Clients will be unaware that this proto
         // was invalid.
+        this.corrupted = true;
         this.value = defaultInstance;
         this.memoizedBytes = ByteString.EMPTY;
       }
@@ -414,4 +417,9 @@ private static void checkArguments(ExtensionRegistryLite extensionRegistry, Byte
       throw new NullPointerException("found null ByteString");
     }
   }
+
+  /** Returns whether the lazy field was corrupted and replaced with an empty message. */
+  boolean isCorrupted() {
+    return corrupted;
+  }
 }

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

@@ -159,4 +159,42 @@ public void testLoadCorruptedLazyField_getsReplacedWithEmptyMessage() throws Exc
 
     assertThat(actualRaw).isEqualTo(expectedRaw);
   }
+
+  @Test
+  public void testLoadCorruptedLazyField_getSerializedSize() throws Exception {
+    ExtensionRegistry extensionRegistry = ExtensionRegistry.newInstance();
+    extensionRegistry.add(TestMessageSetExtension1.messageSetExtension);
+    RawMessageSet inputRaw =
+        RawMessageSet.newBuilder()
+            .addItem(
+                RawMessageSet.Item.newBuilder()
+                    .setTypeId(TYPE_ID_1)
+                    .setMessage(CORRUPTED_MESSAGE_PAYLOAD))
+            .build();
+    ByteString inputData = inputRaw.toByteString();
+    TestMessageSet messageSet = TestMessageSet.parseFrom(inputData, extensionRegistry);
+
+    // Effectively cache the serialized size of the message set.
+    assertThat(messageSet.getSerializedSize()).isEqualTo(9);
+
+    // getExtension should mark the memoized size as "dirty" (i.e. -1).
+    assertThat(messageSet.getExtension(TestMessageSetExtension1.messageSetExtension))
+        .isEqualTo(TestMessageSetExtension1.getDefaultInstance());
+
+    // toByteString calls getSerializedSize() which should re-compute the serialized size as the
+    // message contains lazy fields.
+    ByteString outputData = messageSet.toByteString();
+
+    // Re-parse as RawMessageSet
+    RawMessageSet actualRaw =
+        RawMessageSet.parseFrom(outputData, ExtensionRegistry.getEmptyRegistry());
+
+    RawMessageSet expectedRaw =
+        RawMessageSet.newBuilder()
+            .addItem(
+                RawMessageSet.Item.newBuilder().setTypeId(TYPE_ID_1).setMessage(ByteString.empty()))
+            .build();
+
+    assertThat(actualRaw).isEqualTo(expectedRaw);
+  }
 }