From 2de70d710510ea7c5ad7ec0c72bfed7f411c7b60 Mon Sep 17 00:00:00 2001 From: Protobuf Team Bot Date: Wed, 19 Aug 2026 13:44:36 -0700 Subject: [PATCH] Add allocation failure tests for unknowns and defbuilder PiperOrigin-RevId: 967386997 --- upb/message/BUILD | 1 + upb/message/internal/compare_unknown.c | 36 ++++++-- upb/message/internal/compare_unknown_test.cc | 58 ++++++++++++ upb/reflection/BUILD | 1 + upb/reflection/internal/def_builder.c | 8 +- upb/reflection/internal/def_builder_test.cc | 95 ++++++++++++++++++++ 6 files changed, 191 insertions(+), 8 deletions(-) diff --git a/upb/message/BUILD b/upb/message/BUILD index 174431b0bd..c4aa476f2b 100644 --- a/upb/message/BUILD +++ b/upb/message/BUILD @@ -456,6 +456,7 @@ cc_test( "//upb/wire", "//upb/wire:reader", "//upb/wire/test_util:wire_message", + "@abseil-cpp//absl/cleanup", "@googletest//:gtest", "@googletest//:gtest_main", ], diff --git a/upb/message/internal/compare_unknown.c b/upb/message/internal/compare_unknown.c index c2b25e1396..16bba3535b 100644 --- a/upb/message/internal/compare_unknown.c +++ b/upb/message/internal/compare_unknown.c @@ -21,6 +21,7 @@ #include "upb/message/unknown_fields.h" #include "upb/mini_table/extension.h" #include "upb/mini_table/message.h" +#include "upb/port/overflow.h" #include "upb/wire/encode.h" #include "upb/wire/eps_copy_input_stream.h" #include "upb/wire/internal/back_alloc.h" @@ -80,10 +81,17 @@ static void upb_UnknownFields_Grow(upb_UnknownField_Context* ctx, upb_UnknownField** ptr, upb_UnknownField** end) { size_t old = (*ptr - *base); - size_t new = UPB_MAX(4, old * 2); + size_t new; + if (upb_MulOverflow((uint32_t)2, old, &new)) { + upb_UnknownFields_OutOfMemory(ctx); + } + new = UPB_MAX(4, new); + size_t new_bytes; + if (upb_MulOverflow(new, sizeof(**base), &new_bytes)) { + upb_UnknownFields_OutOfMemory(ctx); + } - *base = upb_Arena_Realloc(ctx->arena, *base, old * sizeof(**base), - new * sizeof(**base)); + *base = upb_Arena_Realloc(ctx->arena, *base, old * sizeof(**base), new_bytes); if (!*base) upb_UnknownFields_OutOfMemory(ctx); *ptr = *base + old; @@ -131,11 +139,25 @@ static void upb_UnknownFields_SortRecursive(upb_UnknownField* arr, size_t start, static void upb_UnknownFields_Sort(upb_UnknownField_Context* ctx, upb_UnknownFields* fields) { if (ctx->tmp_size < fields->size) { - const int oldsize = ctx->tmp_size * sizeof(*ctx->tmp); + size_t oldsize; + if (upb_MulOverflow(sizeof(*ctx->tmp), ctx->tmp_size, &oldsize)) { + upb_UnknownFields_OutOfMemory(ctx); + } ctx->tmp_size = UPB_MAX(8, ctx->tmp_size); - while (ctx->tmp_size < fields->size) ctx->tmp_size *= 2; - const int newsize = ctx->tmp_size * sizeof(*ctx->tmp); - ctx->tmp = upb_grealloc(ctx->tmp, oldsize, newsize); + while (ctx->tmp_size < fields->size) { + if (upb_MulOverflow((uint32_t)2, ctx->tmp_size, &ctx->tmp_size)) { + upb_UnknownFields_OutOfMemory(ctx); + } + } + size_t newsize; + if (upb_MulOverflow(ctx->tmp_size, sizeof(*ctx->tmp), &newsize)) { + upb_UnknownFields_OutOfMemory(ctx); + } + upb_UnknownField* tmp = upb_grealloc(ctx->tmp, oldsize, newsize); + if (!tmp) { + upb_UnknownFields_OutOfMemory(ctx); + } + ctx->tmp = tmp; } upb_UnknownFields_SortRecursive(fields->fields, 0, fields->size, ctx->tmp); } diff --git a/upb/message/internal/compare_unknown_test.cc b/upb/message/internal/compare_unknown_test.cc index 2c0a15f59a..7346111cd8 100644 --- a/upb/message/internal/compare_unknown_test.cc +++ b/upb/message/internal/compare_unknown_test.cc @@ -14,9 +14,12 @@ #include #include +#include "absl/cleanup/cleanup.h" #include "google/protobuf/test_messages_proto2.upb.h" #include "upb/base/string_view.h" #include "upb/base/upcast.h" +#include "upb/mem/alloc.h" +#include "upb/mem/arena.h" #include "upb/mem/arena.hpp" #include "upb/message/compare.h" #include "upb/message/internal/accessors.h" @@ -365,5 +368,60 @@ TEST(CompareTest, MessageIsEqualWithCanonicalAndNonCanonicalExtensions) { } // namespace +TEST(CompareTest, AllocationFailure) { + if (!upb_AllocationCount_IsAvailable()) return; + + auto RunScenario = [&]() -> bool { + upb_Arena* arena1 = upb_Arena_New(); + if (!arena1) return false; + auto cleanup1 = absl::MakeCleanup([arena1] { upb_Arena_Free(arena1); }); + + upb_Arena* arena2 = upb_Arena_New(); + if (!arena2) return false; + auto cleanup2 = absl::MakeCleanup([arena2] { upb_Arena_Free(arena2); }); + + protobuf_test_messages_proto2_TestAllTypesProto2* msg1 = + protobuf_test_messages_proto2_TestAllTypesProto2_new(arena1); + if (!msg1) return false; + protobuf_test_messages_proto2_TestAllTypesProto2* msg2 = + protobuf_test_messages_proto2_TestAllTypesProto2_new(arena2); + if (!msg2) return false; + + WireMessage uf1 = {{2, Fixed32(456)}, {1, Fixed64(123)}}; + WireMessage uf2 = {{2, Fixed32(456)}, {1, Fixed64(123)}}; + std::string buf1 = ToBinaryPayloadWithLongVarints(uf1, 1, 1); + std::string buf2 = ToBinaryPayloadWithLongVarints(uf2, 1, 1); + + if (!UPB_PRIVATE(_upb_Message_AddUnknown)(UPB_UPCAST(msg1), buf1.data(), + buf1.size(), arena1, + kUpb_AddUnknown_Copy)) { + return false; + } + if (!UPB_PRIVATE(_upb_Message_AddUnknown)(UPB_UPCAST(msg2), buf2.data(), + buf2.size(), arena2, + kUpb_AddUnknown_Copy)) { + return false; + } + upb_UnknownCompareResult res = + UPB_PRIVATE(_upb_Message_UnknownFieldsAreEqual)(UPB_UPCAST(msg1), + UPB_UPCAST(msg2), 64); + return res == kUpb_UnknownCompareResult_Equal; + }; + + upb_AllocationCount_Reset(); + if (RunScenario()) { + 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 = RunScenario(); + EXPECT_FALSE(success_with_fail) + << "Unknown comparison unexpectedly succeeded when allocation " + << "number " << i << " was failed."; + } + } + upb_AllocationCount_Reset(); +} + } // namespace test } // namespace upb diff --git a/upb/reflection/BUILD b/upb/reflection/BUILD index 2cfb75b700..5f5091aa5f 100644 --- a/upb/reflection/BUILD +++ b/upb/reflection/BUILD @@ -239,6 +239,7 @@ cc_test( "//upb/mini_descriptor", "//upb/mini_table", "//upb/port", + "@abseil-cpp//absl/cleanup", "@abseil-cpp//absl/strings", "@googletest//:gtest", "@googletest//:gtest_main", diff --git a/upb/reflection/internal/def_builder.c b/upb/reflection/internal/def_builder.c index 5aedb4c1d8..ac4e30be37 100644 --- a/upb/reflection/internal/def_builder.c +++ b/upb/reflection/internal/def_builder.c @@ -21,6 +21,7 @@ #include "upb/mem/alloc.h" #include "upb/mem/arena.h" #include "upb/message/copy.h" +#include "upb/port/overflow.h" #include "upb/reflection/def.h" #include "upb/reflection/def_type.h" #include "upb/reflection/descriptor_bootstrap.h" @@ -138,7 +139,12 @@ const void* _upb_DefBuilder_ResolveAny(upb_DefBuilder* ctx, } else { // Remove components from base until we find an entry or run out. size_t baselen = base ? strlen(base) : 0; - char* tmp = upb_gmalloc(sym.size + baselen + 1); + size_t alloc_size; + if (upb_AddOverflow(sym.size, baselen + 1, &alloc_size)) { + _upb_DefBuilder_OomErr(ctx); + } + char* tmp = upb_gmalloc(alloc_size); + if (!tmp) _upb_DefBuilder_OomErr(ctx); while (1) { char* p = tmp; if (baselen) { diff --git a/upb/reflection/internal/def_builder_test.cc b/upb/reflection/internal/def_builder_test.cc index d8e5a4640b..36ffd37308 100644 --- a/upb/reflection/internal/def_builder_test.cc +++ b/upb/reflection/internal/def_builder_test.cc @@ -7,9 +7,20 @@ #include "upb/reflection/internal/def_builder.h" +#include +#include + #include +#include "absl/cleanup/cleanup.h" #include "absl/strings/string_view.h" +#include "upb/base/status.h" +#include "upb/base/string_view.h" +#include "upb/mem/alloc.h" +#include "upb/mem/arena.h" #include "upb/mem/arena.hpp" +#include "upb/reflection/def.h" +#include "upb/reflection/def_pool.h" +#include "upb/reflection/def_type.h" // Must be last. #include "upb/port/def.inc" @@ -82,3 +93,87 @@ INSTANTIATE_TEST_SUITE_P(PartIdentTest, PartIdentTestBase, {"#", false}, {".", false}, {"", false}})); + +TEST(DefBuilderTest, AllocationFailure) { + if (!upb_AllocationCount_IsAvailable()) return; + + auto RunScenario = [&]() -> bool { + upb_Arena* arena = upb_Arena_New(); + if (!arena) return false; + auto cleanup = absl::MakeCleanup([arena] { upb_Arena_Free(arena); }); + + upb_Status status; + upb_DefBuilder ctx; + ctx.status = &status; + ctx.arena = arena; + upb_Status_Clear(&status); + + if (UPB_SETJMP(ctx.err)) { + return false; + } + + _upb_DefBuilder_MakeFullName(&ctx, "abc", + upb_StringView_FromDataAndSize("foo", 3)); + return true; + }; + + upb_AllocationCount_Reset(); + if (RunScenario()) { + 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 = RunScenario(); + EXPECT_FALSE(success_with_fail) + << "DefBuilder unexpectedly succeeded when allocation " + << "number " << i << " was failed."; + } + } + upb_AllocationCount_Reset(); +} + +TEST(DefBuilderTest, ResolveAllocationFailure) { + if (!upb_AllocationCount_IsAvailable()) return; + + auto RunScenario = [&]() -> bool { + upb_Arena* arena = upb_Arena_New(); + if (!arena) return false; + auto cleanup = absl::MakeCleanup([arena] { upb_Arena_Free(arena); }); + + upb_DefPool* symtab = upb_DefPool_New(); + if (!symtab) return false; + auto cleanup_symtab = + absl::MakeCleanup([symtab] { upb_DefPool_Free(symtab); }); + + upb_Status status; + upb_DefBuilder ctx; + ctx.status = &status; + ctx.arena = arena; + ctx.symtab = symtab; + upb_Status_Clear(&status); + + if (UPB_SETJMP(ctx.err)) { + return false; + } + + upb_deftype_t found_type; + _upb_DefBuilder_ResolveAny(&ctx, "dbg", "base.package", + upb_StringView_FromDataAndSize("foo", 3), + &found_type); + return true; + }; + + upb_AllocationCount_Reset(); + if (RunScenario()) { + 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 = RunScenario(); + EXPECT_FALSE(success_with_fail) + << "DefBuilder Resolve unexpectedly succeeded when allocation " + << "number " << i << " was failed."; + } + } + upb_AllocationCount_Reset(); +}