From a0c4f21114fed42a9dcac48f0becfb6465018288 Mon Sep 17 00:00:00 2001 From: chottokun <29515187+chottokun@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:32:21 +0000 Subject: [PATCH 1/2] =?UTF-8?q?=F0=9F=94=92=20Fix=20SSRF=20vulnerability?= =?UTF-8?q?=20in=20TEI=20proxy=20path=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validate requested paths in _proxy_to_tei against an explicit allowlist (ALLOWED_TEI_PATHS) to prevent SSRF and unauthorized endpoint invocation. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- src/app/dependencies/services.py | 14 ++++++++++++-- src/tests/test_tei_proxy.py | 31 +++++++++++++++++++++++++------ 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/src/app/dependencies/services.py b/src/app/dependencies/services.py index a8f7e2b..89e8be5 100644 --- a/src/app/dependencies/services.py +++ b/src/app/dependencies/services.py @@ -2,6 +2,8 @@ import httpx from fastapi import HTTPException +ALLOWED_TEI_PATHS = {"/v1/embeddings", "/rerank", "/predict", "/embed"} + from app.config import EMBEDDING_MODELS, RERANK_MODELS from app.models import get_model from app.services import ( @@ -15,6 +17,14 @@ 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 13a7ed2..5eb65d6 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 From da5430fea7a5f337d68531c7caea353ecc482556 Mon Sep 17 00:00:00 2001 From: chottokun Date: Mon, 12 Oct 2026 11:32:41 +0900 Subject: [PATCH 2/2] fix(lint): move ALLOWED_TEI_PATHS below imports to fix E402 --- src/app/dependencies/services.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/app/dependencies/services.py b/src/app/dependencies/services.py index 89e8be5..e41e1dd 100644 --- a/src/app/dependencies/services.py +++ b/src/app/dependencies/services.py @@ -2,8 +2,6 @@ import httpx from fastapi import HTTPException -ALLOWED_TEI_PATHS = {"/v1/embeddings", "/rerank", "/predict", "/embed"} - from app.config import EMBEDDING_MODELS, RERANK_MODELS from app.models import get_model from app.services import ( @@ -14,6 +12,8 @@ ) 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."""