molecule-ai/molecule-core
35
0
Fork 2
Code Issues 16 Pull Requests 2 Actions 49 Packages Projects Releases Wiki Activity

fix(terminal): check org_token_id context to allow org-token A2A routing (KI-005 followup)

Browse Source 
PR #1885 introduced a regression: HandleConnect called wsauth.ValidateToken
for any bearer token when X-Workspace-ID ≠ workspaceID. Org-scoped tokens
(org_api_tokens table) are not in workspace_auth_tokens, so ValidateToken
always returned ErrInvalidToken for them → hard 401 for all A2A routing
that uses org tokens.

Fix: if WorkspaceAuth already validated an org token (org_token_id set in
gin context by orgtoken.Validate), skip the workspace_auth_tokens lookup and
trust the X-Workspace-ID claim. Hierarchy enforcement via canCommunicateCheck
is unchanged — org token holders are still subject to the workspace hierarchy.

Workspace-scoped tokens continue to require ValidateToken binding. Invalid
tokens (neither workspace-bound nor org-level) still return 401. This closes
the regression while preserving the KI-005 security property.

Add TestKI005_OrgToken_SkipsValidateToken to terminal_test.go as a regression
guard for this exact path.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Molecule AI Core Platform Lead 2026-04-23 23:36:48 +00:00
parent 6b62391e5d
commit 78f8391f02
2 changed files with 50 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
workspace-server/internal/handlers/terminal.go
View File

@ -77,17 +77,26 @@ func (h *TerminalHandler) HandleConnect(c *gin.Context) {
// A2A message-passing, so we apply the same hierarchy check here.
// GH#756/#1609 security fix: if the caller claims a specific workspace
// identity (X-Workspace-ID header), the bearer token — if present — must
// belong to that claimed workspace. ValidateAnyToken accepted ANY valid org
// token, allowing Workspace A to forge X-Workspace-ID: B and reach B's
// terminal if A held any valid token. ValidateToken binds the token to
// the claimed workspace identity.
// belong to that claimed workspace. Previously ValidateAnyToken accepted
// ANY valid org token, allowing Workspace A to forge X-Workspace-ID: B
// and reach B's terminal if A held any valid token. ValidateToken binds
// the workspace-scoped token to the claimed workspace identity. Org-level
// tokens are handled separately via the org_token_id context key.
callerID := c.GetHeader("X-Workspace-ID")
if callerID != "" && callerID != workspaceID {
tok := wsauth.BearerTokenFromHeader(c.GetHeader("Authorization"))
if tok != "" {
if err := wsauth.ValidateToken(ctx, db.DB, callerID, tok); err != nil {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token for claimed workspace"})
return
// Org-scoped tokens (org_api_tokens) are validated at the org level
// by WorkspaceAuth and do not have a workspace_auth_tokens row, so
// ValidateToken always returns ErrInvalidToken for them. If WorkspaceAuth
// already validated an org token (org_token_id set in context), trust
// the X-Workspace-ID claim — the hierarchy is enforced by
// canCommunicateCheck below. Reject everything else.
if c.GetString("org_token_id") == "" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token for claimed workspace"})
return
}
}
}
if !canCommunicateCheck(callerID, workspaceID) {

35
workspace-server/internal/handlers/terminal_test.go
View File

@ -455,3 +455,38 @@ func TestTerminalConnect_KI005_AllowsSiblingWorkspace(t *testing.T) {
}
}
// TestKI005_OrgToken_SkipsValidateToken verifies that when WorkspaceAuth already
// validated an org token (org_token_id set in gin context), the X-Workspace-ID
// claim is trusted without a workspace_auth_tokens lookup. The hierarchy is still
// enforced by canCommunicateCheck. Regression guard for the A2A routing regression
// introduced in GH#1885: internal routing uses org tokens which are not in
// workspace_auth_tokens, so ValidateToken would always fail for them.
func TestKI005_OrgToken_SkipsValidateToken(t *testing.T) {
setupTestDB(t) // no ValidateToken ExpectQuery — none should fire
prev := canCommunicateCheck
canCommunicateCheck = func(callerID, targetID string) bool {
// Simulate platform agent → target workspace (same org).
return callerID == "ws-platform" && targetID == "ws-target"
}
defer func() { canCommunicateCheck = prev }()
h := NewTerminalHandler(nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: "ws-target"}}
c.Request = httptest.NewRequest("GET", "/workspaces/ws-target/terminal", nil)
c.Request.Header.Set("X-Workspace-ID", "ws-platform")
c.Request.Header.Set("Authorization", "Bearer org-token-abc123")
// Simulate WorkspaceAuth having validated the org token (orgtoken.Validate
// succeeded). HandleConnect must skip ValidateToken and trust the claim.
c.Set("org_token_id", "tok-org-abc")
h.HandleConnect(c)
// Org token path: ValidateToken skipped → canCommunicateCheck=true →
// falls through to Docker path → 503 nil-docker (no Docker client).
if w.Code != http.StatusServiceUnavailable {
t.Errorf("org-token A2A: got %d, want 503 nil-docker (%s)", w.Code, w.Body.String())
}
}