refactor: simplify path validation and connection_init via shared helpers

- Extract _validate_path() in unraid.py — consolidates traversal check + normpath
  + prefix validation used by disk/logs and live/log_tail into one place
- Extract build_connection_init() in subscriptions/utils.py — removes 4 duplicate
  connection_init payload blocks from snapshot.py (×2), manager.py, diagnostics.py;
  also fixes diagnostics.py bug where x-api-key: None was sent when no key configured
- Remove _LIVE_ALLOWED_LOG_PREFIXES alias — direct reference to _ALLOWED_LOG_PREFIXES
- Move import hmac to module level in server.py (was inside verify_token hot path)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jacob Magar
2026-03-23 11:57:00 -04:00
parent dc1e5f18d8
commit e68d4a80e4
7 changed files with 51 additions and 49 deletions

View File

@@ -21,7 +21,12 @@ from ..core.exceptions import ToolError
from ..core.utils import safe_display_url
from .manager import subscription_manager
from .resources import ensure_subscriptions_started
from .utils import _analyze_subscription_status, build_ws_ssl_context, build_ws_url
from .utils import (
_analyze_subscription_status,
build_connection_init,
build_ws_ssl_context,
build_ws_url,
)
# Schema field names that appear inside the selection set of allowed subscriptions.
@@ -125,15 +130,8 @@ def register_diagnostic_tools(mcp: FastMCP) -> None:
ping_interval=30,
ping_timeout=10,
) as websocket:
# Send connection init (using standard X-API-Key format)
await websocket.send(
json.dumps(
{
"type": "connection_init",
"payload": {"x-api-key": _settings.UNRAID_API_KEY},
}
)
)
# Send connection init
await websocket.send(json.dumps(build_connection_init()))
# Wait for ack
response = await websocket.recv()

View File

@@ -19,7 +19,7 @@ from ..config import settings as _settings
from ..config.logging import logger
from ..core.client import redact_sensitive
from ..core.types import SubscriptionData
from .utils import build_ws_ssl_context, build_ws_url
from .utils import build_connection_init, build_ws_ssl_context, build_ws_url
# Resource data size limits to prevent unbounded memory growth
@@ -284,13 +284,9 @@ class SubscriptionManager:
logger.debug(
f"[PROTOCOL:{subscription_name}] Initializing GraphQL-WS protocol..."
)
init_type = "connection_init"
init_payload: dict[str, Any] = {"type": init_type}
if _settings.UNRAID_API_KEY:
init_payload = build_connection_init()
if "payload" in init_payload:
logger.debug(f"[AUTH:{subscription_name}] Adding authentication payload")
# Use graphql-ws connectionParams format (direct key, not nested headers)
init_payload["payload"] = {"x-api-key": _settings.UNRAID_API_KEY}
else:
logger.warning(
f"[AUTH:{subscription_name}] No API key available for authentication"

View File

@@ -18,10 +18,9 @@ from typing import Any
import websockets
from websockets.typing import Subprotocol
from ..config import settings as _settings
from ..config.logging import logger
from ..core.exceptions import ToolError
from .utils import build_ws_ssl_context, build_ws_url
from .utils import build_connection_init, build_ws_ssl_context, build_ws_url
async def subscribe_once(
@@ -48,10 +47,7 @@ async def subscribe_once(
sub_id = "snapshot-1"
# Handshake
init: dict[str, Any] = {"type": "connection_init"}
if _settings.UNRAID_API_KEY:
init["payload"] = {"x-api-key": _settings.UNRAID_API_KEY}
await ws.send(json.dumps(init))
await ws.send(json.dumps(build_connection_init()))
raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
ack = json.loads(raw)
@@ -123,10 +119,7 @@ async def subscribe_collect(
proto = ws.subprotocol or "graphql-transport-ws"
sub_id = "snapshot-1"
init: dict[str, Any] = {"type": "connection_init"}
if _settings.UNRAID_API_KEY:
init["payload"] = {"x-api-key": _settings.UNRAID_API_KEY}
await ws.send(json.dumps(init))
await ws.send(json.dumps(build_connection_init()))
raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
ack = json.loads(raw)

View File

@@ -45,7 +45,7 @@ def build_ws_ssl_context(ws_url: str) -> _ssl.SSLContext | None:
ws_url: The WebSocket URL to connect to.
Returns:
An SSLContext configured per _settings.UNRAID_VERIFY_SSL, or None for non-TLS URLs.
An SSLContext configured per UNRAID_VERIFY_SSL, or None for non-TLS URLs.
"""
if not ws_url.startswith("wss://"):
return None
@@ -60,6 +60,18 @@ def build_ws_ssl_context(ws_url: str) -> _ssl.SSLContext | None:
return ctx
def build_connection_init() -> dict[str, Any]:
"""Build the graphql-ws connection_init message.
Omits the payload key entirely when no API key is configured —
sending {"x-api-key": None} and omitting the key differ for some servers.
"""
msg: dict[str, Any] = {"type": "connection_init"}
if _settings.UNRAID_API_KEY:
msg["payload"] = {"x-api-key": _settings.UNRAID_API_KEY}
return msg
def _analyze_subscription_status(
status: dict[str, Any],
) -> tuple[int, list[dict[str, Any]]]: