agentshield/analysis/interprocedural/
propagate.rs

1use std::collections::HashSet;
2
3use crate::analysis::AnalysisBundle;
4use crate::ir::data_surface::{TaintPath, TaintSource};
5use crate::ir::{ScanTarget, SourceLocation};
6
7use super::types::CallGraph;
8
9pub const MAX_PROPAGATION_DEPTH: usize = 16;
10
11pub(crate) struct TraversalContext<'a> {
12    graph: &'a CallGraph,
13    source: &'a TaintSource,
14    trace: Vec<SourceLocation>,
15    visited_functions: HashSet<String>,
16    visited_paths: &'a mut HashSet<(
17        std::path::PathBuf,
18        usize,
19        std::path::PathBuf,
20        usize,
21        crate::ir::data_surface::TaintSinkType,
22    )>,
23    new_paths: &'a mut Vec<TaintPath>,
24    depth: usize,
25}
26
27impl<'a> TraversalContext<'a> {
28    fn trace_callee(&mut self, callee_name: &str) {
29        if self.depth > MAX_PROPAGATION_DEPTH || self.visited_functions.contains(callee_name) {
30            return;
31        }
32        self.visited_functions.insert(callee_name.to_string());
33
34        if let Some(nodes) = self.graph.functions.get(callee_name) {
35            for node in nodes {
36                self.trace.push(node.location.clone());
37
38                // Check if this callee contains execution sinks
39                for sink in &node.sinks {
40                    let path_key = (
41                        self.source.location.file.clone(),
42                        self.source.location.line,
43                        sink.location.file.clone(),
44                        sink.location.line,
45                        sink.sink_type,
46                    );
47                    if !self.visited_paths.contains(&path_key) {
48                        self.visited_paths.insert(path_key);
49                        self.new_paths.push(TaintPath {
50                            source: self.source.clone(),
51                            sink: sink.clone(),
52                            through: self.trace.clone(),
53                            confidence: 0.9,
54                        });
55                    }
56                }
57
58                // Recurse into calls made by this callee
59                for next_call in &self.graph.call_sites {
60                    if next_call.file_path == node.file_path && next_call.caller_name == node.name {
61                        self.trace.push(next_call.location.clone());
62                        self.depth += 1;
63                        self.trace_callee(&next_call.callee_name);
64                        self.depth -= 1;
65                        self.trace.pop();
66                    }
67                }
68
69                self.trace.pop();
70            }
71        }
72
73        self.visited_functions.remove(callee_name);
74    }
75}
76
77/// Propagates taint across call-graph edges and constructs multi-hop `TaintPath` instances.
78pub fn propagate_interprocedural_taint(target: &ScanTarget, graph: &CallGraph) -> Vec<TaintPath> {
79    let mut new_paths = Vec::new();
80    let mut visited_paths = HashSet::new();
81
82    // Map each taint source to its enclosing function
83    for source in &target.data.sources {
84        let caller_name =
85            graph.find_enclosing_function(&source.location.file, source.location.line);
86
87        // Find calls originating from this function passing a tainted variable
88        for call in &graph.call_sites {
89            if call.file_path == source.location.file && call.caller_name == caller_name {
90                let mut ctx = TraversalContext {
91                    graph,
92                    source,
93                    trace: vec![call.location.clone()],
94                    visited_functions: HashSet::from([caller_name.clone()]),
95                    visited_paths: &mut visited_paths,
96                    new_paths: &mut new_paths,
97                    depth: 0,
98                };
99                ctx.trace_callee(&call.callee_name);
100            }
101        }
102    }
103
104    new_paths
105}
106
107/// Analyze analysis bundles and enrich target DataSurface with interprocedural taint paths.
108pub(crate) fn analyze_and_enrich_targets(bundles: &mut [AnalysisBundle]) {
109    for bundle in bundles {
110        let call_graph = CallGraph::build(&bundle.target);
111        let interprocedural_paths = propagate_interprocedural_taint(&bundle.target, &call_graph);
112
113        for path in interprocedural_paths {
114            if !bundle.target.data.taint_paths.iter().any(|p| {
115                p.source.location.file == path.source.location.file
116                    && p.source.location.line == path.source.location.line
117                    && p.sink.location.file == path.sink.location.file
118                    && p.sink.location.line == path.sink.location.line
119            }) {
120                bundle.target.data.taint_paths.push(path);
121            }
122        }
123    }
124}
agentshield/analysis/interprocedural/propagate.rs

agentshield/analysis/interprocedural/
propagate.rs
