From 21a5a92cd97091c7616d2b432c5d0f71f5c16508 Mon Sep 17 00:00:00 2001
From: Tito0015 <mhiritarek@gmail.com>
Date: Sun, 4 Oct 2026 01:41:53 -0400
Subject: [PATCH] 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 <https://github.co
    Tornado, Web framework
    Twisted, Web framework
    Gradio, Web framework
+   fastmcp, MCP server framework
+   mcp, MCP server framework
    starlette, Asynchronous Server Gateway Interface (ASGI)
    ldap3, Lightweight Directory Access Protocol (LDAP)
    python-ldap, Lightweight Directory Access Protocol (LDAP)
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
new file mode 100644
index 000000000000..98d92dbc49a6
--- /dev/null
+++ b/python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md
@@ -0,0 +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`).
diff --git a/python/ql/lib/semmle/python/Frameworks.qll b/python/ql/lib/semmle/python/Frameworks.qll
index 1c2f0a6e6d41..51efee7b48be 100644
--- a/python/ql/lib/semmle/python/Frameworks.qll
+++ b/python/ql/lib/semmle/python/Frameworks.qll
@@ -50,6 +50,7 @@ private import semmle.python.frameworks.Libxml2
 private import semmle.python.frameworks.Lxml
 private import semmle.python.frameworks.Mako
 private import semmle.python.frameworks.MarkupSafe
+private import semmle.python.frameworks.Mcp
 private import semmle.python.frameworks.Multidict
 private import semmle.python.frameworks.Mysql
 private import semmle.python.frameworks.MySQLdb
diff --git a/python/ql/lib/semmle/python/frameworks/Mcp.qll b/python/ql/lib/semmle/python/frameworks/Mcp.qll
new file mode 100644
index 000000000000..f3b51d8f2dbf
--- /dev/null
+++ b/python/ql/lib/semmle/python/frameworks/Mcp.qll
@@ -0,0 +1,137 @@
+/**
+ * Provides classes modeling security-relevant aspects of the `mcp` and `fastmcp` PyPI packages
+ * (Model Context Protocol server SDKs).
+ */
+
+private import python
+private import semmle.python.dataflow.new.DataFlow
+private import semmle.python.dataflow.new.RemoteFlowSources
+private import semmle.python.ApiGraphs
+
+/**
+ * Provides models for MCP server handler parameters as remote flow sources.
+ */
+module Mcp {
+  private string registrationMethod() { result in ["tool", "prompt", "resource"] }
+
+  /** Gets API nodes for MCP / FastMCP server classes used to register handlers. */
+  private API::Node mcpServerClass() {
+    result =
+      [
+        API::moduleImport("fastmcp").getMember("FastMCP"),
+        API::moduleImport("fastmcp").getMember("server").getMember("FastMCP"),
+        API::moduleImport("fastmcp").getMember("server").getMember("server").getMember("FastMCP"),
+        API::moduleImport("mcp").getMember("server").getMember("fastmcp").getMember("FastMCP"),
+        API::moduleImport("mcp")
+            .getMember("server")
+            .getMember("fastmcp")
+            .getMember("server")
+            .getMember("FastMCP"),
+        API::moduleImport("mcp").getMember("server").getMember("MCPServer"),
+        API::moduleImport("mcp").getMember("server").getMember("mcpserver").getMember("MCPServer"),
+        API::moduleImport("mcp")
+            .getMember("server")
+            .getMember("mcpserver")
+            .getMember("server")
+            .getMember("MCPServer"),
+      ]
+  }
+
+  /** 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) {
+    exists(API::Node cls | cls = mcpServerClass() |
+      exists(API::CallNode call |
+        call = cls.getAnInstance().getMember(registrationMethod()).getACall() and
+        call.getNode().getNode() = handler.getADecorator()
+      )
+      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()
+          )
+        )
+      )
+    )
+    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()
+        )
+        or
+        handler.getADecorator() = reg.getAValueReachableFromSource().asExpr()
+      )
+    )
+  }
+
+  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
+      ann = ctx.getAValueReachableFromSource().asExpr()
+    )
+  }
+
+  private predicate isMcpContextParameter(Parameter p) {
+    exists(Expr ann | ann = p.getAnnotation() |
+      isContextAnnotation(ann)
+      or
+      exists(Subscript sub | sub = ann | isContextAnnotation(sub.getValue()))
+      or
+      exists(BinaryExpr bin | bin = ann and bin.getOp() instanceof BitOr |
+        isContextAnnotation(bin.getLeft()) or isContextAnnotation(bin.getRight())
+      )
+    )
+  }
+
+  private predicate isHandlerParameter(Function handler, Parameter p) {
+    (
+      p = handler.getAnArg()
+      or
+      p = handler.getAKeywordOnlyArg()
+      or
+      p = handler.getVararg()
+      or
+      p = handler.getKwarg()
+    ) and
+    not (handler.isMethod() and p = handler.getArg(0)) and
+    not isMcpContextParameter(p)
+  }
+
+  private class McpHandlerParameter extends RemoteFlowSource::Range, DataFlow::ParameterNode {
+    McpHandlerParameter() {
+      exists(Function handler |
+        isMcpHandler(handler) and
+        isHandlerParameter(handler, this.getParameter())
+      )
+    }
+
+    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
new file mode 100644
index 000000000000..9544d54e00bd
--- /dev/null
+++ b/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected
@@ -0,0 +1,12 @@
+| 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 |
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
new file mode 100644
index 000000000000..7dacb68ee854
--- /dev/null
+++ b/python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql
@@ -0,0 +1,6 @@
+import python
+import semmle.python.dataflow.new.RemoteFlowSources
+
+from RemoteFlowSource src
+where 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
new file mode 100644
index 000000000000..4809aa27a417
--- /dev/null
+++ b/python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py
@@ -0,0 +1,81 @@
+from fastmcp import Context, FastMCP
+from fastmcp import tool as standalone_tool
+from mcp.server.fastmcp import FastMCP as McpFastMCP
+from mcp.server import MCPServer
+
+mcp_app = FastMCP("test")
+mcp_sdk = McpFastMCP("test")
+mcp_v2 = MCPServer("test")
+
+
+def other_decorator(fn):
+    return fn
+
+
+@mcp_app.tool()
+def sync_tool(command: str) -> 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