Add allocation failure tests for unknowns and defbuilder

PiperOrigin-RevId: 967386997

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",
     ],

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);
 }

upb/message/internal/compare_unknown_test.cc

@@ -14,9 +14,12 @@
 #include <string>
 
 #include <gtest/gtest.h>
+#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

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",

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) {

upb/reflection/internal/def_builder_test.cc

@@ -7,9 +7,20 @@
 
 #include "upb/reflection/internal/def_builder.h"
 
+#include <cstddef>
+#include <vector>
+
 #include <gtest/gtest.h>
+#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();
+}