mmo-ref/dep-protobuf
1
0
Fork 0
mirror of https://github.com/protocolbuffers/protobuf synced 2026-08-26 02:23:14 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Check return values and handle alloc failures in python

Browse source 
PiperOrigin-RevId: 955022274
This commit is contained in:
Protobuf Team Bot 2026-07-27 22:08:52 -07:00 • committed by Copybara-Service
parent a0c493ce55
commit ad6a7e8b64
19 changed files with 764 additions and 117 deletions
Download patch file Download diff file Expand all files Collapse all files

1
python/BUILD.bazel
View file

@ -115,6 +115,7 @@ py_extension(
target_compatible_with = select(_message_target_compatible_with),
deps = [
":breaking_changes",
"//src/google/protobuf:descriptor_upb_minitable_proto",
"//src/google/protobuf:descriptor_upb_reflection_proto",
"//third_party/utf8_range",
"//upb/base",

39
python/convert.c
View file

@ -7,9 +7,19 @@
#include "python/convert.h"
#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include "python/message.h"
#include "python/protobuf.h"
#include "upb/base/descriptor_constants.h"
#include "upb/mem/arena.h"
#include "upb/message/array.h"
#include "upb/message/compare.h"
#include "upb/message/message.h"
#include "upb/mini_table/message.h"
#include "upb/reflection/def.h"
#include "utf8_range.h"
@ -134,18 +144,26 @@ static bool PyUpb_GetUint32(PyObject* obj, const upb_FieldDef* f,
// If `arena` is specified, copies the string data into the given arena.
// Otherwise aliases the given data.
static upb_MessageValue PyUpb_MaybeCopyString(const char* ptr, size_t size,
upb_Arena* arena) {
static bool PyUpb_MaybeCopyString(const char* ptr, size_t size,
upb_MessageValue* val, upb_Arena* arena) {
upb_MessageValue ret;
ret.str_val.size = size;
if (arena) {
char* buf = upb_Arena_Malloc(arena, size);
memcpy(buf, ptr, size);
ret.str_val.data = buf;
if (size == 0) {
ret.str_val.data = "";
} else {
char* buf = upb_Arena_Malloc(arena, size);
if (!buf) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
ret.str_val.data = memcpy(buf, ptr, size);
}
} else {
ret.str_val.data = ptr;
}
return ret;
*val = ret;
return true;
}
const char* upb_FieldDef_TypeString(const upb_FieldDef* f) {
@ -295,8 +313,7 @@ bool PyUpb_PyToUpb(PyObject* obj, const upb_FieldDef* f, upb_MessageValue* val,
char* ptr;
Py_ssize_t size;
if (PyBytes_AsStringAndSize(obj, &ptr, &size) < 0) return false;
*val = PyUpb_MaybeCopyString(ptr, size, arena);
return true;
return PyUpb_MaybeCopyString(ptr, size, val, arena);
}
case kUpb_CType_String: {
Py_ssize_t size;
@ -312,13 +329,11 @@ bool PyUpb_PyToUpb(PyObject* obj, const upb_FieldDef* f, upb_MessageValue* val,
assert(!obj);
return false;
}
*val = PyUpb_MaybeCopyString(ptr, size, arena);
return true;
return PyUpb_MaybeCopyString(ptr, size, val, arena);
}
const char* ptr = PyUnicode_AsUTF8AndSize(obj, &size);
if (PyErr_Occurred()) return false;
*val = PyUpb_MaybeCopyString(ptr, size, arena);
return true;
return PyUpb_MaybeCopyString(ptr, size, val, arena);
}
case kUpb_CType_Message:
PyErr_Format(PyExc_ValueError, "Message objects may not be assigned");

54
python/descriptor.c
View file

@ -7,18 +7,34 @@
#include "python/descriptor.h"
// clang-format off
#include "Python.h"
// clang-format on
#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include "google/protobuf/descriptor.upb_minitable.h"
#include "google/protobuf/breaking_changes.h"
#include "python/convert.h"
#include "python/descriptor_containers.h"
#include "python/descriptor_pool.h"
#include "python/message.h"
#include "python/protobuf.h"
#include "upb/base/descriptor_constants.h"
#include "upb/base/upcast.h"
#include "upb/mem/arena.h"
#include "upb/message/message.h"
#include "upb/mini_table/message.h"
#include "upb/reflection/def.h"
#include "upb/reflection/message.h"
#include "upb/util/def_to_proto.h"
// Must be last.
#include "upb/port/def.inc"
#include "upb/wire/decode.h"
#include "upb/wire/encode.h"
// -----------------------------------------------------------------------------
// DescriptorBase
@ -74,6 +90,7 @@ static PyObject* PyUpb_DescriptorBase_Get(PyUpb_DescriptorType type,
if (!base) {
base = PyUpb_DescriptorBase_DoCreate(type, def, file);
if (!base) return NULL;
}
return &base->ob_base;
@ -99,6 +116,7 @@ static PyObject* PyUpb_DescriptorBase_GetCached(PyObject** cached,
const upb_MiniTable* layout,
const char* msg_name,
const char* strip_field) {
PyObject* py_arena = NULL;
if (!*cached) {
// Load descriptors protos if they are not loaded already. We have to do
// this lazily, otherwise, it would lead to circular imports.
@ -118,19 +136,40 @@ static PyObject* PyUpb_DescriptorBase_GetCached(PyObject** cached,
// the descriptor_pb2 that was loaded at runtime has the same members or
// layout as the C types that were compiled in.
size_t size;
PyObject* py_arena = PyUpb_Arena_New();
py_arena = PyUpb_Arena_New();
if (!py_arena) goto err;
upb_Arena* arena = PyUpb_Arena_Get(py_arena);
char* pb;
// TODO: Need to correctly handle failed return codes.
(void)upb_Encode(opts, layout, 0, arena, &pb, &size);
upb_EncodeStatus es = upb_Encode(opts, layout, 0, arena, &pb, &size);
if (es != kUpb_EncodeStatus_Ok) {
if (es == kUpb_EncodeStatus_OutOfMemory) {
PyErr_SetNone(PyExc_MemoryError);
} else {
PyErr_Format(PyUpb_ModuleState_Get()->decode_error_class,
"Error parsing descriptor: %s",
upb_EncodeStatus_String(es));
}
goto err;
}
const upb_MiniTable* opts2_layout = upb_MessageDef_MiniTable(m);
upb_Message* opts2 = upb_Message_New(opts2_layout, arena);
assert(opts2);
if (!opts2) {
PyErr_SetNone(PyExc_MemoryError);
goto err;
}
upb_DecodeStatus ds =
upb_Decode(pb, size, opts2, opts2_layout,
upb_DefPool_ExtensionRegistry(symtab), 0, arena);
(void)ds;
assert(ds == kUpb_DecodeStatus_Ok);
if (ds != kUpb_DecodeStatus_Ok) {
if (ds == kUpb_DecodeStatus_OutOfMemory) {
PyErr_SetNone(PyExc_MemoryError);
} else {
PyErr_Format(PyUpb_ModuleState_Get()->decode_error_class,
"Error parsing descriptor: %s",
upb_DecodeStatus_String(ds));
}
goto err;
}
if (strip_field) {
const upb_FieldDef* field =
@ -150,6 +189,9 @@ static PyObject* PyUpb_DescriptorBase_GetCached(PyObject** cached,
Py_INCREF(*cached);
return *cached;
err:
Py_XDECREF(py_arena);
return NULL;
}
static PyObject* PyUpb_DescriptorBase_GetOptions(PyObject** cached,

12
python/descriptor_pool.c
View file

@ -46,11 +46,21 @@ const upb_MessageDef* PyUpb_DescriptorPool_GetFileProtoDef(void) {
static PyObject* PyUpb_DescriptorPool_DoCreateWithCache(
PyTypeObject* type, PyObject* db, PyUpb_WeakMap* obj_cache) {
PyUpb_DescriptorPool* pool = (void*)PyType_GenericAlloc(type, 0);
if (!pool) goto err;
pool->symtab = upb_DefPool_New();
if (!pool->symtab) {
PyErr_SetNone(PyExc_MemoryError);
goto err;
}
pool->db = db;
Py_XINCREF(pool->db);
PyUpb_KnownObjCache_Add(obj_cache, pool->symtab, &pool->ob_base);
if (!PyUpb_KnownObjCache_Add(obj_cache, pool->symtab, &pool->ob_base)) {
goto err;
}
return &pool->ob_base;
err:
Py_XDECREF(pool);
return NULL;
}
static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,

39
python/google/protobuf/internal/allocation_count.py Normal file
View file

@ -0,0 +1,39 @@
# Protocol Buffers - Google's data interchange format
# Copyright 2026 Google LLC. All rights reserved.
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd
"""allocation_count helper for testing OOM/allocation failures."""
_upb_allocation_count = None
try:
from google.protobuf.internal import api_implementation
if api_implementation.Type() == 'upb':
from google._upb import _message
_upb_allocation_count = _message
except ImportError:
pass
def is_available():
if _upb_allocation_count is not None:
return _upb_allocation_count._AllocationCount_IsAvailable()
return False
def get():
if _upb_allocation_count is not None:
return _upb_allocation_count._AllocationCount_Get()
return 0
def reset():
if _upb_allocation_count is not None:
_upb_allocation_count._AllocationCount_Reset()
def fail_on(n):
if _upb_allocation_count is not None:
_upb_allocation_count._AllocationCount_FailOn(n)

94
python/google/protobuf/internal/message_test.py
View file

@ -31,7 +31,8 @@ import warnings
cmp = lambda x, y: (x > y) - (x < y)
from google.protobuf.internal import message_set_extensions_pb2
from google.protobuf.internal import api_implementation # pylint: disable=g-import-not-at-top
from google.protobuf.internal import allocation_count # pylint: disable=g-import-not-at-top
from google.protobuf.internal import api_implementation
from google.protobuf.internal import decoder
from google.protobuf.internal import encoder
from google.protobuf.internal import enum_type_wrapper
@ -63,6 +64,95 @@ warnings.simplefilter('error', DeprecationWarning)
@testing_refleaks.TestCase
class MessageTest(unittest.TestCase):
@unittest.skipIf(
not allocation_count.is_available(),
'Requires Debug-only allocation_count API',
)
def testOom(self, message_module):
def ManyAllocsScenario():
msg = message_module.TestAllTypes()
test_util.SetAllFields(msg)
msg.repeated_int32.extend(range(100))
msg.repeated_nested_message.add().bb = 123
serialized = msg.SerializeToString()
msg2 = message_module.TestAllTypes()
msg2.ParseFromString(serialized)
msg3 = message_module.TestAllTypes()
msg3.MergeFrom(msg2)
_ = msg3.optional_string
_ = msg3.optional_bytes
if hasattr(message_module, 'TestAllExtensions'):
ext_msg = message_module.TestAllExtensions()
test_util.SetAllExtensions(ext_msg)
ext_serialized = ext_msg.SerializeToString()
ext_msg2 = message_module.TestAllExtensions()
ext_msg2.ParseFromString(ext_serialized)
ext_msg3 = message_module.TestAllExtensions()
ext_msg3.MergeFrom(ext_msg2)
_ = ext_msg3.Extensions[unittest_pb2.optional_int32_extension]
_ = ext_msg3.Extensions[unittest_pb2.optional_nested_message_extension]
_ = ext_msg3.Extensions[
unittest_pb2.optional_nested_message_extension
].bb
# MessageSet extensions coverage
mset_msg = message_set_extensions_pb2.TestMessageSet()
ext1 = (
message_set_extensions_pb2.TestMessageSetExtension1.message_set_extension
)
ext2 = (
message_set_extensions_pb2.TestMessageSetExtension2.message_set_extension
)
mset_msg.Extensions[ext1].i = 123
mset_msg.Extensions[ext2].str = 'hello'
mset_serialized = mset_msg.SerializeToString()
mset_msg2 = message_set_extensions_pb2.TestMessageSet()
mset_msg2.ParseFromString(mset_serialized)
mset_msg3 = message_set_extensions_pb2.TestMessageSet()
mset_msg3.MergeFrom(mset_msg2)
_ = mset_msg3.Extensions[ext1].i
_ = mset_msg3.Extensions[ext2].str
# Unknown fields in MessageSet representation
mset_unknown = message_set_extensions_pb2.TestMessageSet()
mset_unknown.ParseFromString(
b'\x0b\x10\x01\x1a\x03foo\x0c\x0b\x10\x02\x1a\x03bar\x0c'
)
unknown_mset = unknown_fields.UnknownFieldSet(mset_unknown)
_ = len(unknown_mset)
if len(unknown_mset) > 0:
_ = unknown_mset[0].field_number
_ = unknown_mset[0].wire_type
_ = unknown_mset[0].data
if hasattr(message_module, 'TestEmptyMessage'):
empty = message_module.TestEmptyMessage()
empty.ParseFromString(serialized)
unknown = unknown_fields.UnknownFieldSet(empty)
_ = len(unknown)
if len(unknown) > 0:
_ = unknown[0].field_number
_ = unknown[0].wire_type
_ = unknown[0].data
for field in unknown:
_ = field.field_number
_ = field.wire_type
_ = field.data
# Warm up the cache so that subsequent runs do not trigger resizes.
ManyAllocsScenario()
allocation_count.reset()
ManyAllocsScenario()
total = allocation_count.get()
self.assertGreater(total, 0)
for i in range(total):
allocation_count.reset()
allocation_count.fail_on(i)
with self.assertRaises(MemoryError):
ManyAllocsScenario()
allocation_count.reset()
def testBadUtf8String(self, message_module):
if api_implementation.Type() != 'python':
self.skipTest(
@ -3186,6 +3276,7 @@ class Proto3Test(unittest.TestCase):
msg.map_string_string.clear()
with self.assertRaises(RuntimeError):
next(it)
def testSubmessageMap(self):
msg = map_unittest_pb2.TestMap()
@ -3714,6 +3805,7 @@ class MessageMetaGetAttrTest(unittest.TestCase):
def testMessageMetaGetAttrException(self):
class BombDescriptor:
def __get__(self, obj, objtype=None):
raise KeyboardInterrupt('should not be swallowed')

30
python/map.c
View file

@ -7,11 +7,21 @@
#include "python/map.h"
#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include "google/protobuf/breaking_changes.h"
#include "python/convert.h"
#include "python/descriptor.h"
#include "python/message.h"
#include "python/protobuf.h"
#include "upb/mem/arena.h"
#include "upb/message/array.h"
#include "upb/message/map.h"
#include "upb/message/message.h"
#include "upb/mini_table/message.h"
#include "upb/reflection/def.h"
// -----------------------------------------------------------------------------
@ -84,6 +94,7 @@ PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
return NULL;
}
PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
if (map == NULL) return NULL;
map->arena = arena;
map->field = (uintptr_t)f | 1;
map->ptr.parent = parent;
@ -104,6 +115,10 @@ upb_Map* PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map,
const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
map = upb_Map_New(arena, upb_FieldDef_CType(key_f),
upb_FieldDef_CType(val_f));
if (!map) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
}
if (subobj_map) {
PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
@ -114,7 +129,9 @@ upb_Map* PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map,
return NULL;
}
}
PyUpb_ObjCache_Add(map, &self->ob_base);
if (!PyUpb_ObjCache_Add(map, &self->ob_base)) {
return NULL;
}
Py_DECREF(self->ptr.parent);
self->ptr.map = map; // Overwrites self->ptr.parent.
self->field &= ~(uintptr_t)1;
@ -161,6 +178,7 @@ static bool PyUpb_MapContainer_Set(PyUpb_MapContainer* self, upb_Map* map,
self->version--;
return true;
case kUpb_MapInsertStatus_OutOfMemory:
PyErr_SetNone(PyExc_MemoryError);
return false;
}
return false; // Unreachable, silence compiler warning.
@ -209,6 +227,10 @@ static PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
const upb_MessageDef* m = upb_FieldDef_MessageSubDef(val_f);
const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
u_val.msg_val = upb_Message_New(layout, arena);
if (!u_val.msg_val) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
} else {
memset(&u_val, 0, sizeof(u_val));
}
@ -390,12 +412,16 @@ PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
return NULL;
}
ret = (void*)PyType_GenericAlloc(cls, 0);
if (ret == NULL) return NULL;
ret->arena = arena;
ret->field = (uintptr_t)f;
ret->ptr.map = map;
ret->version = 0;
Py_INCREF(arena);
PyUpb_ObjCache_Add(map, &ret->ob_base);
if (!PyUpb_ObjCache_Add(map, &ret->ob_base)) {
Py_DECREF(&ret->ob_base);
return NULL;
}
return &ret->ob_base;
}

211
python/message.c
View file

@ -256,16 +256,28 @@ static PyObject* PyUpb_Message_New(PyObject* cls, PyObject* unused_args,
const upb_MessageDef* msgdef = PyUpb_MessageMeta_GetMsgdef(cls);
const upb_MiniTable* layout = upb_MessageDef_MiniTable(msgdef);
PyUpb_Message* msg = (void*)PyType_GenericAlloc((PyTypeObject*)cls, 0);
if (!msg) return NULL;
msg->def = (uintptr_t)msgdef;
msg->arena = PyUpb_Arena_New();
if (!msg->arena) goto err1;
msg->ptr.msg = upb_Message_New(layout, PyUpb_Arena_Get(msg->arena));
if (!msg->ptr.msg) {
PyErr_SetNone(PyExc_MemoryError);
goto err1;
}
msg->unset_subobj_map = NULL;
msg->ext_dict = NULL;
msg->version = 0;
PyObject* ret = &msg->ob_base;
PyUpb_ObjCache_Add(msg->ptr.msg, ret);
if (!PyUpb_ObjCache_Add(msg->ptr.msg, ret)) goto err2;
return ret;
err2:
Py_DECREF(msg->arena);
err1:
Py_DECREF(msg);
return NULL;
}
/*
@ -404,6 +416,7 @@ static bool PyUpb_Message_InitWKTOrMerge(const upb_MessageDef* msgdef,
// Fall back to init as normal message field.
PyErr_Clear();
PyObject* tmp = PyUpb_Message_Clear((PyUpb_Message*)msg);
if (!tmp) return false;
Py_DECREF(tmp);
ok = PyUpb_Message_InitAttributes(msg, NULL, value) >= 0;
}
@ -493,7 +506,6 @@ static bool PyUpb_Message_InitMessageAttribute(PyObject* _self, PyObject* name,
PyObject* value) {
PyObject* submsg = PyUpb_Message_GetAttr(_self, name);
if (!submsg) return -1;
assert(!PyErr_Occurred());
bool ok;
const upb_MessageDef* m_def = upb_FieldDef_MessageSubDef(field);
if (PyDict_Check(value) &&
@ -511,16 +523,16 @@ static bool PyUpb_Message_InitScalarAttribute(upb_Message* msg,
PyObject* value,
upb_Arena* arena) {
upb_MessageValue msgval;
assert(!PyErr_Occurred());
if (!PyUpb_PyToUpb(value, f, &msgval, arena)) return false;
upb_Message_SetFieldByDef(msg, f, msgval, arena);
if (!upb_Message_SetFieldByDef(msg, f, msgval, arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
return true;
}
int PyUpb_Message_InitAttributes(PyObject* _self, PyObject* args,
PyObject* kwargs) {
assert(!PyErr_Occurred());
if (args != NULL && PyTuple_Size(args) != 0) {
PyErr_SetString(PyExc_TypeError, "No positional arguments allowed");
return -1;
@ -532,22 +544,18 @@ int PyUpb_Message_InitAttributes(PyObject* _self, PyObject* args,
Py_ssize_t pos = 0;
PyObject* name;
PyObject* value;
PyUpb_Message_AssureWritable(self);
if (!PyUpb_Message_AssureWritable(self)) return -1;
upb_Message* msg = PyUpb_Message_GetMsg(self);
upb_Arena* arena = PyUpb_Arena_Get(self->arena);
while (PyDict_Next(kwargs, &pos, &name, &value)) {
assert(!PyErr_Occurred());
const upb_FieldDef* f;
assert(!PyErr_Occurred());
if (!PyUpb_Message_LookupName(self, name, &f, NULL, PyExc_ValueError)) {
return -1;
}
if (value == Py_None) continue; // Ignored.
assert(!PyErr_Occurred());
if (upb_FieldDef_IsMap(f)) {
if (!PyUpb_Message_InitMapAttribute(_self, name, f, value)) return -1;
} else if (upb_FieldDef_IsRepeated(f)) {
@ -581,6 +589,10 @@ static PyObject* PyUpb_Message_NewStub(PyObject* parent, const upb_FieldDef* f,
if (!cls) return NULL;
PyUpb_Message* msg = (void*)PyType_GenericAlloc((PyTypeObject*)cls, 0);
if (msg == NULL) {
Py_DECREF(cls);
return NULL;
}
msg->def = (uintptr_t)f | 1;
msg->arena = arena;
msg->ptr.parent = (PyUpb_Message*)parent;
@ -634,19 +646,32 @@ static const upb_FieldDef* PyUpb_Message_InitAsMsg(PyUpb_Message* m,
upb_Arena* arena) {
const upb_FieldDef* f = PyUpb_Message_GetFieldDef(m);
const upb_MessageDef* m2 = upb_FieldDef_MessageSubDef(f);
m->ptr.msg = upb_Message_New(upb_MessageDef_MiniTable(m2), arena);
upb_Message* msg = upb_Message_New(upb_MessageDef_MiniTable(m2), arena);
if (!msg) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
m->ptr.msg = msg;
m->def = (uintptr_t)m2;
PyUpb_ObjCache_Add(m->ptr.msg, &m->ob_base);
if (!PyUpb_ObjCache_Add(m->ptr.msg, &m->ob_base)) {
return NULL;
}
return f;
}
static void PyUpb_Message_SetField(PyUpb_Message* parent, const upb_FieldDef* f,
static bool PyUpb_Message_SetField(PyUpb_Message* parent, const upb_FieldDef* f,
PyUpb_Message* child, upb_Arena* arena) {
upb_MessageValue msgval = {.msg_val = PyUpb_Message_GetMsg(child)};
upb_Message_SetFieldByDef(PyUpb_Message_GetMsg(parent), f, msgval, arena);
if (!upb_Message_SetFieldByDef(PyUpb_Message_GetMsg(parent), f, msgval,
arena)) {
PyErr_SetNone(PyExc_MemoryError);
Py_DECREF(child);
return false;
}
PyUpb_WeakMap_Delete(parent->unset_subobj_map, f);
// Releases a ref previously owned by child->ptr.parent of our child.
Py_DECREF(child);
return true;
}
/*
@ -688,6 +713,7 @@ bool PyUpb_Message_AssureWritable(PyUpb_Message* self) {
PyUpb_Message* child = self;
PyUpb_Message* parent = self->ptr.parent;
const upb_FieldDef* child_f = PyUpb_Message_InitAsMsg(child, arena);
if (!child_f) return false;
Py_INCREF(child); // To avoid a special-case in PyUpb_Message_SetField().
do {
@ -695,8 +721,9 @@ bool PyUpb_Message_AssureWritable(PyUpb_Message* self) {
const upb_FieldDef* parent_f = NULL;
if (PyUpb_Message_IsStub(parent)) {
parent_f = PyUpb_Message_InitAsMsg(parent, arena);
if (!parent_f) goto err2;
}
PyUpb_Message_SetField(parent, child_f, child, arena);
if (!PyUpb_Message_SetField(parent, child_f, child, arena)) goto err1;
child = parent;
child_f = parent_f;
parent = next_parent;
@ -706,9 +733,15 @@ bool PyUpb_Message_AssureWritable(PyUpb_Message* self) {
Py_DECREF(child);
self->version++;
return true;
err2:
Py_DECREF(child);
err1:
Py_DECREF(parent);
return false;
}
static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self);
static bool PyUpb_Message_SyncSubobjs(PyUpb_Message* self);
/*
* PyUpb_Message_Reify()
@ -717,7 +750,7 @@ static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self);
* the wrapper from the unset state (owning a reference on self->ptr.parent) to
* the set state (having a non-owning pointer to self->ptr.msg).
*/
static void PyUpb_Message_Reify(PyUpb_Message* self, const upb_FieldDef* f,
static bool PyUpb_Message_Reify(PyUpb_Message* self, const upb_FieldDef* f,
upb_Message* msg, PyUpb_WeakMap* subobj_map,
intptr_t iter) {
assert(f == PyUpb_Message_GetFieldDef(self));
@ -726,12 +759,18 @@ static void PyUpb_Message_Reify(PyUpb_Message* self, const upb_FieldDef* f,
const upb_MessageDef* msgdef = PyUpb_Message_GetMsgdef((PyObject*)self);
const upb_MiniTable* layout = upb_MessageDef_MiniTable(msgdef);
msg = upb_Message_New(layout, PyUpb_Arena_Get(self->arena));
if (!msg) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
}
if (!PyUpb_ObjCache_Add(msg, &self->ob_base)) {
return false;
}
PyUpb_ObjCache_Add(msg, &self->ob_base);
Py_DECREF(&self->ptr.parent->ob_base);
self->ptr.msg = msg; // Overwrites self->ptr.parent
self->def = (uintptr_t)upb_FieldDef_MessageSubDef(f);
PyUpb_Message_SyncSubobjs(self);
return PyUpb_Message_SyncSubobjs(self);
}
/*
@ -752,9 +791,9 @@ static void PyUpb_Message_Reify(PyUpb_Message* self, const upb_FieldDef* f,
* This requires that all of the new sub-objects that have appeared are owned
* by `self`'s arena.
*/
static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self) {
static bool PyUpb_Message_SyncSubobjs(PyUpb_Message* self) {
PyUpb_WeakMap* subobj_map = self->unset_subobj_map;
if (!subobj_map) return;
if (!subobj_map) return true;
upb_Message* msg = PyUpb_Message_GetMsg(self);
intptr_t iter = PYUPB_WEAKMAP_BEGIN;
@ -768,6 +807,7 @@ static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self) {
// done iterating.
Py_INCREF(&self->ob_base);
bool ok = true;
while (PyUpb_WeakMap_Next(subobj_map, &key, &obj, &iter)) {
const upb_FieldDef* f = key;
if (upb_FieldDef_HasPresence(f) && !upb_Message_HasFieldByDef(msg, f))
@ -775,23 +815,33 @@ static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self) {
upb_MessageValue msgval = upb_Message_GetFieldByDef(msg, f);
if (upb_FieldDef_IsMap(f)) {
if (!msgval.map_val) continue;
PyUpb_MapContainer_Reify(obj, (upb_Map*)msgval.map_val, subobj_map, iter);
if (!PyUpb_MapContainer_Reify(obj, (upb_Map*)msgval.map_val, subobj_map,
iter)) {
ok = false;
break;
}
} else if (upb_FieldDef_IsRepeated(f)) {
if (!msgval.array_val) continue;
PyUpb_RepeatedContainer_Reify(obj, (upb_Array*)msgval.array_val,
subobj_map, iter);
if (!PyUpb_RepeatedContainer_Reify(obj, (upb_Array*)msgval.array_val,
subobj_map, iter)) {
ok = false;
break;
}
} else {
PyUpb_Message* sub = (void*)obj;
assert(self == sub->ptr.parent);
PyUpb_Message_Reify(sub, f, (upb_Message*)msgval.msg_val, subobj_map,
iter);
if (!PyUpb_Message_Reify(sub, f, (upb_Message*)msgval.msg_val, subobj_map,
iter)) {
ok = false;
break;
}
}
}
Py_DECREF(&self->ob_base);
// TODO: present fields need to be iterated too if they can reach
// a WeakMap.
return ok;
}
static PyObject* PyUpb_Message_ToString(PyUpb_Message* self) {
@ -851,8 +901,11 @@ bool PyUpb_Message_SetConcreteSubobj(PyObject* _self, const upb_FieldDef* f,
PyUpb_Message* self = (void*)_self;
if (!PyUpb_Message_AssureWritable(self)) return false;
PyUpb_Message_CacheDelete(_self, f);
upb_Message_SetFieldByDef(self->ptr.msg, f, subobj,
PyUpb_Arena_Get(self->arena));
if (!upb_Message_SetFieldByDef(self->ptr.msg, f, subobj,
PyUpb_Arena_Get(self->arena))) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
return true;
}
@ -863,7 +916,7 @@ static void PyUpb_Message_Dealloc(PyObject* _self) {
PyUpb_Message_CacheDelete((PyObject*)self->ptr.parent,
PyUpb_Message_GetFieldDef(self));
Py_DECREF(self->ptr.parent);
} else {
} else if (self->ptr.msg) {
PyUpb_ObjCache_Delete(self->ptr.msg);
}
@ -871,7 +924,7 @@ static void PyUpb_Message_Dealloc(PyObject* _self) {
PyUpb_WeakMap_Free(self->unset_subobj_map);
}
Py_DECREF(self->arena);
Py_XDECREF(self->arena);
PyUpb_Dealloc(self);
}
@ -885,6 +938,11 @@ PyObject* PyUpb_Message_Get(upb_Message* u_msg, const upb_MessageDef* m,
// It is not safe to use PyObject_{,GC}_New() due to:
// https://bugs.python.org/issue35810
PyUpb_Message* py_msg = (void*)PyType_GenericAlloc((PyTypeObject*)cls, 0);
if (!py_msg) {
Py_DECREF(ret);
Py_DECREF(cls);
return NULL;
}
py_msg->arena = arena;
py_msg->def = (uintptr_t)m;
py_msg->ptr.msg = u_msg;
@ -894,7 +952,10 @@ PyObject* PyUpb_Message_Get(upb_Message* u_msg, const upb_MessageDef* m,
ret = &py_msg->ob_base;
Py_DECREF(cls);
Py_INCREF(arena);
PyUpb_ObjCache_Add(u_msg, ret);
if (!PyUpb_ObjCache_Add(u_msg, ret)) {
Py_DECREF(ret);
return NULL;
}
return ret;
}
@ -917,6 +978,7 @@ PyObject* PyUpb_Message_GetStub(PyUpb_Message* self,
PyObject* _self = (void*)self;
if (!self->unset_subobj_map) {
self->unset_subobj_map = PyUpb_WeakMap_New();
if (!self->unset_subobj_map) return NULL;
}
PyObject* subobj = PyUpb_WeakMap_Get(self->unset_subobj_map, field);
@ -929,8 +991,10 @@ PyObject* PyUpb_Message_GetStub(PyUpb_Message* self,
} else {
subobj = PyUpb_Message_NewStub(&self->ob_base, field, self->arena);
}
if (!subobj) return NULL;
PyUpb_WeakMap_Add(self->unset_subobj_map, field, subobj);
if (!PyUpb_WeakMap_Add(self->unset_subobj_map, field, subobj)) {
Py_DECREF(subobj);
return NULL;
}
return subobj;
}
@ -949,11 +1013,17 @@ PyObject* PyUpb_Message_GetPresentWrapper(PyUpb_Message* self,
}
if (upb_FieldDef_IsMap(field)) {
assert(val.map_val);
if (!val.map_val) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
return PyUpb_MapContainer_GetOrCreateWrapper((upb_Map*)val.map_val, field,
self->arena);
} else {
assert(val.array_val);
if (!val.array_val) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
return PyUpb_RepeatedContainer_GetOrCreateWrapper((upb_Array*)val.array_val,
field, self->arena);
}
@ -1049,7 +1119,10 @@ int PyUpb_Message_SetFieldValue(PyObject* _self, const upb_FieldDef* field,
return -1;
}
upb_Message_SetFieldByDef(self->ptr.msg, field, val, arena);
if (!upb_Message_SetFieldByDef(self->ptr.msg, field, val, arena)) {
PyErr_SetNone(PyExc_MemoryError);
return -1;
}
return 0;
}
@ -1109,7 +1182,6 @@ __attribute__((flatten)) static PyObject* PyUpb_Message_GetAttr(
}
// Check base class attributes.
assert(!PyErr_Occurred());
PyObject* ret = PyObject_GenericGetAttr(_self, attr);
if (ret) return ret;
@ -1379,15 +1451,19 @@ static PyObject* PyUpb_Message_CopyFrom(PyObject* _self, PyObject* arg) {
const upb_Message* other_msg = PyUpb_Message_GetIfReified((PyObject*)other);
if (other_msg) {
upb_Message_DeepCopy(
self->ptr.msg, other_msg,
upb_MessageDef_MiniTable((const upb_MessageDef*)other->def),
PyUpb_Arena_Get(self->arena));
if (!upb_Message_DeepCopy(
self->ptr.msg, other_msg,
upb_MessageDef_MiniTable((const upb_MessageDef*)other->def),
PyUpb_Arena_Get(self->arena))) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
} else {
PyObject* tmp = PyUpb_Message_Clear(self);
if (!tmp) return NULL;
Py_DECREF(tmp);
}
PyUpb_Message_SyncSubobjs(self);
if (!PyUpb_Message_SyncSubobjs(self)) return NULL;
Py_RETURN_NONE;
}
@ -1471,11 +1547,15 @@ PyObject* PyUpb_Message_MergeFromString(PyObject* _self, PyObject* arg) {
upb_DecodeStatus status =
upb_Decode(buf, size, self->ptr.msg, layout, extreg, options, arena);
Py_XDECREF(mv_contiguous);
PyUpb_Message_SyncSubobjs(self);
if (!PyUpb_Message_SyncSubobjs(self)) return NULL;
if (status != kUpb_DecodeStatus_Ok) {
PyErr_Format(
state->decode_error_class, "Error parsing message with type '%s': %s",
upb_MessageDef_FullName(msgdef), upb_DecodeStatus_String(status));
if (status == kUpb_DecodeStatus_OutOfMemory) {
PyErr_SetNone(PyExc_MemoryError);
} else {
PyErr_Format(
state->decode_error_class, "Error parsing message with type '%s': %s",
upb_MessageDef_FullName(msgdef), upb_DecodeStatus_String(status));
}
return NULL;
}
return PyLong_FromSsize_t(size);
@ -1483,6 +1563,7 @@ PyObject* PyUpb_Message_MergeFromString(PyObject* _self, PyObject* arg) {
static PyObject* PyUpb_Message_ParseFromString(PyObject* self, PyObject* arg) {
PyObject* tmp = PyUpb_Message_Clear((PyUpb_Message*)self);
if (!tmp) return NULL;
Py_DECREF(tmp);
return PyUpb_Message_MergeFromString(self, arg);
}
@ -1516,15 +1597,17 @@ static PyObject* PyUpb_Message_Clear(PyUpb_Message* self) {
const upb_FieldDef* f = key;
if (upb_FieldDef_IsMap(f)) {
assert(upb_Message_GetFieldByDef(msg, f).map_val == NULL);
PyUpb_MapContainer_Reify(obj, NULL, subobj_map, iter);
if (!PyUpb_MapContainer_Reify(obj, NULL, subobj_map, iter)) return NULL;
} else if (upb_FieldDef_IsRepeated(f)) {
assert(upb_Message_GetFieldByDef(msg, f).array_val == NULL);
PyUpb_RepeatedContainer_Reify(obj, NULL, subobj_map, iter);
if (!PyUpb_RepeatedContainer_Reify(obj, NULL, subobj_map, iter)) {
return NULL;
}
} else {
assert(!upb_Message_HasFieldByDef(msg, f));
PyUpb_Message* sub = (void*)obj;
assert(self == sub->ptr.parent);
PyUpb_Message_Reify(sub, f, NULL, subobj_map, iter);
if (!PyUpb_Message_Reify(sub, f, NULL, subobj_map, iter)) return NULL;
}
}
}
@ -1744,6 +1827,10 @@ PyObject* PyUpb_Message_SerializeInternal(PyObject* _self, PyObject* args,
}
upb_Arena* arena = upb_Arena_New();
if (!arena) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
const upb_MiniTable* layout = upb_MessageDef_MiniTable(msgdef);
size_t size = 0;
// Python does not currently have any effective limit on serialization depth.
@ -1756,13 +1843,17 @@ PyObject* PyUpb_Message_SerializeInternal(PyObject* _self, PyObject* args,
PyObject* ret = NULL;
if (status != kUpb_EncodeStatus_Ok) {
PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
PyObject* errors = PyUpb_Message_FindInitializationErrors(_self, NULL);
if (PyList_Size(errors) != 0) {
PyUpb_Message_ReportInitializationErrors(msgdef, errors,
state->encode_error_class);
if (status == kUpb_EncodeStatus_OutOfMemory) {
PyErr_SetNone(PyExc_MemoryError);
} else {
PyErr_Format(state->encode_error_class, "Failed to serialize proto");
PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
PyObject* errors = PyUpb_Message_FindInitializationErrors(_self, NULL);
if (PyList_Size(errors) != 0) {
PyUpb_Message_ReportInitializationErrors(msgdef, errors,
state->encode_error_class);
} else {
PyErr_Format(state->encode_error_class, "Failed to serialize proto");
}
}
goto done;
}
@ -1803,10 +1894,16 @@ PyObject* DeepCopy(PyObject* _self, PyObject* arg) {
const upb_MiniTable* mini_table = upb_MessageDef_MiniTable(def);
upb_Message* msg = PyUpb_Message_GetIfReified(_self);
PyObject* arena = PyUpb_Arena_New();
if (!arena) return NULL;
upb_Arena* upb_arena = PyUpb_Arena_Get(arena);
upb_Message* clone = msg ? upb_Message_DeepClone(msg, mini_table, upb_arena)
: upb_Message_New(mini_table, upb_arena);
if (!clone) {
Py_DECREF(arena);
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
PyObject* ret = PyUpb_Message_Get(clone, def, arena);
Py_DECREF(arena);

94
python/protobuf.c
View file

@ -81,9 +81,47 @@ PyObject* PyUpb_SetAllowOversizeProtos(PyObject* m, PyObject* arg) {
return arg;
}
PyObject* PyUpb_AllocationCount_IsAvailable(PyObject* self, PyObject* args) {
if (upb_AllocationCount_IsAvailable()) {
Py_RETURN_TRUE;
} else {
Py_RETURN_FALSE;
}
}
PyObject* PyUpb_AllocationCount_Get(PyObject* self, PyObject* args) {
return PyLong_FromSize_t(upb_AllocationCount_Get());
}
PyObject* PyUpb_AllocationCount_Reset(PyObject* self, PyObject* args) {
upb_AllocationCount_Reset();
Py_RETURN_NONE;
}
PyObject* PyUpb_AllocationCount_FailOn(PyObject* self, PyObject* arg) {
Py_ssize_t n = PyNumber_AsSsize_t(arg, PyExc_OverflowError);
if (n == -1 && PyErr_Occurred()) {
return NULL;
}
if (n < 0) {
PyErr_SetString(PyExc_ValueError, "FailOn must be non-negative");
return NULL;
}
upb_AllocationCount_FailOn((size_t)n);
Py_RETURN_NONE;
}
static PyMethodDef PyUpb_ModuleMethods[] = {
{"SetAllowOversizeProtos", PyUpb_SetAllowOversizeProtos, METH_O,
"Enable/disable oversize proto parsing."},
{"_AllocationCount_IsAvailable", PyUpb_AllocationCount_IsAvailable,
METH_NOARGS, "Returns whether allocation count debugging is available."},
{"_AllocationCount_Get", PyUpb_AllocationCount_Get, METH_NOARGS,
"Returns the current allocation.count."},
{"_AllocationCount_Reset", PyUpb_AllocationCount_Reset, METH_NOARGS,
"Resets the current allocation count and failure settings."},
{"_AllocationCount_FailOn", PyUpb_AllocationCount_FailOn, METH_O,
"Configures allocation failure at the N-th allocation."},
{NULL, NULL}};
static struct PyModuleDef module_def = {PyModuleDef_HEAD_INIT,
@ -155,12 +193,24 @@ struct PyUpb_WeakMap {
PyUpb_WeakMap* PyUpb_WeakMap_New(void) {
upb_Arena* arena = PyUpb_NewArena();
if (!arena) {
return NULL;
}
PyUpb_WeakMap* map = upb_Arena_Malloc(arena, sizeof(*map));
if (!map) {
upb_Arena_Free(arena);
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
map->arena = arena;
#ifdef ENABLE_MUTEX
pthread_mutex_init(&map->mutex.mutex, NULL);
#endif
upb_inttable_init(&map->table, map->arena);
if (!upb_inttable_init(&map->table, map->arena)) {
upb_Arena_Free(arena);
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
return map;
}
@ -181,14 +231,18 @@ uintptr_t PyUpb_WeakMap_GetKey(const void* key) {
return n >> PyUpb_PtrShift;
}
void PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject* py_obj) {
bool PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject* py_obj) {
#ifdef Py_GIL_DISABLED
PyUnstable_EnableTryIncRef(py_obj);
#endif
FreeThreadingLock(&map->mutex);
upb_inttable_insert(&map->table, PyUpb_WeakMap_GetKey(key),
upb_value_ptr(py_obj), map->arena);
bool ok = upb_inttable_insert(&map->table, PyUpb_WeakMap_GetKey(key),
upb_value_ptr(py_obj), map->arena);
if (!ok) {
PyErr_SetNone(PyExc_MemoryError);
}
FreeThreadingUnlock(&map->mutex);
return ok;
}
void PyUpb_WeakMap_Delete(PyUpb_WeakMap* map, const void* key) {
@ -287,17 +341,17 @@ static PyUpb_WeakMap* PyUpb_ObjCache_MaybeInstance(void) {
return state->obj_cache;
}
void PyUpb_ObjCache_Add(const void* key, PyObject* py_obj) {
bool PyUpb_ObjCache_Add(const void* key, PyObject* py_obj) {
PyUpb_WeakMap* cache = PyUpb_ObjCache_MaybeInstance();
if (!cache) {
return;
return true;
}
PyUpb_WeakMap_Add(cache, key, py_obj);
return PyUpb_WeakMap_Add(cache, key, py_obj);
}
void PyUpb_KnownObjCache_Add(PyUpb_WeakMap* cache, const void* key,
bool PyUpb_KnownObjCache_Add(PyUpb_WeakMap* cache, const void* key,
PyObject* py_obj) {
PyUpb_WeakMap_Add(cache, key, py_obj);
return PyUpb_WeakMap_Add(cache, key, py_obj);
}
void PyUpb_ObjCache_Delete(const void* key) {
@ -381,19 +435,35 @@ static upb_alloc trim_alloc = {&upb_trim_allocfunc};
static upb_alloc* global_alloc = &trim_alloc;
static upb_Arena* PyUpb_NewArena(void) {
return upb_Arena_Init(NULL, 0, global_alloc);
upb_Arena* arena = upb_Arena_Init(NULL, 0, global_alloc);
if (!arena) {
PyErr_SetNone(PyExc_MemoryError);
}
return arena;
}
PyObject* PyUpb_Arena_New(void) {
PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
PyUpb_ModuleState* state = PyUpb_ModuleState_MaybeGet();
if (!state) {
PyErr_SetString(PyExc_RuntimeError, "Interpreter is finalizing");
return NULL;
}
PyUpb_Arena* arena = (void*)PyType_GenericAlloc(state->arena_type, 0);
if (!arena) return NULL;
arena->arena = PyUpb_NewArena();
if (!arena->arena) {
Py_DECREF(arena);
return NULL;
}
arena->frozen = false;
return &arena->ob_base;
}
static void PyUpb_Arena_Dealloc(PyObject* self) {
upb_Arena_Free(PyUpb_Arena_Get(self));
upb_Arena* arena = PyUpb_Arena_Get(self);
if (arena) {
upb_Arena_Free(arena);
}
PyUpb_Dealloc(self);
}

12
python/protobuf.h
View file

@ -125,7 +125,7 @@ PyUpb_WeakMap* PyUpb_WeakMap_New(void);
void PyUpb_WeakMap_Free(PyUpb_WeakMap* map);
// Adds the given object to the map, indexed by the given key.
void PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject* py_obj);
bool PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject* py_obj);
// Removes the given key from the cache. It must exist in the cache currently.
void PyUpb_WeakMap_Delete(PyUpb_WeakMap* map, const void* key);
@ -160,8 +160,8 @@ void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap* map, intptr_t* iter);
// The object cache is a global WeakMap for mapping upb objects to the
// corresponding wrapper.
void PyUpb_ObjCache_Add(const void* key, PyObject* py_obj);
void PyUpb_KnownObjCache_Add(PyUpb_WeakMap* cache, const void* key,
bool PyUpb_ObjCache_Add(const void* key, PyObject* py_obj);
bool PyUpb_KnownObjCache_Add(PyUpb_WeakMap* cache, const void* key,
PyObject* py_obj);
void PyUpb_ObjCache_Delete(const void* key);
PyObject* PyUpb_ObjCache_Get(const void* key); // returns NULL if not present.
@ -207,7 +207,11 @@ static inline void PyUpb_Dealloc(void* self) {
PyTypeObject* tp = Py_TYPE(self);
assert(PyType_GetFlags(tp) & Py_TPFLAGS_HEAPTYPE);
freefunc tp_free = (freefunc)PyType_GetSlot(tp, Py_tp_free);
tp_free(self);
if (tp_free) {
tp_free(self);
} else {
PyObject_Free(self);
}
Py_DECREF(tp);
}

106
python/repeated.c
View file

@ -16,8 +16,15 @@
#include "google/protobuf/breaking_changes.h"
#include "python/buffer_convert.h"
#include "python/convert.h"
#include "python/descriptor.h"
#include "python/message.h"
#include "python/protobuf.h"
#include "upb/base/descriptor_constants.h"
#include "upb/mem/arena.h"
#include "upb/message/array.h"
#include "upb/message/message.h"
#include "upb/mini_table/message.h"
#include "upb/reflection/def.h"
// Must be last.
#include "upb/port/def.inc"
@ -77,6 +84,10 @@ upb_Array* PyUpb_RepeatedContainer_Reify(PyObject* _self, upb_Array* arr,
if (!arr) {
upb_Arena* arena = PyUpb_Arena_Get(self->arena);
arr = upb_Array_New(arena, upb_FieldDef_CType(f));
if (!arr) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
}
if (subobj_map) {
PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
@ -86,7 +97,9 @@ upb_Array* PyUpb_RepeatedContainer_Reify(PyObject* _self, upb_Array* arr,
return NULL;
}
}
PyUpb_ObjCache_Add(arr, &self->ob_base);
if (!PyUpb_ObjCache_Add(arr, &self->ob_base)) {
return NULL;
}
Py_DECREF(self->ptr.parent);
self->ptr.arr = arr; // Overwrites self->ptr.parent.
self->field &= ~(uintptr_t)1;
@ -154,6 +167,7 @@ PyObject* PyUpb_RepeatedContainer_NewStub(PyObject* parent,
return NULL;
}
PyUpb_RepeatedContainer* repeated = (void*)PyType_GenericAlloc(cls, 0);
if (repeated == NULL) return NULL;
repeated->arena = arena;
repeated->field = (uintptr_t)PyUpb_FieldDescriptor_Get(f) | 1;
repeated->ptr.parent = parent;
@ -180,7 +194,10 @@ PyObject* PyUpb_RepeatedContainer_GetOrCreateWrapper(upb_Array* arr,
repeated->ptr.arr = arr;
ret = &repeated->ob_base;
Py_INCREF(arena);
PyUpb_ObjCache_Add(arr, ret);
if (!PyUpb_ObjCache_Add(arr, ret)) {
Py_DECREF(ret);
return NULL;
}
return ret;
}
@ -194,17 +211,27 @@ PyObject* PyUpb_RepeatedContainer_DeepCopy(PyObject* _self, PyObject* value) {
if (clone == NULL) return NULL;
const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
clone->arena = PyUpb_Arena_New();
if (clone->arena == NULL) goto err;
clone->field = (uintptr_t)PyUpb_FieldDescriptor_Get(f);
clone->ptr.arr =
upb_Array_New(PyUpb_Arena_Get(clone->arena), upb_FieldDef_CType(f));
PyUpb_ObjCache_Add(clone->ptr.arr, (PyObject*)clone);
if (clone->ptr.arr == NULL) {
PyErr_SetNone(PyExc_MemoryError);
goto err;
}
if (!PyUpb_ObjCache_Add(clone->ptr.arr, (PyObject*)clone)) {
goto err;
}
PyObject* result = PyUpb_RepeatedContainer_MergeFrom((PyObject*)clone, _self);
if (!result) {
Py_DECREF(clone);
return NULL;
goto err;
}
Py_DECREF(result);
return (PyObject*)clone;
err:
Py_DECREF(clone);
return NULL;
}
#if PyUpb_SUPPORT_BUFFER_VIEW
@ -421,14 +448,21 @@ static bool PyUpb_ExtendSizeCb(Py_ssize_t size, void* vctx) {
ctx->size_hint = size;
size_t old_size = upb_Array_Size(ctx->arr);
if (size > 0 && ((size_t)size <= SIZE_MAX - old_size)) {
upb_Array_Reserve(ctx->arr, old_size + size, ctx->arena);
if (!upb_Array_Reserve(ctx->arr, old_size + size, ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
}
return true;
}
static bool PyUpb_ExtendElemCb(upb_MessageValue val, void* vctx) {
PyUpb_ExtendCtx* ctx = (PyUpb_ExtendCtx*)vctx;
return upb_Array_Append(ctx->arr, val, ctx->arena);
if (!upb_Array_Append(ctx->arr, val, ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
return true;
}
typedef enum {
@ -461,12 +495,18 @@ static bool PyUpb_ExtendBulkCb(const void* data, Py_ssize_t count,
char* dst;
switch (PyUpb_ArrayOverlaps(ctx->arr, data, count, itemsize)) {
case kDisjoint:
upb_Array_Resize(ctx->arr, old_size + count, ctx->arena);
if (!upb_Array_Resize(ctx->arr, old_size + count, ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
dst = upb_Array_MutableDataPtr(ctx->arr);
break;
case kSubset: {
char* old_dst = upb_Array_MutableDataPtr(ctx->arr);
upb_Array_Resize(ctx->arr, old_size + count, ctx->arena);
if (!upb_Array_Resize(ctx->arr, old_size + count, ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
dst = upb_Array_MutableDataPtr(ctx->arr);
if (old_dst != dst) {
data = dst + ((const char*)data - old_dst);
@ -511,7 +551,7 @@ PyObject* PyUpb_RepeatedContainer_Extend(PyObject* _self, PyObject* value) {
PyUpb_ExtendCtx ctx = {arr, arena};
if (!PyUpb_IterInput(value, f, arena, PyUpb_ExtendSizeCb, PyUpb_ExtendElemCb,
PyUpb_ExtendBulkCb, &ctx)) {
upb_Array_Resize(arr, old_size, NULL);
(void)upb_Array_Resize(arr, old_size, NULL);
return NULL;
}
Py_RETURN_NONE;
@ -616,7 +656,11 @@ static bool PyUpb_SetSubscriptSizeCb(Py_ssize_t seq_size, void* vctx) {
if (ctx->step == 1) {
// We must shift the tail elements (either right or left).
size_t tail = upb_Array_Size(ctx->arr) - (ctx->index + ctx->count);
upb_Array_Resize(ctx->arr, ctx->index + seq_size + tail, ctx->arena);
if (!upb_Array_Resize(ctx->arr, ctx->index + seq_size + tail,
ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
upb_Array_Move(ctx->arr, ctx->index + seq_size, ctx->index + ctx->count,
tail);
ctx->count = seq_size;
@ -684,7 +728,10 @@ static bool PyUpb_SetSubscriptBulkCb(const void* data, Py_ssize_t count,
count, ctx->count);
return false;
}
upb_Array_Resize(ctx->arr, ctx->index + count + tail, ctx->arena);
if (!upb_Array_Resize(ctx->arr, ctx->index + count + tail, ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
dst = upb_Array_MutableDataPtr(ctx->arr);
upb_Array_Move(ctx->arr, ctx->index + count, ctx->index + ctx->count,
tail);
@ -718,7 +765,7 @@ static bool PyUpb_SetSubscriptBulkCb(const void* data, Py_ssize_t count,
if (count < ctx->count) {
upb_Array_Move(ctx->arr, ctx->index + count, ctx->index + ctx->count,
tail);
upb_Array_Resize(ctx->arr, ctx->index + count + tail, ctx->arena);
(void)upb_Array_Resize(ctx->arr, ctx->index + count + tail, ctx->arena);
}
return true;
}
@ -728,7 +775,10 @@ static bool PyUpb_SetSubscriptBulkCb(const void* data, Py_ssize_t count,
// Append to the end of the array.
char* old_dst = upb_Array_MutableDataPtr(ctx->arr);
upb_Array_Resize(ctx->arr, old_size + count, ctx->arena);
if (!upb_Array_Resize(ctx->arr, old_size + count, ctx->arena)) {
PyErr_SetNone(PyExc_MemoryError);
return false;
}
dst = upb_Array_MutableDataPtr(ctx->arr);
if (old_dst != dst) {
data = dst + ((const char*)data - old_dst);
@ -756,7 +806,7 @@ static bool PyUpb_SetSubscriptBulkCb(const void* data, Py_ssize_t count,
}
}
upb_Array_Resize(ctx->arr, ctx->index + count + tail, ctx->arena);
(void)upb_Array_Resize(ctx->arr, ctx->index + count + tail, ctx->arena);
return true;
}
@ -828,7 +878,7 @@ static int PyUpb_RepeatedContainer_DeleteSubscript(upb_Array* arr,
size_t new_size = dst + tail;
assert(new_size == upb_Array_Size(arr) - count);
upb_Array_Move(arr, dst, src, tail);
upb_Array_Resize(arr, new_size, NULL);
(void)upb_Array_Resize(arr, new_size, NULL);
return 0;
}
@ -1012,8 +1062,15 @@ static PyObject* PyUpb_RepeatedCompositeContainer_AppendNew(PyObject* _self) {
const upb_MessageDef* m = upb_FieldDef_MessageSubDef(f);
const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
upb_Message* msg = upb_Message_New(layout, arena);
if (!msg) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
upb_MessageValue msgval = {.msg_val = msg};
upb_Array_Append(arr, msgval, arena);
if (!upb_Array_Append(arr, msgval, arena)) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
return PyUpb_Message_Get(msg, m, self->arena);
}
@ -1067,6 +1124,10 @@ static PyObject* PyUpb_RepeatedContainer_Insert(PyObject* _self,
const upb_MessageDef* m = upb_FieldDef_MessageSubDef(f);
const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
upb_Message* msg = upb_Message_New(layout, arena);
if (!msg) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
PyObject* py_msg = PyUpb_Message_Get(msg, m, self->arena);
PyObject* ret = PyUpb_Message_MergeFrom(py_msg, value);
Py_DECREF(py_msg);
@ -1077,7 +1138,10 @@ static PyObject* PyUpb_RepeatedContainer_Insert(PyObject* _self,
if (!PyUpb_PyToUpb(value, f, &msgval, arena)) return NULL;
}
upb_Array_Insert(arr, index, 1, arena);
if (!upb_Array_Insert(arr, index, 1, arena)) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
upb_Array_Set(arr, index, msgval);
Py_RETURN_NONE;
@ -1145,7 +1209,10 @@ static PyObject* PyUpb_RepeatedScalarContainer_Append(PyObject* _self,
if (!PyUpb_PyToUpb(value, f, &msgval, arena)) {
return NULL;
}
upb_Array_Append(arr, msgval, arena);
if (!upb_Array_Append(arr, msgval, arena)) {
PyErr_SetNone(PyExc_MemoryError);
return NULL;
}
Py_RETURN_NONE;
}
@ -1371,7 +1438,6 @@ ret:
Py_XDECREF(default_dtype);
Py_XDECREF(nparray);
Py_DECREF(np_module);
assert(!PyErr_Occurred());
return return_value;
}

1
upb/mem/BUILD
View file

@ -34,6 +34,7 @@ cc_library(
cc_library(
name = "internal",
hdrs = [
"internal/alloc.h",
"internal/arena.h",
],
copts = UPB_DEFAULT_COPTS,

49
upb/mem/alloc.c
View file

@ -7,6 +7,7 @@
#include "upb/mem/alloc.h"
#include <stdint.h>
#include <stdlib.h>
// Must be last.
@ -25,4 +26,52 @@ static void* upb_global_allocfunc(upb_alloc* alloc, void* ptr, size_t oldsize,
}
}
#ifdef UPB_ALLOCATION_COUNT
#if (defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) && \
!defined(__STDC_NO_THREADS__)) || \
UPB_HAS_EXTENSION(c_thread_local)
#define UPB_THREAD_LOCAL _Thread_local
#elif defined(_MSC_VER)
#define UPB_THREAD_LOCAL __declspec(thread)
#elif defined(__GNUC__) || defined(__clang__)
#define UPB_THREAD_LOCAL __thread
#else
#define UPB_THREAD_LOCAL
#endif
UPB_THREAD_LOCAL size_t upb_arena_alloc_count = 0;
UPB_THREAD_LOCAL size_t upb_arena_alloc_fail_on = SIZE_MAX;
#undef UPB_THREAD_LOCAL
#endif
UPB_NODISCARD bool upb_AllocationCount_IsAvailable(void) {
#ifdef UPB_ALLOCATION_COUNT
return true;
#else
return false;
#endif
}
UPB_NODISCARD size_t upb_AllocationCount_Get(void) {
#ifdef UPB_ALLOCATION_COUNT
return upb_arena_alloc_count;
#else
return 0;
#endif
}
void upb_AllocationCount_Reset(void) {
#ifdef UPB_ALLOCATION_COUNT
upb_arena_alloc_count = 0;
upb_arena_alloc_fail_on = SIZE_MAX;
#endif
}
void upb_AllocationCount_FailOn(size_t n) {
#ifdef UPB_ALLOCATION_COUNT
upb_arena_alloc_fail_on = n;
#endif
}
upb_alloc upb_alloc_global = {&upb_global_allocfunc};

26
upb/mem/alloc.h
View file

@ -10,6 +10,8 @@
#include <stddef.h>
#include "upb/mem/internal/alloc.h"
// Must be last.
#include "upb/port/def.inc"
@ -41,6 +43,9 @@ struct upb_alloc {
UPB_NODISCARD UPB_INLINE void* upb_malloc(upb_alloc* alloc, size_t size) {
UPB_ASSERT(alloc);
if (!upb_AllocationCount_IncrementAndCheck()) {
return NULL;
}
return alloc->func(alloc, NULL, 0, size, NULL);
}
@ -52,6 +57,11 @@ typedef struct {
UPB_INLINE upb_SizedPtr upb_SizeReturningMalloc(upb_alloc* alloc, size_t size) {
UPB_ASSERT(alloc);
upb_SizedPtr result;
if (!upb_AllocationCount_IncrementAndCheck()) {
result.p = NULL;
result.n = 0;
return result;
}
result.n = 0;
result.p = alloc->func(alloc, NULL, 0, size, &result.n);
result.n = result.p != NULL ? UPB_MAX(result.n, size) : 0;
@ -61,6 +71,11 @@ UPB_INLINE upb_SizedPtr upb_SizeReturningMalloc(upb_alloc* alloc, size_t size) {
UPB_NODISCARD UPB_INLINE void* upb_realloc(upb_alloc* alloc, void* ptr,
size_t oldsize, size_t size) {
UPB_ASSERT(alloc);
if (size != 0) {
if (!upb_AllocationCount_IncrementAndCheck()) {
return NULL;
}
}
return alloc->func(alloc, ptr, oldsize, size, NULL);
}
@ -94,6 +109,17 @@ UPB_NODISCARD UPB_INLINE void* upb_grealloc(void* ptr, size_t oldsize,
UPB_INLINE void upb_gfree(void* ptr) { upb_free(&upb_alloc_global, ptr); }
// Returns whether thread-local allocation count/ OOM-simulation features
// are supported.
UPB_API UPB_NODISCARD bool upb_AllocationCount_IsAvailable(void);
// Returns the thread-local allocation count since the last reset.
UPB_API UPB_NODISCARD size_t upb_AllocationCount_Get(void);
// Resets the thread-local allocation count and failure threshold.
UPB_API void upb_AllocationCount_Reset(void);
// Artificially triggers memory allocation failure in the thread on the n-th
// allocation.
UPB_API void upb_AllocationCount_FailOn(size_t n);
#ifdef __cplusplus
} /* extern "C" */
#endif

6
upb/mem/arena.c
View file

@ -9,6 +9,7 @@
#include <string.h>
#include "upb/mem/internal/alloc.h"
#include "upb/port/sanitizers.h"
#ifdef UPB_TRACING_ENABLED
@ -503,11 +504,14 @@ void* UPB_PRIVATE(_upb_Arena_SlowMalloc)(upb_Arena* a, size_t span) {
} else {
UPB_PRIVATE(_upb_Arena_UseBlock)(a, block, block_size);
UPB_ASSERT(UPB_PRIVATE(_upb_ArenaHas)(a) >= span);
return upb_Arena_Malloc(a, size);
return _upb_Arena_Malloc_Unchecked(a, size);
}
}
static upb_Arena* _upb_Arena_InitSlow(upb_alloc* alloc, size_t first_size) {
if (!upb_AllocationCount_IncrementAndCheck()) {
return NULL;
}
if (!alloc) return NULL;
// We need to malloc the initial block.

39
upb/mem/arena_test.cc
View file

@ -1020,4 +1020,43 @@ TEST(ArenaDeathTest, ArenaRefFuseCycle) {
#endif // UPB_SUPPRESS_MISSING_ATOMICS
TEST(ArenaTest, AllocationCountFailureInjection) {
if (!upb_AllocationCount_IsAvailable()) {
return;
}
// Try normal scenario
upb_AllocationCount_Reset();
upb_Arena* arena = upb_Arena_New();
EXPECT_NE(arena, nullptr);
// Allocate some blocks
for (int i = 0; i < 10; ++i) {
void* p = upb_Arena_Malloc(arena, 500);
EXPECT_NE(p, nullptr);
}
size_t total = upb_AllocationCount_Get();
EXPECT_GT(total, 0);
upb_Arena_Free(arena);
// Now verify failure after i allocations
for (size_t i = 0; i < total; ++i) {
upb_AllocationCount_Reset();
upb_AllocationCount_FailOn(i);
// The i-th arena-level initial or block allocation should fail.
upb_Arena* fail_arena = upb_Arena_New();
if (fail_arena != nullptr) {
bool failed = false;
for (int j = 0; j < 10; ++j) {
void* p = upb_Arena_Malloc(fail_arena, 500);
if (p == nullptr) {
failed = true;
break;
}
}
upb_Arena_Free(fail_arena);
EXPECT_TRUE(failed);
}
}
upb_AllocationCount_Reset();
}
} // namespace

53
upb/mem/internal/alloc.h Normal file
View file

@ -0,0 +1,53 @@
#ifndef GOOGLE_UPB_UPB_MEM_INTERNAL_ALLOC_H__
#define GOOGLE_UPB_UPB_MEM_INTERNAL_ALLOC_H__
#include <stddef.h>
// Must be last.
#include "upb/port/def.inc"
#ifdef __cplusplus
extern "C" {
#endif
#if (defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) && \
!defined(__STDC_NO_THREADS__)) || \
UPB_HAS_EXTENSION(c_thread_local)
#define UPB_THREAD_LOCAL _Thread_local
#elif defined(_MSC_VER)
#define UPB_THREAD_LOCAL __declspec(thread)
#elif defined(__GNUC__) || defined(__clang__)
#define UPB_THREAD_LOCAL __thread
#else
#define UPB_THREAD_LOCAL
#endif
#if !defined(NDEBUG)
#define UPB_ALLOCATION_COUNT
#endif
#ifdef UPB_ALLOCATION_COUNT
extern UPB_THREAD_LOCAL size_t upb_arena_alloc_count;
extern UPB_THREAD_LOCAL size_t upb_arena_alloc_fail_on;
#undef UPB_THREAD_LOCAL
#endif
UPB_NODISCARD UPB_FORCEINLINE bool upb_AllocationCount_IncrementAndCheck(void) {
#ifdef UPB_ALLOCATION_COUNT
bool ok = (upb_arena_alloc_count != upb_arena_alloc_fail_on);
upb_arena_alloc_count++;
return ok;
#else
return true;
#endif
}
#ifdef __cplusplus
} /* extern "C" */
#endif
#include "upb/port/undef.inc"
#endif // GOOGLE_UPB_UPB_MEM_INTERNAL_ALLOC_H__

11
upb/mem/internal/arena.h
View file

@ -12,6 +12,7 @@
#include <stdint.h>
#include <string.h>
#include "upb/mem/internal/alloc.h"
#include "upb/port/sanitizers.h"
// Must be last.
@ -70,7 +71,8 @@ UPB_INLINE bool UPB_PRIVATE(_upb_Arena_IsAligned)(const void* ptr) {
return (uintptr_t)ptr % UPB_MALLOC_ALIGN == 0;
}
UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size) {
UPB_API_INLINE void* _upb_Arena_Malloc_Unchecked(struct upb_Arena* a,
size_t size) {
UPB_PRIVATE(upb_Xsan_AccessReadWrite)(UPB_XSAN(a));
size_t span = UPB_PRIVATE(_upb_Arena_AllocSpan)(size);
@ -89,6 +91,13 @@ UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size) {
return UPB_PRIVATE(upb_Xsan_NewUnpoisonedRegion)(UPB_XSAN(a), ret, size);
}
UPB_API_INLINE void* upb_Arena_Malloc(struct upb_Arena* a, size_t size) {
if (!upb_AllocationCount_IncrementAndCheck()) {
return NULL;
}
return _upb_Arena_Malloc_Unchecked(a, size);
}
UPB_API_INLINE void upb_Arena_ShrinkLast(struct upb_Arena* a, void* ptr,
size_t oldsize, size_t size) {
UPB_ASSERT(ptr);

4
upb/wire/internal/back_alloc.c
View file

@ -12,6 +12,7 @@
#include <string.h>
#include "upb/base/internal/log2.h"
#include "upb/mem/internal/alloc.h"
#include "upb/mem/internal/arena.h"
// Must be last.
@ -92,6 +93,9 @@ static char* upb_BackAlloc_Realloc(upb_BackAlloc* a, char* ptr, size_t n) {
}
char* upb_BackAlloc_Grow(upb_BackAlloc* a, char* ptr, size_t n) {
if (!upb_AllocationCount_IncrementAndCheck()) {
return NULL;
}
if (a->limit == a->buf) {
// First allocation: try to steal a block.
size_t size = n;