From ef41c82b580f655e0a7fa6fd2ee7592956b333f7 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 10:31:44 +0900 Subject: [PATCH 01/11] Add allocation budget to parsing --- packages/protobuf-py-ext/src/budget.rs | 61 +++++ packages/protobuf-py-ext/src/json_parse.rs | 220 ++++++++++++++---- packages/protobuf-py-ext/src/lib.rs | 1 + packages/protobuf-py-ext/src/marshaler.rs | 20 +- packages/protobuf-py-ext/src/nativemessage.rs | 64 ++++- packages/protobuf-py-ext/src/parser.rs | 128 ++++++++-- packages/protobuf-py-ext/src/wkt_registry.rs | 100 +++++--- src/protobuf/_budget.py | 107 +++++++++ src/protobuf/_from_binary.py | 118 ++++++++-- src/protobuf/_from_json.py | 173 +++++++++++--- src/protobuf/_message.py | 69 +++++- src/protobuf/_wkt_registry.py | 28 ++- tests/test_from_binary.py | 41 ++++ tests/test_json.py | 45 ++++ 14 files changed, 1004 insertions(+), 171 deletions(-) create mode 100644 packages/protobuf-py-ext/src/budget.rs create mode 100644 src/protobuf/_budget.py diff --git a/packages/protobuf-py-ext/src/budget.rs b/packages/protobuf-py-ext/src/budget.rs new file mode 100644 index 0000000..d15d5d8 --- /dev/null +++ b/packages/protobuf-py-ext/src/budget.rs @@ -0,0 +1,61 @@ +use pyo3::{PyResult, exceptions::PyValueError}; + +// Approximate sizes of CPython heap allocations, measured with +// `sys.getsizeof` on 64-bit CPython 3.14. The budget guards against +// unbounded allocation from malicious payloads rather than providing exact +// accounting, so small inaccuracies across versions and builds are fine. + +/// GC header allocated in front of every GC-tracked object. `tp_basicsize` +/// does not include it; `sys.getsizeof` does. +pub(crate) const GC_HEAD_SIZE: usize = 16; +/// A float object. +pub(crate) const FLOAT_SIZE: usize = 24; +/// A 64-bit int object. Smaller ints are slightly smaller. +pub(crate) const INT_SIZE: usize = 36; +/// Header of a compact ASCII str. The UTF-8 byte length is charged on top +/// as an approximation of the payload; non-ASCII headers are slightly larger. +pub(crate) const STR_OVERHEAD: usize = 41; +/// Header of a bytes object. +pub(crate) const BYTES_OVERHEAD: usize = 33; +/// An empty list, as created for repeated field defaults. +pub(crate) const EMPTY_LIST_SIZE: usize = 56; +/// An empty dict, as created for map field defaults. +pub(crate) const EMPTY_DICT_SIZE: usize = 64; +/// One appended list element: an 8-byte pointer slot (amortized growth +/// measures ~8.8 bytes per item). +pub(crate) const LIST_SLOT_SIZE: usize = 8; +/// One inserted dict entry: hash + key + value words plus index table and +/// growth slack (amortized growth measures ~37 bytes per item). +pub(crate) const DICT_ENTRY_SIZE: usize = 40; +/// A `Oneof` wrapper object: `PyObject` header plus two object pointers. +pub(crate) const ONEOF_SIZE: usize = 32; + +/// Tracks the approximate bytes of Python objects allocated while parsing a +/// message, raising an error once a configured limit is exceeded. Charges are +/// made before the corresponding allocation so the limit bounds the actual +/// peak, not just the size observed after the fact. +pub(crate) struct Budget { + current: usize, + max: usize, +} + +impl Budget { + pub(crate) fn new(limit: Option) -> Self { + Self { + current: 0, + max: limit.unwrap_or(usize::MAX), + } + } + + pub(crate) fn charge(&mut self, amount: usize) -> PyResult<()> { + self.current = self.current.saturating_add(amount); + if self.current > self.max { + Err(PyValueError::new_err(format!( + "allocation budget exceeded: needed {} bytes, limit is {}", + self.current, self.max + ))) + } else { + Ok(()) + } + } +} diff --git a/packages/protobuf-py-ext/src/json_parse.rs b/packages/protobuf-py-ext/src/json_parse.rs index 81a900f..b2e6e83 100644 --- a/packages/protobuf-py-ext/src/json_parse.rs +++ b/packages/protobuf-py-ext/src/json_parse.rs @@ -17,6 +17,7 @@ use pyo3::{ }; use crate::{ + budget::{self, Budget}, descriptor::{DescEnum, DescFieldValue, ScalarType}, json_source::{JiterSource, JsonKind, JsonSource, PyTreeSource}, marshaler::MessageMarshaler, @@ -53,9 +54,10 @@ pub(crate) fn merge_from_json<'py>( message: &Bound<'py, NativeMessage>, data: &[u8], opts: &FromJsonOpts, + alloc_budget: &mut Budget, ) -> PyResult<()> { let mut src = JiterSource::new(py, data); - read_message(marshaler, message, &mut src, opts, 0)?; + read_message(marshaler, message, &mut src, opts, 0, alloc_budget)?; src.finish() } @@ -66,9 +68,10 @@ pub(crate) fn read_message_from_tree<'py>( message: &Bound<'py, NativeMessage>, tree: Bound<'py, PyAny>, opts: &FromJsonOpts, + alloc_budget: &mut Budget, ) -> PyResult<()> { let mut src = PyTreeSource::new(py, tree); - read_message(marshaler, message, &mut src, opts, 0) + read_message(marshaler, message, &mut src, opts, 0, alloc_budget) } /// Reads a message value, dispatching on its well-known-type kind. Ordinary @@ -80,6 +83,7 @@ pub(crate) fn read_message<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { if depth > DEPTH_LIMIT { return Err(pyo3::exceptions::PyRecursionError::new_err(format!( @@ -87,8 +91,8 @@ pub(crate) fn read_message<'py, R: JsonSource<'py>>( ))); } match &marshaler.wkt { - Some(wkt) => wkt.read_json(marshaler, message, src, opts, depth), - None => read_generic_object(marshaler, message, src, opts, depth), + Some(wkt) => wkt.read_json(marshaler, message, src, opts, depth, alloc_budget), + None => read_generic_object(marshaler, message, src, opts, depth, alloc_budget), } } @@ -98,10 +102,11 @@ fn read_generic_object<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? != JsonKind::Object { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; let qualname = marshaler.python_type.bind(py).qualname()?; return Err(PyTypeError::new_err(format!( "cannot decode {qualname} from JSON: {}", @@ -132,9 +137,9 @@ fn read_generic_object<'py, R: JsonSource<'py>>( if merges_on_duplicate(parser) && !seen.insert(number) { reset_duplicate_field(py, parser, message)?; } - read_field(marshaler, parser, message, src, opts, depth)?; + read_field(marshaler, parser, message, src, opts, depth, alloc_budget)?; } else { - handle_unknown_key(marshaler, message, key, src, opts, depth)?; + handle_unknown_key(marshaler, message, key, src, opts, depth, alloc_budget)?; } Ok(()) })?; @@ -178,6 +183,7 @@ fn read_field<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { match &parser.type_ { ParserFieldType::Singular { @@ -192,8 +198,11 @@ fn read_field<'py, R: JsonSource<'py>>( depth, oneof_attr.as_ref(), *requires_presence, + alloc_budget, ), - ParserFieldType::List { .. } => read_list(marshaler, parser, message, src, opts, depth), + ParserFieldType::List { .. } => { + read_list(marshaler, parser, message, src, opts, depth, alloc_budget) + } ParserFieldType::Map { key_type, value_parser, @@ -207,6 +216,7 @@ fn read_field<'py, R: JsonSource<'py>>( depth, *key_type, value_parser, + alloc_budget, ), } } @@ -221,6 +231,7 @@ fn read_singular<'py, R: JsonSource<'py>>( depth: usize, oneof_attr: Option<&crate::attribute_access::AttributeAccess>, requires_presence: bool, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); match &parser.value { @@ -238,7 +249,10 @@ fn read_singular<'py, R: JsonSource<'py>>( } let name = parser.name.bind(py); let ctx = FieldContext::Field { marshaler, name }; - let value = read_scalar(&ctx, src, *scalar)?; + let value = read_scalar(&ctx, src, *scalar, alloc_budget)?; + if oneof_attr.is_some() { + alloc_budget.charge(budget::ONEOF_SIZE)?; + } parser.assign_singular(py, message, &value, oneof_attr, requires_presence) } FieldParserValue::Enum(enum_) => { @@ -253,7 +267,10 @@ fn read_singular<'py, R: JsonSource<'py>>( )?; return Ok(()); } - if let Some(value) = read_enum(enum_, src, opts)? { + if let Some(value) = read_enum(enum_, src, opts, alloc_budget)? { + if oneof_attr.is_some() { + alloc_budget.charge(budget::ONEOF_SIZE)?; + } parser.assign_singular(py, message, &value, oneof_attr, requires_presence)?; } Ok(()) @@ -271,9 +288,15 @@ fn read_singular<'py, R: JsonSource<'py>>( let existing = parser.get_field_value(py, message)?; let target = match existing { Some(value) if !value.is_none() => value.cast_into::()?, - _ => inner.new_empty_message(py, msg_desc.get_python_type(py))?, + _ => { + alloc_budget.charge(inner.base_alloc_size)?; + inner.new_empty_message(py, msg_desc.get_python_type(py))? + } }; - read_message(inner, &target, src, opts, depth + 1)?; + read_message(inner, &target, src, opts, depth + 1, alloc_budget)?; + if oneof_attr.is_some() { + alloc_budget.charge(budget::ONEOF_SIZE)?; + } parser.assign_singular(py, message, target.as_any(), oneof_attr, requires_presence) } } @@ -312,6 +335,7 @@ fn read_list<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let name = parser.name.bind(py); @@ -322,7 +346,10 @@ fn read_list<'py, R: JsonSource<'py>>( } if src.peek()? != JsonKind::Array { return Err(ctx.error( - &format!("expected list got {}", read_json_value(src)?.get_type()), + &format!( + "expected list got {}", + read_json_value(src, alloc_budget)?.get_type() + ), Exc::Type, )); } @@ -331,7 +358,10 @@ fn read_list<'py, R: JsonSource<'py>>( .get(py, message.as_any())? .cast_into::()?; src.for_each_array_item(|src| { - if let Some(value) = read_container_item(&ctx, &parser.value, src, opts, depth, false)? { + if let Some(value) = + read_container_item(&ctx, &parser.value, src, opts, depth, false, alloc_budget)? + { + alloc_budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } Ok(()) @@ -349,6 +379,7 @@ fn read_map<'py, R: JsonSource<'py>>( depth: usize, key_type: ScalarType, value_parser: &FieldParser, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let name = parser.name.bind(py); @@ -359,7 +390,10 @@ fn read_map<'py, R: JsonSource<'py>>( } if src.peek()? != JsonKind::Object { return Err(ctx.error( - &format!("expected dict got {}", read_json_value(src)?.get_type()), + &format!( + "expected dict got {}", + read_json_value(src, alloc_budget)?.get_type() + ), Exc::Type, )); } @@ -368,9 +402,17 @@ fn read_map<'py, R: JsonSource<'py>>( .get(py, message.as_any())? .cast_into::()?; src.for_each_object_key(|key, src| { - let map_key = read_map_key(py, &ctx, key_type, key)?; - if let Some(value) = read_container_item(&ctx, &value_parser.value, src, opts, depth, true)? - { + let map_key = read_map_key(py, &ctx, key_type, key, alloc_budget)?; + if let Some(value) = read_container_item( + &ctx, + &value_parser.value, + src, + opts, + depth, + true, + alloc_budget, + )? { + alloc_budget.charge(budget::DICT_ENTRY_SIZE)?; dict.set_item(map_key, value)?; } Ok(()) @@ -387,11 +429,14 @@ fn read_container_item<'py, R: JsonSource<'py>>( opts: &FromJsonOpts, depth: usize, is_map: bool, + alloc_budget: &mut Budget, ) -> PyResult>> { let py = src.py(); let is_null = src.peek()? == JsonKind::Null; match element { - FieldParserValue::Scalar(scalar) if !is_null => Ok(Some(read_scalar(ctx, src, *scalar)?)), + FieldParserValue::Scalar(scalar) if !is_null => { + Ok(Some(read_scalar(ctx, src, *scalar, alloc_budget)?)) + } FieldParserValue::Message { message: msg_desc, .. } => { @@ -401,12 +446,13 @@ fn read_container_item<'py, R: JsonSource<'py>>( src.next_null()?; return Err(container_null_error(ctx, is_map)); } + alloc_budget.charge(inner.base_alloc_size)?; let target = inner.new_empty_message(py, msg_desc.get_python_type(py))?; - read_message(inner, &target, src, opts, depth + 1)?; + read_message(inner, &target, src, opts, depth + 1, alloc_budget)?; Ok(Some(target.into_any())) } FieldParserValue::Enum(enum_) if !is_null || enum_.is_null_value => { - read_enum(enum_, src, opts) + read_enum(enum_, src, opts, alloc_budget) } _ => { // Resetting null for a list item / map value: error. @@ -426,6 +472,7 @@ fn read_map_key<'py>( ctx: &FieldContext<'_, 'py>, key_type: ScalarType, raw_key: &str, + alloc_budget: &mut Budget, ) -> PyResult> { match key_type { ScalarType::Bool => match raw_key { @@ -436,22 +483,30 @@ fn read_map_key<'py>( Exc::Value, )), }, - ScalarType::String => Ok(PyString::new(py, raw_key).into_any()), - _ => parse_int_string(py, ctx, raw_key, key_type), + ScalarType::String => { + alloc_budget.charge(budget::STR_OVERHEAD + raw_key.len())?; + Ok(PyString::new(py, raw_key).into_any()) + } + _ => { + alloc_budget.charge(budget::INT_SIZE)?; + parse_int_string(py, ctx, raw_key, key_type) + } } } -/// Reads a scalar value. +/// Reads a scalar value, charging the allocation budget for the resulting +/// Python object. pub(crate) fn read_scalar<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, scalar: ScalarType, + alloc_budget: &mut Budget, ) -> PyResult> { let py = src.py(); match scalar { ScalarType::Bool => { if src.peek()? != JsonKind::Bool { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; return Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -460,25 +515,39 @@ pub(crate) fn read_scalar<'py, R: JsonSource<'py>>( Ok(PyBool::new(py, src.next_bool()?).to_owned().into_any()) } ScalarType::Float => { - let value = parse_float(ctx, src)?; + let value = parse_float(ctx, src, alloc_budget)?; if value.is_finite() && !(FLOAT32_MIN..=FLOAT32_MAX).contains(&value) { return Err(ctx.error(&format!("float value out of range: {value}"), Exc::Overflow)); } + alloc_budget.charge(budget::FLOAT_SIZE)?; Ok(PyFloat::new(py, value).into_any()) } - ScalarType::Double => Ok(PyFloat::new(py, parse_float(ctx, src)?).into_any()), - ScalarType::String => Ok(read_string(ctx, src)?.into_any()), - ScalarType::Bytes => read_bytes(ctx, src), - _ => read_int(ctx, src, scalar), + ScalarType::Double => { + let value = parse_float(ctx, src, alloc_budget)?; + alloc_budget.charge(budget::FLOAT_SIZE)?; + Ok(PyFloat::new(py, value).into_any()) + } + ScalarType::String => { + let value = read_string(ctx, src, alloc_budget)?; + // Character count approximates the payload size. + alloc_budget.charge(budget::STR_OVERHEAD + value.len()?)?; + Ok(value.into_any()) + } + ScalarType::Bytes => read_bytes(ctx, src, alloc_budget), + _ => { + alloc_budget.charge(budget::INT_SIZE)?; + read_int(ctx, src, scalar, alloc_budget) + } } } fn read_string<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, + alloc_budget: &mut Budget, ) -> PyResult> { if src.peek()? != JsonKind::String { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; return Err(ctx.error( &format!("expected string got: {}", value.get_type()), Exc::Type, @@ -490,10 +559,11 @@ fn read_string<'py, R: JsonSource<'py>>( fn read_bytes<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, + alloc_budget: &mut Budget, ) -> PyResult> { let py = src.py(); if src.peek()? != JsonKind::String { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; return Err(ctx.error( &format!("expected base64-encoded string got: {}", value.get_type()), Exc::Type, @@ -511,6 +581,7 @@ fn read_bytes<'py, R: JsonSource<'py>>( .map_err(|_| ctx.error("invalid base64 data", Exc::Value))?; Ok(decoded) })?; + alloc_budget.charge(budget::BYTES_OVERHEAD + decoded.len())?; Ok(PyBytes::new(py, &decoded).into_any()) } @@ -531,7 +602,11 @@ fn base64_url_safe() -> base64::engine::GeneralPurpose { } /// Parses a float/double. -fn parse_float<'py, R: JsonSource<'py>>(ctx: &FieldContext<'_, 'py>, src: &mut R) -> PyResult { +fn parse_float<'py, R: JsonSource<'py>>( + ctx: &FieldContext<'_, 'py>, + src: &mut R, + alloc_budget: &mut Budget, +) -> PyResult { match src.peek()? { JsonKind::Number => { let value = src.next_float()?; @@ -560,7 +635,7 @@ fn parse_float<'py, R: JsonSource<'py>>(ctx: &FieldContext<'_, 'py>, src: &mut R } }), _ => { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -574,6 +649,7 @@ fn read_int<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, int_type: ScalarType, + alloc_budget: &mut Budget, ) -> PyResult> { let py = src.py(); let value = match src.peek()? { @@ -599,7 +675,7 @@ fn read_int<'py, R: JsonSource<'py>>( }); } _ => { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; return Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -676,6 +752,7 @@ fn read_enum<'py, R: JsonSource<'py>>( enum_desc: &DescEnum, src: &mut R, opts: &FromJsonOpts, + alloc_budget: &mut Budget, ) -> PyResult>> { let py = src.py(); match src.peek()? { @@ -694,6 +771,7 @@ fn read_enum<'py, R: JsonSource<'py>>( if opts.ignore_unknown_fields { return Ok(None); } + alloc_budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; return Ok(Some(enum_desc.py_type.bind(py).call1((number,))?)); }; if let Some(value) = enum_desc.values.get(&int_value) { @@ -701,7 +779,9 @@ fn read_enum<'py, R: JsonSource<'py>>( } else if opts.ignore_unknown_fields { Ok(None) } else { - // Open enum: succeeds; closed enum: raises via Python enum call. + // Open enum: succeeds and allocates a new int-subclass + // instance; closed enum: raises via Python enum call. + alloc_budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; Ok(Some(enum_desc.py_type.bind(py).call1((int_value,))?)) } } @@ -720,7 +800,7 @@ fn read_enum<'py, R: JsonSource<'py>>( } }), _ => { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; Err(decode_enum_error(py, enum_desc, &value)) } } @@ -745,6 +825,7 @@ fn handle_unknown_key<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if raw_key.starts_with('[') @@ -760,7 +841,15 @@ fn handle_unknown_key<'py, R: JsonSource<'py>>( let extendee_name = extendee.getattr(&marshaler.constants.type_name)?; let extendee_name = extendee_name.cast::()?.to_str()?; if extendee_name == &*marshaler.type_name { - read_extension(marshaler, message, &extension, src, opts, depth)?; + read_extension( + marshaler, + message, + &extension, + src, + opts, + depth, + alloc_budget, + )?; } else { src.skip()?; } @@ -785,6 +874,7 @@ fn read_extension<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let ext_type = extension.getattr(&marshaler.constants.type_)?; @@ -803,7 +893,7 @@ fn read_extension<'py, R: JsonSource<'py>>( src.next_null()?; target.del_item(&ext_type)?; } else { - let value = read_scalar(&ctx, src, *scalar_type)?; + let value = read_scalar(&ctx, src, *scalar_type, alloc_budget)?; target.set_item(&ext_type, value)?; } } @@ -811,7 +901,7 @@ fn read_extension<'py, R: JsonSource<'py>>( if src.peek()? == JsonKind::Null && !enum_.is_null_value { src.next_null()?; target.del_item(&ext_type)?; - } else if let Some(value) = read_enum(enum_, src, opts)? { + } else if let Some(value) = read_enum(enum_, src, opts, alloc_budget)? { target.set_item(&ext_type, value)?; } } @@ -824,8 +914,9 @@ fn read_extension<'py, R: JsonSource<'py>>( src.next_null()?; target.del_item(&ext_type)?; } else { + alloc_budget.charge(inner.base_alloc_size)?; let value = inner.new_empty_message(py, msg_desc.get_python_type(py))?; - read_message(inner, &value, src, opts, depth + 1)?; + read_message(inner, &value, src, opts, depth + 1, alloc_budget)?; target.set_item(&ext_type, value)?; } } @@ -837,16 +928,27 @@ fn read_extension<'py, R: JsonSource<'py>>( } if src.peek()? != JsonKind::Array { return Err(ctx.error( - &format!("expected list got {}", read_json_value(src)?.get_type()), + &format!( + "expected list got {}", + read_json_value(src, alloc_budget)?.get_type() + ), Exc::Type, )); } let element_value = FieldParserValue::from_desc_single(element); + alloc_budget.charge(budget::EMPTY_LIST_SIZE)?; let list = PyList::empty(py); src.for_each_array_item(|src| { - if let Some(value) = - read_container_item(&ctx, &element_value, src, opts, depth, false)? - { + if let Some(value) = read_container_item( + &ctx, + &element_value, + src, + opts, + depth, + false, + alloc_budget, + )? { + alloc_budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } Ok(()) @@ -861,8 +963,12 @@ fn read_extension<'py, R: JsonSource<'py>>( Ok(()) } -/// Materializes the next JSON value as a Python object. -pub(crate) fn read_json_value<'py, R: JsonSource<'py>>(src: &mut R) -> PyResult> { +/// Materializes the next JSON value as a Python object, charging the +/// allocation budget for the resulting tree. +pub(crate) fn read_json_value<'py, R: JsonSource<'py>>( + src: &mut R, + alloc_budget: &mut Budget, +) -> PyResult> { let py = src.py(); match src.peek()? { JsonKind::Null => { @@ -870,20 +976,32 @@ pub(crate) fn read_json_value<'py, R: JsonSource<'py>>(src: &mut R) -> PyResult< Ok(py.None().into_bound(py)) } JsonKind::Bool => Ok(PyBool::new(py, src.next_bool()?).to_owned().into_any()), - JsonKind::Number => src.next_number(), - JsonKind::String => Ok(src.next_py_str()?.into_any()), + JsonKind::Number => { + alloc_budget.charge(budget::INT_SIZE)?; + src.next_number() + } + JsonKind::String => { + let value = src.next_py_str()?; + alloc_budget.charge(budget::STR_OVERHEAD + value.len()?)?; + Ok(value.into_any()) + } JsonKind::Array => { + alloc_budget.charge(budget::EMPTY_LIST_SIZE)?; let list = PyList::empty(py); src.for_each_array_item(|src| { - list.append(read_json_value(src)?)?; + let value = read_json_value(src, alloc_budget)?; + alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + list.append(value)?; Ok(()) })?; Ok(list.into_any()) } JsonKind::Object => { + alloc_budget.charge(budget::EMPTY_DICT_SIZE)?; let dict = PyDict::new(py); src.for_each_object_key(|key, src| { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; + alloc_budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; dict.set_item(PyString::new(py, key), value)?; Ok(()) })?; diff --git a/packages/protobuf-py-ext/src/lib.rs b/packages/protobuf-py-ext/src/lib.rs index 1802da2..ef493f6 100644 --- a/packages/protobuf-py-ext/src/lib.rs +++ b/packages/protobuf-py-ext/src/lib.rs @@ -4,6 +4,7 @@ use pyo3::prelude::*; mod alloc; mod attribute_access; mod bitset; +mod budget; mod buffer; mod constants; mod descriptor; diff --git a/packages/protobuf-py-ext/src/marshaler.rs b/packages/protobuf-py-ext/src/marshaler.rs index e8e7287..aca2ced 100644 --- a/packages/protobuf-py-ext/src/marshaler.rs +++ b/packages/protobuf-py-ext/src/marshaler.rs @@ -7,12 +7,13 @@ use pyo3::{ pyclass, types::{ PyAnyMethods as _, PyBytes, PyDict, PyDictMethods as _, PyList, PyStringMethods as _, - PyType, + PyType, PyTypeMethods as _, }, }; use crate::{ attribute_access::AttributeAccess, + budget::{self, Budget}, constants::Constants, descriptor::{DescFieldValue, message_fields}, nativemessage::NativeMessage, @@ -74,6 +75,11 @@ pub(crate) struct MessageMarshalerInner { /// The maximum field number of the message. pub(crate) max_field_number: u32, + /// Approximate heap size of a freshly-initialized instance of this + /// message type: the fixed instance size (all fields are slots) plus the + /// empty containers created for repeated/map field defaults. + pub(crate) base_alloc_size: usize, + /// The Python type of the message. pub(crate) python_type: Py, @@ -198,6 +204,15 @@ impl MessageMarshaler { .push((member.attr.clone_ref(py), default.unbind())); } } + // tp_basicsize is the exact instance allocation size: message types + // only use slots, so instances never grow beyond it. The GC header is + // allocated in front of every instance on top of it. + // SAFETY - the type pointer of a live PyType is always valid. + let basic_size = unsafe { (*python_type.as_type_ptr()).tp_basicsize }; + let base_alloc_size = usize::try_from(basic_size).unwrap_or(0) + + budget::GC_HEAD_SIZE + + defaults.lists.len() * budget::EMPTY_LIST_SIZE + + defaults.dicts.len() * budget::EMPTY_DICT_SIZE; Ok(Self { inner: Arc::new(MessageMarshalerInner { parser, @@ -205,6 +220,7 @@ impl MessageMarshaler { members_by_name: members_by_name.unbind(), members, max_field_number, + base_alloc_size, python_type: python_type.clone().unbind(), type_name, wkt, @@ -222,6 +238,7 @@ impl MessageMarshaler { message: &Bound<'_, NativeMessage>, mut data: Bytes, ignore_unknown_fields: bool, + budget: &mut Budget, ) -> PyResult<()> { self.inner.parser.merge_from_binary( py, @@ -231,6 +248,7 @@ impl MessageMarshaler { ignore_unknown_fields, }, 0, + budget, ) } diff --git a/packages/protobuf-py-ext/src/nativemessage.rs b/packages/protobuf-py-ext/src/nativemessage.rs index dd361b3..df42674 100644 --- a/packages/protobuf-py-ext/src/nativemessage.rs +++ b/packages/protobuf-py-ext/src/nativemessage.rs @@ -13,6 +13,7 @@ use pyo3::{ use crate::{ attribute_access::generic_setattr, bitset::BitSet, + budget::Budget, buffer::Buffer, constants::Constants, json_parse::{FromJsonOpts, merge_from_json, read_message_from_tree}, @@ -100,20 +101,23 @@ impl NativeMessage { } #[classmethod] - #[pyo3(signature = (data, *, ignore_unknown_fields = false))] + #[pyo3(signature = (data, *, ignore_unknown_fields = false, allocation_limit = None))] fn from_binary<'py>( cls: &Bound<'py, PyType>, py: Python<'py>, data: Buffer, ignore_unknown_fields: bool, + allocation_limit: Option, ) -> PyResult> { let data = data.into_inner(); let constants = Constants::get(py)?; let marshaler_any = cls.getattr(&constants.ext_marshaler)?; let marshaler = marshaler_any.cast::()?.get(); + let mut alloc_budget = Budget::new(allocation_limit); + alloc_budget.charge(marshaler.base_alloc_size)?; let message = marshaler.new_empty_message(py, cls)?; let slf = message.cast::()?; - marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields)?; + marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut alloc_budget)?; Ok(message) } @@ -146,17 +150,20 @@ impl NativeMessage { } #[classmethod] - #[pyo3(signature = (json, *, ignore_unknown_fields = false, registry = None))] + #[pyo3(signature = (json, *, ignore_unknown_fields = false, registry = None, allocation_limit = None))] fn from_json<'py>( cls: &Bound<'py, PyType>, py: Python<'py>, json: &Bound<'py, PyAny>, ignore_unknown_fields: bool, registry: Option>, + allocation_limit: Option, ) -> PyResult> { let constants = Constants::get(py)?; let marshaler_any = cls.getattr(&constants.ext_marshaler)?; let marshaler = marshaler_any.cast::()?.get().clone(); + let mut alloc_budget = Budget::new(allocation_limit); + alloc_budget.charge(marshaler.base_alloc_size)?; let message = marshaler.new_empty_message(py, cls)?; parse_json_into( py, @@ -165,20 +172,32 @@ impl NativeMessage { json, ignore_unknown_fields, registry, + &mut alloc_budget, )?; Ok(message) } - #[pyo3(signature = (json, *, ignore_unknown_fields = false, registry = None))] + #[pyo3(signature = (json, *, ignore_unknown_fields = false, registry = None, allocation_limit = None))] fn _merge_from_json( slf: &Bound<'_, Self>, py: Python<'_>, json: &Bound<'_, PyAny>, ignore_unknown_fields: bool, registry: Option>, + allocation_limit: Option, ) -> PyResult<()> { let marshaler = NativeMessage::get_marshaler(slf)?; - parse_json_into(py, &marshaler, slf, json, ignore_unknown_fields, registry) + let mut alloc_budget = Budget::new(allocation_limit); + alloc_budget.charge(marshaler.base_alloc_size)?; + parse_json_into( + py, + &marshaler, + slf, + json, + ignore_unknown_fields, + registry, + &mut alloc_budget, + ) } #[pyo3(signature = (*, registry = None, always_emit_implicit = false, print_enums_as_ints = false, use_proto_field_name = false))] @@ -203,23 +222,26 @@ impl NativeMessage { } #[classmethod] - #[pyo3(signature = (data, *, ignore_unknown_fields = false, registry = None))] + #[pyo3(signature = (data, *, ignore_unknown_fields = false, registry = None, allocation_limit = None))] fn _from_json_value<'py>( cls: &Bound<'py, PyType>, py: Python<'py>, data: Bound<'py, PyAny>, ignore_unknown_fields: bool, registry: Option>, + allocation_limit: Option, ) -> PyResult> { let constants = Constants::get(py)?; let marshaler_any = cls.getattr(&constants.ext_marshaler)?; let marshaler = marshaler_any.cast::()?.get().clone(); + let mut alloc_budget = Budget::new(allocation_limit); + alloc_budget.charge(marshaler.base_alloc_size)?; let message = marshaler.new_empty_message(py, cls)?; let opts = FromJsonOpts { ignore_unknown_fields, registry, }; - read_message_from_tree(py, &marshaler, &message, data, &opts)?; + read_message_from_tree(py, &marshaler, &message, data, &opts, &mut alloc_budget)?; Ok(message) } @@ -264,14 +286,19 @@ impl NativeMessage { Ok(new) } + #[pyo3(signature = (data, ignore_unknown_fields, allocation_limit = None))] fn _merge_from_binary( slf: &Bound<'_, Self>, py: Python<'_>, data: Buffer, ignore_unknown_fields: bool, + allocation_limit: Option, ) -> PyResult<()> { let data = data.into_inner(); - NativeMessage::get_marshaler(slf)?.merge_from_binary(py, slf, data, ignore_unknown_fields) + let marshaler = NativeMessage::get_marshaler(slf)?; + let mut alloc_budget = Budget::new(allocation_limit); + alloc_budget.charge(marshaler.base_alloc_size)?; + marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut alloc_budget) } fn __deepcopy__<'py>( @@ -514,18 +541,33 @@ fn parse_json_into<'py>( json: &Bound<'py, PyAny>, ignore_unknown_fields: bool, registry: Option>, + alloc_budget: &mut Budget, ) -> PyResult<()> { let opts = FromJsonOpts { ignore_unknown_fields, registry, }; if let Ok(text) = json.cast::() { - merge_from_json(py, marshaler, message, text.to_str()?.as_bytes(), &opts) + merge_from_json( + py, + marshaler, + message, + text.to_str()?.as_bytes(), + &opts, + alloc_budget, + ) } else if let Ok(bytes) = json.cast::() { - merge_from_json(py, marshaler, message, bytes.as_bytes(), &opts) + merge_from_json( + py, + marshaler, + message, + bytes.as_bytes(), + &opts, + alloc_budget, + ) } else if let Ok(bytearray) = json.cast::() { let owned = bytearray.to_vec(); - merge_from_json(py, marshaler, message, &owned, &opts) + merge_from_json(py, marshaler, message, &owned, &opts, alloc_budget) } else { Err(PyTypeError::new_err(format!( "json must be str, bytes, or bytearray, got {}", diff --git a/packages/protobuf-py-ext/src/parser.rs b/packages/protobuf-py-ext/src/parser.rs index dde1d80..e48db6a 100644 --- a/packages/protobuf-py-ext/src/parser.rs +++ b/packages/protobuf-py-ext/src/parser.rs @@ -23,6 +23,7 @@ use pyo3::{ use crate::{ attribute_access::AttributeAccess, + budget::{self, Budget}, constants::Constants, descriptor::{ DescEnum, DescField, DescFieldValue, DescMessage, DescSingleValue, FieldPresence, @@ -250,6 +251,7 @@ impl FieldParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { match &self.type_ { ParserFieldType::Singular { @@ -265,6 +267,7 @@ impl FieldParser { depth, oneof_attr.as_ref(), *requires_presence, + alloc_budget, )?, ParserFieldType::List { unpacked_wire_type, @@ -279,6 +282,7 @@ impl FieldParser { depth, *unpacked_wire_type, *packable, + alloc_budget, )?, ParserFieldType::Map { key_type, @@ -296,6 +300,7 @@ impl FieldParser { value_parser, key_default_value, value_default_value, + alloc_budget, )?, } @@ -313,11 +318,23 @@ impl FieldParser { depth: usize, oneof_attr: Option<&AttributeAccess>, requires_presence: bool, + alloc_budget: &mut Budget, ) -> PyResult<()> { - let value = - self.read_single_value(py, tag, wire_type, Some(message), buffer, opts, depth)?; + let value = self.read_single_value( + py, + tag, + wire_type, + Some(message), + buffer, + opts, + depth, + alloc_budget, + )?; match value { SingleValue::Parsed(value) => { + if oneof_attr.is_some() { + alloc_budget.charge(budget::ONEOF_SIZE)?; + } self.assign_singular(py, message, &value, oneof_attr, requires_presence)?; } SingleValue::UnknownEnumValue(number) => { @@ -325,7 +342,7 @@ impl FieldParser { let mut field = BytesMut::new(); encode_varint(tag as u64, &mut field); encode_varint(number as u64, &mut field); - write_unknown_field(py, message, tag >> 3, &field)?; + write_unknown_field(py, message, tag >> 3, &field, alloc_budget)?; } } } @@ -365,6 +382,7 @@ impl FieldParser { depth: usize, unpacked_wire_type: WireType, packable: bool, + alloc_budget: &mut Budget, ) -> PyResult<()> { let python_list = self.attr.get(py, message)?; let list = python_list.cast::()?; @@ -381,9 +399,13 @@ impl FieldParser { &mut list_buffer, opts, depth, + alloc_budget, )?; match value { - SingleValue::Parsed(value) => list.append(value)?, + SingleValue::Parsed(value) => { + alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + list.append(value)?; + } SingleValue::UnknownEnumValue(number) => { if !opts.ignore_unknown_fields { // Always add unknown values as unpacked. @@ -391,21 +413,33 @@ impl FieldParser { let tag = (tag & !0b111) | (unpacked_wire_type as u32); encode_varint(tag as u64, &mut field); encode_varint(number as u64, &mut field); - write_unknown_field(py, message, tag >> 3, &field)?; + write_unknown_field(py, message, tag >> 3, &field, alloc_budget)?; } } } } } else { - let value = self.read_single_value(py, tag, wire_type, None, buffer, opts, depth)?; + let value = self.read_single_value( + py, + tag, + wire_type, + None, + buffer, + opts, + depth, + alloc_budget, + )?; match value { - SingleValue::Parsed(value) => list.append(value)?, + SingleValue::Parsed(value) => { + alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + list.append(value)?; + } SingleValue::UnknownEnumValue(number) => { if !opts.ignore_unknown_fields { let mut field = BytesMut::new(); encode_varint(tag as u64, &mut field); encode_varint(number as u64, &mut field); - write_unknown_field(py, message, tag >> 3, &field)?; + write_unknown_field(py, message, tag >> 3, &field, alloc_budget)?; } } } @@ -425,6 +459,7 @@ impl FieldParser { value_parser: &FieldParser, key_default_value: &Py, value_default_value: &Py, + alloc_budget: &mut Budget, ) -> PyResult<()> { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; @@ -446,9 +481,10 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, + alloc_budget, ); } - key = Some(read_scalar(py, key_type, &mut entry_buffer)?); + key = Some(read_scalar(py, key_type, &mut entry_buffer, alloc_budget)?); } 2 => { if value_parser.wire_type != wire_type { @@ -458,6 +494,7 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, + alloc_budget, ); } value = Some(value_parser.read_single_value( @@ -468,6 +505,7 @@ impl FieldParser { &mut entry_buffer, opts, depth, + alloc_budget, )?); } _ => { @@ -481,14 +519,23 @@ impl FieldParser { value } else if let FieldParserValue::Message { message, .. } = &value_parser.value { // For message values, the default is a new instance of the message, not None. + alloc_budget.charge(message.get_marshaler(py)?.base_alloc_size)?; SingleValue::Parsed(message.get_python_type(py).call0()?) } else { SingleValue::Parsed(value_default_value.bind(py).clone()) }; let SingleValue::Parsed(value) = value else { // Unknown value (i.e., unknown closed enum value), whole entry is an unknown field. - return Self::read_unknown_map_entry(py, message, entry_tag, &entry_checkpoint, opts); + return Self::read_unknown_map_entry( + py, + message, + entry_tag, + &entry_checkpoint, + opts, + alloc_budget, + ); }; + alloc_budget.charge(budget::DICT_ENTRY_SIZE)?; self.assign_map_entry(py, message, key, value)?; Ok(()) } @@ -513,13 +560,14 @@ impl FieldParser { tag: u32, entry_bytes: &[u8], opts: FromBinaryOpts, + alloc_budget: &mut Budget, ) -> PyResult<()> { if !opts.ignore_unknown_fields { let mut field_bytes = BytesMut::new(); encode_varint(tag as u64, &mut field_bytes); encode_varint(entry_bytes.len() as u64, &mut field_bytes); field_bytes.extend_from_slice(entry_bytes); - write_unknown_field(py, message, tag >> 3, &field_bytes)?; + write_unknown_field(py, message, tag >> 3, &field_bytes, alloc_budget)?; } Ok(()) } @@ -534,9 +582,12 @@ impl FieldParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult> { let value = match &self.value { - FieldParserValue::Scalar(scalar_type) => read_scalar(py, *scalar_type, buffer)?, + FieldParserValue::Scalar(scalar_type) => { + read_scalar(py, *scalar_type, buffer, alloc_budget)? + } FieldParserValue::Message { message: message_desc, .. @@ -564,6 +615,7 @@ impl FieldParser { { existing.cast_into::()? } else { + alloc_budget.charge(marshaler.base_alloc_size)?; marshaler.new_empty_message(py, parser.inner.python_type.bind(py))? }; parser.merge_from_binary( @@ -572,6 +624,7 @@ impl FieldParser { &mut message_buffer, opts, depth + 1, + alloc_budget, )?; message_instance.into_any() } @@ -581,6 +634,9 @@ impl FieldParser { if let Some(value) = value { value.bind(py).clone() } else if enum_.open { + // Unknown open enum values allocate a new int-subclass + // instance of the enum type. + alloc_budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; enum_.py_type.bind(py).call1((number,))? } else { return Ok(SingleValue::UnknownEnumValue(number)); @@ -678,6 +734,7 @@ impl MessageParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { check_parse_recursion_depth(depth)?; while buffer.has_remaining() { @@ -691,12 +748,27 @@ impl MessageParser { if let Some(field) = self.inner.fields.get(field_number) && field.wire_type_matches(wire_type) { - field.read_field(py, message, tag, wire_type, buffer, opts, depth)?; + field.read_field( + py, + message, + tag, + wire_type, + buffer, + opts, + depth, + alloc_budget, + )?; } else { skip_field_with_wire_type(field_number, wire_type, buffer, depth + 1)?; if !opts.ignore_unknown_fields { let field_len = checkpoint.len() - buffer.len(); - write_unknown_field(py, message, field_number, &checkpoint[..field_len])?; + write_unknown_field( + py, + message, + field_number, + &checkpoint[..field_len], + alloc_budget, + )?; } } } @@ -704,12 +776,32 @@ impl MessageParser { } } -/// Reads a scalar value from the wire buffer. +/// Reads a scalar value from the wire buffer, charging the allocation budget +/// for the resulting Python object. fn read_scalar<'py>( py: Python<'py>, s: ScalarType, buffer: &mut Bytes, + alloc_budget: &mut Budget, ) -> PyResult> { + match s { + ScalarType::Double | ScalarType::Float => alloc_budget.charge(budget::FLOAT_SIZE)?, + // Small ints are interned and larger ones vary a little in size, but + // a flat charge is close enough. + ScalarType::Int64 + | ScalarType::Uint64 + | ScalarType::Int32 + | ScalarType::Uint32 + | ScalarType::Fixed64 + | ScalarType::Fixed32 + | ScalarType::Sfixed32 + | ScalarType::Sfixed64 + | ScalarType::Sint32 + | ScalarType::Sint64 => alloc_budget.charge(budget::INT_SIZE)?, + // Bools are shared singletons; strings and bytes are charged below + // once the length is known. + ScalarType::Bool | ScalarType::String | ScalarType::Bytes => {} + } let res = match s { ScalarType::Double => { PyFloat::new(py, buffer.try_get_f64_le().map_err(map_try_get_err)?).into_any() @@ -743,12 +835,14 @@ fn read_scalar<'py>( ScalarType::String => { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; + alloc_budget.charge(budget::STR_OVERHEAD + len)?; let bytes = buffer.split_to(len); PyString::from_bytes(py, &bytes)?.into_any() } ScalarType::Bytes => { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; + alloc_budget.charge(budget::BYTES_OVERHEAD + len)?; let bytes = buffer.split_to(len); PyBytes::new(py, &bytes).into_any() } @@ -872,7 +966,11 @@ fn write_unknown_field( message: &Bound<'_, NativeMessage>, field_number: u32, field_bytes: &[u8], + alloc_budget: &mut Budget, ) -> PyResult<()> { + // The bytes copy plus the list slot holding it. The dict/list created for + // the first unknown field of a number are not charged for simplicity. + alloc_budget.charge(budget::BYTES_OVERHEAD + field_bytes.len() + budget::LIST_SLOT_SIZE)?; let unknown_fields_unbound = message.get().get_or_init_unknown_fields(py); let unknown_fields = unknown_fields_unbound.bind(py); let field_list = if let Ok(list) = unknown_fields.get_item(field_number) { diff --git a/packages/protobuf-py-ext/src/wkt_registry.rs b/packages/protobuf-py-ext/src/wkt_registry.rs index 69811f7..1ec0478 100644 --- a/packages/protobuf-py-ext/src/wkt_registry.rs +++ b/packages/protobuf-py-ext/src/wkt_registry.rs @@ -9,13 +9,14 @@ use pyo3::{ Bound, IntoPyObjectExt as _, Py, PyAny, PyResult, Python, exceptions::{PyTypeError, PyValueError}, types::{ - PyAnyMethods as _, PyBool, PyDict, PyDictMethods as _, PyFloat, PyInt, PyList, - PyListMethods as _, PyString, PyStringMethods as _, PyType, + PyAnyMethods as _, PyBool, PyBytesMethods as _, PyDict, PyDictMethods as _, PyFloat, PyInt, + PyList, PyListMethods as _, PyString, PyStringMethods as _, PyType, }, }; use crate::{ attribute_access::AttributeAccess, + budget::{self, Budget}, constants::Constants, descriptor::{DescField, DescFieldValue, DescMessage, DescSingleValue, ScalarType}, json_parse::{FieldContext, FromJsonOpts, read_json_value, read_message, read_scalar}, @@ -58,11 +59,13 @@ impl WktTimestamp { src: &mut R, _opts: &FromJsonOpts, _depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); - let (secs, nanos) = with_wkt_str(marshaler, src, |text| { + let (secs, nanos) = with_wkt_str(marshaler, src, alloc_budget, |text, _| { parse_timestamp(&marshaler.type_name, text) })?; + alloc_budget.charge(2 * budget::INT_SIZE)?; self.seconds .set(message.as_any(), PyInt::new(py, secs).as_any())?; self.nanos @@ -99,11 +102,13 @@ impl WktDuration { src: &mut R, _opts: &FromJsonOpts, _depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); - let (secs, nanos) = with_wkt_str(marshaler, src, |text| { + let (secs, nanos) = with_wkt_str(marshaler, src, alloc_budget, |text, _| { parse_duration(&marshaler.type_name, text) })?; + alloc_budget.charge(2 * budget::INT_SIZE)?; self.seconds .set(message.as_any(), PyInt::new(py, secs).as_any())?; self.nanos @@ -157,7 +162,7 @@ impl WktAny { let value = self.value.get(py, message.as_any())?.extract::()?; let inner_msg = inner_marshaler.new_empty_message(py, &inner_type)?; - inner_marshaler.merge_from_binary(py, &inner_msg, value, false)?; + inner_marshaler.merge_from_binary(py, &inner_msg, value, false, &mut Budget::new(None))?; sink.begin_object()?; if inner_marshaler.wkt.is_none() { @@ -183,12 +188,13 @@ impl WktAny { src: &mut R, opts: &FromJsonOpts, _depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let message_type_name = &marshaler.type_name; // For a string input, this will eagerly parse out a whole dictionary since // we need to first find `@type` before parsing. - let tree = read_json_value(src)?; + let tree = read_json_value(src, alloc_budget)?; let Ok(dict) = tree.cast::() else { return Err(PyTypeError::new_err(format!( "cannot decode {message_type_name} from JSON: {}", @@ -230,6 +236,7 @@ impl WktAny { .getattr(&marshaler.constants.ext_marshaler)? .cast_into::()?; let inner_marshaler = inner_marshaler.get(); + alloc_budget.charge(inner_marshaler.base_alloc_size)?; let inner_msg = inner_marshaler.new_empty_message(py, &inner_type)?; let is_wkt = inner_marshaler.wkt.is_some(); @@ -238,17 +245,19 @@ impl WktAny { .get_item("value")? .unwrap_or_else(|| py.None().into_bound(py)); let mut sub = PyTreeSource::new(py, value); - read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1)?; + read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1, alloc_budget)?; } else { let copy = dict.copy()?; copy.del_item("@type")?; let mut sub = PyTreeSource::new(py, copy.into_any()); - read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1)?; + read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1, alloc_budget)?; } // Any.pack let packed_url = format!("type.googleapis.com/{}", inner_marshaler.type_name); let packed_value = inner_marshaler.to_binary(py, &inner_msg, true)?; + alloc_budget.charge(budget::STR_OVERHEAD + packed_url.len())?; + alloc_budget.charge(budget::BYTES_OVERHEAD + packed_value.as_bytes().len())?; self.type_url .set(message.as_any(), &PyString::new(py, &packed_url).into_any())?; self.value.set(message.as_any(), packed_value.as_any())?; @@ -300,13 +309,14 @@ impl WktFieldMask { src: &mut R, _opts: &FromJsonOpts, _depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let paths = self .paths .get(py, message.as_any())? .cast_into::()?; - with_wkt_str(marshaler, src, |text| { + with_wkt_str(marshaler, src, alloc_budget, |text, alloc_budget| { if text.is_empty() { return Ok(()); } @@ -317,6 +327,7 @@ impl WktFieldMask { marshaler.type_name ))); } + alloc_budget.charge(budget::STR_OVERHEAD + part.len() + budget::LIST_SLOT_SIZE)?; paths.append(buffa_wkt::camel_to_snake(part))?; } Ok(()) @@ -360,10 +371,11 @@ impl WktStruct { src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? != JsonKind::Object { - let json = read_json_value(src)?; + let json = read_json_value(src, alloc_budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, @@ -377,9 +389,18 @@ impl WktStruct { .cast_into::()?; // Duplicate keys use last-in-wins semantics (per the ProtoJSON spec). src.for_each_object_key(|key, src| { + alloc_budget.charge(value_marshaler.base_alloc_size)?; let value_msg = value_marshaler.new_empty_message(py, self.value.get_python_type(py))?; - read_message(value_marshaler, &value_msg, src, opts, depth + 1)?; + read_message( + value_marshaler, + &value_msg, + src, + opts, + depth + 1, + alloc_budget, + )?; + alloc_budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; dict.set_item(key, value_msg)?; Ok(()) })?; @@ -422,10 +443,11 @@ impl WktListValue { src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? != JsonKind::Array { - let json = read_json_value(src)?; + let json = read_json_value(src, alloc_budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, @@ -438,9 +460,18 @@ impl WktListValue { .get(py, message.as_any())? .cast_into::()?; src.for_each_array_item(|src| { + alloc_budget.charge(element_marshaler.base_alloc_size)?; let value_msg = element_marshaler.new_empty_message(py, self.element.get_python_type(py))?; - read_message(element_marshaler, &value_msg, src, opts, depth + 1)?; + read_message( + element_marshaler, + &value_msg, + src, + opts, + depth + 1, + alloc_budget, + )?; + alloc_budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value_msg)?; Ok(()) })?; @@ -510,8 +541,10 @@ impl WktValue { src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); + alloc_budget.charge(budget::ONEOF_SIZE)?; let oneof = match src.peek()? { JsonKind::Null => { src.next_null()?; @@ -523,6 +556,7 @@ impl WktValue { ), JsonKind::Number => { let number = src.next_float()?; + alloc_budget.charge(budget::FLOAT_SIZE)?; Oneof::new( self.number_name.bind(py), &PyFloat::new(py, number).into_any(), @@ -530,20 +564,23 @@ impl WktValue { } JsonKind::String => { let string = src.next_py_str()?; + alloc_budget.charge(budget::STR_OVERHEAD + string.len()?)?; Oneof::new(self.string_name.bind(py), &string) } JsonKind::Array => { let desc = &self.list_message; let inner = desc.get_marshaler(py)?; + alloc_budget.charge(inner.base_alloc_size)?; let list_msg = inner.new_empty_message(py, desc.get_python_type(py))?; - read_message(inner, &list_msg, src, opts, depth + 1)?; + read_message(inner, &list_msg, src, opts, depth + 1, alloc_budget)?; Oneof::new(self.list_name.bind(py), &list_msg.into_any()) } JsonKind::Object => { let desc = &self.struct_message; let inner = desc.get_marshaler(py)?; + alloc_budget.charge(inner.base_alloc_size)?; let struct_msg = inner.new_empty_message(py, desc.get_python_type(py))?; - read_message(inner, &struct_msg, src, opts, depth + 1)?; + read_message(inner, &struct_msg, src, opts, depth + 1, alloc_budget)?; Oneof::new(self.struct_name.bind(py), &struct_msg.into_any()) } }; @@ -581,6 +618,7 @@ impl WktWrapper { src: &mut R, _opts: &FromJsonOpts, _depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? == JsonKind::Null { @@ -591,7 +629,7 @@ impl WktWrapper { } let name = self.name.bind(py); let ctx = FieldContext::Field { marshaler, name }; - let value = read_scalar(&ctx, src, self.scalar)?; + let value = read_scalar(&ctx, src, self.scalar, alloc_budget)?; self.field.set(message.as_any(), &value) } } @@ -639,16 +677,23 @@ impl WktKind { src: &mut R, opts: &FromJsonOpts, depth: usize, + alloc_budget: &mut Budget, ) -> PyResult<()> { match self { - WktKind::Timestamp(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::Duration(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::Any(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::FieldMask(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::Struct(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::ListValue(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::Value(w) => w.read_json(marshaler, message, src, opts, depth), - WktKind::Wrapper(w) => w.read_json(marshaler, message, src, opts, depth), + WktKind::Timestamp(w) => { + w.read_json(marshaler, message, src, opts, depth, alloc_budget) + } + WktKind::Duration(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), + WktKind::Any(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), + WktKind::FieldMask(w) => { + w.read_json(marshaler, message, src, opts, depth, alloc_budget) + } + WktKind::Struct(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), + WktKind::ListValue(w) => { + w.read_json(marshaler, message, src, opts, depth, alloc_budget) + } + WktKind::Value(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), + WktKind::Wrapper(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), } } @@ -887,17 +932,18 @@ fn match_wrapper(fields: &[DescField], by_name: &HashMap) -> Opti fn with_wkt_str<'py, S: JsonSource<'py>, R>( marshaler: &MessageMarshaler, src: &mut S, - f: impl FnOnce(&str) -> PyResult, + alloc_budget: &mut Budget, + f: impl FnOnce(&str, &mut Budget) -> PyResult, ) -> PyResult { if src.peek()? != JsonKind::String { - let value = read_json_value(src)?; + let value = read_json_value(src, alloc_budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, value.str()? ))); } - src.with_next_str(f) + src.with_next_str(|text| f(text, alloc_budget)) } fn type_url_to_name(url: &str) -> PyResult<&str> { diff --git a/src/protobuf/_budget.py b/src/protobuf/_budget.py new file mode 100644 index 0000000..2ed0ebb --- /dev/null +++ b/src/protobuf/_budget.py @@ -0,0 +1,107 @@ +# Copyright (c) 2025-2026 Buf Technologies, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from functools import cache +from typing import TYPE_CHECKING, cast + +from ._descriptors import ScalarType + +if TYPE_CHECKING: + from collections.abc import Sized + + from ._message import Message + +# Approximate sizes of CPython heap allocations, measured with `sys.getsizeof` +# on 64-bit CPython 3.14. The budget guards against unbounded allocation from +# malicious payloads rather than providing exact accounting, so small +# inaccuracies across versions and builds are fine. These values mirror the +# constants in the native extension's budget.rs. + +GC_HEAD_SIZE = 16 +"""GC header allocated in front of every GC-tracked object.""" +FLOAT_SIZE = 24 +"""A float object.""" +INT_SIZE = 36 +"""A 64-bit int object. Smaller ints are slightly smaller.""" +STR_OVERHEAD = 41 +"""Header of a compact ASCII str; the length is charged on top.""" +BYTES_OVERHEAD = 33 +"""Header of a bytes object.""" +EMPTY_LIST_SIZE = 56 +"""An empty list, as created for repeated field defaults.""" +EMPTY_DICT_SIZE = 64 +"""An empty dict, as created for map field defaults.""" +LIST_SLOT_SIZE = 8 +"""One appended list element: an 8-byte pointer slot.""" +DICT_ENTRY_SIZE = 40 +"""One inserted dict entry: hash + key + value words plus growth slack.""" +ONEOF_SIZE = 32 +"""A Oneof wrapper object: object header plus two references.""" + + +@cache +def _base_alloc_size(message_type: type[Message]) -> int: + """Approximate heap size of a freshly-initialized message instance. + + The fixed instance size (all fields are slots) plus the empty containers + created for repeated/map field defaults. + """ + size = message_type.__basicsize__ + GC_HEAD_SIZE + for _, default in message_type._desc._defaults: + if isinstance(default, list): + size += EMPTY_LIST_SIZE + elif isinstance(default, dict): + size += EMPTY_DICT_SIZE + return size + + +class Budget: + """Tracks approximate allocations while parsing a message. + + Raises an error once the configured limit is exceeded. Charges are made + before the corresponding allocation where practical so the limit bounds + the actual peak. + """ + + __slots__ = ("current", "max") + + def __init__(self, limit: int) -> None: + self.current = 0 + self.max = limit + + def charge(self, amount: int) -> None: + self.current += amount + if self.current > self.max: + msg = f"allocation budget exceeded: needed {self.current} bytes, limit is {self.max}" + raise ValueError(msg) + + def charge_scalar(self, scalar_type: ScalarType, value: object) -> None: + """Charges for one parsed scalar value of the given type.""" + if scalar_type == ScalarType.STRING: + self.charge(STR_OVERHEAD + len(cast("Sized", value))) + elif scalar_type == ScalarType.BYTES: + self.charge(BYTES_OVERHEAD + len(cast("Sized", value))) + elif scalar_type == ScalarType.BOOL: + # Bools are shared singletons; nothing is allocated. + pass + elif scalar_type in (ScalarType.DOUBLE, ScalarType.FLOAT): + self.charge(FLOAT_SIZE) + else: + self.charge(INT_SIZE) + + def charge_message(self, message_type: type[Message]) -> None: + """Charges the base size of a new instance of the message type.""" + self.charge(_base_alloc_size(message_type)) diff --git a/src/protobuf/_from_binary.py b/src/protobuf/_from_binary.py index 39968a9..ce28bad 100644 --- a/src/protobuf/_from_binary.py +++ b/src/protobuf/_from_binary.py @@ -19,6 +19,15 @@ from typing_extensions import Buffer, assert_never +from ._budget import ( + BYTES_OVERHEAD, + DICT_ENTRY_SIZE, + GC_HEAD_SIZE, + INT_SIZE, + LIST_SLOT_SIZE, + ONEOF_SIZE, + Budget, +) from ._descriptors import ( DescEnum, DescFieldValueEnum, @@ -47,9 +56,12 @@ class FromBinaryOptions: Args: ignore_unknown_fields: If `True`, unknown fields are ignored instead of being added to the message. + budget: If set, tracks approximate allocations during the parse and + raises once the configured limit is exceeded. """ ignore_unknown_fields: bool = False + budget: Budget | None = None # Dispatch table for reading scalar values. CPython currently does not generate @@ -79,10 +91,15 @@ class FromBinaryOptions: ) -def read_scalar(scalar_type: ScalarType, reader: BinaryReader) -> Any: +def read_scalar( + scalar_type: ScalarType, reader: BinaryReader, budget: Budget | None = None +) -> Any: reader_method = _SCALAR_READERS[scalar_type.value] assert reader_method is not None # noqa: S101 - return reader_method(reader) + value = reader_method(reader) + if budget is not None: + budget.charge_scalar(scalar_type, value) + return value # TODO delete this, and either: @@ -128,19 +145,31 @@ def read_message( field_raw = reader.skip(tag.wire_type, depth + 1, field_number=tag.number) if not opts.ignore_unknown_fields: key_raw = _encode_varint((tag.number << 3) | tag.wire_type) + if (budget := opts.budget) is not None: + budget.charge( + BYTES_OVERHEAD + len(key_raw) + len(field_raw) + LIST_SLOT_SIZE + ) message._get_or_init_unknown_fields().setdefault(tag.number, []).append( key_raw + bytes(field_raw) ) continue + budget = opts.budget match field_value := desc_field.value: case DescFieldValueScalar(): - message._set_member(desc_field, read_scalar(field_value.scalar, reader)) + value = read_scalar(field_value.scalar, reader, budget) + if budget is not None and field_value.oneof is not None: + budget.charge(ONEOF_SIZE) + message._set_member(desc_field, value) case DescFieldValueMessage( message=desc_nested_message, delimited_encoding=delimited_encoding ): existing: Message | None = message._get_member(desc_field) if existing is None: + if budget is not None: + budget.charge_message(desc_nested_message.type) + if field_value.oneof is not None: + budget.charge(ONEOF_SIZE) existing = desc_nested_message.type() message._set_member(desc_field, existing) if delimited_encoding: @@ -152,11 +181,13 @@ def read_message( existing, reader, opts, depth + 1, length=reader.varint() ) case DescFieldValueEnum(): - value = read_enum(field_value.enum, reader) + value = read_enum(field_value.enum, reader, budget) if isinstance(value, Enum): + if budget is not None and field_value.oneof is not None: + budget.charge(ONEOF_SIZE) message._set_member(desc_field, value) elif not opts.ignore_unknown_fields: - _write_unknown_enum_field(message, desc_field.number, value) + _write_unknown_enum_field(message, desc_field.number, value, budget) case DescFieldValueList(): read_list( message, @@ -174,6 +205,8 @@ def read_message( ) if entry: key, value = entry + if budget is not None: + budget.charge(DICT_ENTRY_SIZE) message._get_member(desc_field)[key] = value case _: assert_never(desc_field) @@ -204,15 +237,22 @@ def read_list( field_bytes = reader.skip(wire_type, depth + 1, field_number=field_number) if not opts.ignore_unknown_fields and message: key_raw = _encode_varint((field_number << 3) | wire_type) + if (budget := opts.budget) is not None: + budget.charge( + BYTES_OVERHEAD + len(key_raw) + len(field_bytes) + LIST_SLOT_SIZE + ) message._get_or_init_unknown_fields().setdefault(field_number, []).append( key_raw + bytes(field_bytes) ) return + budget = opts.budget match element_type: case ScalarType(): - value = read_scalar(element_type, reader) + value = read_scalar(element_type, reader, budget) case DescMessage(): + if budget is not None: + budget.charge_message(element_type.type) if field_value.delimited_encoding: value = read_message( element_type.type(), @@ -226,13 +266,15 @@ def read_list( element_type.type(), reader, opts, depth + 1, length=reader.varint() ) case DescEnum(): - value = read_enum(element_type, reader) + value = read_enum(element_type, reader, budget) if not isinstance(value, Enum): if not opts.ignore_unknown_fields: - _write_unknown_enum_field(message, field_number, value) + _write_unknown_enum_field(message, field_number, value, budget) return case _: assert_never(element_type) + if budget is not None: + budget.charge(LIST_SLOT_SIZE) list_.append(value) @@ -246,38 +288,53 @@ def _read_packed_list( ) -> None: length = reader.varint() end = reader.offset + length + budget = opts.budget while reader.offset < end: match element_type: case ScalarType(): - list_.append(read_scalar(element_type, reader)) + value = read_scalar(element_type, reader, budget) + if budget is not None: + budget.charge(LIST_SLOT_SIZE) + list_.append(value) case DescEnum(): - value = read_enum(element_type, reader) + value = read_enum(element_type, reader, budget) if isinstance(value, Enum): + if budget is not None: + budget.charge(LIST_SLOT_SIZE) list_.append(value) elif not opts.ignore_unknown_fields: # Even for packed fields we write unknown enum values as unpacked. - _write_unknown_enum_field(message, field_number, value) + _write_unknown_enum_field(message, field_number, value, budget) case _: assert_never(element_type) -def read_enum(desc_enum: DescEnum, reader: BinaryReader) -> Enum | int: +def read_enum( + desc_enum: DescEnum, reader: BinaryReader, budget: Budget | None = None +) -> Enum | int: value = reader.int32() - if not desc_enum.open and not desc_enum._values_by_number.get(value): - return value + if not desc_enum._values_by_number.get(value): + if not desc_enum.open: + return value + # Unknown open enum values allocate a new int-subclass instance. + if budget is not None: + budget.charge(INT_SIZE + GC_HEAD_SIZE) return desc_enum.type(value) def _write_unknown_enum_field( - message: Message | None, field_number: int, value: int + message: Message | None, field_number: int, value: int, budget: Budget | None = None ) -> None: if message is None: return writer = BinaryWriter() writer.tag(field_number, WireType.VARINT) writer.int32(value) + field_bytes = writer.finish() + if budget is not None: + budget.charge(BYTES_OVERHEAD + len(field_bytes) + LIST_SLOT_SIZE) message._get_or_init_unknown_fields().setdefault(field_number, []).append( - writer.finish() + field_bytes ) @@ -296,6 +353,7 @@ def read_map_entry( key: Any = None value: Any = None + budget = opts.budget while reader.offset < end: tag = reader.tag() if tag.number == 1: # key @@ -304,7 +362,7 @@ def read_map_entry( message, field_number, reader, start_offset, opts, depth ) return None - key = read_scalar(field_value.key, reader) + key = read_scalar(field_value.key, reader, budget) elif tag.number == 2: # value if tag.wire_type != field_value._value_wire_type: _read_unknown_map_entry( @@ -313,15 +371,17 @@ def read_map_entry( return None match field_value.value: case ScalarType() as scalar_type: - value = read_scalar(scalar_type, reader) + value = read_scalar(scalar_type, reader, budget) case DescEnum() as desc_enum: - value = read_enum(desc_enum, reader) + value = read_enum(desc_enum, reader, budget) if not isinstance(value, Enum): _read_unknown_map_entry( message, field_number, reader, start_offset, opts, depth ) return None case DescMessage(): + if budget is not None: + budget.charge_message(field_value.value.type) value = field_value.value.type() read_message(value, reader, opts, depth + 1, length=reader.varint()) case _: @@ -338,6 +398,8 @@ def read_map_entry( case DescEnum() as desc_enum: value = desc_enum.type(desc_enum.values[0].number) case DescMessage(): + if budget is not None: + budget.charge_message(field_value.value.type) value = field_value.value.type() case _: assert_never(field_value.value) @@ -361,13 +423,21 @@ def _read_unknown_map_entry( ) if not opts.ignore_unknown_fields and message: key_raw = _encode_varint((field_number << 3) | WireType.LENGTH_DELIMITED) + if (budget := opts.budget) is not None: + budget.charge( + BYTES_OVERHEAD + len(key_raw) + len(entry_bytes) + LIST_SLOT_SIZE + ) message._get_or_init_unknown_fields().setdefault(field_number, []).append( key_raw + bytes(entry_bytes) ) def merge_from_binary( - message: Message, data: Buffer, *, ignore_unknown_fields: bool = False + message: Message, + data: Buffer, + *, + ignore_unknown_fields: bool = False, + allocation_limit: int | None = None, ) -> None: """Parse serialized binary data, merging fields into an existing message. @@ -383,5 +453,11 @@ def merge_from_binary( message: The message instance to merge into. data: Serialized binary protobuf data. Must not be mutated during parsing. ignore_unknown_fields: If `True`, unknown fields in the binary data are silently discarded. + allocation_limit: If set, the approximate number of bytes of Python + objects the parse may allocate before raising a ValueError. Guards + against malicious payloads that expand into unexpectedly large + messages. """ - message._merge_from_binary(data, ignore_unknown_fields) + message._merge_from_binary( + data, ignore_unknown_fields, allocation_limit=allocation_limit + ) diff --git a/src/protobuf/_from_json.py b/src/protobuf/_from_json.py index 0aca051..66cd853 100644 --- a/src/protobuf/_from_json.py +++ b/src/protobuf/_from_json.py @@ -21,6 +21,17 @@ from typing_extensions import assert_never +from ._budget import ( + DICT_ENTRY_SIZE, + EMPTY_LIST_SIZE, + FLOAT_SIZE, + GC_HEAD_SIZE, + INT_SIZE, + LIST_SLOT_SIZE, + ONEOF_SIZE, + STR_OVERHEAD, + Budget, +) from ._descriptors import ( DescEnum, DescExtension, @@ -70,6 +81,8 @@ class FromJsonOptions: registry: Registry | None depth: int = 0 """Current message nesting depth. Mutated during a parse.""" + budget: Budget | None = None + """Optional allocation budget shared across the whole parse.""" def _enter_message(opts: FromJsonOptions) -> FromJsonOptions: @@ -86,6 +99,7 @@ def merge_from_json( *, ignore_unknown_fields: bool = False, registry: Registry | None = None, + allocation_limit: int | None = None, ) -> None: """Parse a ProtoJSON string, merging fields into an existing message. @@ -103,6 +117,10 @@ def merge_from_json( silently discarded instead of raising an error. registry: Required to read google.protobuf.Any and extensions from JSON format. + allocation_limit: If set, the approximate number of bytes of Python + objects the parse may allocate before raising a ValueError. Guards + against malicious payloads that expand into unexpectedly large + messages. Raises: json.JSONDecodeError: If json_source is not valid JSON. @@ -111,7 +129,10 @@ def merge_from_json( through the registry. """ message._merge_from_json( - json, ignore_unknown_fields=ignore_unknown_fields, registry=registry + json, + ignore_unknown_fields=ignore_unknown_fields, + registry=registry, + allocation_limit=allocation_limit, ) @@ -179,7 +200,7 @@ def _read_field( ) -> None: match field_value := field.value: case DescFieldValueScalar(): - _read_scalar_field(msg, field, field_value, json) + _read_scalar_field(msg, field, field_value, json, opts.budget) case DescFieldValueEnum(): _read_enum_field(msg, field, field_value, json, opts) case DescFieldValueMessage(): @@ -197,7 +218,7 @@ def _read_extension( ) -> None: match field_value := ext.value: case DescFieldValueScalar(): - _read_scalar_extension(msg, ext, field_value, json) + _read_scalar_extension(msg, ext, field_value, json, opts.budget) case DescFieldValueEnum(): _read_enum_extension(msg, ext, field_value, json, opts) case DescFieldValueMessage(): @@ -209,12 +230,16 @@ def _read_extension( def _read_scalar_extension( - msg: Message, ext: DescExtension, field_value: DescFieldValueScalar, json: JsonValue + msg: Message, + ext: DescExtension, + field_value: DescFieldValueScalar, + json: JsonValue, + budget: Budget | None, ) -> None: if json is None: del msg[ext.type] else: - msg[ext.type] = _read_scalar(ext, field_value.scalar, json) + msg[ext.type] = _read_scalar(ext, field_value.scalar, json, budget) def _read_enum_extension( @@ -227,7 +252,9 @@ def _read_enum_extension( if _is_resetting_null(field_value.enum, json): del msg[ext.type] else: - value = _read_enum(field_value.enum, json, opts.ignore_unknown_fields) + value = _read_enum( + field_value.enum, json, opts.ignore_unknown_fields, opts.budget + ) if value is not None: msg[ext.type] = value @@ -242,6 +269,8 @@ def _read_message_extension( if _is_resetting_null(field_value.message, json): del msg[ext.type] else: + if (budget := opts.budget) is not None: + budget.charge_message(field_value.message.type) value = field_value.message.type() _read_message(value, json, opts) msg[ext.type] = value @@ -258,21 +287,33 @@ def _read_list_extension( return if not isinstance(json, list): raise _field_error(ext, f"expected list got {type(json)}", TypeError) - msg[ext.type] = [ - v - for element in json - if (v := _read_container_item(ext, field_value.element, element, opts)) - is not None - ] + budget = opts.budget + if budget is not None: + budget.charge(EMPTY_LIST_SIZE) + values = [] + for element in json: + v = _read_container_item(ext, field_value.element, element, opts) + if v is not None: + if budget is not None: + budget.charge(LIST_SLOT_SIZE) + values.append(v) + msg[ext.type] = values def _read_scalar_field( - msg: Message, field: DescField, field_value: DescFieldValueScalar, json: JsonValue + msg: Message, + field: DescField, + field_value: DescFieldValueScalar, + json: JsonValue, + budget: Budget | None, ) -> None: if json is None: msg._del_member(field) return - msg._set_member(field, _read_scalar(field, field_value.scalar, json)) + value = _read_scalar(field, field_value.scalar, json, budget) + if budget is not None and field_value.oneof is not None: + budget.charge(ONEOF_SIZE) + msg._set_member(field, value) def _read_enum_field( @@ -285,8 +326,10 @@ def _read_enum_field( if _is_resetting_null(field_value.enum, json): msg._del_member(field) return - value = _read_enum(field_value.enum, json, opts.ignore_unknown_fields) + value = _read_enum(field_value.enum, json, opts.ignore_unknown_fields, opts.budget) if value is not None: + if (budget := opts.budget) is not None and field_value.oneof is not None: + budget.charge(ONEOF_SIZE) msg._set_member(field, value) @@ -301,12 +344,13 @@ def _read_list_field( return if not isinstance(json, list): raise _field_error(field, f"expected list got {type(json)}", TypeError) - list_.extend( - value - for element in json - if (value := _read_container_item(field, field_value.element, element, opts)) - is not None - ) + budget = opts.budget + for element in json: + value = _read_container_item(field, field_value.element, element, opts) + if value is not None: + if budget is not None: + budget.charge(LIST_SLOT_SIZE) + list_.append(value) def _read_map_field( @@ -320,10 +364,13 @@ def _read_map_field( return if not isinstance(json, dict): raise _field_error(field, f"expected dict got {type(json)}", TypeError) + budget = opts.budget for json_key, json_value in json.items(): - key = _read_map_key(field, field_value, json_key) + key = _read_map_key(field, field_value, json_key, budget) value = _read_container_item(field, field_value.value, json_value, opts) if value is not None: + if budget is not None: + budget.charge(DICT_ENTRY_SIZE) dict_[key] = value @@ -337,17 +384,24 @@ def _read_message_field( if _is_resetting_null(field_value.message, json): msg._del_member(field) return - value = ( - msg._get_member(field) - if msg._contains_member(field) - else field_value.message.type() - ) + budget = opts.budget + if msg._contains_member(field): + value = msg._get_member(field) + else: + if budget is not None: + budget.charge_message(field_value.message.type) + value = field_value.message.type() _read_message(value, json, opts) + if budget is not None and field_value.oneof is not None: + budget.charge(ONEOF_SIZE) msg._set_member(field, value) def _read_map_key( - field: DescField, field_value: DescFieldValueMap, json: JsonValue + field: DescField, + field_value: DescFieldValueMap, + json: JsonValue, + budget: Budget | None, ) -> bool | int | str: match field_value.key: case ScalarType.BOOL: @@ -357,13 +411,18 @@ def _read_map_key( return False raise _field_error(field, f"unexpected bool map key value {json}") case ScalarType.STRING: - return _read_string(field, json) + key = _read_string(field, json) + if budget is not None: + budget.charge(STR_OVERHEAD + len(key)) + return key case ( ScalarType.DOUBLE | ScalarType.FLOAT | ScalarType.BYTES ): # This is because the Map key is not narrow enough msg = f"invalid map key type: {field_value.key}" raise AssertionError(msg) case _: + if budget is not None: + budget.charge(INT_SIZE) return _read_int(field, field_value.key, json) @@ -374,13 +433,15 @@ def _read_container_item( opts: FromJsonOptions, ) -> bool | int | float | str | bytes | Message | Enum | None: if isinstance(element, ScalarType) and json is not None: - return _read_scalar(field, element, json) + return _read_scalar(field, element, json, opts.budget) if isinstance(element, DescMessage) and not _is_resetting_null(element, json): + if (budget := opts.budget) is not None: + budget.charge_message(element.type) msg = element.type() _read_message(msg, json, opts) return msg if isinstance(element, DescEnum) and not _is_resetting_null(element, json): - return _read_enum(element, json, opts.ignore_unknown_fields) + return _read_enum(element, json, opts.ignore_unknown_fields, opts.budget) raise _field_error( field, f"unexpected null value for {'map value' if isinstance(field, DescField) and isinstance(field.value, DescFieldValueMap) else 'list item'}", @@ -391,6 +452,7 @@ def _read_enum( desc: DescEnum, json: JsonValue, ignore_unknown_fields: bool, # noqa: FBT001 + budget: Budget | None = None, ) -> Enum | None: if json is None: return desc.type(desc.values[0].number) @@ -401,7 +463,10 @@ def _read_enum( return desc.type(value.number) if ignore_unknown_fields: return None - # Succeeds for open enum, raises an error for closed + # Succeeds for open enum (allocating a new int-subclass instance), + # raises an error for closed. + if budget is not None: + budget.charge(INT_SIZE + GC_HEAD_SIZE) return desc.type(json) if isinstance(json, str): if value := desc._values_by_name.get(json): @@ -416,6 +481,18 @@ def _read_enum( def _read_scalar( + desc: DescField | DescExtension, + scalar_type: ScalarType, + json: JsonValue, + budget: Budget | None = None, +) -> bool | int | float | str | bytes: + value = _read_scalar_value(desc, scalar_type, json) + if budget is not None: + budget.charge_scalar(scalar_type, value) + return value + + +def _read_scalar_value( desc: DescField | DescExtension, scalar_type: ScalarType, json: JsonValue ) -> bool | int | float | str | bytes: match scalar_type: @@ -605,9 +682,14 @@ def _struct_from_json( assert isinstance(value_desc, DescMessage) # noqa: S101 value_wkt = match_wkt(value_desc) assert isinstance(value_wkt, WktValue) # noqa: S101 + budget = opts.budget for k, v in json.items(): + if budget is not None: + budget.charge_message(value_desc.type) val = cast("Value", value_desc.type()) _value_from_json(val, v, opts, value_wkt) + if budget is not None: + budget.charge(DICT_ENTRY_SIZE + STR_OVERHEAD + len(k)) msg.fields[k] = val @@ -624,9 +706,14 @@ def _list_value_from_json( assert isinstance(element_desc, DescMessage) # noqa: S101 element_wkt = match_wkt(element_desc) assert isinstance(element_wkt, WktValue) # noqa: S101 + budget = opts.budget for e in json: + if budget is not None: + budget.charge_message(element_desc.type) val = cast("Value", element_desc.type()) _value_from_json(val, e, opts, element_wkt) + if budget is not None: + budget.charge(LIST_SLOT_SIZE) msg.values.append(val) @@ -642,6 +729,22 @@ def _value_from_json( def _value_from_json_inner( msg: Value, json: JsonValue, opts: FromJsonOptions, wkt: WktValue ) -> None: + budget = opts.budget + if budget is not None: + budget.charge(ONEOF_SIZE) + match json: + case bool() | None: + pass + case int() | float(): + budget.charge(FLOAT_SIZE) + case str(): + budget.charge(STR_OVERHEAD + len(json)) + case list(): + budget.charge_message(wkt.list_value.message.type) + case dict(): + budget.charge_message(wkt.struct_value.message.type) + case _: + pass match json: case None: msg.kind = Oneof( @@ -677,6 +780,7 @@ def message_from_json_value( *, ignore_unknown_fields: bool = False, registry: Registry | None = None, + allocation_limit: int | None = None, ) -> T: """Converts the Python value parsed from JSON data to a new Message of the given type. @@ -700,5 +804,8 @@ def message_from_json_value( ``` """ return message_type._from_json_value( - data, ignore_unknown_fields=ignore_unknown_fields, registry=registry + data, + ignore_unknown_fields=ignore_unknown_fields, + registry=registry, + allocation_limit=allocation_limit, ) diff --git a/src/protobuf/_message.py b/src/protobuf/_message.py index e5ab7b2..5955451 100644 --- a/src/protobuf/_message.py +++ b/src/protobuf/_message.py @@ -584,6 +584,7 @@ def from_json( *, ignore_unknown_fields: bool = False, registry: Registry | None = None, + allocation_limit: int | None = None, ) -> Self: """Create a new message from a ProtoJSON string. @@ -596,6 +597,10 @@ def from_json( registry: This option is required to read `google.protobuf.Any` and extensions from JSON format. + allocation_limit: + If set, the approximate number of bytes of Python objects the + parse may allocate before raising a ValueError. Guards against + malicious payloads that expand into unexpectedly large messages. Raises: json.JSONDecodeError: If json_source is not valid JSON. @@ -605,13 +610,20 @@ def from_json( """ msg = cls() msg._merge_from_json( - json, ignore_unknown_fields=ignore_unknown_fields, registry=registry + json, + ignore_unknown_fields=ignore_unknown_fields, + registry=registry, + allocation_limit=allocation_limit, ) return msg @classmethod def from_binary( - cls: type[Self], data: Buffer, *, ignore_unknown_fields: bool = False + cls: type[Self], + data: Buffer, + *, + ignore_unknown_fields: bool = False, + allocation_limit: int | None = None, ) -> Self: """Create a new message by parsing serialized binary data. @@ -620,9 +632,17 @@ def from_binary( Args: data: Serialized binary protobuf data. Must not be mutated during parsing. ignore_unknown_fields: If `True`, unknown fields in the binary data are silently discarded. + allocation_limit: + If set, the approximate number of bytes of Python objects the + parse may allocate before raising a ValueError. Guards against + malicious payloads that expand into unexpectedly large messages. """ message = cls() - message._merge_from_binary(data, ignore_unknown_fields=ignore_unknown_fields) + message._merge_from_binary( + data, + ignore_unknown_fields=ignore_unknown_fields, + allocation_limit=allocation_limit, + ) return message @classmethod @@ -671,8 +691,24 @@ def _resolve_field(self, key: str) -> DescField: # Marshaling methods overridden in native code when available. - def _merge_from_binary(self, data: Buffer, ignore_unknown_fields: bool) -> None: # noqa: FBT001 - opts = FromBinaryOptions(ignore_unknown_fields=ignore_unknown_fields) + def _merge_from_binary( + self, + data: Buffer, + ignore_unknown_fields: bool, # noqa: FBT001 + allocation_limit: int | None = None, + ) -> None: + from ._budget import Budget # noqa: PLC0415 + + budget = None + if allocation_limit is not None: + budget = Budget(allocation_limit) + # The root is charged here rather than in from_binary so both + # entry points share one budget semantic; a merge into an + # existing message overcharges by one base size. + budget.charge_message(type(self)) + opts = FromBinaryOptions( + ignore_unknown_fields=ignore_unknown_fields, budget=budget + ) view = memoryview(data) read_message(self, BinaryReader(view), opts, depth=0, length=len(view)) @@ -682,15 +718,25 @@ def _merge_from_json( *, ignore_unknown_fields: bool = False, registry: Registry | None = None, + allocation_limit: int | None = None, ) -> None: from json import loads as parse_json # noqa: PLC0415 + from ._budget import Budget # noqa: PLC0415 + # Needs to be lazy import since JSON specially handles many WKTs. from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 json_value = parse_json(json) + budget = None + if allocation_limit is not None: + budget = Budget(allocation_limit) + # See _merge_from_binary for why the root is charged here. + budget.charge_message(type(self)) opts = FromJsonOptions( - ignore_unknown_fields=ignore_unknown_fields, registry=registry + ignore_unknown_fields=ignore_unknown_fields, + registry=registry, + budget=budget, ) _read_message(self, json_value, opts) @@ -730,16 +776,25 @@ def _from_json_value( *, ignore_unknown_fields: bool = False, registry: Registry | None = None, + allocation_limit: int | None = None, ) -> Self: + from ._budget import Budget # noqa: PLC0415 + # Needs to be lazy import since JSON specially handles many WKTs. from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 + budget = None + if allocation_limit is not None: + budget = Budget(allocation_limit) + budget.charge_message(cls) message = cls() _read_message( message, data, FromJsonOptions( - ignore_unknown_fields=ignore_unknown_fields, registry=registry + ignore_unknown_fields=ignore_unknown_fields, + registry=registry, + budget=budget, ), ) return message diff --git a/src/protobuf/_wkt_registry.py b/src/protobuf/_wkt_registry.py index e954183..397e999 100644 --- a/src/protobuf/_wkt_registry.py +++ b/src/protobuf/_wkt_registry.py @@ -98,7 +98,8 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: nanos_str = nanos_str[:6] return f"{iso_secs}.{nanos_str}Z" - def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bool: + def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: + from ._budget import INT_SIZE # noqa: PLC0415 from .wkt import Timestamp # noqa: PLC0415 value = cast("Timestamp", msg) @@ -128,6 +129,8 @@ def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bo " 0001-01-01T00:00:00Z to 9999-12-31T23:59:59Z inclusive" ) raise ValueError(err) + if opts.budget is not None: + opts.budget.charge(2 * INT_SIZE) value.seconds = Timestamp.from_datetime(dt).seconds value.nanos = nanos return True @@ -170,7 +173,8 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: text = "-" + text return text + "s" - def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bool: + def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: + from ._budget import INT_SIZE # noqa: PLC0415 from .wkt._mixin._const import DURATION_SECONDS_MAX # noqa: PLC0415 value = cast("Duration", msg) @@ -190,6 +194,8 @@ def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bo nanos = int(duration_match[2] + "0" * (9 - len(duration_match[2]))) if seconds < 0 or duration_match[1] == "-0": nanos = -nanos + if opts.budget is not None: + opts.budget.charge(2 * INT_SIZE) value.seconds = seconds value.nanos = nanos return True @@ -248,6 +254,8 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo if not desc: err = f"cannot decode {Any._desc.type_name} from JSON: {type_url} is not in the type registry" raise ValueError(err) + if (budget := opts.budget) is not None: + budget.charge_message(desc.type) message = desc.type() if _has_custom_json(desc) and "value" in json: _read_message(message, json["value"], opts) @@ -256,6 +264,11 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo del json["@type"] _read_message(message, json, opts) any_ = Any.pack(message) + if budget is not None: + from ._budget import BYTES_OVERHEAD, STR_OVERHEAD # noqa: PLC0415 + + budget.charge(STR_OVERHEAD + len(any_.type_url)) + budget.charge(BYTES_OVERHEAD + len(any_.value)) value.type_url = any_.type_url value.value = any_.value return True @@ -283,7 +296,8 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: parts.append(proto_camel_path) return ",".join(parts) - def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bool: + def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: + from ._budget import LIST_SLOT_SIZE, STR_OVERHEAD # noqa: PLC0415 from ._names import proto_snake_case # noqa: PLC0415 value = cast("FieldMask", msg) @@ -296,6 +310,8 @@ def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bo if "_" in path: err = f"cannot decode {value._desc.type_name} from JSON: path names must be lowerCamelCase" raise ValueError(err) + if opts.budget is not None: + opts.budget.charge(STR_OVERHEAD + len(path) + LIST_SLOT_SIZE) value.paths.append(proto_snake_case(path)) return True @@ -378,13 +394,15 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: return _scalar_to_json_value(self.value.scalar, msg[self.field]) - def from_json(self, msg: Message, json: JsonValue, _opts: FromJsonOptions) -> bool: + def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: from ._from_json import _read_scalar # noqa: PLC0415 if json is None: del msg[self.field] else: - msg[self.field] = _read_scalar(self.field, self.value.scalar, json) + msg[self.field] = _read_scalar( + self.field, self.value.scalar, json, opts.budget + ) return True def mixin(self) -> type | None: diff --git a/tests/test_from_binary.py b/tests/test_from_binary.py index 363e2b2..1d141fe 100644 --- a/tests/test_from_binary.py +++ b/tests/test_from_binary.py @@ -398,3 +398,44 @@ def test_delimited_encoding(self) -> None: # Should be enough to check every field is there without asserting the content. assert msg._unknown_fields is not None assert {2, 3} == set(msg._unknown_fields.keys()) + + +class TestAllocationLimit: + def test_within_limit(self) -> None: + data = Scalars(string_field="a" * 1000).to_binary() + msg = Scalars.from_binary(data, allocation_limit=1024 * 1024) + assert msg.string_field == "a" * 1000 + + def test_string_exceeds_limit(self) -> None: + data = Scalars(string_field="a" * 1000).to_binary() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Scalars.from_binary(data, allocation_limit=500) + + def test_repeated_exceeds_limit(self) -> None: + data = Lists(string_list=["a" * 100] * 100).to_binary() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Lists.from_binary(data, allocation_limit=5000) + + def test_nested_messages_exceed_limit(self) -> None: + data = Lists(msg_list=[Lists.Msg() for _ in range(1000)]).to_binary() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Lists.from_binary(data, allocation_limit=10_000) + + def test_unknown_fields_exceed_limit(self) -> None: + w = BinaryWriter() + w.tag(1000, WireType.LENGTH_DELIMITED) + w.bytes_(b"x" * 10_000) + data = w.finish() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Scalars.from_binary(data, allocation_limit=5000) + + def test_no_limit_by_default(self) -> None: + data = Lists(string_list=["a" * 100] * 100).to_binary() + msg = Lists.from_binary(data) + assert len(msg.string_list) == 100 + + def test_merge_from_binary_limit(self) -> None: + data = Scalars(string_field="a" * 1000).to_binary() + msg = Scalars() + with pytest.raises(ValueError, match="allocation budget exceeded"): + merge_from_binary(msg, data, allocation_limit=500) diff --git a/tests/test_json.py b/tests/test_json.py index f15c6b2..6b38ea5 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -25,6 +25,7 @@ message_from_json_value, message_to_json_value, ) +from protobuf.wkt import Struct from .gen.enums_pb import ClosedColor, Color, EnumMessage from .gen.json_enum_names_pb import JsonEnumNames, Season @@ -824,3 +825,47 @@ def test_to_json_cyclic_message() -> None: match="exceeded maximum recursion depth 100 while serializing message", ): message_to_json_value(msg) + + +class TestAllocationLimit: + def test_within_limit(self) -> None: + data = Scalars(string_field="a" * 1000).to_json() + msg = Scalars.from_json(data, allocation_limit=1024 * 1024) + assert msg.string_field == "a" * 1000 + + def test_string_exceeds_limit(self) -> None: + data = Scalars(string_field="a" * 1000).to_json() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Scalars.from_json(data, allocation_limit=500) + + def test_repeated_exceeds_limit(self) -> None: + data = Lists(string_list=["a" * 100] * 100).to_json() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Lists.from_json(data, allocation_limit=5000) + + def test_nested_messages_exceed_limit(self) -> None: + data = Lists(msg_list=[Lists.Msg() for _ in range(1000)]).to_json() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Lists.from_json(data, allocation_limit=10_000) + + def test_map_exceeds_limit(self) -> None: + data = Maps( + string_to_string={f"key{i}": "x" * 50 for i in range(100)} + ).to_json() + with pytest.raises(ValueError, match="allocation budget exceeded"): + Maps.from_json(data, allocation_limit=2000) + + def test_struct_exceeds_limit(self) -> None: + data = json.dumps({f"key{i}": "x" * 50 for i in range(100)}) + with pytest.raises(ValueError, match="allocation budget exceeded"): + Struct.from_json(data, allocation_limit=2000) + + def test_from_json_value_limit(self) -> None: + data = cast("dict[str, Any]", {"stringList": ["a" * 100] * 100}) + with pytest.raises(ValueError, match="allocation budget exceeded"): + message_from_json_value(Lists, data, allocation_limit=5000) + + def test_no_limit_by_default(self) -> None: + data = Lists(string_list=["a" * 100] * 100).to_json() + msg = Lists.from_json(data) + assert len(msg.string_list) == 100 From 77c9f22cf5ae0e91419bb410d4d8227828112146 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 10:45:51 +0900 Subject: [PATCH 02/11] cleanups --- packages/protobuf-py-ext/src/budget.rs | 19 +- packages/protobuf-py-ext/src/json_parse.rs | 205 ++++++++---------- packages/protobuf-py-ext/src/marshaler.rs | 4 +- packages/protobuf-py-ext/src/nativemessage.rs | 45 ++-- packages/protobuf-py-ext/src/parser.rs | 106 ++++----- packages/protobuf-py-ext/src/wkt_registry.rs | 118 +++++----- src/protobuf/_budget.py | 11 +- src/protobuf/_from_binary.py | 5 +- src/protobuf/_from_json.py | 7 +- 9 files changed, 218 insertions(+), 302 deletions(-) diff --git a/packages/protobuf-py-ext/src/budget.rs b/packages/protobuf-py-ext/src/budget.rs index d15d5d8..b769e89 100644 --- a/packages/protobuf-py-ext/src/budget.rs +++ b/packages/protobuf-py-ext/src/budget.rs @@ -5,35 +5,30 @@ use pyo3::{PyResult, exceptions::PyValueError}; // unbounded allocation from malicious payloads rather than providing exact // accounting, so small inaccuracies across versions and builds are fine. -/// GC header allocated in front of every GC-tracked object. `tp_basicsize` -/// does not include it; `sys.getsizeof` does. +/// GC header allocated in front of every GC-tracked object. pub(crate) const GC_HEAD_SIZE: usize = 16; /// A float object. pub(crate) const FLOAT_SIZE: usize = 24; /// A 64-bit int object. Smaller ints are slightly smaller. pub(crate) const INT_SIZE: usize = 36; /// Header of a compact ASCII str. The UTF-8 byte length is charged on top -/// as an approximation of the payload; non-ASCII headers are slightly larger. +/// as an approximation of the payload. pub(crate) const STR_OVERHEAD: usize = 41; /// Header of a bytes object. pub(crate) const BYTES_OVERHEAD: usize = 33; -/// An empty list, as created for repeated field defaults. +/// An empty list. pub(crate) const EMPTY_LIST_SIZE: usize = 56; -/// An empty dict, as created for map field defaults. +/// An empty dict. pub(crate) const EMPTY_DICT_SIZE: usize = 64; -/// One appended list element: an 8-byte pointer slot (amortized growth -/// measures ~8.8 bytes per item). +/// One appended list element: an 8-byte pointer slot. pub(crate) const LIST_SLOT_SIZE: usize = 8; -/// One inserted dict entry: hash + key + value words plus index table and -/// growth slack (amortized growth measures ~37 bytes per item). +/// One inserted dict entry: hash + key + value words plus growth slack. pub(crate) const DICT_ENTRY_SIZE: usize = 40; /// A `Oneof` wrapper object: `PyObject` header plus two object pointers. pub(crate) const ONEOF_SIZE: usize = 32; /// Tracks the approximate bytes of Python objects allocated while parsing a -/// message, raising an error once a configured limit is exceeded. Charges are -/// made before the corresponding allocation so the limit bounds the actual -/// peak, not just the size observed after the fact. +/// message, raising an error once a configured limit is exceeded. pub(crate) struct Budget { current: usize, max: usize, diff --git a/packages/protobuf-py-ext/src/json_parse.rs b/packages/protobuf-py-ext/src/json_parse.rs index b2e6e83..5166c49 100644 --- a/packages/protobuf-py-ext/src/json_parse.rs +++ b/packages/protobuf-py-ext/src/json_parse.rs @@ -54,10 +54,10 @@ pub(crate) fn merge_from_json<'py>( message: &Bound<'py, NativeMessage>, data: &[u8], opts: &FromJsonOpts, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let mut src = JiterSource::new(py, data); - read_message(marshaler, message, &mut src, opts, 0, alloc_budget)?; + read_message(marshaler, message, &mut src, opts, 0, budget)?; src.finish() } @@ -68,10 +68,10 @@ pub(crate) fn read_message_from_tree<'py>( message: &Bound<'py, NativeMessage>, tree: Bound<'py, PyAny>, opts: &FromJsonOpts, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let mut src = PyTreeSource::new(py, tree); - read_message(marshaler, message, &mut src, opts, 0, alloc_budget) + read_message(marshaler, message, &mut src, opts, 0, budget) } /// Reads a message value, dispatching on its well-known-type kind. Ordinary @@ -83,7 +83,7 @@ pub(crate) fn read_message<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { if depth > DEPTH_LIMIT { return Err(pyo3::exceptions::PyRecursionError::new_err(format!( @@ -91,8 +91,8 @@ pub(crate) fn read_message<'py, R: JsonSource<'py>>( ))); } match &marshaler.wkt { - Some(wkt) => wkt.read_json(marshaler, message, src, opts, depth, alloc_budget), - None => read_generic_object(marshaler, message, src, opts, depth, alloc_budget), + Some(wkt) => wkt.read_json(marshaler, message, src, opts, depth, budget), + None => read_generic_object(marshaler, message, src, opts, depth, budget), } } @@ -102,11 +102,11 @@ fn read_generic_object<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? != JsonKind::Object { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; let qualname = marshaler.python_type.bind(py).qualname()?; return Err(PyTypeError::new_err(format!( "cannot decode {qualname} from JSON: {}", @@ -137,9 +137,9 @@ fn read_generic_object<'py, R: JsonSource<'py>>( if merges_on_duplicate(parser) && !seen.insert(number) { reset_duplicate_field(py, parser, message)?; } - read_field(marshaler, parser, message, src, opts, depth, alloc_budget)?; + read_field(marshaler, parser, message, src, opts, depth, budget)?; } else { - handle_unknown_key(marshaler, message, key, src, opts, depth, alloc_budget)?; + handle_unknown_key(marshaler, message, key, src, opts, depth, budget)?; } Ok(()) })?; @@ -183,7 +183,7 @@ fn read_field<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { match &parser.type_ { ParserFieldType::Singular { @@ -198,10 +198,10 @@ fn read_field<'py, R: JsonSource<'py>>( depth, oneof_attr.as_ref(), *requires_presence, - alloc_budget, + budget, ), ParserFieldType::List { .. } => { - read_list(marshaler, parser, message, src, opts, depth, alloc_budget) + read_list(marshaler, parser, message, src, opts, depth, budget) } ParserFieldType::Map { key_type, @@ -216,7 +216,7 @@ fn read_field<'py, R: JsonSource<'py>>( depth, *key_type, value_parser, - alloc_budget, + budget, ), } } @@ -231,7 +231,7 @@ fn read_singular<'py, R: JsonSource<'py>>( depth: usize, oneof_attr: Option<&crate::attribute_access::AttributeAccess>, requires_presence: bool, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); match &parser.value { @@ -249,9 +249,9 @@ fn read_singular<'py, R: JsonSource<'py>>( } let name = parser.name.bind(py); let ctx = FieldContext::Field { marshaler, name }; - let value = read_scalar(&ctx, src, *scalar, alloc_budget)?; + let value = read_scalar(&ctx, src, *scalar, budget)?; if oneof_attr.is_some() { - alloc_budget.charge(budget::ONEOF_SIZE)?; + budget.charge(budget::ONEOF_SIZE)?; } parser.assign_singular(py, message, &value, oneof_attr, requires_presence) } @@ -267,9 +267,9 @@ fn read_singular<'py, R: JsonSource<'py>>( )?; return Ok(()); } - if let Some(value) = read_enum(enum_, src, opts, alloc_budget)? { + if let Some(value) = read_enum(enum_, src, opts, budget)? { if oneof_attr.is_some() { - alloc_budget.charge(budget::ONEOF_SIZE)?; + budget.charge(budget::ONEOF_SIZE)?; } parser.assign_singular(py, message, &value, oneof_attr, requires_presence)?; } @@ -289,13 +289,13 @@ fn read_singular<'py, R: JsonSource<'py>>( let target = match existing { Some(value) if !value.is_none() => value.cast_into::()?, _ => { - alloc_budget.charge(inner.base_alloc_size)?; + budget.charge(inner.alloc_size)?; inner.new_empty_message(py, msg_desc.get_python_type(py))? } }; - read_message(inner, &target, src, opts, depth + 1, alloc_budget)?; + read_message(inner, &target, src, opts, depth + 1, budget)?; if oneof_attr.is_some() { - alloc_budget.charge(budget::ONEOF_SIZE)?; + budget.charge(budget::ONEOF_SIZE)?; } parser.assign_singular(py, message, target.as_any(), oneof_attr, requires_presence) } @@ -335,7 +335,7 @@ fn read_list<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let name = parser.name.bind(py); @@ -348,7 +348,7 @@ fn read_list<'py, R: JsonSource<'py>>( return Err(ctx.error( &format!( "expected list got {}", - read_json_value(src, alloc_budget)?.get_type() + read_json_value(src, budget)?.get_type() ), Exc::Type, )); @@ -359,9 +359,9 @@ fn read_list<'py, R: JsonSource<'py>>( .cast_into::()?; src.for_each_array_item(|src| { if let Some(value) = - read_container_item(&ctx, &parser.value, src, opts, depth, false, alloc_budget)? + read_container_item(&ctx, &parser.value, src, opts, depth, false, budget)? { - alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } Ok(()) @@ -379,7 +379,7 @@ fn read_map<'py, R: JsonSource<'py>>( depth: usize, key_type: ScalarType, value_parser: &FieldParser, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let name = parser.name.bind(py); @@ -392,7 +392,7 @@ fn read_map<'py, R: JsonSource<'py>>( return Err(ctx.error( &format!( "expected dict got {}", - read_json_value(src, alloc_budget)?.get_type() + read_json_value(src, budget)?.get_type() ), Exc::Type, )); @@ -402,17 +402,11 @@ fn read_map<'py, R: JsonSource<'py>>( .get(py, message.as_any())? .cast_into::()?; src.for_each_object_key(|key, src| { - let map_key = read_map_key(py, &ctx, key_type, key, alloc_budget)?; - if let Some(value) = read_container_item( - &ctx, - &value_parser.value, - src, - opts, - depth, - true, - alloc_budget, - )? { - alloc_budget.charge(budget::DICT_ENTRY_SIZE)?; + let map_key = read_map_key(py, &ctx, key_type, key, budget)?; + if let Some(value) = + read_container_item(&ctx, &value_parser.value, src, opts, depth, true, budget)? + { + budget.charge(budget::DICT_ENTRY_SIZE)?; dict.set_item(map_key, value)?; } Ok(()) @@ -429,13 +423,13 @@ fn read_container_item<'py, R: JsonSource<'py>>( opts: &FromJsonOpts, depth: usize, is_map: bool, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult>> { let py = src.py(); let is_null = src.peek()? == JsonKind::Null; match element { FieldParserValue::Scalar(scalar) if !is_null => { - Ok(Some(read_scalar(ctx, src, *scalar, alloc_budget)?)) + Ok(Some(read_scalar(ctx, src, *scalar, budget)?)) } FieldParserValue::Message { message: msg_desc, .. @@ -446,13 +440,13 @@ fn read_container_item<'py, R: JsonSource<'py>>( src.next_null()?; return Err(container_null_error(ctx, is_map)); } - alloc_budget.charge(inner.base_alloc_size)?; + budget.charge(inner.alloc_size)?; let target = inner.new_empty_message(py, msg_desc.get_python_type(py))?; - read_message(inner, &target, src, opts, depth + 1, alloc_budget)?; + read_message(inner, &target, src, opts, depth + 1, budget)?; Ok(Some(target.into_any())) } FieldParserValue::Enum(enum_) if !is_null || enum_.is_null_value => { - read_enum(enum_, src, opts, alloc_budget) + read_enum(enum_, src, opts, budget) } _ => { // Resetting null for a list item / map value: error. @@ -472,7 +466,7 @@ fn read_map_key<'py>( ctx: &FieldContext<'_, 'py>, key_type: ScalarType, raw_key: &str, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { match key_type { ScalarType::Bool => match raw_key { @@ -484,29 +478,28 @@ fn read_map_key<'py>( )), }, ScalarType::String => { - alloc_budget.charge(budget::STR_OVERHEAD + raw_key.len())?; + budget.charge(budget::STR_OVERHEAD + raw_key.len())?; Ok(PyString::new(py, raw_key).into_any()) } _ => { - alloc_budget.charge(budget::INT_SIZE)?; + budget.charge(budget::INT_SIZE)?; parse_int_string(py, ctx, raw_key, key_type) } } } -/// Reads a scalar value, charging the allocation budget for the resulting -/// Python object. +/// Reads a scalar value. pub(crate) fn read_scalar<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, scalar: ScalarType, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { let py = src.py(); match scalar { ScalarType::Bool => { if src.peek()? != JsonKind::Bool { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; return Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -515,28 +508,28 @@ pub(crate) fn read_scalar<'py, R: JsonSource<'py>>( Ok(PyBool::new(py, src.next_bool()?).to_owned().into_any()) } ScalarType::Float => { - let value = parse_float(ctx, src, alloc_budget)?; + let value = parse_float(ctx, src, budget)?; if value.is_finite() && !(FLOAT32_MIN..=FLOAT32_MAX).contains(&value) { return Err(ctx.error(&format!("float value out of range: {value}"), Exc::Overflow)); } - alloc_budget.charge(budget::FLOAT_SIZE)?; + budget.charge(budget::FLOAT_SIZE)?; Ok(PyFloat::new(py, value).into_any()) } ScalarType::Double => { - let value = parse_float(ctx, src, alloc_budget)?; - alloc_budget.charge(budget::FLOAT_SIZE)?; + let value = parse_float(ctx, src, budget)?; + budget.charge(budget::FLOAT_SIZE)?; Ok(PyFloat::new(py, value).into_any()) } ScalarType::String => { - let value = read_string(ctx, src, alloc_budget)?; + let value = read_string(ctx, src, budget)?; // Character count approximates the payload size. - alloc_budget.charge(budget::STR_OVERHEAD + value.len()?)?; + budget.charge(budget::STR_OVERHEAD + value.len()?)?; Ok(value.into_any()) } - ScalarType::Bytes => read_bytes(ctx, src, alloc_budget), + ScalarType::Bytes => read_bytes(ctx, src, budget), _ => { - alloc_budget.charge(budget::INT_SIZE)?; - read_int(ctx, src, scalar, alloc_budget) + budget.charge(budget::INT_SIZE)?; + read_int(ctx, src, scalar, budget) } } } @@ -544,10 +537,10 @@ pub(crate) fn read_scalar<'py, R: JsonSource<'py>>( fn read_string<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { if src.peek()? != JsonKind::String { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; return Err(ctx.error( &format!("expected string got: {}", value.get_type()), Exc::Type, @@ -559,11 +552,11 @@ fn read_string<'py, R: JsonSource<'py>>( fn read_bytes<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { let py = src.py(); if src.peek()? != JsonKind::String { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; return Err(ctx.error( &format!("expected base64-encoded string got: {}", value.get_type()), Exc::Type, @@ -581,7 +574,7 @@ fn read_bytes<'py, R: JsonSource<'py>>( .map_err(|_| ctx.error("invalid base64 data", Exc::Value))?; Ok(decoded) })?; - alloc_budget.charge(budget::BYTES_OVERHEAD + decoded.len())?; + budget.charge(budget::BYTES_OVERHEAD + decoded.len())?; Ok(PyBytes::new(py, &decoded).into_any()) } @@ -605,7 +598,7 @@ fn base64_url_safe() -> base64::engine::GeneralPurpose { fn parse_float<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult { match src.peek()? { JsonKind::Number => { @@ -635,7 +628,7 @@ fn parse_float<'py, R: JsonSource<'py>>( } }), _ => { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -649,7 +642,7 @@ fn read_int<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, int_type: ScalarType, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { let py = src.py(); let value = match src.peek()? { @@ -675,7 +668,7 @@ fn read_int<'py, R: JsonSource<'py>>( }); } _ => { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; return Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -752,7 +745,7 @@ fn read_enum<'py, R: JsonSource<'py>>( enum_desc: &DescEnum, src: &mut R, opts: &FromJsonOpts, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult>> { let py = src.py(); match src.peek()? { @@ -771,7 +764,7 @@ fn read_enum<'py, R: JsonSource<'py>>( if opts.ignore_unknown_fields { return Ok(None); } - alloc_budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; + budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; return Ok(Some(enum_desc.py_type.bind(py).call1((number,))?)); }; if let Some(value) = enum_desc.values.get(&int_value) { @@ -779,9 +772,8 @@ fn read_enum<'py, R: JsonSource<'py>>( } else if opts.ignore_unknown_fields { Ok(None) } else { - // Open enum: succeeds and allocates a new int-subclass - // instance; closed enum: raises via Python enum call. - alloc_budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; + // Open enum: succeeds; closed enum: raises via Python enum call. + budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; Ok(Some(enum_desc.py_type.bind(py).call1((int_value,))?)) } } @@ -800,7 +792,7 @@ fn read_enum<'py, R: JsonSource<'py>>( } }), _ => { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; Err(decode_enum_error(py, enum_desc, &value)) } } @@ -825,7 +817,7 @@ fn handle_unknown_key<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if raw_key.starts_with('[') @@ -841,15 +833,7 @@ fn handle_unknown_key<'py, R: JsonSource<'py>>( let extendee_name = extendee.getattr(&marshaler.constants.type_name)?; let extendee_name = extendee_name.cast::()?.to_str()?; if extendee_name == &*marshaler.type_name { - read_extension( - marshaler, - message, - &extension, - src, - opts, - depth, - alloc_budget, - )?; + read_extension(marshaler, message, &extension, src, opts, depth, budget)?; } else { src.skip()?; } @@ -874,7 +858,7 @@ fn read_extension<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let ext_type = extension.getattr(&marshaler.constants.type_)?; @@ -893,7 +877,7 @@ fn read_extension<'py, R: JsonSource<'py>>( src.next_null()?; target.del_item(&ext_type)?; } else { - let value = read_scalar(&ctx, src, *scalar_type, alloc_budget)?; + let value = read_scalar(&ctx, src, *scalar_type, budget)?; target.set_item(&ext_type, value)?; } } @@ -901,7 +885,7 @@ fn read_extension<'py, R: JsonSource<'py>>( if src.peek()? == JsonKind::Null && !enum_.is_null_value { src.next_null()?; target.del_item(&ext_type)?; - } else if let Some(value) = read_enum(enum_, src, opts, alloc_budget)? { + } else if let Some(value) = read_enum(enum_, src, opts, budget)? { target.set_item(&ext_type, value)?; } } @@ -914,9 +898,9 @@ fn read_extension<'py, R: JsonSource<'py>>( src.next_null()?; target.del_item(&ext_type)?; } else { - alloc_budget.charge(inner.base_alloc_size)?; + budget.charge(inner.alloc_size)?; let value = inner.new_empty_message(py, msg_desc.get_python_type(py))?; - read_message(inner, &value, src, opts, depth + 1, alloc_budget)?; + read_message(inner, &value, src, opts, depth + 1, budget)?; target.set_item(&ext_type, value)?; } } @@ -930,25 +914,19 @@ fn read_extension<'py, R: JsonSource<'py>>( return Err(ctx.error( &format!( "expected list got {}", - read_json_value(src, alloc_budget)?.get_type() + read_json_value(src, budget)?.get_type() ), Exc::Type, )); } let element_value = FieldParserValue::from_desc_single(element); - alloc_budget.charge(budget::EMPTY_LIST_SIZE)?; + budget.charge(budget::EMPTY_LIST_SIZE)?; let list = PyList::empty(py); src.for_each_array_item(|src| { - if let Some(value) = read_container_item( - &ctx, - &element_value, - src, - opts, - depth, - false, - alloc_budget, - )? { - alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + if let Some(value) = + read_container_item(&ctx, &element_value, src, opts, depth, false, budget)? + { + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } Ok(()) @@ -963,11 +941,10 @@ fn read_extension<'py, R: JsonSource<'py>>( Ok(()) } -/// Materializes the next JSON value as a Python object, charging the -/// allocation budget for the resulting tree. +/// Materializes the next JSON value as a Python object. pub(crate) fn read_json_value<'py, R: JsonSource<'py>>( src: &mut R, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { let py = src.py(); match src.peek()? { @@ -977,31 +954,31 @@ pub(crate) fn read_json_value<'py, R: JsonSource<'py>>( } JsonKind::Bool => Ok(PyBool::new(py, src.next_bool()?).to_owned().into_any()), JsonKind::Number => { - alloc_budget.charge(budget::INT_SIZE)?; + budget.charge(budget::INT_SIZE)?; src.next_number() } JsonKind::String => { let value = src.next_py_str()?; - alloc_budget.charge(budget::STR_OVERHEAD + value.len()?)?; + budget.charge(budget::STR_OVERHEAD + value.len()?)?; Ok(value.into_any()) } JsonKind::Array => { - alloc_budget.charge(budget::EMPTY_LIST_SIZE)?; + budget.charge(budget::EMPTY_LIST_SIZE)?; let list = PyList::empty(py); src.for_each_array_item(|src| { - let value = read_json_value(src, alloc_budget)?; - alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + let value = read_json_value(src, budget)?; + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; Ok(()) })?; Ok(list.into_any()) } JsonKind::Object => { - alloc_budget.charge(budget::EMPTY_DICT_SIZE)?; + budget.charge(budget::EMPTY_DICT_SIZE)?; let dict = PyDict::new(py); src.for_each_object_key(|key, src| { - let value = read_json_value(src, alloc_budget)?; - alloc_budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; + let value = read_json_value(src, budget)?; + budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; dict.set_item(PyString::new(py, key), value)?; Ok(()) })?; diff --git a/packages/protobuf-py-ext/src/marshaler.rs b/packages/protobuf-py-ext/src/marshaler.rs index aca2ced..2e5c28b 100644 --- a/packages/protobuf-py-ext/src/marshaler.rs +++ b/packages/protobuf-py-ext/src/marshaler.rs @@ -78,7 +78,7 @@ pub(crate) struct MessageMarshalerInner { /// Approximate heap size of a freshly-initialized instance of this /// message type: the fixed instance size (all fields are slots) plus the /// empty containers created for repeated/map field defaults. - pub(crate) base_alloc_size: usize, + pub(crate) alloc_size: usize, /// The Python type of the message. pub(crate) python_type: Py, @@ -220,7 +220,7 @@ impl MessageMarshaler { members_by_name: members_by_name.unbind(), members, max_field_number, - base_alloc_size, + alloc_size: base_alloc_size, python_type: python_type.clone().unbind(), type_name, wkt, diff --git a/packages/protobuf-py-ext/src/nativemessage.rs b/packages/protobuf-py-ext/src/nativemessage.rs index df42674..0fe8d08 100644 --- a/packages/protobuf-py-ext/src/nativemessage.rs +++ b/packages/protobuf-py-ext/src/nativemessage.rs @@ -113,11 +113,11 @@ impl NativeMessage { let constants = Constants::get(py)?; let marshaler_any = cls.getattr(&constants.ext_marshaler)?; let marshaler = marshaler_any.cast::()?.get(); - let mut alloc_budget = Budget::new(allocation_limit); - alloc_budget.charge(marshaler.base_alloc_size)?; + let mut budget = Budget::new(allocation_limit); + budget.charge(marshaler.alloc_size)?; let message = marshaler.new_empty_message(py, cls)?; let slf = message.cast::()?; - marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut alloc_budget)?; + marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut budget)?; Ok(message) } @@ -162,8 +162,8 @@ impl NativeMessage { let constants = Constants::get(py)?; let marshaler_any = cls.getattr(&constants.ext_marshaler)?; let marshaler = marshaler_any.cast::()?.get().clone(); - let mut alloc_budget = Budget::new(allocation_limit); - alloc_budget.charge(marshaler.base_alloc_size)?; + let mut budget = Budget::new(allocation_limit); + budget.charge(marshaler.alloc_size)?; let message = marshaler.new_empty_message(py, cls)?; parse_json_into( py, @@ -172,7 +172,7 @@ impl NativeMessage { json, ignore_unknown_fields, registry, - &mut alloc_budget, + &mut budget, )?; Ok(message) } @@ -187,8 +187,8 @@ impl NativeMessage { allocation_limit: Option, ) -> PyResult<()> { let marshaler = NativeMessage::get_marshaler(slf)?; - let mut alloc_budget = Budget::new(allocation_limit); - alloc_budget.charge(marshaler.base_alloc_size)?; + let mut budget = Budget::new(allocation_limit); + budget.charge(marshaler.alloc_size)?; parse_json_into( py, &marshaler, @@ -196,7 +196,7 @@ impl NativeMessage { json, ignore_unknown_fields, registry, - &mut alloc_budget, + &mut budget, ) } @@ -234,14 +234,14 @@ impl NativeMessage { let constants = Constants::get(py)?; let marshaler_any = cls.getattr(&constants.ext_marshaler)?; let marshaler = marshaler_any.cast::()?.get().clone(); - let mut alloc_budget = Budget::new(allocation_limit); - alloc_budget.charge(marshaler.base_alloc_size)?; + let mut budget = Budget::new(allocation_limit); + budget.charge(marshaler.alloc_size)?; let message = marshaler.new_empty_message(py, cls)?; let opts = FromJsonOpts { ignore_unknown_fields, registry, }; - read_message_from_tree(py, &marshaler, &message, data, &opts, &mut alloc_budget)?; + read_message_from_tree(py, &marshaler, &message, data, &opts, &mut budget)?; Ok(message) } @@ -296,9 +296,9 @@ impl NativeMessage { ) -> PyResult<()> { let data = data.into_inner(); let marshaler = NativeMessage::get_marshaler(slf)?; - let mut alloc_budget = Budget::new(allocation_limit); - alloc_budget.charge(marshaler.base_alloc_size)?; - marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut alloc_budget) + let mut budget = Budget::new(allocation_limit); + budget.charge(marshaler.alloc_size)?; + marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut budget) } fn __deepcopy__<'py>( @@ -541,7 +541,7 @@ fn parse_json_into<'py>( json: &Bound<'py, PyAny>, ignore_unknown_fields: bool, registry: Option>, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let opts = FromJsonOpts { ignore_unknown_fields, @@ -554,20 +554,13 @@ fn parse_json_into<'py>( message, text.to_str()?.as_bytes(), &opts, - alloc_budget, + budget, ) } else if let Ok(bytes) = json.cast::() { - merge_from_json( - py, - marshaler, - message, - bytes.as_bytes(), - &opts, - alloc_budget, - ) + merge_from_json(py, marshaler, message, bytes.as_bytes(), &opts, budget) } else if let Ok(bytearray) = json.cast::() { let owned = bytearray.to_vec(); - merge_from_json(py, marshaler, message, &owned, &opts, alloc_budget) + merge_from_json(py, marshaler, message, &owned, &opts, budget) } else { Err(PyTypeError::new_err(format!( "json must be str, bytes, or bytearray, got {}", diff --git a/packages/protobuf-py-ext/src/parser.rs b/packages/protobuf-py-ext/src/parser.rs index e48db6a..e75ed27 100644 --- a/packages/protobuf-py-ext/src/parser.rs +++ b/packages/protobuf-py-ext/src/parser.rs @@ -251,7 +251,7 @@ impl FieldParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { match &self.type_ { ParserFieldType::Singular { @@ -267,7 +267,7 @@ impl FieldParser { depth, oneof_attr.as_ref(), *requires_presence, - alloc_budget, + budget, )?, ParserFieldType::List { unpacked_wire_type, @@ -282,7 +282,7 @@ impl FieldParser { depth, *unpacked_wire_type, *packable, - alloc_budget, + budget, )?, ParserFieldType::Map { key_type, @@ -300,7 +300,7 @@ impl FieldParser { value_parser, key_default_value, value_default_value, - alloc_budget, + budget, )?, } @@ -318,7 +318,7 @@ impl FieldParser { depth: usize, oneof_attr: Option<&AttributeAccess>, requires_presence: bool, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let value = self.read_single_value( py, @@ -328,12 +328,12 @@ impl FieldParser { buffer, opts, depth, - alloc_budget, + budget, )?; match value { SingleValue::Parsed(value) => { if oneof_attr.is_some() { - alloc_budget.charge(budget::ONEOF_SIZE)?; + budget.charge(budget::ONEOF_SIZE)?; } self.assign_singular(py, message, &value, oneof_attr, requires_presence)?; } @@ -342,7 +342,7 @@ impl FieldParser { let mut field = BytesMut::new(); encode_varint(tag as u64, &mut field); encode_varint(number as u64, &mut field); - write_unknown_field(py, message, tag >> 3, &field, alloc_budget)?; + write_unknown_field(py, message, tag >> 3, &field, budget)?; } } } @@ -382,7 +382,7 @@ impl FieldParser { depth: usize, unpacked_wire_type: WireType, packable: bool, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let python_list = self.attr.get(py, message)?; let list = python_list.cast::()?; @@ -399,11 +399,11 @@ impl FieldParser { &mut list_buffer, opts, depth, - alloc_budget, + budget, )?; match value { SingleValue::Parsed(value) => { - alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } SingleValue::UnknownEnumValue(number) => { @@ -413,25 +413,17 @@ impl FieldParser { let tag = (tag & !0b111) | (unpacked_wire_type as u32); encode_varint(tag as u64, &mut field); encode_varint(number as u64, &mut field); - write_unknown_field(py, message, tag >> 3, &field, alloc_budget)?; + write_unknown_field(py, message, tag >> 3, &field, budget)?; } } } } } else { - let value = self.read_single_value( - py, - tag, - wire_type, - None, - buffer, - opts, - depth, - alloc_budget, - )?; + let value = + self.read_single_value(py, tag, wire_type, None, buffer, opts, depth, budget)?; match value { SingleValue::Parsed(value) => { - alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } SingleValue::UnknownEnumValue(number) => { @@ -439,7 +431,7 @@ impl FieldParser { let mut field = BytesMut::new(); encode_varint(tag as u64, &mut field); encode_varint(number as u64, &mut field); - write_unknown_field(py, message, tag >> 3, &field, alloc_budget)?; + write_unknown_field(py, message, tag >> 3, &field, budget)?; } } } @@ -459,7 +451,7 @@ impl FieldParser { value_parser: &FieldParser, key_default_value: &Py, value_default_value: &Py, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; @@ -481,10 +473,10 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, - alloc_budget, + budget, ); } - key = Some(read_scalar(py, key_type, &mut entry_buffer, alloc_budget)?); + key = Some(read_scalar(py, key_type, &mut entry_buffer, budget)?); } 2 => { if value_parser.wire_type != wire_type { @@ -494,7 +486,7 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, - alloc_budget, + budget, ); } value = Some(value_parser.read_single_value( @@ -505,7 +497,7 @@ impl FieldParser { &mut entry_buffer, opts, depth, - alloc_budget, + budget, )?); } _ => { @@ -519,7 +511,7 @@ impl FieldParser { value } else if let FieldParserValue::Message { message, .. } = &value_parser.value { // For message values, the default is a new instance of the message, not None. - alloc_budget.charge(message.get_marshaler(py)?.base_alloc_size)?; + budget.charge(message.get_marshaler(py)?.alloc_size)?; SingleValue::Parsed(message.get_python_type(py).call0()?) } else { SingleValue::Parsed(value_default_value.bind(py).clone()) @@ -532,10 +524,10 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, - alloc_budget, + budget, ); }; - alloc_budget.charge(budget::DICT_ENTRY_SIZE)?; + budget.charge(budget::DICT_ENTRY_SIZE)?; self.assign_map_entry(py, message, key, value)?; Ok(()) } @@ -560,14 +552,14 @@ impl FieldParser { tag: u32, entry_bytes: &[u8], opts: FromBinaryOpts, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { if !opts.ignore_unknown_fields { let mut field_bytes = BytesMut::new(); encode_varint(tag as u64, &mut field_bytes); encode_varint(entry_bytes.len() as u64, &mut field_bytes); field_bytes.extend_from_slice(entry_bytes); - write_unknown_field(py, message, tag >> 3, &field_bytes, alloc_budget)?; + write_unknown_field(py, message, tag >> 3, &field_bytes, budget)?; } Ok(()) } @@ -582,12 +574,10 @@ impl FieldParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { let value = match &self.value { - FieldParserValue::Scalar(scalar_type) => { - read_scalar(py, *scalar_type, buffer, alloc_budget)? - } + FieldParserValue::Scalar(scalar_type) => read_scalar(py, *scalar_type, buffer, budget)?, FieldParserValue::Message { message: message_desc, .. @@ -615,7 +605,7 @@ impl FieldParser { { existing.cast_into::()? } else { - alloc_budget.charge(marshaler.base_alloc_size)?; + budget.charge(marshaler.alloc_size)?; marshaler.new_empty_message(py, parser.inner.python_type.bind(py))? }; parser.merge_from_binary( @@ -624,7 +614,7 @@ impl FieldParser { &mut message_buffer, opts, depth + 1, - alloc_budget, + budget, )?; message_instance.into_any() } @@ -634,9 +624,7 @@ impl FieldParser { if let Some(value) = value { value.bind(py).clone() } else if enum_.open { - // Unknown open enum values allocate a new int-subclass - // instance of the enum type. - alloc_budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; + budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; enum_.py_type.bind(py).call1((number,))? } else { return Ok(SingleValue::UnknownEnumValue(number)); @@ -734,7 +722,7 @@ impl MessageParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { check_parse_recursion_depth(depth)?; while buffer.has_remaining() { @@ -748,16 +736,7 @@ impl MessageParser { if let Some(field) = self.inner.fields.get(field_number) && field.wire_type_matches(wire_type) { - field.read_field( - py, - message, - tag, - wire_type, - buffer, - opts, - depth, - alloc_budget, - )?; + field.read_field(py, message, tag, wire_type, buffer, opts, depth, budget)?; } else { skip_field_with_wire_type(field_number, wire_type, buffer, depth + 1)?; if !opts.ignore_unknown_fields { @@ -767,7 +746,7 @@ impl MessageParser { message, field_number, &checkpoint[..field_len], - alloc_budget, + budget, )?; } } @@ -776,16 +755,15 @@ impl MessageParser { } } -/// Reads a scalar value from the wire buffer, charging the allocation budget -/// for the resulting Python object. +/// Reads a scalar value from the wire buffer. fn read_scalar<'py>( py: Python<'py>, s: ScalarType, buffer: &mut Bytes, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult> { match s { - ScalarType::Double | ScalarType::Float => alloc_budget.charge(budget::FLOAT_SIZE)?, + ScalarType::Double | ScalarType::Float => budget.charge(budget::FLOAT_SIZE)?, // Small ints are interned and larger ones vary a little in size, but // a flat charge is close enough. ScalarType::Int64 @@ -797,7 +775,7 @@ fn read_scalar<'py>( | ScalarType::Sfixed32 | ScalarType::Sfixed64 | ScalarType::Sint32 - | ScalarType::Sint64 => alloc_budget.charge(budget::INT_SIZE)?, + | ScalarType::Sint64 => budget.charge(budget::INT_SIZE)?, // Bools are shared singletons; strings and bytes are charged below // once the length is known. ScalarType::Bool | ScalarType::String | ScalarType::Bytes => {} @@ -835,14 +813,14 @@ fn read_scalar<'py>( ScalarType::String => { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; - alloc_budget.charge(budget::STR_OVERHEAD + len)?; + budget.charge(budget::STR_OVERHEAD + len)?; let bytes = buffer.split_to(len); PyString::from_bytes(py, &bytes)?.into_any() } ScalarType::Bytes => { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; - alloc_budget.charge(budget::BYTES_OVERHEAD + len)?; + budget.charge(budget::BYTES_OVERHEAD + len)?; let bytes = buffer.split_to(len); PyBytes::new(py, &bytes).into_any() } @@ -966,11 +944,11 @@ fn write_unknown_field( message: &Bound<'_, NativeMessage>, field_number: u32, field_bytes: &[u8], - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { // The bytes copy plus the list slot holding it. The dict/list created for // the first unknown field of a number are not charged for simplicity. - alloc_budget.charge(budget::BYTES_OVERHEAD + field_bytes.len() + budget::LIST_SLOT_SIZE)?; + budget.charge(budget::BYTES_OVERHEAD + field_bytes.len() + budget::LIST_SLOT_SIZE)?; let unknown_fields_unbound = message.get().get_or_init_unknown_fields(py); let unknown_fields = unknown_fields_unbound.bind(py); let field_list = if let Ok(list) = unknown_fields.get_item(field_number) { diff --git a/packages/protobuf-py-ext/src/wkt_registry.rs b/packages/protobuf-py-ext/src/wkt_registry.rs index 1ec0478..a017128 100644 --- a/packages/protobuf-py-ext/src/wkt_registry.rs +++ b/packages/protobuf-py-ext/src/wkt_registry.rs @@ -59,13 +59,13 @@ impl WktTimestamp { src: &mut R, _opts: &FromJsonOpts, _depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); - let (secs, nanos) = with_wkt_str(marshaler, src, alloc_budget, |text, _| { + let (secs, nanos) = with_wkt_str(marshaler, src, budget, |text, _| { parse_timestamp(&marshaler.type_name, text) })?; - alloc_budget.charge(2 * budget::INT_SIZE)?; + budget.charge(2 * budget::INT_SIZE)?; self.seconds .set(message.as_any(), PyInt::new(py, secs).as_any())?; self.nanos @@ -102,13 +102,13 @@ impl WktDuration { src: &mut R, _opts: &FromJsonOpts, _depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); - let (secs, nanos) = with_wkt_str(marshaler, src, alloc_budget, |text, _| { + let (secs, nanos) = with_wkt_str(marshaler, src, budget, |text, _| { parse_duration(&marshaler.type_name, text) })?; - alloc_budget.charge(2 * budget::INT_SIZE)?; + budget.charge(2 * budget::INT_SIZE)?; self.seconds .set(message.as_any(), PyInt::new(py, secs).as_any())?; self.nanos @@ -162,6 +162,8 @@ impl WktAny { let value = self.value.get(py, message.as_any())?.extract::()?; let inner_msg = inner_marshaler.new_empty_message(py, &inner_type)?; + // While we go through parsing to serialize a any to JSON, we don't need to track budget since this is + // not a real parse operation. inner_marshaler.merge_from_binary(py, &inner_msg, value, false, &mut Budget::new(None))?; sink.begin_object()?; @@ -188,13 +190,13 @@ impl WktAny { src: &mut R, opts: &FromJsonOpts, _depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let message_type_name = &marshaler.type_name; // For a string input, this will eagerly parse out a whole dictionary since // we need to first find `@type` before parsing. - let tree = read_json_value(src, alloc_budget)?; + let tree = read_json_value(src, budget)?; let Ok(dict) = tree.cast::() else { return Err(PyTypeError::new_err(format!( "cannot decode {message_type_name} from JSON: {}", @@ -236,7 +238,7 @@ impl WktAny { .getattr(&marshaler.constants.ext_marshaler)? .cast_into::()?; let inner_marshaler = inner_marshaler.get(); - alloc_budget.charge(inner_marshaler.base_alloc_size)?; + budget.charge(inner_marshaler.alloc_size)?; let inner_msg = inner_marshaler.new_empty_message(py, &inner_type)?; let is_wkt = inner_marshaler.wkt.is_some(); @@ -245,19 +247,19 @@ impl WktAny { .get_item("value")? .unwrap_or_else(|| py.None().into_bound(py)); let mut sub = PyTreeSource::new(py, value); - read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1, alloc_budget)?; + read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1, budget)?; } else { let copy = dict.copy()?; copy.del_item("@type")?; let mut sub = PyTreeSource::new(py, copy.into_any()); - read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1, alloc_budget)?; + read_message(inner_marshaler, &inner_msg, &mut sub, opts, 1, budget)?; } // Any.pack let packed_url = format!("type.googleapis.com/{}", inner_marshaler.type_name); let packed_value = inner_marshaler.to_binary(py, &inner_msg, true)?; - alloc_budget.charge(budget::STR_OVERHEAD + packed_url.len())?; - alloc_budget.charge(budget::BYTES_OVERHEAD + packed_value.as_bytes().len())?; + budget.charge(budget::STR_OVERHEAD + packed_url.len())?; + budget.charge(budget::BYTES_OVERHEAD + packed_value.as_bytes().len())?; self.type_url .set(message.as_any(), &PyString::new(py, &packed_url).into_any())?; self.value.set(message.as_any(), packed_value.as_any())?; @@ -309,14 +311,14 @@ impl WktFieldMask { src: &mut R, _opts: &FromJsonOpts, _depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let paths = self .paths .get(py, message.as_any())? .cast_into::()?; - with_wkt_str(marshaler, src, alloc_budget, |text, alloc_budget| { + with_wkt_str(marshaler, src, budget, |text, budget| { if text.is_empty() { return Ok(()); } @@ -327,7 +329,7 @@ impl WktFieldMask { marshaler.type_name ))); } - alloc_budget.charge(budget::STR_OVERHEAD + part.len() + budget::LIST_SLOT_SIZE)?; + budget.charge(budget::STR_OVERHEAD + part.len() + budget::LIST_SLOT_SIZE)?; paths.append(buffa_wkt::camel_to_snake(part))?; } Ok(()) @@ -371,11 +373,11 @@ impl WktStruct { src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? != JsonKind::Object { - let json = read_json_value(src, alloc_budget)?; + let json = read_json_value(src, budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, @@ -389,18 +391,11 @@ impl WktStruct { .cast_into::()?; // Duplicate keys use last-in-wins semantics (per the ProtoJSON spec). src.for_each_object_key(|key, src| { - alloc_budget.charge(value_marshaler.base_alloc_size)?; + budget.charge(value_marshaler.alloc_size)?; let value_msg = value_marshaler.new_empty_message(py, self.value.get_python_type(py))?; - read_message( - value_marshaler, - &value_msg, - src, - opts, - depth + 1, - alloc_budget, - )?; - alloc_budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; + read_message(value_marshaler, &value_msg, src, opts, depth + 1, budget)?; + budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; dict.set_item(key, value_msg)?; Ok(()) })?; @@ -443,11 +438,11 @@ impl WktListValue { src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? != JsonKind::Array { - let json = read_json_value(src, alloc_budget)?; + let json = read_json_value(src, budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, @@ -460,18 +455,11 @@ impl WktListValue { .get(py, message.as_any())? .cast_into::()?; src.for_each_array_item(|src| { - alloc_budget.charge(element_marshaler.base_alloc_size)?; + budget.charge(element_marshaler.alloc_size)?; let value_msg = element_marshaler.new_empty_message(py, self.element.get_python_type(py))?; - read_message( - element_marshaler, - &value_msg, - src, - opts, - depth + 1, - alloc_budget, - )?; - alloc_budget.charge(budget::LIST_SLOT_SIZE)?; + read_message(element_marshaler, &value_msg, src, opts, depth + 1, budget)?; + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value_msg)?; Ok(()) })?; @@ -541,10 +529,10 @@ impl WktValue { src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); - alloc_budget.charge(budget::ONEOF_SIZE)?; + budget.charge(budget::ONEOF_SIZE)?; let oneof = match src.peek()? { JsonKind::Null => { src.next_null()?; @@ -556,7 +544,7 @@ impl WktValue { ), JsonKind::Number => { let number = src.next_float()?; - alloc_budget.charge(budget::FLOAT_SIZE)?; + budget.charge(budget::FLOAT_SIZE)?; Oneof::new( self.number_name.bind(py), &PyFloat::new(py, number).into_any(), @@ -564,23 +552,23 @@ impl WktValue { } JsonKind::String => { let string = src.next_py_str()?; - alloc_budget.charge(budget::STR_OVERHEAD + string.len()?)?; + budget.charge(budget::STR_OVERHEAD + string.len()?)?; Oneof::new(self.string_name.bind(py), &string) } JsonKind::Array => { let desc = &self.list_message; let inner = desc.get_marshaler(py)?; - alloc_budget.charge(inner.base_alloc_size)?; + budget.charge(inner.alloc_size)?; let list_msg = inner.new_empty_message(py, desc.get_python_type(py))?; - read_message(inner, &list_msg, src, opts, depth + 1, alloc_budget)?; + read_message(inner, &list_msg, src, opts, depth + 1, budget)?; Oneof::new(self.list_name.bind(py), &list_msg.into_any()) } JsonKind::Object => { let desc = &self.struct_message; let inner = desc.get_marshaler(py)?; - alloc_budget.charge(inner.base_alloc_size)?; + budget.charge(inner.alloc_size)?; let struct_msg = inner.new_empty_message(py, desc.get_python_type(py))?; - read_message(inner, &struct_msg, src, opts, depth + 1, alloc_budget)?; + read_message(inner, &struct_msg, src, opts, depth + 1, budget)?; Oneof::new(self.struct_name.bind(py), &struct_msg.into_any()) } }; @@ -618,7 +606,7 @@ impl WktWrapper { src: &mut R, _opts: &FromJsonOpts, _depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? == JsonKind::Null { @@ -629,7 +617,7 @@ impl WktWrapper { } let name = self.name.bind(py); let ctx = FieldContext::Field { marshaler, name }; - let value = read_scalar(&ctx, src, self.scalar, alloc_budget)?; + let value = read_scalar(&ctx, src, self.scalar, budget)?; self.field.set(message.as_any(), &value) } } @@ -677,23 +665,17 @@ impl WktKind { src: &mut R, opts: &FromJsonOpts, depth: usize, - alloc_budget: &mut Budget, + budget: &mut Budget, ) -> PyResult<()> { match self { - WktKind::Timestamp(w) => { - w.read_json(marshaler, message, src, opts, depth, alloc_budget) - } - WktKind::Duration(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), - WktKind::Any(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), - WktKind::FieldMask(w) => { - w.read_json(marshaler, message, src, opts, depth, alloc_budget) - } - WktKind::Struct(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), - WktKind::ListValue(w) => { - w.read_json(marshaler, message, src, opts, depth, alloc_budget) - } - WktKind::Value(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), - WktKind::Wrapper(w) => w.read_json(marshaler, message, src, opts, depth, alloc_budget), + WktKind::Timestamp(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::Duration(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::Any(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::FieldMask(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::Struct(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::ListValue(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::Value(w) => w.read_json(marshaler, message, src, opts, depth, budget), + WktKind::Wrapper(w) => w.read_json(marshaler, message, src, opts, depth, budget), } } @@ -932,18 +914,18 @@ fn match_wrapper(fields: &[DescField], by_name: &HashMap) -> Opti fn with_wkt_str<'py, S: JsonSource<'py>, R>( marshaler: &MessageMarshaler, src: &mut S, - alloc_budget: &mut Budget, + budget: &mut Budget, f: impl FnOnce(&str, &mut Budget) -> PyResult, ) -> PyResult { if src.peek()? != JsonKind::String { - let value = read_json_value(src, alloc_budget)?; + let value = read_json_value(src, budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, value.str()? ))); } - src.with_next_str(|text| f(text, alloc_budget)) + src.with_next_str(|text| f(text, budget)) } fn type_url_to_name(url: &str) -> PyResult<&str> { diff --git a/src/protobuf/_budget.py b/src/protobuf/_budget.py index 2ed0ebb..b244987 100644 --- a/src/protobuf/_budget.py +++ b/src/protobuf/_budget.py @@ -27,8 +27,7 @@ # Approximate sizes of CPython heap allocations, measured with `sys.getsizeof` # on 64-bit CPython 3.14. The budget guards against unbounded allocation from # malicious payloads rather than providing exact accounting, so small -# inaccuracies across versions and builds are fine. These values mirror the -# constants in the native extension's budget.rs. +# inaccuracies across versions and builds are fine. GC_HEAD_SIZE = 16 """GC header allocated in front of every GC-tracked object.""" @@ -41,9 +40,9 @@ BYTES_OVERHEAD = 33 """Header of a bytes object.""" EMPTY_LIST_SIZE = 56 -"""An empty list, as created for repeated field defaults.""" +"""An empty list.""" EMPTY_DICT_SIZE = 64 -"""An empty dict, as created for map field defaults.""" +"""An empty dict.""" LIST_SLOT_SIZE = 8 """One appended list element: an 8-byte pointer slot.""" DICT_ENTRY_SIZE = 40 @@ -71,9 +70,7 @@ def _base_alloc_size(message_type: type[Message]) -> int: class Budget: """Tracks approximate allocations while parsing a message. - Raises an error once the configured limit is exceeded. Charges are made - before the corresponding allocation where practical so the limit bounds - the actual peak. + Raises an error once the configured limit is exceeded. """ __slots__ = ("current", "max") diff --git a/src/protobuf/_from_binary.py b/src/protobuf/_from_binary.py index ce28bad..5a9e1b1 100644 --- a/src/protobuf/_from_binary.py +++ b/src/protobuf/_from_binary.py @@ -316,7 +316,6 @@ def read_enum( if not desc_enum._values_by_number.get(value): if not desc_enum.open: return value - # Unknown open enum values allocate a new int-subclass instance. if budget is not None: budget.charge(INT_SIZE + GC_HEAD_SIZE) return desc_enum.type(value) @@ -454,9 +453,7 @@ def merge_from_binary( data: Serialized binary protobuf data. Must not be mutated during parsing. ignore_unknown_fields: If `True`, unknown fields in the binary data are silently discarded. allocation_limit: If set, the approximate number of bytes of Python - objects the parse may allocate before raising a ValueError. Guards - against malicious payloads that expand into unexpectedly large - messages. + objects the parse may allocate before raising a ValueError. """ message._merge_from_binary( data, ignore_unknown_fields, allocation_limit=allocation_limit diff --git a/src/protobuf/_from_json.py b/src/protobuf/_from_json.py index 66cd853..887f5e1 100644 --- a/src/protobuf/_from_json.py +++ b/src/protobuf/_from_json.py @@ -118,9 +118,7 @@ def merge_from_json( registry: Required to read google.protobuf.Any and extensions from JSON format. allocation_limit: If set, the approximate number of bytes of Python - objects the parse may allocate before raising a ValueError. Guards - against malicious payloads that expand into unexpectedly large - messages. + objects the parse may allocate before raising a ValueError. Raises: json.JSONDecodeError: If json_source is not valid JSON. @@ -463,8 +461,7 @@ def _read_enum( return desc.type(value.number) if ignore_unknown_fields: return None - # Succeeds for open enum (allocating a new int-subclass instance), - # raises an error for closed. + # Succeeds for open enum, raises an error for closed if budget is not None: budget.charge(INT_SIZE + GC_HEAD_SIZE) return desc.type(json) From 010fb8f1c88ead5020dbdadd1724a81aeef153a3 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 10:59:59 +0900 Subject: [PATCH 03/11] cleanup --- src/protobuf/_budget.py | 5 +- src/protobuf/_from_binary.py | 81 ++++++++++++------------------ src/protobuf/_from_json.py | 94 +++++++++++++---------------------- src/protobuf/_message.py | 33 ++++-------- src/protobuf/_unknown.py | 4 +- src/protobuf/_wkt_registry.py | 23 +++------ 6 files changed, 89 insertions(+), 151 deletions(-) diff --git a/src/protobuf/_budget.py b/src/protobuf/_budget.py index b244987..d8a6a6e 100644 --- a/src/protobuf/_budget.py +++ b/src/protobuf/_budget.py @@ -14,6 +14,7 @@ from __future__ import annotations +import sys from functools import cache from typing import TYPE_CHECKING, cast @@ -75,9 +76,9 @@ class Budget: __slots__ = ("current", "max") - def __init__(self, limit: int) -> None: + def __init__(self, limit: int | None = None) -> None: self.current = 0 - self.max = limit + self.max = sys.maxsize if limit is None else limit def charge(self, amount: int) -> None: self.current += amount diff --git a/src/protobuf/_from_binary.py b/src/protobuf/_from_binary.py index 5a9e1b1..db4cf96 100644 --- a/src/protobuf/_from_binary.py +++ b/src/protobuf/_from_binary.py @@ -14,7 +14,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, TypeVar from typing_extensions import Buffer, assert_never @@ -56,12 +56,12 @@ class FromBinaryOptions: Args: ignore_unknown_fields: If `True`, unknown fields are ignored instead of being added to the message. - budget: If set, tracks approximate allocations during the parse and - raises once the configured limit is exceeded. + budget: Tracks approximate allocations during the parse and raises + once the configured limit is exceeded. Unlimited by default. """ ignore_unknown_fields: bool = False - budget: Budget | None = None + budget: Budget = field(default_factory=Budget) # Dispatch table for reading scalar values. CPython currently does not generate @@ -91,14 +91,11 @@ class FromBinaryOptions: ) -def read_scalar( - scalar_type: ScalarType, reader: BinaryReader, budget: Budget | None = None -) -> Any: +def read_scalar(scalar_type: ScalarType, reader: BinaryReader, budget: Budget) -> Any: reader_method = _SCALAR_READERS[scalar_type.value] assert reader_method is not None # noqa: S101 value = reader_method(reader) - if budget is not None: - budget.charge_scalar(scalar_type, value) + budget.charge_scalar(scalar_type, value) return value @@ -145,10 +142,10 @@ def read_message( field_raw = reader.skip(tag.wire_type, depth + 1, field_number=tag.number) if not opts.ignore_unknown_fields: key_raw = _encode_varint((tag.number << 3) | tag.wire_type) - if (budget := opts.budget) is not None: - budget.charge( - BYTES_OVERHEAD + len(key_raw) + len(field_raw) + LIST_SLOT_SIZE - ) + budget = opts.budget + budget.charge( + BYTES_OVERHEAD + len(key_raw) + len(field_raw) + LIST_SLOT_SIZE + ) message._get_or_init_unknown_fields().setdefault(tag.number, []).append( key_raw + bytes(field_raw) ) @@ -158,7 +155,7 @@ def read_message( match field_value := desc_field.value: case DescFieldValueScalar(): value = read_scalar(field_value.scalar, reader, budget) - if budget is not None and field_value.oneof is not None: + if field_value.oneof is not None: budget.charge(ONEOF_SIZE) message._set_member(desc_field, value) case DescFieldValueMessage( @@ -166,10 +163,9 @@ def read_message( ): existing: Message | None = message._get_member(desc_field) if existing is None: - if budget is not None: - budget.charge_message(desc_nested_message.type) - if field_value.oneof is not None: - budget.charge(ONEOF_SIZE) + budget.charge_message(desc_nested_message.type) + if field_value.oneof is not None: + budget.charge(ONEOF_SIZE) existing = desc_nested_message.type() message._set_member(desc_field, existing) if delimited_encoding: @@ -183,7 +179,7 @@ def read_message( case DescFieldValueEnum(): value = read_enum(field_value.enum, reader, budget) if isinstance(value, Enum): - if budget is not None and field_value.oneof is not None: + if field_value.oneof is not None: budget.charge(ONEOF_SIZE) message._set_member(desc_field, value) elif not opts.ignore_unknown_fields: @@ -205,8 +201,7 @@ def read_message( ) if entry: key, value = entry - if budget is not None: - budget.charge(DICT_ENTRY_SIZE) + budget.charge(DICT_ENTRY_SIZE) message._get_member(desc_field)[key] = value case _: assert_never(desc_field) @@ -237,10 +232,10 @@ def read_list( field_bytes = reader.skip(wire_type, depth + 1, field_number=field_number) if not opts.ignore_unknown_fields and message: key_raw = _encode_varint((field_number << 3) | wire_type) - if (budget := opts.budget) is not None: - budget.charge( - BYTES_OVERHEAD + len(key_raw) + len(field_bytes) + LIST_SLOT_SIZE - ) + budget = opts.budget + budget.charge( + BYTES_OVERHEAD + len(key_raw) + len(field_bytes) + LIST_SLOT_SIZE + ) message._get_or_init_unknown_fields().setdefault(field_number, []).append( key_raw + bytes(field_bytes) ) @@ -251,8 +246,7 @@ def read_list( case ScalarType(): value = read_scalar(element_type, reader, budget) case DescMessage(): - if budget is not None: - budget.charge_message(element_type.type) + budget.charge_message(element_type.type) if field_value.delimited_encoding: value = read_message( element_type.type(), @@ -273,8 +267,7 @@ def read_list( return case _: assert_never(element_type) - if budget is not None: - budget.charge(LIST_SLOT_SIZE) + budget.charge(LIST_SLOT_SIZE) list_.append(value) @@ -293,14 +286,12 @@ def _read_packed_list( match element_type: case ScalarType(): value = read_scalar(element_type, reader, budget) - if budget is not None: - budget.charge(LIST_SLOT_SIZE) + budget.charge(LIST_SLOT_SIZE) list_.append(value) case DescEnum(): value = read_enum(element_type, reader, budget) if isinstance(value, Enum): - if budget is not None: - budget.charge(LIST_SLOT_SIZE) + budget.charge(LIST_SLOT_SIZE) list_.append(value) elif not opts.ignore_unknown_fields: # Even for packed fields we write unknown enum values as unpacked. @@ -309,20 +300,17 @@ def _read_packed_list( assert_never(element_type) -def read_enum( - desc_enum: DescEnum, reader: BinaryReader, budget: Budget | None = None -) -> Enum | int: +def read_enum(desc_enum: DescEnum, reader: BinaryReader, budget: Budget) -> Enum | int: value = reader.int32() if not desc_enum._values_by_number.get(value): if not desc_enum.open: return value - if budget is not None: - budget.charge(INT_SIZE + GC_HEAD_SIZE) + budget.charge(INT_SIZE + GC_HEAD_SIZE) return desc_enum.type(value) def _write_unknown_enum_field( - message: Message | None, field_number: int, value: int, budget: Budget | None = None + message: Message | None, field_number: int, value: int, budget: Budget ) -> None: if message is None: return @@ -330,8 +318,7 @@ def _write_unknown_enum_field( writer.tag(field_number, WireType.VARINT) writer.int32(value) field_bytes = writer.finish() - if budget is not None: - budget.charge(BYTES_OVERHEAD + len(field_bytes) + LIST_SLOT_SIZE) + budget.charge(BYTES_OVERHEAD + len(field_bytes) + LIST_SLOT_SIZE) message._get_or_init_unknown_fields().setdefault(field_number, []).append( field_bytes ) @@ -379,8 +366,7 @@ def read_map_entry( ) return None case DescMessage(): - if budget is not None: - budget.charge_message(field_value.value.type) + budget.charge_message(field_value.value.type) value = field_value.value.type() read_message(value, reader, opts, depth + 1, length=reader.varint()) case _: @@ -397,8 +383,7 @@ def read_map_entry( case DescEnum() as desc_enum: value = desc_enum.type(desc_enum.values[0].number) case DescMessage(): - if budget is not None: - budget.charge_message(field_value.value.type) + budget.charge_message(field_value.value.type) value = field_value.value.type() case _: assert_never(field_value.value) @@ -422,10 +407,8 @@ def _read_unknown_map_entry( ) if not opts.ignore_unknown_fields and message: key_raw = _encode_varint((field_number << 3) | WireType.LENGTH_DELIMITED) - if (budget := opts.budget) is not None: - budget.charge( - BYTES_OVERHEAD + len(key_raw) + len(entry_bytes) + LIST_SLOT_SIZE - ) + budget = opts.budget + budget.charge(BYTES_OVERHEAD + len(key_raw) + len(entry_bytes) + LIST_SLOT_SIZE) message._get_or_init_unknown_fields().setdefault(field_number, []).append( key_raw + bytes(entry_bytes) ) diff --git a/src/protobuf/_from_json.py b/src/protobuf/_from_json.py index 887f5e1..1ce7fb3 100644 --- a/src/protobuf/_from_json.py +++ b/src/protobuf/_from_json.py @@ -16,7 +16,7 @@ import math from base64 import b64decode -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Literal, TypeVar, cast from typing_extensions import assert_never @@ -81,8 +81,8 @@ class FromJsonOptions: registry: Registry | None depth: int = 0 """Current message nesting depth. Mutated during a parse.""" - budget: Budget | None = None - """Optional allocation budget shared across the whole parse.""" + budget: Budget = field(default_factory=Budget) + """Allocation budget shared across the whole parse. Unlimited by default.""" def _enter_message(opts: FromJsonOptions) -> FromJsonOptions: @@ -232,7 +232,7 @@ def _read_scalar_extension( ext: DescExtension, field_value: DescFieldValueScalar, json: JsonValue, - budget: Budget | None, + budget: Budget, ) -> None: if json is None: del msg[ext.type] @@ -267,8 +267,8 @@ def _read_message_extension( if _is_resetting_null(field_value.message, json): del msg[ext.type] else: - if (budget := opts.budget) is not None: - budget.charge_message(field_value.message.type) + budget = opts.budget + budget.charge_message(field_value.message.type) value = field_value.message.type() _read_message(value, json, opts) msg[ext.type] = value @@ -286,14 +286,12 @@ def _read_list_extension( if not isinstance(json, list): raise _field_error(ext, f"expected list got {type(json)}", TypeError) budget = opts.budget - if budget is not None: - budget.charge(EMPTY_LIST_SIZE) + budget.charge(EMPTY_LIST_SIZE) values = [] for element in json: v = _read_container_item(ext, field_value.element, element, opts) if v is not None: - if budget is not None: - budget.charge(LIST_SLOT_SIZE) + budget.charge(LIST_SLOT_SIZE) values.append(v) msg[ext.type] = values @@ -303,13 +301,13 @@ def _read_scalar_field( field: DescField, field_value: DescFieldValueScalar, json: JsonValue, - budget: Budget | None, + budget: Budget, ) -> None: if json is None: msg._del_member(field) return value = _read_scalar(field, field_value.scalar, json, budget) - if budget is not None and field_value.oneof is not None: + if field_value.oneof is not None: budget.charge(ONEOF_SIZE) msg._set_member(field, value) @@ -326,8 +324,8 @@ def _read_enum_field( return value = _read_enum(field_value.enum, json, opts.ignore_unknown_fields, opts.budget) if value is not None: - if (budget := opts.budget) is not None and field_value.oneof is not None: - budget.charge(ONEOF_SIZE) + if field_value.oneof is not None: + opts.budget.charge(ONEOF_SIZE) msg._set_member(field, value) @@ -346,8 +344,7 @@ def _read_list_field( for element in json: value = _read_container_item(field, field_value.element, element, opts) if value is not None: - if budget is not None: - budget.charge(LIST_SLOT_SIZE) + budget.charge(LIST_SLOT_SIZE) list_.append(value) @@ -367,8 +364,7 @@ def _read_map_field( key = _read_map_key(field, field_value, json_key, budget) value = _read_container_item(field, field_value.value, json_value, opts) if value is not None: - if budget is not None: - budget.charge(DICT_ENTRY_SIZE) + budget.charge(DICT_ENTRY_SIZE) dict_[key] = value @@ -386,20 +382,16 @@ def _read_message_field( if msg._contains_member(field): value = msg._get_member(field) else: - if budget is not None: - budget.charge_message(field_value.message.type) + budget.charge_message(field_value.message.type) value = field_value.message.type() _read_message(value, json, opts) - if budget is not None and field_value.oneof is not None: + if field_value.oneof is not None: budget.charge(ONEOF_SIZE) msg._set_member(field, value) def _read_map_key( - field: DescField, - field_value: DescFieldValueMap, - json: JsonValue, - budget: Budget | None, + field: DescField, field_value: DescFieldValueMap, json: JsonValue, budget: Budget ) -> bool | int | str: match field_value.key: case ScalarType.BOOL: @@ -410,8 +402,7 @@ def _read_map_key( raise _field_error(field, f"unexpected bool map key value {json}") case ScalarType.STRING: key = _read_string(field, json) - if budget is not None: - budget.charge(STR_OVERHEAD + len(key)) + budget.charge(STR_OVERHEAD + len(key)) return key case ( ScalarType.DOUBLE | ScalarType.FLOAT | ScalarType.BYTES @@ -419,8 +410,7 @@ def _read_map_key( msg = f"invalid map key type: {field_value.key}" raise AssertionError(msg) case _: - if budget is not None: - budget.charge(INT_SIZE) + budget.charge(INT_SIZE) return _read_int(field, field_value.key, json) @@ -433,8 +423,8 @@ def _read_container_item( if isinstance(element, ScalarType) and json is not None: return _read_scalar(field, element, json, opts.budget) if isinstance(element, DescMessage) and not _is_resetting_null(element, json): - if (budget := opts.budget) is not None: - budget.charge_message(element.type) + budget = opts.budget + budget.charge_message(element.type) msg = element.type() _read_message(msg, json, opts) return msg @@ -450,7 +440,7 @@ def _read_enum( desc: DescEnum, json: JsonValue, ignore_unknown_fields: bool, # noqa: FBT001 - budget: Budget | None = None, + budget: Budget, ) -> Enum | None: if json is None: return desc.type(desc.values[0].number) @@ -462,8 +452,7 @@ def _read_enum( if ignore_unknown_fields: return None # Succeeds for open enum, raises an error for closed - if budget is not None: - budget.charge(INT_SIZE + GC_HEAD_SIZE) + budget.charge(INT_SIZE + GC_HEAD_SIZE) return desc.type(json) if isinstance(json, str): if value := desc._values_by_name.get(json): @@ -481,11 +470,10 @@ def _read_scalar( desc: DescField | DescExtension, scalar_type: ScalarType, json: JsonValue, - budget: Budget | None = None, + budget: Budget, ) -> bool | int | float | str | bytes: value = _read_scalar_value(desc, scalar_type, json) - if budget is not None: - budget.charge_scalar(scalar_type, value) + budget.charge_scalar(scalar_type, value) return value @@ -681,12 +669,10 @@ def _struct_from_json( assert isinstance(value_wkt, WktValue) # noqa: S101 budget = opts.budget for k, v in json.items(): - if budget is not None: - budget.charge_message(value_desc.type) + budget.charge_message(value_desc.type) val = cast("Value", value_desc.type()) _value_from_json(val, v, opts, value_wkt) - if budget is not None: - budget.charge(DICT_ENTRY_SIZE + STR_OVERHEAD + len(k)) + budget.charge(DICT_ENTRY_SIZE + STR_OVERHEAD + len(k)) msg.fields[k] = val @@ -705,12 +691,10 @@ def _list_value_from_json( assert isinstance(element_wkt, WktValue) # noqa: S101 budget = opts.budget for e in json: - if budget is not None: - budget.charge_message(element_desc.type) + budget.charge_message(element_desc.type) val = cast("Value", element_desc.type()) _value_from_json(val, e, opts, element_wkt) - if budget is not None: - budget.charge(LIST_SLOT_SIZE) + budget.charge(LIST_SLOT_SIZE) msg.values.append(val) @@ -727,21 +711,7 @@ def _value_from_json_inner( msg: Value, json: JsonValue, opts: FromJsonOptions, wkt: WktValue ) -> None: budget = opts.budget - if budget is not None: - budget.charge(ONEOF_SIZE) - match json: - case bool() | None: - pass - case int() | float(): - budget.charge(FLOAT_SIZE) - case str(): - budget.charge(STR_OVERHEAD + len(json)) - case list(): - budget.charge_message(wkt.list_value.message.type) - case dict(): - budget.charge_message(wkt.struct_value.message.type) - case _: - pass + budget.charge(ONEOF_SIZE) match json: case None: msg.kind = Oneof( @@ -750,10 +720,13 @@ def _value_from_json_inner( case bool(): msg.kind = Oneof("bool_value", json) case int() | float(): + budget.charge(FLOAT_SIZE) msg.kind = Oneof("number_value", float(json)) case str(): + budget.charge(STR_OVERHEAD + len(json)) msg.kind = Oneof("string_value", json) case list(): + budget.charge_message(wkt.list_value.message.type) lv_desc = wkt.list_value.message lv_wkt = match_wkt(lv_desc) assert isinstance(lv_wkt, WktListValue) # noqa: S101 @@ -761,6 +734,7 @@ def _value_from_json_inner( _list_value_from_json(lv, json, opts, lv_wkt.values) msg.kind = Oneof("list_value", lv) case dict(): + budget.charge_message(wkt.struct_value.message.type) struct_desc = wkt.struct_value.message struct_wkt = match_wkt(struct_desc) assert isinstance(struct_wkt, WktStruct) # noqa: S101 diff --git a/src/protobuf/_message.py b/src/protobuf/_message.py index 5955451..49da302 100644 --- a/src/protobuf/_message.py +++ b/src/protobuf/_message.py @@ -19,6 +19,7 @@ from typing_extensions import Buffer, TypeVar from . import _native_message +from ._budget import Budget from ._descriptors import ( DescField, DescFieldValueEnum, @@ -697,15 +698,11 @@ def _merge_from_binary( ignore_unknown_fields: bool, # noqa: FBT001 allocation_limit: int | None = None, ) -> None: - from ._budget import Budget # noqa: PLC0415 - - budget = None - if allocation_limit is not None: - budget = Budget(allocation_limit) - # The root is charged here rather than in from_binary so both - # entry points share one budget semantic; a merge into an - # existing message overcharges by one base size. - budget.charge_message(type(self)) + budget = Budget(allocation_limit) + # The root is charged here rather than in from_binary so both entry + # points share one budget semantic; a merge into an existing message + # overcharges by one base size. + budget.charge_message(type(self)) opts = FromBinaryOptions( ignore_unknown_fields=ignore_unknown_fields, budget=budget ) @@ -722,17 +719,13 @@ def _merge_from_json( ) -> None: from json import loads as parse_json # noqa: PLC0415 - from ._budget import Budget # noqa: PLC0415 - # Needs to be lazy import since JSON specially handles many WKTs. from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 json_value = parse_json(json) - budget = None - if allocation_limit is not None: - budget = Budget(allocation_limit) - # See _merge_from_binary for why the root is charged here. - budget.charge_message(type(self)) + budget = Budget(allocation_limit) + # See _merge_from_binary for why the root is charged here. + budget.charge_message(type(self)) opts = FromJsonOptions( ignore_unknown_fields=ignore_unknown_fields, registry=registry, @@ -778,15 +771,11 @@ def _from_json_value( registry: Registry | None = None, allocation_limit: int | None = None, ) -> Self: - from ._budget import Budget # noqa: PLC0415 - # Needs to be lazy import since JSON specially handles many WKTs. from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 - budget = None - if allocation_limit is not None: - budget = Budget(allocation_limit) - budget.charge_message(cls) + budget = Budget(allocation_limit) + budget.charge_message(cls) message = cls() _read_message( message, diff --git a/src/protobuf/_unknown.py b/src/protobuf/_unknown.py index e9b5d42..4a18dad 100644 --- a/src/protobuf/_unknown.py +++ b/src/protobuf/_unknown.py @@ -79,11 +79,11 @@ def get_unknown_field( case DescFieldValueScalar(): reader = BinaryReader(memoryview(binary_fields[-1])) reader.tag() - return read_scalar(field_value.scalar, reader) + return read_scalar(field_value.scalar, reader, opts.budget) case DescFieldValueEnum(): reader = BinaryReader(memoryview(binary_fields[-1])) reader.tag() - enum_value = read_enum(field_value.enum, reader) + enum_value = read_enum(field_value.enum, reader, opts.budget) if isinstance(enum_value, Enum): return enum_value return None diff --git a/src/protobuf/_wkt_registry.py b/src/protobuf/_wkt_registry.py index 397e999..1d44749 100644 --- a/src/protobuf/_wkt_registry.py +++ b/src/protobuf/_wkt_registry.py @@ -26,6 +26,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, TypeAlias, cast +from ._budget import BYTES_OVERHEAD, INT_SIZE, LIST_SLOT_SIZE, STR_OVERHEAD from ._descriptors import ( DescEnum, DescField, @@ -99,7 +100,6 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: return f"{iso_secs}.{nanos_str}Z" def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: - from ._budget import INT_SIZE # noqa: PLC0415 from .wkt import Timestamp # noqa: PLC0415 value = cast("Timestamp", msg) @@ -129,8 +129,7 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo " 0001-01-01T00:00:00Z to 9999-12-31T23:59:59Z inclusive" ) raise ValueError(err) - if opts.budget is not None: - opts.budget.charge(2 * INT_SIZE) + opts.budget.charge(2 * INT_SIZE) value.seconds = Timestamp.from_datetime(dt).seconds value.nanos = nanos return True @@ -174,7 +173,6 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: return text + "s" def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: - from ._budget import INT_SIZE # noqa: PLC0415 from .wkt._mixin._const import DURATION_SECONDS_MAX # noqa: PLC0415 value = cast("Duration", msg) @@ -194,8 +192,7 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo nanos = int(duration_match[2] + "0" * (9 - len(duration_match[2]))) if seconds < 0 or duration_match[1] == "-0": nanos = -nanos - if opts.budget is not None: - opts.budget.charge(2 * INT_SIZE) + opts.budget.charge(2 * INT_SIZE) value.seconds = seconds value.nanos = nanos return True @@ -254,8 +251,7 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo if not desc: err = f"cannot decode {Any._desc.type_name} from JSON: {type_url} is not in the type registry" raise ValueError(err) - if (budget := opts.budget) is not None: - budget.charge_message(desc.type) + opts.budget.charge_message(desc.type) message = desc.type() if _has_custom_json(desc) and "value" in json: _read_message(message, json["value"], opts) @@ -264,11 +260,8 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo del json["@type"] _read_message(message, json, opts) any_ = Any.pack(message) - if budget is not None: - from ._budget import BYTES_OVERHEAD, STR_OVERHEAD # noqa: PLC0415 - - budget.charge(STR_OVERHEAD + len(any_.type_url)) - budget.charge(BYTES_OVERHEAD + len(any_.value)) + opts.budget.charge(STR_OVERHEAD + len(any_.type_url)) + opts.budget.charge(BYTES_OVERHEAD + len(any_.value)) value.type_url = any_.type_url value.value = any_.value return True @@ -297,7 +290,6 @@ def to_json_value(self, msg: Message, _opts: ToJsonOptions) -> JsonValue: return ",".join(parts) def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> bool: - from ._budget import LIST_SLOT_SIZE, STR_OVERHEAD # noqa: PLC0415 from ._names import proto_snake_case # noqa: PLC0415 value = cast("FieldMask", msg) @@ -310,8 +302,7 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo if "_" in path: err = f"cannot decode {value._desc.type_name} from JSON: path names must be lowerCamelCase" raise ValueError(err) - if opts.budget is not None: - opts.budget.charge(STR_OVERHEAD + len(path) + LIST_SLOT_SIZE) + opts.budget.charge(STR_OVERHEAD + len(path) + LIST_SLOT_SIZE) value.paths.append(proto_snake_case(path)) return True From 324f6049c80e516bb4256942bcfb8dc1c8981936 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 11:17:30 +0900 Subject: [PATCH 04/11] demonstrate amplification in test --- src/protobuf/_message.py | 6 ++---- tests/test_from_binary.py | 12 ++++++++++++ tests/test_json.py | 11 ++++++++++- 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/protobuf/_message.py b/src/protobuf/_message.py index 49da302..86301b7 100644 --- a/src/protobuf/_message.py +++ b/src/protobuf/_message.py @@ -600,8 +600,7 @@ def from_json( from JSON format. allocation_limit: If set, the approximate number of bytes of Python objects the - parse may allocate before raising a ValueError. Guards against - malicious payloads that expand into unexpectedly large messages. + parse may allocate before raising a ValueError. Raises: json.JSONDecodeError: If json_source is not valid JSON. @@ -635,8 +634,7 @@ def from_binary( ignore_unknown_fields: If `True`, unknown fields in the binary data are silently discarded. allocation_limit: If set, the approximate number of bytes of Python objects the - parse may allocate before raising a ValueError. Guards against - malicious payloads that expand into unexpectedly large messages. + parse may allocate before raising a ValueError. """ message = cls() message._merge_from_binary( diff --git a/tests/test_from_binary.py b/tests/test_from_binary.py index 1d141fe..fc36735 100644 --- a/tests/test_from_binary.py +++ b/tests/test_from_binary.py @@ -29,6 +29,7 @@ merge_from_binary, ) from protobuf._wire import BinaryWriter, WireType +from protobuf.wkt import FileDescriptorSet from .gen.delimited_encoding_pb import DelimitedEncoding from .gen.enums_pb import ClosedColor, EnumMessage @@ -439,3 +440,14 @@ def test_merge_from_binary_limit(self) -> None: msg = Scalars() with pytest.raises(ValueError, match="allocation budget exceeded"): merge_from_binary(msg, data, allocation_limit=500) + + def test_descriptor_set_amplification_rejected(self) -> None: + # Each empty FileDescriptorProto is two wire bytes (b"\x0a\x00": field 1, + # LEN, length 0) yet allocates a full descriptor instance with all its + # default containers -- ~700 bytes each, a ~350x blow-up. 1.9M entries + # is 3.8 MB on the wire but ~1.3 GB parsed. + rpc_read_limit = 4 * 1024 * 1024 + data = b"\x0a\x00" * 1_900_000 + assert len(data) < rpc_read_limit + with pytest.raises(ValueError, match="allocation budget exceeded"): + FileDescriptorSet.from_binary(data, allocation_limit=64 * 1024 * 1024) diff --git a/tests/test_json.py b/tests/test_json.py index 6b38ea5..129e8fc 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -25,7 +25,7 @@ message_from_json_value, message_to_json_value, ) -from protobuf.wkt import Struct +from protobuf.wkt import FileDescriptorSet, Struct from .gen.enums_pb import ClosedColor, Color, EnumMessage from .gen.json_enum_names_pb import JsonEnumNames, Season @@ -869,3 +869,12 @@ def test_no_limit_by_default(self) -> None: data = Lists(string_list=["a" * 100] * 100).to_json() msg = Lists.from_json(data) assert len(msg.string_list) == 100 + + def test_descriptor_set_amplification_rejected(self) -> None: + # Each empty "{}" in "file" for FileDescriptorSet + # parses to a full FileDescriptorProto (~700 bytes) from ~3 JSON bytes. + rpc_read_limit = 4 * 1024 * 1024 + data = '{"file":[' + ",".join(["{}"] * 1_200_000) + "]}" + assert len(data) < rpc_read_limit + with pytest.raises(ValueError, match="allocation budget exceeded"): + FileDescriptorSet.from_json(data, allocation_limit=64 * 1024 * 1024) From 42862bcbda00c36907a9374f82e10cccccdff875 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 11:20:23 +0900 Subject: [PATCH 05/11] fix stable ABI --- packages/protobuf-py-ext/src/constants.rs | 3 +++ packages/protobuf-py-ext/src/marshaler.rs | 15 ++++++++------- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/packages/protobuf-py-ext/src/constants.rs b/packages/protobuf-py-ext/src/constants.rs index 229bd32..7ccbfff 100644 --- a/packages/protobuf-py-ext/src/constants.rs +++ b/packages/protobuf-py-ext/src/constants.rs @@ -93,6 +93,8 @@ pub(crate) struct ConstantsInner { /// The string `__new__`. pub(crate) dunder_new: Py, + /// The string `__basicsize__`. + pub(crate) dunder_basicsize: Py, /// Python types. pub(crate) types: Types, @@ -146,6 +148,7 @@ impl Constants { values: PyString::new(py, "values").unbind(), dunder_new: PyString::new(py, "__new__").unbind(), + dunder_basicsize: PyString::new(py, "__basicsize__").unbind(), types: Types { desc_field: mod_descriptors diff --git a/packages/protobuf-py-ext/src/marshaler.rs b/packages/protobuf-py-ext/src/marshaler.rs index 2e5c28b..31e912d 100644 --- a/packages/protobuf-py-ext/src/marshaler.rs +++ b/packages/protobuf-py-ext/src/marshaler.rs @@ -7,7 +7,7 @@ use pyo3::{ pyclass, types::{ PyAnyMethods as _, PyBytes, PyDict, PyDictMethods as _, PyList, PyStringMethods as _, - PyType, PyTypeMethods as _, + PyType, }, }; @@ -204,12 +204,13 @@ impl MessageMarshaler { .push((member.attr.clone_ref(py), default.unbind())); } } - // tp_basicsize is the exact instance allocation size: message types - // only use slots, so instances never grow beyond it. The GC header is - // allocated in front of every instance on top of it. - // SAFETY - the type pointer of a live PyType is always valid. - let basic_size = unsafe { (*python_type.as_type_ptr()).tp_basicsize }; - let base_alloc_size = usize::try_from(basic_size).unwrap_or(0) + // __basicsize__ (tp_basicsize) is the exact instance allocation size: + // message types only use slots, so instances never grow beyond it. The + // GC header is allocated in front of every instance on top of it. + let basic_size = python_type + .getattr(&constants.dunder_basicsize)? + .extract::()?; + let base_alloc_size = basic_size + budget::GC_HEAD_SIZE + defaults.lists.len() * budget::EMPTY_LIST_SIZE + defaults.dicts.len() * budget::EMPTY_DICT_SIZE; From fc00070aa3cce521706a95d1b9dd8fcc4c31ba74 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 11:30:08 +0900 Subject: [PATCH 06/11] Fix scalar charging in pure python --- src/protobuf/_from_binary.py | 46 +++++++++++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 6 deletions(-) diff --git a/src/protobuf/_from_binary.py b/src/protobuf/_from_binary.py index db4cf96..9d172d2 100644 --- a/src/protobuf/_from_binary.py +++ b/src/protobuf/_from_binary.py @@ -22,10 +22,12 @@ from ._budget import ( BYTES_OVERHEAD, DICT_ENTRY_SIZE, + FLOAT_SIZE, GC_HEAD_SIZE, INT_SIZE, LIST_SLOT_SIZE, ONEOF_SIZE, + STR_OVERHEAD, Budget, ) from ._descriptors import ( @@ -78,10 +80,10 @@ class FromBinaryOptions: BinaryReader.fixed64, # 6: FIXED64 BinaryReader.fixed32, # 7: FIXED32 BinaryReader.bool_, # 8: BOOL - BinaryReader.string, # 9: STRING + None, # 9: STRING (length-delimited, handled in read_scalar) None, # 10: GROUP None, # 11: MESSAGE - BinaryReader.bytes_, # 12: BYTES + None, # 12: BYTES (length-delimited, handled in read_scalar) BinaryReader.uint32, # 13: UINT32 None, # 14: ENUM BinaryReader.sfixed32, # 15: SFIXED32 @@ -90,13 +92,45 @@ class FromBinaryOptions: BinaryReader.sint64, # 18: SINT64 ) +# Allocation charged for each fixed-size scalar before it is read, indexed like +# _SCALAR_READERS. Bools are shared singletons and allocate nothing. +_SCALAR_CHARGES = ( + None, # 0: unused + FLOAT_SIZE, # 1: DOUBLE + FLOAT_SIZE, # 2: FLOAT + INT_SIZE, # 3: INT64 + INT_SIZE, # 4: UINT64 + INT_SIZE, # 5: INT32 + INT_SIZE, # 6: FIXED64 + INT_SIZE, # 7: FIXED32 + 0, # 8: BOOL + None, # 9: STRING + None, # 10: GROUP + None, # 11: MESSAGE + None, # 12: BYTES + INT_SIZE, # 13: UINT32 + None, # 14: ENUM + INT_SIZE, # 15: SFIXED32 + INT_SIZE, # 16: SFIXED64 + INT_SIZE, # 17: SINT32 + INT_SIZE, # 18: SINT64 +) + def read_scalar(scalar_type: ScalarType, reader: BinaryReader, budget: Budget) -> Any: + if scalar_type == ScalarType.STRING: + length = reader.varint() + budget.charge(STR_OVERHEAD + length) + return str(reader.read(length), "utf-8") + if scalar_type == ScalarType.BYTES: + length = reader.varint() + budget.charge(BYTES_OVERHEAD + length) + return bytes(reader.read(length)) + charge = _SCALAR_CHARGES[scalar_type.value] reader_method = _SCALAR_READERS[scalar_type.value] - assert reader_method is not None # noqa: S101 - value = reader_method(reader) - budget.charge_scalar(scalar_type, value) - return value + assert charge is not None and reader_method is not None # noqa: S101, PT018 + budget.charge(charge) + return reader_method(reader) # TODO delete this, and either: From d19e8a6f4bcf4df1d19289e1a737637dedb241fd Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 11:41:03 +0900 Subject: [PATCH 07/11] Fix inconsistency in json charging --- packages/protobuf-py-ext/src/json_parse.rs | 5 +++-- packages/protobuf-py-ext/src/wkt_registry.rs | 7 ++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/packages/protobuf-py-ext/src/json_parse.rs b/packages/protobuf-py-ext/src/json_parse.rs index 5166c49..5e6fbf4 100644 --- a/packages/protobuf-py-ext/src/json_parse.rs +++ b/packages/protobuf-py-ext/src/json_parse.rs @@ -478,7 +478,7 @@ fn read_map_key<'py>( )), }, ScalarType::String => { - budget.charge(budget::STR_OVERHEAD + raw_key.len())?; + budget.charge(budget::STR_OVERHEAD + raw_key.chars().count())?; Ok(PyString::new(py, raw_key).into_any()) } _ => { @@ -978,7 +978,8 @@ pub(crate) fn read_json_value<'py, R: JsonSource<'py>>( let dict = PyDict::new(py); src.for_each_object_key(|key, src| { let value = read_json_value(src, budget)?; - budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; + budget + .charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.chars().count())?; dict.set_item(PyString::new(py, key), value)?; Ok(()) })?; diff --git a/packages/protobuf-py-ext/src/wkt_registry.rs b/packages/protobuf-py-ext/src/wkt_registry.rs index a017128..4570d43 100644 --- a/packages/protobuf-py-ext/src/wkt_registry.rs +++ b/packages/protobuf-py-ext/src/wkt_registry.rs @@ -258,7 +258,7 @@ impl WktAny { // Any.pack let packed_url = format!("type.googleapis.com/{}", inner_marshaler.type_name); let packed_value = inner_marshaler.to_binary(py, &inner_msg, true)?; - budget.charge(budget::STR_OVERHEAD + packed_url.len())?; + budget.charge(budget::STR_OVERHEAD + packed_url.chars().count())?; budget.charge(budget::BYTES_OVERHEAD + packed_value.as_bytes().len())?; self.type_url .set(message.as_any(), &PyString::new(py, &packed_url).into_any())?; @@ -329,7 +329,8 @@ impl WktFieldMask { marshaler.type_name ))); } - budget.charge(budget::STR_OVERHEAD + part.len() + budget::LIST_SLOT_SIZE)?; + budget + .charge(budget::STR_OVERHEAD + part.chars().count() + budget::LIST_SLOT_SIZE)?; paths.append(buffa_wkt::camel_to_snake(part))?; } Ok(()) @@ -395,7 +396,7 @@ impl WktStruct { let value_msg = value_marshaler.new_empty_message(py, self.value.get_python_type(py))?; read_message(value_marshaler, &value_msg, src, opts, depth + 1, budget)?; - budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.len())?; + budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.chars().count())?; dict.set_item(key, value_msg)?; Ok(()) })?; From c71e55adf66b01520f8b87287d60f0bc11ad42b4 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 7 Sep 2026 11:56:54 +0900 Subject: [PATCH 08/11] Fix pypy --- src/protobuf/_budget.py | 38 ++++++++++++++++++++++++++--------- src/protobuf/_descriptors.py | 4 ++++ src/protobuf/_from_binary.py | 8 ++++---- src/protobuf/_from_json.py | 14 ++++++------- src/protobuf/_message.py | 6 +++--- src/protobuf/_wkt_registry.py | 2 +- 6 files changed, 47 insertions(+), 25 deletions(-) diff --git a/src/protobuf/_budget.py b/src/protobuf/_budget.py index d8a6a6e..3249fb1 100644 --- a/src/protobuf/_budget.py +++ b/src/protobuf/_budget.py @@ -15,21 +15,22 @@ from __future__ import annotations import sys -from functools import cache from typing import TYPE_CHECKING, cast -from ._descriptors import ScalarType +from ._descriptors import DescMessage, ScalarType if TYPE_CHECKING: from collections.abc import Sized - from ._message import Message - # Approximate sizes of CPython heap allocations, measured with `sys.getsizeof` # on 64-bit CPython 3.14. The budget guards against unbounded allocation from # malicious payloads rather than providing exact accounting, so small # inaccuracies across versions and builds are fine. +OBJECT_HEADER_SIZE = 16 +"""Refcount and type pointer at the start of every object.""" +SLOT_SIZE = 8 +"""One `__slots__` entry: a pointer.""" GC_HEAD_SIZE = 16 """GC header allocated in front of every GC-tracked object.""" FLOAT_SIZE = 24 @@ -52,15 +53,28 @@ """A Oneof wrapper object: object header plus two references.""" -@cache -def _base_alloc_size(message_type: type[Message]) -> int: +def _instance_size(message_type: type) -> int: + """The fixed allocation size of an instance of a slots-only class.""" + basicsize = getattr(message_type, "__basicsize__", None) + if basicsize is not None: + return basicsize + # PyPy doesn't provide basicsize so we approximate it based on slot count. + slots = 0 + for cls in message_type.__mro__: + names = cls.__dict__.get("__slots__", ()) + names = (names,) if isinstance(names, str) else names + slots += sum(1 for name in names if name not in ("__dict__", "__weakref__")) + return OBJECT_HEADER_SIZE + SLOT_SIZE * slots + + +def _message_alloc_size(desc: DescMessage) -> int: """Approximate heap size of a freshly-initialized message instance. The fixed instance size (all fields are slots) plus the empty containers created for repeated/map field defaults. """ - size = message_type.__basicsize__ + GC_HEAD_SIZE - for _, default in message_type._desc._defaults: + size = _instance_size(desc.type) + GC_HEAD_SIZE + for _, default in desc._defaults: if isinstance(default, list): size += EMPTY_LIST_SIZE elif isinstance(default, dict): @@ -100,6 +114,10 @@ def charge_scalar(self, scalar_type: ScalarType, value: object) -> None: else: self.charge(INT_SIZE) - def charge_message(self, message_type: type[Message]) -> None: + def charge_message(self, desc: DescMessage) -> None: """Charges the base size of a new instance of the message type.""" - self.charge(_base_alloc_size(message_type)) + size = desc._alloc_size + if size is None: + size = _message_alloc_size(desc) + object.__setattr__(desc, "_alloc_size", size) + self.charge(size) diff --git a/src/protobuf/_descriptors.py b/src/protobuf/_descriptors.py index ef2d61d..bbd5a1c 100644 --- a/src/protobuf/_descriptors.py +++ b/src/protobuf/_descriptors.py @@ -310,6 +310,10 @@ class DescMessage: repr=False, compare=False, hash=False, init=False ) """Whether this message has any fields that require separate presence tracking.""" + _alloc_size: int | None = dataclassfield( + default=None, repr=False, compare=False, hash=False, init=False + ) + """Approximate heap size of a new instance, for the parsing allocation budget.""" def _finish_init(self) -> None: """Finish initialization of private attributes. diff --git a/src/protobuf/_from_binary.py b/src/protobuf/_from_binary.py index 9d172d2..86186b2 100644 --- a/src/protobuf/_from_binary.py +++ b/src/protobuf/_from_binary.py @@ -197,7 +197,7 @@ def read_message( ): existing: Message | None = message._get_member(desc_field) if existing is None: - budget.charge_message(desc_nested_message.type) + budget.charge_message(desc_nested_message) if field_value.oneof is not None: budget.charge(ONEOF_SIZE) existing = desc_nested_message.type() @@ -280,7 +280,7 @@ def read_list( case ScalarType(): value = read_scalar(element_type, reader, budget) case DescMessage(): - budget.charge_message(element_type.type) + budget.charge_message(element_type) if field_value.delimited_encoding: value = read_message( element_type.type(), @@ -400,7 +400,7 @@ def read_map_entry( ) return None case DescMessage(): - budget.charge_message(field_value.value.type) + budget.charge_message(field_value.value) value = field_value.value.type() read_message(value, reader, opts, depth + 1, length=reader.varint()) case _: @@ -417,7 +417,7 @@ def read_map_entry( case DescEnum() as desc_enum: value = desc_enum.type(desc_enum.values[0].number) case DescMessage(): - budget.charge_message(field_value.value.type) + budget.charge_message(field_value.value) value = field_value.value.type() case _: assert_never(field_value.value) diff --git a/src/protobuf/_from_json.py b/src/protobuf/_from_json.py index 1ce7fb3..2950a01 100644 --- a/src/protobuf/_from_json.py +++ b/src/protobuf/_from_json.py @@ -268,7 +268,7 @@ def _read_message_extension( del msg[ext.type] else: budget = opts.budget - budget.charge_message(field_value.message.type) + budget.charge_message(field_value.message) value = field_value.message.type() _read_message(value, json, opts) msg[ext.type] = value @@ -382,7 +382,7 @@ def _read_message_field( if msg._contains_member(field): value = msg._get_member(field) else: - budget.charge_message(field_value.message.type) + budget.charge_message(field_value.message) value = field_value.message.type() _read_message(value, json, opts) if field_value.oneof is not None: @@ -424,7 +424,7 @@ def _read_container_item( return _read_scalar(field, element, json, opts.budget) if isinstance(element, DescMessage) and not _is_resetting_null(element, json): budget = opts.budget - budget.charge_message(element.type) + budget.charge_message(element) msg = element.type() _read_message(msg, json, opts) return msg @@ -669,7 +669,7 @@ def _struct_from_json( assert isinstance(value_wkt, WktValue) # noqa: S101 budget = opts.budget for k, v in json.items(): - budget.charge_message(value_desc.type) + budget.charge_message(value_desc) val = cast("Value", value_desc.type()) _value_from_json(val, v, opts, value_wkt) budget.charge(DICT_ENTRY_SIZE + STR_OVERHEAD + len(k)) @@ -691,7 +691,7 @@ def _list_value_from_json( assert isinstance(element_wkt, WktValue) # noqa: S101 budget = opts.budget for e in json: - budget.charge_message(element_desc.type) + budget.charge_message(element_desc) val = cast("Value", element_desc.type()) _value_from_json(val, e, opts, element_wkt) budget.charge(LIST_SLOT_SIZE) @@ -726,7 +726,7 @@ def _value_from_json_inner( budget.charge(STR_OVERHEAD + len(json)) msg.kind = Oneof("string_value", json) case list(): - budget.charge_message(wkt.list_value.message.type) + budget.charge_message(wkt.list_value.message) lv_desc = wkt.list_value.message lv_wkt = match_wkt(lv_desc) assert isinstance(lv_wkt, WktListValue) # noqa: S101 @@ -734,7 +734,7 @@ def _value_from_json_inner( _list_value_from_json(lv, json, opts, lv_wkt.values) msg.kind = Oneof("list_value", lv) case dict(): - budget.charge_message(wkt.struct_value.message.type) + budget.charge_message(wkt.struct_value.message) struct_desc = wkt.struct_value.message struct_wkt = match_wkt(struct_desc) assert isinstance(struct_wkt, WktStruct) # noqa: S101 diff --git a/src/protobuf/_message.py b/src/protobuf/_message.py index 86301b7..384b30a 100644 --- a/src/protobuf/_message.py +++ b/src/protobuf/_message.py @@ -700,7 +700,7 @@ def _merge_from_binary( # The root is charged here rather than in from_binary so both entry # points share one budget semantic; a merge into an existing message # overcharges by one base size. - budget.charge_message(type(self)) + budget.charge_message(self._desc) opts = FromBinaryOptions( ignore_unknown_fields=ignore_unknown_fields, budget=budget ) @@ -723,7 +723,7 @@ def _merge_from_json( json_value = parse_json(json) budget = Budget(allocation_limit) # See _merge_from_binary for why the root is charged here. - budget.charge_message(type(self)) + budget.charge_message(self._desc) opts = FromJsonOptions( ignore_unknown_fields=ignore_unknown_fields, registry=registry, @@ -773,7 +773,7 @@ def _from_json_value( from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 budget = Budget(allocation_limit) - budget.charge_message(cls) + budget.charge_message(cls._desc) message = cls() _read_message( message, diff --git a/src/protobuf/_wkt_registry.py b/src/protobuf/_wkt_registry.py index 1d44749..9993ea2 100644 --- a/src/protobuf/_wkt_registry.py +++ b/src/protobuf/_wkt_registry.py @@ -251,7 +251,7 @@ def from_json(self, msg: Message, json: JsonValue, opts: FromJsonOptions) -> boo if not desc: err = f"cannot decode {Any._desc.type_name} from JSON: {type_url} is not in the type registry" raise ValueError(err) - opts.budget.charge_message(desc.type) + opts.budget.charge_message(desc) message = desc.type() if _has_custom_json(desc) and "value" in json: _read_message(message, json["value"], opts) From 64e97370414debbcc4f95c345eb00d743796f89e Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Thu, 10 Sep 2026 11:00:07 +0900 Subject: [PATCH 09/11] git add --- tests/test_budget.py | 67 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 tests/test_budget.py diff --git a/tests/test_budget.py b/tests/test_budget.py new file mode 100644 index 0000000..8c45972 --- /dev/null +++ b/tests/test_budget.py @@ -0,0 +1,67 @@ +# Copyright (c) 2025-2026 Buf Technologies, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import sys +from typing import Any + +import pytest + +from protobuf._budget import SLOT_SIZE, Budget, _instance_size +from protobuf._native_message import NativeMessageClass +from protobuf.wkt import FileDescriptorSet, Struct, Timestamp, Value + +from .gen.lists_pb import Lists +from .gen.maps_pb import Maps +from .gen.oneofs_pb import Oneofs +from .gen.scalars_pb import Scalars + + +class _NoBasicsize: + """Simulates a PyPy type, by being an instantiated object that includes the mro of a type.""" + + def __init__(self, real: type) -> None: + self.__mro__ = real.__mro__ + + +@pytest.mark.skipif( + NativeMessageClass is not None, reason="fallback only needed for pure Python" +) +@pytest.mark.parametrize( + "message_type", + [Scalars, Lists, Maps, Oneofs, Struct, Value, Timestamp, FileDescriptorSet], +) +def test_instance_size_fallback_matches_basicsize(message_type: type[Any]) -> None: + # Ensure the type (and its __slots__) exists before inspecting it. + message_type() + expected = message_type.__basicsize__ + if sys.version_info < (3, 12): + # Before managed weakrefs, the __weakref__ slot occupied a pointer + # inside the instance, which the fallback deliberately excludes. + expected -= SLOT_SIZE + mt = _NoBasicsize(message_type) + assert not hasattr(mt, "__basicsize__") + assert _instance_size(mt) == expected # type: ignore[arg-type] + + +def test_alloc_size_cached_on_descriptor() -> None: + desc = Scalars._desc + object.__setattr__(desc, "_alloc_size", None) + budget = Budget() + budget.charge_message(desc) + assert desc._alloc_size is not None + assert budget.current == desc._alloc_size + budget.charge_message(desc) + assert budget.current == 2 * desc._alloc_size From b588dec1158d21fbf316b3f33898f3d0cec81e6c Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 28 Sep 2026 15:59:35 +0900 Subject: [PATCH 10/11] Add doc --- docs/serialization.md | 31 +++++++++++++++++++++++++++++++ tests/test_budget.py | 33 ++++++++++++++++++++++++++++++--- 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/docs/serialization.md b/docs/serialization.md index d756211..d33bec8 100644 --- a/docs/serialization.md +++ b/docs/serialization.md @@ -22,6 +22,8 @@ user = User.from_binary(data) ### Options +#### Unknown Fields + When a message is parsed from binary data containing field numbers it doesn't recognize, the unknown fields are stored internally and re-emitted during serialization. This means a message can pass through an intermediary that doesn't know about newer fields without losing data. @@ -37,6 +39,35 @@ data = user.to_binary(write_unknown_fields=False) user = User.from_binary(data, ignore_unknown_fields=True) ``` +#### Allocation Limit + +Because protobuf serialization can create very compact binary payloads, it is possible for the memory usage of a +parsed message to differ drastically from the input number of bytes. When parsing untrusted payloads, +such as in an external-facing API server, this can allow malicious users to send small messages that take +a large amount of memory or potentially crashing the server. This is most pronounced in schemas with +repeated fields of message type with a large number of fields. + +You can mitigate this using the `allocation_limit` option in `from_binary`. When set, an estimate of the +memory usage of a message is maintained while it is parsed, and if it goes over the limit, the parse fails +immediately before processing the entire payload. The limit is based on the schema and content, not the environment, +i.e., it charges a value for a string field based on the number of characters with a fixed overhead for a Python +string. The overhead can change between Python versions, so the allocation budget should not be considered +a precise value. Such changes are generally relatively small and fixed so an allocation limit determined for +one environment should generally work fine when i.e., updating Python. + +The allocation budget is an upper bound, so it is possible that a message that would be under the limit is +rejected. For example, when parsing binary the string is charged with the number of utf8 bytes in the payload. +For ASCII strings, this will be the precise number, but for i.e. certain CJK characters, it overcharges by +~30%. + +It is recommended to set an allocation limit for applications parsing untrusted payloads based on your target +memory usage. You may need to experiment with values to see the effective memory utilization due to potential +overcharging. + +```python +user = User.from_binary(data, allocation_limit=8 * 1024 * 1024) # Roughly cap memory usage of parsed message to 8MB +``` + ## JSON ```python diff --git a/tests/test_budget.py b/tests/test_budget.py index 8c45972..f995300 100644 --- a/tests/test_budget.py +++ b/tests/test_budget.py @@ -15,11 +15,11 @@ from __future__ import annotations import sys -from typing import Any +from typing import Any, cast import pytest -from protobuf._budget import SLOT_SIZE, Budget, _instance_size +from protobuf._budget import OBJECT_HEADER_SIZE, SLOT_SIZE, Budget, _instance_size from protobuf._native_message import NativeMessageClass from protobuf.wkt import FileDescriptorSet, Struct, Timestamp, Value @@ -36,9 +36,21 @@ def __init__(self, real: type) -> None: self.__mro__ = real.__mro__ +def _slot_count(message_type: type) -> int: + slots = 0 + for cls in message_type.__mro__: + names = cls.__dict__.get("__slots__", ()) + names = (names,) if isinstance(names, str) else names + slots += sum(1 for n in names if n not in ("__dict__", "__weakref__")) + return slots + + @pytest.mark.skipif( NativeMessageClass is not None, reason="fallback only needed for pure Python" ) +@pytest.mark.skipif( + sys.implementation.name != "cpython", reason="no __basicsize__ to compare against" +) @pytest.mark.parametrize( "message_type", [Scalars, Lists, Maps, Oneofs, Struct, Value, Timestamp, FileDescriptorSet], @@ -51,9 +63,24 @@ def test_instance_size_fallback_matches_basicsize(message_type: type[Any]) -> No # Before managed weakrefs, the __weakref__ slot occupied a pointer # inside the instance, which the fallback deliberately excludes. expected -= SLOT_SIZE + if not getattr(sys, "_is_gil_enabled", lambda: True)(): + # Free-threaded builds fold the GC header into a 32-byte object + # header, which basicsize includes and the formula does not. + expected -= 16 mt = _NoBasicsize(message_type) assert not hasattr(mt, "__basicsize__") - assert _instance_size(mt) == expected # type: ignore[arg-type] + assert _instance_size(cast("type", mt)) == expected + + +@pytest.mark.skipif( + sys.implementation.name != "pypy", reason="exercises the real fallback path" +) +def test_instance_size_fallback_on_pypy() -> None: + Scalars() + assert not hasattr(Scalars, "__basicsize__") + assert _instance_size(Scalars) == OBJECT_HEADER_SIZE + SLOT_SIZE * _slot_count( + Scalars + ) def test_alloc_size_cached_on_descriptor() -> None: From 2f01eff5f6425573430b6e800a4b62ff70ba1982 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 28 Sep 2026 16:02:27 +0900 Subject: [PATCH 11/11] format --- docs/serialization.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/serialization.md b/docs/serialization.md index d33bec8..06bcfb9 100644 --- a/docs/serialization.md +++ b/docs/serialization.md @@ -65,7 +65,9 @@ memory usage. You may need to experiment with values to see the effective memory overcharging. ```python -user = User.from_binary(data, allocation_limit=8 * 1024 * 1024) # Roughly cap memory usage of parsed message to 8MB +user = User.from_binary( + data, allocation_limit=8 * 1024 * 1024 +) # Roughly cap memory usage of parsed message to 8MB ``` ## JSON