From 734cfa958e41a70f0ac8395863d0cdcd704dc129 Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Thu, 17 Sep 2026 12:55:04 +0200 Subject: [PATCH] Authenticate project server communication with the shared secret - Require a valid Bearer token on heartbeat and project-query requests. - Send the configured auth_secret from clients and the query-project tool. - Test accepted and rejected credentials, including real HTTP client requests. --- src/serena/project_server.py | 33 ++++++++++++---- src/serena/repl/external_project.py | 2 +- src/serena/tools/query_project_tools.py | 2 +- test/serena/test_external_projects.py | 5 ++- test/serena/test_project_server.py | 52 ++++++++++++++++++++++++- 5 files changed, 82 insertions(+), 12 deletions(-) diff --git a/src/serena/project_server.py b/src/serena/project_server.py index 4e8a69c8..e00100d3 100644 --- a/src/serena/project_server.py +++ b/src/serena/project_server.py @@ -3,11 +3,12 @@ import json import logging import pickle +import secrets import threading from typing import TYPE_CHECKING, Any import requests as requests_lib -from flask import Flask, Response, request +from flask import Flask, Response, abort, request from pydantic import BaseModel from sensai.util.logging import LogTime @@ -90,7 +91,22 @@ class ProjectServer: self._setup_routes() + def get_serena_config(self) -> SerenaConfig: + return self._agent.serena_config + + def get_auth_secret(self) -> str: + """Returns the authentication secret used by the server.""" + return self._agent.serena_config.auth_secret + def _setup_routes(self) -> None: + @self._app.before_request + def authenticate() -> None: + # authenticate every request before parsing input or accessing projects + secret = self.get_auth_secret() + provided = request.headers.get("Authorization", "") + if not secret or not secrets.compare_digest(provided.encode("utf-8"), f"Bearer {secret}".encode()): + abort(401) + @self._app.route("/heartbeat", methods=["GET"]) def heartbeat() -> dict[str, str]: return {"status": "alive"} @@ -195,20 +211,23 @@ class ProjectServerClient: :class:`ConnectionError` is raised. """ - def __init__(self, host: str = "127.0.0.1", port: int | None = None, timeout: int = 300) -> None: + def __init__(self, serena_config: SerenaConfig, host: str = "127.0.0.1", port: int | None = None) -> None: """ :param host: the host address of the project server. - :param port: the port of the project server; None for the default port. + :param port: the port of the project server; if None, use default. + :param auth_secret: the shared authentication secret; defaults to the secret in Serena's configuration. :raises ConnectionError: if the project server is not reachable. """ if port is None: port = ProjectServer.PORT self._base_url = f"http://{host}:{port}" - self._timeout = timeout + self._timeout = serena_config.tool_timeout - 1 + auth_secret = serena_config.auth_secret + self._headers = {"Authorization": f"Bearer {auth_secret}"} # verify that the server is running try: - response = requests_lib.get(f"{self._base_url}/heartbeat", timeout=5) + response = requests_lib.get(f"{self._base_url}/heartbeat", headers=self._headers, timeout=5) response.raise_for_status() except requests_lib.ConnectionError: raise ConnectionError(f"ProjectServer is not reachable at {self._base_url}. Make sure the server is running.") @@ -233,7 +252,7 @@ class ProjectServerClient: tool_params_json=tool_params_json, ).model_dump() - response = requests_lib.post(f"{self._base_url}/query_project", json=payload, timeout=self._timeout) + response = requests_lib.post(f"{self._base_url}/query_project", json=payload, headers=self._headers, timeout=self._timeout) response.raise_for_status() return response.text @@ -251,7 +270,7 @@ class ProjectServerClient: payload = CallFacadeMethodRequest( project_name=project_name, facade_name=facade_name, method_name=method_name, args=args, kwargs=kwargs ).model_dump() - response = requests_lib.post(f"{self._base_url}/call_facade_method", json=payload, timeout=self._timeout) + response = requests_lib.post(f"{self._base_url}/call_facade_method", json=payload, headers=self._headers, timeout=self._timeout) if not response.ok: raise ValueError(f"Project server error ({response.status_code}): {response.text[:2000]}") return pickle.loads(response.content) diff --git a/src/serena/repl/external_project.py b/src/serena/repl/external_project.py index 1e909b39..401c2d19 100644 --- a/src/serena/repl/external_project.py +++ b/src/serena/repl/external_project.py @@ -63,5 +63,5 @@ class ExternalProjectExecution: if self._client is None: from serena.project_server import ProjectServerClient - self._client = ProjectServerClient() + self._client = ProjectServerClient(self._agent.serena_config) return self._client.call_facade_method(self.project_name, facade_name, method_name, list(args), kwargs) diff --git a/src/serena/tools/query_project_tools.py b/src/serena/tools/query_project_tools.py index 514454b8..68e4eb3e 100644 --- a/src/serena/tools/query_project_tools.py +++ b/src/serena/tools/query_project_tools.py @@ -57,7 +57,7 @@ class QueryProjectTool(Tool, ToolMarkerOptional, ToolMarkerDoesNotRequireActiveP tool = self.agent.get_tool_by_name(tool_name) assert tool.is_readonly(), f"Tool {tool_name} is not read-only and cannot be executed in another project." if self._is_project_server_required(tool): - client = ProjectServerClient() + client = ProjectServerClient(self.agent.serena_config) return client.query_project(project_name, tool_name, tool_params_json) else: registered_project = self.agent.serena_config.get_registered_project(project_name) diff --git a/test/serena/test_external_projects.py b/test/serena/test_external_projects.py index a2bb73d7..ce74565c 100644 --- a/test/serena/test_external_projects.py +++ b/test/serena/test_external_projects.py @@ -61,8 +61,8 @@ def project_server(serena_config: SerenaConfig) -> Iterator[tuple[ProjectServer, @pytest.mark.python @pytest.mark.skipif(not language_server_tests_enabled(LanguageServerId.PYTHON), reason="python tests are disabled in this environment") def test_facade_method_results_are_transferred_from_the_project_server(project_server: tuple[ProjectServer, int]) -> None: - _, port = project_server - client = ProjectServerClient(port=port) + server, port = project_server + client = ProjectServerClient(server.get_serena_config(), port=port) result = client.call_facade_method("test_repo_python", "lsp", "find_symbol", ["create_user"], {"include_body": True}) # the result is a self-contained object which can be processed and rendered locally @@ -86,6 +86,7 @@ def test_external_project_context_in_repl( serena_config.included_apis = ["ext"] # the querying agent has another project active and queries the python test project + serena_config.auth_secret = server.get_auth_secret() agent = SerenaAgent(project="test_repo_typescript", serena_config=serena_config) agent.execute_task(lambda: None) try: diff --git a/test/serena/test_project_server.py b/test/serena/test_project_server.py index 1288e177..845603b1 100644 --- a/test/serena/test_project_server.py +++ b/test/serena/test_project_server.py @@ -8,8 +8,11 @@ from typing import Any, cast from unittest.mock import MagicMock import pytest +from flask import Flask +from werkzeug.serving import make_server -from serena.project_server import ProjectServer, QueryProjectRequest +from serena.config.serena_config import SerenaConfig +from serena.project_server import ProjectServer, ProjectServerClient, QueryProjectRequest @pytest.fixture @@ -23,6 +26,53 @@ def project_server() -> ProjectServer: return server +@pytest.fixture +def authenticated_server(project_server: ProjectServer, monkeypatch: pytest.MonkeyPatch) -> ProjectServer: + # expose the real HTTP routes with a query handler that needs no language servers + project_server._agent.serena_config.auth_secret = "test-shared-secret" + project_server._app = Flask(__name__) + monkeypatch.setattr(project_server, "_query_project", lambda req: req.project_name) + project_server._setup_routes() + return project_server + + +@pytest.mark.parametrize("authorization", [None, "Bearer wrong-secret", "test-shared-secret", "Bearer café"]) +@pytest.mark.parametrize("path", ["/heartbeat", "/query_project"]) +def test_project_server_rejects_invalid_credentials(authenticated_server: ProjectServer, authorization: str | None, path: str) -> None: + # unauthorized requests are rejected even before query payload validation + headers = {} if authorization is None else {"Authorization": authorization} + with authenticated_server._app.test_client() as client: + response = client.open(path, method="GET" if path == "/heartbeat" else "POST", headers=headers) + assert response.status_code == 401 + + +@pytest.mark.parametrize("use_wrong_password", [True, False]) +def test_project_server_client_authenticates_requests(authenticated_server: ProjectServer, use_wrong_password: bool) -> None: + # run the authenticated endpoints on an ephemeral local port + http_server = make_server("127.0.0.1", 0, authenticated_server._app) + thread = threading.Thread(target=http_server.serve_forever, daemon=True) + thread.start() + try: + + def check_client(): + serena_config = SerenaConfig() + serena_config.auth_secret = "wrong-secret" if use_wrong_password else authenticated_server.get_auth_secret() + client = ProjectServerClient(serena_config, port=http_server.server_port) + assert client.query_project("other", "find_symbol", "{}") == "other" + + # construction authenticates the heartbeat, raising a Connection error if using the wrong password + if use_wrong_password: + with pytest.raises(expected_exception=ConnectionError, match="401"): + check_client() + else: + check_client() + + finally: + http_server.shutdown() + thread.join(timeout=5) + http_server.server_close() + + def test_cached_project_lookup_is_not_blocked_by_unrelated_cold_load(project_server: ProjectServer) -> None: cached_root = Path("/cached") cold_root = Path("/cold")