diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b43dee..fe4ccf3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -53,6 +53,22 @@ format changes. - **`Violation.kind` values `"default_option"`, `"retired_number"` and `"well_known_shape"`,** for the bind-time checks above. They came back as `""`. +- **Per-call decode limits** (#36), the six of HARDENING.md § Mandatory + limits, as keyword arguments on `pxf.unmarshal`, `unmarshal_full`, their + `_bytes` forms, `bind_row` and `DatasetReader.from_bytes`: + `max_message_size`, `max_nesting_depth`, `max_numeric_literal_digits`, + `max_bytes_literal_length`, `max_repeated_count` and + `max_big_float_exponent`. 0, the default, keeps protowire-cpp's limit; a + negative value is a `ValueError`. `sbe.Codec` and `Codec.from_message` take + `max_message_size` and `max_repeated_count` for every call, as + protowire-cpp's `sbe::CodecOptions` does. `DatasetReader.scan` binds each + row under the reader's limits. Two caveats: the decoded message still loads + through the protobuf runtime, which stops at 100 levels of submessages (#40), + and the reader parses rows under the default `max_bytes_literal_length` + until protowire-cpp#160 is fixed (`scan` still applies it when binding). +- **`scripts/check_decode.py --limit Name=N`**, so the corpus rows that lower + a limit run on this port. pb rows take none: the protobuf runtime has no + such settings, and such a row exits 2. - **`max_nesting_depth` on `pxf.marshal` and `pxf.marshal_bytes`** (#34, protowire#502), the deepest PXF they write, counted as a decoder counts it. 0, the default, means 100. Raise it only together with the limit of every diff --git a/scripts/check_decode.py b/scripts/check_decode.py index d3cf391..8fc56ce 100644 --- a/scripts/check_decode.py +++ b/scripts/check_decode.py @@ -11,7 +11,8 @@ check_decode.py --format \\ --schema \\ --proto \\ - --input + --input \\ + [--limit = ...] Exit 0 -> input was accepted (decode succeeded) Exit 1 -> input was rejected (decode returned a clean error; the first @@ -42,8 +43,14 @@ Schemas come from the ``adversarial.binpb`` FileDescriptorSet next to ``--proto``, since the Python runtime cannot compile a .proto. The corpus generator writes both, so there are no hand-written bindings to keep in step -with the corpus. ``--limit`` is not implemented: rows that need it list -"python" in their manifest "skip". +with the corpus. + +``--limit MaxMessageSize=1024`` sets one HARDENING limit for the row, under +the Go name the manifest uses. pxf rows take all six, sbe rows +MaxMessageSize and MaxRepeatedCount, and pb-render MaxNestingDepth. A limit +this port cannot apply to the row's format exits 2: running the row without +it would not test the limit. pb rows take none, since the protobuf runtime +has no such settings. """ from __future__ import annotations @@ -79,16 +86,47 @@ def _message_class(fds_bytes: bytes, schema: str): return message_factory.GetMessageClass(pool.FindMessageTypeByName(schema)) -def _decode_pxf(fds_bytes: bytes, schema: str, data: bytes) -> int: +# --limit names, as the manifest spells them, and the keywords they set. +_LIMITS = { + "MaxMessageSize": "max_message_size", + "MaxNestingDepth": "max_nesting_depth", + "MaxNumericLiteralDigits": "max_numeric_literal_digits", + "MaxBytesLiteralLength": "max_bytes_literal_length", + "MaxRepeatedCount": "max_repeated_count", + "MaxBigFloatExponent": "max_big_float_exponent", +} +# The limits each format can apply. pb has none: see the module docstring. +_FORMAT_LIMITS = { + "pxf": set(_LIMITS.values()), + "sbe": {"max_message_size", "max_repeated_count"}, + "pb-render": {"max_nesting_depth"}, + "pb": set(), + "envelope": set(), +} + + +def _parse_limits(fmt: str, specs: list[str]) -> dict[str, int]: + out = {} + for spec in specs: + name, _, value = spec.partition("=") + if name not in _LIMITS: + raise SystemExit(f"check_decode: unknown limit {name!r}") + if _LIMITS[name] not in _FORMAT_LIMITS[fmt]: + raise SystemExit(f"check_decode: --limit {name} is not supported for {fmt}") + out[_LIMITS[name]] = int(value) + return out + + +def _decode_pxf(fds_bytes: bytes, schema: str, data: bytes, limits: dict[str, int]) -> int: pxf.unmarshal_bytes(b"", fds_bytes, schema) # Binds, or the harness fails. try: - pxf.unmarshal_bytes(data, fds_bytes, schema) + pxf.unmarshal_bytes(data, fds_bytes, schema, **limits) except ValueError as e: return _reject(f"pxf: {e}") return 0 -def _decode_pb(fds_bytes: bytes, schema: str, data: bytes) -> int: +def _decode_pb(fds_bytes: bytes, schema: str, data: bytes, limits: dict[str, int]) -> int: msg = _message_class(fds_bytes, schema)() try: msg.ParseFromString(data) @@ -103,10 +141,10 @@ def _decode_pb(fds_bytes: bytes, schema: str, data: bytes) -> int: _PB_READ_FAILED = "pxf.marshal: proto parse failed" -def _pb_render(fds_bytes: bytes, schema: str, data: bytes) -> int: +def _pb_render(fds_bytes: bytes, schema: str, data: bytes, limits: dict[str, int]) -> int: _message_class(fds_bytes, schema) # Or the harness fails. try: - pxf.marshal_bytes(data, fds_bytes, schema) + pxf.marshal_bytes(data, fds_bytes, schema, **limits) except ValueError as e: if str(e).startswith(_PB_READ_FAILED): raise # The encoder never ran: exit 2, not a reject. @@ -114,9 +152,9 @@ def _pb_render(fds_bytes: bytes, schema: str, data: bytes) -> int: return 0 -def _decode_sbe(fds_bytes: bytes, schema: str, data: bytes) -> int: +def _decode_sbe(fds_bytes: bytes, schema: str, data: bytes, limits: dict[str, int]) -> int: cls = _message_class(fds_bytes, schema) - codec = sbe.Codec.from_message(cls) + codec = sbe.Codec.from_message(cls, **limits) try: codec.unmarshal(data, cls()) except ValueError as e: @@ -133,7 +171,8 @@ def _run(args: argparse.Namespace) -> int: "pb-render": _pb_render, "sbe": _decode_sbe, }[args.format] - return decode(_load_fds(args.proto), args.schema, Path(args.input).read_bytes()) + limits = _parse_limits(args.format, args.limit) + return decode(_load_fds(args.proto), args.schema, Path(args.input).read_bytes(), limits) def main(argv: list[str] | None = None) -> int: @@ -142,6 +181,7 @@ def main(argv: list[str] | None = None) -> int: p.add_argument("--schema", required=True) p.add_argument("--proto", required=True) p.add_argument("--input", required=True) + p.add_argument("--limit", action="append", default=[], metavar="NAME=N") args = p.parse_args(argv) try: return _run(args) diff --git a/src/_protowire/module.cc b/src/_protowire/module.cc index ae30a5d..a319fe8 100644 --- a/src/_protowire/module.cc +++ b/src/_protowire/module.cc @@ -132,16 +132,42 @@ nb::object CellToPyTuple(const std::optional& cell) { *cell); } +// The decode limits of HARDENING.md § Mandatory limits, as Python passes them: +// (max_message_size, max_nesting_depth, max_numeric_literal_digits, +// max_bytes_literal_length, max_repeated_count, max_big_float_exponent). +// 0 selects the default in protowire/limits.h. +using PxfLimits = std::tuple; + +void ApplyLimits(const PxfLimits& l, protowire::pxf::UnmarshalOptions& opts) { + static constexpr const char* kNames[] = { + "max_message_size", "max_nesting_depth", "max_numeric_literal_digits", + "max_bytes_literal_length", "max_repeated_count", "max_big_float_exponent"}; + const int v[] = {std::get<0>(l), std::get<1>(l), std::get<2>(l), + std::get<3>(l), std::get<4>(l), std::get<5>(l)}; + for (int i = 0; i < 6; ++i) { + if (v[i] < 0) { + throw nb::value_error((std::string(kNames[i]) + " must not be negative").c_str()); + } + } + opts.max_message_size = v[0]; + opts.max_nesting_depth = v[1]; + opts.max_numeric_literal_digits = v[2]; + opts.max_bytes_literal_length = v[3]; + opts.max_repeated_count = v[4]; + opts.max_big_float_exponent = v[5]; +} + // PXF text -> binary proto bytes. nb::bytes PxfUnmarshal(nb::bytes text, nb::bytes fds_bytes, const std::string& full_name, bool discard_unknown, - bool skip_validate) { + bool skip_validate, const PxfLimits& limits) { auto schema = BuildSchema(std::string_view(fds_bytes.c_str(), fds_bytes.size())); const auto* desc = FindDescriptor(schema, full_name); std::unique_ptr msg( schema.factory->GetPrototype(desc)->New()); protowire::pxf::UnmarshalOptions opts; + ApplyLimits(limits, opts); opts.discard_unknown = discard_unknown; opts.skip_validate = skip_validate; auto st = protowire::pxf::Unmarshal( @@ -174,13 +200,14 @@ std::tuple, std::vector, std::vector> PxfUnmarshalFull(nb::bytes text, nb::bytes fds_bytes, const std::string& full_name, bool discard_unknown, - bool skip_validate) { + bool skip_validate, const PxfLimits& limits) { auto schema = BuildSchema(std::string_view(fds_bytes.c_str(), fds_bytes.size())); const auto* desc = FindDescriptor(schema, full_name); std::unique_ptr msg( schema.factory->GetPrototype(desc)->New()); protowire::pxf::UnmarshalOptions opts; + ApplyLimits(limits, opts); opts.discard_unknown = discard_unknown; opts.skip_validate = skip_validate; auto r = protowire::pxf::UnmarshalFull( @@ -273,11 +300,13 @@ PxfValidateDescriptor(nb::bytes fds_bytes, const std::string& full_name) { // streambuf bridge is a possible follow-up. class PyDatasetReader { public: - static std::unique_ptr FromBytes(nb::bytes data) { + static std::unique_ptr FromBytes(nb::bytes data, const PxfLimits& limits) { + protowire::pxf::UnmarshalOptions opts; + ApplyLimits(limits, opts); auto out = std::unique_ptr(new PyDatasetReader()); out->stream_ = std::make_unique( std::string(data.c_str(), data.size())); - auto tr = protowire::pxf::DatasetReader::Create(out->stream_.get()); + auto tr = protowire::pxf::DatasetReader::Create(out->stream_.get(), opts); if (!tr.ok()) { throw nb::value_error(("pxf.DatasetReader: " + tr.status().ToString()).c_str()); } @@ -388,8 +417,17 @@ std::string PxfMarshal(nb::bytes msg_bytes, nb::bytes fds_bytes, // SbeCodec wraps protowire::sbe::Codec along with the descriptor pool. class SbeCodec { public: + // max_message_size and max_repeated_count are sbe::CodecOptions'; 0 selects + // the default in protowire/limits.h. static std::unique_ptr Create( - nb::bytes fds_bytes, std::vector file_names) { + nb::bytes fds_bytes, std::vector file_names, + int max_message_size, int max_repeated_count) { + if (max_message_size < 0) { + throw nb::value_error("max_message_size must not be negative"); + } + if (max_repeated_count < 0) { + throw nb::value_error("max_repeated_count must not be negative"); + } auto out = std::make_unique(); out->schema_ = BuildSchema(std::string_view(fds_bytes.c_str(), fds_bytes.size())); @@ -408,7 +446,10 @@ class SbeCodec { } if (selected.empty()) selected = out->schema_.files; - auto codec = protowire::sbe::Codec::New(selected); + protowire::sbe::CodecOptions opts; + opts.max_message_size = max_message_size; + opts.max_repeated_count = max_repeated_count; + auto codec = protowire::sbe::Codec::New(selected, opts); if (!codec.ok()) { throw nb::value_error( ("sbe.Codec: " + codec.status().ToString()).c_str()); @@ -541,16 +582,18 @@ inline SbeCodec::PyGroupView SbeCodec::PyView::Group(const std::string& name) co NB_MODULE(_protowire, m) { m.doc() = "protowire native extension (nanobind shim around protowire-cpp)"; + const PxfLimits no_limits{0, 0, 0, 0, 0, 0}; m.def("pxf_unmarshal", &PxfUnmarshal, "text"_a, "fds"_a, "full_name"_a, - "discard_unknown"_a = false, "skip_validate"_a = false); + "discard_unknown"_a = false, "skip_validate"_a = false, "limits"_a = no_limits); m.def("pxf_unmarshal_full", &PxfUnmarshalFull, "text"_a, "fds"_a, - "full_name"_a, "discard_unknown"_a = false, "skip_validate"_a = false); + "full_name"_a, "discard_unknown"_a = false, "skip_validate"_a = false, + "limits"_a = no_limits); m.def("pxf_marshal", &PxfMarshal, "msg_bytes"_a, "fds"_a, "full_name"_a, "max_nesting_depth"_a = 0); m.def("pxf_validate_descriptor", &PxfValidateDescriptor, "fds"_a, "full_name"_a); nb::class_(m, "PxfDatasetReader") - .def_static("from_bytes", &PyDatasetReader::FromBytes, "data"_a) + .def_static("from_bytes", &PyDatasetReader::FromBytes, "data"_a, "limits"_a = no_limits) .def_prop_ro("type", &PyDatasetReader::Type) .def_prop_ro("columns", &PyDatasetReader::Columns) .def_prop_ro("directives", &PyDatasetReader::Directives) @@ -561,7 +604,8 @@ NB_MODULE(_protowire, m) { .def("__next__", &PyDatasetReader::Next); nb::class_(m, "SbeCodec") - .def_static("create", &SbeCodec::Create, "fds"_a, "file_names"_a) + .def_static("create", &SbeCodec::Create, "fds"_a, "file_names"_a, + "max_message_size"_a = 0, "max_repeated_count"_a = 0) .def("marshal", &SbeCodec::Marshal, "msg_bytes"_a, "full_name"_a) .def("unmarshal", &SbeCodec::Unmarshal, "data"_a, "full_name"_a) .def("new_view", &SbeCodec::NewView, "data"_a); diff --git a/src/protowire/pxf.py b/src/protowire/pxf.py index a314876..3bc8649 100644 --- a/src/protowire/pxf.py +++ b/src/protowire/pxf.py @@ -194,6 +194,38 @@ def marshal(msg: Message, *, max_nesting_depth: int = 0) -> str: ) +# --- Decode limits ---------------------------------------------------------- +# +# Every decoder takes the limits of HARDENING.md § Mandatory limits as keyword +# arguments, passed to protowire-cpp's pxf::UnmarshalOptions. 0 selects the +# default; a negative value is a ValueError. +# +# max_message_size input length in bytes (default 64 MiB) +# max_nesting_depth `{` / `[` nesting, the root is depth 0 (100) +# max_numeric_literal_digits digits of a literal bound to pxf.BigInt, +# Decimal or BigFloat, and a Decimal's scale (4096) +# max_bytes_literal_length decoded length of a b"..." literal (64 MiB) +# max_repeated_count elements of one repeated or map field (64 Mi) +# max_big_float_exponent a pxf.BigFloat's binary exponent (16384) +# +# The decoded message reaches Python through the protobuf runtime, which stops +# at 100 levels of submessages whatever max_nesting_depth says (#40). + + +def _limits( + max_message_size: int, + max_nesting_depth: int, + max_numeric_literal_digits: int, + max_bytes_literal_length: int, + max_repeated_count: int, + max_big_float_exponent: int, +) -> tuple[int, int, int, int, int, int]: + return ( + max_message_size, max_nesting_depth, max_numeric_literal_digits, + max_bytes_literal_length, max_repeated_count, max_big_float_exponent, + ) + + # --- bytes-only helpers (used by the CLI / advanced callers) ------------- @@ -216,11 +248,24 @@ def unmarshal_bytes( *, discard_unknown: bool = False, skip_validate: bool = False, + max_message_size: int = 0, + max_nesting_depth: int = 0, + max_numeric_literal_digits: int = 0, + max_bytes_literal_length: int = 0, + max_repeated_count: int = 0, + max_big_float_exponent: int = 0, ) -> bytes: - """Decode PXF text into raw proto-binary bytes against an explicit FDS.""" + """Decode PXF text into raw proto-binary bytes against an explicit FDS. + + The limits are described under "Decode limits" above. + """ text = data.encode("utf-8") if isinstance(data, str) else bytes(data) return _protowire.pxf_unmarshal( - text, bytes(fds), full_name, discard_unknown, skip_validate + text, bytes(fds), full_name, discard_unknown, skip_validate, + _limits( + max_message_size, max_nesting_depth, max_numeric_literal_digits, + max_bytes_literal_length, max_repeated_count, max_big_float_exponent, + ), ) @@ -231,10 +276,20 @@ def unmarshal_full_bytes( *, discard_unknown: bool = False, skip_validate: bool = False, + max_message_size: int = 0, + max_nesting_depth: int = 0, + max_numeric_literal_digits: int = 0, + max_bytes_literal_length: int = 0, + max_repeated_count: int = 0, + max_big_float_exponent: int = 0, ) -> tuple[bytes, Result]: text = data.encode("utf-8") if isinstance(data, str) else bytes(data) raw, set_paths, null_paths, dirs, tables, protos = _protowire.pxf_unmarshal_full( - text, bytes(fds), full_name, discard_unknown, skip_validate + text, bytes(fds), full_name, discard_unknown, skip_validate, + _limits( + max_message_size, max_nesting_depth, max_numeric_literal_digits, + max_bytes_literal_length, max_repeated_count, max_big_float_exponent, + ), ) return raw, _wrap_result(set_paths, null_paths, dirs, tables, protos) @@ -248,12 +303,25 @@ def unmarshal( *, discard_unknown: bool = False, skip_validate: bool = False, + max_message_size: int = 0, + max_nesting_depth: int = 0, + max_numeric_literal_digits: int = 0, + max_bytes_literal_length: int = 0, + max_repeated_count: int = 0, + max_big_float_exponent: int = 0, ) -> None: - """Decode PXF text into `msg` (in place). Mirrors Go pxf.Unmarshal.""" + """Decode PXF text into `msg` (in place). Mirrors Go pxf.Unmarshal. + + The limits are described under "Decode limits" above. + """ text = data.encode("utf-8") if isinstance(data, str) else bytes(data) fds = fds_for_message(msg) raw = _protowire.pxf_unmarshal( - text, fds, msg.DESCRIPTOR.full_name, discard_unknown, skip_validate + text, fds, msg.DESCRIPTOR.full_name, discard_unknown, skip_validate, + _limits( + max_message_size, max_nesting_depth, max_numeric_literal_digits, + max_bytes_literal_length, max_repeated_count, max_big_float_exponent, + ), ) msg.Clear() msg.MergeFromString(raw) @@ -265,16 +333,27 @@ def unmarshal_full( *, discard_unknown: bool = False, skip_validate: bool = False, + max_message_size: int = 0, + max_nesting_depth: int = 0, + max_numeric_literal_digits: int = 0, + max_bytes_literal_length: int = 0, + max_repeated_count: int = 0, + max_big_float_exponent: int = 0, ) -> Result: """Decode PXF + return per-field presence (set/null) metadata and any `@` / `@dataset` directives the decoder saw at the document root. - Mirrors Go pxf.UnmarshalFull. + Mirrors Go pxf.UnmarshalFull. The limits are described under "Decode + limits" above. """ text = data.encode("utf-8") if isinstance(data, str) else bytes(data) fds = fds_for_message(msg) raw, set_paths, null_paths, dirs, tables, protos = _protowire.pxf_unmarshal_full( - text, fds, msg.DESCRIPTOR.full_name, discard_unknown, skip_validate + text, fds, msg.DESCRIPTOR.full_name, discard_unknown, skip_validate, + _limits( + max_message_size, max_nesting_depth, max_numeric_literal_digits, + max_bytes_literal_length, max_repeated_count, max_big_float_exponent, + ), ) msg.Clear() msg.MergeFromString(raw) @@ -364,16 +443,39 @@ class DatasetReader: file-like / chunked-IO bridge is a possible follow-up. """ - def __init__(self, native: object) -> None: + def __init__(self, native: object, limits: dict[str, int] | None = None) -> None: # Private — construct via `from_bytes`. We keep the native handle - # only; type/columns/directives are forwarded on demand. + # only; type/columns/directives are forwarded on demand. `limits` are + # the from_bytes keywords, which scan passes on to bind_row. self._native = native + self._limits = limits or {} @classmethod - def from_bytes(cls, data: bytes | str) -> "DatasetReader": + def from_bytes( + cls, + data: bytes | str, + *, + max_message_size: int = 0, + max_nesting_depth: int = 0, + max_numeric_literal_digits: int = 0, + max_bytes_literal_length: int = 0, + max_repeated_count: int = 0, + max_big_float_exponent: int = 0, + ) -> "DatasetReader": + """Read `data`'s `@dataset` header. The limits are described under + "Decode limits" above; `scan` binds each row under the same ones.""" if isinstance(data, str): data = data.encode("utf-8") - return cls(_protowire.PxfDatasetReader.from_bytes(bytes(data))) + limits = dict( + max_message_size=max_message_size, + max_nesting_depth=max_nesting_depth, + max_numeric_literal_digits=max_numeric_literal_digits, + max_bytes_literal_length=max_bytes_literal_length, + max_repeated_count=max_repeated_count, + max_big_float_exponent=max_big_float_exponent, + ) + native = _protowire.PxfDatasetReader.from_bytes(bytes(data), _limits(**limits)) + return cls(native, limits) @property def type(self) -> str: @@ -437,7 +539,7 @@ def scan(self, msg: Message, *, skip_validate: bool = True) -> bool: cells = self.next_or_none() if cells is None: return False - bind_row(msg, self.columns, cells, skip_validate=skip_validate) + bind_row(msg, self.columns, cells, skip_validate=skip_validate, **self._limits) return True @@ -450,6 +552,12 @@ def bind_row( row: tuple[Cell, ...] | list[Cell], *, skip_validate: bool = True, + max_message_size: int = 0, + max_nesting_depth: int = 0, + max_numeric_literal_digits: int = 0, + max_bytes_literal_length: int = 0, + max_repeated_count: int = 0, + max_big_float_exponent: int = 0, ) -> None: """Bind a `@dataset` row's cells to `msg`'s fields by column name. @@ -470,7 +578,8 @@ def bind_row( `skip_validate` defaults to True: the descriptor was validated once when the caller constructed the message factory, and re-running the - reserved-name check per row is wasteful in tight loops. + reserved-name check per row is wasteful in tight loops. The limits apply + to the synthetic body, as for `unmarshal`. """ if len(columns) != len(row): raise ValueError( @@ -481,7 +590,15 @@ def bind_row( if cell is None: continue parts.append(f"{col} = {_cell_to_pxf(cell)}\n") - unmarshal("".join(parts), msg, skip_validate=skip_validate) + unmarshal( + "".join(parts), msg, skip_validate=skip_validate, + max_message_size=max_message_size, + max_nesting_depth=max_nesting_depth, + max_numeric_literal_digits=max_numeric_literal_digits, + max_bytes_literal_length=max_bytes_literal_length, + max_repeated_count=max_repeated_count, + max_big_float_exponent=max_big_float_exponent, + ) def _cell_to_pxf(cell: tuple[CellKind, object]) -> str: diff --git a/src/protowire/sbe.py b/src/protowire/sbe.py index 1740de5..5af4a98 100644 --- a/src/protowire/sbe.py +++ b/src/protowire/sbe.py @@ -73,9 +73,23 @@ def entry(self, i: int) -> View: class Codec: - """SBE codec built from one or more proto FileDescriptors with SBE annotations.""" + """SBE codec built from one or more proto FileDescriptors with SBE annotations. + + The codec's decoders (`unmarshal` and `view`) enforce two limits of + HARDENING.md § Mandatory limits, set here for every call, as + protowire-cpp's sbe::CodecOptions sets them. `max_message_size` caps the + input to one call (default 64 MiB). `max_repeated_count` caps a group's + entry count, checked before any entry is read (default 64 Mi). 0 selects + the default; a negative value is a ValueError. + """ - def __init__(self, files: Iterable[FileDescriptor]) -> None: + def __init__( + self, + files: Iterable[FileDescriptor], + *, + max_message_size: int = 0, + max_repeated_count: int = 0, + ) -> None: files = list(files) if not files: raise ValueError("sbe.Codec requires at least one FileDescriptor") @@ -83,11 +97,23 @@ def __init__(self, files: Iterable[FileDescriptor]) -> None: # Pass the *selected* file names so the C++ Codec only registers those # — transitive dep files (descriptor.proto, sbe/annotations.proto) # don't carry an (sbe.schema_id) option and would otherwise fail. - self._impl = _protowire.SbeCodec.create(fds, [f.name for f in files]) + self._impl = _protowire.SbeCodec.create( + fds, [f.name for f in files], max_message_size, max_repeated_count + ) @classmethod - def from_message(cls, msg_type: type[Message]) -> "Codec": - return cls([msg_type.DESCRIPTOR.file]) + def from_message( + cls, + msg_type: type[Message], + *, + max_message_size: int = 0, + max_repeated_count: int = 0, + ) -> "Codec": + return cls( + [msg_type.DESCRIPTOR.file], + max_message_size=max_message_size, + max_repeated_count=max_repeated_count, + ) def marshal(self, msg: Message) -> bytes: return self._impl.marshal(msg.SerializeToString(), msg.DESCRIPTOR.full_name) diff --git a/tests/test_pxf_decode_limits.py b/tests/test_pxf_decode_limits.py new file mode 100644 index 0000000..b9ed908 --- /dev/null +++ b/tests/test_pxf_decode_limits.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 TrendVidia, LLC. +"""Per-call decode limits (HARDENING.md § Mandatory limits, #36). + +Each limit is a keyword argument passed to protowire-cpp's +pxf::UnmarshalOptions. For each one, a value below the default refuses an +input it would accept, a value at the input's size accepts it again, and a +value above the default accepts an input the default refuses. + +max_repeated_count is not raised here: its default is 64 Mi elements, too +many for a unit test. +""" + +from __future__ import annotations + +import base64 + +import pytest +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory + +from protowire import pxf + +F = descriptor_pb2.FieldDescriptorProto +MiB64 = 64 << 20 + + +def _schema(): + # pxf/bignum.proto, as protowire-cpp ships it. + bignum = descriptor_pb2.FileDescriptorProto(name="pxf/bignum.proto", package="pxf", + syntax="proto3") + bi = bignum.message_type.add(name="BigInt") + bi.field.add(name="abs", number=1, type=F.TYPE_BYTES) + bi.field.add(name="negative", number=2, type=F.TYPE_BOOL) + bf = bignum.message_type.add(name="BigFloat") + bf.field.add(name="mantissa", number=1, type=F.TYPE_BYTES) + bf.field.add(name="exponent", number=2, type=F.TYPE_INT32) + bf.field.add(name="prec", number=3, type=F.TYPE_UINT32) + bf.field.add(name="negative", number=4, type=F.TYPE_BOOL) + + fd = descriptor_pb2.FileDescriptorProto(name="limits.proto", package="lim", syntax="proto3") + fd.dependency.append("pxf/bignum.proto") + h = fd.message_type.add(name="H") + h.field.add(name="s", number=1, type=F.TYPE_STRING) + h.field.add(name="b", number=2, type=F.TYPE_BYTES) + h.field.add(name="l", number=3, type=F.TYPE_INT32, label=F.LABEL_REPEATED) + h.field.add(name="c", number=4, type=F.TYPE_MESSAGE, type_name=".lim.H") + h.field.add(name="kids", number=5, type=F.TYPE_MESSAGE, label=F.LABEL_REPEATED, + type_name=".lim.H") + h.field.add(name="i", number=6, type=F.TYPE_MESSAGE, type_name=".pxf.BigInt") + h.field.add(name="f", number=7, type=F.TYPE_MESSAGE, type_name=".pxf.BigFloat") + + pool = descriptor_pool.DescriptorPool() + pool.Add(bignum) + pool.Add(fd) + return message_factory.GetMessageClass(pool.FindMessageTypeByName("lim.H")) + + +H = _schema() + + +def _decode(text: str, **limits) -> None: + pxf.unmarshal(text, H(), **limits) + + +def _refused(text: str, limit_name: str, **limits) -> None: + with pytest.raises(ValueError, match=limit_name): + _decode(text, **limits) + + +def _chain(levels: int) -> str: + """PXF nested `levels` deep through the singular field c.""" + return "c { " * levels + "s = \"x\"" + " }" * levels + "\n" + + +def _list_chain(elements: int) -> str: + """PXF nested 2 * `elements` deep (each `[` and `{` is a level), but only + `elements` protobuf levels, which the protobuf runtime can load (#40).""" + return "kids = [ { " * elements + "s = \"x\"" + " } ]" * elements + "\n" + + +def _bytes_doc(n: int) -> str: + return 'b = b"' + base64.b64encode(b"\0" * n).decode() + '"\n' + + +# --- max_message_size -------------------------------------------------------- + + +def test_message_size_lowered(): + doc = 's = "' + "a" * 2000 + '"\n' + _refused(doc, "MaxMessageSize=1024", max_message_size=1024) + _decode(doc, max_message_size=len(doc)) + + +def test_message_size_raised(): + doc = 's = "' + "a" * MiB64 + '"\n' + _refused(doc, "MaxMessageSize") + _decode(doc, max_message_size=len(doc)) + + +# --- max_nesting_depth ------------------------------------------------------- + + +def test_nesting_depth_lowered(): + _refused(_chain(5), "MaxNestingDepth=4", max_nesting_depth=4) + _decode(_chain(5), max_nesting_depth=5) + + +def test_nesting_depth_raised(): + doc = _list_chain(60) # depth 120 + _refused(doc, "MaxNestingDepth") + _decode(doc, max_nesting_depth=120) + _refused(doc, "MaxNestingDepth=119", max_nesting_depth=119) + + +# --- max_numeric_literal_digits ---------------------------------------------- + + +def test_numeric_literal_digits_lowered(): + doc = "i = " + "1" * 10 + "\n" + _refused(doc, "MaxNumericLiteralDigits", max_numeric_literal_digits=9) + _decode(doc, max_numeric_literal_digits=10) + + +def test_numeric_literal_digits_raised(): + doc = "i = " + "1" * 5000 + "\n" + _refused(doc, "MaxNumericLiteralDigits") + _decode(doc, max_numeric_literal_digits=5000) + + +# --- max_bytes_literal_length ------------------------------------------------ + + +def test_bytes_literal_length_lowered(): + doc = _bytes_doc(192) # unpadded base64: see trendvidia/protowire#588 + _refused(doc, "MaxBytesLiteralLength=191", max_bytes_literal_length=191) + _decode(doc, max_bytes_literal_length=192) + + +def test_bytes_literal_length_raised(): + n = MiB64 + 2 # a multiple of 3, so its base64 has no padding (protowire#588) + doc = _bytes_doc(n) + _refused(doc, "MaxBytesLiteralLength", max_message_size=2 * MiB64) + _decode(doc, max_message_size=2 * MiB64, max_bytes_literal_length=n) + + +# --- max_repeated_count ------------------------------------------------------ + + +def test_repeated_count_lowered(): + doc = "l = [" + ", ".join(["1"] * 16) + "]\n" + _refused(doc, "MaxRepeatedCount=15", max_repeated_count=15) + _decode(doc, max_repeated_count=16) + + +# --- max_big_float_exponent -------------------------------------------------- + + +def test_big_float_exponent_lowered(): + doc = "f = 1e100\n" # binary exponent about 333 + _refused(doc, "exponent", max_big_float_exponent=300) + _decode(doc, max_big_float_exponent=400) + + +def test_big_float_exponent_raised(): + doc = "f = 1e5000\n" # binary exponent about 16610, past the default 16384 + _refused(doc, "exponent") + _decode(doc, max_big_float_exponent=17000) + + +# --- every entry point -------------------------------------------------------- + + +def _fds() -> bytes: + fds = descriptor_pb2.FileDescriptorSet() + for name in ("pxf/bignum.proto", "limits.proto"): + fds.file.add().MergeFromString(H.DESCRIPTOR.file.pool.FindFileByName(name).serialized_pb) + return fds.SerializeToString() + + +_DOC = 's = "' + "x" * 40 + '"\n' +_ENTRY_POINTS = { + "unmarshal": lambda **k: pxf.unmarshal(_DOC, H(), **k), + "unmarshal_full": lambda **k: pxf.unmarshal_full(_DOC, H(), **k), + "unmarshal_bytes": lambda **k: pxf.unmarshal_bytes(_DOC, _fds(), "lim.H", **k), + "unmarshal_full_bytes": lambda **k: pxf.unmarshal_full_bytes(_DOC, _fds(), "lim.H", **k), + # The synthetic body bind_row decodes is the same 47 bytes. + "bind_row": lambda **k: pxf.bind_row(H(), ["s"], [("string", "x" * 40)], **k), +} + + +@pytest.mark.parametrize("name", list(_ENTRY_POINTS)) +def test_each_entry_point_takes_the_limits(name): + decode = _ENTRY_POINTS[name] + with pytest.raises(ValueError, match="MaxMessageSize=16"): + decode(max_message_size=16) + decode(max_message_size=64) + + +@pytest.mark.parametrize("name", [ + "max_message_size", "max_nesting_depth", "max_numeric_literal_digits", + "max_bytes_literal_length", "max_repeated_count", "max_big_float_exponent", +]) +def test_a_negative_limit_is_an_error(name): + with pytest.raises(ValueError, match=f"{name} must not be negative"): + _decode('s = "x"\n', **{name: -1}) + + +# --- DatasetReader ------------------------------------------------------------- + + +_BYTES_ROW = "@dataset lim.H (b)\n(" + _bytes_doc(192).split(" = ", 1)[1].strip() + ")\n" + + +@pytest.mark.xfail(strict=True, reason="protowire-cpp#160: rows are read under the default " + "MaxBytesLiteralLength") +def test_dataset_reader_reads_rows_under_its_limits(): + reader = pxf.DatasetReader.from_bytes(_BYTES_ROW, max_bytes_literal_length=191) + with pytest.raises(ValueError, match="MaxBytesLiteralLength=191"): + reader.next_or_none() + + +def test_scan_refuses_a_cell_past_the_readers_limit(): + # Until protowire-cpp#160, the reader returns the row and binding refuses it. + with pytest.raises(ValueError, match="MaxBytesLiteralLength=191"): + pxf.DatasetReader.from_bytes(_BYTES_ROW, max_bytes_literal_length=191).scan(H()) + msg = H() + assert pxf.DatasetReader.from_bytes(_BYTES_ROW, max_bytes_literal_length=192).scan(msg) + assert len(msg.b) == 192 + + +def test_scan_binds_under_the_readers_limits(): + # The reader does not know the column is a BigInt, so only binding checks + # the digits. Under the default of 4096 the cell would be refused. + digits = "1" * 5000 + data = f"@dataset lim.H (i)\n({digits})\n" + with pytest.raises(ValueError, match="MaxNumericLiteralDigits"): + pxf.DatasetReader.from_bytes(data).scan(H()) + msg = H() + assert pxf.DatasetReader.from_bytes(data, max_numeric_literal_digits=5000).scan(msg) + assert msg.HasField("i") diff --git a/tests/test_sbe_limits.py b/tests/test_sbe_limits.py new file mode 100644 index 0000000..fa06290 --- /dev/null +++ b/tests/test_sbe_limits.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 TrendVidia, LLC. +"""sbe.Codec's decode limits (HARDENING.md § Mandatory limits, #36). + +protowire-cpp sets them per codec, in sbe::CodecOptions, and so does this +port. They are tested lowered only. A group's entry count is a uint16 on the +wire, far below the default of 64 Mi, and these messages are a few dozen +bytes, against a default of 64 MiB. +""" + +from __future__ import annotations + +import shutil +import subprocess + +import pytest +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory + +from protowire import sbe +from test_sbe import CPP_PROTO_DIR + +BOOK_PROTO_SRC = """ +syntax = "proto3"; +package sbe.limits; + +import "sbe/annotations.proto"; + +option (sbe.schema_id) = 3; +option (sbe.version) = 0; + +message Book { + option (sbe.template_id) = 1; + uint64 id = 1; + + message Level { + int64 price = 1; + uint32 qty = 2; + } + repeated Level levels = 2; +} +""" + + +@pytest.fixture(scope="module") +def book(tmp_path_factory): + if not CPP_PROTO_DIR.exists(): + pytest.skip(f"protowire-cpp proto/ tree not found at {CPP_PROTO_DIR}") + protoc = shutil.which("protoc") + if not protoc: + pytest.skip("protoc not available") + tmp = tmp_path_factory.mktemp("book") + (tmp / "book.proto").write_text(BOOK_PROTO_SRC) + out = tmp / "book.fds" + res = subprocess.run( + [protoc, "-I", str(CPP_PROTO_DIR), "-I", str(tmp), "--include_imports", + f"--descriptor_set_out={out}", "book.proto"], + capture_output=True, text=True, + ) + if res.returncode != 0: + pytest.fail(f"protoc failed: {res.stderr}") + fds = descriptor_pb2.FileDescriptorSet.FromString(out.read_bytes()) + pool = descriptor_pool.DescriptorPool() + for f in fds.file: + pool.Add(f) + return message_factory.GetMessageClass(pool.FindMessageTypeByName("sbe.limits.Book")) + + +def _encoded(Book, levels: int) -> bytes: + msg = Book(id=1) + for i in range(levels): + msg.levels.add(price=i, qty=1) + return sbe.Codec.from_message(Book).marshal(msg) + + +def test_message_size(book): + data = _encoded(book, 2) + with pytest.raises(ValueError, match="MaxMessageSize"): + sbe.Codec.from_message(book, max_message_size=len(data) - 1).unmarshal(data, book()) + with pytest.raises(ValueError, match="MaxMessageSize"): + sbe.Codec.from_message(book, max_message_size=len(data) - 1).view(data) + got = book() + sbe.Codec.from_message(book, max_message_size=len(data)).unmarshal(data, got) + assert len(got.levels) == 2 + + +def test_repeated_count(book): + data = _encoded(book, 16) + with pytest.raises(ValueError, match="MaxRepeatedCount=8"): + sbe.Codec.from_message(book, max_repeated_count=8).unmarshal(data, book()) + got = book() + sbe.Codec.from_message(book, max_repeated_count=16).unmarshal(data, got) + assert len(got.levels) == 16 + + +@pytest.mark.parametrize("name", ["max_message_size", "max_repeated_count"]) +def test_a_negative_limit_is_an_error(book, name): + with pytest.raises(ValueError, match=f"{name} must not be negative"): + sbe.Codec.from_message(book, **{name: -1})