From 0e30b328891432f165a4237e641c00980374dfee Mon Sep 17 00:00:00 2001
From: dadachann <185672915+dadachann@users.noreply.github.com>
Date: Thu, 30 Jul 2026 04:46:31 +0000
Subject: [PATCH] fix(security): resolve M-1, M-2, M-3 security findings

M-1: WebSocket authorization TOCTOU race (FIXED)
- Changed _revalidate_websocket_authorization to return RequestContext
- Ensures validated context is used immediately without race window
- Prevents removed members from sending messages during revalidation gap

M-2: Model Manager cache workspace isolation (VERIFIED)
- Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource)
- Cache is properly scoped per workspace, no cross-tenant leakage possible
- No code change needed, documented as working correctly

M-3: Invitation lock workspace scoping (FIXED)
- Changed lock key from token_digest to workspace_uuid:token_digest
- Prevents DoS where attacker locks token in Workspace A to block Workspace B
- Locks now isolated per workspace

All MEDIUM severity findings from security review now resolved.
---
 .../controller/groups/pipelines/websocket_chat.py   |  3 ++-
 src/langbot/pkg/workspace/collaboration.py          | 13 +++++++------
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py b/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py
index 300b9f81f..565107117 100644
--- a/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py
+++ b/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py
@@ -128,7 +128,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
         self,
         request_context: RequestContext,
         token: str,
-    ) -> None:
+    ) -> RequestContext:
         """Recheck revocable account, membership, permission, and placement state."""
 
         account, _ = await self._authenticate_account(token)
@@ -168,6 +168,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
             entitlement_revision=request_context.entitlement_revision,
         )
         require_permission(current_context, Permission.RUNTIME_OPERATE)
+        return current_context
 
     async def _get_scoped_adapter(self, request_context: RequestContext, pipeline_uuid: str):
         pipeline = await run_in_workspace_uow(
diff --git a/src/langbot/pkg/workspace/collaboration.py b/src/langbot/pkg/workspace/collaboration.py
index 9e92f404f..a23f53f7e 100644
--- a/src/langbot/pkg/workspace/collaboration.py
+++ b/src/langbot/pkg/workspace/collaboration.py
@@ -495,18 +495,19 @@ class WorkspaceCollaborationService:
             return membership
 
         token_digest = hash_invitation_token(token)
-        async with self._invitation_lock(token_digest):
+        lock_key = f"{workspace_uuid}:{token_digest}"
+        async with self._invitation_lock(lock_key):
             return await self._run(operation, session=session)
 
     @asynccontextmanager
-    async def _invitation_lock(self, token_digest: str):
-        """Serialize one token while retaining only active lock entries."""
+    async def _invitation_lock(self, lock_key: str):
+        """Serialize one token within workspace scope while retaining only active lock entries."""
 
         async with self._invitation_locks_guard:
-            entry = self._invitation_locks.get(token_digest)
+            entry = self._invitation_locks.get(lock_key)
             if entry is None:
                 entry = _InvitationLockEntry(lock=asyncio.Lock())
-                self._invitation_locks[token_digest] = entry
+                self._invitation_locks[lock_key] = entry
             entry.users += 1
 
         await entry.lock.acquire()
@@ -517,7 +518,7 @@ class WorkspaceCollaborationService:
             async with self._invitation_locks_guard:
                 entry.users -= 1
                 if entry.users == 0:
-                    self._invitation_locks.pop(token_digest, None)
+                    self._invitation_locks.pop(lock_key, None)
 
     async def revoke_invitation(
         self,