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

Merge pull request #1791 from Molecule-AI/fix/memory-poisoning-GH1610

Browse Source 
fix(security): cross-tenant memory poisoning — GLOBAL scope isolation (GH#1610)
This commit is contained in:
Hongming Wang 2026-04-23 10:26:02 -07:00 committed by GitHub
commit 6904a8c448
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 236 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

127
workspace/a2a_tools.py
View File

@ -5,6 +5,7 @@ Imports shared client functions and constants from a2a_client.
import hashlib
import json
import os
import uuid
import httpx
@ -22,6 +23,83 @@ from a2a_client import (
from builtin_tools.security import _redact_secrets
# ---------------------------------------------------------------------------
# RBAC helpers (mirror builtin_tools/audit.py for a2a_tools isolation)
# ---------------------------------------------------------------------------
_ROLE_PERMISSIONS = {
"admin": {"delegate", "approve", "memory.read", "memory.write"},
"operator": {"delegate", "approve", "memory.read", "memory.write"},
"read-only": {"memory.read"},
"no-delegation": {"approve", "memory.read", "memory.write"},
"no-approval": {"delegate", "memory.read", "memory.write"},
"memory-readonly": {"memory.read"},
}
def _get_workspace_tier() -> int:
"""Return the workspace tier from config (0 = root, 1+ = tenant)."""
try:
from config import load_config
cfg = load_config()
return getattr(cfg, "tier", 1)
except Exception:
return int(os.environ.get("WORKSPACE_TIER", 1))
def _check_memory_write_permission() -> bool:
"""Return True if this workspace's RBAC roles grant memory.write."""
try:
from config import load_config
cfg = load_config()
roles = list(getattr(cfg, "rbac", None).roles or ["operator"])
allowed = dict(getattr(cfg, "rbac", None).allowed_actions or {})
except Exception:
# Fail closed: deny when config is unavailable
roles = ["operator"]
allowed = {}
for role in roles:
if role == "admin":
return True
if role in allowed:
if "memory.write" in allowed[role]:
return True
elif role in _ROLE_PERMISSIONS and "memory.write" in _ROLE_PERMISSIONS[role]:
return True
return False
def _check_memory_read_permission() -> bool:
"""Return True if this workspace's RBAC roles grant memory.read."""
try:
from config import load_config
cfg = load_config()
roles = list(getattr(cfg, "rbac", None).roles or ["operator"])
allowed = dict(getattr(cfg, "rbac", None).allowed_actions or {})
except Exception:
roles = ["operator"]
allowed = {}
for role in roles:
if role == "admin":
return True
if role in allowed:
if "memory.read" in allowed[role]:
return True
elif role in _ROLE_PERMISSIONS and "memory.read" in _ROLE_PERMISSIONS[role]:
return True
return False
def _is_root_workspace() -> bool:
"""Return True if this workspace is tier 0 (root/root-org)."""
return _get_workspace_tier() == 0
def _auth_headers_for_heartbeat() -> dict[str, str]:
"""Return Phase 30.1 auth headers; tolerate platform_auth being absent
in older installs (e.g. during rolling upgrade)."""
@ -228,18 +306,46 @@ async def tool_get_workspace_info() -> str:
async def tool_commit_memory(content: str, scope: str = "LOCAL") -> str:
"""Save important information to persistent memory."""
"""Save important information to persistent memory.
GLOBAL scope is writable only by root workspaces (tier == 0).
RBAC memory.write permission is required for all scope levels.
The source workspace_id is embedded in every record so the platform
can enforce cross-workspace isolation and audit trail.
"""
if not content:
return "Error: content is required"
content = _redact_secrets(content)
scope = scope.upper()
if scope not in ("LOCAL", "TEAM", "GLOBAL"):
scope = "LOCAL"
# RBAC: require memory.write permission (mirrors builtin_tools/memory.py)
if not _check_memory_write_permission():
return (
"Error: RBAC — this workspace does not have the 'memory.write' "
"permission for this operation."
)
# Scope enforcement: only root workspaces (tier 0) can write GLOBAL memory.
# This prevents tenant workspaces from poisoning org-wide memory (GH#1610).
if scope == "GLOBAL" and not _is_root_workspace():
return (
"Error: RBAC — only root workspaces (tier 0) can write to GLOBAL scope. "
"Non-root workspaces may use LOCAL or TEAM scope."
)
try:
async with httpx.AsyncClient(timeout=10.0) as client:
resp = await client.post(
f"{PLATFORM_URL}/workspaces/{WORKSPACE_ID}/memories",
json={"content": content, "scope": scope},
json={
"content": content,
"scope": scope,
# Embed source workspace so the platform can namespace-isolate
# and audit cross-workspace writes (GH#1610 fix).
"workspace_id": WORKSPACE_ID,
},
headers=_auth_headers_for_heartbeat(),
)
data = resp.json()
@ -251,8 +357,21 @@ async def tool_commit_memory(content: str, scope: str = "LOCAL") -> str:
async def tool_recall_memory(query: str = "", scope: str = "") -> str:
"""Search persistent memory for previously saved information."""
params = {}
"""Search persistent memory for previously saved information.
RBAC memory.read permission is required (mirrors builtin_tools/memory.py).
The workspace_id is sent as a query parameter so the platform can
cross-validate it against the auth token and defend against any future
path traversal / cross-tenant read bugs in the platform itself.
"""
# RBAC: require memory.read permission (mirrors builtin_tools/memory.py)
if not _check_memory_read_permission():
return (
"Error: RBAC — this workspace does not have the 'memory.read' "
"permission for this operation."
)
params: dict[str, str] = {"workspace_id": WORKSPACE_ID}
if query:
params["q"] = query
if scope:

132
workspace/tests/test_a2a_tools_impl.py
View File

@ -469,7 +469,9 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_resp=_resp(201, {"id": "mem-1"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("Remember this", scope="local")
data = json.loads(result)
@ -481,7 +483,9 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_resp=_resp(200, {"id": "mem-2"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("Remember this", scope="INVALID")
data = json.loads(result)
@ -491,17 +495,22 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_resp=_resp(200, {"id": "mem-3"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("Team info", scope="TEAM")
data = json.loads(result)
assert data["scope"] == "TEAM"
async def test_global_scope_accepted(self):
async def test_global_scope_accepted_for_root_workspace(self):
"""GLOBAL scope succeeds only when _is_root_workspace() returns True."""
import a2a_tools
mc = _make_http_mock(post_resp=_resp(201, {"id": "mem-4"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=True):
result = await a2a_tools.tool_commit_memory("Global info", scope="GLOBAL")
data = json.loads(result)
@ -511,7 +520,9 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_resp=_resp(200, {"id": "mem-5"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("info")
data = json.loads(result)
@ -522,7 +533,9 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_resp=_resp(201, {"id": "mem-6"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("info")
data = json.loads(result)
@ -533,7 +546,9 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_resp=_resp(400, {"error": "bad request payload"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("info")
assert "Error" in result
@ -543,12 +558,65 @@ class TestToolCommitMemory:
import a2a_tools
mc = _make_http_mock(post_exc=RuntimeError("storage failure"))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("info")
assert "Error saving memory" in result
assert "storage failure" in result
# -----------------------------------------------------------------------
# GH#1610 — cross-tenant memory poisoning security regression tests
# -----------------------------------------------------------------------
async def test_global_scope_denied_for_non_root_workspace(self):
"""Tenant (tier > 0) cannot write to GLOBAL scope (GH#1610)."""
import a2a_tools
mc = _make_http_mock(post_resp=_resp(201, {"id": "mem-poison"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("poisoned GLOBAL memory", scope="GLOBAL")
# Must NOT have called the platform — early rejection
mc.post.assert_not_called()
assert "Error" in result
assert "GLOBAL" in result
assert "tier 0" in result
async def test_rbac_deny_blocks_all_scopes_including_local(self):
"""RBAC memory.write denial blocks all scope levels (GH#1610)."""
import a2a_tools
mc = _make_http_mock(post_resp=_resp(201, {"id": "mem-7"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=False), \
patch("a2a_tools._is_root_workspace", return_value=False):
result = await a2a_tools.tool_commit_memory("should be denied", scope="LOCAL")
mc.post.assert_not_called()
assert "Error" in result
assert "memory.write" in result
async def test_post_includes_workspace_id_in_body(self):
"""POST body includes workspace_id so platform can audit/namespace (GH#1610)."""
import a2a_tools
mc = _make_http_mock(post_resp=_resp(201, {"id": "mem-8"}))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_write_permission", return_value=True), \
patch("a2a_tools._is_root_workspace", return_value=False):
await a2a_tools.tool_commit_memory("test content", scope="LOCAL")
call_kwargs = mc.post.call_args.kwargs
payload = call_kwargs.get("json")
assert payload is not None
assert "workspace_id" in payload
# Value should be the module's WORKSPACE_ID constant
assert payload["workspace_id"] == a2a_tools.WORKSPACE_ID
# ---------------------------------------------------------------------------
# tool_recall_memory
@ -564,7 +632,8 @@ class TestToolRecallMemory:
{"scope": "TEAM", "content": "We use Python 3.11"},
]
mc = _make_http_mock(get_resp=_resp(200, memories))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
result = await a2a_tools.tool_recall_memory(query="capital")
assert "[LOCAL]" in result
@ -576,7 +645,8 @@ class TestToolRecallMemory:
import a2a_tools
mc = _make_http_mock(get_resp=_resp(200, []))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
result = await a2a_tools.tool_recall_memory(query="anything")
assert result == "No memories found."
@ -587,7 +657,8 @@ class TestToolRecallMemory:
payload = {"error": "search unavailable"}
mc = _make_http_mock(get_resp=_resp(200, payload))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
result = await a2a_tools.tool_recall_memory()
parsed = json.loads(result)
@ -597,7 +668,8 @@ class TestToolRecallMemory:
import a2a_tools
mc = _make_http_mock(get_exc=RuntimeError("search service down"))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
result = await a2a_tools.tool_recall_memory(query="test")
assert "Error recalling memory" in result
@ -608,35 +680,57 @@ class TestToolRecallMemory:
import a2a_tools
mc = _make_http_mock(get_resp=_resp(200, []))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
await a2a_tools.tool_recall_memory(query="paris", scope="local")
call_kwargs = mc.get.call_args.kwargs
params = call_kwargs.get("params", {})
assert params.get("q") == "paris"
assert params.get("scope") == "LOCAL" # uppercased
assert params.get("workspace_id") == a2a_tools.WORKSPACE_ID
async def test_no_query_or_scope_sends_empty_params(self):
"""With no query/scope, params dict is empty (no keys added)."""
async def test_recall_includes_workspace_id_in_params(self):
"""workspace_id is always included in params for platform cross-validation (GH#1610)."""
import a2a_tools
mc = _make_http_mock(get_resp=_resp(200, []))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
await a2a_tools.tool_recall_memory()
call_kwargs = mc.get.call_args.kwargs
params = call_kwargs.get("params", {})
assert params == {}
assert "workspace_id" in params
assert params["workspace_id"] == a2a_tools.WORKSPACE_ID
async def test_scope_only_uppercased_in_params(self):
"""scope without query → only 'scope' key in params, uppercased."""
import a2a_tools
mc = _make_http_mock(get_resp=_resp(200, []))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc):
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=True):
await a2a_tools.tool_recall_memory(scope="team")
call_kwargs = mc.get.call_args.kwargs
params = call_kwargs.get("params", {})
assert "q" not in params
assert params.get("scope") == "TEAM"
# -----------------------------------------------------------------------
# GH#1610 — cross-tenant memory poisoning security regression tests
# -----------------------------------------------------------------------
async def test_rbac_deny_blocks_recall(self):
"""RBAC memory.read denial blocks recall entirely (GH#1610)."""
import a2a_tools
mc = _make_http_mock(get_resp=_resp(200, [{"scope": "GLOBAL", "content": "secret"}]))
with patch("a2a_tools.httpx.AsyncClient", return_value=mc), \
patch("a2a_tools._check_memory_read_permission", return_value=False):
result = await a2a_tools.tool_recall_memory(query="secret")
mc.get.assert_not_called()
assert "Error" in result
assert "memory.read" in result