feat: basic fuzzing protection for model name and ios, android tokens

src/mlpa/core/auth/authorize.py

@@ -12,11 +12,15 @@
     SearchRequest,
     ServiceType,
 )
-from mlpa.core.config import env
+from mlpa.core.config import ERROR_CODE_INVALID_MODEL_NAME, env
 from mlpa.core.metrics import record_chat_availability_for
 from mlpa.core.prometheus_metrics import AvailabilityReason
 from mlpa.core.routers.appattest import app_attest_auth
-from mlpa.core.utils import extract_user_from_play_integrity_jwt, parse_app_attest_jwt
+from mlpa.core.utils import (
+    extract_user_from_play_integrity_jwt,
+    is_valid_model_name,
+    parse_app_attest_jwt,
+)
 
 TAuthorizedRequest = TypeVar(
     "TAuthorizedRequest", AuthorizedChatRequest, AuthorizedSearchRequest
@@ -120,6 +124,20 @@ async def authorize_chat_request(
     use_qa_certificates: Annotated[bool | None, Header()] = None,
     use_play_integrity: Annotated[bool | None, Header()] = None,
 ) -> AuthorizedChatRequest:
+    # Charset/length check runs before any auth, DB, or LiteLLM work below, so
+    # fuzzed/scanner model values are rejected without that cost.
+    if not is_valid_model_name(chat_request.model):
+        record_chat_availability_for(
+            AvailabilityReason.INVALID_MODEL_NAME,
+            model=chat_request.model,
+            service_type=service_type.value,
+            purpose=purpose or "",
+        )
+        raise HTTPException(
+            status_code=400,
+            detail={"error": ERROR_CODE_INVALID_MODEL_NAME},
+        )
+
     is_service_type_valid = env.valid_service_type_for_model(
         service_type.value, chat_request.model
     )

src/mlpa/core/prometheus_metrics.py

@@ -44,6 +44,8 @@ class AvailabilityReason(StrEnum):
     AUTH_SYSTEM_FAILURE = "auth_system_failure"  # failure
 
     # --- completion-stage reasons (recorded inside stream_completion / get_completion) ---
+    # INVALID_MODEL_NAME is the exception: also recorded pre-completion, in
+    # authorize_chat_request, when the model fails the charset/length check.
     VALID_RESPONSE = "valid_response"  # success
     UPSTREAM_ERROR = "upstream_error"  # failure
     EMPTY_RESPONSE = "empty_response"  # failure

src/mlpa/core/routers/appattest/middleware.py

@@ -14,7 +14,11 @@
     verify_assert,
     verify_attest,
 )
-from mlpa.core.utils import b64decode_safe, parse_app_attest_jwt
+from mlpa.core.utils import (
+    b64decode_safe,
+    is_plausible_base64_key_id,
+    parse_app_attest_jwt,
+)
 
 router = APIRouter()
 
@@ -41,6 +45,8 @@ async def get_challenge(
         raise HTTPException(status_code=400, detail="Bad Request: missing key_id_b64")
     # iOS key id generation not urlsafe workaround
     key_id_b64 = key_id_b64.replace(" ", "+")
+    if not is_plausible_base64_key_id(key_id_b64):
+        raise HTTPException(status_code=400, detail="Bad Request: invalid key_id_b64")
     return {"challenge": await generate_client_challenge(key_id_b64)}
 
 
@@ -79,6 +85,8 @@ async def attest(
     ] = None,
 ):
     attestationAuth = parse_app_attest_jwt(authorization, "attest")
