1#![allow(missing_docs)]
2use crate::cfg::ControlFlowGraph;
11use crate::errors::AnalysisResult;
12use serde::{Deserialize, Serialize};
13use std::collections::{HashMap, HashSet};
14use tracing::{debug, info};
15
16#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub enum ValueOrigin {
19 Constant { value: String },
21 Parameter { index: usize, tainted: bool },
23 StateVar { name: String, tainted: bool },
25 FunctionCall { function: String, tainted: bool },
27 ExternalCall { address: String },
29 Arithmetic {
31 operands: Vec<String>,
32 tainted: bool,
33 },
34 Untrusted { source: String },
36 Unknown,
38}
39
40impl ValueOrigin {
41 pub fn is_tainted(&self) -> bool {
43 matches!(
44 self,
45 ValueOrigin::Parameter { tainted: true, .. }
46 | ValueOrigin::StateVar { tainted: true, .. }
47 | ValueOrigin::FunctionCall { tainted: true, .. }
48 | ValueOrigin::ExternalCall { .. }
49 | ValueOrigin::Arithmetic { tainted: true, .. }
50 | ValueOrigin::Untrusted { .. }
51 )
52 }
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct DefUse {
58 pub variable: String,
60 pub definition: Location,
62 pub uses: Vec<Location>,
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct Location {
69 pub block: usize,
71 pub statement_idx: usize,
73 pub line: Option<usize>,
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct DataFlow {
80 pub origins: HashMap<String, ValueOrigin>,
82 pub def_use_chains: Vec<DefUse>,
84 pub tainted_vars: HashSet<String>,
86 pub dependencies: HashMap<String, HashSet<String>>,
88}
89
90impl DataFlow {
91 pub fn new() -> Self {
93 Self {
94 origins: HashMap::new(),
95 def_use_chains: Vec::new(),
96 tainted_vars: HashSet::new(),
97 dependencies: HashMap::new(),
98 }
99 }
100
101 pub fn is_tainted(&self, var: &str) -> bool {
103 self.tainted_vars.contains(var)
104 }
105
106 pub fn transitive_dependents(&self, var: &str) -> HashSet<String> {
108 let mut result = HashSet::new();
109 let mut queue = vec![var.to_string()];
110
111 while let Some(current) = queue.pop() {
112 if !result.insert(current.clone()) {
113 continue;
114 }
115
116 for (dependent, depends_on) in &self.dependencies {
118 if depends_on.contains(¤t) {
119 queue.push(dependent.clone());
120 }
121 }
122 }
123
124 result.remove(var);
125 result
126 }
127
128 pub fn can_affect(&self, a: &str, b: &str) -> bool {
130 self.transitive_dependents(a).contains(b)
131 }
132}
133
134impl Default for DataFlow {
135 fn default() -> Self {
136 Self::new()
137 }
138}
139
140pub struct DataFlowAnalyzer;
142
143impl DataFlowAnalyzer {
144 pub fn analyze(cfg: &ControlFlowGraph) -> AnalysisResult<DataFlow> {
146 info!("Performing data flow analysis");
147
148 let mut analysis = DataFlow::new();
149
150 let reaching_defs = Self::compute_reaching_definitions(cfg)?;
152 debug!(
153 "Computed reaching definitions for {} blocks",
154 reaching_defs.len()
155 );
156
157 let _live_vars = Self::compute_live_variables(cfg)?;
159 debug!("Computed live variables");
160
161 analysis = Self::extract_def_use_chains(cfg, analysis)?;
163 debug!("Extracted {} def-use chains", analysis.def_use_chains.len());
164
165 analysis = Self::identify_tainted(cfg, analysis)?;
167 debug!(
168 "Identified {} tainted variables",
169 analysis.tainted_vars.len()
170 );
171
172 analysis = Self::build_dependencies(cfg, analysis)?;
174
175 Ok(analysis)
176 }
177
178 fn compute_reaching_definitions(
180 cfg: &ControlFlowGraph,
181 ) -> AnalysisResult<HashMap<usize, HashSet<String>>> {
182 let mut reaching = HashMap::new();
183
184 let blocks = vec![cfg.entry()];
187 for block in blocks {
188 reaching.insert(block, HashSet::new());
189 }
190
191 Ok(reaching)
192 }
193
194 fn compute_live_variables(
196 cfg: &ControlFlowGraph,
197 ) -> AnalysisResult<HashMap<usize, HashSet<String>>> {
198 let mut live = HashMap::new();
199
200 let blocks = vec![cfg.entry()];
203 for block in blocks {
204 live.insert(block, HashSet::new());
205 }
206
207 Ok(live)
208 }
209
210 fn extract_def_use_chains(
212 _cfg: &ControlFlowGraph,
213 analysis: DataFlow,
214 ) -> AnalysisResult<DataFlow> {
215 Ok(analysis)
218 }
219
220 fn identify_tainted(_cfg: &ControlFlowGraph, analysis: DataFlow) -> AnalysisResult<DataFlow> {
222 info!("Identifying tainted variables");
229
230 Ok(analysis)
233 }
234
235 fn build_dependencies(_cfg: &ControlFlowGraph, analysis: DataFlow) -> AnalysisResult<DataFlow> {
237 Ok(analysis)
240 }
241}
242
243pub struct TaintAnalyzer {
245 tainted_sources: HashSet<String>,
247}
248
249impl TaintAnalyzer {
250 pub fn new() -> Self {
252 let mut sources = HashSet::new();
253 sources.insert("msg.sender".to_string());
255 sources.insert("tx.origin".to_string());
256 sources.insert("calldata".to_string());
257 sources.insert("input".to_string());
258
259 Self {
260 tainted_sources: sources,
261 }
262 }
263
264 pub fn add_tainted_source(&mut self, source: String) {
266 self.tainted_sources.insert(source);
267 }
268
269 pub fn is_tainted(&self, origin: &ValueOrigin) -> bool {
271 origin.is_tainted()
272 || (match origin {
273 ValueOrigin::FunctionCall { function, .. } => {
274 self.tainted_sources.contains(function)
275 }
276 _ => false,
277 })
278 }
279
280 pub fn propagate(&self, source_tainted: bool, _assignment_op: &str) -> bool {
282 source_tainted
285 }
286}
287
288impl Default for TaintAnalyzer {
289 fn default() -> Self {
290 Self::new()
291 }
292}
293
294pub struct UseDefAnalyzer;
296
297impl UseDefAnalyzer {
298 pub fn compute_chains(_cfg: &ControlFlowGraph) -> AnalysisResult<HashMap<String, Vec<DefUse>>> {
300 let chains: HashMap<String, Vec<DefUse>> = HashMap::new();
301
302 Ok(chains)
305 }
306
307 pub fn find_uses(def_use: &DefUse) -> Vec<Location> {
309 def_use.uses.clone()
310 }
311
312 pub fn find_definition(def_uses: &[DefUse], var: &str) -> Option<Location> {
314 def_uses
315 .iter()
316 .find(|du| du.variable == var)
317 .map(|du| du.definition.clone())
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 #[test]
326 fn test_value_origin_taint() {
327 let constant = ValueOrigin::Constant {
328 value: "42".to_string(),
329 };
330 assert!(!constant.is_tainted());
331
332 let untrusted = ValueOrigin::Untrusted {
333 source: "user_input".to_string(),
334 };
335 assert!(untrusted.is_tainted());
336
337 let external = ValueOrigin::ExternalCall {
338 address: "0x...".to_string(),
339 };
340 assert!(external.is_tainted());
341 }
342
343 #[test]
344 fn test_taint_analyzer_creation() {
345 let analyzer = TaintAnalyzer::new();
346 assert!(analyzer.tainted_sources.contains("msg.sender"));
347 assert!(analyzer.tainted_sources.contains("calldata"));
348 }
349
350 #[test]
351 fn test_data_flow_taint_tracking() {
352 let mut data_flow = DataFlow::new();
353 data_flow.tainted_vars.insert("user_input".to_string());
354
355 assert!(data_flow.is_tainted("user_input"));
356 assert!(!data_flow.is_tainted("safe_var"));
357 }
358
359 #[test]
360 fn test_data_flow_dependencies() {
361 let mut data_flow = DataFlow::new();
362 let mut deps = HashSet::new();
363 deps.insert("a".to_string());
364 deps.insert("b".to_string());
365
366 data_flow.dependencies.insert("c".to_string(), deps);
367
368 assert!(data_flow.dependencies.contains_key("c"));
369 assert_eq!(data_flow.dependencies["c"].len(), 2);
370 }
371
372 #[test]
373 fn test_transitive_dependents() {
374 let mut data_flow = DataFlow::new();
375
376 let mut deps_b = HashSet::new();
377 deps_b.insert("a".to_string());
378 data_flow.dependencies.insert("b".to_string(), deps_b);
379
380 let mut deps_c = HashSet::new();
381 deps_c.insert("b".to_string());
382 data_flow.dependencies.insert("c".to_string(), deps_c);
383
384 let dependents = data_flow.transitive_dependents("a");
385 assert!(dependents.contains("b"));
386 assert!(dependents.contains("c"));
387 }
388}