From 3c2c77e9c654c082981616e571f59f00e1fa92eb Mon Sep 17 00:00:00 2001 From: Tito0015 Date: Sun, 4 Oct 2026 01:41:53 -0400 Subject: [PATCH 1/2] Python: model MCP and FastMCP server handler parameters as remote flow sources --- .../codeql/reusables/supported-frameworks.rst | 2 + .../2026-10-04-mcp-remote-sources.md | 4 + python/ql/lib/semmle/python/Frameworks.qll | 1 + .../ql/lib/semmle/python/frameworks/Mcp.qll | 137 ++++++++++++++++++ .../RemoteFlowSources.expected | 12 ++ .../remote-flow-sources/RemoteFlowSources.ql | 6 + .../dataflow/remote-flow-sources/mcp_test.py | 81 +++++++++++ 7 files changed, 243 insertions(+) create mode 100644 python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md create mode 100644 python/ql/lib/semmle/python/frameworks/Mcp.qll create mode 100644 python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected create mode 100644 python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql create mode 100644 python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py diff --git a/docs/codeql/reusables/supported-frameworks.rst b/docs/codeql/reusables/supported-frameworks.rst index 6fa32547d05a..78ca5754d7e2 100644 --- a/docs/codeql/reusables/supported-frameworks.rst +++ b/docs/codeql/reusables/supported-frameworks.rst @@ -226,6 +226,8 @@ and the CodeQL library pack ``codeql/python-all`` (`changelog str: + return command + + +@mcp_app.tool +def bare_tool(command: str) -> str: + return command + + +@mcp_sdk.prompt() +def sdk_prompt(topic: str, *, style: str = "short") -> str: + return topic + style + + +@mcp_v2.resource("notes://{name}") +def v2_resource(name: str) -> str: + return name + + +async def async_fetch(url: str) -> str: + return url + + +mcp_app.add_tool(async_fetch) + + +def added_via_add(url: str) -> str: + return url + + +mcp_sdk.add_tool(fn=added_via_add) + + +@standalone_tool +def module_tool(cmd: str) -> str: + return cmd + + +@standalone_tool() +def module_tool_called(cmd: str) -> str: + return cmd + + +@other_decorator +@mcp_app.tool() +def stacked_decorator(cmd: str) -> str: + return cmd + + +def not_registered(value: str) -> str: + return value + + +@mcp_app.tool() +def with_context(ctx: Context, data: str) -> str: + return data + + +class Service: + def __init__(self): + self.mcp = FastMCP("svc") + + def register(self): + @self.mcp.tool() + def instance_lookup(host: str) -> str: + return host From b33c67602ec9560e0797723a5fee181a32e25f81 Mon Sep 17 00:00:00 2001 From: Tito0015 Date: Thu, 15 Oct 2026 14:59:30 -0400 Subject: [PATCH 2/2] Python: restrict MCP remote sources to stacked and instance handlers --- .../2026-10-04-mcp-remote-sources.md | 2 +- .../ql/lib/semmle/python/frameworks/Mcp.qll | 78 ++++++++++--------- .../RemoteFlowSources.expected | 11 +-- .../dataflow/remote-flow-sources/mcp_test.py | 6 ++ 4 files changed, 48 insertions(+), 49 deletions(-) diff --git a/python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md b/python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md index 98d92dbc49a6..d51085c184a2 100644 --- a/python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md +++ b/python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md @@ -1,4 +1,4 @@ --- category: minorAnalysis --- -* Added `RemoteFlowSource` modeling for parameters of MCP server handlers registered via the `mcp` and `fastmcp` PyPI packages (`tool`, `prompt`, `resource`, and `add_tool`). +* Added complementary `RemoteFlowSource` modeling for MCP/`fastmcp` handlers that Models-as-Data misses: stacked custom decorators and instance registrations such as `@self.mcp.tool()`. diff --git a/python/ql/lib/semmle/python/frameworks/Mcp.qll b/python/ql/lib/semmle/python/frameworks/Mcp.qll index f3b51d8f2dbf..18f8b0e0ab0c 100644 --- a/python/ql/lib/semmle/python/frameworks/Mcp.qll +++ b/python/ql/lib/semmle/python/frameworks/Mcp.qll @@ -1,6 +1,10 @@ /** * Provides classes modeling security-relevant aspects of the `mcp` and `fastmcp` PyPI packages * (Model Context Protocol server SDKs). + * + * Baseline `@mcp.tool` / `@mcp.prompt` / `@mcp.resource` handlers on standalone functions are + * covered by Models-as-Data. This module only captures programmatic edge cases MaD misses: + * stacked/wrapped custom decorators and instance registrations such as `@self.mcp.tool()`. */ private import python @@ -37,57 +41,55 @@ module Mcp { ] } - /** Gets API nodes for module-level `fastmcp` registration decorators (`@tool`, etc.). */ - private API::Node fastmcpModuleRegistration(string name) { - name = registrationMethod() and result = API::moduleImport("fastmcp").getMember(name) - or - name = "tool" and result = API::moduleImport("fastmcp").getMember("tools").getMember("tool") - } - - private predicate isMcpHandler(Function handler) { + /** An MCP `@tool`/`@prompt`/`@resource` decorator expression on `handler`. */ + private predicate mcpRegistrationDecorator(Function handler, Expr dec) { + dec = handler.getADecorator() and exists(API::Node cls | cls = mcpServerClass() | exists(API::CallNode call | call = cls.getAnInstance().getMember(registrationMethod()).getACall() and - call.getNode().getNode() = handler.getADecorator() + call.getNode().getNode() = dec ) or exists(string name | name = registrationMethod() | - handler.getADecorator() = - cls.getAnInstance().getMember(name).getAValueReachableFromSource().asExpr() - ) - or - exists(API::CallNode addCall | - addCall = cls.getAnInstance().getMember("add_tool").getACall() and - ( - exists(DataFlow::LocalSourceNode funcSrc | - ( - funcSrc.flowsTo(addCall.getArg(0)) or - funcSrc.flowsTo(addCall.getArgByName("fn")) or - funcSrc.flowsTo(addCall.getArgByName("tool")) - ) and - funcSrc.asExpr() = handler.getDefinition() - ) - or - exists(DataFlow::Node arg | - arg in [addCall.getArg(0), addCall.getArgByName("fn"), addCall.getArgByName("tool")] and - arg.asExpr().(Name).getId() = handler.getName() and - arg.getScope() = handler.getScope() - ) - ) + dec = cls.getAnInstance().getMember(name).getAValueReachableFromSource().asExpr() ) ) - or - exists(string name | name = registrationMethod() | - exists(API::Node reg | reg = fastmcpModuleRegistration(name) | - exists(API::CallNode call | - call = reg.getACall() and call.getNode().getNode() = handler.getADecorator() - ) + } + + /** + * Holds if `handler` is an MCP registration wrapped by another decorator. + * MaD YAML currently drops these stacked/custom project wrappers. + */ + private predicate isWrappedMcpHandler(Function handler) { + exists(Expr mcpDec, Expr extra | + mcpRegistrationDecorator(handler, mcpDec) and + extra = handler.getADecorator() and + extra != mcpDec + ) + } + + /** + * Holds if `handler` is registered via a nested instance lookup such as `@self.mcp.tool()`. + * The callee object is itself an `Attribute`, not a simple name (`mcp_app.tool()`). + */ + private predicate isInstanceLookupHandler(Function handler) { + exists(Expr dec | mcpRegistrationDecorator(handler, dec) | + exists(Attribute callee | + callee = dec.(Call).getFunc() or - handler.getADecorator() = reg.getAValueReachableFromSource().asExpr() + callee = dec + | + callee.getObject() instanceof Attribute ) ) } + private predicate isMcpHandler(Function handler) { + isWrappedMcpHandler(handler) + or + isInstanceLookupHandler(handler) + } + private predicate isContextAnnotation(Expr ann) { exists(API::Node ctx | ctx in [ diff --git a/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected b/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected index 9544d54e00bd..bc01256b1652 100644 --- a/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected +++ b/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected @@ -1,12 +1,3 @@ -| mcp_test.py:16:15:16:21 | ControlFlowNode for command | MCP server handler parameter | -| mcp_test.py:21:15:21:21 | ControlFlowNode for command | MCP server handler parameter | -| mcp_test.py:26:16:26:20 | ControlFlowNode for topic | MCP server handler parameter | -| mcp_test.py:26:31:26:35 | ControlFlowNode for style | MCP server handler parameter | -| mcp_test.py:31:17:31:20 | ControlFlowNode for name | MCP server handler parameter | -| mcp_test.py:35:23:35:25 | ControlFlowNode for url | MCP server handler parameter | -| mcp_test.py:42:19:42:21 | ControlFlowNode for url | MCP server handler parameter | -| mcp_test.py:50:17:50:19 | ControlFlowNode for cmd | MCP server handler parameter | -| mcp_test.py:55:24:55:26 | ControlFlowNode for cmd | MCP server handler parameter | | mcp_test.py:61:23:61:25 | ControlFlowNode for cmd | MCP server handler parameter | -| mcp_test.py:70:32:70:35 | ControlFlowNode for data | MCP server handler parameter | | mcp_test.py:80:29:80:32 | ControlFlowNode for host | MCP server handler parameter | +| mcp_test.py:86:40:86:46 | ControlFlowNode for payload | MCP server handler parameter | diff --git a/python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py b/python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py index 4809aa27a417..1a6a3a1ce940 100644 --- a/python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py +++ b/python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py @@ -79,3 +79,9 @@ def register(self): @self.mcp.tool() def instance_lookup(host: str) -> str: return host + + +@other_decorator +@mcp_app.tool() +def stacked_with_context(ctx: Context, payload: str) -> str: + return payload