From a0957a24f854597fa4a9777490c625fd410b5373 Mon Sep 17 00:00:00 2001 From: "Benedicto Franco Jr." Date: Sat, 7 Nov 2026 16:54:43 -0800 Subject: [PATCH] pxf, sbe: load a decode result past upb's 100-level limit The native decoders return PB bytes, which Python loads with upb. 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, and upb refused it with DecodeError. protowire-go decodes it, and protowire#537 decided that it decodes. A max_nesting_depth raised past 100 hit the same wall. _upb.load replaces the five Clear + MergeFromString sites. When upb refuses for depth alone, it parses again with SetAllowOversizeProtos on, upb's only depth switch (it sets upb_DecodeOptions_MaxDepth to 65535 and nothing else, python/message.c). The switch is process-wide with no getter, so load first probes it with a 101-level message and turns it back off only if it found it off. The native decoder has already enforced the caller's limit, so this lifts only upb's second, stricter check. Fixes #40 --- CHANGELOG.md | 13 ++++ src/protowire/_upb.py | 117 +++++++++++++++++++++++++++++++++++ src/protowire/pxf.py | 14 ++--- src/protowire/sbe.py | 5 +- tests/test_upb_depth.py | 133 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 270 insertions(+), 12 deletions(-) create mode 100644 src/protowire/_upb.py create mode 100644 tests/test_upb_depth.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f73c12..ef0c21c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -144,6 +144,19 @@ 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. + ## [1.1.0] — 2026-10-31 ### Added 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