fix missed oom handling in unset required

PiperOrigin-RevId: 956610684

python/google/protobuf/internal/message_test.py

@@ -85,10 +85,21 @@ class MessageTest(unittest.TestCase):
       _ = msg3.SerializePartialToString()
       _ = msg3.ListFields()
       _ = msg3.DiscardUnknownFields()
+      _ = msg3.FindInitializationErrors()
 
       msg4 = message_module.TestAllTypes()
       msg4.CopyFrom(msg3)
 
+      # Try deepcopy
+      _ = copy.deepcopy(msg3)
+
+      # Try other upb message APIs
+      _ = msg3.ByteSize()
+      _ = msg3.SerializePartialToString()
+      _ = msg3.ListFields()
+      _ = msg3.FindInitializationErrors()
+      _ = msg3.DiscardUnknownFields()
+
       if hasattr(message_module, 'TestAllExtensions'):
         ext_msg = message_module.TestAllExtensions()
         test_util.SetAllExtensions(ext_msg)

python/message.c

@@ -20,6 +20,7 @@
 #include "python/map.h"
 #include "python/protobuf.h"
 #include "python/repeated.h"
+#include "upb/base/status.h"
 #include "upb/mem/alloc.h"
 #include "upb/mem/arena.h"
 #include "upb/message/array.h"
@@ -1315,7 +1316,14 @@ static PyObject* PyUpb_Message_IsInitialized(PyObject* _self, PyObject* args) {
     upb_Message* msg = PyUpb_Message_GetIfReified(_self);
     const upb_MessageDef* m = PyUpb_Message_GetMsgdef(_self);
     const upb_DefPool* symtab = upb_FileDef_Pool(upb_MessageDef_File(m));
-    bool initialized = !upb_util_HasUnsetRequired(msg, m, symtab, NULL);
+    upb_Status status;
+    upb_Status_Clear(&status);
+    bool initialized =
+        !upb_util_HasUnsetRequired(msg, m, symtab, NULL, &status);
+    if (!upb_Status_IsOk(&status)) {
+      PyErr_SetNone(PyExc_MemoryError);
+      return NULL;
+    }
     return PyBool_FromLong(initialized);
   }
 }
@@ -1702,9 +1710,18 @@ static PyObject* PyUpb_Message_FindInitializationErrors(PyObject* _self,
   const upb_MessageDef* msgdef = _PyUpb_Message_GetMsgdef(self);
   const upb_DefPool* ext_pool = upb_FileDef_Pool(upb_MessageDef_File(msgdef));
   upb_FieldPathEntry* fields_base;
+  upb_Status status;
+  upb_Status_Clear(&status);
+  bool has_unset_required =
+      upb_util_HasUnsetRequired(msg, msgdef, ext_pool, &fields_base, &status);
+  if (!upb_Status_IsOk(&status)) {
+    PyErr_SetString(PyExc_MemoryError, upb_Status_ErrorMessage(&status));
+    return NULL;
+  }
   PyObject* ret = PyList_New(0);
   if (!ret) return NULL;
-  if (!upb_util_HasUnsetRequired(msg, msgdef, ext_pool, &fields_base)) {
+
+  if (!has_unset_required) {
     return ret;
   }
   upb_FieldPathEntry* fields = fields_base;

upb/util/required_fields.c

@@ -17,6 +17,7 @@
 #include <string.h>
 
 #include "upb/base/descriptor_constants.h"
+#include "upb/base/status.h"
 #include "upb/mem/alloc.h"
 #include "upb/message/array.h"
 #include "upb/message/map.h"
@@ -282,24 +283,47 @@ static void upb_util_FindUnsetRequiredInternal(upb_FindContext* ctx,
   }
 }
 
+static bool upb_util_DoHasUnsetRequired(upb_FindContext* ctx,
+                                        const upb_Message* msg,
+                                        const upb_MessageDef* m) {
+  if (UPB_SETJMP(ctx->err) != 0) {
+    return false;
+  }
+  upb_util_FindUnsetRequiredInternal(ctx, msg, m);
+  return true;
+}
+
 bool upb_util_HasUnsetRequired(const upb_Message* msg, const upb_MessageDef* m,
                                const upb_DefPool* ext_pool,
-                               upb_FieldPathEntry** fields) {
+                               upb_FieldPathEntry** fields,
+                               upb_Status* status) {
   upb_FindContext ctx;
   ctx.has_unset_required = false;
   ctx.save_paths = fields != NULL;
   ctx.ext_pool = ext_pool;
   upb_FieldPathVector_Init(&ctx.stack);
   upb_FieldPathVector_Init(&ctx.out_fields);
-  upb_util_FindUnsetRequiredInternal(&ctx, msg, m);
-  upb_gfree(ctx.stack.path);
 
-  if (ctx.has_unset_required && fields) {
-    upb_FieldPathVector_Reserve(&ctx, &ctx.out_fields, 1);
-    ctx.out_fields.path[ctx.out_fields.size] =
-        (upb_FieldPathEntry){.field = NULL};
-    *fields = ctx.out_fields.path;
+  if (upb_util_DoHasUnsetRequired(&ctx, msg, m)) {
+    upb_gfree(ctx.stack.path);
+
+    if (ctx.has_unset_required && fields) {
+      upb_FieldPathVector_Reserve(&ctx, &ctx.out_fields, 1);
+      ctx.out_fields.path[ctx.out_fields.size] =
+          (upb_FieldPathEntry){.field = NULL};
+      *fields = ctx.out_fields.path;
+    }
+
+    return ctx.has_unset_required;
+  } else {
+    upb_gfree(ctx.stack.path);
+    upb_gfree(ctx.out_fields.path);
+    if (fields) {
+      *fields = NULL;
+    }
+    if (status) {
+      upb_Status_SetErrorMessage(status, "out of memory");
+    }
+    return false;
   }
-
-  return ctx.has_unset_required;
 }

upb/util/required_fields.h

@@ -10,8 +10,10 @@
 
 #include <stddef.h>
 
+#include "upb/base/status.h"
+#include "upb/message/array.h"
+#include "upb/message/message.h"
 #include "upb/reflection/def.h"
-#include "upb/reflection/message.h"
 
 // Must be last.
 #include "upb/port/def.inc"
@@ -65,7 +67,7 @@ size_t upb_FieldPath_ToText(upb_FieldPathEntry** path, char* buf, size_t size);
 // freeing this array.
 bool upb_util_HasUnsetRequired(const upb_Message* msg, const upb_MessageDef* m,
                                const upb_DefPool* ext_pool,
-                               upb_FieldPathEntry** fields);
+                               upb_FieldPathEntry** fields, upb_Status* status);
 
 #ifdef __cplusplus
 } /* extern "C" */

upb/util/required_fields_test.cc

@@ -61,7 +61,7 @@ class RequiredFieldsTest : public testing::Test {
     upb_FieldPathEntry* entries = nullptr;
     EXPECT_EQ(!missing.empty(),
               upb_util_HasUnsetRequired(UPB_UPCAST(test_msg), m.ptr(),
-                                        defpool.ptr(), &entries));
+                                        defpool.ptr(), &entries, status.ptr()));
     if (entries) {
       EXPECT_EQ(missing, PathsToText(entries));
       free(entries);
@@ -71,7 +71,7 @@ class RequiredFieldsTest : public testing::Test {
     // about them.
     EXPECT_EQ(!missing.empty(),
               upb_util_HasUnsetRequired(UPB_UPCAST(test_msg), m.ptr(),
-                                        defpool.ptr(), nullptr));
+                                        defpool.ptr(), nullptr, status.ptr()));
   }
 };