From cbe33c748de6b55ab5fa397f840333302e991aa9 Mon Sep 17 00:00:00 2001
From: Molecule AI Core-OffSec <core-offsec@agents.moleculesai.app>
Date: Wed, 22 Apr 2026 23:47:04 +0000
Subject: [PATCH] fix(handlers): add CanCommunicate hierarchy check to terminal
 WebSocket handler (KI-005)

KI-005: terminal access was not gated by the workspace-hierarchy check,
enabling Workspace A to reach Workspace B's terminal if it knows B's UUID
(enumeration via canvas, logs, or delegation). Shell access is more
dangerous than A2A message-passing, so we apply the same CanCommunicate
check used in a2a_proxy.

Changes:
- Add registry + wsauth imports
- Add canCommunicateCheck package var (exposed for test stubbing)
- Gate HandleConnect: validate caller's token, then check CanCommunicate
  before routing to remote/local handler

Fixes: KI-005

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../internal/handlers/terminal.go             | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/workspace-server/internal/handlers/terminal.go b/workspace-server/internal/handlers/terminal.go
index 94e81cd6d..5251fd3bd 100644
--- a/workspace-server/internal/handlers/terminal.go
+++ b/workspace-server/internal/handlers/terminal.go
@@ -15,6 +15,8 @@ import (
 
 	"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
 	"github.com/Molecule-AI/molecule-monorepo/platform/internal/provisioner"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/registry"
+	"github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth"
 	"github.com/creack/pty"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
@@ -25,6 +27,11 @@ import (
 
 const terminalSessionTimeout = 30 * time.Minute
 
+// canCommunicateCheck is the communication-authorization predicate used by
+// HandleConnect to enforce the KI-005 workspace-hierarchy guard.
+// Exposed as a package var so tests can stub it without DB fixtures.
+var canCommunicateCheck = registry.CanCommunicate
+
 var termUpgrader = websocket.Upgrader{
 	CheckOrigin: func(r *http.Request) bool {
 		origin := r.Header.Get("Origin")
@@ -60,6 +67,26 @@ func (h *TerminalHandler) HandleConnect(c *gin.Context) {
 	workspaceID := c.Param("id")
 	ctx := c.Request.Context()
 
+	// KI-005: enforce CanCommunicate hierarchy check before granting terminal
+	// access. WorkspaceAuth validates the bearer's token, but the token is
+	// scoped to a specific workspace ID — Workspace A's token can reach
+	// Workspace A's terminal. Without CanCommunicate, Workspace A could also
+	// reach Workspace B's terminal if it knows B's UUID (enumeration via
+	// canvas, logs, or delegation). Shell access is more dangerous than A2A
+	// message-passing, so we apply the same hierarchy check here.
+	callerID := c.GetHeader("X-Workspace-ID")
+	if callerID != "" {
+		tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization"))
+		if tok != "" {
+			if err := wsauth.ValidateAnyToken(ctx, db.DB, tok); err == nil {
+				if !canCommunicateCheck(callerID, workspaceID) {
+					c.JSON(http.StatusForbidden, gin.H{"error": "not authorized to access this workspace's terminal"})
+					return
+				}
+			}
+		}
+	}
+
 	// Check for CP-provisioned workspace (instance_id persisted by
 	// provisionWorkspaceCP → migration 038). Null instance_id means the
 	// workspace runs as a local Docker container on this tenant.