diff --git a/muxplex/main.py b/muxplex/main.py index d86d7ee..5a32878 100644 --- a/muxplex/main.py +++ b/muxplex/main.py @@ -725,50 +725,19 @@ async def create_session(payload: CreateSessionPayload) -> dict: return {"name": name, "ok": True} -@app.post("/api/sessions/{name}/connect") -async def connect_session(name: str) -> dict: - """Connect to a tmux session via ttyd. - - Uses the process pool: if the session already has a live ttyd, returns - immediately. Otherwise spawns a new ttyd on a free port, polls until - listening, and returns the port. - - Returns {active_session: name, ttyd_port: }. - Raises HTTP 404 if *name* is not in the known session list (when non-empty). - """ - known = get_session_list() - if known and name not in known: - raise HTTPException(status_code=404, detail=f"Session '{name}' not found") - - _log.info("Connecting to session '%s'", name) - port = await get_or_spawn(name) - - async with state_lock: - state = load_state() - state["active_session"] = name - save_state(state) - - return {"active_session": name, "ttyd_port": port} - - @app.delete("/api/sessions/current") async def delete_current_session() -> dict: - """Disconnect the current ttyd session. + """Clear the active session. - Kills the pooled ttyd for the active session and clears active_session - in persistent state. + Sets active_session to None in persistent state. Returns {active_session: None}. """ async with state_lock: state = load_state() - active = state.get("active_session") state["active_session"] = None save_state(state) - if active: - await kill_session(active) - return {"active_session": None} @@ -816,10 +785,6 @@ async def delete_session(name: str) -> dict: except Exception: _log.warning("Delete command failed: %r", command) - # Clean up the pool entry (ttyd will die when tmux session dies, - # but clean up immediately so the port is freed). - await kill_session(name) - return {"ok": True, "name": name} @@ -1017,127 +982,6 @@ async def instance_info() -> dict: } -# --------------------------------------------------------------------------- -# WebSocket proxy — bridges browser to ttyd (eliminates Caddy dependency) -# --------------------------------------------------------------------------- - - -# _ttyd_is_listening is imported from muxplex.ttyd (pool-aware, accepts port param) - - -async def _ws_auth_check(websocket: WebSocket) -> bool: - """Return True if the WebSocket caller is authorized. - - Closes the WebSocket with code 4001 and returns False if the caller - is not authorized. Localhost connections (127.0.0.1 / ::1) are - unconditionally trusted. Remote callers must present a valid - ``muxplex_session`` cookie OR a Bearer token matching ``_federation_key``. - """ - host = websocket.client.host if websocket.client else "" - if host in ("127.0.0.1", "::1"): - return True - session_cookie = websocket.cookies.get("muxplex_session") - cookie_ok = session_cookie and verify_session_cookie( - _auth_secret, session_cookie, _auth_ttl - ) - bearer_ok = False - if _federation_key: - auth_header = websocket.headers.get("authorization", "") - if auth_header.lower().startswith("bearer "): - bearer_ok = hmac.compare_digest(auth_header[7:], _federation_key) - if not cookie_ok and not bearer_ok: - await websocket.close(code=4001) - return False - return True - - -@app.websocket("/terminal/ws") -async def terminal_ws_proxy(websocket: WebSocket) -> None: - """Proxy WebSocket frames between the browser and ttyd. - - Resolves the target ttyd port from the process pool BEFORE accepting the - browser WebSocket. If ttyd is not in the pool (e.g. after a service - restart), auto-spawns it using the session from query params or - active_session from state, then waits for it to bind its port. - - Only after ttyd is confirmed reachable does the function call - websocket.accept() — so the browser's 'open' event only fires once a real - relay is possible. This prevents the reconnect-counter bounce bug where - the proxy accepted immediately (resetting _reconnectAttempts to 0) and - then closed as soon as it couldn't reach the dead ttyd. - """ - # Auth check before accepting — BaseHTTPMiddleware doesn't cover WebSocket scope - if not await _ws_auth_check(websocket): - return - - # Determine which session (and port) to connect to. - # Priority: ?session= query param > active_session from state. - session = websocket.query_params.get("session", "") - if not session: - try: - async with state_lock: - state = load_state() - session = state.get("active_session") or "" - except Exception: - session = "" - - # Look up port from pool, or auto-spawn if not listening - port = pool_port(session) if session else None - - if port is None and session: - if not _ttyd_is_listening(pool_port(session) or TTYD_PORT): - _log.info( - "WS proxy: ttyd not listening, auto-spawning for '%s'", - session, - ) - try: - port = await get_or_spawn(session) - except Exception as exc: - _log.warning("WS proxy: failed to auto-spawn ttyd: %s", exc) - - if port is None: - # Fallback to TTYD_PORT (legacy single-process or externally managed ttyd) - port = TTYD_PORT - - await websocket.accept(subprotocol="tty") - - ttyd_url = f"ws://localhost:{port}/ws" - try: - async with websockets.connect( - ttyd_url, subprotocols=[Subprotocol("tty")] - ) as ttyd_ws: - - async def client_to_ttyd() -> None: - try: - while True: - msg = await websocket.receive() - if msg.get("bytes"): - await ttyd_ws.send(msg["bytes"]) - elif msg.get("text"): - await ttyd_ws.send(msg["text"]) - except Exception as exc: - _log.debug("ws relay closed (client_to_ttyd): %s", exc) - - async def ttyd_to_client() -> None: - try: - async for message in ttyd_ws: - if isinstance(message, bytes): - await websocket.send_bytes(message) - else: - await websocket.send_text(message) - except Exception as exc: - _log.debug("ws relay closed (ttyd_to_client): %s", exc) - - await asyncio.gather(client_to_ttyd(), ttyd_to_client()) - except Exception as exc: - _log.debug("ws proxy closed: %s", exc) - finally: - try: - await websocket.close() - except Exception: - pass - - # --------------------------------------------------------------------------- # Federation helper utilities # --------------------------------------------------------------------------- @@ -1191,9 +1035,21 @@ async def federation_terminal_ws_proxy(websocket: WebSocket, device_id: str) -> Auth check uses the same cookie + bearer pattern as terminal_ws_proxy. Closes with code 4004 if device_id does not match any remote. """ - # Auth check before accepting — same pattern as terminal_ws_proxy - if not await _ws_auth_check(websocket): - return + # Auth check before accepting — BaseHTTPMiddleware doesn't cover WS scope + host = websocket.client.host if websocket.client else "" + if host not in ("127.0.0.1", "::1"): + session_cookie = websocket.cookies.get("muxplex_session") + cookie_ok = session_cookie and verify_session_cookie( + _auth_secret, session_cookie, _auth_ttl + ) + bearer_ok = False + if _federation_key: + auth_header = websocket.headers.get("authorization", "") + if auth_header.lower().startswith("bearer "): + bearer_ok = hmac.compare_digest(auth_header[7:], _federation_key) + if not cookie_ok and not bearer_ok: + await websocket.close(code=4001) + return # Look up remote instance by device_id remote = _lookup_remote_by_device_id(device_id) diff --git a/muxplex/tests/test_api.py b/muxplex/tests/test_api.py index 10d573d..6f770c2 100644 --- a/muxplex/tests/test_api.py +++ b/muxplex/tests/test_api.py @@ -356,76 +356,13 @@ def test_get_sessions_returns_empty_list_when_no_sessions(client, monkeypatch): assert response.json() == [] -# --------------------------------------------------------------------------- -# POST /api/sessions/{name}/connect -# --------------------------------------------------------------------------- - - -def test_connect_session_returns_200(client, monkeypatch): - """POST /api/sessions/{name}/connect returns 200 and correct body when session exists.""" - monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha"]) - - async def mock_get_or_spawn(name): - return 7682 - - monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn) - - response = client.post("/api/sessions/alpha/connect") - assert response.status_code == 200 - data = response.json() - assert data["active_session"] == "alpha" - assert data["ttyd_port"] == 7682 - - -def test_connect_session_sets_active_session(client, monkeypatch): - """POST /api/sessions/{name}/connect persists active_session to state.""" - from muxplex.state import load_state - - monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha"]) - - async def mock_get_or_spawn(name): - return 7682 - - monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn) - - client.post("/api/sessions/alpha/connect") - - state = load_state() - assert state["active_session"] == "alpha" - - -def test_connect_session_calls_get_or_spawn(client, monkeypatch): - """POST /api/sessions/{name}/connect calls get_or_spawn with the session name.""" - monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha"]) - - spawn_calls = [] - - async def mock_get_or_spawn(name): - spawn_calls.append(name) - return 7682 - - monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn) - - response = client.post("/api/sessions/alpha/connect") - assert response.status_code == 200 - assert spawn_calls == ["alpha"] - - -def test_connect_nonexistent_session_returns_404(client, monkeypatch): - """POST /api/sessions/{name}/connect returns 404 when session is not in list.""" - monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["alpha", "beta"]) - - response = client.post("/api/sessions/gamma/connect") - assert response.status_code == 404 - - # --------------------------------------------------------------------------- # DELETE /api/sessions/current # --------------------------------------------------------------------------- -def test_delete_current_kills_ttyd_and_clears_active(client, monkeypatch): - """DELETE /api/sessions/current kills the session's ttyd and clears active_session.""" +def test_delete_current_clears_active_session(client, monkeypatch): + """DELETE /api/sessions/current clears active_session in state.""" from muxplex.state import load_state, save_state # Set up initial state with active session @@ -438,19 +375,10 @@ def test_delete_current_kills_ttyd_and_clears_active(client, monkeypatch): } ) - kill_session_calls = [] - - async def mock_kill_session(name): - kill_session_calls.append(name) - return True - - monkeypatch.setattr("muxplex.main.kill_session", mock_kill_session) - response = client.delete("/api/sessions/current") assert response.status_code == 200 data = response.json() assert data["active_session"] is None - assert kill_session_calls == ["alpha"] # Verify state was persisted state = load_state() @@ -800,20 +728,6 @@ def test_api_routes_not_shadowed(client): assert isinstance(response.json(), list) -def test_terminal_ws_route_exists(): - """The app must have a WebSocket route registered at /terminal/ws.""" - from fastapi.routing import APIRoute, APIWebSocketRoute - - from muxplex.main import app - - ws_routes = [ - r - for r in app.routes - if isinstance(r, (APIRoute, APIWebSocketRoute)) and r.path == "/terminal/ws" - ] - assert len(ws_routes) == 1, "Expected exactly one /terminal/ws route" - - # --------------------------------------------------------------------------- # Auth middleware integration # --------------------------------------------------------------------------- @@ -967,90 +881,6 @@ def test_logout_clears_session_cookie(monkeypatch): assert "max-age=0" in set_cookie.lower() -# --------------------------------------------------------------------------- -# WebSocket auth tests -# --------------------------------------------------------------------------- - - -def _wrap_with_client_host(wrapped_app, host: str): - """Return an ASGI wrapper that forces websocket scope client to `host`. - - This lets tests simulate a WebSocket connection appearing to originate - from a specific IP without touching Starlette internals. - """ - - async def _middleware(scope, receive, send): - if scope.get("type") == "websocket": - scope = {**scope, "client": (host, 50000)} - await wrapped_app(scope, receive, send) - - return _middleware - - -def test_ws_localhost_no_cookie_bypasses_auth(): - """WebSocket from 127.0.0.1 is accepted even without a session cookie.""" - from starlette.websockets import WebSocketDisconnect - - # Force scope to look like localhost so auth check is bypassed - localhost_app = _wrap_with_client_host(app, "127.0.0.1") - - with TestClient(localhost_app) as c: - try: - with c.websocket_connect("/terminal/ws") as _: - pass # connection was accepted — auth bypassed for localhost - except WebSocketDisconnect as e: - # The websocket was accepted (auth bypassed); ttyd is not running so - # the proxy fails and closes with a non-4001 code. - assert e.code != 4001, ( - f"Localhost WebSocket should not be rejected; got close code {e.code}" - ) - - -def test_ws_valid_cookie_non_localhost_not_rejected_4001(): - """WebSocket from non-localhost with a valid cookie is not rejected with 4001.""" - from starlette.websockets import WebSocketDisconnect - from muxplex.auth import create_session_cookie - from muxplex.main import _auth_secret, _auth_ttl - - cookie = create_session_cookie(_auth_secret, _auth_ttl) - - # TestClient default host is "testclient" — treated as non-localhost - with TestClient(app) as c: - c.cookies["muxplex_session"] = cookie - try: - with c.websocket_connect("/terminal/ws") as _: - pass # connection was accepted — auth passed - except WebSocketDisconnect as e: - # Auth passed; ttyd not running → proxy fails → close with code != 4001 - assert e.code != 4001, ( - f"Valid-cookie WebSocket should not be rejected; got close code {e.code}" - ) - - -def test_ws_no_cookie_non_localhost_rejected_4001(): - """WebSocket from non-localhost without a cookie is closed with code 4001.""" - from starlette.websockets import WebSocketDisconnect - - # TestClient default host "testclient" is treated as non-localhost - with TestClient(app) as c: - with pytest.raises(WebSocketDisconnect) as exc_info: - with c.websocket_connect("/terminal/ws") as _: - pass - assert exc_info.value.code == 4001 - - -def test_ws_invalid_cookie_non_localhost_rejected_4001(): - """WebSocket from non-localhost with a tampered cookie is closed with code 4001.""" - from starlette.websockets import WebSocketDisconnect - - with TestClient(app) as c: - c.cookies["muxplex_session"] = "tampered.invalid.cookie.value" - with pytest.raises(WebSocketDisconnect) as exc_info: - with c.websocket_connect("/terminal/ws") as _: - pass - assert exc_info.value.code == 4001 - - # --------------------------------------------------------------------------- # GET /api/settings # --------------------------------------------------------------------------- @@ -2464,30 +2294,6 @@ def test_create_session_logs_command(client, monkeypatch, tmp_path, caplog): ) -def test_connect_session_logs_session_name(client, monkeypatch, caplog): - """POST /api/sessions/{name}/connect must log the session name at INFO level.""" - import logging - - monkeypatch.setattr("muxplex.main.get_session_list", lambda: ["target-session"]) - - async def mock_kill_ttyd(): - pass - - async def mock_spawn_ttyd(name): - pass - - monkeypatch.setattr("muxplex.main.kill_ttyd", mock_kill_ttyd) - monkeypatch.setattr("muxplex.main.spawn_ttyd", mock_spawn_ttyd) - - with caplog.at_level(logging.INFO, logger="muxplex.main"): - client.post("/api/sessions/target-session/connect") - - log_messages = "\n".join(caplog.messages) - assert "target-session" in log_messages, ( - f"connect_session must log session name at INFO level, got logs:\n{log_messages}" - ) - - def test_cli_uvicorn_log_level_is_info(): """cli.py serve() must pass log_level='info' to uvicorn.run so logs appear in journalctl.""" import inspect diff --git a/muxplex/tests/test_ttyd.py b/muxplex/tests/test_ttyd.py deleted file mode 100644 index d79a8eb..0000000 --- a/muxplex/tests/test_ttyd.py +++ /dev/null @@ -1,448 +0,0 @@ -""" -Tests for coordinator/ttyd.py — ttyd process lifecycle management. -All 11 acceptance-criteria tests are defined here, plus 5 tests for the -_pids_on_port() lsof→fuser→ss fallback chain (the root-cause guard for systems -where lsof is not installed). -""" - -import signal -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -import muxplex.ttyd as ttyd_mod -from muxplex.ttyd import ( - _kill_pids_on_port, - _pids_on_port, - kill_orphan_ttyd, - kill_ttyd, - spawn_ttyd, -) - - -# --------------------------------------------------------------------------- -# autouse fixture — redirect PID paths to tmp_path for every test -# --------------------------------------------------------------------------- - - -@pytest.fixture(autouse=True) -def use_tmp_pid_dir(tmp_path, monkeypatch): - """Redirect PID file I/O to a fresh temp directory for every test.""" - tmp_pid_dir = tmp_path / "ttyd" - tmp_pid_path = tmp_pid_dir / "ttyd.pid" - monkeypatch.setattr("muxplex.ttyd.TTYD_PID_DIR", tmp_pid_dir) - monkeypatch.setattr("muxplex.ttyd.TTYD_PID_PATH", tmp_pid_path) - # Clear the pool between tests - ttyd_mod._pool.clear() - - -# --------------------------------------------------------------------------- -# Helper for mocking asyncio.create_subprocess_exec -# --------------------------------------------------------------------------- - - -def _make_mock_ttyd_process(pid: int = 12345): - """Return a mock ttyd process with the given PID.""" - proc = MagicMock() - proc.pid = pid - return proc - - -# --------------------------------------------------------------------------- -# spawn_ttyd tests -# --------------------------------------------------------------------------- - - -async def test_spawn_ttyd_writes_pid_file(): - """spawn_ttyd() must write the process PID to TTYD_PID_PATH.""" - mock_proc = _make_mock_ttyd_process(pid=99999) - - with patch( - "asyncio.create_subprocess_exec", - new=AsyncMock(return_value=mock_proc), - ): - await spawn_ttyd("my-session") - - pid_path = ttyd_mod.TTYD_PID_PATH - assert pid_path.exists(), "PID file was not created" - assert pid_path.read_text().strip() == "99999" - - -async def test_spawn_ttyd_uses_correct_command(): - """spawn_ttyd() must call ttyd with args: -W -m 3 -p 7682 tmux attach -t .""" - mock_proc = _make_mock_ttyd_process(pid=54321) - - with patch( - "asyncio.create_subprocess_exec", - new=AsyncMock(return_value=mock_proc), - ) as mock_create: - await spawn_ttyd("test-session") - - call_args = mock_create.call_args[0] - assert list(call_args) == [ - "ttyd", - "-W", - "-m", - "3", - "-p", - "7682", - "tmux", - "attach", - "-t", - "test-session", - ] - - -async def test_spawn_ttyd_returns_process_object(): - """spawn_ttyd() must return the process object from create_subprocess_exec.""" - mock_proc = _make_mock_ttyd_process(pid=11111) - - with patch( - "asyncio.create_subprocess_exec", - new=AsyncMock(return_value=mock_proc), - ): - result = await spawn_ttyd("another-session") - - assert result is mock_proc - - -# --------------------------------------------------------------------------- -# kill_ttyd tests -# --------------------------------------------------------------------------- - - -async def test_kill_ttyd_returns_false_when_no_pid_file(): - """kill_ttyd() returns False when no PID file exists.""" - # autouse fixture ensures no PID file is present - result = await kill_ttyd() - assert result is False - - -async def test_kill_ttyd_reads_pid_file_and_sends_sigterm(): - """kill_ttyd() reads the PID file and sends SIGTERM to the running process.""" - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text("12345") - - kill_calls = [] - - def mock_os_kill(pid, sig): - kill_calls.append((pid, sig)) - # First existence-check (sig=0) succeeds; subsequent sig=0 calls raise - if sig == 0 and sum(1 for _, s in kill_calls if s == 0) > 1: - raise ProcessLookupError - - with patch("os.kill", side_effect=mock_os_kill): - result = await kill_ttyd() - - assert result is True - sigterm_calls = [(pid, sig) for pid, sig in kill_calls if sig == signal.SIGTERM] - assert len(sigterm_calls) == 1 - assert sigterm_calls[0][0] == 12345 - - -async def test_kill_ttyd_removes_pid_file(): - """kill_ttyd() removes the PID file regardless of whether process was alive.""" - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text("12345") - - def mock_os_kill(pid, sig): - # All os.kill calls raise ProcessLookupError — simulates already-dead process - raise ProcessLookupError - - with patch("os.kill", side_effect=mock_os_kill): - result = await kill_ttyd() - - assert result is True - assert not pid_path.exists(), "PID file should be removed after kill_ttyd()" - - -async def test_kill_ttyd_handles_process_already_dead(): - """kill_ttyd() returns True and clears state when process is already gone.""" - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text("99999") - - # Simulate process already dead: os.kill(pid, 0) raises ProcessLookupError - with patch("os.kill", side_effect=ProcessLookupError): - result = await kill_ttyd() - - assert result is True - assert not pid_path.exists(), ( - "PID file should be removed when process was already dead" - ) - assert ttyd_mod._active_process is None - - -# --------------------------------------------------------------------------- -# kill_orphan_ttyd tests -# -# kill_orphan_ttyd() is a thin delegation to kill_ttyd(). These tests verify -# both the delegation wiring and that behaviour is consistent with kill_ttyd(). -# --------------------------------------------------------------------------- - - -async def test_kill_orphan_ttyd_returns_false_when_no_pid_file(): - """kill_orphan_ttyd() returns False when no PID file exists (no orphan).""" - # autouse fixture ensures no PID file is present - result = await kill_orphan_ttyd() - assert result is False - - -async def test_kill_orphan_ttyd_kills_running_process(): - """kill_orphan_ttyd() kills a running orphan process and returns True.""" - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text("55555") - - kill_calls = [] - - def mock_os_kill(pid, sig): - kill_calls.append((pid, sig)) - # First existence-check (sig=0) succeeds; subsequent sig=0 calls raise - if sig == 0 and sum(1 for _, s in kill_calls if s == 0) > 1: - raise ProcessLookupError - - with patch("os.kill", side_effect=mock_os_kill): - result = await kill_orphan_ttyd() - - assert result is True - assert not pid_path.exists(), "PID file should be removed after kill_orphan_ttyd()" - sigterm_calls = [(pid, sig) for pid, sig in kill_calls if sig == signal.SIGTERM] - assert len(sigterm_calls) == 1 - assert sigterm_calls[0][0] == 55555 - - -async def test_kill_orphan_ttyd_handles_pid_file_with_dead_process(): - """kill_orphan_ttyd() handles a stale PID file whose process is already gone.""" - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text("77777") - - with patch("os.kill", side_effect=ProcessLookupError): - result = await kill_orphan_ttyd() - - assert result is True - assert not pid_path.exists(), "PID file should be removed after orphan cleanup" - - -async def test_kill_ttyd_kills_orphan_on_port_when_pid_file_desynced(): - """kill_ttyd() must also kill orphaned ttyd processes on TTYD_PORT. - - If the PID file points to a dead process (desynced), but the REAL ttyd is - still running on TTYD_PORT (orphaned from a previous spawn), kill_ttyd() - must find and kill it via lsof -ti :. This is the belt-and-suspenders - fallback that prevents the session-switching bug where a new ttyd cannot bind - the port because the old one is still running. - """ - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - # PID file with a dead process (desynced) - pid_path.write_text("99999") - - killed_pids: list[tuple[int, int]] = [] - - def mock_os_kill(pid: int, sig: int) -> None: - if pid == 99999 and sig == 0: - raise ProcessLookupError # PID file process is already dead - killed_pids.append((pid, sig)) - - mock_lsof_result = MagicMock() - mock_lsof_result.returncode = 0 - mock_lsof_result.stdout = "12345\n" # orphan PID occupying TTYD_PORT - - def mock_subprocess_run(cmd, **kwargs): # noqa: ANN001 - result = MagicMock() - if "lsof" in cmd and "-ti" in cmd: - return mock_lsof_result - result.returncode = 1 - result.stdout = "" - return result - - with ( - patch("os.kill", side_effect=mock_os_kill), - patch("muxplex.ttyd._subprocess.run", side_effect=mock_subprocess_run), - ): - result = await kill_ttyd() - - assert result is True, "kill_ttyd must return True when orphan found on port" - orphan_killed = any( - pid == 12345 and sig == signal.SIGTERM for pid, sig in killed_pids - ) - assert orphan_killed, ( - "kill_ttyd must send SIGTERM to orphan process (12345) found via " - "lsof on TTYD_PORT when PID file is desynced" - ) - - -async def test_spawn_ttyd_force_kills_process_on_port_before_binding(): - """spawn_ttyd() must force-kill any process occupying TTYD_PORT before spawning. - - If kill_ttyd() completed but the port is still occupied (race condition), - spawn_ttyd() must do a final SIGKILL cleanup so the new ttyd can bind. - """ - mock_proc = _make_mock_ttyd_process(pid=22222) - - # First lsof call returns an occupant; second call (after kill) returns empty - call_count = 0 - - def mock_subprocess_run(cmd, **kwargs): # noqa: ANN001 - nonlocal call_count - result = MagicMock() - if "lsof" in cmd and "-ti" in cmd: - call_count += 1 - if call_count == 1: - result.returncode = 0 - result.stdout = "54321\n" # port still occupied - return result - result.returncode = 1 - result.stdout = "" - return result - - killed_pids: list[tuple[int, int]] = [] - - def mock_os_kill(pid: int, sig: int) -> None: - killed_pids.append((pid, sig)) - - with ( - patch("asyncio.create_subprocess_exec", new=AsyncMock(return_value=mock_proc)), - patch("muxplex.ttyd._subprocess.run", side_effect=mock_subprocess_run), - patch("os.kill", side_effect=mock_os_kill), - ): - await spawn_ttyd("test-session") - - force_killed = any( - pid == 54321 and sig == signal.SIGKILL for pid, sig in killed_pids - ) - assert force_killed, ( - "spawn_ttyd must SIGKILL any process occupying TTYD_PORT before spawning " - "to prevent 'address already in use' errors" - ) - - -# --------------------------------------------------------------------------- -# _pids_on_port fallback chain tests -# -# These cover the exact production failure: lsof not installed → all kill -# strategies silently returned False → stale ttyd survived restarts. -# --------------------------------------------------------------------------- - - -def _make_subprocess_result(returncode: int, stdout: str) -> MagicMock: - r = MagicMock() - r.returncode = returncode - r.stdout = stdout - return r - - -def test_pids_on_port_uses_lsof_when_available(): - """_pids_on_port() returns PID list from lsof output when lsof is available.""" - - def mock_run(cmd, **kwargs): # noqa: ANN001 - if cmd[0] == "lsof": - return _make_subprocess_result(0, "3555095\n") - return _make_subprocess_result(1, "") - - with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run): - result = _pids_on_port(7682) - - assert result == [3555095], "Should return PID from lsof" - - -def test_pids_on_port_falls_back_to_fuser_when_lsof_missing(): - """_pids_on_port() uses fuser when lsof is not installed (FileNotFoundError). - - This is the exact failure mode that caused the yazi-stuck-ttyd bug: - lsof was absent, _kill_pids_on_port silently returned False, and the - stale ttyd was never killed. - """ - - def mock_run(cmd, **kwargs): # noqa: ANN001 - if cmd[0] == "lsof": - raise FileNotFoundError("lsof not found") - if cmd[0] == "fuser": - return _make_subprocess_result(0, " 3555095") - return _make_subprocess_result(1, "") - - with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run): - result = _pids_on_port(7682) - - assert result == [3555095], "Should fall back to fuser when lsof is missing" - - -def test_pids_on_port_falls_back_to_ss_when_lsof_and_fuser_missing(): - """_pids_on_port() uses ss when both lsof and fuser are unavailable.""" - ss_output = ( - 'LISTEN 0 128 0.0.0.0:7682 0.0.0.0:* users:(("ttyd",pid=3555095,fd=12))\n' - ) - - def mock_run(cmd, **kwargs): # noqa: ANN001 - if cmd[0] in ("lsof", "fuser"): - raise FileNotFoundError(f"{cmd[0]} not found") - if cmd[0] == "ss": - return _make_subprocess_result(0, ss_output) - return _make_subprocess_result(1, "") - - with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run): - result = _pids_on_port(7682) - - assert result == [3555095], "Should fall back to ss when lsof and fuser missing" - - -def test_pids_on_port_returns_empty_when_all_tools_fail(): - """_pids_on_port() returns [] when no discovery tool is available.""" - - def mock_run(cmd, **kwargs): # noqa: ANN001 - raise FileNotFoundError(f"{cmd[0]} not found") - - with patch("muxplex.ttyd._subprocess.run", side_effect=mock_run): - result = _pids_on_port(7682) - - assert result == [], "Should return empty list when all tools fail" - - -def test_kill_pids_on_port_sends_signal_via_fuser_fallback(): - """_kill_pids_on_port() sends the correct signal using the fuser fallback. - - End-to-end regression guard: verifies that the full chain from - _kill_pids_on_port → _pids_on_port → fuser → os.kill works correctly - so that stale ttyd processes are killed even when lsof is absent. - """ - killed: list[tuple[int, int]] = [] - - def mock_run(cmd, **kwargs): # noqa: ANN001 - if cmd[0] == "lsof": - raise FileNotFoundError("lsof not found") - if cmd[0] == "fuser": - return _make_subprocess_result(0, " 3555095") - return _make_subprocess_result(1, "") - - def mock_kill(pid: int, sig: int) -> None: - killed.append((pid, sig)) - - with ( - patch("muxplex.ttyd._subprocess.run", side_effect=mock_run), - patch("os.kill", side_effect=mock_kill), - ): - result = _kill_pids_on_port(7682, signal.SIGKILL) - - assert result is True, "_kill_pids_on_port must return True when fuser finds a PID" - assert (3555095, signal.SIGKILL) in killed, ( - "Must send SIGKILL to PID 3555095 discovered via fuser fallback" - ) - - -async def test_kill_orphan_ttyd_handles_invalid_pid_file_content(): - """kill_orphan_ttyd() gracefully handles a PID file with non-integer content.""" - pid_path = ttyd_mod.TTYD_PID_PATH - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text("not-a-pid") - - # Should not raise and should clean up the file. - # kill_ttyd() calls pid_path.unlink(missing_ok=True) before returning False - # on invalid content, so the file is removed even though no kill occurred. - result = await kill_orphan_ttyd() - - assert result is False - assert not pid_path.exists(), "Invalid PID file should be removed" diff --git a/muxplex/tests/test_ws_proxy.py b/muxplex/tests/test_ws_proxy.py deleted file mode 100644 index 8aecb04..0000000 --- a/muxplex/tests/test_ws_proxy.py +++ /dev/null @@ -1,447 +0,0 @@ -""" -Comprehensive tests for the WebSocket proxy in muxplex/main.py. -""" - -import inspect -import threading -import time - -import pytest -from fastapi.testclient import TestClient -from starlette.websockets import WebSocketDisconnect - -from muxplex.auth import create_session_cookie -from muxplex.main import app, terminal_ws_proxy - - -# --------------------------------------------------------------------------- -# Polling helper — deterministic alternative to time.sleep() for async relay -# --------------------------------------------------------------------------- - - -def _wait_for(condition, timeout: float = 2.0, interval: float = 0.01) -> bool: - """Poll *condition()* until it returns True or *timeout* seconds elapses. - - Returns True if the condition was met, False on timeout. - Using a polling loop instead of a fixed sleep makes relay tests deterministic: - on fast machines the loop exits as soon as the relay completes; on slow machines - it waits up to *timeout* seconds rather than racing against a fixed 200ms budget. - """ - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if condition(): - return True - time.sleep(interval) - return False # pragma: no cover — timeout branch only on pathological machines - - -# --------------------------------------------------------------------------- -# autouse fixture — redirect state/PID files to tmp_path, mock startup side-effects -# --------------------------------------------------------------------------- - - -@pytest.fixture(autouse=True) -def patch_startup_and_state(tmp_path, monkeypatch): - """Redirect state files to tmp_path, mock muxterm start/stop, replace _poll_loop with no-op.""" - # Redirect state files - tmp_state_dir = tmp_path / "state" - tmp_state_path = tmp_state_dir / "state.json" - monkeypatch.setattr("muxplex.state.STATE_DIR", tmp_state_dir) - monkeypatch.setattr("muxplex.state.STATE_PATH", tmp_state_path) - - # Mock start_muxterm/stop_muxterm so startup doesn't touch real processes - from unittest.mock import AsyncMock - - monkeypatch.setattr("muxplex.main.start_muxterm", AsyncMock()) - monkeypatch.setattr("muxplex.main.stop_muxterm", AsyncMock()) - - # Replace _poll_loop with a no-op so tests don't spin up real poll cycles - async def noop_poll_loop() -> None: - pass - - monkeypatch.setattr("muxplex.main._poll_loop", noop_poll_loop) - - -# --------------------------------------------------------------------------- -# Helper — create TestClient with valid session cookie -# --------------------------------------------------------------------------- - - -def _make_authed_client(): - """Creates TestClient with valid session cookie.""" - from muxplex.main import _auth_secret, _auth_ttl - - cookie = create_session_cookie(_auth_secret, _auth_ttl) - client = TestClient(app) - client.cookies.set("muxplex_session", cookie) - return client - - -# --------------------------------------------------------------------------- -# FakeTtydWs — mock ttyd WebSocket for relay testing -# --------------------------------------------------------------------------- - - -class FakeTtydWs: - """Mock ttyd WebSocket that stores sent messages and yields pre-loaded responses. - - Supports send(), close(), async iterator, and async context manager. - """ - - def __init__(self, responses=None): - self.sent = [] - self._responses = list(responses or []) - self._closed = False - - async def send(self, message): - self.sent.append(message) - - async def close(self): - self._closed = True - - def __aiter__(self): - return self._async_gen() - - async def _async_gen(self): - for msg in self._responses: - yield msg - - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.close() - return False - - -# --------------------------------------------------------------------------- -# Test 1: regression — proxy source must use receive(), not receive_bytes() -# --------------------------------------------------------------------------- - - -# --------------------------------------------------------------------------- -# Tests: ttyd liveness check before websocket.accept -# --------------------------------------------------------------------------- - - -def test_ttyd_is_listening_function_exists(): - """_ttyd_is_listening() must exist in main.py (TCP probe helper).""" - # Import will fail if function doesn't exist — that IS the failing test - from muxplex.main import _ttyd_is_listening # noqa: F401 - - assert callable(_ttyd_is_listening) - - -def test_ws_proxy_checks_ttyd_before_accepting(): - """terminal_ws_proxy must check _ttyd_is_listening BEFORE websocket.accept. - - Root cause of the reconnect loop: the proxy called websocket.accept() before - checking if ttyd was alive. The browser's 'open' event fired immediately, - resetting _reconnectAttempts to 0. The counter bounced 0→1→0→1 forever so - the client-side /connect POST (at >= 2 attempts) never fired. - - Fix: check _ttyd_is_listening() first. If not listening, auto-spawn ttyd - THEN accept — so the browser only gets 'open' when ttyd is actually ready. - """ - source = inspect.getsource(terminal_ws_proxy) - # Use "await websocket.accept" to avoid matching the docstring mention - accept_idx = source.index("await websocket.accept") - ttyd_check_idx = source.index("_ttyd_is_listening") - assert ttyd_check_idx < accept_idx, ( - "_ttyd_is_listening() must be checked BEFORE await websocket.accept() — " - "proxy must not accept the browser WS until ttyd is confirmed alive" - ) - - -def test_ws_proxy_auto_spawns_ttyd_when_dead(monkeypatch): - """WS proxy must call get_or_spawn when pool has no entry for the session.""" - spawn_calls = [] - - async def mock_get_or_spawn(name: str): - spawn_calls.append(name) - return 7682 - - # pool_port returns None → no existing ttyd for this session - monkeypatch.setattr("muxplex.main.pool_port", lambda name: None) - # _ttyd_is_listening returns False → nothing listening on fallback port either - monkeypatch.setattr("muxplex.main._ttyd_is_listening", lambda port=7682: False) - # get_or_spawn is what the proxy should call to auto-spawn - monkeypatch.setattr("muxplex.main.get_or_spawn", mock_get_or_spawn) - - # Provide a fake websockets.connect that immediately closes (no real ttyd) - fake_ws = FakeTtydWs(responses=[]) - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - # Patch load_state to return state with active_session - monkeypatch.setattr( - "muxplex.main.load_state", - lambda: {"active_session": "test-session", "sessions": {}, "session_order": []}, - ) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as _: - pass - - assert spawn_calls == ["test-session"], ( - "get_or_spawn must be called with active_session when ttyd is not in pool" - ) - - -def test_terminal_ws_proxy_does_not_use_receive_bytes(): - """Regression: receive_bytes() silently drops TEXT frames (like the ttyd auth token). - - terminal.js sends {"AuthToken": ""} as a TEXT WebSocket frame. The original - proxy used receive_bytes() which fails on text frames, swallowed the exception, - and exited — meaning ttyd never received the auth token, never started - streaming, resulting in a permanent black screen and reconnect loop. - - The proxy MUST use receive() and dispatch on message type to handle both - binary and text frames correctly. - """ - source = inspect.getsource(terminal_ws_proxy) - assert "receive_bytes" not in source, ( - "client_to_ttyd must not use receive_bytes() — silently drops text frames " - 'like the ttyd auth token {"AuthToken": ""}' - ) - assert ".receive()" in source, ( - "client_to_ttyd must use receive() to handle both text and binary frames" - ) - - -# --------------------------------------------------------------------------- -# Tests 2–3: auth rejection -# --------------------------------------------------------------------------- - - -def test_ws_auth_rejection_no_cookie(): - """WebSocket from non-localhost without cookie is closed with code 4001.""" - # TestClient default host is "testclient" which is treated as non-localhost - with TestClient(app) as c: - with pytest.raises(WebSocketDisconnect) as exc_info: - with c.websocket_connect("/terminal/ws") as _: - pass - assert exc_info.value.code == 4001 - - -def test_ws_auth_rejection_invalid_cookie(): - """WebSocket from non-localhost with a tampered cookie is closed with code 4001.""" - with TestClient(app) as c: - c.cookies.set("muxplex_session", "tampered.invalid.cookie.value") - with pytest.raises(WebSocketDisconnect) as exc_info: - with c.websocket_connect("/terminal/ws") as _: - pass - assert exc_info.value.code == 4001 - - -# --------------------------------------------------------------------------- -# Test: Bearer token auth accepted -# --------------------------------------------------------------------------- - - -def test_ws_bearer_auth_accepted(monkeypatch): - """WebSocket from non-localhost with valid Bearer federation key is NOT rejected with 4001. - - When a valid federation key is provided as 'Authorization: Bearer ', - the WebSocket connection must be accepted (not closed with code 4001). - """ - fed_key = "test-federation-secret-key" - monkeypatch.setattr("muxplex.main._federation_key", fed_key) - - fake_ws = FakeTtydWs(responses=[]) - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - # Connect without a session cookie but with a valid Bearer token. - # Should NOT raise WebSocketDisconnect with code 4001. - with TestClient(app) as c: - # If Bearer auth is not implemented, this raises WebSocketDisconnect(code=4001) - with c.websocket_connect( - "/terminal/ws", - headers={"Authorization": f"Bearer {fed_key}"}, - ) as _ws: - pass # Successfully connected — auth was accepted - - -# --------------------------------------------------------------------------- -# Tests 4–5: browser → ttyd relay -# --------------------------------------------------------------------------- - - -def test_browser_text_relayed_to_ttyd(monkeypatch): - """Text message from browser is forwarded to ttyd via FakeTtydWs.send().""" - fake_ws = FakeTtydWs() - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as ws: - ws.send_text("hello from browser") - _wait_for(lambda: "hello from browser" in fake_ws.sent) - - assert "hello from browser" in fake_ws.sent - - -def test_browser_bytes_relayed_to_ttyd(monkeypatch): - """Binary message from browser is forwarded to ttyd via FakeTtydWs.send().""" - fake_ws = FakeTtydWs() - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as ws: - ws.send_bytes(b"\x00\x01\x02 binary data") - _wait_for(lambda: b"\x00\x01\x02 binary data" in fake_ws.sent) - - assert b"\x00\x01\x02 binary data" in fake_ws.sent - - -# --------------------------------------------------------------------------- -# Tests 6–7: ttyd → browser relay -# --------------------------------------------------------------------------- - - -def test_ttyd_text_relayed_to_browser(monkeypatch): - """Text message from ttyd is forwarded to browser via websocket.send_text().""" - fake_ws = FakeTtydWs(responses=["hello from ttyd"]) - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as ws: - msg = ws.receive_text() - assert msg == "hello from ttyd" - - -def test_ttyd_bytes_relayed_to_browser(monkeypatch): - """Binary message from ttyd is forwarded to browser via websocket.send_bytes().""" - fake_ws = FakeTtydWs(responses=[b"\xde\xad\xbe\xef binary"]) - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as ws: - msg = ws.receive_bytes() - assert msg == b"\xde\xad\xbe\xef binary" - - -# --------------------------------------------------------------------------- -# Test 8: ttyd close propagates to browser -# --------------------------------------------------------------------------- - - -def test_ttyd_close_propagates_to_browser(monkeypatch): - """When ttyd exhausts its messages, the proxy cleans up and closes the browser WS.""" - fake_ws = FakeTtydWs(responses=[]) # no responses — exhausts immediately - monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as _: - # FakeTtydWs has no responses so ttyd_to_client exhausts immediately. - # Exiting the context manager closes the browser WS, which causes - # client_to_ttyd to complete, gather finishes, and the proxy - # finally-block calls fake_ws.close(). - pass - - # fake_ws should have been closed when the async-with block exited - assert fake_ws._closed - - -# --------------------------------------------------------------------------- -# Test 9: ttyd unreachable closes browser WS -# --------------------------------------------------------------------------- - - -def test_ttyd_unreachable_closes_browser_ws(monkeypatch): - """OSError on ttyd connect closes the browser WebSocket (no hang, no 4001).""" - - def mock_connect_raises(*args, **kwargs): - raise OSError("Connection refused — ttyd not running") - - monkeypatch.setattr("muxplex.main.websockets.connect", mock_connect_raises) - - with _make_authed_client() as c: - with c.websocket_connect("/terminal/ws") as ws: - # Proxy accepts, then closes after failing to reach ttyd. - # Receive the close frame — proves the proxy closed (no hang) - # and that auth was not rejected (which would use code 4001). - close_frame = ws.receive() - assert close_frame.get("type") == "websocket.close", ( - "Proxy must close the WebSocket" - ) - assert close_frame.get("code") != 4001, "Must not be an auth rejection (4001)" - - -# --------------------------------------------------------------------------- -# Test 10: concurrent sessions don't interfere -# --------------------------------------------------------------------------- - - -def test_concurrent_ws_sessions(monkeypatch): - """Two simultaneous proxy sessions relay to separate FakeTtydWs instances.""" - # Create two separate FakeTtydWs instances, one per connection - ws_pool = [FakeTtydWs(), FakeTtydWs()] - call_count = 0 - lock = threading.Lock() - - def mock_connect(*args, **kwargs): - nonlocal call_count - with lock: - idx = call_count % len(ws_pool) - call_count += 1 - return ws_pool[idx] - - monkeypatch.setattr("muxplex.main.websockets.connect", mock_connect) - - errors = [] - - with _make_authed_client() as c: - - def send_msg(text): - try: - with c.websocket_connect("/terminal/ws") as ws: - ws.send_text(text) - _wait_for(lambda: text in ws_pool[0].sent + ws_pool[1].sent) - except Exception as exc: - errors.append(exc) - - t1 = threading.Thread(target=send_msg, args=("session_one_msg",)) - t2 = threading.Thread(target=send_msg, args=("session_two_msg",)) - t1.start() - t2.start() - t1.join(timeout=10) - t2.join(timeout=10) - - assert not errors, f"Concurrent sessions raised errors: {errors}" - - # Both messages must have been relayed (one to each fake_ws) - all_sent = ws_pool[0].sent + ws_pool[1].sent - assert "session_one_msg" in all_sent - assert "session_two_msg" in all_sent - - -# --------------------------------------------------------------------------- -# Task-11: federation WebSocket proxy route -# --------------------------------------------------------------------------- - - -def test_federation_ws_proxy_route_exists(): - """App must have a WebSocket route at /federation/{device_id}/terminal/ws.""" - from starlette.routing import WebSocketRoute - - ws_routes = [r for r in app.routes if isinstance(r, WebSocketRoute)] - paths = [r.path for r in ws_routes] - assert "/federation/{device_id}/terminal/ws" in paths, ( - "App must have a WebSocket route at /federation/{device_id}/terminal/ws" - ) - - -def test_federation_ws_proxy_uses_ssl_context_for_wss(): - """Federation WS proxy must pass an SSL context when connecting via wss://. - - Self-signed certs on remote instances (cortex, spark-2, etc.) fail the - default SSL verification in websockets.connect(). The proxy must build an - SSLContext with CERT_NONE for wss:// URLs — the same fix already applied - to the httpx client (verify=False) but for the websockets library. - """ - from muxplex.main import federation_terminal_ws_proxy - - source = inspect.getsource(federation_terminal_ws_proxy) - assert "ssl" in source and ("CERT_NONE" in source or "ssl_context" in source), ( - "Federation WS proxy must configure an SSL context (CERT_NONE / ssl_context) " - "for self-signed cert support on wss:// connections" - ) diff --git a/muxplex/ttyd.py b/muxplex/ttyd.py deleted file mode 100644 index 18be467..0000000 --- a/muxplex/ttyd.py +++ /dev/null @@ -1,471 +0,0 @@ -""" -ttyd process lifecycle management for the tmux-web muxplex. - -Constants: - TTYD_PID_DIR — directory for PID files (default: ~/.local/share/tmux-web/) - TTYD_PID_PATH — full path to the legacy PID file (TTYD_PID_DIR / 'ttyd.pid') - TTYD_PORT — base port for the ttyd pool (7682) - -Module state: - _active_process — the currently running ttyd subprocess (or None) [backward compat] - _pool — session→PoolEntry mapping for the process pool - -Pool API: - get_or_spawn(session_name) — Return port; spawn if needed, poll until listening - pool_port(session_name) — Look up port for session without spawning - kill_session(session_name) — Kill ttyd for one session, remove from pool - kill_all() — Kill all pooled ttyd processes - kill_orphans() — Startup cleanup: PID files, port range scan - -Backward-compat API: - spawn_ttyd(session_name) — spawn ttyd attached to a tmux session, write PID file - kill_ttyd() — kill the running ttyd process, clean up PID file - kill_orphan_ttyd() — kill any orphaned ttyd from a previous coordinator run -""" - -import asyncio -import os -import re as _re -import signal -import socket as _socket -import subprocess as _subprocess -import time -from dataclasses import dataclass -from pathlib import Path - -# --------------------------------------------------------------------------- -# Paths and constants -# --------------------------------------------------------------------------- - -_default_ttyd_pid_dir = Path.home() / ".local" / "share" / "tmux-web" -TTYD_PID_DIR: Path = Path(os.environ.get("TMUX_WEB_STATE_DIR", _default_ttyd_pid_dir)) -TTYD_PID_PATH: Path = TTYD_PID_DIR / "ttyd.pid" - -TTYD_PORT: int = 7682 - -_PORT_MIN: int = 7682 -_PORT_MAX: int = 7701 - -# --------------------------------------------------------------------------- -# Module state -# --------------------------------------------------------------------------- - -_active_process: asyncio.subprocess.Process | None = None - - -@dataclass -class PoolEntry: - """A pooled ttyd process and its assigned port.""" - - process: asyncio.subprocess.Process - port: int - - -_pool: dict[str, PoolEntry] = {} - -# --------------------------------------------------------------------------- -# Internal helpers -# --------------------------------------------------------------------------- - - -def _pid_path_for_port(port: int) -> Path: - """Return the per-port PID file path: ``ttyd-{port}.pid`` in TTYD_PID_DIR.""" - return TTYD_PID_DIR / f"ttyd-{port}.pid" - - -def _is_alive(entry: PoolEntry) -> bool: - """Return True if the pool entry's process has not exited.""" - return entry.process.returncode is None - - -def _ttyd_is_listening(port: int = TTYD_PORT) -> bool: - """Return True if something is accepting TCP connections on *port*. - - Uses a raw socket connect (no WebSocket handshake, no PTY spawned). - Takes < 1 ms on localhost when ttyd is running; fails immediately with - ConnectionRefusedError when it's not. OSError/TimeoutError are also - caught so the caller always gets a bool. - """ - try: - with _socket.create_connection(("127.0.0.1", port), timeout=0.5): - return True - except (ConnectionRefusedError, OSError, TimeoutError): - return False - - -def _allocate_port() -> int: - """Return the next free port in the pool range [_PORT_MIN, _PORT_MAX].""" - used_ports = {e.port for e in _pool.values()} - for port in range(_PORT_MIN, _PORT_MAX + 1): - if port not in used_ports: - return port - raise RuntimeError("ttyd pool exhausted: all ports in range are in use") - - -def _pids_on_port(port: int) -> list[int]: - """Return PIDs of all processes listening on *port*. - - Tries three tools in order, stopping at the first that returns results: - - 1. ``lsof -ti :`` — one PID per line, most widely available. - 2. ``fuser /tcp`` — space-separated PIDs on stdout (psmisc). - 3. ``ss -Hnltp`` — parses ``pid=N`` from users field (iproute2). - - Returns an empty list if none of the tools are available or find anything. - Silently swallows all errors so callers always get a list. - """ - pids: list[int] = [] - - # --- Tool 1: lsof --- - try: - result = _subprocess.run( - ["lsof", "-ti", f":{port}"], - capture_output=True, - text=True, - timeout=5, - ) - if result.returncode == 0 and result.stdout.strip(): - for tok in result.stdout.split(): - try: - pids.append(int(tok)) - except ValueError: - pass - except Exception: # noqa: BLE001 - pass - - if pids: - return pids - - # --- Tool 2: fuser (psmisc) --- - try: - result = _subprocess.run( - ["fuser", f"{port}/tcp"], - capture_output=True, - text=True, - timeout=5, - ) - # fuser writes the port label to stderr and PIDs to stdout. - if result.stdout.strip(): - for tok in result.stdout.split(): - try: - pids.append(int(tok)) - except ValueError: - pass - except Exception: # noqa: BLE001 - pass - - if pids: - return pids - - # --- Tool 3: ss (iproute2) --- - try: - result = _subprocess.run( - ["ss", "-Hnltp", f"sport = :{port}"], - capture_output=True, - text=True, - timeout=5, - ) - if result.returncode == 0 and result.stdout.strip(): - for match in _re.finditer(r"pid=(\d+)", result.stdout): - try: - pids.append(int(match.group(1))) - except ValueError: - pass - except Exception: # noqa: BLE001 - pass - - return pids - - -def _kill_pids_on_port(port: int, sig: int) -> bool: - """Find and signal all processes listening on *port*. - - Uses :func:`_pids_on_port` (lsof → fuser → ss) to locate PIDs, then - signals each with *sig*. - - Returns True if at least one PID was found and signalled. - Silently ignores unavailable tools and already-dead processes. - """ - pids = _pids_on_port(port) - if not pids: - return False - sent = False - for pid in pids: - try: - os.kill(pid, sig) - sent = True - except (ProcessLookupError, PermissionError): - pass - return sent - - -# --------------------------------------------------------------------------- -# Pool API -# --------------------------------------------------------------------------- - - -async def get_or_spawn(session_name: str) -> int: - """Return the port for *session_name*, spawning ttyd if needed. - - If the session already has a live ttyd in the pool (process alive AND - port is listening), returns immediately. Otherwise allocates the next - free port from [_PORT_MIN, _PORT_MAX], spawns ttyd, polls until - listening (up to ~1 s), and returns the port. - """ - # Check existing pool entry — require both alive process AND listening port - entry = _pool.get(session_name) - if entry is not None: - if _is_alive(entry) and _ttyd_is_listening(entry.port): - return entry.port - # Stale entry — clean up - try: - entry.process.terminate() - except ProcessLookupError: - pass - _pid_path_for_port(entry.port).unlink(missing_ok=True) - del _pool[session_name] - - # Allocate next free port - port = _allocate_port() - - # Force-free the port (catches races) - if _kill_pids_on_port(port, signal.SIGKILL): - await asyncio.sleep(0.3) - - # Spawn ttyd - proc = await asyncio.create_subprocess_exec( - "ttyd", - "-W", - "-m", - "3", - "-p", - str(port), - "tmux", - "attach", - "-t", - session_name, - stdout=asyncio.subprocess.DEVNULL, - stderr=asyncio.subprocess.DEVNULL, - start_new_session=True, - ) - - # Write per-port PID file - TTYD_PID_DIR.mkdir(parents=True, exist_ok=True) - _pid_path_for_port(port).write_text(str(proc.pid)) - - # Add to pool - _pool[session_name] = PoolEntry(process=proc, port=port) - - # Poll until listening (up to ~1 s) - for _ in range(20): - if _ttyd_is_listening(port): - break - await asyncio.sleep(0.05) - - return port - - -def pool_port(session_name: str) -> int | None: - """Look up the port for *session_name* without spawning. - - Returns the port if the session has a live pool entry, else None. - """ - entry = _pool.get(session_name) - if entry is not None and _is_alive(entry): - return entry.port - return None - - -async def kill_session(session_name: str) -> bool: - """Kill the ttyd process for *session_name* and remove from pool. - - Returns True if a pool entry was found and cleaned up, False otherwise. - """ - entry = _pool.pop(session_name, None) - if entry is None: - return False - try: - entry.process.terminate() - try: - await asyncio.wait_for(entry.process.wait(), timeout=2.0) - except asyncio.TimeoutError: - entry.process.kill() - except ProcessLookupError: - pass - _pid_path_for_port(entry.port).unlink(missing_ok=True) - return True - - -async def kill_all() -> None: - """Kill all pooled ttyd processes and clear the pool.""" - for name in list(_pool): - await kill_session(name) - - -async def kill_orphans() -> None: - """Startup cleanup: read PID files, SIGTERM, scan port range for stragglers. - - Handles both legacy ``ttyd.pid`` and per-port ``ttyd-{port}.pid`` files. - """ - await kill_orphan_ttyd() - - -# --------------------------------------------------------------------------- -# Backward-compat public API -# --------------------------------------------------------------------------- - - -async def kill_ttyd() -> bool: - """Kill the running ttyd process and clean up the PID file. - - Belt-and-suspenders strategy: - - Strategy 1 — PID file: - Reads the PID from TTYD_PID_PATH. If no PID file exists, returns False. - If the file content is not a valid integer, removes the file and returns - False. Checks whether the process is alive via ``os.kill(pid, 0)``. If - already gone (ProcessLookupError), cleans up and proceeds. Otherwise - sends SIGTERM and polls every 0.1 s for up to 2 s. - - Strategy 2 — port-based fallback: - After the PID-file kill, finds and kills any process still listening on - TTYD_PORT via ``_pids_on_port()`` (lsof → fuser → ss). This catches - orphaned ttyd processes whose PID was never recorded in the file (e.g. - after a coordinator crash). A brief 0.3 s wait is added to let the OS - release the port. - - The PID file and ``_active_process`` are cleared in all cases before - returning. - - Returns: - True — a process was killed (or was already dead) via either strategy. - False — no PID file found and no process was listening on the port. - """ - global _active_process - - killed = False - - # ------------------------------------------------------------------- - # Strategy 1: PID file - # ------------------------------------------------------------------- - if TTYD_PID_PATH.exists(): - try: - pid = int(TTYD_PID_PATH.read_text().strip()) - except ValueError: - TTYD_PID_PATH.unlink(missing_ok=True) - pid = None - else: - # Check whether the process is still alive. - try: - os.kill(pid, 0) - except ProcessLookupError: - # Already dead — clean up and note success. - TTYD_PID_PATH.unlink(missing_ok=True) - killed = True - pid = None - else: - # Process is alive — ask it to terminate. - os.kill(pid, signal.SIGTERM) - - # Poll up to 2 s for the process to exit. - deadline = time.time() + 2.0 - while time.time() < deadline: - try: - os.kill(pid, 0) - except (ProcessLookupError, PermissionError): - break - await asyncio.sleep(0.1) - - TTYD_PID_PATH.unlink(missing_ok=True) - killed = True - pid = None # noqa: F841 (intentional) - - # ------------------------------------------------------------------- - # Strategy 2: port-based fallback — catch orphans not in PID file - # ------------------------------------------------------------------- - if _kill_pids_on_port(TTYD_PORT, signal.SIGTERM): - killed = True - # Brief pause so the OS can release the port before the next spawn. - await asyncio.sleep(0.3) - - _active_process = None - return killed - - -async def kill_orphan_ttyd() -> bool: - """Kill any orphaned ttyd process left over from a previous coordinator run. - - Handles both legacy ``ttyd.pid`` and pool-era ``ttyd-{port}.pid`` files, - plus scans the full port range for stragglers. - - Returns: - True — an orphan was found (process was dead or alive). - False — no PID file found, or PID file contained invalid content. - """ - # Pool-era cleanup: per-port PID files - if TTYD_PID_DIR.exists(): - for pid_file in TTYD_PID_DIR.glob("ttyd-*.pid"): - try: - pid = int(pid_file.read_text().strip()) - os.kill(pid, signal.SIGTERM) - except (ValueError, ProcessLookupError, PermissionError): - pass - pid_file.unlink(missing_ok=True) - - # Scan additional ports beyond TTYD_PORT - for port in range(_PORT_MIN + 1, _PORT_MAX + 1): - _kill_pids_on_port(port, signal.SIGTERM) - - # Original behavior: legacy PID file + TTYD_PORT cleanup - return await kill_ttyd() - - -async def spawn_ttyd(session_name: str) -> asyncio.subprocess.Process: - """Spawn a ttyd process attached to *session_name* via ``tmux attach``. - - Runs:: - - ttyd -W -m 3 -p 7682 tmux attach -t - - Before spawning, verifies that TTYD_PORT is free. If any process is still - listening on the port (e.g. a race between kill_ttyd() and spawn_ttyd()), - it sends SIGKILL to force-free the port immediately. - - stdout and stderr are discarded (DEVNULL). The PID is written to - TTYD_PID_PATH. The process handle is stored in ``_active_process``. - - Args: - session_name: The tmux session name to attach to. - - Returns: - The asyncio.subprocess.Process object for the spawned ttyd. - """ - global _active_process - - # Final port-free guard — catches races where kill_ttyd() returned but - # the old ttyd hasn't fully released the socket yet. - if _kill_pids_on_port(TTYD_PORT, signal.SIGKILL): - await asyncio.sleep(0.3) - - proc = await asyncio.create_subprocess_exec( - "ttyd", - "-W", - "-m", - "3", - "-p", - str(TTYD_PORT), - "tmux", - "attach", - "-t", - session_name, - stdout=asyncio.subprocess.DEVNULL, - stderr=asyncio.subprocess.DEVNULL, - start_new_session=True, # detach from parent process group so ttyd survives independently - ) - - # Write PID file (create parent dirs if needed) - TTYD_PID_DIR.mkdir(parents=True, exist_ok=True) - TTYD_PID_PATH.write_text(str(proc.pid)) - - _active_process = proc - return proc