From c78129dba282e57944b8eeb735edf3c5b65077a8 Mon Sep 17 00:00:00 2001 From: Protobuf Team Bot Date: Thu, 19 Dec 2024 09:49:10 -0800 Subject: [PATCH] Fix bug in deterministic extension encoding when empty extensions are present PiperOrigin-RevId: 707943424 --- upb/message/internal/map_sorter.h | 2 +- upb/message/map_sorter.c | 4 +- upb/message/test.cc | 92 +++++++++++++++++++++++++++++++ upb/wire/encode.c | 41 +++++++------- 4 files changed, 115 insertions(+), 24 deletions(-) diff --git a/upb/message/internal/map_sorter.h b/upb/message/internal/map_sorter.h index bd997bc476fcb..4044a55781067 100644 --- a/upb/message/internal/map_sorter.h +++ b/upb/message/internal/map_sorter.h @@ -82,7 +82,7 @@ bool _upb_mapsorter_pushmap(_upb_mapsorter* s, upb_FieldType key_type, const struct upb_Map* map, _upb_sortedmap* sorted); bool _upb_mapsorter_pushexts(_upb_mapsorter* s, const upb_Message_Internal* in, - size_t count, _upb_sortedmap* sorted); + _upb_sortedmap* sorted); #ifdef __cplusplus } /* extern "C" */ diff --git a/upb/message/map_sorter.c b/upb/message/map_sorter.c index 029eabc5ee5e3..8cf68c095bbb8 100644 --- a/upb/message/map_sorter.c +++ b/upb/message/map_sorter.c @@ -146,8 +146,10 @@ static int _upb_mapsorter_cmpext(const void* _a, const void* _b) { } bool _upb_mapsorter_pushexts(_upb_mapsorter* s, const upb_Message_Internal* in, - size_t count, _upb_sortedmap* sorted) { + _upb_sortedmap* sorted) { + size_t count = (in->size - in->ext_begin) / sizeof(upb_Extension); if (!_upb_mapsorter_resize(s, sorted, count)) return false; + if (count == 0) return true; const upb_Extension* exts = UPB_PTR_AT(in, in->ext_begin, const upb_Extension); diff --git a/upb/message/test.cc b/upb/message/test.cc index ed46c8c56e17b..c18835005a2aa 100644 --- a/upb/message/test.cc +++ b/upb/message/test.cc @@ -119,6 +119,98 @@ TEST(MessageTest, Extensions) { VerifyMessage(ext_msg4); } +TEST(MessageTest, ExtensionsDeterministic) { + upb::Arena arena; + upb_test_TestExtensions* ext_msg = upb_test_TestExtensions_new(arena.ptr()); + + EXPECT_FALSE(upb_test_TestExtensions_has_optional_int32_ext(ext_msg)); + // EXPECT_FALSE(upb_test_TestExtensions_Nested_has_optional_int32_ext(ext_msg)); + EXPECT_FALSE(upb_test_has_optional_msg_ext(ext_msg)); + + upb::DefPool defpool; + upb::MessageDefPtr m(upb_test_TestExtensions_getmsgdef(defpool.ptr())); + EXPECT_TRUE(m.ptr() != nullptr); + + std::string json = R"json( + { + "[upb_test.TestExtensions.optional_int32_ext]": 123, + "[upb_test.TestExtensions.Nested.repeated_int32_ext]": [], + "[upb_test.optional_msg_ext]": {"optional_int32": 456} + } + )json"; + upb::Status status; + EXPECT_TRUE(upb_JsonDecode(json.data(), json.size(), UPB_UPCAST(ext_msg), + m.ptr(), defpool.ptr(), 0, arena.ptr(), + status.ptr())) + << status.error_message(); + + VerifyMessage(ext_msg); + + size_t size; + char* serialized = + upb_test_TestExtensions_serialize(ext_msg, arena.ptr(), &size); + ASSERT_TRUE(serialized != nullptr); + ASSERT_GE(size, 0); + + size_t deterministic_size; + char* deterministic_serialized = upb_test_TestExtensions_serialize_ex( + ext_msg, kUpb_EncodeOption_Deterministic, arena.ptr(), + &deterministic_size); + ASSERT_TRUE(deterministic_serialized != nullptr); + ASSERT_EQ(deterministic_size, size); +} + +TEST(MessageTest, ExtensionsEmpty) { + upb::Arena arena; + + upb::DefPool defpool; + upb::MessageDefPtr m(upb_test_TestExtensions_getmsgdef(defpool.ptr())); + EXPECT_TRUE(m.ptr() != nullptr); + + for (int options : {0, int{kUpb_EncodeOption_Deterministic}}) { + std::string json_with_empty = R"json( + { + "[upb_test.TestExtensions.optional_int32_ext]": 123, + "[upb_test.TestExtensions.Nested.repeated_int32_ext]": [] + } + )json"; + upb::Status status_empty; + upb_test_TestExtensions* ext_msg_with_empty = + upb_test_TestExtensions_new(arena.ptr()); + EXPECT_TRUE(upb_JsonDecode(json_with_empty.data(), json_with_empty.size(), + UPB_UPCAST(ext_msg_with_empty), m.ptr(), + defpool.ptr(), 0, arena.ptr(), + status_empty.ptr())) + << status_empty.error_message(); + + std::string json = R"json( + { + "[upb_test.TestExtensions.optional_int32_ext]": 123 + } + )json"; + upb::Status status; + upb_test_TestExtensions* ext_msg = upb_test_TestExtensions_new(arena.ptr()); + EXPECT_TRUE(upb_JsonDecode(json.data(), json.size(), UPB_UPCAST(ext_msg), + m.ptr(), defpool.ptr(), 0, arena.ptr(), + status.ptr())) + << status.error_message(); + + size_t size_with_empty; + char* serialized = upb_test_TestExtensions_serialize_ex( + ext_msg_with_empty, options, arena.ptr(), &size_with_empty); + ASSERT_TRUE(serialized != nullptr); + ASSERT_GE(size_with_empty, 0); + + size_t size; + serialized = upb_test_TestExtensions_serialize_ex(ext_msg, options, + arena.ptr(), &size); + ASSERT_TRUE(serialized != nullptr); + // Presence or absence of an empty extension should not affect the + // serialized output. + ASSERT_EQ(size_with_empty, size); + } +} + void VerifyMessageSet(const upb_test_TestMessageSet* mset_msg) { ASSERT_TRUE(mset_msg != nullptr); bool has = upb_test_MessageSetMember_has_message_set_extension(mset_msg); diff --git a/upb/wire/encode.c b/upb/wire/encode.c index 55c628955237e..9ed482ee5dffc 100644 --- a/upb/wire/encode.c +++ b/upb/wire/encode.c @@ -593,28 +593,25 @@ static void encode_message(upb_encstate* e, const upb_Message* msg, /* Encode all extensions together. Unlike C++, we do not attempt to keep * these in field number order relative to normal fields or even to each * other. */ - size_t ext_count = upb_Message_ExtensionCount(msg); - if (ext_count) { - if (e->options & kUpb_EncodeOption_Deterministic) { - _upb_sortedmap sorted; - if (!_upb_mapsorter_pushexts(&e->sorter, in, ext_count, &sorted)) { - // TODO: b/378744096 - handle alloc failure - } - const upb_Extension* ext; - while (_upb_sortedmap_nextext(&e->sorter, &sorted, &ext)) { - encode_ext(e, ext->ext, ext->data, - m->UPB_PRIVATE(ext) == kUpb_ExtMode_IsMessageSet); - } - _upb_mapsorter_popmap(&e->sorter, &sorted); - } else { - const upb_MiniTableExtension* ext; - upb_MessageValue ext_val; - uintptr_t iter = kUpb_Message_ExtensionBegin; - while (UPB_PRIVATE(_upb_Message_NextExtensionReverse)( - msg, &ext, &ext_val, &iter)) { - encode_ext(e, ext, ext_val, - m->UPB_PRIVATE(ext) == kUpb_ExtMode_IsMessageSet); - } + if (e->options & kUpb_EncodeOption_Deterministic) { + _upb_sortedmap sorted; + if (!_upb_mapsorter_pushexts(&e->sorter, in, &sorted)) { + // TODO: b/378744096 - handle alloc failure + } + const upb_Extension* ext; + while (_upb_sortedmap_nextext(&e->sorter, &sorted, &ext)) { + encode_ext(e, ext->ext, ext->data, + m->UPB_PRIVATE(ext) == kUpb_ExtMode_IsMessageSet); + } + _upb_mapsorter_popmap(&e->sorter, &sorted); + } else { + const upb_MiniTableExtension* ext; + upb_MessageValue ext_val; + uintptr_t iter = kUpb_Message_ExtensionBegin; + while (UPB_PRIVATE(_upb_Message_NextExtensionReverse)( + msg, &ext, &ext_val, &iter)) { + encode_ext(e, ext, ext_val, + m->UPB_PRIVATE(ext) == kUpb_ExtMode_IsMessageSet); } } }