From 8e8af1bf282fca7335d928e883fb7d253ee205cd Mon Sep 17 00:00:00 2001 From: Alexander Osipenko <11722602+subpath@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:00:06 +0200 Subject: [PATCH 1/3] implement model name and ios anndroid tokens sanity checks --- src/mlpa/core/auth/authorize.py | 20 ++++- src/mlpa/core/prometheus_metrics.py | 2 + src/mlpa/core/routers/appattest/middleware.py | 8 +- src/mlpa/core/routers/play/play.py | 8 +- src/mlpa/core/utils.py | 41 +++++++++ src/tests/integration/test_app_attest.py | 8 ++ src/tests/integration/test_play_integrity.py | 17 ++++ src/tests/integration/test_request_country.py | 10 +-- src/tests/unit/test_authorize.py | 29 +++++- .../unit/test_precompletion_availability.py | 2 +- src/tests/unit/test_utils.py | 88 +++++++++++++++++++ 11 files changed, 220 insertions(+), 13 deletions(-) diff --git a/src/mlpa/core/auth/authorize.py b/src/mlpa/core/auth/authorize.py index 6adcf23c..e1f158a9 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,18 @@ async def authorize_chat_request( use_qa_certificates: Annotated[bool | None, Header()] = None, use_play_integrity: Annotated[bool | None, Header()] = None, ) -> AuthorizedChatRequest: + 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..7adf12bd 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)} 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..e54da67e 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,46 @@ def clamp_model(model: str) -> str: return model if model in env.valid_model_labels else "invalid" +_ALLOWED_MODEL_CHARS = frozenset(string.ascii_lowercase + string.digits + ".-") +_ALLOWED_MODEL_EDGE_CHARS = frozenset(string.ascii_lowercase + string.digits) +MAX_MODEL_NAME_LEN = 64 + + +def is_valid_model_name(model: str) -> bool: + if not model or len(model) > MAX_MODEL_NAME_LEN: + return False + if ( + model[0] not in _ALLOWED_MODEL_EDGE_CHARS + or model[-1] not in _ALLOWED_MODEL_EDGE_CHARS + ): + return False + return set(model) <= _ALLOWED_MODEL_CHARS + + +# 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..9e84ae21 100644 --- a/src/tests/integration/test_app_attest.py +++ b/src/tests/integration/test_app_attest.py @@ -42,6 +42,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", 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..ee3d65a5 100644 --- a/src/tests/unit/test_utils.py +++ b/src/tests/unit/test_utils.py @@ -11,9 +11,97 @@ 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, ) +VALID_CLIENT_FACING_MODEL_NAMES = [ + "exa", + "exa-search", + "gemini-3.1-flash-lite", + "gpt-oss-120b", + "qwen3-235b-a22b-instruct-2507-maas", +] + +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_120b", + ], +) +def test_is_valid_model_name_rejects_edge_cases(model): + assert is_valid_model_name(model) is False + + +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 From 5bdce4ca4fa88336574fa2209281798975d33f9f Mon Sep 17 00:00:00 2001 From: Alexander Osipenko <11722602+subpath@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:02:16 +0200 Subject: [PATCH 2/3] fix model charset and key_id gaps --- src/mlpa/core/auth/authorize.py | 2 + src/mlpa/core/routers/appattest/middleware.py | 4 ++ src/mlpa/core/utils.py | 4 +- src/tests/integration/test_app_attest.py | 53 ++++++++++++++++++- src/tests/unit/test_utils.py | 21 +++++--- 5 files changed, 74 insertions(+), 10 deletions(-) diff --git a/src/mlpa/core/auth/authorize.py b/src/mlpa/core/auth/authorize.py index e1f158a9..edb17956 100644 --- a/src/mlpa/core/auth/authorize.py +++ b/src/mlpa/core/auth/authorize.py @@ -124,6 +124,8 @@ 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, diff --git a/src/mlpa/core/routers/appattest/middleware.py b/src/mlpa/core/routers/appattest/middleware.py index 7adf12bd..1814d8f7 100644 --- a/src/mlpa/core/routers/appattest/middleware.py +++ b/src/mlpa/core/routers/appattest/middleware.py @@ -85,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 @@ -114,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/utils.py b/src/mlpa/core/utils.py index e54da67e..c27de4a7 100644 --- a/src/mlpa/core/utils.py +++ b/src/mlpa/core/utils.py @@ -40,7 +40,9 @@ def clamp_model(model: str) -> str: return model if model in env.valid_model_labels else "invalid" -_ALLOWED_MODEL_CHARS = frozenset(string.ascii_lowercase + string.digits + ".-") +# "/" and "_" are required for real, configured model names like +# "openai/gpt-4o" and "vertex_ai/mistral-small-2503" (see litellm_config.yaml). +_ALLOWED_MODEL_CHARS = frozenset(string.ascii_lowercase + string.digits + "./_-") _ALLOWED_MODEL_EDGE_CHARS = frozenset(string.ascii_lowercase + string.digits) MAX_MODEL_NAME_LEN = 64 diff --git a/src/tests/integration/test_app_attest.py b/src/tests/integration/test_app_attest.py index 9e84ae21..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, @@ -86,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() @@ -334,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/unit/test_utils.py b/src/tests/unit/test_utils.py index ee3d65a5..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, @@ -17,13 +18,9 @@ is_valid_model_name, ) -VALID_CLIENT_FACING_MODEL_NAMES = [ - "exa", - "exa-search", - "gemini-3.1-flash-lite", - "gpt-oss-120b", - "qwen3-235b-a22b-instruct-2507-maas", -] +# 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 --", @@ -56,13 +53,21 @@ def test_is_valid_model_name_rejects_sampled_attack_payloads(model): "gpt-oss.", "a" * 65, "GPT-OSS-120B", - "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 From 9fb450db554c8a73bd70d699a2e61e20e51310f8 Mon Sep 17 00:00:00 2001 From: Alexander Osipenko <11722602+subpath@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:20:24 +0200 Subject: [PATCH 3/3] switch model name check to regex --- src/mlpa/core/utils.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/src/mlpa/core/utils.py b/src/mlpa/core/utils.py index c27de4a7..8eb32502 100644 --- a/src/mlpa/core/utils.py +++ b/src/mlpa/core/utils.py @@ -42,20 +42,12 @@ def clamp_model(model: str) -> str: # "/" and "_" are required for real, configured model names like # "openai/gpt-4o" and "vertex_ai/mistral-small-2503" (see litellm_config.yaml). -_ALLOWED_MODEL_CHARS = frozenset(string.ascii_lowercase + string.digits + "./_-") -_ALLOWED_MODEL_EDGE_CHARS = frozenset(string.ascii_lowercase + string.digits) -MAX_MODEL_NAME_LEN = 64 +# 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: - if not model or len(model) > MAX_MODEL_NAME_LEN: - return False - if ( - model[0] not in _ALLOWED_MODEL_EDGE_CHARS - or model[-1] not in _ALLOWED_MODEL_EDGE_CHARS - ): - return False - return set(model) <= _ALLOWED_MODEL_CHARS + return bool(_VALID_MODEL_NAME_PATTERN.match(model)) # Real App Attest key IDs are base64(SHA-256 digest) = 44 chars; this bound is