#![allow(missing_docs)]
use crate::cfg::ControlFlowGraph;
use crate::errors::AnalysisResult;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use tracing::{debug, info};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ValueOrigin {
Constant { value: String },
Parameter { index: usize, tainted: bool },
StateVar { name: String, tainted: bool },
FunctionCall { function: String, tainted: bool },
ExternalCall { address: String },
Arithmetic {
operands: Vec<String>,
tainted: bool,
},
Untrusted { source: String },
Unknown,
}
impl ValueOrigin {
pub fn is_tainted(&self) -> bool {
matches!(
self,
ValueOrigin::Parameter { tainted: true, .. }
| ValueOrigin::StateVar { tainted: true, .. }
| ValueOrigin::FunctionCall { tainted: true, .. }
| ValueOrigin::ExternalCall { .. }
| ValueOrigin::Arithmetic { tainted: true, .. }
| ValueOrigin::Untrusted { .. }
)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DefUse {
pub variable: String,
pub definition: Location,
pub uses: Vec<Location>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Location {
pub block: usize,
pub statement_idx: usize,
pub line: Option<usize>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataFlow {
pub origins: HashMap<String, ValueOrigin>,
pub def_use_chains: Vec<DefUse>,
pub tainted_vars: HashSet<String>,
pub dependencies: HashMap<String, HashSet<String>>,
}
impl DataFlow {
pub fn new() -> Self {
Self {
origins: HashMap::new(),
def_use_chains: Vec::new(),
tainted_vars: HashSet::new(),
dependencies: HashMap::new(),
}
}
pub fn is_tainted(&self, var: &str) -> bool {
self.tainted_vars.contains(var)
}
pub fn transitive_dependents(&self, var: &str) -> HashSet<String> {
let mut result = HashSet::new();
let mut queue = vec![var.to_string()];
while let Some(current) = queue.pop() {
if !result.insert(current.clone()) {
continue;
}
for (dependent, depends_on) in &self.dependencies {
if depends_on.contains(¤t) {
queue.push(dependent.clone());
}
}
}
result.remove(var);
result
}
pub fn can_affect(&self, a: &str, b: &str) -> bool {
self.transitive_dependents(a).contains(b)
}
}
impl Default for DataFlow {
fn default() -> Self {
Self::new()
}
}
pub struct DataFlowAnalyzer;
impl DataFlowAnalyzer {
pub fn analyze(cfg: &ControlFlowGraph) -> AnalysisResult<DataFlow> {
info!("Performing data flow analysis");
let mut analysis = DataFlow::new();
let reaching_defs = Self::compute_reaching_definitions(cfg)?;
debug!(
"Computed reaching definitions for {} blocks",
reaching_defs.len()
);
let _live_vars = Self::compute_live_variables(cfg)?;
debug!("Computed live variables");
analysis = Self::extract_def_use_chains(cfg, analysis)?;
debug!("Extracted {} def-use chains", analysis.def_use_chains.len());
analysis = Self::identify_tainted(cfg, analysis)?;
debug!(
"Identified {} tainted variables",
analysis.tainted_vars.len()
);
analysis = Self::build_dependencies(cfg, analysis)?;
Ok(analysis)
}
fn compute_reaching_definitions(
cfg: &ControlFlowGraph,
) -> AnalysisResult<HashMap<usize, HashSet<String>>> {
let mut reaching = HashMap::new();
let blocks = vec![cfg.entry()];
for block in blocks {
reaching.insert(block, HashSet::new());
}
Ok(reaching)
}
fn compute_live_variables(
cfg: &ControlFlowGraph,
) -> AnalysisResult<HashMap<usize, HashSet<String>>> {
let mut live = HashMap::new();
let blocks = vec![cfg.entry()];
for block in blocks {
live.insert(block, HashSet::new());
}
Ok(live)
}
fn extract_def_use_chains(
_cfg: &ControlFlowGraph,
analysis: DataFlow,
) -> AnalysisResult<DataFlow> {
Ok(analysis)
}
fn identify_tainted(_cfg: &ControlFlowGraph, analysis: DataFlow) -> AnalysisResult<DataFlow> {
info!("Identifying tainted variables");
Ok(analysis)
}
fn build_dependencies(_cfg: &ControlFlowGraph, analysis: DataFlow) -> AnalysisResult<DataFlow> {
Ok(analysis)
}
}
pub struct TaintAnalyzer {
tainted_sources: HashSet<String>,
}
impl TaintAnalyzer {
pub fn new() -> Self {
let mut sources = HashSet::new();
sources.insert("msg.sender".to_string());
sources.insert("tx.origin".to_string());
sources.insert("calldata".to_string());
sources.insert("input".to_string());
Self {
tainted_sources: sources,
}
}
pub fn add_tainted_source(&mut self, source: String) {
self.tainted_sources.insert(source);
}
pub fn is_tainted(&self, origin: &ValueOrigin) -> bool {
origin.is_tainted()
|| (match origin {
ValueOrigin::FunctionCall { function, .. } => {
self.tainted_sources.contains(function)
}
_ => false,
})
}
pub fn propagate(&self, source_tainted: bool, _assignment_op: &str) -> bool {
source_tainted
}
}
impl Default for TaintAnalyzer {
fn default() -> Self {
Self::new()
}
}
pub struct UseDefAnalyzer;
impl UseDefAnalyzer {
pub fn compute_chains(_cfg: &ControlFlowGraph) -> AnalysisResult<HashMap<String, Vec<DefUse>>> {
let chains: HashMap<String, Vec<DefUse>> = HashMap::new();
Ok(chains)
}
pub fn find_uses(def_use: &DefUse) -> Vec<Location> {
def_use.uses.clone()
}
pub fn find_definition(def_uses: &[DefUse], var: &str) -> Option<Location> {
def_uses
.iter()
.find(|du| du.variable == var)
.map(|du| du.definition.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_value_origin_taint() {
let constant = ValueOrigin::Constant {
value: "42".to_string(),
};
assert!(!constant.is_tainted());
let untrusted = ValueOrigin::Untrusted {
source: "user_input".to_string(),
};
assert!(untrusted.is_tainted());
let external = ValueOrigin::ExternalCall {
address: "0x...".to_string(),
};
assert!(external.is_tainted());
}
#[test]
fn test_taint_analyzer_creation() {
let analyzer = TaintAnalyzer::new();
assert!(analyzer.tainted_sources.contains("msg.sender"));
assert!(analyzer.tainted_sources.contains("calldata"));
}
#[test]
fn test_data_flow_taint_tracking() {
let mut data_flow = DataFlow::new();
data_flow.tainted_vars.insert("user_input".to_string());
assert!(data_flow.is_tainted("user_input"));
assert!(!data_flow.is_tainted("safe_var"));
}
#[test]
fn test_data_flow_dependencies() {
let mut data_flow = DataFlow::new();
let mut deps = HashSet::new();
deps.insert("a".to_string());
deps.insert("b".to_string());
data_flow.dependencies.insert("c".to_string(), deps);
assert!(data_flow.dependencies.contains_key("c"));
assert_eq!(data_flow.dependencies["c"].len(), 2);
}
#[test]
fn test_transitive_dependents() {
let mut data_flow = DataFlow::new();
let mut deps_b = HashSet::new();
deps_b.insert("a".to_string());
data_flow.dependencies.insert("b".to_string(), deps_b);
let mut deps_c = HashSet::new();
deps_c.insert("b".to_string());
data_flow.dependencies.insert("c".to_string(), deps_c);
let dependents = data_flow.transitive_dependents("a");
assert!(dependents.contains("b"));
assert!(dependents.contains("c"));
}
}