diff --git a/src/mlpa/core/auth/authorize.py b/src/mlpa/core/auth/authorize.py index 6adcf23c..edb17956 100644 --- a/src/mlpa/core/auth/authorize.py +++ b/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 ) diff --git a/src/mlpa/core/prometheus_metrics.py b/src/mlpa/core/prometheus_metrics.py index 12c19667..b70964ff 100644 --- a/src/mlpa/core/prometheus_metrics.py +++ b/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 diff --git a/src/mlpa/core/routers/appattest/middleware.py b/src/mlpa/core/routers/appattest/middleware.py index fa264b9a..1814d8f7 100644 --- a/src/mlpa/core/routers/appattest/middleware.py +++ b/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") diff --git a/src/mlpa/core/routers/play/play.py b/src/mlpa/core/routers/play/play.py index d030a01e..d13aa653 100644 --- a/src/mlpa/core/routers/play/play.py +++ b/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 ) diff --git a/src/mlpa/core/utils.py b/src/mlpa/core/utils.py index 810cc63e..8eb32502 100644 --- a/src/mlpa/core/utils.py +++ b/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" diff --git a/src/tests/integration/test_app_attest.py b/src/tests/integration/test_app_attest.py index acbd73ec..c3f33ba3 100644 --- a/src/tests/integration/test_app_attest.py +++ b/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) diff --git a/src/tests/integration/test_play_integrity.py b/src/tests/integration/test_play_integrity.py index e7223a2a..0932005d 100644 --- a/src/tests/integration/test_play_integrity.py +++ b/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", diff --git a/src/tests/integration/test_request_country.py b/src/tests/integration/test_request_country.py index 1fec3ba7..55278c44 100644 --- a/src/tests/integration/test_request_country.py +++ b/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 diff --git a/src/tests/unit/test_authorize.py b/src/tests/unit/test_authorize.py index a6da04b8..212670d5 100644 --- a/src/tests/unit/test_authorize.py +++ b/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() diff --git a/src/tests/unit/test_precompletion_availability.py b/src/tests/unit/test_precompletion_availability.py index d25d7cfb..0d772e4a 100644 --- a/src/tests/unit/test_precompletion_availability.py +++ b/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 diff --git a/src/tests/unit/test_utils.py b/src/tests/unit/test_utils.py index 89160079..e997b46c 100644 --- a/src/tests/unit/test_utils.py +++ b/src/tests/unit/test_utils.py @@ -3,6 +3,7 @@ import pytest from fastapi import HTTPException +from mlpa.core.config import env from mlpa.core.utils import ( b64decode_safe, clamp_country, @@ -11,9 +12,101 @@ is_context_window_error, is_invalid_model_name_error, is_invalid_request_error, + is_plausible_base64_key_id, + is_plausible_integrity_token, is_rate_limit_error, + is_valid_model_name, ) +# Sourced from env.valid_model_labels (not hand-copied) so a future model name +# with an unexpected character is caught here instead of silently 400ing. +VALID_CLIENT_FACING_MODEL_NAMES = sorted(env.valid_model_labels) + +SAMPLED_ATTACK_PAYLOADS = [ + "' OR (SELECT pg_sleep(6)) IS NULL --", + "1'||sleep(27*1000)*ugxuhb||'", + "'\"()&%", + "str(__import__('time').sleep(9))+__import__('socket').gethostbyname(...)", + "-1 OR 5*5=25 --", + "1'||DBMS_PIPE.RECEIVE_MESSAGE(CHR(98)...", + "|echo culqcs$()\\ wbjwmr\nz^xyu||a #", +] + + +@pytest.mark.parametrize("model", VALID_CLIENT_FACING_MODEL_NAMES) +def test_is_valid_model_name_accepts_real_model_names(model): + assert is_valid_model_name(model) is True + + +@pytest.mark.parametrize("model", SAMPLED_ATTACK_PAYLOADS) +def test_is_valid_model_name_rejects_sampled_attack_payloads(model): + assert is_valid_model_name(model) is False + + +@pytest.mark.parametrize( + "model", + [ + "", + "-gpt-oss", + "gpt-oss-", + ".gpt-oss", + "gpt-oss.", + "a" * 65, + "GPT-OSS-120B", + "_gpt-oss", + "gpt-oss_", + "/gpt-oss", + "gpt-oss/", + ], +) +def test_is_valid_model_name_rejects_edge_cases(model): + assert is_valid_model_name(model) is False + + +@pytest.mark.parametrize("model", ["openai/gpt-4o", "vertex_ai/mistral-small-2503"]) +def test_is_valid_model_name_accepts_slash_and_underscore_namespaced_models(model): + assert is_valid_model_name(model) is True + + +def test_is_valid_model_name_accepts_max_length(): + assert is_valid_model_name("a" * 64) is True + + +@pytest.mark.parametrize("key_id_b64", ["dGVzdC1rZXktaWQ=", "A" * 44, "a+b/c==", "x"]) +def test_is_plausible_base64_key_id_accepts_valid(key_id_b64): + assert is_plausible_base64_key_id(key_id_b64) is True + + +@pytest.mark.parametrize( + "key_id_b64", + [ + "", + "a" * 129, + *SAMPLED_ATTACK_PAYLOADS, + ], +) +def test_is_plausible_base64_key_id_rejects_invalid(key_id_b64): + assert is_plausible_base64_key_id(key_id_b64) is False + + +@pytest.mark.parametrize( + "integrity_token", ["test-token", "a" * 10000, "header.payload.sig.iv.tag"] +) +def test_is_plausible_integrity_token_accepts_valid(integrity_token): + assert is_plausible_integrity_token(integrity_token) is True + + +@pytest.mark.parametrize( + "integrity_token", + [ + "", + "a" * 10001, + *SAMPLED_ATTACK_PAYLOADS, + ], +) +def test_is_plausible_integrity_token_rejects_invalid(integrity_token): + assert is_plausible_integrity_token(integrity_token) is False + def test_b64decode_safe(): # Valid base64 string