From dc549181555d3a26395fffe9db736e25dc3c97cb Mon Sep 17 00:00:00 2001 From: bolbol Date: Thu, 30 Jul 2026 07:42:32 +0000 Subject: [PATCH] fix(security): lier autorisation MCP a l'agent authentifie (ContextVar) au lieu des arguments client --- app/routes/mcp.py | 66 ++++++++++++++++++++++++----------------------- 1 file changed, 34 insertions(+), 32 deletions(-) diff --git a/app/routes/mcp.py b/app/routes/mcp.py index ae671fd..e50a529 100644 --- a/app/routes/mcp.py +++ b/app/routes/mcp.py @@ -1,5 +1,6 @@ import json import logging +import contextvars from fastapi import APIRouter, Depends, Request, HTTPException from sse_starlette.sse import EventSourceResponse from mcp.server import Server @@ -14,6 +15,8 @@ logger = logging.getLogger(__name__) router = APIRouter() mcp_server = Server("context-hub") +current_agent_var: contextvars.ContextVar[str] = contextvars.ContextVar("current_agent", default="") + @mcp_server.list_tools() async def handle_list_tools() -> list[Tool]: return [ @@ -91,72 +94,67 @@ async def handle_list_tools() -> list[Tool]: @mcp_server.call_tool() async def handle_call_tool(name: str, arguments: dict) -> list[TextContent]: - # We will pass the agent via context, but MCP python SDK doesn't easily pass request context - # to tool handlers without custom Context object. - # For this implementation, since it's a single server instance, we will rely on - # the client to provide the agent name in arguments or we inject it. - # To keep things simple and secure, we will just use the arguments for now. - # A robust solution would tie the SSE connection to the auth session. - - # In a real scenario we'd use the SSE connection's tied agent. - # We will assume agent name is injected or we just enforce it via the API key. - # For now, let's just return the data. - + # Autorisation liee a l'identite reelle de la connexion SSE (current_agent_var), + # jamais a un champ "agent"/"name" declare par le client dans les arguments. + caller_agent = current_agent_var.get() + if name == "get_rules": scope = arguments.get("scope") - # without context of WHICH agent is calling, we'll return the rule if it exists. - # (Security note: in production, the scope check must happen here using the connection's agent) + if not has_scope_access(caller_agent, scope): + return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden on scope {scope}"}))] rule = await get_rule(scope) if not rule: - return [TextContent(type="text", text=json.dumps({"error": "Scope not found or forbidden"}))] + return [TextContent(type="text", text=json.dumps({"error": "Scope not found"}))] return [TextContent(type="text", text=json.dumps(rule))] - + elif name == "get_ports": + if not has_scope_access(caller_agent, "infra"): + return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden: requires infra scope"}))] ports = await get_all_ports() return [TextContent(type="text", text=json.dumps(ports))] - + elif name == "search_rules": query = arguments.get("query", "").lower() all_rules = await get_all_rules() results = [] - for scope, content in all_rules.items(): - if query in json.dumps(content).lower(): - results.append({"scope": scope, "content": content}) + for scope, content_ in all_rules.items(): + if not has_scope_access(caller_agent, scope): + continue + if query in json.dumps(content_).lower(): + results.append({"scope": scope, "content": content_}) return [TextContent(type="text", text=json.dumps(results))] - + elif name == "get_agent_config": - agent_name = arguments.get("name", "") - # Compile config all_rules = await get_all_rules() config = {} for scope in ALL_SCOPES: - if has_scope_access(agent_name, scope): + if has_scope_access(caller_agent, scope): config[scope] = all_rules.get(scope, {}) return [TextContent(type="text", text=json.dumps(config))] - + elif name == "check_port": + if not has_scope_access(caller_agent, "infra"): + return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden: requires infra scope"}))] port = arguments.get("port") port_info = await get_port(port) if port_info: return [TextContent(type="text", text=json.dumps({"status": "occupied", "info": port_info}))] else: return [TextContent(type="text", text=json.dumps({"status": "free", "port": port}))] - + elif name == "get_my_context": - agent_name = arguments.get("agent", "") all_rules = await get_all_rules() result = {} for scope in ALL_SCOPES: - if has_scope_access(agent_name, scope): + if has_scope_access(caller_agent, scope): result[scope] = all_rules.get(scope, {}) - return [TextContent(type="text", text=json.dumps({"agent": agent_name, "scopes": result}))] + return [TextContent(type="text", text=json.dumps({"agent": caller_agent, "scopes": result}))] elif name == "update_rule": scope = arguments.get("scope") - agent_name = arguments.get("agent", "") data = arguments.get("data", {}) - if not has_scope_access(agent_name, scope): - return [TextContent(type="text", text=json.dumps({"error": f"Agent {agent_name} forbidden on scope {scope}"}))] + if not has_scope_access(caller_agent, scope): + return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden on scope {scope}"}))] if not isinstance(data, dict): return [TextContent(type="text", text=json.dumps({"error": "data must be a JSON object"}))] existing = await get_rule(scope) or {} @@ -176,7 +174,11 @@ async def mcp_sse(request: Request, agent: str = Depends(get_current_agent)): sse_transports[agent] = transport async def run_server(): - await mcp_server.run(transport.read_stream(), transport.write_stream(), mcp_server.create_initialization_options()) + token = current_agent_var.set(agent) + try: + await mcp_server.run(transport.read_stream(), transport.write_stream(), mcp_server.create_initialization_options()) + finally: + current_agent_var.reset(token) import asyncio asyncio.create_task(run_server())