feat: improved data flow analysis with move semantics

untrusted_value_taint_checker/Cargo.toml

@@ -14,6 +14,7 @@ tempfile = "3.10.1"
 serde = { version = "1.0.204", features = ["derive"] }
 semver = "1.0.23"
 petgraph = "0.6.5"
+itertools = "0.13.0"
 #owo-colors = { version = "4", features = ["supports-colors"] }
 #cargo_metadata = "0.18.1"
 #bpaf = { version = "0.9.12", features = ["bpaf_derive", "autocomplete"] }

untrusted_value_taint_checker/sample/src/main.rs

@@ -20,6 +20,60 @@ fn fails() {
     println!("{:?}", insecure_env)
 }
 
+#[no_mangle]
+fn complicated_nonsense_function(mut arg: String) -> String {
+    // dont mind the horrible code, used to check the data flow algorithm ;)
+
+    let count = arg.chars().filter(|c| *c == '.').count();
+
+    loop {
+        for i in 0..arg.len() {
+            if let Some(index) = arg.find(".") {
+                arg.remove(index);
+                arg.push('+');
+            }
+            let ith = arg.chars().nth(i).unwrap();
+            if ith == '#' {
+                return arg.to_lowercase();
+            } else if ith == '-' {
+                arg = arg.to_uppercase();
+            } else if ith == '_' {
+                arg = arg.to_lowercase();
+            }
+        }
+        if !arg.contains(".") {
+            break;
+        }
+    }
+
+    arg.repeat(count)
+}
+
+#[no_mangle]
+fn nested_function(mut input: usize) -> usize {
+    #[no_mangle]
+    fn inner_function(unsafe_arg: &mut usize) {
+        *unsafe_arg = 10;
+    }
+
+    inner_function(&mut input);
+
+    input
+}
+
+struct ABC {
+    value: i32,
+}
+#[no_mangle]
+fn test_rename(q: ABC) -> ABC {
+    let mut x = q;
+    let mut y = x;
+    y.value = 10;
+    let z = y;
+
+    z
+}
+
 fn main() {
     works();
     fails();

untrusted_value_taint_checker/src/analysis/callbacks.rs

@@ -1,12 +1,15 @@
 use std::path::PathBuf;
 
-use petgraph::dot::{Config, Dot};
+use crate::analysis::taint_source::TaintSource;
+use petgraph::dot::Dot;
 use rustc_driver::Compilation;
-use rustc_middle::{mir::{visit::Visitor as VisitorMir}, ty::{TyCtxt}};
+use rustc_middle::{mir::visit::Visitor as VisitorMir, ty::TyCtxt};
 use tracing::{event, span, Level};
-use crate::analysis::taint_source::TaintSource;
 
-use super::{hir::crate_function_finder::{CrateFunctionFinder, FunctionInfo}, mir::data_flow::DataFlowTaintTracker};
+use super::{
+    hir::crate_function_finder::{CrateFunctionFinder, FunctionInfo},
+    mir::data_flow::DataFlowTaintTracker,
+};
 
 pub struct TaintCompilerCallbacks<'tsrc> {
     pub package_name: String,
@@ -66,18 +69,20 @@ pub fn mir_analysis(tcx: TyCtxt, callback_data: &mut TaintCompilerCallbacks) {
         );
     }
 
-    for function in &functions { 
-        let body = tcx.optimized_mir(function.local_def_id);
-        let mut tracker = DataFlowTaintTracker::new(tcx, body);
+    for function in &functions {
+        if callback_data.package_name == "sample"
+            || callback_data.package_name.starts_with("untrusted_value")
+        {
+            let body = tcx.optimized_mir(function.local_def_id);
+            let mut tracker = DataFlowTaintTracker::new(tcx, body);
 
-        println!("{}", function.function_name);
-        tracker.visit_body(body);
-        println!("\n\n\n");
-        if callback_data.package_name == "sample" || callback_data.package_name.starts_with("untrusted_value") {
+            println!("{}", function.function_name);
+            tracker.visit_body(body);
+            println!("\n\n\n");
             let dir_path = PathBuf::from("/tmp/taint/").join(&callback_data.package_name);
             std::fs::create_dir_all(&dir_path).expect("Failed to create directory");
             let dot_file = dir_path.join(&function.function_name).with_extension("dot");
-            let dot = Dot::with_config(&tracker.data_dependency_graph, &[Config::EdgeNoLabel]);
+            let dot = Dot::with_config(&tracker.data_dependency_graph, &[]);
             std::fs::write(&dot_file, format!("{:?}", dot)).expect("Failed to write dot file");
             let pdf_file = dot_file.with_extension("pdf");
             std::process::Command::new("dot")

untrusted_value_taint_checker/src/analysis/hir/crate_function_finder.rs

@@ -1,6 +1,5 @@
-use rustc_middle::ty;
 use rustc_hir::intravisit::Visitor as HirVisitor;
-
+use rustc_middle::ty;
 
 pub struct FunctionInfo {
     pub function_name: String,