diff --git a/upb/decode.c b/upb/decode.c index 683e4ec32b..2d5bb49525 100644 --- a/upb/decode.c +++ b/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; } diff --git a/upb/msg_test.cc b/upb/msg_test.cc index cd6fdf67fb..2d1f8e991b 100644 --- a/upb/msg_test.cc +++ b/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)); +} diff --git a/upb/msg_test.proto b/upb/msg_test.proto index 81c7dee079..1f80f6918c 100644 --- a/upb/msg_test.proto +++ b/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 map_field = 1; +} + +message TestMapFieldExtra { + enum EnumMap { + ZERO = 0; + ONE = 1; + TWO = 2; + THREE = 3; + } + map map_field = 1; +}