diff --git a/CHANGELOG.md b/CHANGELOG.md index 8021d33..725a003 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -144,6 +144,18 @@ format changes. the corpus's 101-level rows never reached the encoder. A read that fails still exits 2. +- **A document at the depth limit with a well-known-type literal at the + bottom now decodes** (#40), as protowire-go decodes it and as + trendvidia/protowire#537 decided. Before, `pxf.unmarshal` and the other + decoders raised upb's `DecodeError` on it. The native decoder returns PB + bytes, and a Timestamp written as a literal is no PXF level but one PB + level, so 100 PXF levels can be 101 PB levels, one past upb's limit. The + same fix makes a `max_nesting_depth` raised past 100 work. When upb refuses + the bytes for depth alone, they are parsed again with upb's limit lifted + (`SetAllowOversizeProtos`), which is then restored to what it was. The + caller's `max_nesting_depth` is still enforced by the native decoder. + Caveat: upb's switch is process-wide, so another thread's parse during that + retry also gets the lifted limit. - **The macOS wheel's support floor no longer depends on the CI runner's macOS** (#20). protobuf and abseil are now built from source at `MACOSX_DEPLOYMENT_TARGET` and linked statically diff --git a/src/protowire/_upb.py b/src/protowire/_upb.py new file mode 100644 index 0000000..f3cba38 --- /dev/null +++ b/src/protowire/_upb.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 TrendVidia, LLC. +"""Load the PB bytes the native decoders return into a Python message. + +The native decoders (protowire-cpp) hand a decoded message back as PB bytes, +and Python loads them with `MergeFromString`. upb, the protobuf runtime's +parser, refuses PB nested past 100 levels, and that can be deeper than the +PXF the caller's limit allowed: +- A well-known type written as a literal (`t = 2024-01-01T00:00:00Z`) is no + PXF level but one PB level. So a document at the default PXF limit of 100 + can be 101 PB levels. The spec says it decodes (protowire#537), and Go + decodes it. +- A caller may raise `max_nesting_depth` past 100. + +The native decoder has already enforced the caller's limit, so upb's own +limit is a second, stricter check on bytes we produced. When upb refuses the +bytes for depth alone, `load` parses them again with upb's limit lifted +(#40). + +upb has one switch for that limit, `SetAllowOversizeProtos`. It is +process-wide and has no getter. So `load` probes the current setting first, +and turns the switch back off only if it found it off. Another thread +parsing during that window gets the lifted limit too; one calling +`SetAllowOversizeProtos(True)` during it may find it turned back off. +""" + +from __future__ import annotations + +import threading + +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from google.protobuf.internal import api_implementation +from google.protobuf.message import DecodeError, Message + +# upb's message for a parse refused for depth. +_DEPTH_ERROR = "upb_DecodeOptions_MaxDepth" +# upb's depth limit with the switch on (UINT16_MAX, python/message.c). +_LIFTED_LIMIT = 65535 + +if api_implementation.Type() == "upb": + from google._upb import _message as _upb_module + + _set_allow_oversize = _upb_module.SetAllowOversizeProtos +else: # pragma: no cover - the pure-Python runtime has no such switch + _set_allow_oversize = None + +_lock = threading.Lock() +_probe = None # (message class, bytes 101 levels deep), built on first use + + +def _varint(n: int) -> bytes: + out = bytearray() + while n >= 0x80: + out.append((n & 0x7F) | 0x80) + n >>= 7 + out.append(n) + return bytes(out) + + +def _probe_message(): + global _probe + if _probe is None: + fd = descriptor_pb2.FileDescriptorProto( + name="protowire_upb_probe.proto", package="protowire_upb_probe", syntax="proto3" + ) + n = fd.message_type.add(name="N") + n.field.add( + name="c", number=1, type=descriptor_pb2.FieldDescriptorProto.TYPE_MESSAGE, + type_name=".protowire_upb_probe.N", + ) + pool = descriptor_pool.DescriptorPool() + pool.Add(fd) + cls = message_factory.GetMessageClass(pool.FindMessageTypeByName("protowire_upb_probe.N")) + # 101 submessages below the root: one past what upb accepts by + # default (it takes 100). Each level is field 1 (tag 0x0a), a varint + # length, then the level below. + data = b"" + for _ in range(101): + data = b"\x0a" + _varint(len(data)) + data + _probe = (cls, data) + return _probe + + +def _allow_oversize_is_on() -> bool: + cls, data = _probe_message() + try: + cls().MergeFromString(data) + except DecodeError: + return False + return True + + +def load(msg: Message, raw: bytes) -> None: + """Replace `msg`'s contents with the message `raw` encodes.""" + msg.Clear() + try: + msg.MergeFromString(raw) + return + except DecodeError as e: + if _set_allow_oversize is None or _DEPTH_ERROR not in str(e): + raise + msg.Clear() + with _lock: + was_on = _allow_oversize_is_on() + _set_allow_oversize(True) + try: + msg.MergeFromString(raw) + except DecodeError as e: + if _DEPTH_ERROR not in str(e): + raise + raise ValueError( + f"{msg.DESCRIPTOR.full_name} nests past {_LIFTED_LIMIT} levels, " + "the protobuf runtime's deepest" + ) from e + finally: + if not was_on: + _set_allow_oversize(False) diff --git a/src/protowire/pxf.py b/src/protowire/pxf.py index 57dac74..dac2a79 100644 --- a/src/protowire/pxf.py +++ b/src/protowire/pxf.py @@ -13,7 +13,7 @@ from google.protobuf.message import Message -from . import _protowire +from . import _protowire, _upb from ._schema import fds_for_message @@ -322,8 +322,7 @@ def unmarshal( max_bytes_literal_length, max_repeated_count, max_big_float_exponent, ), ) - msg.Clear() - msg.MergeFromString(raw) + _upb.load(msg, raw) def unmarshal_full( @@ -354,8 +353,7 @@ def unmarshal_full( max_bytes_literal_length, max_repeated_count, max_big_float_exponent, ), ) - msg.Clear() - msg.MergeFromString(raw) + _upb.load(msg, raw) return _wrap_result(set_paths, null_paths, dirs, tables, protos) @@ -556,8 +554,7 @@ def scan( ) if raw is None: return False - msg.Clear() - msg.MergeFromString(raw) + _upb.load(msg, raw) return True @@ -615,8 +612,7 @@ def bind_row( max_bytes_literal_length, max_repeated_count, max_big_float_exponent, ), ) - msg.Clear() - msg.MergeFromString(raw) + _upb.load(msg, raw) def _native_cell(cell: Cell | None): diff --git a/src/protowire/sbe.py b/src/protowire/sbe.py index 5af4a98..a267927 100644 --- a/src/protowire/sbe.py +++ b/src/protowire/sbe.py @@ -9,7 +9,7 @@ from google.protobuf.descriptor import FileDescriptor from google.protobuf.message import Message -from . import _protowire +from . import _protowire, _upb from ._schema import fds_for_files @@ -120,8 +120,7 @@ def marshal(self, msg: Message) -> bytes: def unmarshal(self, data: bytes, msg: Message) -> None: raw = self._impl.unmarshal(bytes(data), msg.DESCRIPTOR.full_name) - msg.Clear() - msg.MergeFromString(raw) + _upb.load(msg, raw) def view(self, data: bytes) -> View: return View(self._impl.new_view(bytes(data))) diff --git a/tests/test_upb_depth.py b/tests/test_upb_depth.py new file mode 100644 index 0000000..88d29ce --- /dev/null +++ b/tests/test_upb_depth.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 TrendVidia, LLC. +"""A decode result deeper than upb's 100-level limit still loads (#40). + +The native decoder returns PB bytes, which upb loads. A well-known type +written as a literal is no PXF level but one PB level, so a document at the +PXF limit of 100 with a Timestamp at the bottom is 101 PB levels. The spec +says it decodes (protowire#537), and protowire-go decodes it. A +`max_nesting_depth` raised past 100 needs the same. +""" + +from __future__ import annotations + +import pytest +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory, timestamp_pb2 +from google.protobuf.internal import api_implementation + +from protowire import pxf + +F = descriptor_pb2.FieldDescriptorProto + +pytestmark = pytest.mark.skipif( + api_implementation.Type() != "upb", reason="the depth limit and its switch are upb's" +) + + +def _schema(): + fd = descriptor_pb2.FileDescriptorProto(name="upbdepth.proto", package="ud", syntax="proto3") + fd.dependency.append("google/protobuf/timestamp.proto") + n = fd.message_type.add(name="N") + n.field.add(name="c", number=1, type=F.TYPE_MESSAGE, type_name=".ud.N") + n.field.add(name="v", number=2, type=F.TYPE_INT32) + n.field.add(name="t", number=4, type=F.TYPE_MESSAGE, type_name=".google.protobuf.Timestamp") + pool = descriptor_pool.DescriptorPool() + pool.AddSerializedFile(timestamp_pb2.DESCRIPTOR.serialized_pb) + pool.Add(fd) + return message_factory.GetMessageClass(pool.FindMessageTypeByName("ud.N")) + + +N = _schema() +TS = "t = 2024-01-01T00:00:00Z" + + +def _doc(levels: int, leaf: str) -> str: + return "c { " * levels + leaf + " }" * levels + + +def _bottom(msg): + depth = 0 + while msg.HasField("c"): + msg = msg.c + depth += 1 + return depth, msg + + +def _switch(): + from google._upb import _message + + return _message.SetAllowOversizeProtos + + +def _switch_is_on() -> bool: + from protowire import _upb + + return _upb._allow_oversize_is_on() + + +@pytest.fixture(autouse=True) +def _switch_off_after(): + yield + _switch()(False) + + +def test_the_switch_still_exists(): + # The fix rests on this private upb function. If a protobuf release + # drops it, this fails first. + assert callable(_switch()) + + +def test_upb_alone_refuses_the_document(): + # What #40 measured: without the switch, upb refuses the bytes. If upb + # ever raises its own limit, this test says the fix is no longer needed. + from protowire import _protowire + from protowire._schema import fds_for_message + + out = _protowire.pxf_unmarshal( + _doc(100, TS).encode(), fds_for_message(N()), "ud.N", False, False, (0,) * 6 + ) + with pytest.raises(Exception, match="upb_DecodeOptions_MaxDepth"): + N().MergeFromString(out) + + +@pytest.mark.parametrize("decode", [pxf.unmarshal, pxf.unmarshal_full]) +def test_a_literal_at_the_limit_decodes(decode): + msg = N() + decode(_doc(100, TS), msg) + depth, bottom = _bottom(msg) + assert depth == 100 + assert bottom.t.seconds == 1704067200 + + +def test_the_limit_itself_is_still_enforced(): + with pytest.raises(ValueError, match="MaxNestingDepth=100"): + pxf.unmarshal(_doc(101, "v = 1"), N()) + + +def test_a_raised_limit_decodes_past_100(): + msg = N() + pxf.unmarshal(_doc(300, "v = 1"), msg, max_nesting_depth=300) + assert _bottom(msg)[0] == 300 + with pytest.raises(ValueError, match="MaxNestingDepth=299"): + pxf.unmarshal(_doc(300, "v = 1"), N(), max_nesting_depth=299) + + +def test_the_switch_is_turned_back_off(): + assert not _switch_is_on() + pxf.unmarshal(_doc(100, TS), N()) + assert not _switch_is_on() + + +def test_a_callers_switch_is_left_on(): + _switch()(True) + pxf.unmarshal(_doc(100, TS), N()) + assert _switch_is_on() + + +def test_the_message_is_replaced_not_merged(): + # The first parse fails part-way; what it left must not survive. + msg = N() + msg.v = 7 + pxf.unmarshal(_doc(100, TS), msg) + assert msg.v == 0 + assert _bottom(msg)[0] == 100