Skip to main content

agentshield/analysis/interprocedural/
types.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use once_cell::sync::Lazy;
5use regex::Regex;
6
7use crate::ir::data_surface::TaintSink;
8use crate::ir::{ScanTarget, SourceLocation};
9
10pub(crate) static CALL_EXPR_RE: Lazy<Regex> =
11    Lazy::new(|| Regex::new(r#"\b([A-Za-z0-9_]+)\s*\(([^)]*)\)"#).expect("valid regex"));
12
13/// A representation of a function definition in the project.
14#[derive(Debug, Clone)]
15pub struct FunctionNode {
16    pub name: String,
17    pub file_path: PathBuf,
18    pub params: Vec<String>,
19    pub start_line: usize,
20    pub end_line: usize,
21    pub location: SourceLocation,
22    pub sinks: Vec<TaintSink>,
23}
24
25/// A call-site invoking a function.
26#[derive(Debug, Clone)]
27pub struct CallSite {
28    pub caller_name: String,
29    pub callee_name: String,
30    pub file_path: PathBuf,
31    pub line_number: usize,
32    pub args: Vec<String>,
33    pub location: SourceLocation,
34}
35
36/// Interprocedural Call Graph.
37#[derive(Debug, Default)]
38pub struct CallGraph {
39    pub functions: HashMap<String, Vec<FunctionNode>>,
40    pub call_sites: Vec<CallSite>,
41}
42
43impl CallGraph {
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    /// Build call graph from all source files in a ScanTarget.
49    pub fn build(target: &ScanTarget) -> Self {
50        let mut graph = CallGraph::new();
51
52        for sf in &target.source_files {
53            let ext = sf.path.extension().and_then(|e| e.to_str()).unwrap_or("");
54            if ext == "py" {
55                super::python::parse_python_file(&mut graph, &sf.path, &sf.content, target);
56            } else if matches!(ext, "ts" | "js" | "tsx" | "jsx" | "mjs") {
57                super::typescript::parse_typescript_file(&mut graph, &sf.path, &sf.content, target);
58            }
59        }
60
61        graph
62    }
63
64    pub fn find_enclosing_function(&self, file_path: &Path, line: usize) -> String {
65        let mut best_match: Option<(&FunctionNode, usize)> = None;
66        for nodes in self.functions.values() {
67            for node in nodes {
68                let start = node.start_line.saturating_sub(1);
69                if node.file_path == file_path && line >= start && line <= node.end_line {
70                    let span = node.end_line - start;
71                    if best_match.is_none_or(|(_, best_span)| span < best_span) {
72                        best_match = Some((node, span));
73                    }
74                }
75            }
76        }
77        best_match
78            .map(|(n, _)| n.name.clone())
79            .unwrap_or_else(|| "<global>".to_string())
80    }
81}