Gan, Jimmy
b742ab8fef
Deploy Dashboard (Dev) / deploy-dev (push) Successful in 2m44s
- Wrap tryRefreshSession in try-finally to ensure authRecoveryInFlight is always reset - Fix variable scope issue: use const tabState instead of reassigning d - Add error handling for process.stdin.write() to catch BrokenPipeError/OSError - Verify WebSocket instance before closing in ping timeout callback - Properly await conn.wait_closed() to prevent SSH connection leaks - Gracefully cancel and wait for read_task with 2s timeout These fixes prevent race conditions, resource leaks, and improve error recovery.
221 lines
8.0 KiB
Python
221 lines
8.0 KiB
Python
import asyncio
|
|
import struct
|
|
import logging
|
|
import shlex
|
|
import jwt
|
|
from collections import Counter
|
|
from fastapi import WebSocket, Query
|
|
import asyncssh
|
|
import config
|
|
import auth
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MAX_SESSIONS = 5
|
|
MAX_SESSIONS_PER_USER = 5
|
|
_session_lock = asyncio.Lock()
|
|
_active_sessions: dict[int, str] = {}
|
|
PERSISTENCE_STATUS_PREFIX = "__DASHBOARD_PERSISTENCE__:"
|
|
|
|
HOSTS = {
|
|
"nas": {"host": config.SSH_HOST, "user": config.SSH_USER, "key": config.SSH_KEY_PATH},
|
|
"vps": {"host": config.VPS_SSH_HOST, "user": config.VPS_SSH_USER, "key": config.VPS_SSH_KEY_PATH},
|
|
"caddy-vps": {"host": config.CADDY_VPS_SSH_HOST, "user": config.CADDY_VPS_SSH_USER, "key": config.VPS_SSH_KEY_PATH},
|
|
"mac": {"host": config.MAC_SSH_HOST, "user": config.MAC_SSH_USER, "key": config.MAC_SSH_KEY_PATH},
|
|
"claude-dev": {
|
|
"host": config.SSH_HOST,
|
|
"user": config.SSH_USER,
|
|
"key": config.SSH_KEY_PATH,
|
|
"persistent_session": {
|
|
"container": "claude-dev",
|
|
"session_name": "dashboard-terminal",
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
# Load known_hosts once at import time — fail closed if missing
|
|
_known_hosts_available = False
|
|
try:
|
|
_known_hosts = asyncssh.read_known_hosts(config.SSH_KNOWN_HOSTS)
|
|
_known_hosts_available = True
|
|
except Exception:
|
|
_known_hosts = None
|
|
logger.error("SSH known_hosts file not found at %s — terminal connections will be refused", config.SSH_KNOWN_HOSTS)
|
|
|
|
|
|
def build_host_command(host_config: dict) -> str | None:
|
|
persistent = host_config.get("persistent_session")
|
|
if not persistent:
|
|
return host_config.get("command")
|
|
|
|
docker_bin = "/volume1/@appstore/ContainerManager/usr/bin/docker"
|
|
container = persistent["container"]
|
|
session_name = persistent["session_name"]
|
|
quoted_container = shlex.quote(container)
|
|
quoted_session = shlex.quote(session_name)
|
|
quoted_marker = shlex.quote(PERSISTENCE_STATUS_PREFIX)
|
|
attach_cmd = (
|
|
f"{docker_bin} exec -it {quoted_container} "
|
|
f"tmux new-session -A -s {quoted_session}"
|
|
)
|
|
|
|
script = " ".join([
|
|
"if",
|
|
f"{docker_bin} exec {quoted_container} tmux has-session -t {quoted_session} >/dev/null 2>&1;",
|
|
"then",
|
|
f"printf '%sattached\\n' {quoted_marker};",
|
|
"else",
|
|
f"printf '%screated\\n' {quoted_marker};",
|
|
"fi;",
|
|
f"exec {attach_cmd}",
|
|
])
|
|
return f"bash -lc {shlex.quote(script)}"
|
|
|
|
|
|
async def _reserve_session(session_id: int, username: str) -> tuple[bool, str | None, int | None]:
|
|
async with _session_lock:
|
|
if len(_active_sessions) >= MAX_SESSIONS:
|
|
return False, "Too many sessions", 1013
|
|
user_sessions = Counter(_active_sessions.values())
|
|
if user_sessions[username] >= MAX_SESSIONS_PER_USER:
|
|
return False, "Too many sessions for user", 1013
|
|
_active_sessions[session_id] = username
|
|
return True, None, session_id
|
|
|
|
|
|
async def _release_session(session_id: int | None):
|
|
if session_id is None:
|
|
return
|
|
async with _session_lock:
|
|
_active_sessions.pop(session_id, None)
|
|
|
|
|
|
async def ws_endpoint(websocket: WebSocket, host: str = Query("nas")):
|
|
if host not in HOSTS:
|
|
await websocket.close(code=1008, reason="Unknown host")
|
|
return
|
|
|
|
access_token = websocket.cookies.get(auth.ACCESS_COOKIE_NAME)
|
|
if not access_token:
|
|
logger.warning("WebSocket auth failed for %s: no access token cookie", host)
|
|
await websocket.close(code=1008, reason="Unauthorized")
|
|
return
|
|
|
|
try:
|
|
user = await auth.get_current_user_ws(access_token)
|
|
except jwt.ExpiredSignatureError as e:
|
|
logger.warning("WebSocket auth failed for %s: token expired - %s", host, e)
|
|
await websocket.close(code=1008, reason="Token expired")
|
|
return
|
|
except jwt.InvalidTokenError as e:
|
|
logger.warning("WebSocket auth failed for %s: invalid token - %s", host, e)
|
|
await websocket.close(code=1008, reason="Invalid token")
|
|
return
|
|
except Exception as e:
|
|
logger.error("WebSocket auth failed for %s: unexpected error - %s", host, e)
|
|
await websocket.close(code=1008, reason="Unauthorized")
|
|
return
|
|
|
|
if user.pages != "*" and "terminal" not in user.pages:
|
|
await websocket.close(code=1008, reason="Access denied")
|
|
return
|
|
|
|
if not _known_hosts_available:
|
|
await websocket.close(code=1011, reason="SSH host verification unavailable")
|
|
return
|
|
|
|
session_id = id(websocket)
|
|
reserved, reason, session_or_code = await _reserve_session(session_id, user.username)
|
|
if not reserved:
|
|
await websocket.close(code=session_or_code, reason=reason)
|
|
return
|
|
session_id = session_or_code
|
|
|
|
await websocket.accept()
|
|
logger.info("WebSocket accepted for %s (user: %s)", host, user.username)
|
|
|
|
h = HOSTS[host]
|
|
try:
|
|
conn = await asyncio.wait_for(
|
|
asyncssh.connect(
|
|
h["host"], username=h["user"],
|
|
client_keys=[h["key"]], known_hosts=_known_hosts,
|
|
keepalive_interval=15,
|
|
keepalive_count_max=8,
|
|
),
|
|
timeout=10.0
|
|
)
|
|
logger.info("SSH connection established to %s", host)
|
|
except asyncio.TimeoutError:
|
|
logger.error("SSH connection to %s timed out after 10s", host)
|
|
await _release_session(session_id)
|
|
await websocket.close(code=1011, reason="SSH connection timeout")
|
|
return
|
|
except Exception as e:
|
|
logger.error("SSH connection to %s failed: %s", host, e)
|
|
await _release_session(session_id)
|
|
await websocket.close(code=1011, reason="SSH connection failed")
|
|
return
|
|
|
|
try:
|
|
process = await conn.create_process(
|
|
build_host_command(h), term_type="xterm-256color", term_size=(80, 24),
|
|
encoding=None,
|
|
)
|
|
|
|
async def read_ssh():
|
|
try:
|
|
while True:
|
|
data = await process.stdout.read(4096)
|
|
if not data:
|
|
break
|
|
await websocket.send_bytes(data)
|
|
except asyncssh.Error as e:
|
|
logger.warning("SSH read error for %s: %s", host, e)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error("Unexpected error reading SSH for %s: %s", host, e)
|
|
|
|
read_task = asyncio.create_task(read_ssh())
|
|
|
|
try:
|
|
while True:
|
|
msg = await websocket.receive()
|
|
if "bytes" in msg and msg["bytes"]:
|
|
data = msg["bytes"]
|
|
if data[0:1] == b"\x01" and len(data) == 5:
|
|
cols = struct.unpack("!H", data[1:3])[0]
|
|
rows = struct.unpack("!H", data[3:5])[0]
|
|
process.channel.change_terminal_size(cols, rows)
|
|
else:
|
|
try:
|
|
process.stdin.write(data)
|
|
except (BrokenPipeError, OSError) as e:
|
|
logger.warning("Failed to write to stdin for %s: %s", host, e)
|
|
break
|
|
elif "text" in msg and msg["text"]:
|
|
if msg["text"] == "__ping__":
|
|
logger.debug("Received ping from %s, sending pong", host)
|
|
await websocket.send_text("__pong__")
|
|
continue
|
|
try:
|
|
process.stdin.write(msg["text"].encode())
|
|
except (BrokenPipeError, OSError) as e:
|
|
logger.warning("Failed to write to stdin for %s: %s", host, e)
|
|
break
|
|
except Exception as e:
|
|
logger.warning("WebSocket receive loop ended for %s: %s", host, e)
|
|
finally:
|
|
read_task.cancel()
|
|
try:
|
|
await asyncio.wait_for(read_task, timeout=2.0)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
process.close()
|
|
finally:
|
|
conn.close()
|
|
await conn.wait_closed()
|
|
await _release_session(session_id)
|