From d96dc20e868e5bfb628c139a4d9466e372469b49 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sun, 7 Jun 2026 15:18:49 +0200 Subject: [PATCH] secure endpoint --- routstr/payment/models.py | 53 ++++- .../test_model_test_endpoint_security.py | 224 ++++++++++++++++++ 2 files changed, 270 insertions(+), 7 deletions(-) create mode 100644 tests/integration/test_model_test_endpoint_security.py diff --git a/routstr/payment/models.py b/routstr/payment/models.py index f7fcb3e0..1c1ee52a 100644 --- a/routstr/payment/models.py +++ b/routstr/payment/models.py @@ -3,7 +3,7 @@ import json import random import httpx -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException, Request from pydantic import BaseModel as V2BaseModel from pydantic.v1 import BaseModel from sqlmodel.ext.asyncio.session import AsyncSession @@ -17,6 +17,24 @@ logger = get_logger(__name__) models_router = APIRouter() +_MODEL_TEST_ENDPOINT_PATHS = { + "chat-completions": "chat/completions", + "completions": "completions", + "embeddings": "embeddings", + "responses": "responses", +} + +# Cap the caller-supplied test payload to avoid forwarding oversized bodies +# upstream on the operator's credentials. +_MODEL_TEST_MAX_REQUEST_BYTES = 64 * 1024 + + +async def _require_admin_api(request: Request) -> None: + """Require admin auth without creating an import-time cycle with core.admin.""" + from ..core.admin import require_admin_api + + await require_admin_api(request) + class Architecture(BaseModel): modality: str @@ -418,7 +436,9 @@ class ModelTestRequest(V2BaseModel): request_data: dict -@models_router.post("/api/models/test") +@models_router.post( + "/api/models/test", dependencies=[Depends(_require_admin_api)] +) async def test_model( payload: ModelTestRequest, session: AsyncSession = Depends(get_session), @@ -446,16 +466,35 @@ async def test_model( "status_code": 404, } - base_url = provider.base_url.rstrip("/") - if payload.endpoint_type == "chat-completions": - url = f"{base_url}/chat/completions" - else: - url = f"{base_url}/{payload.endpoint_type}" + endpoint_path = _MODEL_TEST_ENDPOINT_PATHS.get(payload.endpoint_type) + if endpoint_path is None: + raise HTTPException(status_code=400, detail="Unsupported endpoint_type") actual_model_id = model_row.forwarded_model_id or model_row.id request_data = dict(payload.request_data) request_data["model"] = actual_model_id + try: + request_size = len(json.dumps(request_data).encode("utf-8")) + except (TypeError, ValueError): + raise HTTPException(status_code=400, detail="Invalid request_data") + if request_size > _MODEL_TEST_MAX_REQUEST_BYTES: + raise HTTPException(status_code=413, detail="request_data too large") + + base_url = provider.base_url.rstrip("/") + url = f"{base_url}/{endpoint_path}" + + logger.info( + "admin model test", + extra={ + "model_id": payload.model_id, + "forwarded_model_id": actual_model_id, + "endpoint_type": payload.endpoint_type, + "upstream_provider_id": model_row.upstream_provider_id, + "request_bytes": request_size, + }, + ) + headers = { "Content-Type": "application/json", "Authorization": f"Bearer {provider.api_key}", diff --git a/tests/integration/test_model_test_endpoint_security.py b/tests/integration/test_model_test_endpoint_security.py new file mode 100644 index 00000000..dca7e9fa --- /dev/null +++ b/tests/integration/test_model_test_endpoint_security.py @@ -0,0 +1,224 @@ +import time +from types import TracebackType +from typing import Any +from unittest.mock import patch + +import pytest +from httpx import AsyncClient + +from routstr.core.admin import admin_sessions +from routstr.core.db import ModelRow, UpstreamProviderRow + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_model_test_endpoint_requires_admin_auth( + integration_client: AsyncClient, +) -> None: + with patch("httpx.AsyncClient") as mock_async_client: + response = await integration_client.post( + "/api/models/test", + json={ + "model_id": "model-a", + "endpoint_type": "chat-completions", + "request_data": {"messages": []}, + }, + ) + + assert response.status_code == 403 + mock_async_client.assert_not_called() + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_model_test_endpoint_rejects_unsupported_endpoint_type( + integration_client: AsyncClient, + integration_session: Any, +) -> None: + admin_token = "test-admin-token-model-test" + admin_sessions[admin_token] = int(time.time()) + 3600 + integration_client.headers["Authorization"] = f"Bearer {admin_token}" + + provider = UpstreamProviderRow( + provider_type="custom", + base_url="https://api.example.com/v1", + api_key="sk-upstream-test", + enabled=True, + provider_fee=1.01, + ) + integration_session.add(provider) + await integration_session.commit() + await integration_session.refresh(provider) + assert provider.id is not None + + model = ModelRow( + id="model-a", + name="Model A", + created=1, + description="desc", + context_length=100, + architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}', + pricing='{"prompt": 1.0, "completion": 1.0}', + upstream_provider_id=provider.id, + enabled=True, + ) + integration_session.add(model) + await integration_session.commit() + + try: + with patch("httpx.AsyncClient") as mock_async_client: + response = await integration_client.post( + "/api/models/test", + json={ + "model_id": "model-a", + "endpoint_type": "../../abuse", + "request_data": {"messages": []}, + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "Unsupported endpoint_type" + mock_async_client.assert_not_called() + finally: + admin_sessions.pop(admin_token, None) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_model_test_endpoint_rejects_oversized_request_data( + integration_client: AsyncClient, + integration_session: Any, +) -> None: + admin_token = "test-admin-token-model-test-oversized" + admin_sessions[admin_token] = int(time.time()) + 3600 + integration_client.headers["Authorization"] = f"Bearer {admin_token}" + + provider = UpstreamProviderRow( + provider_type="custom", + base_url="https://api.example.com/v1", + api_key="sk-upstream-test", + enabled=True, + provider_fee=1.01, + ) + integration_session.add(provider) + await integration_session.commit() + await integration_session.refresh(provider) + assert provider.id is not None + + model = ModelRow( + id="model-a", + name="Model A", + created=1, + description="desc", + context_length=100, + architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}', + pricing='{"prompt": 1.0, "completion": 1.0}', + upstream_provider_id=provider.id, + enabled=True, + ) + integration_session.add(model) + await integration_session.commit() + + oversized = "x" * (64 * 1024 + 1) + + try: + with patch("httpx.AsyncClient") as mock_async_client: + response = await integration_client.post( + "/api/models/test", + json={ + "model_id": "model-a", + "endpoint_type": "chat-completions", + "request_data": {"blob": oversized}, + }, + ) + + assert response.status_code == 413 + assert response.json()["detail"] == "request_data too large" + mock_async_client.assert_not_called() + finally: + admin_sessions.pop(admin_token, None) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_model_test_endpoint_admin_uses_allowed_upstream_path( + integration_client: AsyncClient, + integration_session: Any, +) -> None: + admin_token = "test-admin-token-model-test-success" + admin_sessions[admin_token] = int(time.time()) + 3600 + integration_client.headers["Authorization"] = f"Bearer {admin_token}" + + provider = UpstreamProviderRow( + provider_type="custom", + base_url="https://api.example.com/v1", + api_key="sk-upstream-test", + enabled=True, + provider_fee=1.01, + ) + integration_session.add(provider) + await integration_session.commit() + await integration_session.refresh(provider) + assert provider.id is not None + + model = ModelRow( + id="model-a", + name="Model A", + created=1, + description="desc", + context_length=100, + architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}', + pricing='{"prompt": 1.0, "completion": 1.0}', + upstream_provider_id=provider.id, + enabled=True, + forwarded_model_id="upstream-model-a", + ) + integration_session.add(model) + await integration_session.commit() + + class MockResponse: + status_code = 200 + text = '{"ok": true}' + + def json(self) -> dict[str, bool]: + return {"ok": True} + + class MockAsyncClient: + async def __aenter__(self) -> "MockAsyncClient": + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + async def post( + self, url: str, json: dict[str, Any], headers: dict[str, str] + ) -> MockResponse: + assert url == "https://api.example.com/v1/chat/completions" + assert json["model"] == "upstream-model-a" + assert headers["Authorization"] == "Bearer sk-upstream-test" + return MockResponse() + + try: + with patch("httpx.AsyncClient", return_value=MockAsyncClient()): + response = await integration_client.post( + "/api/models/test", + json={ + "model_id": "model-a", + "endpoint_type": "chat-completions", + "request_data": {"messages": []}, + }, + ) + + assert response.status_code == 200 + assert response.json() == { + "success": True, + "data": {"ok": True}, + "status_code": 200, + } + finally: + admin_sessions.pop(admin_token, None)