diff --git a/docs/serialization.md b/docs/serialization.md index d756211..06bcfb9 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,37 @@ 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/packages/protobuf-py-ext/src/budget.rs b/packages/protobuf-py-ext/src/budget.rs new file mode 100644 index 0000000..b769e89 --- /dev/null +++ b/packages/protobuf-py-ext/src/budget.rs @@ -0,0 +1,56 @@ +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. +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. +pub(crate) const STR_OVERHEAD: usize = 41; +/// Header of a bytes object. +pub(crate) const BYTES_OVERHEAD: usize = 33; +/// An empty list. +pub(crate) const EMPTY_LIST_SIZE: usize = 56; +/// An empty dict. +pub(crate) const EMPTY_DICT_SIZE: usize = 64; +/// 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 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. +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/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/json_parse.rs b/packages/protobuf-py-ext/src/json_parse.rs index 81a900f..5e6fbf4 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, + 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, 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, + 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, 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, + 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, budget), + None => read_generic_object(marshaler, message, src, opts, depth, budget), } } @@ -98,10 +102,11 @@ fn read_generic_object<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + 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, 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, budget)?; } else { - handle_unknown_key(marshaler, message, key, src, opts, depth)?; + handle_unknown_key(marshaler, message, key, src, opts, depth, budget)?; } Ok(()) })?; @@ -178,6 +183,7 @@ fn read_field<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + 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, + budget, ), - ParserFieldType::List { .. } => read_list(marshaler, parser, message, src, opts, depth), + ParserFieldType::List { .. } => { + read_list(marshaler, parser, message, src, opts, depth, budget) + } ParserFieldType::Map { key_type, value_parser, @@ -207,6 +216,7 @@ fn read_field<'py, R: JsonSource<'py>>( depth, *key_type, value_parser, + budget, ), } } @@ -221,6 +231,7 @@ fn read_singular<'py, R: JsonSource<'py>>( depth: usize, oneof_attr: Option<&crate::attribute_access::AttributeAccess>, requires_presence: bool, + 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, budget)?; + if oneof_attr.is_some() { + 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, budget)? { + if oneof_attr.is_some() { + 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))?, + _ => { + budget.charge(inner.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, budget)?; + if oneof_attr.is_some() { + 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, + 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, 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, budget)? + { + 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, + 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, budget)?.get_type() + ), Exc::Type, )); } @@ -368,9 +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)?; - 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, 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(()) @@ -387,11 +423,14 @@ fn read_container_item<'py, R: JsonSource<'py>>( opts: &FromJsonOpts, depth: usize, is_map: bool, + 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, budget)?)) + } FieldParserValue::Message { message: msg_desc, .. } => { @@ -401,12 +440,13 @@ fn read_container_item<'py, R: JsonSource<'py>>( src.next_null()?; return Err(container_null_error(ctx, is_map)); } + 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)?; + 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) + read_enum(enum_, src, opts, budget) } _ => { // Resetting null for a list item / map value: error. @@ -426,6 +466,7 @@ fn read_map_key<'py>( ctx: &FieldContext<'_, 'py>, key_type: ScalarType, raw_key: &str, + budget: &mut Budget, ) -> PyResult> { match key_type { ScalarType::Bool => match raw_key { @@ -436,8 +477,14 @@ 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 => { + budget.charge(budget::STR_OVERHEAD + raw_key.chars().count())?; + Ok(PyString::new(py, raw_key).into_any()) + } + _ => { + budget.charge(budget::INT_SIZE)?; + parse_int_string(py, ctx, raw_key, key_type) + } } } @@ -446,12 +493,13 @@ pub(crate) fn read_scalar<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, scalar: ScalarType, + 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, budget)?; return Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -460,25 +508,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, budget)?; if value.is_finite() && !(FLOAT32_MIN..=FLOAT32_MAX).contains(&value) { return Err(ctx.error(&format!("float value out of range: {value}"), Exc::Overflow)); } + budget.charge(budget::FLOAT_SIZE)?; + Ok(PyFloat::new(py, value).into_any()) + } + ScalarType::Double => { + let value = parse_float(ctx, src, budget)?; + 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::String => { + let value = read_string(ctx, src, budget)?; + // Character count approximates the payload size. + budget.charge(budget::STR_OVERHEAD + value.len()?)?; + Ok(value.into_any()) + } + ScalarType::Bytes => read_bytes(ctx, src, budget), + _ => { + budget.charge(budget::INT_SIZE)?; + read_int(ctx, src, scalar, budget) + } } } fn read_string<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, + budget: &mut Budget, ) -> PyResult> { if src.peek()? != JsonKind::String { - let value = read_json_value(src)?; + let value = read_json_value(src, budget)?; return Err(ctx.error( &format!("expected string got: {}", value.get_type()), Exc::Type, @@ -490,10 +552,11 @@ fn read_string<'py, R: JsonSource<'py>>( fn read_bytes<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, + 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, budget)?; return Err(ctx.error( &format!("expected base64-encoded string got: {}", value.get_type()), Exc::Type, @@ -511,6 +574,7 @@ fn read_bytes<'py, R: JsonSource<'py>>( .map_err(|_| ctx.error("invalid base64 data", Exc::Value))?; Ok(decoded) })?; + budget.charge(budget::BYTES_OVERHEAD + decoded.len())?; Ok(PyBytes::new(py, &decoded).into_any()) } @@ -531,7 +595,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, + budget: &mut Budget, +) -> PyResult { match src.peek()? { JsonKind::Number => { let value = src.next_float()?; @@ -560,7 +628,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, budget)?; Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -574,6 +642,7 @@ fn read_int<'py, R: JsonSource<'py>>( ctx: &FieldContext<'_, 'py>, src: &mut R, int_type: ScalarType, + budget: &mut Budget, ) -> PyResult> { let py = src.py(); let value = match src.peek()? { @@ -599,7 +668,7 @@ fn read_int<'py, R: JsonSource<'py>>( }); } _ => { - let value = read_json_value(src)?; + let value = read_json_value(src, budget)?; return Err(ctx.error( &format!("unexpected json type: {}", value.get_type()), Exc::Type, @@ -676,6 +745,7 @@ fn read_enum<'py, R: JsonSource<'py>>( enum_desc: &DescEnum, src: &mut R, opts: &FromJsonOpts, + budget: &mut Budget, ) -> PyResult>> { let py = src.py(); match src.peek()? { @@ -694,6 +764,7 @@ fn read_enum<'py, R: JsonSource<'py>>( if opts.ignore_unknown_fields { return Ok(None); } + 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) { @@ -702,6 +773,7 @@ fn read_enum<'py, R: JsonSource<'py>>( Ok(None) } else { // 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,))?)) } } @@ -720,7 +792,7 @@ fn read_enum<'py, R: JsonSource<'py>>( } }), _ => { - let value = read_json_value(src)?; + let value = read_json_value(src, budget)?; Err(decode_enum_error(py, enum_desc, &value)) } } @@ -745,6 +817,7 @@ fn handle_unknown_key<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if raw_key.starts_with('[') @@ -760,7 +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)?; + read_extension(marshaler, message, &extension, src, opts, depth, budget)?; } else { src.skip()?; } @@ -785,6 +858,7 @@ fn read_extension<'py, R: JsonSource<'py>>( src: &mut R, opts: &FromJsonOpts, depth: usize, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); let ext_type = extension.getattr(&marshaler.constants.type_)?; @@ -803,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)?; + let value = read_scalar(&ctx, src, *scalar_type, budget)?; target.set_item(&ext_type, value)?; } } @@ -811,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)? { + } else if let Some(value) = read_enum(enum_, src, opts, budget)? { target.set_item(&ext_type, value)?; } } @@ -824,8 +898,9 @@ fn read_extension<'py, R: JsonSource<'py>>( src.next_null()?; target.del_item(&ext_type)?; } else { + 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)?; + read_message(inner, &value, src, opts, depth + 1, budget)?; target.set_item(&ext_type, value)?; } } @@ -837,16 +912,21 @@ 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, budget)?.get_type() + ), Exc::Type, )); } let element_value = FieldParserValue::from_desc_single(element); + 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)? + read_container_item(&ctx, &element_value, src, opts, depth, false, budget)? { + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value)?; } Ok(()) @@ -862,7 +942,10 @@ fn read_extension<'py, R: JsonSource<'py>>( } /// Materializes the next JSON value as a Python object. -pub(crate) fn read_json_value<'py, R: JsonSource<'py>>(src: &mut R) -> PyResult> { +pub(crate) fn read_json_value<'py, R: JsonSource<'py>>( + src: &mut R, + budget: &mut Budget, +) -> PyResult> { let py = src.py(); match src.peek()? { JsonKind::Null => { @@ -870,20 +953,33 @@ 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 => { + budget.charge(budget::INT_SIZE)?; + src.next_number() + } + JsonKind::String => { + let value = src.next_py_str()?; + budget.charge(budget::STR_OVERHEAD + value.len()?)?; + Ok(value.into_any()) + } JsonKind::Array => { + 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, budget)?; + budget.charge(budget::LIST_SLOT_SIZE)?; + list.append(value)?; Ok(()) })?; Ok(list.into_any()) } JsonKind::Object => { + 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, budget)?; + 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/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..31e912d 100644 --- a/packages/protobuf-py-ext/src/marshaler.rs +++ b/packages/protobuf-py-ext/src/marshaler.rs @@ -13,6 +13,7 @@ use pyo3::{ 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) alloc_size: usize, + /// The Python type of the message. pub(crate) python_type: Py, @@ -198,6 +204,16 @@ impl MessageMarshaler { .push((member.attr.clone_ref(py), default.unbind())); } } + // __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; Ok(Self { inner: Arc::new(MessageMarshalerInner { parser, @@ -205,6 +221,7 @@ impl MessageMarshaler { members_by_name: members_by_name.unbind(), members, max_field_number, + alloc_size: base_alloc_size, python_type: python_type.clone().unbind(), type_name, wkt, @@ -222,6 +239,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 +249,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..0fe8d08 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 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)?; + marshaler.merge_from_binary(py, slf, data, ignore_unknown_fields, &mut 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 budget = Budget::new(allocation_limit); + budget.charge(marshaler.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 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 budget = Budget::new(allocation_limit); + budget.charge(marshaler.alloc_size)?; + parse_json_into( + py, + &marshaler, + slf, + json, + ignore_unknown_fields, + registry, + &mut 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 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)?; + read_message_from_tree(py, &marshaler, &message, data, &opts, &mut 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 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>( @@ -514,18 +541,26 @@ fn parse_json_into<'py>( json: &Bound<'py, PyAny>, ignore_unknown_fields: bool, registry: Option>, + 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, + 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, 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, 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..e75ed27 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, + budget: &mut Budget, ) -> PyResult<()> { match &self.type_ { ParserFieldType::Singular { @@ -265,6 +267,7 @@ impl FieldParser { depth, oneof_attr.as_ref(), *requires_presence, + budget, )?, ParserFieldType::List { unpacked_wire_type, @@ -279,6 +282,7 @@ impl FieldParser { depth, *unpacked_wire_type, *packable, + budget, )?, ParserFieldType::Map { key_type, @@ -296,6 +300,7 @@ impl FieldParser { value_parser, key_default_value, value_default_value, + budget, )?, } @@ -313,11 +318,23 @@ impl FieldParser { depth: usize, oneof_attr: Option<&AttributeAccess>, requires_presence: bool, + 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, + budget, + )?; match value { SingleValue::Parsed(value) => { + if oneof_attr.is_some() { + 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, budget)?; } } } @@ -365,6 +382,7 @@ impl FieldParser { depth: usize, unpacked_wire_type: WireType, packable: bool, + 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, + budget, )?; match value { - SingleValue::Parsed(value) => list.append(value)?, + SingleValue::Parsed(value) => { + 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,25 @@ 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, 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, budget)?; match value { - SingleValue::Parsed(value) => list.append(value)?, + SingleValue::Parsed(value) => { + 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, budget)?; } } } @@ -425,6 +451,7 @@ impl FieldParser { value_parser: &FieldParser, key_default_value: &Py, value_default_value: &Py, + budget: &mut Budget, ) -> PyResult<()> { let len = decode_varint(buffer).map_err(map_varint_err)? as usize; check_buffer_remaining(buffer, len)?; @@ -446,9 +473,10 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, + budget, ); } - key = Some(read_scalar(py, key_type, &mut entry_buffer)?); + key = Some(read_scalar(py, key_type, &mut entry_buffer, budget)?); } 2 => { if value_parser.wire_type != wire_type { @@ -458,6 +486,7 @@ impl FieldParser { entry_tag, &entry_checkpoint, opts, + budget, ); } value = Some(value_parser.read_single_value( @@ -468,6 +497,7 @@ impl FieldParser { &mut entry_buffer, opts, depth, + budget, )?); } _ => { @@ -481,14 +511,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. + 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()) }; 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, + budget, + ); }; + budget.charge(budget::DICT_ENTRY_SIZE)?; self.assign_map_entry(py, message, key, value)?; Ok(()) } @@ -513,13 +552,14 @@ impl FieldParser { tag: u32, entry_bytes: &[u8], opts: FromBinaryOpts, + 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, budget)?; } Ok(()) } @@ -534,9 +574,10 @@ impl FieldParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, + 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, budget)?, FieldParserValue::Message { message: message_desc, .. @@ -564,6 +605,7 @@ impl FieldParser { { existing.cast_into::()? } else { + budget.charge(marshaler.alloc_size)?; marshaler.new_empty_message(py, parser.inner.python_type.bind(py))? }; parser.merge_from_binary( @@ -572,6 +614,7 @@ impl FieldParser { &mut message_buffer, opts, depth + 1, + budget, )?; message_instance.into_any() } @@ -581,6 +624,7 @@ impl FieldParser { if let Some(value) = value { value.bind(py).clone() } else if enum_.open { + budget.charge(budget::INT_SIZE + budget::GC_HEAD_SIZE)?; enum_.py_type.bind(py).call1((number,))? } else { return Ok(SingleValue::UnknownEnumValue(number)); @@ -678,6 +722,7 @@ impl MessageParser { buffer: &mut Bytes, opts: FromBinaryOpts, depth: usize, + budget: &mut Budget, ) -> PyResult<()> { check_parse_recursion_depth(depth)?; while buffer.has_remaining() { @@ -691,12 +736,18 @@ 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, 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], + budget, + )?; } } } @@ -709,7 +760,26 @@ fn read_scalar<'py>( py: Python<'py>, s: ScalarType, buffer: &mut Bytes, + budget: &mut Budget, ) -> PyResult> { + match s { + 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 + | ScalarType::Uint64 + | ScalarType::Int32 + | ScalarType::Uint32 + | ScalarType::Fixed64 + | ScalarType::Fixed32 + | ScalarType::Sfixed32 + | ScalarType::Sfixed64 + | ScalarType::Sint32 + | 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 => {} + } let res = match s { ScalarType::Double => { PyFloat::new(py, buffer.try_get_f64_le().map_err(map_try_get_err)?).into_any() @@ -743,12 +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)?; + 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)?; + budget.charge(budget::BYTES_OVERHEAD + len)?; let bytes = buffer.split_to(len); PyBytes::new(py, &bytes).into_any() } @@ -872,7 +944,11 @@ fn write_unknown_field( message: &Bound<'_, NativeMessage>, field_number: u32, field_bytes: &[u8], + 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. + 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..4570d43 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, + 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, budget, |text, _| { parse_timestamp(&marshaler.type_name, text) })?; + 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, + 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, budget, |text, _| { parse_duration(&marshaler.type_name, text) })?; + budget.charge(2 * budget::INT_SIZE)?; self.seconds .set(message.as_any(), PyInt::new(py, secs).as_any())?; self.nanos @@ -157,7 +162,9 @@ 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)?; + // 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()?; if inner_marshaler.wkt.is_none() { @@ -183,12 +190,13 @@ impl WktAny { src: &mut R, opts: &FromJsonOpts, _depth: usize, + 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, budget)?; let Ok(dict) = tree.cast::() else { return Err(PyTypeError::new_err(format!( "cannot decode {message_type_name} from JSON: {}", @@ -230,6 +238,7 @@ impl WktAny { .getattr(&marshaler.constants.ext_marshaler)? .cast_into::()?; let inner_marshaler = inner_marshaler.get(); + 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(); @@ -238,17 +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)?; + 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)?; + 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)?; + 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())?; self.value.set(message.as_any(), packed_value.as_any())?; @@ -300,13 +311,14 @@ impl WktFieldMask { src: &mut R, _opts: &FromJsonOpts, _depth: usize, + 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, budget, |text, budget| { if text.is_empty() { return Ok(()); } @@ -317,6 +329,8 @@ impl WktFieldMask { marshaler.type_name ))); } + budget + .charge(budget::STR_OVERHEAD + part.chars().count() + budget::LIST_SLOT_SIZE)?; paths.append(buffa_wkt::camel_to_snake(part))?; } Ok(()) @@ -360,10 +374,11 @@ impl WktStruct { src: &mut R, opts: &FromJsonOpts, depth: usize, + 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, budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, @@ -377,9 +392,11 @@ impl WktStruct { .cast_into::()?; // Duplicate keys use last-in-wins semantics (per the ProtoJSON spec). src.for_each_object_key(|key, src| { + 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)?; + read_message(value_marshaler, &value_msg, src, opts, depth + 1, budget)?; + budget.charge(budget::DICT_ENTRY_SIZE + budget::STR_OVERHEAD + key.chars().count())?; dict.set_item(key, value_msg)?; Ok(()) })?; @@ -422,10 +439,11 @@ impl WktListValue { src: &mut R, opts: &FromJsonOpts, depth: usize, + 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, budget)?; return Err(PyTypeError::new_err(format!( "cannot decode {} from JSON: {}", marshaler.type_name, @@ -438,9 +456,11 @@ impl WktListValue { .get(py, message.as_any())? .cast_into::()?; src.for_each_array_item(|src| { + 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)?; + read_message(element_marshaler, &value_msg, src, opts, depth + 1, budget)?; + budget.charge(budget::LIST_SLOT_SIZE)?; list.append(value_msg)?; Ok(()) })?; @@ -510,8 +530,10 @@ impl WktValue { src: &mut R, opts: &FromJsonOpts, depth: usize, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); + budget.charge(budget::ONEOF_SIZE)?; let oneof = match src.peek()? { JsonKind::Null => { src.next_null()?; @@ -523,6 +545,7 @@ impl WktValue { ), JsonKind::Number => { let number = src.next_float()?; + budget.charge(budget::FLOAT_SIZE)?; Oneof::new( self.number_name.bind(py), &PyFloat::new(py, number).into_any(), @@ -530,20 +553,23 @@ impl WktValue { } JsonKind::String => { let string = src.next_py_str()?; + 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)?; + 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)?; + 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)?; + 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)?; + read_message(inner, &struct_msg, src, opts, depth + 1, budget)?; Oneof::new(self.struct_name.bind(py), &struct_msg.into_any()) } }; @@ -581,6 +607,7 @@ impl WktWrapper { src: &mut R, _opts: &FromJsonOpts, _depth: usize, + budget: &mut Budget, ) -> PyResult<()> { let py = src.py(); if src.peek()? == JsonKind::Null { @@ -591,7 +618,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, budget)?; self.field.set(message.as_any(), &value) } } @@ -639,16 +666,17 @@ impl WktKind { src: &mut R, opts: &FromJsonOpts, depth: usize, + 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, 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), } } @@ -887,17 +915,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, + 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, 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, 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..3249fb1 --- /dev/null +++ b/src/protobuf/_budget.py @@ -0,0 +1,123 @@ +# 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 TYPE_CHECKING, cast + +from ._descriptors import DescMessage, ScalarType + +if TYPE_CHECKING: + from collections.abc import Sized + +# 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 +"""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.""" +EMPTY_DICT_SIZE = 64 +"""An empty dict.""" +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.""" + + +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 = _instance_size(desc.type) + GC_HEAD_SIZE + for _, default in 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. + """ + + __slots__ = ("current", "max") + + def __init__(self, limit: int | None = None) -> None: + self.current = 0 + self.max = sys.maxsize if limit is None else 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, desc: DescMessage) -> None: + """Charges the base size of a new instance of the 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 39968a9..86186b2 100644 --- a/src/protobuf/_from_binary.py +++ b/src/protobuf/_from_binary.py @@ -14,11 +14,22 @@ 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 +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 ( DescEnum, DescFieldValueEnum, @@ -47,9 +58,12 @@ class FromBinaryOptions: Args: ignore_unknown_fields: If `True`, unknown fields are ignored instead of being added to the message. + 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 = field(default_factory=Budget) # Dispatch table for reading scalar values. CPython currently does not generate @@ -66,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 @@ -78,10 +92,44 @@ 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) -> Any: +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 + assert charge is not None and reader_method is not None # noqa: S101, PT018 + budget.charge(charge) return reader_method(reader) @@ -128,19 +176,30 @@ 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) + 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) ) 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 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: + budget.charge_message(desc_nested_message) + 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 +211,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 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 +235,7 @@ def read_message( ) if entry: key, value = entry + budget.charge(DICT_ENTRY_SIZE) message._get_member(desc_field)[key] = value case _: assert_never(desc_field) @@ -204,15 +266,21 @@ 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) + 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) ) return + budget = opts.budget match element_type: case ScalarType(): - value = read_scalar(element_type, reader) + value = read_scalar(element_type, reader, budget) case DescMessage(): + budget.charge_message(element_type) if field_value.delimited_encoding: value = read_message( element_type.type(), @@ -226,13 +294,14 @@ 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) + budget.charge(LIST_SLOT_SIZE) list_.append(value) @@ -246,38 +315,46 @@ 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) + 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): + 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) -> 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 + 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: if message is None: return writer = BinaryWriter() writer.tag(field_number, WireType.VARINT) writer.int32(value) + field_bytes = writer.finish() + 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 +373,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 +382,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 +391,16 @@ 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(): + budget.charge_message(field_value.value) value = field_value.value.type() read_message(value, reader, opts, depth + 1, length=reader.varint()) case _: @@ -338,6 +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) value = field_value.value.type() case _: assert_never(field_value.value) @@ -361,13 +441,19 @@ def _read_unknown_map_entry( ) if not opts.ignore_unknown_fields and message: key_raw = _encode_varint((field_number << 3) | WireType.LENGTH_DELIMITED) + 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) ) 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 +469,9 @@ 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. """ - 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..2950a01 100644 --- a/src/protobuf/_from_json.py +++ b/src/protobuf/_from_json.py @@ -16,11 +16,22 @@ 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 +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 = field(default_factory=Budget) + """Allocation budget shared across the whole parse. Unlimited by default.""" 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,8 @@ 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. Raises: json.JSONDecodeError: If json_source is not valid JSON. @@ -111,7 +127,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 +198,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 +216,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 +228,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: 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 +250,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 +267,8 @@ def _read_message_extension( if _is_resetting_null(field_value.message, json): del msg[ext.type] else: + budget = opts.budget + budget.charge_message(field_value.message) value = field_value.message.type() _read_message(value, json, opts) msg[ext.type] = value @@ -258,21 +285,31 @@ 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 + 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: + 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: 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 field_value.oneof is not None: + budget.charge(ONEOF_SIZE) + msg._set_member(field, value) def _read_enum_field( @@ -285,8 +322,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 field_value.oneof is not None: + opts.budget.charge(ONEOF_SIZE) msg._set_member(field, value) @@ -301,12 +340,12 @@ 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: + budget.charge(LIST_SLOT_SIZE) + list_.append(value) def _read_map_field( @@ -320,10 +359,12 @@ 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: + budget.charge(DICT_ENTRY_SIZE) dict_[key] = value @@ -337,17 +378,20 @@ 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: + budget.charge_message(field_value.message) + value = field_value.message.type() _read_message(value, json, opts) + 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 + field: DescField, field_value: DescFieldValueMap, json: JsonValue, budget: Budget ) -> bool | int | str: match field_value.key: case ScalarType.BOOL: @@ -357,13 +401,16 @@ 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) + 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 _: + budget.charge(INT_SIZE) return _read_int(field, field_value.key, json) @@ -374,13 +421,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): + budget = opts.budget + budget.charge_message(element) 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 +440,7 @@ def _read_enum( desc: DescEnum, json: JsonValue, ignore_unknown_fields: bool, # noqa: FBT001 + budget: Budget, ) -> Enum | None: if json is None: return desc.type(desc.values[0].number) @@ -402,6 +452,7 @@ def _read_enum( if ignore_unknown_fields: return None # Succeeds for open enum, raises an error for closed + 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 +467,17 @@ def _read_enum( def _read_scalar( + desc: DescField | DescExtension, + scalar_type: ScalarType, + json: JsonValue, + budget: Budget, +) -> bool | int | float | str | bytes: + value = _read_scalar_value(desc, scalar_type, json) + 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 +667,12 @@ 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(): + 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)) msg.fields[k] = val @@ -624,9 +689,12 @@ 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: + budget.charge_message(element_desc) val = cast("Value", element_desc.type()) _value_from_json(val, e, opts, element_wkt) + budget.charge(LIST_SLOT_SIZE) msg.values.append(val) @@ -642,6 +710,8 @@ def _value_from_json( def _value_from_json_inner( msg: Value, json: JsonValue, opts: FromJsonOptions, wkt: WktValue ) -> None: + budget = opts.budget + budget.charge(ONEOF_SIZE) match json: case None: msg.kind = Oneof( @@ -650,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) lv_desc = wkt.list_value.message lv_wkt = match_wkt(lv_desc) assert isinstance(lv_wkt, WktListValue) # noqa: S101 @@ -661,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) struct_desc = wkt.struct_value.message struct_wkt = match_wkt(struct_desc) assert isinstance(struct_wkt, WktStruct) # noqa: S101 @@ -677,6 +751,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 +775,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..384b30a 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, @@ -584,6 +585,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 +598,9 @@ 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. 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,16 @@ 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. """ 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 +690,20 @@ 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: + 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(self._desc) + 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,6 +713,7 @@ 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 @@ -689,8 +721,13 @@ def _merge_from_json( from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 json_value = parse_json(json) + budget = Budget(allocation_limit) + # See _merge_from_binary for why the root is charged here. + budget.charge_message(self._desc) 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 +767,21 @@ def _from_json_value( *, ignore_unknown_fields: bool = False, registry: Registry | None = None, + allocation_limit: int | None = None, ) -> Self: # Needs to be lazy import since JSON specially handles many WKTs. from ._from_json import FromJsonOptions, _read_message # noqa: PLC0415 + budget = Budget(allocation_limit) + budget.charge_message(cls._desc) 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/_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 e954183..9993ea2 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, @@ -98,7 +99,7 @@ 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 .wkt import Timestamp # noqa: PLC0415 value = cast("Timestamp", msg) @@ -128,6 +129,7 @@ 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) + opts.budget.charge(2 * INT_SIZE) value.seconds = Timestamp.from_datetime(dt).seconds value.nanos = nanos return True @@ -170,7 +172,7 @@ 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 .wkt._mixin._const import DURATION_SECONDS_MAX # noqa: PLC0415 value = cast("Duration", msg) @@ -190,6 +192,7 @@ 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 + opts.budget.charge(2 * INT_SIZE) value.seconds = seconds value.nanos = nanos return True @@ -248,6 +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) message = desc.type() if _has_custom_json(desc) and "value" in json: _read_message(message, json["value"], opts) @@ -256,6 +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) + 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 @@ -283,7 +289,7 @@ 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 ._names import proto_snake_case # noqa: PLC0415 value = cast("FieldMask", msg) @@ -296,6 +302,7 @@ 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) + opts.budget.charge(STR_OVERHEAD + len(path) + LIST_SLOT_SIZE) value.paths.append(proto_snake_case(path)) return True @@ -378,13 +385,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_budget.py b/tests/test_budget.py new file mode 100644 index 0000000..f995300 --- /dev/null +++ b/tests/test_budget.py @@ -0,0 +1,94 @@ +# 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, cast + +import pytest + +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 + +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__ + + +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], +) +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 + 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(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: + 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 diff --git a/tests/test_from_binary.py b/tests/test_from_binary.py index 363e2b2..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 @@ -398,3 +399,55 @@ 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) + + 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 f15c6b2..129e8fc 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 FileDescriptorSet, Struct from .gen.enums_pb import ClosedColor, Color, EnumMessage from .gen.json_enum_names_pb import JsonEnumNames, Season @@ -824,3 +825,56 @@ 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 + + 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)