diff --git a/upb/json/encode.c b/upb/json/encode.c index c0cd54f82f..81fdc3366e 100644 --- a/upb/json/encode.c +++ b/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) { diff --git a/upb/json/fuzz_test.cc b/upb/json/fuzz_test.cc index d430d15f40..17680d1625 100644 --- a/upb/json/fuzz_test.cc +++ b/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); diff --git a/upb/message/test.cc b/upb/message/test.cc index e7df132e5c..595ff65780 100644 --- a/upb/message/test.cc +++ b/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" diff --git a/upb/mini_descriptor/decode.c b/upb/mini_descriptor/decode.c index db15ae9a0e..3e10de570e 100644 --- a/upb/mini_descriptor/decode.c +++ b/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 diff --git a/upb/test/fuzz_util.cc b/upb/test/fuzz_util.cc index 58056d1e1a..f22b44de72 100644 --- a/upb/test/fuzz_util.cc +++ b/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_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(&ext), 1); - UPB_ASSERT(status == kUpb_ExtensionRegistryStatus_Ok); + if (status != kUpb_ExtensionRegistryStatus_Ok) { + *exts = nullptr; + return; + } } } } diff --git a/upb/wire/internal/encoder.c b/upb/wire/internal/encoder.c index 78f2b901c1..c9e5e7f600 100644 --- a/upb/wire/internal/encoder.c +++ b/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)) { diff --git a/upb_generator/c/generator.cc b/upb_generator/c/generator.cc index 5c2caf4ce8..33f5a203d7 100644 --- a/upb_generator/c/generator.cc +++ b/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; }