Allocation fault injection tests for upb encode and decode

PiperOrigin-RevId: 960627263

upb/json/encode.c

@@ -73,6 +73,9 @@ static upb_Arena* jsonenc_arena(jsonenc* e) {
   /* Create lazily, since it's only needed for Any */
   if (!e->arena) {
     e->arena = upb_Arena_New();
+    if (!e->arena) {
+      jsonenc_err(e, "Out of memory");
+    }
   }
   return e->arena;
 }
@@ -390,6 +393,10 @@ static void jsonenc_any(jsonenc* e, const upb_Message* msg,
   const upb_MiniTable* any_layout = upb_MessageDef_MiniTable(any_m);
   upb_Arena* arena = jsonenc_arena(e);
   upb_Message* any = upb_Message_New(any_layout, arena);
+  if (!any) {
+    jsonenc_err(e, "Out of memory");
+    return;
+  }
 
   if (upb_Decode(value.data, value.size, any, any_layout, NULL, 0, arena) !=
       kUpb_DecodeStatus_Ok) {

upb/json/fuzz_test.cc

@@ -17,6 +17,7 @@
 #include "upb/json/encode.h"
 #include "upb/json/test.upb.h"
 #include "upb/json/test.upbdefs.h"
+#include "upb/mem/alloc.h"
 #include "upb/mem/arena.h"
 #include "upb/mem/arena.hpp"
 #include "upb/reflection/def.hpp"
@@ -50,6 +51,68 @@ void DecodeEncodeArbitraryJson(std::string_view json) {
   size_t written = upb_JsonEncode(UPB_UPCAST(box), m.ptr(), defpool.ptr(),
                                   options, json_buf, size + 1, status.ptr());
   EXPECT_EQ(written, size);
+
+  if (upb_AllocationCount_IsAvailable()) {
+    auto RunJsonScenario = [&]() -> bool {
+      upb_Arena* local_arena = upb_Arena_New();
+      if (!local_arena) return false;
+
+      upb_test_Box* local_box = upb_test_Box_new(local_arena);
+      if (!local_box) {
+        upb_Arena_Free(local_arena);
+        return false;
+      }
+
+      upb::Status local_status;
+      bool local_ok = upb_JsonDecode(
+          json.data(), json.size(), UPB_UPCAST(local_box), m.ptr(),
+          defpool.ptr(), options, local_arena, local_status.ptr());
+      if (!local_ok) {
+        upb_Arena_Free(local_arena);
+        return false;
+      }
+
+      size_t local_sz =
+          upb_JsonEncode(UPB_UPCAST(local_box), m.ptr(), defpool.ptr(), options,
+                         nullptr, 0, local_status.ptr());
+      if (local_sz == (size_t)-1 || !local_status.ok()) {
+        upb_Arena_Free(local_arena);
+        return false;
+      }
+
+      char* local_json_buf = (char*)upb_Arena_Malloc(local_arena, local_sz + 1);
+      if (!local_json_buf) {
+        upb_Arena_Free(local_arena);
+        return false;
+      }
+
+      size_t local_written =
+          upb_JsonEncode(UPB_UPCAST(local_box), m.ptr(), defpool.ptr(), options,
+                         local_json_buf, local_sz + 1, local_status.ptr());
+      if (local_written == (size_t)-1 || !local_status.ok()) {
+        upb_Arena_Free(local_arena);
+        return false;
+      }
+
+      upb_Arena_Free(local_arena);
+      return true;
+    };
+
+    upb_AllocationCount_Reset();
+    if (RunJsonScenario()) {
+      size_t total_allocations = upb_AllocationCount_Get();
+      for (size_t i = 0; i < total_allocations; ++i) {
+        upb_AllocationCount_Reset();
+        upb_AllocationCount_FailOn(i);
+        bool success_with_fail = RunJsonScenario();
+        EXPECT_FALSE(success_with_fail)
+            << "Fuzzed JSON scenario unexpectedly succeeded when allocation "
+            << "number " << i << " was failed, with "
+            << upb_AllocationCount_Get() << " total.";
+      }
+    }
+    upb_AllocationCount_Reset();
+  }
 }
 FUZZ_TEST(FuzzTest, DecodeEncodeArbitraryJson);
 

upb/message/test.cc

@@ -29,6 +29,7 @@
 #include "upb/base/upcast.h"
 #include "upb/json/decode.h"
 #include "upb/json/encode.h"
+#include "upb/mem/alloc.h"
 #include "upb/mem/arena.h"
 #include "upb/mem/arena.hpp"
 #include "upb/message/accessors.h"

upb/mini_descriptor/decode.c

@@ -832,6 +832,7 @@ done:
 #endif
 
   upb_MiniTable* ret = upb_Arena_Malloc(decoder->arena, mt_size);
+  upb_MdDecoder_CheckOutOfMemory(&decoder->base, ret);
   memcpy(ret, &decoder->table, sizeof(*ret));
 
 #if UPB_FASTTABLE

upb/test/fuzz_util.cc

@@ -38,8 +38,17 @@ class Builder {
 
   const upb_MiniTable* Build(upb_ExtensionRegistry** exts) {
     BuildMessages();
+    if (mini_tables_.size() < input_->mini_descriptors.size()) {
+      return nullptr;
+    }
     BuildEnums();
+    if (enum_tables_.size() < input_->enum_mini_descriptors.size()) {
+      return nullptr;
+    }
     BuildExtensions(exts);
+    if (!input_->extensions.empty() && (!exts || !*exts)) {
+      return nullptr;
+    }
     if (!LinkMessages()) return nullptr;
     return mini_tables_.empty() ? nullptr : mini_tables_.front();
   }
@@ -117,18 +126,26 @@ void Builder::BuildExtensions(upb_ExtensionRegistry** exts) {
     *exts = nullptr;
   } else {
     *exts = upb_ExtensionRegistry_New(arena_);
+    if (!*exts) return;
     const char* ptr = input_->extensions.data();
     const char* end = ptr + input_->extensions.size();
     // Iterate through the buffer, building extensions as long as we can.
     while (ptr < end) {
       upb_MiniTableExtension* ext = reinterpret_cast<upb_MiniTableExtension*>(
           upb_Arena_Malloc(arena_, sizeof(*ext)));
+      if (!ext) {
+        *exts = nullptr;
+        return;
+      }
       upb_MiniTableSub sub;
       const upb_MiniTable* extendee = NextMiniTable();
       if (!extendee) break;
       ptr = upb_MiniTableExtension_Init(ptr, end - ptr, ext, extendee, sub,
                                         status.ptr());
-      if (!ptr) break;
+      if (!ptr) {
+        *exts = nullptr;
+        return;
+      }
       if (!LinkExtension(ext)) continue;
       if (upb_MiniTable_FindFieldByNumber(
               extendee, upb_MiniTableExtension_Number(ext)) != nullptr) {
@@ -139,7 +156,10 @@ void Builder::BuildExtensions(upb_ExtensionRegistry** exts) {
         continue;
       auto status = upb_ExtensionRegistry_AddArray(
           *exts, const_cast<const upb_MiniTableExtension**>(&ext), 1);
-      UPB_ASSERT(status == kUpb_ExtensionRegistryStatus_Ok);
+      if (status != kUpb_ExtensionRegistryStatus_Ok) {
+        *exts = nullptr;
+        return;
+      }
     }
   }
 }

upb/wire/internal/encoder.c

@@ -601,9 +601,12 @@ static char* encode_map(char* ptr, upb_encstate* e, const upb_Message* msg,
 
   if (e->options & kUpb_EncodeOption_Deterministic) {
     _upb_sortedmap sorted;
-    _upb_mapsorter_pushmap(
-        &e->sorter, layout->UPB_PRIVATE(fields)[0].UPB_PRIVATE(descriptortype),
-        map, &sorted);
+    if (!_upb_mapsorter_pushmap(
+            &e->sorter,
+            layout->UPB_PRIVATE(fields)[0].UPB_PRIVATE(descriptortype), map,
+            &sorted)) {
+      encode_err(e, kUpb_EncodeStatus_OutOfMemory);
+    }
     upb_MapEntry ent;
     while (_upb_sortedmap_next(&e->sorter, map, &sorted, &ent)) {
       ptr = encode_mapentry(ptr, e, upb_MiniTableField_Number(f), layout, &ent);
@@ -739,7 +742,7 @@ static char* encode_exts(char* ptr, upb_encstate* e, const upb_MiniTable* m,
   if (e->options & kUpb_EncodeOption_Deterministic) {
     _upb_sortedmap sorted;
     if (!_upb_mapsorter_pushexts(&e->sorter, in, &sorted)) {
-      // TODO: b/378744096 - handle alloc failure
+      encode_err(e, kUpb_EncodeStatus_OutOfMemory);
     }
     const upb_Extension* ext;
     while (_upb_sortedmap_nextext(&e->sorter, &sorted, &ext)) {

upb_generator/c/generator.cc

@@ -786,6 +786,7 @@ void GenerateMapSetters(Context& c, upb::FieldDefPtr field,
           const upb_MiniTableField field = $field_init$;
           upb_Map* map = _upb_Message_GetOrCreateMutableMap(
               UPB_UPCAST(msg), &field, $key_size$, $val_size$, a);
+          if (!map) return false;
           return _upb_Map_Insert(map, &key, $key_size$, &val, $val_size$, a) !=
                  kUpb_MapInsertStatus_OutOfMemory;
         }