pxf, sbe: load a decode result past upb's 100-level limit

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

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)

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):

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)))

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