Merge pull request #544 from haberman/unknown_map_enum

Put map entries into unknown fields when unknown enum values are encountered

upb/decode.c

@@ -385,6 +385,18 @@ static char* encode_varint32(uint32_t val, char* ptr) {
   return ptr;
 }
 
+static void upb_Decode_AddUnknownVarints(upb_Decoder* d, upb_Message* msg,
+                                         uint32_t val1, uint32_t val2) {
+  char buf[20];
+  char* end = buf;
+  end = encode_varint32(val1, end);
+  end = encode_varint32(val2, end);
+
+  if (!_upb_Message_AddUnknown(msg, buf, end - buf, &d->arena)) {
+    decode_err(d, kUpb_DecodeStatus_OutOfMemory);
+  }
+}
+
 UPB_NOINLINE
 static bool decode_checkenum_slow(upb_Decoder* d, const char* ptr,
                                   upb_Message* msg, const upb_MiniTable_Enum* e,
@@ -398,17 +410,9 @@ static bool decode_checkenum_slow(upb_Decoder* d, const char* ptr,
 
   // Unrecognized enum goes into unknown fields.
   // For packed fields the tag could be arbitrarily far in the past, so we
-  // just re-encode the tag here.
-  char buf[20];
-  char* end = buf;
+  // just re-encode the tag and value here.
   uint32_t tag = ((uint32_t)field->number << 3) | kUpb_WireType_Varint;
-  end = encode_varint32(tag, end);
-  end = encode_varint32(v, end);
-
-  if (!_upb_Message_AddUnknown(msg, buf, end - buf, &d->arena)) {
-    decode_err(d, kUpb_DecodeStatus_OutOfMemory);
-  }
-
+  upb_Decode_AddUnknownVarints(d, msg, tag, v);
   return false;
 }
 
@@ -627,8 +631,20 @@ static const char* decode_tomap(upb_Decoder* d, const char* ptr,
         upb_value_ptr(_upb_Message_New(entry->subs[0].submsg, &d->arena));
   }
 
+  const char* start = ptr;
   ptr = decode_tosubmsg(d, ptr, &ent.k, subs, field, val->size);
-  _upb_Map_Set(map, &ent.k, map->key_size, &ent.v, map->val_size, &d->arena);
+  // check if ent had any unknown fields
+  size_t size;
+  upb_Message_GetUnknown(&ent.k, &size);
+  if (size != 0) {
+    uint32_t tag = ((uint32_t)field->number << 3) | kUpb_WireType_Delimited;
+    upb_Decode_AddUnknownVarints(d, msg, tag, (uint32_t)(ptr - start));
+    if (!_upb_Message_AddUnknown(msg, start, ptr - start, &d->arena)) {
+      decode_err(d, kUpb_DecodeStatus_OutOfMemory);
+    }
+  } else {
+    _upb_Map_Set(map, &ent.k, map->key_size, &ent.v, map->val_size, &d->arena);
+  }
   return ptr;
 }
 

upb/msg_test.cc

@@ -398,3 +398,32 @@ TEST(MessageTest, MaxRequiredFields) {
       test_msg, kUpb_Encode_CheckRequired, arena.ptr(), &size);
   ASSERT_TRUE(serialized != nullptr);
 }
+
+TEST(MessageTest, MapField) {
+  upb::Arena arena;
+  upb_test_TestMapFieldExtra* test_msg_extra =
+      upb_test_TestMapFieldExtra_new(arena.ptr());
+
+  ASSERT_TRUE(upb_test_TestMapFieldExtra_map_field_set(
+      test_msg_extra, 0, upb_test_TestMapFieldExtra_THREE, arena.ptr()));
+
+  size_t size;
+  char* serialized = upb_test_TestMapFieldExtra_serialize_ex(
+      test_msg_extra, 0, arena.ptr(), &size);
+  ASSERT_NE(nullptr, serialized);
+  ASSERT_NE(0, size);
+
+  upb_test_TestMapField* test_msg =
+      upb_test_TestMapField_parse(serialized, size, arena.ptr());
+  ASSERT_NE(nullptr, test_msg);
+
+  ASSERT_FALSE(upb_test_TestMapField_map_field_get(test_msg, 0, nullptr));
+  serialized =
+      upb_test_TestMapField_serialize_ex(test_msg, 0, arena.ptr(), &size);
+  ASSERT_NE(0, size);
+  // parse into second instance
+  upb_test_TestMapFieldExtra* test_msg_extra2 =
+      upb_test_TestMapFieldExtra_parse(serialized, size, arena.ptr());
+  ASSERT_TRUE(
+      upb_test_TestMapFieldExtra_map_field_get(test_msg_extra2, 0, nullptr));
+}

upb/msg_test.proto

@@ -158,3 +158,22 @@ message TestMaxRequiredFields {
   required int32 required_int32_61 = 61;
   required int32 required_int32_62 = 62;
 }
+
+message TestMapField {
+  enum EnumMap {
+    ZERO = 0;
+    ONE = 1;
+    TWO = 2;
+  }
+  map<int32, EnumMap> map_field = 1;
+}
+
+message TestMapFieldExtra {
+  enum EnumMap {
+    ZERO = 0;
+    ONE = 1;
+    TWO = 2;
+    THREE = 3;
+  }
+  map<int32, EnumMap> map_field = 1;
+}