Python: Model complex MCP AST handler parameters as remote sources (Complementary to #22703) by Tito0015 · Pull Request #22750 · github/codeql

Tito0015 · 2026-10-04T05:43:29Z

Summary

This PR provides the native QL AST modeling for Python MCP and fastmcp frameworks to handle complex programmatic edge cases that currently fall outside the scope of Models-as-Data (MaD) YAML.

Context & Positioning

Following the excellent global engine update in #22822 (which upgraded legacy queries to accept ActiveThreatModelSource), MaD YAML PR #22703 is now perfectly positioned to handle the 80% baseline of standard MCP handlers and headers.

However, as noted in the mcp-vulnbench benchmarks, YAML currently drops handlers if they are wrapped behind stacked custom project decorators (e.g., @custom_telemetry above @mcp.tool()).

This PR is designed to be merged alongside #22703. It acts as a lightweight AST extension extending RemoteFlowSource::Range to capture:

Handlers hidden under wrapped/stacked custom decorators (via getADecorator()).
Complex instance method registrations (e.g., @self.mcp.tool()).
Strict exclusions for Context parameters to prevent false positives.

Testing

Rebased onto main. All local unit tests under python/ql/test/library-tests/dataflow/remote-flow-sources/ pass cleanly.

Both proposals for Python MCP sources on the 28 Python cases: the class detects 9, the rows 13, both 13. Records what each covers that the other does not, the method (derived image, checks) and the limits of the comparison.

Tito0015 · 2026-10-05T19:52:47Z

Hi @geoffw0 — hope you are well.

You reviewed and merged my earlier C++ MMIO PR (#22438); thank you again for that. When you have a moment, I would appreciate a look at this Python follow-up for #22702.

It models MCP / FastMCP server handler parameters as RemoteFlowSources in QL (Mcp.qll), mainly so queries that still take RemoteFlowSource directly (e.g. py/xxe, py/xml-bomb, py/nosql-injection) can see MCP tool/prompt/resource args, and so stacked custom decorators on handlers are covered via getADecorator() (similar to Flask routing).

There is already open MaD coverage in #22703; this PR is intentionally QL-only and complementary. Library tests are under python/ql/test/library-tests/dataflow/remote-flow-sources/ and pass locally with codeql test run.

No rush — happy to adjust based on team preference (QL vs MaD, scope, naming). Thanks for any pointers.

+ */
+
+private import python
+private import semmle.python.dataflow.new.DataFlow


geoffw0 · 2026-10-06T08:56:00Z

Hi, this looks reasonable at a glance (though I was a little surprised to see QL-modelled sources rather than MaD models). In any case I've started the tests and am handing over to another member of the team.

owen-mc · 2026-10-14T15:46:40Z

I was also surprised by the assertion that QL-modelling was needed. I think it's correct though. We should fix it - I think changing those queries (and any others) to use ActiveThreatModelSource instead of RemoteFlowSource would solve it. I've put up this PR to fix it: #22822.

Tito0015 · 2026-10-14T19:42:56Z

Thanks for the review and for driving this upstream architectural fix in #22822, @owen-mc — migrating those queries to ActiveThreatModelSource is definitely the right long-term move for the Python dataflow pipeline.

I would love to collaborate directly on #22822. I can help benchmark query recall across the affected legacy rules, run test suites against edge cases, or review changes alongside the team.

Once #22822 lands, I will rebase and align this PR (#22750):

MaD Alignment: Shift standard MCP tool, prompt, and resource handler definitions into Models-as-Data (MaD) YAML (coordinating with Python: Add models for the MCP server SDKs (mcp, fastmcp) #22703) so they feed cleanly into ActiveThreatModelSource.
Targeted QL Modeling: Keep the QL logic in Mcp.qll focused strictly on dynamic AST patterns that MaD cannot currently express—specifically stacked and custom wrapper decorators (getADecorator()), class/instance registrations (@self.mcp.tool(), add_tool), and strict parameter exclusions (e.g., Context types).

With @Alex-Hofer already referencing this work in mcp-vulnbench, there is clear momentum to establish robust, high-recall taint tracking for Agentic AI and MCP server interfaces. I’m very keen to help the CodeQL team establish standard security coverage across this emerging attack surface.

For background on my work in AST/CPG query lowering and static analysis compiler infrastructure:

GitHub: https://github.com/Tito0015
Technical Research Write-up: https://medium.com/@mhiritarek/the-terraform-moment-for-threat-detection-how-meta-compilation-ends-the-appsec-rule-crisis-ebbdf676d2b1

Let me know the best way to coordinate on #22822—I'm happy to push validation runs or draft test cases against your branch whenever helpful!

owen-mc · 2026-10-15T08:37:03Z

Thanks for the offer, @Tito0015 , but it isn't needed. #22822 has now been merged.

Tito0015 · 2026-10-15T20:01:45Z

Awesome to see #22822 land, @owen-mc! Great work resolving that across the query suite.

I have rebased this PR onto main and narrowed Mcp.qll strictly to the dynamic AST patterns that MaD cannot currently express—specifically stacked custom project decorators (getADecorator()) and instance registrations (@self.mcp.tool()).

This keeps the PR lightweight and complementary to #22703 without duplicating the baseline YAML definitions. Ready for review whenever the team has bandwidth!

…w sources

owen-mc · 2026-10-20T09:38:03Z

Thanks for updating the PR. Now that it's ready for review, please don't force-push again, as that makes it harder to review.

Tito0015 · 2026-10-20T16:39:10Z

Understood, @owen-mc! Leaving git history as-is and won't force-push. Thanks for taking a look whenever you have the bandwidth.

Tito0015 · 2026-11-02T16:56:25Z

@owen-mc Gentle nudge — merged latest main in 40805f5 without force-pushing, per your request. Tests still green. No rush; thanks for reviewing when you can

owen-mc · 2026-11-04T14:40:00Z

+    not isMcpContextParameter(p)
+  }
+
+  private class McpHandlerParameter extends RemoteFlowSource::Range, DataFlow::ParameterNode {


Suggested change

private class McpHandlerParameter extends RemoteFlowSource::Range, DataFlow::ParameterNode {

private class McpHandlerParameter extends ThreatModelSource::Range, DataFlow::ParameterNode {

I assume we want this to have kind ai, instead of remote. So, instead you should extend ThreatModelSource::Range and add the predicate override string getThreatModel() { result = "ai" }.

owen-mc · 2026-11-04T14:45:21Z

+      or
+      p = handler.getKwarg()
+    ) and
+    not (handler.isMethod() and p = handler.getArg(0)) and


We do want the first argument to static methods. (isMethod holds if handler is declared inside a class.) Add private import semmle.python.dataflow.new.internal.DataFlowDispatch at the top and make this change:

Suggested change

not (handler.isMethod() and p = handler.getArg(0)) and

not (

handler.isMethod() and

not isStaticmethod(handler) and

p = handler.getArg(0)

) and

Ideally you'd add a test for this as well.

owen-mc · 2026-11-04T15:00:01Z

+    exists(API::Node ctx |
+      ctx in [
+          API::moduleImport("fastmcp").getMember("Context"),
+          API::moduleImport("mcp").getMember("server").getMember("fastmcp").getMember("Context"),


Suggested change

API::moduleImport("mcp").getMember("server").getMember("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"),

Apply review feedback: ThreatModelSource kind ai, expanded Context API paths, isStaticmethod parameter handling, and staticmethod test coverage.

Tito0015 · 2026-11-04T17:00:36Z

Thanks for the clear feedback, @owen-mc! All items have been addressed in 99cdc47 (pushed normally without force-pushing):

Threat Model Source: McpHandlerParameter now extends ThreatModelSource::Range with override string getThreatModel() { result = "ai" }.
Static Methods: Imported DataFlowDispatch and adjusted parameter exclusions with not isStaticmethod(handler).
Context API Paths: Added the full alternative import paths for Context (fastmcp.server.context.Context and mcp.server.fastmcp.server.Context).
Tests & Docs: Added a @staticmethod tool test case in mcp_test.py, retargeted RemoteFlowSources.ql to select ThreatModelSource with kind "ai", verified the test suite passes locally (codeql test run), and updated the change note.

Also merged latest upstream/main cleanly. Ready for your review whenever you have time!

Copilot

🟡 Changes recommended

The MCP model misclassifies injected Context parameters and drops valid handler inputs, while the regression query would produce unexpected rows when merged with the companion AI-source changes.

6 open issues

What changed in this PR

Adds native Python AST modeling for complex mcp and fastmcp handler registrations, complementing #22703’s Models-as-Data coverage under the opt-in ai threat model.

Changes:

Recognizes stacked decorators, instance lookups, and static handlers.
Models handler parameters as sources while excluding recognized Context annotations.
Adds regression fixtures, expected results, and framework-support documentation.

File	Description
`python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.ql`	Queries AI sources for regression testing.
`python/ql/test/library-tests/dataflow/remote-flow-sources/RemoteFlowSources.expected`	Records four expected native parameter sources.
`python/ql/test/library-tests/dataflow/remote-flow-sources/mcp_test.py`	Exercises handler registrations and Context exclusions.
`python/ql/lib/semmle/python/frameworks/Mcp.qll`	Implements handler recognition and parameter-source modeling.
`python/ql/lib/semmle/python/Frameworks.qll`	Loads the MCP model with existing framework models.
`python/ql/lib/change-notes/2026-10-04-mcp-remote-sources.md`	Announces complementary opt-in AI-source coverage.
`docs/codeql/reusables/supported-frameworks.rst`	Lists both MCP packages as supported frameworks.

_{Review effort: Max (auto)}

+from ThreatModelSource src
+where src.getThreatModel() = "ai" and exists(src.getLocation().getFile().getRelativePath())
+select src, src.getSourceType()


+            .getMember("server")
+            .getMember("fastmcp")
+            .getMember("server")
+            .getMember("Context"),


+    exists(Expr ann | ann = p.getAnnotation() |
+      isContextAnnotation(ann)


+    exists(Expr ann | ann = p.getAnnotation() |
+      isContextAnnotation(ann)
+      or
+      exists(Subscript sub | sub = ann | isContextAnnotation(sub.getValue()))


+    ) and
+    not (
+      handler.isMethod() and
+      not isStaticmethod(handler) and
+      p = handler.getArg(0)
+    ) and
+    not isMcpContextParameter(p)


+def other_decorator(fn):
+    return fn


Tito0015 requested review from a team as code owners October 4, 2026 05:43

github-actions Bot added the external-contribution label Oct 4, 2026

Alex-Hofer mentioned this pull request Oct 4, 2026

Python: model MCP server handler parameters (mcp, fastmcp) as remote sources #22702

Open

github-actions Bot added documentation Python labels Oct 4, 2026

Tito0015 force-pushed the feature/python-mcp-remote-sources branch from 2771182 to ec38c23 Compare October 5, 2026 19:46

github-advanced-security Bot found potential problems Oct 6, 2026

View reviewed changes

Comment thread python/ql/lib/semmle/python/frameworks/Mcp.qll

*/

private import python

private import semmle.python.dataflow.new.DataFlow

Tito0015 force-pushed the feature/python-mcp-remote-sources branch 4 times, most recently from c909145 to 1ae3030 Compare October 13, 2026 04:14

owen-mc mentioned this pull request Oct 14, 2026

Python: fix 3 queries to use sources for the active threat model #22822

Merged

Tito0015 force-pushed the feature/python-mcp-remote-sources branch from 410350c to b33c676 Compare October 15, 2026 19:00

Tito0015 changed the title ~~Python: Model MCP and FastMCP server handler parameters as remote sources (#22702)~~ Python: Model complex MCP AST handler parameters as remote sources (Complementary to #22703) Oct 15, 2026

Tito0015 force-pushed the feature/python-mcp-remote-sources branch from b33c676 to 86d2f8e Compare October 16, 2026 23:47

Tito0015 added 2 commits October 19, 2026 14:46

Python: model MCP and FastMCP server handler parameters as remote flo…

5e30bf2

…w sources

Python: restrict MCP remote sources to stacked and instance handlers

d0a798c

Tito0015 force-pushed the feature/python-mcp-remote-sources branch from 86d2f8e to d0a798c Compare October 19, 2026 18:46

Merge upstream/main into feature/python-mcp-remote-sources

40805f5

owen-mc requested changes Nov 4, 2026

View reviewed changes

Python: model MCP handler params as ai threat-model sources

99cdc47

Apply review feedback: ThreatModelSource kind ai, expanded Context API paths, isStaticmethod parameter handling, and staticmethod test coverage.

Copilot AI auto review requested due to automatic review settings November 4, 2026 16:54

Copilot started reviewing on behalf of Tito0015 November 4, 2026 16:55

View session

Merge upstream/main into feature/python-mcp-remote-sources

9caec6c

Copilot AI reviewed Nov 4, 2026

View reviewed changes

	private class McpHandlerParameter extends RemoteFlowSource::Range, DataFlow::ParameterNode {
	private class McpHandlerParameter extends ThreatModelSource::Range, DataFlow::ParameterNode {

		exists(Expr ann \| ann = p.getAnnotation() \|
		isContextAnnotation(ann)

Conversation

Tito0015 commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Context & Positioning

Testing

Uh oh!

Tito0015 commented Oct 5, 2026

Uh oh!

geoffw0 commented Oct 6, 2026

Uh oh!

owen-mc commented Oct 14, 2026

Uh oh!

Tito0015 commented Oct 14, 2026

Uh oh!

owen-mc commented Oct 15, 2026

Uh oh!

Tito0015 commented Oct 15, 2026

Uh oh!

owen-mc commented Oct 20, 2026

Uh oh!

Tito0015 commented Oct 20, 2026

Uh oh!

Tito0015 commented Nov 2, 2026

Uh oh!

owen-mc Nov 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

owen-mc Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

owen-mc Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

Tito0015 commented Nov 4, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Tito0015 commented Oct 4, 2026 •

edited

Loading

owen-mc Nov 4, 2026 •

edited

Loading