+    if not is_plausible_base64_key_id(attestationAuth.key_id_b64):
+        raise HTTPException(status_code=400, detail="Bad Request: invalid key_id_b64")
     challenge_bytes = b64decode_safe(attestationAuth.challenge_b64, "challenge_b64")
     if not await validate_challenge(
         challenge_bytes.decode(), attestationAuth.key_id_b64
@@ -108,6 +116,8 @@ async def app_attest_auth(
     expected_hash: bytes,
     use_qa_certificates: bool,
 ):
+    if not is_plausible_base64_key_id(assertionAuth.key_id_b64):
+        raise HTTPException(status_code=401, detail="Invalid App Attest assertion")
     challenge_bytes = b64decode_safe(assertionAuth.challenge_b64, "challenge_b64")
     if not await validate_challenge(challenge_bytes.decode(), assertionAuth.key_id_b64):
         raise HTTPException(status_code=401, detail="Invalid or expired challenge")

src/mlpa/core/routers/play/play.py

@@ -13,7 +13,11 @@
 from mlpa.core.config import PLAY_VERIFY_RESPONSES, env
 from mlpa.core.http_client import get_http_client
 from mlpa.core.prometheus_metrics import PrometheusResult, metrics
-from mlpa.core.utils import issue_mlpa_access_token, raise_and_log
+from mlpa.core.utils import (
+    is_plausible_integrity_token,
+    issue_mlpa_access_token,
+    raise_and_log,
+)
 
 router = APIRouter()
 
@@ -113,6 +117,8 @@ async def verify_play_integrity(payload: PlayIntegrityRequest):
     start_time = time.perf_counter()
     try:
         result = PrometheusResult.ERROR
+        if not is_plausible_integrity_token(payload.integrity_token):
+            raise HTTPException(status_code=400, detail="Invalid integrity_token")
         decoded = await _decode_integrity_token(
             payload.integrity_token, payload.package_name
         )

src/mlpa/core/utils.py

@@ -2,6 +2,7 @@
 import base64
 import json
 import re
+import string
 import time
 from functools import lru_cache
 from typing import Any, Literal, NoReturn, cast, overload
@@ -39,6 +40,40 @@ def clamp_model(model: str) -> str:
     return model if model in env.valid_model_labels else "invalid"
 
 
+# "/" and "_" are required for real, configured model names like
+# "openai/gpt-4o" and "vertex_ai/mistral-small-2503" (see litellm_config.yaml).
+# Max length 64: edge chars (2) + up to 62 interior chars.
+_VALID_MODEL_NAME_PATTERN = re.compile(r"^[a-z0-9](?:[a-z0-9./_-]{0,62}[a-z0-9])?$")
+
+
+def is_valid_model_name(model: str) -> bool:
+    return bool(_VALID_MODEL_NAME_PATTERN.match(model))
+
+
+# Real App Attest key IDs are base64(SHA-256 digest) = 44 chars; this bound is
+# generous slack, not a hand-tuned exact length.
+_BASE64_STANDARD_CHARS = frozenset(string.ascii_letters + string.digits + "+/=")
+MAX_KEY_ID_B64_LEN = 128
+
+
+def is_plausible_base64_key_id(key_id_b64: str) -> bool:
+    if not key_id_b64 or len(key_id_b64) > MAX_KEY_ID_B64_LEN:
+        return False
+    return set(key_id_b64) <= _BASE64_STANDARD_CHARS
+
+
+# Play Integrity tokens are compact JWT/JWE strings (base64url segments joined
+# by "."); this bound is generous slack for the real token size, not exact.
+_JWT_LIKE_CHARS = frozenset(string.ascii_letters + string.digits + "-_.")
+MAX_INTEGRITY_TOKEN_LEN = 10000
+
+
+def is_plausible_integrity_token(integrity_token: str) -> bool:
+    if not integrity_token or len(integrity_token) > MAX_INTEGRITY_TOKEN_LEN:
+        return False
+    return set(integrity_token) <= _JWT_LIKE_CHARS
+
+
 def clamp_service_type(raw: str) -> str:
     return raw if raw == "" or raw in env.valid_service_types_set else "other"
 

src/tests/integration/test_app_attest.py

@@ -10,8 +10,9 @@
 from fastapi import HTTPException
 from pyattest.testutils.factories.attestation import apple as apple_factory
 
+from mlpa.core.classes import AssertionAuth
 from mlpa.core.config import env
-from mlpa.core.routers.appattest import appattest
+from mlpa.core.routers.appattest import appattest, middleware
 from tests.consts import (
     SAMPLE_CHAT_REQUEST,
     SUCCESSFUL_CHAT_RESPONSE,
@@ -42,6 +43,14 @@ def test_get_challenge(mocked_client_integration):
     assert len(response.json().get("challenge")) > 0
 
 
+def test_get_challenge_rejects_garbage_key_id(mocked_client_integration):
+    response = mocked_client_integration.get(
+        "/verify/challenge",
+        params={"key_id_b64": "' OR (SELECT pg_sleep(6)) IS NULL --"},
+    )
+    assert response.status_code == 400
+
+
 def test_invalid_methods(mocked_client_integration):
     response = mocked_client_integration.post(
         "/verify/challenge",
@@ -78,6 +87,35 @@ def test_invalid_methods(mocked_client_integration):
     assert response.status_code == 405
 
 
+def test_attest_rejects_garbage_key_id(mocker, mocked_client_integration):
+    get_challenge_mock = mocker.patch(
+        "mlpa.core.routers.appattest.appattest.app_attest_pg.get_challenge",
+    )
+    app_attest_jwt = jwt.encode(
+        {
+            "challenge_b64": base64.b64encode(b"whatever").decode(),
+            "key_id_b64": "' OR (SELECT pg_sleep(6)) IS NULL --",
+            "attestation_obj_b64": "VEVTVF9BVFRFU1RBVElPTl9CQVNFNjRVUkw=",
+            "bundle_id": TEST_BUNDLE_ID,
+        },
+        key=jwt_secret,
+        algorithm="HS256",
+    )
+    headers = {
+        "Authorization": f"Bearer {app_attest_jwt}",
+        "use-app-attest": "true",
+        "service-type": "ai",
+        "purpose": "chat",
+    }
+    response = mocked_client_integration.post(
+        "/verify/attest",
+        headers=headers,
+        json=sample_chat_request,
+    )
+    assert response.status_code == 400
+    get_challenge_mock.assert_not_called()
+
+
 def test_invalid_challenge(mocked_client_integration):
     challenge = "bad_challenge"
     challenge_b64 = base64.b64encode(challenge.encode()).decode()
@@ -326,6 +364,27 @@ def _build_fake_assertion(counter: int) -> bytes:
     return cbor2.dumps({"authenticatorData": bytes(auth_data)})
 
 
+async def test_app_attest_auth_rejects_garbage_key_id(mocker):
+    # app_attest_auth is mocked out entirely by mocked_client_integration, so
+    # this exercises the real function directly to prove the DB lookup is
+    # never reached for a garbage key_id_b64.
+    get_challenge_mock = mocker.patch(
+        "mlpa.core.routers.appattest.appattest.app_attest_pg.get_challenge",
+    )
+    assertion_auth = AssertionAuth(
+        key_id_b64="' OR (SELECT pg_sleep(6)) IS NULL --",
+        challenge_b64=base64.b64encode(b"whatever").decode(),
+        assertion_obj_b64="VEVTVF9BU1NFUlRJT05fQkFTRTY0VVJM",
+        bundle_id=TEST_BUNDLE_ID,
+    )
+
+    with pytest.raises(HTTPException) as exc:
+        await middleware.app_attest_auth(assertion_auth, b"expected_hash", False)
+
+    assert exc.value.status_code == 401
+    get_challenge_mock.assert_not_called()
+
+
 async def test_verify_assert_rejects_non_monotonic_counter(mocker):
     mock_pg = MockAppAttestPGService(MockLiteLLMPGService())
     mocker.patch("mlpa.core.routers.appattest.appattest.app_attest_pg", mock_pg)

src/tests/integration/test_play_integrity.py

@@ -37,6 +37,23 @@ def test_verify_play_integrity_success(mocked_client_integration, mocker):
     assert data["expires_in"] == env.MLPA_ACCESS_TOKEN_TTL_SECONDS
 
 
+def test_verify_play_integrity_rejects_garbage_token(mocked_client_integration, mocker):
+    decode_mock = mocker.patch(
+        "mlpa.core.routers.play.play._decode_integrity_token",
+    )
+
+    response = mocked_client_integration.post(
+        "/verify/play",
+        json={
+            "integrity_token": "' OR (SELECT pg_sleep(6)) IS NULL --",
+            "user_id": TEST_USER_ID,
+        },
+    )
+
+    assert response.status_code == 400
+    decode_mock.assert_not_called()
+
+
 def test_verify_play_integrity_invalid_hash(mocked_client_integration, mocker):
     mocker.patch(
         "mlpa.core.routers.play.play._decode_integrity_token",

src/tests/integration/test_request_country.py

@@ -9,7 +9,7 @@ def _country_count(metrics_spy, **labels):
     return metrics_spy.value("requests_by_country_total", **labels)
 
 
-def _chat_payload(model: str = "openai/gpt-4o"):
+def _chat_payload(model: str = "gpt-oss-120b"):
     return SAMPLE_REQUEST.model_copy(update={"model": model}).model_dump()
 
 
@@ -29,7 +29,7 @@ def test_chat_records_country(mocked_client_integration, metrics_spy):
         _country_count(
             metrics_spy,
             service_type="ai",
-            model="openai/gpt-4o",
+            model="gpt-oss-120b",
             client_country="DE",
         )
         == 1.0
@@ -51,7 +51,7 @@ def test_chat_missing_geo_header_is_unknown(mocked_client_integration, metrics_s
         _country_count(
             metrics_spy,
             service_type="ai",
-            model="openai/gpt-4o",
+            model="gpt-oss-120b",
             client_country="unknown",
         )
         == 1.0
@@ -73,7 +73,7 @@ def test_chat_spoofed_geo_header_is_clamped(mocked_client_integration, metrics_s
         _country_count(
             metrics_spy,
             service_type="ai",
-            model="openai/gpt-4o",
+            model="gpt-oss-120b",
             client_country="unknown",
         )
         == 1.0
@@ -82,7 +82,7 @@ def test_chat_spoofed_geo_header_is_clamped(mocked_client_integration, metrics_s
         _country_count(
             metrics_spy,
             service_type="ai",
-            model="openai/gpt-4o",
+            model="gpt-oss-120b",
             client_country="ZZZ",
         )
         == 0.0

src/tests/unit/test_authorize.py

@@ -30,7 +30,7 @@ async def test_authorize_chat_request_returns_authorized_chat_request(mocker):
     result = await authorize_module.authorize_chat_request(
         request=_make_request("/v1/chat/completions"),
         chat_request=ChatRequest(
-            model="openai/gpt-4o", messages=[{"role": "user", "content": "hello"}]
+            model="gpt-oss-120b", messages=[{"role": "user", "content": "hello"}]
         ),
         authorization="Bearer token",
         service_type=authorize_module.ServiceType.ai,
@@ -94,7 +94,7 @@ async def test_authorize_chat_request_rejects_answer_for_non_exa_model():
     with pytest.raises(HTTPException) as exc_info:
         await authorize_module.authorize_chat_request(
             request=_make_request("/v1/chat/completions"),
-            chat_request=ChatRequest(model="openai/gpt-4o", messages=[]),
+            chat_request=ChatRequest(model="gpt-oss-120b", messages=[]),
             authorization="Bearer token",
             service_type=authorize_module.ServiceType.answer,
             purpose=None,
@@ -103,6 +103,29 @@ async def test_authorize_chat_request_rejects_answer_for_non_exa_model():
     assert exc_info.value.status_code == 400
     assert (
         exc_info.value.detail
-        == "Invalid service-type value answer for model openai/gpt-4o. "
+        == "Invalid service-type value answer for model gpt-oss-120b. "
         "Service type answer is only valid for models ['exa']"
     )
+
+
+async def test_authorize_chat_request_rejects_invalid_model_name(mocker):
+    valid_service_type_mock = mocker.patch.object(
+        type(authorize_module.env), "valid_service_type_for_model"
+    )
+    fxa_auth_mock = mocker.patch.object(authorize_module, "fxa_auth")
+
+    with pytest.raises(HTTPException) as exc_info:
+        await authorize_module.authorize_chat_request(
+            request=_make_request("/v1/chat/completions"),
+            chat_request=ChatRequest(
+                model="' OR (SELECT pg_sleep(6)) IS NULL --", messages=[]
+            ),
+            authorization="Bearer token",
+            service_type=authorize_module.ServiceType.ai,
+            purpose="chat",
+        )
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.detail == {"error": 8}
+    valid_service_type_mock.assert_not_called()
+    fxa_auth_mock.assert_not_called()

src/tests/unit/test_precompletion_availability.py

@@ -28,7 +28,7 @@
 
 # A model/service-type pair that is valid together, so the wrapper passes its own
 # check and reaches the shared auth call.
-_VALID_MODEL = "openai/gpt-4o"
+_VALID_MODEL = "gpt-oss-120b"
 _AI = authorize_module.ServiceType.ai