pxf, sbe: take the decode limits per call; check_decode --limit

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

scripts/check_decode.py

@@ -11,7 +11,8 @@
     check_decode.py --format <pxf|pb|pb-render|sbe|envelope> \\
                     --schema <fully.qualified.MessageType> \\
                     --proto  <path-to-adversarial.proto> \\
-                    --input  <path>
+                    --input  <path> \\
+                    [--limit <Name>=<N> ...]
 
     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,20 +141,20 @@ 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.
         return _reject(f"pb-render: {e}")
     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)

src/_protowire/module.cc

@@ -132,16 +132,42 @@ nb::object CellToPyTuple(const std::optional<protowire::pxf::ValuePtr>& 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<int, int, int, int, int, int>;
+
+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<pbuf::Message> 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<nb::bytes, std::vector<std::string>, std::vector<std::string>,
            std::vector<PyProtoDirective>>
 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<pbuf::Message> 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<PyDatasetReader> FromBytes(nb::bytes data) {
+  static std::unique_ptr<PyDatasetReader> FromBytes(nb::bytes data, const PxfLimits& limits) {
+    protowire::pxf::UnmarshalOptions opts;
+    ApplyLimits(limits, opts);
     auto out = std::unique_ptr<PyDatasetReader>(new PyDatasetReader());
     out->stream_ = std::make_unique<std::istringstream>(
         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<SbeCodec> Create(
-      nb::bytes fds_bytes, std::vector<std::string> file_names) {
+      nb::bytes fds_bytes, std::vector<std::string> 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<SbeCodec>();
     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_<PyDatasetReader>(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_<SbeCodec>(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);