Fix bug in deterministic extension encoding when empty extensions are…

… present PiperOrigin-RevId: 707943424
protocolbuffers · Dec 19, 2024 · c78129d · c78129d
1 parent cac05fe
commit c78129d
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 24 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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);
         }
       }
     }