1pub mod propagate;
8pub mod python;
9pub mod types;
10pub mod typescript;
11
12pub(crate) use propagate::analyze_and_enrich_targets;
13pub use propagate::propagate_interprocedural_taint;
14pub use types::{CallGraph, CallSite, FunctionNode};
15
16#[cfg(test)]
17mod tests {
18 use super::*;
19 use crate::ir::data_surface::{
20 DataSurface, TaintSink, TaintSinkType, TaintSource, TaintSourceType,
21 };
22 use crate::ir::dependency_surface::DependencySurface;
23 use crate::ir::execution_surface::ExecutionSurface;
24 use crate::ir::provenance_surface::ProvenanceSurface;
25 use crate::ir::tool_surface::ToolSurface;
26 use crate::ir::{Language, ScanTarget, SourceLocation};
27 use std::path::PathBuf;
28
29 #[test]
30 fn test_cross_function_command_injection_propagation() {
31 let py_code = r#"
32def helper_run(cmd):
33 import subprocess
34 subprocess.run(cmd, shell=True)
35
36@mcp.tool()
37def execute_tool(user_query: str):
38 helper_run(user_query)
39"#;
40 let file_path = PathBuf::from("server.py");
41
42 let target = ScanTarget {
43 name: "test-mcp".into(),
44 framework: crate::ir::Framework::Mcp,
45 root_path: PathBuf::from("/test"),
46 source_files: vec![crate::ir::SourceFile {
47 path: file_path.clone(),
48 language: Language::Python,
49 content: py_code.into(),
50 size_bytes: py_code.len() as u64,
51 content_hash: "hash".into(),
52 }],
53 dependencies: DependencySurface::default(),
54 data: DataSurface {
55 sources: vec![TaintSource {
56 source_type: TaintSourceType::ToolArgument,
57 description: "Tool 'execute_tool' parameter 'user_query'".into(),
58 location: SourceLocation {
59 file: file_path.clone(),
60 line: 7,
61 column: 0,
62 end_line: None,
63 end_column: None,
64 },
65 }],
66 sinks: vec![TaintSink {
67 sink_type: TaintSinkType::ProcessExec,
68 description: "Process execution via subprocess.run".into(),
69 location: SourceLocation {
70 file: file_path.clone(),
71 line: 4,
72 column: 4,
73 end_line: None,
74 end_column: None,
75 },
76 }],
77 taint_paths: vec![],
78 },
79 tools: vec![ToolSurface {
80 name: "execute_tool".into(),
81 description: None,
82 input_schema: None,
83 output_schema: None,
84 declared_permissions: vec![],
85 defined_at: Some(SourceLocation {
86 file: file_path.clone(),
87 line: 7,
88 column: 0,
89 end_line: None,
90 end_column: None,
91 }),
92 declared_capabilities: std::collections::BTreeSet::new(),
93 capability_declarations: vec![],
94 observed_capabilities: std::collections::BTreeSet::new(),
95 capability_observation_complete: false,
96 capability_evidence: vec![],
97 }],
98 execution: ExecutionSurface::default(),
99 provenance: ProvenanceSurface::default(),
100 };
101
102 let graph = CallGraph::build(&target);
103 assert!(graph.functions.contains_key("helper_run"));
104 assert!(graph.functions.contains_key("execute_tool"));
105
106 let paths = propagate_interprocedural_taint(&target, &graph);
107 assert_eq!(paths.len(), 1);
108 assert_eq!(paths[0].source.source_type, TaintSourceType::ToolArgument);
109 assert_eq!(paths[0].sink.sink_type, TaintSinkType::ProcessExec);
110 assert!(!paths[0].through.is_empty());
111 }
112
113 #[test]
114 fn test_cross_function_typescript_ssrf_propagation() {
115 let ts_code = r#"
116async function sendHttpRequest(targetUrl: string) {
117 return await fetch(targetUrl);
118}
119
120export async function handleApiRequest(urlInput: string) {
121 return await sendHttpRequest(urlInput);
122}
123"#;
124 let file_path = PathBuf::from("index.ts");
125
126 let target = ScanTarget {
127 name: "test-ts-agent".into(),
128 framework: crate::ir::Framework::OpenClaw,
129 root_path: PathBuf::from("/test-ts"),
130 source_files: vec![crate::ir::SourceFile {
131 path: file_path.clone(),
132 language: Language::TypeScript,
133 content: ts_code.into(),
134 size_bytes: ts_code.len() as u64,
135 content_hash: "hash-ts".into(),
136 }],
137 dependencies: DependencySurface::default(),
138 data: DataSurface {
139 sources: vec![TaintSource {
140 source_type: TaintSourceType::ToolArgument,
141 description: "Tool 'handleApiRequest' parameter 'urlInput'".into(),
142 location: SourceLocation {
143 file: file_path.clone(),
144 line: 6,
145 column: 0,
146 end_line: None,
147 end_column: None,
148 },
149 }],
150 sinks: vec![TaintSink {
151 sink_type: TaintSinkType::HttpRequest,
152 description: "HTTP fetch request".into(),
153 location: SourceLocation {
154 file: file_path.clone(),
155 line: 3,
156 column: 4,
157 end_line: None,
158 end_column: None,
159 },
160 }],
161 taint_paths: vec![],
162 },
163 tools: vec![ToolSurface {
164 name: "handleApiRequest".into(),
165 description: None,
166 input_schema: None,
167 output_schema: None,
168 declared_permissions: vec![],
169 defined_at: Some(SourceLocation {
170 file: file_path.clone(),
171 line: 6,
172 column: 0,
173 end_line: None,
174 end_column: None,
175 }),
176 declared_capabilities: std::collections::BTreeSet::new(),
177 capability_declarations: vec![],
178 observed_capabilities: std::collections::BTreeSet::new(),
179 capability_observation_complete: false,
180 capability_evidence: vec![],
181 }],
182 execution: ExecutionSurface::default(),
183 provenance: ProvenanceSurface::default(),
184 };
185
186 let graph = CallGraph::build(&target);
187 assert!(graph.functions.contains_key("sendHttpRequest"));
188 assert!(graph.functions.contains_key("handleApiRequest"));
189
190 let paths = propagate_interprocedural_taint(&target, &graph);
191 assert_eq!(paths.len(), 1);
192 assert_eq!(paths[0].source.source_type, TaintSourceType::ToolArgument);
193 assert_eq!(paths[0].sink.sink_type, TaintSinkType::HttpRequest);
194 assert!(!paths[0].through.is_empty());
195 }
196
197 #[test]
198 fn test_allman_style_typescript_function_brace() {
199 let ts_code =
200 "function executeCommand(cmd: string)\n{\n return child_process.execSync(cmd);\n}\n";
201 let file_path = PathBuf::from("exec.ts");
202
203 let target = ScanTarget {
204 name: "test-allman-ts".into(),
205 framework: crate::ir::Framework::Mcp,
206 root_path: PathBuf::from("/test-allman"),
207 source_files: vec![crate::ir::SourceFile {
208 path: file_path.clone(),
209 language: Language::TypeScript,
210 content: ts_code.into(),
211 size_bytes: ts_code.len() as u64,
212 content_hash: "hash-allman".into(),
213 }],
214 dependencies: DependencySurface::default(),
215 data: DataSurface {
216 sources: vec![TaintSource {
217 source_type: TaintSourceType::ToolArgument,
218 description: "Tool parameter 'cmd'".into(),
219 location: SourceLocation {
220 file: file_path.clone(),
221 line: 1,
222 column: 0,
223 end_line: None,
224 end_column: None,
225 },
226 }],
227 sinks: vec![TaintSink {
228 sink_type: TaintSinkType::ProcessExec,
229 description: "ExecSync command".into(),
230 location: SourceLocation {
231 file: file_path.clone(),
232 line: 3,
233 column: 4,
234 end_line: None,
235 end_column: None,
236 },
237 }],
238 taint_paths: vec![],
239 },
240 tools: vec![],
241 execution: ExecutionSurface::default(),
242 provenance: ProvenanceSurface::default(),
243 };
244
245 let graph = CallGraph::build(&target);
246 assert!(graph.functions.contains_key("executeCommand"));
247 let node = &graph.functions["executeCommand"][0];
248 assert_eq!(
249 node.sinks.len(),
250 1,
251 "Sink on line 3 must be captured inside Allman brace function boundary"
252 );
253 }
254
255 #[test]
256 fn test_multi_file_line_collision_preserves_distinct_sinks() {
257 let py_code_a = "def run_task(q):\n return helper(q)\n";
258 let py_code_b = "def helper(q):\n subprocess.run(q, shell=True)\n";
259 let file_a = PathBuf::from("pkg/a.py");
260 let file_b = PathBuf::from("pkg/b.py");
261
262 let target = ScanTarget {
263 name: "test-multi-file-collision".into(),
264 framework: crate::ir::Framework::Mcp,
265 root_path: PathBuf::from("/test-multi"),
266 source_files: vec![
267 crate::ir::SourceFile {
268 path: file_a.clone(),
269 language: Language::Python,
270 content: py_code_a.into(),
271 size_bytes: py_code_a.len() as u64,
272 content_hash: "hash-a".into(),
273 },
274 crate::ir::SourceFile {
275 path: file_b.clone(),
276 language: Language::Python,
277 content: py_code_b.into(),
278 size_bytes: py_code_b.len() as u64,
279 content_hash: "hash-b".into(),
280 },
281 ],
282 dependencies: DependencySurface::default(),
283 data: DataSurface {
284 sources: vec![TaintSource {
285 source_type: TaintSourceType::ToolArgument,
286 description: "Param 'q'".into(),
287 location: SourceLocation {
288 file: file_a.clone(),
289 line: 1,
290 column: 0,
291 end_line: None,
292 end_column: None,
293 },
294 }],
295 sinks: vec![TaintSink {
296 sink_type: TaintSinkType::ProcessExec,
297 description: "Subprocess sink".into(),
298 location: SourceLocation {
299 file: file_b.clone(),
300 line: 2,
301 column: 4,
302 end_line: None,
303 end_column: None,
304 },
305 }],
306 taint_paths: vec![],
307 },
308 tools: vec![],
309 execution: ExecutionSurface::default(),
310 provenance: ProvenanceSurface::default(),
311 };
312
313 let graph = CallGraph::build(&target);
314 let paths = propagate_interprocedural_taint(&target, &graph);
315 assert_eq!(paths.len(), 1);
316 assert_eq!(paths[0].source.location.file, file_a);
317 assert_eq!(paths[0].sink.location.file, file_b);
318 }
319
320 #[test]
321 fn test_typed_arrow_function_and_no_spurious_call_sites() {
322 let ts_code = r#"
323const helperExec = async (cmd: string): Promise<void> => {
324 child_process.execSync(cmd);
325};
326
327export async function runTool(userInput: string) {
328 return await helperExec(userInput);
329}
330"#;
331 let file_path = PathBuf::from("index.ts");
332 let target = ScanTarget {
333 name: "test-ts-arrow".into(),
334 framework: crate::ir::Framework::Mcp,
335 root_path: PathBuf::from("/test-ts-arrow"),
336 source_files: vec![crate::ir::SourceFile {
337 path: file_path.clone(),
338 language: Language::TypeScript,
339 content: ts_code.into(),
340 size_bytes: ts_code.len() as u64,
341 content_hash: "hash-arrow".into(),
342 }],
343 dependencies: DependencySurface::default(),
344 data: DataSurface {
345 sources: vec![TaintSource {
346 source_type: TaintSourceType::ToolArgument,
347 description: "Tool param 'userInput'".into(),
348 location: SourceLocation {
349 file: file_path.clone(),
350 line: 6,
351 column: 0,
352 end_line: None,
353 end_column: None,
354 },
355 }],
356 sinks: vec![TaintSink {
357 sink_type: TaintSinkType::ProcessExec,
358 description: "ExecSync sink".into(),
359 location: SourceLocation {
360 file: file_path.clone(),
361 line: 3,
362 column: 4,
363 end_line: None,
364 end_column: None,
365 },
366 }],
367 taint_paths: vec![],
368 },
369 tools: vec![],
370 execution: ExecutionSurface::default(),
371 provenance: ProvenanceSurface::default(),
372 };
373
374 let graph = CallGraph::build(&target);
375 assert!(
376 graph.functions.contains_key("helperExec"),
377 "Typed arrow function must be captured in CallGraph"
378 );
379 assert!(graph.functions.contains_key("runTool"));
380
381 let self_calls = graph
383 .call_sites
384 .iter()
385 .filter(|c| c.caller_name == "runTool" && c.callee_name == "runTool")
386 .count();
387 assert_eq!(
388 self_calls, 0,
389 "Function signature must not create spurious self-call"
390 );
391
392 let paths = propagate_interprocedural_taint(&target, &graph);
393 assert_eq!(paths.len(), 1);
394 assert_eq!(paths[0].sink.sink_type, TaintSinkType::ProcessExec);
395 }
396}