Skip to content

MCP Surface Extraction

Overview

The surface extraction system (internal/surface/) identifies the MCP-specific attack surface in analyzed code. It detects MCP tools, resources, prompts, transport types, and authentication signals from parsed AST files.

Architecture

┌─────────────────────────────────────────────────────────────┐
│                   Surface Extractor                          │
├─────────────────────────────────────────────────────────────┤
│                                                               │
│  ┌───────────┐    ┌───────────┐    ┌─────────────────────┐  │
│  │ AST Files │───▶│ Extractor │───▶│    MCPSurface       │  │
│  │           │    │           │    │ - Tools             │  │
│  └───────────┘    └───────────┘    │ - Resources         │  │
│                        │           │ - Prompts           │  │
│                        │           │ - Transport         │  │
│                        ▼           │ - AuthSignals       │  │
│                  ┌───────────┐     └─────────────────────┘  │
│                  │SDK Detect │                               │
│                  └───────────┘                               │
│                                                               │
└─────────────────────────────────────────────────────────────┘

Core Data Structures

MCPSurface

The main structure containing all detected MCP components:

type MCPSurface struct {
    Tools       []Tool        // MCP tool definitions
    Resources   []Resource    // MCP resource definitions
    Prompts     []Prompt      // MCP prompt definitions
    Transport   TransportType // Detected transport type
    AuthSignals []AuthSignal  // Authentication patterns
}

Tool

Represents a detected MCP tool:

type Tool struct {
    Name        string                 // Tool name
    Description string                 // Tool description
    Schema      map[string]interface{} // Input schema (JSON Schema)
    Handler     *HandlerRef            // Reference to handler function
    Decorators  []string               // Applied decorators
    Location    types.Location         // Source location
    Parameters  []Parameter            // Tool parameters
}

HandlerRef

Reference to a function that handles an MCP component:

type HandlerRef struct {
    FunctionName string         // Function name
    FileName     string         // Source file path
    Location     types.Location // Definition location
}

Resource

Represents a detected MCP resource:

type Resource struct {
    Name        string
    Description string
    Handler     *HandlerRef
    Location    types.Location
}

Prompt

Represents a detected MCP prompt:

type Prompt struct {
    Name        string
    Description string
    Handler     *HandlerRef
    Location    types.Location
}

AuthSignal

Detected authentication-related patterns:

type AuthSignal struct {
    Type     string         // "cookie", "header", "oauth_state"
    Name     string         // Function/variable name
    Location types.Location // Source location
}

Transport Types

const (
    TransportSTDIO     TransportType = "stdio"     // Standard I/O
    TransportHTTP      TransportType = "http"      // HTTP server
    TransportWebSocket TransportType = "websocket" // WebSocket
    TransportUnknown   TransportType = "unknown"   // Undetermined
)

SDK Detection

The extractor detects which MCP SDK is being used:

SDK Type Import Pattern Language
Python Official mcp, mcp.* Python
FastMCP fastmcp, fastmcp.* Python
TypeScript Official @modelcontextprotocol/sdk TypeScript
type sdkType int

const (
    sdkUnknown sdkType = iota
    sdkPythonOfficial
    sdkPythonFastMCP
    sdkTypeScriptOfficial
)

Tool Detection

Decorator-Based Detection

Tools are primarily detected via decorators:

# Python Official SDK
@server.tool()
def search_files(query: str):
    """Search for files matching the query."""
    pass

# FastMCP
@mcp.tool()
def execute_command(command: str):
    pass

# Generic tool decorator
@tool
def custom_tool(param: str):
    pass

Recognized decorators: - tool - server.tool - mcp.tool - app.tool - fastmcp.tool

Heuristic-Based Detection

Functions matching naming conventions are flagged as potential tools:

Prefixes: - handle_* - tool_* - execute_* - run_* - do_*

Suffixes: - *_tool - *_handler - *_action - *_command

# Detected by heuristics
def handle_search(query: str):
    pass

def execute_command(cmd: str):
    pass

def file_search_tool(pattern: str):
    pass

Transport Detection

Import-Based Detection

Import Pattern Transport
Contains stdio STDIO
Contains http, express, fastapi, flask HTTP
Contains websocket, ws WebSocket

Code Pattern Detection

# Detected as STDIO
server.run_stdio()

# Detected as HTTP
app.listen(8080)
server.serve()

Authentication Signal Detection

The extractor identifies authentication patterns:

response.set_cookie("session", value)  # Detected
set_cookie("auth_token", token)        # Detected

Header Detection

request.headers["Authorization"]  # Detected
set_header("X-API-Key", key)     # Detected

OAuth State Detection

oauth_state = generate_state()  # Detected
nonce = create_nonce()          # Detected

API Usage

Basic Extraction

