Skip to content

Commit

Permalink
feat: finished data tainting
Browse files Browse the repository at this point in the history
  • Loading branch information
@0xCCF4
0xCCF4 committed Jul 27, 2024
1 parent 0337dea commit 5022fff
Show file tree
Hide file tree
Showing 4 changed files with 229 additions and 43 deletions.
5 changes: 5 additions & 0 deletions untrusted_value_taint_checker/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,15 @@ serde = { version = "1.0.204", features = ["derive"] }
semver = "1.0.23"
petgraph = "0.6.5"
itertools = "0.13.0"
graphviz-rust = { version = "0.9.0", optional = true }
#owo-colors = { version = "4", features = ["supports-colors"] }
#cargo_metadata = "0.18.1"
#bpaf = { version = "0.9.12", features = ["bpaf_derive", "autocomplete"] }

[features]
graphs = ["dep:graphviz-rust"]
default = ["graphs"]

[package.metadata.rust-analyzer]
rustc_private = true

Expand Down
30 changes: 21 additions & 9 deletions untrusted_value_taint_checker/src/analysis/callbacks.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use std::path::PathBuf;

use crate::analysis::taint_source::TaintSource;
use crate::analysis::{mir::data_flow_checker, taint_source::TaintSource};
use graphviz_rust::printer::{DotPrinter, PrinterContext};
use petgraph::dot::Dot;
use rustc_driver::Compilation;
use rustc_middle::{mir::visit::Visitor as VisitorMir, ty::TyCtxt};
Expand Down Expand Up @@ -71,30 +72,41 @@ pub fn mir_analysis(tcx: TyCtxt, callback_data: &mut TaintCompilerCallbacks) {

for function in &functions {
if callback_data.package_name == "sample"
|| callback_data.package_name.starts_with("untrusted_value")
// || 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");
//println!("\n\n\n");

let check = data_flow_checker::check_data_flow(
tracker.data_dependency_graph,
tcx,
callback_data.taint_sources,
true,
);

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, &[]);
std::fs::write(&dot_file, format!("{:?}", dot)).expect("Failed to write dot file");
let dot = check.data_dependency_graph.unwrap().print(&mut PrinterContext::default());
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")
let mut cmd = std::process::Command::new("dot")
.arg("-Tpdf")
.arg("-o")
.arg(&pdf_file)
.arg(&dot_file)
.output()
.spawn()
.expect("Failed to execute dot command");
std::fs::remove_file(&dot_file).expect("Failed to delete dot file");
let exit = cmd.wait().expect("Failed to wait for dot command");
if exit.success() {
std::fs::remove_file(&dot_file).expect("Failed to delete dot file");
}
}
println!("\n\n\n");
//println!("\n\n\n");
}
let mut functions = functions;
callback_data.internal_interface_functions = std::mem::take(&mut functions);
Expand Down
69 changes: 39 additions & 30 deletions untrusted_value_taint_checker/src/analysis/mir/data_flow.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,18 +11,27 @@ use rustc_span::Span;

