|
1
|
""" |
|
2
|
MCP server security — query validation and complexity checks. |
|
3
|
|
|
4
|
Two layers of protection: |
|
5
|
1. validate_cypher — blocks write operations and injection patterns |
|
6
|
2. check_complexity — enforces depth and result-set size limits |
|
7
|
""" |
|
8
|
|
|
9
|
from __future__ import annotations |
|
10
|
|
|
11
|
import re |
|
12
|
|
|
13
|
|
|
14
|
class QueryValidationError(Exception): |
|
15
|
"""Raised when a Cypher query contains a disallowed pattern.""" |
|
16
|
|
|
17
|
|
|
18
|
class QueryComplexityError(Exception): |
|
19
|
"""Raised when a Cypher query exceeds complexity limits.""" |
|
20
|
|
|
21
|
|
|
22
|
# ── Write-operation keywords ─────────────────────────────────────────────────── |
|
23
|
|
|
24
|
_WRITE_KEYWORDS: tuple[str, ...] = ( |
|
25
|
"CREATE", |
|
26
|
"MERGE", |
|
27
|
"SET", |
|
28
|
"DELETE", |
|
29
|
"REMOVE", |
|
30
|
"DROP", |
|
31
|
) |
|
32
|
|
|
33
|
# ── Injection patterns ──────────────────────────────────────────────────────── |
|
34
|
|
|
35
|
# CALL procedure invocations (e.g. CALL db.labels()) |
|
36
|
_CALL_RE = re.compile(r"\bCALL\b", re.IGNORECASE) |
|
37
|
|
|
38
|
# Nested / sub-queries introduced by { ... } preceded by a Cypher keyword. |
|
39
|
# We detect the presence of balanced braces that follow MATCH/WITH/WHERE/RETURN |
|
40
|
# as a heuristic for sub-query injection. |
|
41
|
_SUBQUERY_RE = re.compile( |
|
42
|
r"\b(?:MATCH|WITH|WHERE|RETURN)\s*\{", |
|
43
|
re.IGNORECASE, |
|
44
|
) |
|
45
|
|
|
46
|
# ── Variable-length path pattern ───────────────────────────────────────────── |
|
47
|
|
|
48
|
# Matches relationship depth specifiers like *1..100 or *..50 or *5.. |
|
49
|
# Group 1 = lower bound (may be empty), Group 2 = upper bound (may be empty). |
|
50
|
_VARLEN_RE = re.compile(r"\*(\d*)\.\.((\d*))|\*(\d+)") |
|
51
|
|
|
52
|
|
|
53
|
def validate_cypher(query: str) -> None: |
|
54
|
""" |
|
55
|
Validate *query* for dangerous or disallowed patterns. |
|
56
|
|
|
57
|
Raises: |
|
58
|
QueryValidationError: if the query contains write operations or |
|
59
|
injection patterns (CALL, nested sub-queries). |
|
60
|
""" |
|
61
|
# Strip single-line comments before analysis |
|
62
|
stripped = re.sub(r"//[^\n]*", "", query) |
|
63
|
|
|
64
|
upper = stripped.upper() |
|
65
|
|
|
66
|
# Check for write-operation keywords as whole words |
|
67
|
for kw in _WRITE_KEYWORDS: |
|
68
|
pattern = re.compile(rf"\b{kw}\b") |
|
69
|
if pattern.search(upper): |
|
70
|
raise QueryValidationError(f"Write operation '{kw}' is not allowed in read-only mode.") |
|
71
|
|
|
72
|
# Check for CALL procedure injection |
|
73
|
if _CALL_RE.search(stripped): |
|
74
|
raise QueryValidationError("CALL procedures are not allowed in read-only mode.") |
|
75
|
|
|
76
|
# Check for nested / sub-query patterns |
|
77
|
if _SUBQUERY_RE.search(stripped): |
|
78
|
raise QueryValidationError("Nested sub-queries are not allowed in read-only mode.") |
|
79
|
|
|
80
|
|
|
81
|
def check_complexity( |
|
82
|
query: str, |
|
83
|
max_depth: int = 5, |
|
84
|
max_results: int = 1000, |
|
85
|
) -> None: |
|
86
|
""" |
|
87
|
Check *query* for complexity issues. |
|
88
|
|
|
89
|
Raises: |
|
90
|
QueryComplexityError: if variable-length paths exceed *max_depth* or |
|
91
|
if the query has no LIMIT and could be unbounded. |
|
92
|
""" |
|
93
|
# ── 1. Variable-length path depth check ─────────────────────────────────── |
|
94
|
for m in _VARLEN_RE.finditer(query): |
|
95
|
# Pattern: *lower..upper |
|
96
|
if m.group(1) is not None or m.group(2) is not None: |
|
97
|
upper_str = m.group(2) # may be empty string (open-ended) |
|
98
|
if upper_str == "": |
|
99
|
# Open-ended: *1.. — treat as unbounded |
|
100
|
raise QueryComplexityError( |
|
101
|
f"Variable-length path with no upper bound is not allowed " |
|
102
|
f"(max depth: {max_depth})." |
|
103
|
) |
|
104
|
upper_val = int(upper_str) |
|
105
|
if upper_val > max_depth: |
|
106
|
raise QueryComplexityError( |
|
107
|
f"Variable-length path depth {upper_val} exceeds maximum " |
|
108
|
f"allowed depth of {max_depth}." |
|
109
|
) |
|
110
|
else: |
|
111
|
# Pattern: *N (exact repetition — check against max_depth) |
|
112
|
exact_str = m.group(4) |
|
113
|
if exact_str: |
|
114
|
exact_val = int(exact_str) |
|
115
|
if exact_val > max_depth: |
|
116
|
raise QueryComplexityError( |
|
117
|
f"Variable-length path depth {exact_val} exceeds maximum " |
|
118
|
f"allowed depth of {max_depth}." |
|
119
|
) |
|
120
|
|
|
121
|
# ── 2. Unbounded result check ────────────────────────────────────────────── |
|
122
|
# Queries that contain MATCH/RETURN but no LIMIT clause may return huge |
|
123
|
# result sets. Skip the check for queries that look purely structural |
|
124
|
# (e.g. COUNT aggregations that are inherently bounded). |
|
125
|
upper = query.upper() |
|
126
|
has_match = bool(re.search(r"\bMATCH\b", upper)) |
|
127
|
has_return = bool(re.search(r"\bRETURN\b", upper)) |
|
128
|
has_limit = bool(re.search(r"\bLIMIT\b", upper)) |
|
129
|
has_count = bool(re.search(r"\bCOUNT\s*\(", upper)) |
|
130
|
|
|
131
|
if has_match and has_return and not has_limit and not has_count: |
|
132
|
raise QueryComplexityError( |
|
133
|
f"Query has no LIMIT clause and could return unbounded results " |
|
134
|
f"(max: {max_results}). Add 'LIMIT {max_results}' or fewer." |
|
135
|
) |
|
136
|
|