From 44bd03c609425ea29f7c52ad8b251563e4f70698 Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Tue, 14 Mar 2023 16:02:35 -0700
Subject: [PATCH] Fix `FieldSet#writeMessageSetTo` to delegate to LazyFields
 when appropriate.

This fixes the following issue:
- In `FieldSet#getMessageSetSerializedSize` we delegate to `LazyField` for getting its size.
- Before this change, in `FieldSet#writeMessageTo` we do not delegate to LazyField's `byteString` method. Instead, we call `getValue` on the `FieldSet`.
- It's not guaranteed that `fieldSet.getSerializedSize() == fieldSet.getValue().getSerializedSize()`.

In particular, this can lead us to run into the message "Serializing com.google.protobuf.DynamicMessage to a byte array threw an IOException (should never happen)" because of mispredicted output size.

We fix the issue by following similar logic to `MessageSetSchema`- if "getSerializedSize" delegates to the LazyField, then the serializer should as well.

PiperOrigin-RevId: 516658372
---
 .../java/com/google/protobuf/FieldSet.java    |  6 ++-
 .../google/protobuf/DynamicMessageTest.java   | 38 +++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/java/core/src/main/java/com/google/protobuf/FieldSet.java b/java/core/src/main/java/com/google/protobuf/FieldSet.java
index 7bb3c9aadf..af7ce87e27 100644
--- a/java/core/src/main/java/com/google/protobuf/FieldSet.java
+++ b/java/core/src/main/java/com/google/protobuf/FieldSet.java
@@ -602,9 +602,11 @@ final class FieldSet<T extends FieldSet.FieldDescriptorLite<T>> {
         && !descriptor.isPacked()) {
       Object value = entry.getValue();
       if (value instanceof LazyField) {
-        value = ((LazyField) value).getValue();
+        ByteString valueBytes = ((LazyField) value).toByteString();
+        output.writeRawMessageSetExtension(entry.getKey().getNumber(), valueBytes);
+      } else {
+        output.writeMessageSetExtension(entry.getKey().getNumber(), (MessageLite) value);
       }
-      output.writeMessageSetExtension(entry.getKey().getNumber(), (MessageLite) value);
     } else {
       writeField(descriptor, entry.getValue(), output);
     }
diff --git a/java/core/src/test/java/com/google/protobuf/DynamicMessageTest.java b/java/core/src/test/java/com/google/protobuf/DynamicMessageTest.java
index abe7cafa79..3afdc1b4d2 100644
--- a/java/core/src/test/java/com/google/protobuf/DynamicMessageTest.java
+++ b/java/core/src/test/java/com/google/protobuf/DynamicMessageTest.java
@@ -38,12 +38,14 @@ import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.Descriptors.OneofDescriptor;
 import dynamicmessagetest.DynamicMessageTestProto.EmptyMessage;
 import dynamicmessagetest.DynamicMessageTestProto.MessageWithMapFields;
+import protobuf_unittest.UnittestMset.TestMessageSetExtension2;
 import protobuf_unittest.UnittestProto;
 import protobuf_unittest.UnittestProto.TestAllExtensions;
 import protobuf_unittest.UnittestProto.TestAllTypes;
 import protobuf_unittest.UnittestProto.TestAllTypes.NestedMessage;
 import protobuf_unittest.UnittestProto.TestEmptyMessage;
 import protobuf_unittest.UnittestProto.TestPackedTypes;
+import proto2_wireformat_unittest.UnittestMsetWireFormat.TestMessageSet;
 import java.util.ArrayList;
 import org.junit.Test;
 import org.junit.function.ThrowingRunnable;
@@ -401,4 +403,40 @@ public class DynamicMessageTest {
           }
         });
   }
+
+  @Test
+  public void serialize_lazyFieldInMessageSet() throws Exception {
+    ExtensionRegistry extensionRegistry = ExtensionRegistry.newInstance();
+    extensionRegistry.add(TestMessageSetExtension2.messageSetExtension);
+    TestMessageSetExtension2 messageSetExtension =
+        TestMessageSetExtension2.newBuilder().setStr("foo").build();
+    // This is a valid serialization of the above message.
+    ByteString suboptimallySerializedMessageSetExtension =
+        messageSetExtension.toByteString().concat(messageSetExtension.toByteString());
+    DynamicMessage expectedMessage =
+        DynamicMessage.newBuilder(TestMessageSet.getDescriptor())
+            .setField(
+                TestMessageSetExtension2.messageSetExtension.getDescriptor(), messageSetExtension)
+            .build();
+    // Constructed with a LazyField, for whom roundtripping the serialized form will shorten the
+    // encoded form.
+    // In particular, this checks matching between lazy field encoding size.
+    DynamicMessage complicatedlyBuiltMessage =
+        DynamicMessage.newBuilder(TestMessageSet.getDescriptor())
+            .setField(
+                TestMessageSetExtension2.messageSetExtension.getDescriptor(),
+                new LazyField(
+                    DynamicMessage.getDefaultInstance(TestMessageSetExtension2.getDescriptor()),
+                    extensionRegistry,
+                    suboptimallySerializedMessageSetExtension))
+            .build();
+
+    DynamicMessage roundtrippedMessage =
+        DynamicMessage.newBuilder(TestMessageSet.getDescriptor())
+            .mergeFrom(complicatedlyBuiltMessage.toByteString(), extensionRegistry)
+            .build();
+
+    assertThat(complicatedlyBuiltMessage).isEqualTo(expectedMessage);
+    assertThat(roundtrippedMessage).isEqualTo(expectedMessage);
+  }
 }