#[derive(PartialEq, Eq, Hash, Clone)]
pub enum DependencyGraphNode<'tcx> {
Local { dst: mir::Local },
Local { dst: mir::Local, ty: ty::Ty<'tcx> },
FunctionCall { function: ty::Ty<'tcx>, span: Span },
ReturnedToCaller,
FunctionInput,
Assembly { spans: Vec<Span> },
ControlFlow,
}

impl<'tcx> DependencyGraphNode<'tcx> {
pub fn from_local(local: mir::Local, body: &mir::Body<'tcx>) -> Self {
DependencyGraphNode::Local {
dst: local,
ty: body.local_decls[local].ty,
}
}
}

impl<'tcx> Debug for DependencyGraphNode<'tcx> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self {
DependencyGraphNode::Local { dst } => write!(f, "{}", dst.index()),
DependencyGraphNode::Local { dst, ty: _ } => write!(f, "{}", dst.index()),
DependencyGraphNode::FunctionCall { function, span: _ } => {
write!(f, "Call: {}", function.to_string())
}
Expand Down Expand Up @@ -57,6 +66,12 @@ pub struct GraphEdge {
pub instances: Vec<EdgeInstance>,
}

impl GraphEdge {
pub fn is_move_only(&self) -> bool {
self.instances.iter().all(|x| x.data_flow_type == EdgeDataFlowType::Move)
}
}

impl Debug for GraphEdge {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(
Expand Down Expand Up @@ -231,7 +246,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
for arg in 1..body.arg_count + 1 {
self.add_data_dependency_edge(
DependencyGraphNode::FunctionInput,
DependencyGraphNode::Local { dst: arg.into() },
DependencyGraphNode::from_local(arg.into(), self.current_body),
body.span,
EdgeDataFlowType::Move,
);
Expand All @@ -240,15 +255,15 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
self.super_body(body);
}
fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) {
println!("\n - Basic Block Data: {}", block.index());
// println!("\n - Basic Block Data: {}", block.index());
self.super_basic_block_data(block, data);
}
fn visit_local_decl(&mut self, local: mir::Local, local_decl: &mir::LocalDecl<'tcx>) {
println!(" let {:?} = {:?}", local, local_decl.ty);
// println!(" let {:?} = {:?}", local, local_decl.ty);
self.super_local_decl(local, local_decl);
}
fn visit_statement(&mut self, statement: &mir::Statement<'tcx>, location: mir::Location) {
println!(" Statement: {:?} @ {:?}", statement, location);
// println!(" Statement: {:?} @ {:?}", statement, location);

let mut dependencies: Vec<(mir::Local, EdgeDataFlowType)> = Vec::with_capacity(2);

Expand All @@ -257,8 +272,8 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
self.get_rvalue_dependency(&rvalue, &mut dependencies);
for (dep, data_flow) in dependencies.iter() {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: *dep },
DependencyGraphNode::Local { dst: place.local },
DependencyGraphNode::from_local(*dep, self.current_body),
DependencyGraphNode::from_local(place.local, self.current_body),
statement.source_info.span,
*data_flow,
);
Expand All @@ -271,16 +286,16 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
match &copy.src {
mir::Operand::Copy(src) => {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: src.local },
DependencyGraphNode::Local { dst: dst.local },
DependencyGraphNode::from_local(src.local, self.current_body),
DependencyGraphNode::from_local(dst.local, self.current_body),
statement.source_info.span,
EdgeDataFlowType::Used,
);
}
mir::Operand::Move(src) => {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: src.local },
DependencyGraphNode::Local { dst: dst.local },
DependencyGraphNode::from_local(src.local, self.current_body),
DependencyGraphNode::from_local(dst.local, self.current_body),
statement.source_info.span,
EdgeDataFlowType::Move,
);
Expand Down Expand Up @@ -310,11 +325,11 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
rvalue: &mir::Rvalue<'tcx>,
location: mir::Location,
) {
println!(" {:?} <- {:?} @ {:?}", place, rvalue, location);
// println!(" {:?} <- {:?} @ {:?}", place, rvalue, location);
self.super_assign(place, rvalue, location);
}
fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, location: mir::Location) {
println!(" stop {:?}", terminator.kind);
// println!(" stop {:?}", terminator.kind);

let mut dependencies: Vec<(mir::Local, EdgeDataFlowType)> = Vec::with_capacity(2);

Expand All @@ -324,7 +339,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
self.get_operand_dependency(&discr, &mut dependencies, None);
for (dep, flow) in dependencies {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: dep },
DependencyGraphNode::from_local(dep, self.current_body),
DependencyGraphNode::ControlFlow,
terminator.source_info.span,
flow,
Expand All @@ -335,9 +350,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
mir::TerminatorKind::UnwindTerminate(_reason) => {}
mir::TerminatorKind::Return => {
self.add_data_dependency_edge(
DependencyGraphNode::Local {
dst: mir::Local::ZERO,
},
DependencyGraphNode::from_local(mir::Local::ZERO, self.current_body),
DependencyGraphNode::ReturnedToCaller,
terminator.source_info.span,
EdgeDataFlowType::Move,
Expand Down Expand Up @@ -370,7 +383,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {

for (dep, flow) in dependencies {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: dep },
DependencyGraphNode::from_local(dep, self.current_body),
func_call.clone(),
fn_span.to_owned(),
flow,
Expand All @@ -379,9 +392,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {

self.add_data_dependency_edge(
func_call,
DependencyGraphNode::Local {
dst: destination.local,
},
DependencyGraphNode::from_local(destination.local, self.current_body),
fn_span.to_owned(),
EdgeDataFlowType::Move,
);
Expand All @@ -402,7 +413,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {

for (dep, flow) in dependencies {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: dep },
DependencyGraphNode::from_local(dep, self.current_body),
func_call.clone(),
fn_span.to_owned(),
flow,
Expand All @@ -427,7 +438,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {

for (dep, flow) in dependencies {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: dep },
DependencyGraphNode::from_local(dep, self.current_body),
DependencyGraphNode::ReturnedToCaller,
terminator.source_info.span,
flow,
Expand All @@ -436,9 +447,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {

self.add_data_dependency_edge(
DependencyGraphNode::FunctionInput,
DependencyGraphNode::Local {
dst: resume_arg.local,
},
DependencyGraphNode::from_local(resume_arg.local, self.current_body),
terminator.source_info.span,
EdgeDataFlowType::Move,
);
Expand Down Expand Up @@ -510,7 +519,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {

for (dep, flow) in inputs {
self.add_data_dependency_edge(
DependencyGraphNode::Local { dst: dep },
DependencyGraphNode::from_local(dep, self.current_body),
asm_block.clone(),
terminator.source_info.span,
flow,
Expand All @@ -519,7 +528,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
for (dep, flow) in outputs {
self.add_data_dependency_edge(
asm_block.clone(),
DependencyGraphNode::Local { dst: dep },
DependencyGraphNode::from_local(dep, self.current_body),
terminator.source_info.span,
flow,
);
Expand All @@ -530,7 +539,7 @@ impl<'tcx, 'a> VisitorMir<'tcx> for DataFlowTaintTracker<'tcx, 'a> {
self.super_terminator(terminator, location);
}
fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, location: mir::Location) {
println!(" op {:?}", operand);
// println!(" op {:?}", operand);
self.super_operand(operand, location);
}
}
Loading

0 comments on commit 5022fff

Please sign in to comment.