From acb55ab0116061ba6652f02ca2611f274ff20cfd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 18 May 2026 18:49:17 +0000 Subject: [PATCH] fix(tests): normalize Google OAuth JWT in VCR matcher to stop cassette overflow Vertex auth POSTs a freshly minted service-account JWT to oauth2.googleapis.com/token on every call. The JWT's iat/exp claims (and resulting RSA signature) change every run, so the safe_body matcher misses, vcrpy records a new episode, and the cassette grows unboundedly until MAX_EPISODES_PER_CASSETTE refuses the save - leaving the two Vertex image-gen tests hitting the live provider on every CI run. Normalize JWT bodies on oauth2.googleapis.com/token (and accounts.google.com/.../token): drop iat/exp/jti/nbf claims and replace the signature with a fixed placeholder. The hook is idempotent, so repeated invocations from vcrpy's matching pipeline remain stable. Different scopes/issuers still produce distinct normalized bodies, so meaningful auth differences aren't collapsed. --- tests/_vcr_conftest_common.py | 124 +++++++++++- tests/llm_translation/test_vcr_filters.py | 231 ++++++++++++++++++++++ 2 files changed, 354 insertions(+), 1 deletion(-) diff --git a/tests/_vcr_conftest_common.py b/tests/_vcr_conftest_common.py index cb43f1abbdd..eea8715007b 100644 --- a/tests/_vcr_conftest_common.py +++ b/tests/_vcr_conftest_common.py @@ -7,14 +7,17 @@ from __future__ import annotations import ast import atexit +import base64 +import binascii import hashlib import json import os import re import socket import sys +import urllib.parse from collections import defaultdict -from typing import Iterable +from typing import Iterable, Optional import pytest @@ -159,6 +162,17 @@ VCR_IMAGE_B64_PLACEHOLDER = "dGVzdA==" # which otherwise turns every multipart cassette into a permanent miss. VCR_FIXED_MULTIPART_BOUNDARY = "vcr-static-boundary" +# Google service-account auth posts a freshly minted JWT assertion to +# ``oauth2.googleapis.com/token`` on every call. The JWT's ``iat``/``exp`` +# claims (and the resulting RSA signature) change every run, so without +# normalization the ``safe_body`` matcher misses and the cassette grows +# unboundedly until ``MAX_EPISODES_PER_CASSETTE`` refuses the save. +GOOGLE_OAUTH_TOKEN_HOSTS = frozenset({"oauth2.googleapis.com", "accounts.google.com"}) +GOOGLE_OAUTH_TOKEN_PATH_SUFFIX = "/token" +GOOGLE_OAUTH_JWT_ASSERTION_FIELD = "assertion" +JWT_VOLATILE_CLAIMS = frozenset({"iat", "exp", "jti", "nbf"}) +JWT_NORMALIZED_SIGNATURE = "vcr-normalized-signature" + def pin_httpx_multipart_boundary(monkeypatch) -> None: try: @@ -542,6 +556,113 @@ def _normalize_multipart_boundary(request) -> None: pass +def _b64url_decode(segment: str) -> bytes: + padding = "=" * (-len(segment) % 4) + return base64.urlsafe_b64decode(segment + padding) + + +def _b64url_encode(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") + + +def _normalize_jwt(token: str) -> Optional[str]: + parts = token.split(".") + if len(parts) != 3: + return None + header_seg, payload_seg, _signature = parts + try: + header = json.loads(_b64url_decode(header_seg)) + payload = json.loads(_b64url_decode(payload_seg)) + except (ValueError, binascii.Error, UnicodeDecodeError): + return None + if not isinstance(header, dict) or not isinstance(payload, dict): + return None + stable_payload = {k: v for k, v in payload.items() if k not in JWT_VOLATILE_CLAIMS} + new_header = _b64url_encode( + json.dumps(header, sort_keys=True, separators=(",", ":")).encode("utf-8") + ) + new_payload = _b64url_encode( + json.dumps(stable_payload, sort_keys=True, separators=(",", ":")).encode( + "utf-8" + ) + ) + return f"{new_header}.{new_payload}.{JWT_NORMALIZED_SIGNATURE}" + + +def _is_google_oauth_token_request(request) -> bool: + method = (getattr(request, "method", "") or "").upper() + if method != "POST": + return False + uri = getattr(request, "uri", None) or getattr(request, "url", None) or "" + parsed = urllib.parse.urlsplit(uri) + if (parsed.hostname or "").lower() not in GOOGLE_OAUTH_TOKEN_HOSTS: + return False + return parsed.path.endswith(GOOGLE_OAUTH_TOKEN_PATH_SUFFIX) + + +def _normalize_google_oauth_jwt(request) -> None: + if not _is_google_oauth_token_request(request): + return + body = getattr(request, "body", None) + if body is None: + return + if isinstance(body, (bytes, bytearray)): + try: + text = bytes(body).decode("utf-8") + except UnicodeDecodeError: + return + was_bytes = True + elif isinstance(body, str): + text = body + was_bytes = False + else: + return + + try: + fields = urllib.parse.parse_qsl(text, keep_blank_values=True) + except ValueError: + return + if not any(name == GOOGLE_OAUTH_JWT_ASSERTION_FIELD for name, _ in fields): + return + + new_fields = [] + changed = False + for name, value in fields: + if name == GOOGLE_OAUTH_JWT_ASSERTION_FIELD: + normalized = _normalize_jwt(value) + if normalized is not None and normalized != value: + new_fields.append((name, normalized)) + changed = True + continue + new_fields.append((name, value)) + + if not changed: + return + + new_text = urllib.parse.urlencode(new_fields) + try: + request.body = new_text.encode("utf-8") if was_bytes else new_text + except (AttributeError, TypeError): + return + + headers = getattr(request, "headers", None) + if headers is not None: + new_len_value = str(len(new_text.encode("utf-8"))) + try: + keys = list(headers.keys()) + except AttributeError: + return + for key in keys: + if str(key).lower() == "content-length": + value = headers[key] + try: + headers[key] = ( + [new_len_value] if isinstance(value, list) else new_len_value + ) + except (TypeError, AttributeError): + return + + def _before_record_request(request): """Fingerprint API keys, scrub them, and normalize multipart boundaries. @@ -572,6 +693,7 @@ def _before_record_request(request): pass _strip_headers(headers, FILTERED_REQUEST_HEADERS) _normalize_multipart_boundary(request) + _normalize_google_oauth_jwt(request) return request diff --git a/tests/llm_translation/test_vcr_filters.py b/tests/llm_translation/test_vcr_filters.py index 03891682781..b46ad34b076 100644 --- a/tests/llm_translation/test_vcr_filters.py +++ b/tests/llm_translation/test_vcr_filters.py @@ -6,21 +6,31 @@ Covers: - ``_normalize_multipart_boundary`` — rewrites random multipart boundaries to a fixed string so audio-transcription request bodies match across record and replay. +- ``_normalize_google_oauth_jwt`` — drops volatile JWT claims and the + signature from Google service-account token requests so the + ``safe_body`` matcher stays stable across runs. """ from __future__ import annotations +import base64 import json import os import sys +import urllib.parse from vcr.request import Request sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) from tests._vcr_conftest_common import ( # noqa: E402 + GOOGLE_OAUTH_JWT_ASSERTION_FIELD, + JWT_NORMALIZED_SIGNATURE, + JWT_VOLATILE_CLAIMS, VCR_FIXED_MULTIPART_BOUNDARY, VCR_IMAGE_B64_PLACEHOLDER, + _normalize_google_oauth_jwt, + _normalize_jwt, _normalize_multipart_boundary, _strip_image_b64_payloads, ) @@ -218,3 +228,224 @@ def test_normalize_multipart_handles_quoted_boundary(): _normalize_multipart_boundary(req) assert b"quoted-boundary" not in req.body assert VCR_FIXED_MULTIPART_BOUNDARY.encode("utf-8") in req.body + + +# --------------------------------------------------------------------------- +# Google OAuth JWT normalizer +# --------------------------------------------------------------------------- + + +def _b64url_encode(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") + + +def _make_jwt( + header: dict, payload: dict, signature: str = "raw-signature-bytes" +) -> str: + header_seg = _b64url_encode(json.dumps(header).encode("utf-8")) + payload_seg = _b64url_encode(json.dumps(payload).encode("utf-8")) + sig_seg = _b64url_encode(signature.encode("utf-8")) + return f"{header_seg}.{payload_seg}.{sig_seg}" + + +def _service_account_jwt(iat: int, exp: int, *, jti: str = "tok-1") -> str: + return _make_jwt( + header={"typ": "JWT", "alg": "RS256", "kid": "abc123"}, + payload={ + "iss": "test-sa@litellm-ci.iam.gserviceaccount.com", + "scope": "https://www.googleapis.com/auth/cloud-platform", + "aud": "https://oauth2.googleapis.com/token", + "iat": iat, + "exp": exp, + "jti": jti, + }, + signature=f"sig-bytes-for-iat-{iat}", + ) + + +def _token_request(jwt_assertion: str) -> Request: + body_text = urllib.parse.urlencode( + [ + (GOOGLE_OAUTH_JWT_ASSERTION_FIELD, jwt_assertion), + ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"), + ] + ) + return Request( + method="POST", + uri="https://oauth2.googleapis.com/token", + body=body_text.encode("utf-8"), + headers={ + "content-type": "application/x-www-form-urlencoded", + "content-length": str(len(body_text.encode("utf-8"))), + }, + ) + + +def test_normalize_jwt_strips_volatile_claims_and_signature(): + token = _service_account_jwt(iat=1779121027, exp=1779124627, jti="abc") + normalized = _normalize_jwt(token) + assert normalized is not None + header_seg, payload_seg, sig_seg = normalized.split(".") + padding = "=" * (-len(payload_seg) % 4) + payload = json.loads(base64.urlsafe_b64decode(payload_seg + padding)) + for claim in JWT_VOLATILE_CLAIMS: + assert claim not in payload + assert payload["iss"] == "test-sa@litellm-ci.iam.gserviceaccount.com" + assert payload["scope"] == "https://www.googleapis.com/auth/cloud-platform" + assert payload["aud"] == "https://oauth2.googleapis.com/token" + assert sig_seg == JWT_NORMALIZED_SIGNATURE + + +def test_normalize_jwt_returns_none_for_malformed_token(): + assert _normalize_jwt("not-a-jwt") is None + assert _normalize_jwt("only.two") is None + assert _normalize_jwt("aaa.bbb.ccc") is None + + +def test_normalize_google_oauth_jwt_makes_two_requests_match(): + """Same logical auth call with different timestamps must produce + byte-identical bodies after normalization.""" + req_record = _token_request(_service_account_jwt(iat=1779055107, exp=1779058707)) + req_replay = _token_request(_service_account_jwt(iat=1779121027, exp=1779124627)) + assert req_record.body != req_replay.body + _normalize_google_oauth_jwt(req_record) + _normalize_google_oauth_jwt(req_replay) + assert req_record.body == req_replay.body + + +def test_normalize_google_oauth_jwt_updates_content_length(): + req = _token_request(_service_account_jwt(iat=1, exp=2)) + _normalize_google_oauth_jwt(req) + expected_len = str(len(req.body)) + assert req.headers["content-length"] == expected_len + + +def test_normalize_google_oauth_jwt_is_idempotent(): + req = _token_request(_service_account_jwt(iat=1, exp=2)) + _normalize_google_oauth_jwt(req) + body_first = req.body + headers_first = dict(req.headers) + _normalize_google_oauth_jwt(req) + assert req.body == body_first + assert dict(req.headers) == headers_first + + +def test_normalize_google_oauth_jwt_preserves_other_form_fields(): + req = _token_request(_service_account_jwt(iat=1, exp=2)) + _normalize_google_oauth_jwt(req) + fields = dict(urllib.parse.parse_qsl(req.body.decode("utf-8"))) + assert fields["grant_type"] == "urn:ietf:params:oauth:grant-type:jwt-bearer" + assert fields[GOOGLE_OAUTH_JWT_ASSERTION_FIELD].endswith( + f".{JWT_NORMALIZED_SIGNATURE}" + ) + + +def test_normalize_google_oauth_jwt_handles_str_body(): + from types import SimpleNamespace + + body_text = urllib.parse.urlencode( + [(GOOGLE_OAUTH_JWT_ASSERTION_FIELD, _service_account_jwt(iat=1, exp=2))] + ) + req = SimpleNamespace( + method="POST", + uri="https://oauth2.googleapis.com/token", + body=body_text, + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + _normalize_google_oauth_jwt(req) + assert isinstance(req.body, str) + assert f".{JWT_NORMALIZED_SIGNATURE}" in req.body + + +def test_normalize_google_oauth_jwt_distinct_scopes_remain_distinct(): + """Different scopes are meaningful auth differences — they must NOT + collapse after normalization.""" + jwt_a = _make_jwt( + header={"typ": "JWT", "alg": "RS256"}, + payload={"iss": "a@x.iam", "scope": "scope-a", "iat": 1, "exp": 2}, + ) + jwt_b = _make_jwt( + header={"typ": "JWT", "alg": "RS256"}, + payload={"iss": "a@x.iam", "scope": "scope-b", "iat": 1, "exp": 2}, + ) + req_a = _token_request(jwt_a) + req_b = _token_request(jwt_b) + _normalize_google_oauth_jwt(req_a) + _normalize_google_oauth_jwt(req_b) + assert req_a.body != req_b.body + + +def test_normalize_google_oauth_jwt_skips_non_oauth_request(): + body = b'{"prompt":"hi"}' + req = Request( + method="POST", + uri="https://api.openai.com/v1/chat/completions", + body=body, + headers={"content-type": "application/json"}, + ) + _normalize_google_oauth_jwt(req) + assert req.body == body + + +def test_normalize_google_oauth_jwt_skips_get_request(): + req = Request( + method="GET", + uri="https://oauth2.googleapis.com/token", + body=None, + headers={}, + ) + _normalize_google_oauth_jwt(req) + assert req.body is None + + +def test_normalize_google_oauth_jwt_skips_when_assertion_missing(): + body_text = urllib.parse.urlencode( + [ + ("grant_type", "refresh_token"), + ("refresh_token", "rt-abc"), + ("client_id", "cid"), + ] + ) + req = Request( + method="POST", + uri="https://oauth2.googleapis.com/token", + body=body_text.encode("utf-8"), + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + original_body = req.body + _normalize_google_oauth_jwt(req) + assert req.body == original_body + + +def test_normalize_google_oauth_jwt_skips_malformed_assertion(): + body_text = urllib.parse.urlencode( + [ + (GOOGLE_OAUTH_JWT_ASSERTION_FIELD, "not-a-jwt"), + ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"), + ] + ) + req = Request( + method="POST", + uri="https://oauth2.googleapis.com/token", + body=body_text.encode("utf-8"), + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + original_body = req.body + _normalize_google_oauth_jwt(req) + assert req.body == original_body + + +def test_normalize_google_oauth_jwt_normalizes_accounts_google_host(): + body_text = urllib.parse.urlencode( + [ + (GOOGLE_OAUTH_JWT_ASSERTION_FIELD, _service_account_jwt(iat=1, exp=2)), + ] + ) + req = Request( + method="POST", + uri="https://accounts.google.com/o/oauth2/token", + body=body_text.encode("utf-8"), + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + _normalize_google_oauth_jwt(req) + assert f".{JWT_NORMALIZED_SIGNATURE}".encode("utf-8") in req.body