From 9176744735eb7458f3ccdb5a302ceb53915bc82e Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 13:56:49 -0700 Subject: [PATCH] perf(mcp): O(1) character precheck before the exact byte size guard in open_envelope --- .../outbound_credentials/envelope.py | 5 +++++ .../outbound_credentials/test_envelope.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index dde69527875..298bc8d98cc 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -257,6 +257,11 @@ def open_envelope( """ if not is_envelope(candidate): return NotAnEnvelope() + # UTF-8 byte length is never below character length, so a character count already over the + # cap rejects an oversize candidate in O(1) without encoding it; the exact byte check then + # runs only on candidates already bounded to <= MAX_ENVELOPE_BYTES characters. + if len(candidate) > MAX_ENVELOPE_BYTES: + return MalformedPayload() if len(candidate.encode("utf-8", "surrogatepass")) > MAX_ENVELOPE_BYTES: return MalformedPayload() claims = _decode_claims(candidate.removeprefix(ENVELOPE_PREFIX), keys.signing_key) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index d479adc3f37..71de206aa1e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -387,6 +387,24 @@ def test_open_size_guard_measures_bytes_not_characters(): decode.assert_not_called() +def test_open_size_guard_rejects_oversize_character_count_before_decode(): + """A candidate whose character count already exceeds the cap is rejected up front, before the + decode path, so an arbitrarily long hostile string is not run through HMAC/decrypt. The cheap + character precheck makes this O(1) since UTF-8 byte length is never below character length.""" + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.outbound_credentials import envelope + + candidate = ENVELOPE_PREFIX + ("a" * (MAX_ENVELOPE_BYTES + 1)) + assert len(candidate) > MAX_ENVELOPE_BYTES + + with patch.object(envelope, "_decode_claims", side_effect=AssertionError("decode reached")) as decode: + result = open_envelope(candidate, _KEYS, _NOW) + + assert isinstance(result, MalformedPayload) + decode.assert_not_called() + + def test_is_envelope_detects_only_prefixed_values(): assert is_envelope(_sealed_token(_full_grant())) raw_jwt = jwt.encode({"sub": "user-123"}, _SIGNING_KEY, algorithm="HS256")