blob: 316b1f6603e0528de4b0aee49766beb01cd912e1 [file] [log] [blame]
// Protocol Buffers - Google's data interchange format
// Copyright 2023 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
#include "python/protobuf.h"
#include "python/descriptor.h"
#include "python/descriptor_containers.h"
#include "python/descriptor_pool.h"
#include "python/extension_dict.h"
#include "python/map.h"
#include "python/message.h"
#include "python/repeated.h"
#include "python/unknown_fields.h"
static upb_Arena* PyUpb_NewArena(void);
static void PyUpb_ModuleDealloc(void* module) {
PyUpb_ModuleState* s = PyModule_GetState(module);
PyUpb_WeakMap_Free(s->obj_cache);
if (s->c_descriptor_symtab) {
upb_DefPool_Free(s->c_descriptor_symtab);
}
}
PyObject* PyUpb_SetAllowOversizeProtos(PyObject* m, PyObject* arg) {
if (!arg || !PyBool_Check(arg)) {
PyErr_SetString(PyExc_TypeError,
"Argument to SetAllowOversizeProtos must be boolean");
return NULL;
}
PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
state->allow_oversize_protos = PyObject_IsTrue(arg);
Py_INCREF(arg);
return arg;
}
static PyMethodDef PyUpb_ModuleMethods[] = {
{"SetAllowOversizeProtos", PyUpb_SetAllowOversizeProtos, METH_O,
"Enable/disable oversize proto parsing."},
{NULL, NULL}};
static struct PyModuleDef module_def = {PyModuleDef_HEAD_INIT,
PYUPB_MODULE_NAME,
"Protobuf Module",
sizeof(PyUpb_ModuleState),
PyUpb_ModuleMethods, // m_methods
NULL, // m_slots
NULL, // m_traverse
NULL, // m_clear
PyUpb_ModuleDealloc};
// -----------------------------------------------------------------------------
// ModuleState
// -----------------------------------------------------------------------------
PyUpb_ModuleState* PyUpb_ModuleState_MaybeGet(void) {
PyObject* module = PyState_FindModule(&module_def);
return module ? PyModule_GetState(module) : NULL;
}
PyUpb_ModuleState* PyUpb_ModuleState_GetFromModule(PyObject* module) {
PyUpb_ModuleState* state = PyModule_GetState(module);
assert(state);
assert(PyModule_GetDef(module) == &module_def);
return state;
}
PyUpb_ModuleState* PyUpb_ModuleState_Get(void) {
PyObject* module = PyState_FindModule(&module_def);
assert(module);
return PyUpb_ModuleState_GetFromModule(module);
}
PyObject* PyUpb_GetWktBases(PyUpb_ModuleState* state) {
if (!state->wkt_bases) {
PyObject* wkt_module = PyImport_ImportModule(PYUPB_PROTOBUF_INTERNAL_PACKAGE
".well_known_types");
if (wkt_module == NULL) {
return false;
}
state->wkt_bases = PyObject_GetAttrString(wkt_module, "WKTBASES");
PyObject* m = PyState_FindModule(&module_def);
// Reparent ownership to m.
PyModule_AddObject(m, "__internal_wktbases", state->wkt_bases);
Py_DECREF(wkt_module);
}
return state->wkt_bases;
}
// -----------------------------------------------------------------------------
// WeakMap
// -----------------------------------------------------------------------------
struct PyUpb_WeakMap {
upb_inttable table;
upb_Arena* arena;
};
PyUpb_WeakMap* PyUpb_WeakMap_New(void) {
upb_Arena* arena = PyUpb_NewArena();
PyUpb_WeakMap* map = upb_Arena_Malloc(arena, sizeof(*map));
map->arena = arena;
upb_inttable_init(&map->table, map->arena);
return map;
}
void PyUpb_WeakMap_Free(PyUpb_WeakMap* map) { upb_Arena_Free(map->arena); }
// To give better entropy in the table key, we shift away low bits that are
// always zero.
static const int PyUpb_PtrShift = (sizeof(void*) == 4) ? 2 : 3;
uintptr_t PyUpb_WeakMap_GetKey(const void* key) {
uintptr_t n = (uintptr_t)key;
assert((n & ((1 << PyUpb_PtrShift) - 1)) == 0);
return n >> PyUpb_PtrShift;
}
void PyUpb_WeakMap_Add(PyUpb_WeakMap* map, const void* key, PyObject* py_obj) {
upb_inttable_insert(&map->table, PyUpb_WeakMap_GetKey(key),
upb_value_ptr(py_obj), map->arena);
}
void PyUpb_WeakMap_Delete(PyUpb_WeakMap* map, const void* key) {
upb_value val;
bool removed =
upb_inttable_remove(&map->table, PyUpb_WeakMap_GetKey(key), &val);
(void)removed;
assert(removed);
}
void PyUpb_WeakMap_TryDelete(PyUpb_WeakMap* map, const void* key) {
upb_inttable_remove(&map->table, PyUpb_WeakMap_GetKey(key), NULL);
}
PyObject* PyUpb_WeakMap_Get(PyUpb_WeakMap* map, const void* key) {
upb_value val;
if (upb_inttable_lookup(&map->table, PyUpb_WeakMap_GetKey(key), &val)) {
PyObject* ret = upb_value_getptr(val);
Py_INCREF(ret);
return ret;
} else {
return NULL;
}
}
bool PyUpb_WeakMap_Next(PyUpb_WeakMap* map, const void** key, PyObject** obj,
intptr_t* iter) {
uintptr_t u_key;
upb_value val;
if (!upb_inttable_next(&map->table, &u_key, &val, iter)) return false;
*key = (void*)(u_key << PyUpb_PtrShift);
*obj = upb_value_getptr(val);
return true;
}
void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap* map, intptr_t* iter) {
upb_inttable_removeiter(&map->table, iter);
}
// -----------------------------------------------------------------------------
// ObjCache
// -----------------------------------------------------------------------------
PyUpb_WeakMap* PyUpb_ObjCache_Instance(void) {
PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
return state->obj_cache;
}
void PyUpb_ObjCache_Add(const void* key, PyObject* py_obj) {
PyUpb_WeakMap_Add(PyUpb_ObjCache_Instance(), key, py_obj);
}
void PyUpb_ObjCache_Delete(const void* key) {
PyUpb_ModuleState* state = PyUpb_ModuleState_MaybeGet();
if (!state) {
// During the shutdown sequence, our object's Dealloc() methods can be
// called *after* our module Dealloc() method has been called. At that
// point our state will be NULL and there is nothing to delete out of the
// map.
return;
}
PyUpb_WeakMap_Delete(state->obj_cache, key);
}
PyObject* PyUpb_ObjCache_Get(const void* key) {
return PyUpb_WeakMap_Get(PyUpb_ObjCache_Instance(), key);
}
// -----------------------------------------------------------------------------
// Arena
// -----------------------------------------------------------------------------
typedef struct {
PyObject_HEAD;
upb_Arena* arena;
} PyUpb_Arena;
// begin:google_only
// static upb_alloc* global_alloc = &upb_alloc_global;
// end:google_only
// begin:github_only
#ifdef __GLIBC__
#include <malloc.h> // malloc_trim()
#endif
// A special allocator that calls malloc_trim() periodically to release
// memory to the OS. Without this call, we appear to leak memory, at least
// as measured in RSS.
//
// We opt not to use this instead of PyMalloc (which would also solve the
// problem) because the latter requires the GIL to be held. This would make
// our messages unsafe to share with other languages that could free at
// unpredictable
// times.
static void* upb_trim_allocfunc(upb_alloc* alloc, void* ptr, size_t oldsize,
size_t size) {
(void)alloc;
(void)oldsize;
if (size == 0) {
free(ptr);
#ifdef __GLIBC__
static int count = 0;
if (++count == 10000) {
malloc_trim(0);
count = 0;
}
#endif
return NULL;
} else {
return realloc(ptr, size);
}
}
static upb_alloc trim_alloc = {&upb_trim_allocfunc};
static const upb_alloc* global_alloc = &trim_alloc;
// end:github_only
static upb_Arena* PyUpb_NewArena(void) {
return upb_Arena_Init(NULL, 0, global_alloc);
}
PyObject* PyUpb_Arena_New(void) {
PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
PyUpb_Arena* arena = (void*)PyType_GenericAlloc(state->arena_type, 0);
arena->arena = PyUpb_NewArena();
return &arena->ob_base;
}
static void PyUpb_Arena_Dealloc(PyObject* self) {
upb_Arena_Free(PyUpb_Arena_Get(self));
PyUpb_Dealloc(self);
}
upb_Arena* PyUpb_Arena_Get(PyObject* arena) {
return ((PyUpb_Arena*)arena)->arena;
}
static PyType_Slot PyUpb_Arena_Slots[] = {
{Py_tp_dealloc, PyUpb_Arena_Dealloc},
{0, NULL},
};
static PyType_Spec PyUpb_Arena_Spec = {
PYUPB_MODULE_NAME ".Arena",
sizeof(PyUpb_Arena),
0, // itemsize
Py_TPFLAGS_DEFAULT,
PyUpb_Arena_Slots,
};
static bool PyUpb_InitArena(PyObject* m) {
PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
state->arena_type = PyUpb_AddClass(m, &PyUpb_Arena_Spec);
return state->arena_type;
}
// -----------------------------------------------------------------------------
// Utilities
// -----------------------------------------------------------------------------
PyTypeObject* AddObject(PyObject* m, const char* name, PyType_Spec* spec) {
PyObject* type = PyType_FromSpec(spec);
return type && PyModule_AddObject(m, name, type) == 0 ? (PyTypeObject*)type
: NULL;
}
static const char* PyUpb_GetClassName(PyType_Spec* spec) {
// spec->name contains a fully-qualified name, like:
// google.protobuf.pyext._message.FooBar
//
// Find the rightmost '.' to get "FooBar".
const char* name = strrchr(spec->name, '.');
assert(name);
return name + 1;
}
PyTypeObject* PyUpb_AddClass(PyObject* m, PyType_Spec* spec) {
PyObject* type = PyType_FromSpec(spec);
const char* name = PyUpb_GetClassName(spec);
if (PyModule_AddObject(m, name, type) < 0) {
Py_XDECREF(type);
return NULL;
}
return (PyTypeObject*)type;
}
PyTypeObject* PyUpb_AddClassWithBases(PyObject* m, PyType_Spec* spec,
PyObject* bases) {
PyObject* type = PyType_FromSpecWithBases(spec, bases);
const char* name = PyUpb_GetClassName(spec);
if (PyModule_AddObject(m, name, type) < 0) {
Py_XDECREF(type);
return NULL;
}
return (PyTypeObject*)type;
}
PyTypeObject* PyUpb_AddClassWithRegister(PyObject* m, PyType_Spec* spec,
PyObject* virtual_base,
const char** methods) {
PyObject* type = PyType_FromSpec(spec);
PyObject* ret1 = PyObject_CallMethod(virtual_base, "register", "O", type);
if (!ret1) {
Py_XDECREF(type);
return NULL;
}
for (size_t i = 0; methods[i] != NULL; i++) {
PyObject* method = PyObject_GetAttrString(virtual_base, methods[i]);
if (!method) {
Py_XDECREF(type);
return NULL;
}
int ret2 = PyObject_SetAttrString(type, methods[i], method);
if (ret2 < 0) {
Py_XDECREF(type);
return NULL;
}
}
return (PyTypeObject*)type;
}
const char* PyUpb_GetStrData(PyObject* obj) {
if (PyUnicode_Check(obj)) {
return PyUnicode_AsUTF8AndSize(obj, NULL);
} else if (PyBytes_Check(obj)) {
return PyBytes_AsString(obj);
} else {
return NULL;
}
}
const char* PyUpb_VerifyStrData(PyObject* obj) {
const char* ret = PyUpb_GetStrData(obj);
if (ret) return ret;
PyErr_Format(PyExc_TypeError, "Expected string: %S", obj);
return NULL;
}
PyObject* PyUpb_Forbidden_New(PyObject* cls, PyObject* args, PyObject* kwds) {
PyObject* name = PyObject_GetAttrString(cls, "__name__");
PyErr_Format(PyExc_RuntimeError,
"Objects of type %U may not be created directly.", name);
Py_XDECREF(name);
return NULL;
}
bool PyUpb_IndexToRange(PyObject* index, Py_ssize_t size, Py_ssize_t* i,
Py_ssize_t* count, Py_ssize_t* step) {
assert(i && count && step);
if (PySlice_Check(index)) {
Py_ssize_t start, stop;
if (PySlice_Unpack(index, &start, &stop, step) < 0) return false;
*count = PySlice_AdjustIndices(size, &start, &stop, *step);
*i = start;
} else {
*i = PyNumber_AsSsize_t(index, PyExc_IndexError);
if (*i == -1 && PyErr_Occurred()) {
PyErr_SetString(PyExc_TypeError, "list indices must be integers");
return false;
}
if (*i < 0) *i += size;
*step = 0;
*count = 1;
if (*i < 0 || size <= *i) {
PyErr_Format(PyExc_IndexError, "list index out of range");
return false;
}
}
return true;
}
// -----------------------------------------------------------------------------
// Module Entry Point
// -----------------------------------------------------------------------------
__attribute__((visibility("default"))) PyMODINIT_FUNC PyInit__message(void) {
PyObject* m = PyModule_Create(&module_def);
if (!m) return NULL;
PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
state->allow_oversize_protos = false;
state->wkt_bases = NULL;
state->obj_cache = PyUpb_WeakMap_New();
state->c_descriptor_symtab = NULL;
if (!PyUpb_InitDescriptorContainers(m) || !PyUpb_InitDescriptorPool(m) ||
!PyUpb_InitDescriptor(m) || !PyUpb_InitArena(m) ||
!PyUpb_InitExtensionDict(m) || !PyUpb_Map_Init(m) ||
!PyUpb_InitMessage(m) || !PyUpb_Repeated_Init(m) ||
!PyUpb_UnknownFields_Init(m)) {
Py_DECREF(m);
return NULL;
}
// Temporary: an cookie we can use in the tests to ensure we are testing upb
// and not another protobuf library on the system.
PyModule_AddIntConstant(m, "_IS_UPB", 1);
return m;
}