extractor := surface.New()
files := []*ast.File{...}

surface, err := extractor.Extract(files)
if err != nil {
    log.Fatal(err)
}

// Access extracted components
for _, tool := range surface.Tools {
    fmt.Printf("Tool: %s at %s:%d\n",
        tool.Name,
        tool.Handler.FileName,
        tool.Location.StartLine)
}

Integration with Taint Engine

// Extract surface for taint analysis
extractor := surface.New()
surface, _ := extractor.Extract(files)

// Use surface to identify tool handlers
engine := taint.New(catalog, cfg)
findings := engine.Analyze(files, surface)

Integration with Pattern Engine

// Use surface for prompt injection detection
surface, _ := extractor.Extract(files)

patternEngine := pattern.New()
findings := patternEngine.Analyze(files, surface)
// Pattern rules can inspect tool descriptions, etc.

Tool Parameter Extraction

When a tool is detected, its parameters are extracted from the function signature:

@server.tool()
def search_files(
    query: str,           # Parameter: name=query, type=str
    max_results: int = 10 # Parameter: name=max_results, type=int, default=10
):
    pass

Parameter structure:

type Parameter struct {
    Name     string
    Type     string
    Default  interface{}
    Required bool
    Location types.Location
}

Description Extraction

Tool descriptions are extracted from:

  1. Decorator arguments:

    @server.tool("Search for files in the filesystem")
    def search_files(query: str):
        pass
    

  2. Docstrings:

    @server.tool()
    def search_files(query: str):
        """Search for files matching the query pattern."""
        pass
    

Multi-File Extraction

The extractor processes multiple files and aggregates results:

files := []*ast.File{
    file1,  // Contains tool A
    file2,  // Contains tools B, C
    file3,  // Contains resource D
}

surface, _ := extractor.Extract(files)
// surface.Tools contains A, B, C
// surface.Resources contains D

Surface Analysis for Security

The extracted surface is used for security analysis:

Tool Input Tainting

All tool parameters are marked as tainted:

if e.isToolHandler(fn, surface) {
    for _, param := range fn.Parameters {
        state.SetTaint(param.Name, &TaintInfo{
            Source:     fn.Location,
            SourceType: types.SourceToolInput,
        })
    }
}

Prompt Injection Detection

Tool descriptions are scanned for injection patterns:

for _, tool := range surface.Tools {
    if containsInjectionMarker(tool.Description) {
        // Report finding
    }
}

Transport Security Analysis

Transport type affects security recommendations:

Transport Security Considerations
STDIO Local only, lower risk
HTTP Need TLS, authentication
WebSocket Need secure handshake, authentication

Configuration

# mcp-scan.yaml
surface:
  # Enable heuristic tool detection
  heuristic_detection: true

  # Custom tool decorators
  tool_decorators:
    - "custom.tool"
    - "mylib.register_tool"

  # Tool name patterns (regex)
  tool_patterns:
    - "^handle_.*"
    - ".*_handler$"

Extending the Extractor

Adding Custom Tool Decorators

func (e *Extractor) isToolDecorator(name string, sdk sdkType) bool {
    toolDecorators := []string{
        "tool",
        "server.tool",
        "mcp.tool",
        // Add custom decorators
        "custom.tool",
        "mylib.register",
    }
    // ...
}

Adding Transport Detection

func (e *Extractor) detectTransport(file *ast.File) TransportType {
    for _, imp := range file.Imports {
        // Add custom transport detection
        if strings.Contains(imp.Module, "grpc") {
            return TransportGRPC // Custom transport type
        }
    }
    // ...
}

Example: Complete Surface

# server.py
from fastmcp import FastMCP

mcp = FastMCP()

@mcp.tool("Search files by pattern")
def search_files(pattern: str, limit: int = 10):
    """Search for files matching the given pattern."""
    pass

@mcp.tool("Execute a shell command")
def run_command(command: str):
    pass

@mcp.resource("/files/{path}")
def get_file(path: str):
    pass

# Run with stdio transport
mcp.run_stdio()

Extracted Surface:

Tools:
  - Name: search_files
    Description: "Search files by pattern"
    Handler:
      FunctionName: search_files
      FileName: server.py
      Location: {StartLine: 6}
    Parameters:
      - Name: pattern
        Type: str
        Required: true
      - Name: limit
        Type: int
        Required: false
        Default: 10

  - Name: run_command
    Description: "Execute a shell command"
    Handler:
      FunctionName: run_command
      FileName: server.py
      Location: {StartLine: 11}
    Parameters:
      - Name: command
        Type: str
        Required: true

Resources:
  - Name: /files/{path}
    Handler:
      FunctionName: get_file
      FileName: server.py
      Location: {StartLine: 15}

Transport: stdio

AuthSignals: []