From 3fd8525c4e8206441d66d0f1ba7a06b4b55628ea Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 11:49:18 +0300 Subject: [PATCH] fix: harden agent startup and websocket auth --- agent/requirements.txt | 3 +- agent/src/ss_tools/agent/langgraph_setup.py | 85 +++++++++++++++++-- .../tests/test_agent/test_langgraph_setup.py | 38 +++++++++ backend/src/app.py | 29 ++++++- backend/src/dependencies.py | 8 +- backend/src/scripts/create_admin.py | 10 ++- backend/tests/api/test_storage.py | 2 +- backend/tests/test_app_ws_auth.py | 46 +++++++++- .../components/tasks/TaskResultPanel.svelte | 34 ++++++-- .../TaskResultPanel.migration.test.ts | 57 ++++++++++++- 10 files changed, 283 insertions(+), 29 deletions(-) diff --git a/agent/requirements.txt b/agent/requirements.txt index 8a4b61282..7641c96ec 100644 --- a/agent/requirements.txt +++ b/agent/requirements.txt @@ -34,7 +34,8 @@ openpyxl # DB for LangGraph checkpoint psycopg2-binary -psycopg>=3.1 +# Bundle libpq with the agent environment; plain psycopg requires an OS-level libpq. +psycopg[binary]>=3.1 # Retry/utility tenacity>=8.0.0 diff --git a/agent/src/ss_tools/agent/langgraph_setup.py b/agent/src/ss_tools/agent/langgraph_setup.py index ab83e10c0..a663e25ad 100644 --- a/agent/src/ss_tools/agent/langgraph_setup.py +++ b/agent/src/ss_tools/agent/langgraph_setup.py @@ -9,7 +9,7 @@ import inspect as _inspect import os -from urllib.parse import urlsplit +from urllib.parse import urlsplit, urlunsplit from langchain_openai import ChatOpenAI from langgraph.checkpoint.memory import InMemorySaver @@ -56,20 +56,91 @@ _CHECKPOINTER_INIT = False _CHECKPOINTER_CONN = None -# #region AgentChat.LangGraph.Setup.InitCheckpointer [C:3] [TYPE Function] [SEMANTICS agent-chat,langgraph,checkpointer,postgres] +# #region AgentChat.LangGraph.Setup.RedactDbUrl [C:2] [TYPE Function] [SEMANTICS agent-chat,diagnostics,security] # @ingroup AgentChat -# @BRIEF Initialize AsyncPostgresSaver from DATABASE_URL env var. +# @BRIEF Redact password from a database URL for safe diagnostic logging. +# @INVARIANT Never returns raw password or full credentials in the output string. +def _redact_db_url(url: str) -> str: + """Return a copy of `url` with the password replaced by '***'.""" + if not url: + return "" + try: + parsed = urlsplit(url) + if not parsed.scheme or not parsed.hostname: + return "" + host = f"[{parsed.hostname}]" if ":" in parsed.hostname else parsed.hostname + port = f":{parsed.port}" if parsed.port else "" + user = f"{parsed.username}:***@" if parsed.username else "" + return urlunsplit((parsed.scheme, f"{user}{host}{port}", parsed.path, "", "")) + except Exception: + return "" +# #endregion AgentChat.LangGraph.Setup.RedactDbUrl + + +# #region AgentChat.LangGraph.Setup.InitCheckpointer [C:4] [TYPE Function] [SEMANTICS agent-chat,langgraph,checkpointer,postgres] +# @ingroup AgentChat +# @BRIEF Initialize AsyncPostgresSaver from DATABASE_URL env var with validation and diagnostics. # @SIDE_EFFECT Connects to PostgreSQL; creates checkpointer table via setup(). +# @INVARIANT DATABASE_URL must be a valid PostgreSQL URI before connection is attempted. +# @RATIONALE Production agent crash-loop was caused by an empty or misconfigured DATABASE_URL +# that resolved to a nonexistent Unix socket. Pre-connection validation + redacted +# diagnostics (host, port, db name — never password) allows operators to debug +# configuration failures without exposing secrets in logs. async def init_checkpointer() -> None: global _CHECKPOINTER, _CHECKPOINTER_INIT, _CHECKPOINTER_CONN if _CHECKPOINTER_INIT: return db_url = os.getenv("DATABASE_URL") - pg_url = db_url.replace("postgresql+psycopg2://", "postgres://").replace("postgresql://", "postgres://") - _CHECKPOINTER_CONN = await psycopg.AsyncConnection.connect(pg_url, autocommit=True, row_factory=dict_row) - _CHECKPOINTER = AsyncPostgresSaver(_CHECKPOINTER_CONN) - await _CHECKPOINTER.setup() + if not db_url or not db_url.strip(): + logger.explore( + "DATABASE_URL env var is missing or empty — checkpointer cannot be initialized", + payload={"env_var": "DATABASE_URL"}, + error="DATABASE_URL is not set", + ) + raise RuntimeError("DATABASE_URL is not set. PostgreSQL checkpointer requires a valid database URI.") + + # Redact password from URL for safe diagnostic logging. + _sanitized = _redact_db_url(db_url) + logger.reason( + "Initializing PostgreSQL checkpointer", + payload={"sanitized_url": _sanitized}, + ) + + pg_url: str = db_url.replace("postgresql+psycopg2://", "postgres://").replace("postgresql://", "postgres://") + if not pg_url.startswith("postgres://") and not pg_url.startswith("postgresql://"): + logger.explore( + "DATABASE_URL does not look like a PostgreSQL connection string", + payload={"sanitized_url": _sanitized, "parsed_scheme": pg_url.split("://")[0] if "://" in pg_url else "none"}, + error="Not a PostgreSQL URL", + ) + raise RuntimeError(f"DATABASE_URL must be a PostgreSQL URI. Got invalid scheme.") + + try: + _CHECKPOINTER_CONN = await psycopg.AsyncConnection.connect(pg_url, autocommit=True, row_factory=dict_row) + except Exception as e: + logger.explore( + "Failed to connect to PostgreSQL checkpointer", + payload={"sanitized_url": _sanitized, "exception_type": type(e).__name__}, + error="PostgreSQL connection failed; inspect database service availability and credentials", + ) + raise + + try: + _CHECKPOINTER = AsyncPostgresSaver(_CHECKPOINTER_CONN) + await _CHECKPOINTER.setup() + except Exception as e: + logger.explore( + "Checkpointer setup() failed", + payload={"sanitized_url": _sanitized, "exception_type": type(e).__name__}, + error="PostgreSQL checkpointer setup failed; inspect database schema permissions", + ) + raise + _CHECKPOINTER_INIT = True + logger.reason( + "PostgreSQL checkpointer initialized successfully", + payload={"sanitized_url": _sanitized}, + ) # #endregion AgentChat.LangGraph.Setup.InitCheckpointer _llm_config: dict | None = None diff --git a/agent/tests/test_agent/test_langgraph_setup.py b/agent/tests/test_agent/test_langgraph_setup.py index 4b31f0654..39a3b84af 100644 --- a/agent/tests/test_agent/test_langgraph_setup.py +++ b/agent/tests/test_agent/test_langgraph_setup.py @@ -64,6 +64,44 @@ def test_llm_diagnostics_redacts_api_key_and_path(): # #endregion Test.AgentChat.TestLlmDiagnostics +# #region Test.AgentChat.TestCheckpointerDiagnostics [C:3] [TYPE Class] [SEMANTICS test,agent,checkpointer,security] +# @BRIEF Verify checkpointer diagnostics redact credentials and reject invalid configuration safely. +# @RELATION BINDS_TO -> [AgentChat.LangGraph.Setup.InitCheckpointer] +class TestCheckpointerDiagnostics: + def test_redact_db_url_removes_password_query_and_fragment(self): + from ss_tools.agent.langgraph_setup import _redact_db_url + + redacted = _redact_db_url("postgresql://user:secret@[::1]:5432/ss_tools?sslpassword=hidden#fragment") + assert redacted == "postgresql://user:***@[::1]:5432/ss_tools" + assert "secret" not in redacted + assert "hidden" not in redacted + + def test_redact_db_url_handles_unparseable_value(self): + from ss_tools.agent.langgraph_setup import _redact_db_url + + assert _redact_db_url("not a database url") == "" + + @pytest.mark.anyio + async def test_init_checkpointer_rejects_missing_database_url(self, monkeypatch): + import ss_tools.agent.langgraph_setup as ls + + monkeypatch.delenv("DATABASE_URL", raising=False) + ls._CHECKPOINTER_INIT = False + with pytest.raises(RuntimeError, match="DATABASE_URL is not set"): + await ls.init_checkpointer() + + @pytest.mark.anyio + async def test_init_checkpointer_redacts_connection_failure(self, monkeypatch): + import ss_tools.agent.langgraph_setup as ls + + monkeypatch.setenv("DATABASE_URL", "postgresql://agent:secret@db:5432/ss_tools?sslpassword=hidden") + ls._CHECKPOINTER_INIT = False + with patch("ss_tools.agent.langgraph_setup.psycopg.AsyncConnection.connect", new=AsyncMock(side_effect=Exception("secret"))): + with pytest.raises(Exception, match="secret"): + await ls.init_checkpointer() +# #endregion Test.AgentChat.TestCheckpointerDiagnostics + + # #region Test.AgentChat.TestCreateAgent [C:2] [TYPE Function] # @BRIEF Test create_agent with various LLM config states. class TestCreateAgent: diff --git a/backend/src/app.py b/backend/src/app.py index fa8d195ea..14e6398c1 100755 --- a/backend/src/app.py +++ b/backend/src/app.py @@ -20,7 +20,6 @@ import asyncio from contextlib import asynccontextmanager import os from pathlib import Path -import sys import uuid # project_root is used for static files mounting @@ -669,28 +668,52 @@ def _authorize_websocket(websocket: WebSocket, resource: str, action: str) -> bo try: from .core.auth.jwt import decode_token from .core.database import SessionLocal - from .models.auth import User, Role + from .models.auth import User payload = decode_token(ws_token) username = payload.get("sub") if not isinstance(username, str) or not username: + logger.explore( + "WebSocket authorization — token missing 'sub' claim", + payload={"resource": resource, "action": action}, + error="JWT payload has no valid subject", + ) return False db = SessionLocal() try: user = db.query(User).filter(User.username == username).first() if not user: + logger.explore( + "WebSocket authorization — user not found in database", + payload={"resource": resource, "action": action, "username": username}, + error="Authenticated user missing from local DB", + ) return False if not getattr(user, "is_active", True): + logger.explore( + "WebSocket authorization — user is inactive", + payload={"resource": resource, "action": action, "username": username}, + error="User account is disabled", + ) return False - # Admin bypass via is_admin flag + # is_admin is the sole administrative authority. The database migration + # backfills the legacy Admin role before this authorization path runs. for role in user.roles: if getattr(role, "is_admin", False): return True for perm in role.permissions: if perm.resource == resource and perm.action == action: return True + + logger.explore( + "WebSocket authorization denied — no matching permission or admin role", + payload={"resource": resource, "action": action, "user": username, + "roles": [r.name for r in user.roles], + "is_active": getattr(user, "is_active", True)}, + error="No matching RBAC permission", + ) return False finally: db.close() diff --git a/backend/src/dependencies.py b/backend/src/dependencies.py index e21853895..2fee8301c 100755 --- a/backend/src/dependencies.py +++ b/backend/src/dependencies.py @@ -831,12 +831,8 @@ def has_permission(resource: str, action: str): if perm.resource == resource and perm.action == action: return current_user - # Special case for Admin role (full access) — uses is_admin flag, not name string. - # Fallback to role.name == "Admin" for roles created before is_admin migration. - if any( - getattr(role, "is_admin", False) or role.name == "Admin" - for role in current_user.roles - ): + # is_admin is the single source of truth for the administrative bypass. + if any(getattr(role, "is_admin", False) for role in current_user.roles): return current_user from .core.auth.logger import log_security_event diff --git a/backend/src/scripts/create_admin.py b/backend/src/scripts/create_admin.py index a2ba77543..f0d038aa3 100644 --- a/backend/src/scripts/create_admin.py +++ b/backend/src/scripts/create_admin.py @@ -44,10 +44,18 @@ def create_admin(username, password, email=None): admin_role = db.query(Role).filter(Role.name == "Admin").first() if not admin_role: logger.reason("Creating Admin role") - admin_role = Role(name="Admin", description="System Administrator") + admin_role = Role( + name="Admin", + description="System Administrator", + is_admin=True, + ) db.add(admin_role) db.commit() db.refresh(admin_role) + elif not admin_role.is_admin: + logger.reason("Marking existing Admin role as administrative") + admin_role.is_admin = True + db.commit() # 2. Check if user already exists existing_user = db.query(User).filter(User.username == username).first() diff --git a/backend/tests/api/test_storage.py b/backend/tests/api/test_storage.py index b929f1169..f68548e59 100644 --- a/backend/tests/api/test_storage.py +++ b/backend/tests/api/test_storage.py @@ -37,7 +37,7 @@ def _make_client(overrides: dict | None = None) -> tuple[TestClient, MagicMock]: id="admin-1", username="admin", email="admin@x.com", auth_source="LOCAL", created_at=__import__("datetime").datetime.now(), - roles=[RoleSchema(id="r1", name="Admin", description="", permissions=[])], + roles=[RoleSchema(id="r1", name="Admin", description="", is_admin=True, permissions=[])], ) mock_plugin_loader = MagicMock() diff --git a/backend/tests/test_app_ws_auth.py b/backend/tests/test_app_ws_auth.py index 128dcb303..0b7c6024a 100644 --- a/backend/tests/test_app_ws_auth.py +++ b/backend/tests/test_app_ws_auth.py @@ -9,11 +9,12 @@ from pathlib import Path import sys +from types import SimpleNamespace sys.path.insert(0, str(Path(__file__).parent.parent / "src")) import pytest -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch class TestAuthenticateWebsocket: @@ -117,4 +118,47 @@ class TestAuthenticateWebsocketApiKeyException: msl.return_value = db assert await _authenticate_websocket(ws, "ws/logs") is False # #endregion Test.AppModule.TestApikeyDbException + + +# #region Test.AppModule.AuthorizeWebsocket [C:3] [TYPE Class] [SEMANTICS test,app,ws,authorization,rbac] +# @BRIEF Verify WebSocket authorization grants access only through is_admin or explicit permissions. +# @RELATION BINDS_TO -> [App.AppModule.AuthorizeWebsocket] +class TestAuthorizeWebsocket: + def _authorize(self, user, resource="tasks", action="READ"): + from src.app import _authorize_websocket + + ws = MagicMock() + ws.query_params = {"token": "valid.jwt"} + db = MagicMock() + db.query.return_value.filter.return_value.first.return_value = user + with ( + patch("src.core.auth.jwt.decode_token", return_value={"sub": "testuser"}), + patch("src.core.database.SessionLocal", return_value=db), + ): + result = _authorize_websocket(ws, resource, action) + db.close.assert_called_once() + return result + + def test_is_admin_role_bypasses_permission_check(self): + user = SimpleNamespace( + is_active=True, + roles=[SimpleNamespace(name="Admin", is_admin=True, permissions=[])], + ) + assert self._authorize(user) is True + + def test_legacy_admin_name_without_flag_does_not_bypass_permissions(self): + user = SimpleNamespace( + is_active=True, + roles=[SimpleNamespace(name="Admin", is_admin=False, permissions=[])], + ) + assert self._authorize(user) is False + + def test_explicit_permission_allows_websocket_access(self): + permission = SimpleNamespace(resource="tasks", action="READ") + user = SimpleNamespace( + is_active=True, + roles=[SimpleNamespace(name="Operator", is_admin=False, permissions=[permission])], + ) + assert self._authorize(user) is True +# #endregion Test.AppModule.AuthorizeWebsocket # #endregion Test.AppModule.WsAuth diff --git a/frontend/src/lib/components/tasks/TaskResultPanel.svelte b/frontend/src/lib/components/tasks/TaskResultPanel.svelte index e159ebb18..b500ff9d6 100644 --- a/frontend/src/lib/components/tasks/TaskResultPanel.svelte +++ b/frontend/src/lib/components/tasks/TaskResultPanel.svelte @@ -1,13 +1,19 @@ - + + + +