Refactor: Improve map entry parsing in GeneratedMessage.mm

This change adds validation for map key types and ensures that default values are used for missing key or value fields when parsing map entries from a stream. Error handling is also improved by checking the return value of `ReadMapEntryField`.

By adding a couple of judicious autoreleases, it also protects us from memory leaks in case of bad maps.

PiperOrigin-RevId: 895935026

Author: Dave MacLachlan

protobuf/runtime/src/com/google/protobuf/GeneratedMessage.mm

@@ -831,8 +831,7 @@ static BOOL AddMapGetWithKeyMethod(Class cls, SEL sel, CGPFieldDescriptor *field
           key.CGPValueField_##KEY_NAME = pKey;                                        \
           CGPValue value;                                                             \
           value.CGPValueField_##VALUE_NAME = pValue;                                  \
-          CGPMapFieldPut(MAP_FIELD_PTR(msg, offset), key, keyType, value, valueType,  \
-                         /* retainedKeyAndValue */ false);                            \
+          CGPMapFieldPut(MAP_FIELD_PTR(msg, offset), key, keyType, value, valueType); \
           return msg;                                                                 \
         });                                                                           \
   }
@@ -1674,28 +1673,37 @@ static BOOL MergeMapEntryFromStream(CGPMapField *field, CGPCodedInputStream *str
   if (!CGPReadInt32(stream, &length)) return NO;
   CGPCodedInputStream::Limit limit = stream->PushLimit(length);
   CGPFieldDescriptor *keyField = entry->fields_->buffer_[0];
+  CGPFieldJavaType keyType = CGPFieldGetJavaType(keyField);
   CGPFieldDescriptor *valueField = entry->fields_->buffer_[1];
-  BOOL hasKey = NO;
-  BOOL hasValue = NO;
+  CGPFieldJavaType valueType = CGPFieldGetJavaType(valueField);
   CGPValue key;
+  key.valueId = nil;
   CGPValue value;
+  if (CGPJavaTypeIsEnum(valueType)) {
+    // Note that even though value.valueId is an id, in the case of enums it is not retained.
+    // It is a "constant" value from the value descriptor that owns it.
+    value.valueId =
+        ((ComGoogleProtobufDescriptors_EnumValueDescriptor *)[valueField getDefaultValue])->enum_;
+  } else {
+    value.valueId = nil;
+  }
   while (YES) {
     uint32_t tag = stream->ReadTag();
     if (tag == 0) break;
     switch (CGPWireFormatGetTagFieldNumber(tag)) {
       case 1:
-        if (hasKey && CGPIsRetainedType(CGPFieldGetJavaType(keyField))) {
-          RELEASE_(key.valueId);
+        if (!ReadMapEntryField(stream, keyField, tag, registry, &key)) {
+          return NO;
+        } else if (CGPIsRetainedType(keyType)) {
+          AUTORELEASE(key.valueId);
         }
-        ReadMapEntryField(stream, keyField, tag, registry, &key);
-        hasKey = YES;
         break;
       case 2:
-        if (hasValue && CGPIsRetainedType(CGPFieldGetJavaType(valueField))) {
-          RELEASE_(value.valueId);
+        if (!ReadMapEntryField(stream, valueField, tag, registry, &value)) {
+          return NO;
+        } else if (CGPIsRetainedType(valueType)) {
+          AUTORELEASE(value.valueId);
         }
-        ReadMapEntryField(stream, valueField, tag, registry, &value);
-        hasValue = YES;
         break;
       default:
         if (!CGPWireFormatSkipField(stream, tag)) return NO;
@@ -1704,8 +1712,33 @@ static BOOL MergeMapEntryFromStream(CGPMapField *field, CGPCodedInputStream *str
   }
   if (!stream->ConsumedEntireMessage()) return NO;
   stream->PopLimit(limit);
-  CGPMapFieldPut(field, key, CGPFieldGetJavaType(keyField), value, CGPFieldGetJavaType(valueField),
-                 /* retainedKeyAndValue */ true);
+  if ((keyType == ComGoogleProtobufDescriptors_FieldDescriptor_JavaType_Enum_STRING) &&
+      (key.valueId == nil)) {
+    key.valueId = @"";
+  }
+  if ((CGPIsRetainedType(valueType)) && (value.valueId == nil)) {
+    switch (valueType) {
+      case ComGoogleProtobufDescriptors_FieldDescriptor_JavaType_Enum_STRING:
+        value.valueId = @"";
+        break;
+      case ComGoogleProtobufDescriptors_FieldDescriptor_JavaType_Enum_BYTE_STRING:
+        value.valueId = [CGPByteString empty];
+        break;
+      case ComGoogleProtobufDescriptors_FieldDescriptor_JavaType_Enum_MESSAGE:
+        if (valueField->valueType_ == nil) {
+          // Should not happen, but protect against crash.
+          return NO;
+        }
+        value.valueId = AUTORELEASE(CGPNewMessage(valueField->valueType_));
+        break;
+      default:
+        // Should not happen, but we don't trust potential undefined behavior
+        // in CGPIsRetainedType.
+        return NO;
+        break;
+    }
+  }
+  CGPMapFieldPut(field, key, keyType, value, valueType);
   return YES;
 }
 

protobuf/runtime/src/com/google/protobuf/MapField.h

@@ -106,10 +106,10 @@ CGP_ALWAYS_INLINE uint32_t CGPMapFieldMapSize(
 CGPMapFieldEntry *CGPMapFieldGetWithKey(
     CGPMapField *field, CGPValue key, CGPFieldJavaType keyType, CGPFieldJavaType valueType);
 
-// The caller will indicate whether the key and value are already retained.
+// The key and value are assumed to not be retained.
 void CGPMapFieldPut(
     CGPMapField *field, CGPValue key, CGPFieldJavaType keyType, CGPValue value,
-    CGPFieldJavaType valueType, bool retainedKeyAndValue);
+    CGPFieldJavaType valueType);
 
 void CGPMapFieldRemove(
     CGPMapField *field, CGPValue key, CGPFieldJavaType keyType, CGPFieldJavaType valueType);

protobuf/runtime/src/com/google/protobuf/MapField.m

@@ -322,11 +322,11 @@ static void EnsureAdditionalHashMapCapacity(
 
 void CGPMapFieldPut(
     CGPMapField *field, CGPValue key, CGPFieldJavaType keyType, CGPValue value,
-    CGPFieldJavaType valueType, bool retainedKeyAndValue) {
+    CGPFieldJavaType valueType) {
   BOOL keyTypeIsRetainable = CGPIsRetainedType(keyType);
   BOOL valueTypeIsRetainable = CGPIsRetainedType(valueType);
   // The value is always added to the map so make sure it's retained.
-  if (valueTypeIsRetainable && !retainedKeyAndValue) {
+  if (valueTypeIsRetainable) {
     RETAIN_(value.valueId);
   }
   uint32_t hash = Hash(key, keyType);
@@ -337,18 +337,14 @@ void CGPMapFieldPut(
     entry = GetFromHashArray(data, key, keyType, hash);
   }
   if (entry) {
-    // Existing entry so the key is not added to the map and must not be retained.
-    if (keyTypeIsRetainable && retainedKeyAndValue) {
-      [key.valueId autorelease];
-    }
     // Release the previous value.
     if (valueTypeIsRetainable) {
       [entry->value.valueId autorelease];
     }
     entry->value = value;
   } else {
     // Creating a new entry using the passed in key which must be retained.
-    if (keyTypeIsRetainable && !retainedKeyAndValue) {
+    if (keyTypeIsRetainable) {
       RETAIN_(key.valueId);
     }
     EnsureAdditionalHashMapCapacity(field, 1, keyType, valueType);

protobuf/tests/MapsTest.java

@@ -13,6 +13,7 @@
  */
 
 import com.google.j2objc.annotations.AutoreleasePool;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.Descriptors.FieldDescriptor.Type;
@@ -23,13 +24,25 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import protos.FakeScalarBytesMap;
+import protos.FakeScalarBytesMapFieldEntry;
+import protos.FakeScalarEnumMap;
+import protos.FakeScalarEnumMapFieldEntry;
+import protos.FakeScalarMsgMap;
+import protos.FakeScalarMsgMapFieldEntry;
+import protos.FakeStringStringMap;
+import protos.FakeStringStringMapFieldEntry;
 import protos.MapMsg;
 import protos.MapMsgOrBuilder;
 import protos.MapValue;
-
-/**
- * Tests for correct serialization and deserialization of map fields.
- */
+import protos.RandomMessage;
+import protos.RealScalarBytesMap;
+import protos.RealScalarEnumMap;
+import protos.RealScalarMsgMap;
+import protos.RealStringBytesMap;
+import protos.RealStringStringMap;
+
+/** Tests for correct serialization and deserialization of map fields. */
 public class MapsTest extends ProtobufTest {
 
   @AutoreleasePool
@@ -272,6 +285,119 @@ public void testEquals() throws Exception {
     assertEquals(msg1.hashCode(), msg2.hashCode());
   }
 
+  public void testStringStringRepeatedFieldsToMapConversions() throws Exception {
+    // According to https://protobuf.dev/programming-guides/proto3/#backwards repeated fields are
+    // converted to maps by using the first field as the key and the second field as the value.
+    // This also verifies that we can parse maps with missing fields by using default values.
+    FakeStringStringMap fakeMap =
+        FakeStringStringMap.newBuilder()
+            .addMapField(FakeStringStringMapFieldEntry.newBuilder().setKey("duck").build())
+            .addMapField(FakeStringStringMapFieldEntry.newBuilder().setValue("quack").build())
+            .addMapField(
+                FakeStringStringMapFieldEntry.newBuilder().setKey("cat").setValue("meow").build())
+            .build();
+    byte[] bytes = fakeMap.toByteArray();
+    RealStringStringMap realMap = RealStringStringMap.parseFrom(bytes);
+    assertEquals("", realMap.getMapFieldOrThrow("duck"));
+    assertEquals("quack", realMap.getMapFieldOrThrow(""));
+    assertEquals("meow", realMap.getMapFieldOrThrow("cat"));
+  }
+
+  public void testScalarBytesRepeatedFieldsToMapConversions() throws Exception {
+    // According to https://protobuf.dev/programming-guides/proto3/#backwards repeated fields are
+    // converted to maps by using the first field as the key and the second field as the value.
+    // This also verifies that we can parse maps with missing fields by using default values.
+    FakeScalarBytesMap fakeScalarBytesMap =
+        FakeScalarBytesMap.newBuilder()
+            .addMapField(FakeScalarBytesMapFieldEntry.newBuilder().setKey(42).build())
+            .addMapField(
+                FakeScalarBytesMapFieldEntry.newBuilder().setValue(ByteString.EMPTY).build())
+            .build();
+    byte[] bytes = fakeScalarBytesMap.toByteArray();
+    RealScalarBytesMap realScalarBytesMap = RealScalarBytesMap.parseFrom(bytes);
+    assertEquals(ByteString.EMPTY, realScalarBytesMap.getMapFieldOrThrow(42));
+    assertEquals(ByteString.EMPTY, realScalarBytesMap.getMapFieldOrThrow(0));
+  }
+
+  public void testScalarMsgRepeatedFieldsToMapConversions() throws Exception {
+    // According to https://protobuf.dev/programming-guides/proto3/#backwards repeated fields are
+    // converted to maps by using the first field as the key and the second field as the value.
+    // This also verifies that we can parse maps with missing fields by using default values.
+    FakeScalarMsgMap fakeScalarMsgMap =
+        FakeScalarMsgMap.newBuilder()
+            .addMapField(FakeScalarMsgMapFieldEntry.newBuilder().setKey(777).build())
+            .addMapField(
+                FakeScalarMsgMapFieldEntry.newBuilder()
+                    .setValue(RandomMessage.getDefaultInstance())
+                    .build())
+            .build();
+    byte[] bytes = fakeScalarMsgMap.toByteArray();
+    RealScalarMsgMap realScalarMsgMap = RealScalarMsgMap.parseFrom(bytes);
+    assertEquals(RandomMessage.getDefaultInstance(), realScalarMsgMap.getMapFieldOrThrow(777));
+    assertEquals(RandomMessage.getDefaultInstance(), realScalarMsgMap.getMapFieldOrThrow(0));
+  }
+
+  public void testScalarEnumRepeatedFieldsToMapConversions() throws Exception {
+    // According to https://protobuf.dev/programming-guides/proto3/#backwards repeated fields are
+    // converted to maps by using the first field as the key and the second field as the value.
+    // This also verifies that we can parse maps with missing fields by using default values.
+    // Note that the default value for an enum is the first enum value.
+    FakeScalarEnumMap fakeScalarEnumMap =
+        FakeScalarEnumMap.newBuilder()
+            .addMapField(FakeScalarEnumMapFieldEntry.newBuilder().setKey(123).build())
+            .addMapField(
+                FakeScalarEnumMapFieldEntry.newBuilder()
+                    .setKey(456)
+                    .setValue(MapMsg.Color.YELLOW)
+                    .build())
+            .build();
+    byte[] bytes = fakeScalarEnumMap.toByteArray();
+    RealScalarEnumMap realScalarEnumMap = RealScalarEnumMap.parseFrom(bytes);
+    assertEquals(MapMsg.Color.GREEN, realScalarEnumMap.getMapFieldOrThrow(123));
+    assertEquals(MapMsg.Color.YELLOW, realScalarEnumMap.getMapFieldOrThrow(456));
+  }
+
+  public void testBadValueType() throws Exception {
+    // Verifies that maps fail to parse if the value type is not the expected type.
+    FakeScalarBytesMap fakeMap =
+        FakeScalarBytesMap.newBuilder()
+            .addMapField(
+                FakeScalarBytesMapFieldEntry.newBuilder()
+                    .setKey(7)
+                    .setValue(ByteString.copyFromUtf8("hello"))
+                    .build())
+            .build();
+    byte[] bytes = fakeMap.toByteArray();
+    try {
+      RealScalarMsgMap realMap = RealScalarMsgMap.parseFrom(bytes);
+      fail("Expected exception instead of map: " + realMap);
+    } catch (Exception e) {
+      // Expected.
+    }
+  }
+
+  // This test is disabled because the java runtime doesn't throw an exception when the key type is
+  // not the expected type. Instead, it has undefined behavior with regards to what map you actually
+  // get. The j2objc runtime throws an exception.
+  public void disabledTestBadKeyType() throws Exception {
+    // Verifies that maps fail to parse if the key type is not the expected type.
+    FakeScalarBytesMap fakeMap =
+        FakeScalarBytesMap.newBuilder()
+            .addMapField(
+                FakeScalarBytesMapFieldEntry.newBuilder()
+                    .setKey(7)
+                    .setValue(ByteString.copyFromUtf8("hello"))
+                    .build())
+            .build();
+    byte[] bytes = fakeMap.toByteArray();
+    try {
+      RealStringBytesMap realMap = RealStringBytesMap.parseFrom(bytes);
+      fail("Expected exception instead of map: " + realMap);
+    } catch (Exception e) {
+      // Expected.
+    }
+  }
+
   public void testToString() throws Exception {
     String result = getFilledMessage().toString();
     assertTrue(result.contains("int_int"));

protobuf/tests/protos/map_fields.proto

@@ -38,7 +38,71 @@ message MapMsg {
 }
 
 message MapValue {
-  string foo = 1 [
-    features.field_presence = LEGACY_REQUIRED
-  ];
+  string foo = 1 [features.field_presence = LEGACY_REQUIRED];
+}
+
+// Verifying conversion from RepeatedFields to Maps for string, string.
+message FakeStringStringMapFieldEntry {
+  string key = 1;
+  string value = 2;
+}
+
+message FakeStringStringMap {
+  repeated FakeStringStringMapFieldEntry map_field = 1;
+}
+
+message RealStringStringMap {
+  map<string, string> map_field = 1;
+}
+
+// Verifying conversion from RepeatedFields to Maps for scalar, bytes.
+message FakeScalarBytesMapFieldEntry {
+  int32 key = 1;
+  bytes value = 2;
+}
+
+message FakeScalarBytesMap {
+  repeated FakeScalarBytesMapFieldEntry map_field = 1;
+}
+
+message RealScalarBytesMap {
+  map<int32, bytes> map_field = 1;
+}
+
+// Verifying conversion from RepeatedFields to Maps for scalar, msg.
+message RandomMessage {
+  string foo = 1;
+}
+
+message FakeScalarMsgMapFieldEntry {
+  int32 key = 1;
+  RandomMessage value = 2;
+}
+
+message FakeScalarMsgMap {
+  repeated FakeScalarMsgMapFieldEntry map_field = 1;
+}
+
+message RealScalarMsgMap {
+  map<int32, RandomMessage> map_field = 1;
+}
+
+// Verifying conversion from RepeatedFields to Maps for scalar, enum.
+// Specifically making sure the default value is used when the value is
+// missing.
+message FakeScalarEnumMapFieldEntry {
+  int32 key = 1;
+  MapMsg.Color value = 2;
+}
+
+message FakeScalarEnumMap {
+  repeated FakeScalarEnumMapFieldEntry map_field = 1;
+}
+
+message RealScalarEnumMap {
+  map<int32, MapMsg.Color> map_field = 1;
+}
+
+message RealStringBytesMap {
+  map<string, bytes> map_field = 1;
 }