molecule-core/workspace-template/events.py
Dev Lead Agent bea0e96a86 fix(security): Cycle 5 — auth middleware, injection hardening, skill sandbox
Fix A — platform/internal/middleware/wsauth_middleware.go (NEW):
  WorkspaceAuth() gin middleware enforces per-workspace bearer-token auth on
  ALL /workspaces/:id/* sub-routes. Same lazy-bootstrap contract as
  secrets.Values: workspaces with no live token are grandfathered through.
  Blocks C2, C3, C4, C5, C7, C8, C9, C12, C13 simultaneously.

Fix A — platform/internal/router/router.go:
  Reorganised route registration: bare CRUD (/workspaces, /workspaces/:id)
  and /a2a remain on root router; all other /workspaces/:id/* sub-routes
  moved into wsAuth = r.Group("/workspaces/:id", middleware.WorkspaceAuth(db.DB)).
  CORS AllowHeaders updated to include Authorization so browser/agent callers
  can send the bearer token cross-origin.

Fix B — workspace-template/heartbeat.py:
  _check_delegations(): validate source_id == self.workspace_id before
  accepting a delegation result. Attacker-crafted records with a foreign
  source_id are silently skipped with a WARNING log (injection attempt).
  trigger_msg no longer embeds raw response_preview text; references
  delegation_id + status only — removes the prompt-injection vector.

Fix C — workspace-template/skill_loader/loader.py:
  load_skill_tools(): before exec_module(), verify script is within
  scripts_dir (path traversal guard) and temporarily scrub sensitive env
  vars (CLAUDE_CODE_OAUTH_TOKEN, ANTHROPIC_API_KEY, OPENAI_API_KEY,
  WORKSPACE_AUTH_TOKEN, GITHUB_TOKEN, GH_TOKEN) from os.environ; restore
  in finally block. Defence-in-depth even if /plugins auth gate is bypassed.

Fix D — platform/internal/handlers/socket.go:
  HandleConnect(): agent connections (X-Workspace-ID present) validated via
  wsauth.HasAnyLiveToken + wsauth.ValidateToken before WebSocket upgrade.
  Canvas clients (no X-Workspace-ID) remain unauthenticated.

Fix D — workspace-template/events.py:
  PlatformEventSubscriber._connect(): include platform_auth bearer token in
  WebSocket upgrade headers alongside X-Workspace-ID.

Fix E — workspace-template/executor_helpers.py:
  recall_memories() and commit_memory() now pass platform_auth bearer token
  in Authorization header so WorkspaceAuth middleware allows access.

Fix F — workspace-template/a2a_client.py:
  send_a2a_message(): timeout=None → httpx.Timeout(connect=30, read=300,
  write=30, pool=30). Resolves H2 flagged across 5 consecutive audits.

Tests: 149/149 Python tests pass (test_heartbeat + test_events updated to
assert new source_id validation behaviour and allow Authorization header).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-14 04:44:42 +00:00

97 lines
3.3 KiB
Python

"""WebSocket subscriber for platform events.
Subscribes to the platform WebSocket with X-Workspace-ID header
so the workspace only receives events about reachable peers.
Triggers system prompt rebuild on relevant peer changes.
"""
import asyncio
import json
import logging
import httpx
logger = logging.getLogger(__name__)
# Events that should trigger a system prompt rebuild
REBUILD_EVENTS = {
"WORKSPACE_ONLINE",
"WORKSPACE_OFFLINE",
"WORKSPACE_EXPANDED",
"WORKSPACE_COLLAPSED",
"WORKSPACE_REMOVED",
"AGENT_CARD_UPDATED",
}
class PlatformEventSubscriber:
"""Subscribes to platform WebSocket for peer events."""
def __init__(
self,
platform_url: str,
workspace_id: str,
on_peer_change=None,
):
self.ws_url = platform_url.replace("http://", "ws://").replace("https://", "wss://") + "/ws"
self.workspace_id = workspace_id
self.on_peer_change = on_peer_change
self._running = False
self._reconnect_delay = 1.0
async def start(self):
"""Connect to platform WebSocket with exponential backoff reconnect."""
self._running = True
while self._running:
try:
await self._connect()
except Exception as e:
if not self._running:
break
logger.warning("WebSocket disconnected: %s. Reconnecting in %.0fs...", e, self._reconnect_delay)
await asyncio.sleep(self._reconnect_delay)
self._reconnect_delay = min(self._reconnect_delay * 2, 30.0)
async def _connect(self):
"""Establish WebSocket connection and process events."""
try:
import websockets
except ImportError:
logger.warning("websockets package not installed, skipping event subscription")
self._running = False
return
# Fix D (Cycle 5): include bearer token in WebSocket upgrade so the
# server's new auth check can validate this agent connection.
# Graceful fallback for workspaces that have no token yet.
headers = {"X-Workspace-ID": self.workspace_id}
try:
from platform_auth import auth_headers as _auth_headers
headers.update(_auth_headers())
except Exception:
pass # No token available — connect unauthenticated (grandfathered)
logger.info("Connecting to platform WebSocket: %s", self.ws_url)
async with websockets.connect(self.ws_url, additional_headers=headers) as ws:
self._reconnect_delay = 1.0 # Reset on successful connect
logger.info("Platform WebSocket connected")
async for message in ws:
try:
event = json.loads(message)
event_type = event.get("event", "")
if event_type in REBUILD_EVENTS:
logger.info("Peer event: %s for workspace %s",
event_type, event.get("workspace_id", ""))
if self.on_peer_change:
await self.on_peer_change(event)
except json.JSONDecodeError:
continue
except Exception as e:
logger.warning("Error processing event: %s", e)
def stop(self):
self._running = False