From 5e30bf27732527f4e13e37c4a0a0b4ba3bc8e475 Mon Sep 17 00:00:00 2001 From: Tito0015 Date: Sun, 4 Oct 2026 01:41:53 -0400 Subject: [PATCH 1/3] 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 d0a798cf100de1a9bf765dd7f477f78734514e76 Mon Sep 17 00:00:00 2001 From: Tito0015 Date: Thu, 15 Oct 2026 14:59:30 -0400 Subject: [PATCH 2/3] 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 From 99cdc47cbd1b2514805a6df32adc7c9b382b4c85 Mon Sep 17 00:00:00 2001 From: Tito0015 Date: Wed, 4 Nov 2026 11:54:37 -0500 Subject: [PATCH 3/3] Python: model MCP handler params as ai threat-model sources Apply review feedback: ThreatModelSource kind ai, expanded Context API paths, isStaticmethod parameter handling, and staticmethod test coverage. --- .../2026-10-04-mcp-remote-sources.md | 2 +- .../ql/lib/semmle/python/frameworks/Mcp.qll | 27 ++++++++++++++----- .../RemoteFlowSources.expected | 1 + .../remote-flow-sources/RemoteFlowSources.ql | 6 ++--- .../dataflow/remote-flow-sources/mcp_test.py | 7 +++++ 5 files changed, 32 insertions(+), 11 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 d51085c184a2..6cd78637b26c 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 complementary `RemoteFlowSource` modeling for MCP/`fastmcp` handlers that Models-as-Data misses: stacked custom decorators and instance registrations such as `@self.mcp.tool()`. +* Added complementary **`ai`** threat-model sources for MCP/`fastmcp` handlers that Models-as-Data misses: stacked custom decorators, instance registrations such as `@self.mcp.tool()`, and `@staticmethod` handler parameters. diff --git a/python/ql/lib/semmle/python/frameworks/Mcp.qll b/python/ql/lib/semmle/python/frameworks/Mcp.qll index 18f8b0e0ab0c..36cb4921f0e8 100644 --- a/python/ql/lib/semmle/python/frameworks/Mcp.qll +++ b/python/ql/lib/semmle/python/frameworks/Mcp.qll @@ -9,11 +9,12 @@ private import python private import semmle.python.dataflow.new.DataFlow -private import semmle.python.dataflow.new.RemoteFlowSources +private import semmle.python.dataflow.new.internal.DataFlowDispatch +private import semmle.python.Concepts private import semmle.python.ApiGraphs /** - * Provides models for MCP server handler parameters as remote flow sources. + * Provides models for MCP server handler parameters as AI threat-model sources. */ module Mcp { private string registrationMethod() { result in ["tool", "prompt", "resource"] } @@ -93,9 +94,15 @@ module Mcp { private predicate isContextAnnotation(Expr ann) { exists(API::Node ctx | ctx in [ - API::moduleImport("fastmcp").getMember("Context"), - API::moduleImport("mcp").getMember("server").getMember("fastmcp").getMember("Context"), - ] and + API::moduleImport("fastmcp").getMember("Context"), + API::moduleImport("fastmcp").getMember("server").getMember("context").getMember("Context"), + API::moduleImport("mcp").getMember("server").getMember("fastmcp").getMember("Context"), + API::moduleImport("mcp") + .getMember("server") + .getMember("fastmcp") + .getMember("server") + .getMember("Context"), + ] and ann = ctx.getAValueReachableFromSource().asExpr() ) } @@ -122,11 +129,15 @@ module Mcp { or p = handler.getKwarg() ) and - not (handler.isMethod() and p = handler.getArg(0)) and + not ( + handler.isMethod() and + not isStaticmethod(handler) and + p = handler.getArg(0) + ) and not isMcpContextParameter(p) } - private class McpHandlerParameter extends RemoteFlowSource::Range, DataFlow::ParameterNode { + private class McpHandlerParameter extends ThreatModelSource::Range, DataFlow::ParameterNode { McpHandlerParameter() { exists(Function handler | isMcpHandler(handler) and @@ -134,6 +145,8 @@ module Mcp { ) } + override string getThreatModel() { result = "ai" } + override string getSourceType() { result = "MCP server handler parameter" } } } 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 bc01256b1652..a649f4e8d3e8 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,3 +1,4 @@ | mcp_test.py:61:23:61:25 | ControlFlowNode for cmd | 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 | +| mcp_test.py:93:25:93:28 | ControlFlowNode for host | MCP server handler parameter | diff --git a/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql b/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql index 7dacb68ee854..ad8fcc1bc656 100644 --- a/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql +++ b/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql @@ -1,6 +1,6 @@ import python -import semmle.python.dataflow.new.RemoteFlowSources +import semmle.python.Concepts -from RemoteFlowSource src -where exists(src.getLocation().getFile().getRelativePath()) +from ThreatModelSource src +where src.getThreatModel() = "ai" and exists(src.getLocation().getFile().getRelativePath()) select src, src.getSourceType() 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 1a6a3a1ce940..800b975314c3 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 @@ -85,3 +85,10 @@ def instance_lookup(host: str) -> str: @mcp_app.tool() def stacked_with_context(ctx: Context, payload: str) -> str: return payload + + +class StaticMcpHandlers: + @mcp_app.tool() + @staticmethod + def static_mcp_tool(host: str) -> str: + return host