mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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.
This commit is contained in:
mateo-berri
parent
bb448b0031
commit
acb55ab011
2 changed files with 354 additions and 1 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue