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>
97 lines
2.8 KiB
Go
97 lines
2.8 KiB
Go
package handlers
|
|
|
|
import (
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
|
|
"github.com/Molecule-AI/molecule-monorepo/platform/internal/metrics"
|
|
"github.com/Molecule-AI/molecule-monorepo/platform/internal/ws"
|
|
"github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
var upgrader = websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
// In production, validate against CORS_ORIGINS. In dev, allow all.
|
|
origins := os.Getenv("CORS_ORIGINS")
|
|
if origins == "" {
|
|
return true // dev mode — no restriction
|
|
}
|
|
origin := r.Header.Get("Origin")
|
|
for _, allowed := range strings.Split(origins, ",") {
|
|
if strings.EqualFold(strings.TrimSpace(allowed), origin) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
},
|
|
}
|
|
|
|
type SocketHandler struct {
|
|
hub *ws.Hub
|
|
}
|
|
|
|
func NewSocketHandler(hub *ws.Hub) *SocketHandler {
|
|
return &SocketHandler{hub: hub}
|
|
}
|
|
|
|
// HandleConnect handles WebSocket upgrade at GET /ws.
|
|
// Canvas clients connect without X-Workspace-ID — they receive all events.
|
|
// Workspace agents send X-Workspace-ID — events are filtered by CanCommunicate.
|
|
//
|
|
// Fix D (Cycle 5): agent connections (X-Workspace-ID present) are now validated
|
|
// via bearer token before the WebSocket upgrade. Canvas clients (no X-Workspace-ID)
|
|
// remain unauthenticated. Pre-token workspaces are grandfathered through.
|
|
func (h *SocketHandler) HandleConnect(c *gin.Context) {
|
|
workspaceID := c.GetHeader("X-Workspace-ID")
|
|
|
|
// Authenticate workspace agents (not canvas browser clients).
|
|
if workspaceID != "" {
|
|
ctx := c.Request.Context()
|
|
hasLive, err := wsauth.HasAnyLiveToken(ctx, db.DB, workspaceID)
|
|
if err != nil {
|
|
log.Printf("wsauth: WebSocket HasAnyLiveToken(%s) failed: %v", workspaceID, err)
|
|
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "auth check failed"})
|
|
return
|
|
}
|
|
if hasLive {
|
|
tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization"))
|
|
if tok == "" {
|
|
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "missing workspace auth token"})
|
|
return
|
|
}
|
|
if err := wsauth.ValidateToken(ctx, db.DB, workspaceID, tok); err != nil {
|
|
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid workspace auth token"})
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
|
if err != nil {
|
|
log.Printf("WebSocket upgrade error: %v", err)
|
|
return
|
|
}
|
|
|
|
client := &ws.Client{
|
|
Conn: conn,
|
|
WorkspaceID: workspaceID,
|
|
Send: make(chan []byte, 256),
|
|
}
|
|
|
|
h.hub.Register <- client
|
|
metrics.TrackWSConnect()
|
|
|
|
// Wrap WritePump and ReadPump so the gauge is decremented exactly once
|
|
// when the client's write goroutine exits (WritePump owns conn lifetime).
|
|
go func() {
|
|
ws.WritePump(client)
|
|
metrics.TrackWSDisconnect()
|
|
}()
|
|
go ws.ReadPump(client, h.hub)
|
|
}
|