diff --git a/src/app/dependencies/services.py b/src/app/dependencies/services.py index a8f7e2b..e41e1dd 100644 --- a/src/app/dependencies/services.py +++ b/src/app/dependencies/services.py @@ -12,9 +12,19 @@ ) from app.services.base import get_validated_model +ALLOWED_TEI_PATHS = {"/v1/embeddings", "/rerank", "/predict", "/embed"} + async def _proxy_to_tei(tei_url: str, path: str, json_data: dict) -> Any: """Helper to send an async POST request to TEI and return the JSON response.""" + if path not in ALLOWED_TEI_PATHS: + raise HTTPException( + status_code=400, + detail=f"Invalid proxy path target: '{path}' is not an allowed endpoint.", + ) + + target_url = f"{tei_url.rstrip('/')}{path}" + try: # Avoid importing app globally to prevent circular imports from app.main import app @@ -24,10 +34,10 @@ async def _proxy_to_tei(tei_url: str, path: str, json_data: dict) -> Any: shared_client is not None and getattr(shared_client, "is_closed", False) is not True ): - response = await shared_client.post(f"{tei_url}{path}", json=json_data) + response = await shared_client.post(target_url, json=json_data) else: async with httpx.AsyncClient(timeout=30.0) as client: - response = await client.post(f"{tei_url}{path}", json=json_data) + response = await client.post(target_url, json=json_data) if response.status_code != 200: error_msg = response.text diff --git a/src/tests/test_tei_proxy.py b/src/tests/test_tei_proxy.py index 8d4abd9..538a989 100644 --- a/src/tests/test_tei_proxy.py +++ b/src/tests/test_tei_proxy.py @@ -148,7 +148,7 @@ async def post(self, url, json): with patch("httpx.AsyncClient", MockAsyncClient): with pytest.raises(HTTPException) as exc_info: - await _proxy_to_tei("http://tei-url", "/path", {"data": "test"}) + await _proxy_to_tei("http://tei-url", "/v1/embeddings", {"data": "test"}) assert exc_info.value.status_code == 500 expected_truncated_text = "A" * 200 + "..." @@ -171,7 +171,7 @@ async def post(self, url, json): with patch("httpx.AsyncClient", MockAsyncClientShort): with pytest.raises(HTTPException) as exc_info: - await _proxy_to_tei("http://tei-url", "/path", {"data": "test"}) + await _proxy_to_tei("http://tei-url", "/v1/embeddings", {"data": "test"}) assert exc_info.value.status_code == 500 assert "Short error" in exc_info.value.detail @@ -195,10 +195,10 @@ async def test_tei_proxy_uses_pooled_client(): mock_client_instance = mock_client_class.return_value.__aenter__.return_value mock_client_instance.post = AsyncMock(return_value=mock_response) - res = await _proxy_to_tei("http://tei-url", "/path", {"test": "data"}) + res = await _proxy_to_tei("http://tei-url", "/v1/embeddings", {"test": "data"}) assert res == {"success": True} mock_client_instance.post.assert_called_once_with( - "http://tei-url/path", json={"test": "data"} + "http://tei-url/v1/embeddings", json={"test": "data"} ) # 2. Test when app.state.tei_client IS initialized @@ -207,13 +207,32 @@ async def test_tei_proxy_uses_pooled_client(): app.state.tei_client = mock_pooled_client with patch("httpx.AsyncClient") as mock_client_class: - res = await _proxy_to_tei("http://tei-url", "/path", {"test": "data"}) + res = await _proxy_to_tei("http://tei-url", "/v1/embeddings", {"test": "data"}) assert res == {"success": True} mock_client_class.assert_not_called() mock_pooled_client.post.assert_called_once_with( - "http://tei-url/path", json={"test": "data"} + "http://tei-url/v1/embeddings", json={"test": "data"} ) # Clean up if hasattr(app.state, "tei_client"): delattr(app.state, "tei_client") + + +@pytest.mark.anyio +async def test_tei_proxy_disallows_unauthorized_paths(): + """Verify that _proxy_to_tei rejects paths not in ALLOWED_TEI_PATHS with HTTP 400.""" + unauthorized_paths = [ + "/admin", + "/../secret", + "http://evil.com", + "/v1/embeddings/../../etc/passwd", + "/health", + "/metrics", + ] + + for invalid_path in unauthorized_paths: + with pytest.raises(HTTPException) as exc_info: + await _proxy_to_tei("http://tei-url", invalid_path, {"test": "data"}) + assert exc_info.value.status_code == 400 + assert "Invalid proxy path target" in exc_info.value.detail