Skip to content

Commit

Permalink
refinement of the switchint traversal
Browse files Browse the repository at this point in the history
  • Loading branch information
@VaynNecol
VaynNecol committed Jan 5, 2025
1 parent ff5f497 commit 3f5dfeb
Show file tree
Hide file tree
Showing 3 changed files with 218 additions and 2 deletions.
72 changes: 72 additions & 0 deletions rap/src/analysis/safedrop/graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,10 @@ use rustc_middle::ty;
use rustc_middle::ty::TyCtxt;
use rustc_span::def_id::DefId;
use rustc_span::Span;
use std::cell::RefCell;
use std::cmp::min;
use std::vec::Vec;

//use crate::rap_info;

#[derive(PartialEq, Debug, Copy, Clone)]
Expand Down Expand Up @@ -67,6 +69,10 @@ pub struct BlockNode<'tcx> {
pub const_value: Vec<(usize, usize)>,
//store switch stmts in current block for the path filtering in path-sensitive analysis.
pub switch_stmts: Vec<Terminator<'tcx>>,

pub modified_value: FxHashSet<usize>,
// (SwitchInt target, enum index) -> outside nodes.
pub scc_outer: RefCell<Option<FxHashMap<(usize, usize), Vec<usize>>>>,
}

impl<'tcx> BlockNode<'tcx> {
Expand All @@ -81,6 +87,8 @@ impl<'tcx> BlockNode<'tcx> {
scc_sub_blocks: Vec::<usize>::new(),
const_value: Vec::<(usize, usize)>::new(),
switch_stmts: Vec::<Terminator<'tcx>>::new(),
modified_value: FxHashSet::<usize>::default(),
scc_outer: RefCell::new(None),
}
}

Expand Down Expand Up @@ -165,6 +173,15 @@ pub struct SafeDropGraph<'tcx> {
pub dead_record: Vec<bool>,
// analysis of heap item
pub adt_owner: AdtOwner,

pub child_scc: FxHashMap<
usize,
(
BlockNode<'tcx>,
rustc_middle::mir::SwitchTargets,
FxHashSet<usize>,
),
>,
}

impl<'tcx> SafeDropGraph<'tcx> {
Expand Down Expand Up @@ -214,6 +231,7 @@ impl<'tcx> SafeDropGraph<'tcx> {
if let StatementKind::Assign(ref assign) = stmt.kind {
let lv_local = assign.0.local.as_usize(); // assign.0 is a Place
let lv = assign.0.clone();
cur_bb.modified_value.insert(lv_local);
match assign.1 {
// assign.1 is a Rvalue
Rvalue::Use(ref x) => {
Expand Down Expand Up @@ -479,6 +497,7 @@ impl<'tcx> SafeDropGraph<'tcx> {
alias_set: alias,
dead_record: dead,
adt_owner,
child_scc: FxHashMap::default(),
}
}

Expand Down Expand Up @@ -507,8 +526,13 @@ impl<'tcx> SafeDropGraph<'tcx> {
}
}
}

// generate SCC
if dfn[index] == low[index] {
let mut modified_set = FxHashSet::<usize>::default();
let mut switch_target = Vec::new();
let mut scc_block_set = FxHashSet::<usize>::default();
let init_block = self.blocks[index].clone();
loop {
let node = stack.pop().unwrap();
self.scc_indices[node] = index;
Expand All @@ -518,11 +542,34 @@ impl<'tcx> SafeDropGraph<'tcx> {
break;
}
self.blocks[index].scc_sub_blocks.push(node);
scc_block_set.insert(node);

for value in &self.blocks[index].modified_value {
modified_set.insert(*value);
}
if let Some(target) = self.switch_target(self.tcx, node) {
switch_target.push((target, self.blocks[index].switch_stmts[0].clone()));
}

let nexts = self.blocks[node].next.clone();
for i in nexts {
self.blocks[index].next.insert(i);
}
}
switch_target.retain(|v| !modified_set.contains(&(v.0)));

if !switch_target.is_empty() && switch_target.len() == 1 {
//let target_index = switch_target[0].0;
let target_terminator = switch_target[0].1.clone();

let TerminatorKind::SwitchInt { discr: _, targets } = target_terminator.kind else {
unreachable!();
};

self.child_scc
.insert(index, (init_block, targets, scc_block_set));
}

/* remove next nodes which are already in the current SCC */
let mut to_remove = Vec::new();
for i in self.blocks[index].next.iter() {
Expand Down Expand Up @@ -577,4 +624,29 @@ impl<'tcx> SafeDropGraph<'tcx> {

paths
}

pub fn switch_target(&mut self, tcx: TyCtxt<'tcx>, block_index: usize) -> Option<usize> {
let block = &self.blocks[block_index];
if block.switch_stmts.is_empty() {
return None;
}

let res = if let TerminatorKind::SwitchInt {
ref discr,
ref targets,
} = &block.switch_stmts[0].kind
{
match discr {
Operand::Copy(p) | Operand::Move(p) => {
let place = self.projection(tcx, false, p.clone());
Some(place)
}
_ => None,
}
} else {
None
};

res
}
}
142 changes: 142 additions & 0 deletions rap/src/analysis/safedrop/safedrop.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use rustc_middle::ty::{TyCtxt, TyKind};
use crate::analysis::core::alias::FnMap;
use crate::analysis::safedrop::SafeDropGraph;
use crate::rap_error;
use rustc_data_structures::fx::FxHashSet;

pub const VISIT_LIMIT: usize = 1000;

Expand Down Expand Up @@ -85,6 +86,7 @@ impl<'tcx> SafeDropGraph<'tcx> {
self.alias_set = backup_alias_set;
self.dead_record = backup_dead;
}

pub fn split_check_with_cond(
&mut self,
bb_index: usize,
Expand Down Expand Up @@ -119,6 +121,146 @@ impl<'tcx> SafeDropGraph<'tcx> {
self.alias_bbcall(self.scc_indices[bb_index], tcx, fn_map);
self.drop_check(self.scc_indices[bb_index], tcx);

if self.child_scc.get(&self.scc_indices[bb_index]).is_some() {
let init_index = self.scc_indices[bb_index];
let (init_block, cur_targets, scc_block_set) =
self.child_scc.get(&init_index).unwrap().clone();

for enum_index in cur_targets.all_targets() {
let backup_values = self.values.clone();
let backup_constant = self.constant.clone();

let mut block_node = if bb_index == init_index {
init_block.clone()
} else {
self.blocks[bb_index].clone()
};

if !block_node.switch_stmts.is_empty() {
let TerminatorKind::SwitchInt { discr, targets } =
block_node.switch_stmts[0].kind.clone()
else {
unreachable!();
};
if cur_targets == targets {
block_node.next = FxHashSet::default();
block_node.next.insert(enum_index.index());
}
}

let mut work_list = Vec::new();
let mut work_set = FxHashSet::<usize>::default();
work_list.push(bb_index);
work_set.insert(bb_index);
while !work_list.is_empty() {
let current_node = work_list.pop().unwrap();
block_node.scc_sub_blocks.push(current_node);
let real_node = if current_node != init_index {
self.blocks[current_node].clone()
} else {
init_block.clone()
};

if real_node.switch_stmts.is_empty() {
for next in &real_node.next {
block_node.next.insert(*next);
}
} else {
let TerminatorKind::SwitchInt {
ref discr,
ref targets,
} = real_node.switch_stmts[0].kind
else {
unreachable!();
};

if cur_targets == *targets {
block_node.next.insert(enum_index.index());
} else {
for next in &real_node.next {
block_node.next.insert(*next);
}
}
}

if real_node.switch_stmts.is_empty() {
for next in &real_node.next {
if scc_block_set.contains(next) && !work_set.contains(next) {
work_set.insert(*next);
work_list.push(*next);
}
}
} else {
let TerminatorKind::SwitchInt {
ref discr,
ref targets,
} = real_node.switch_stmts[0].kind
else {
unreachable!();
};

if cur_targets == *targets {
let next_index = enum_index.index();
if scc_block_set.contains(&next_index)
&& !work_set.contains(&next_index)
{
work_set.insert(next_index);
work_list.push(next_index);
}
} else {
for next in &real_node.next {
if scc_block_set.contains(next) && !work_set.contains(next) {
work_set.insert(*next);
work_list.push(*next);
}
}
}
}
}

/* remove next nodes which are already in the current SCC */
let mut to_remove = Vec::new();
for i in block_node.next.iter() {
if self.scc_indices[*i] == init_index {
to_remove.push(*i);
}
}
for i in to_remove {
block_node.next.remove(&i);
}

for i in block_node.scc_sub_blocks.clone() {
self.alias_bb(i, tcx);
self.alias_bbcall(i, tcx, fn_map);
self.drop_check(i, tcx);
}
/* Reach a leaf node, check bugs */
match block_node.next.len() {
0 => {
// check the bugs.
if Self::should_check(self.def_id) {
self.dp_check(&cur_block);
}
return;
}
_ => {
/*
* Equivalent to self.check(cur_block.next[0]..);
* We cannot use [0] for FxHashSet.
*/
for next in block_node.next {
self.check(next, tcx, fn_map);
}
}
}

self.values = backup_values;
self.constant = backup_constant;
}

return;
}

/* Handle cases if the current block is a merged scc block with sub block */
if cur_block.scc_sub_blocks.len() > 0 {
for i in cur_block.scc_sub_blocks.clone() {
Expand Down
6 changes: 4 additions & 2 deletions tests/support/uaf/uaf_swithint_FP/src/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::mem::ManuallyDrop;
enum A {
A1,
A2
A2,
}
fn evil_test(a: A) {
let mut count = 0;
Expand All @@ -12,6 +12,7 @@ fn evil_test(a: A) {
if count < 2 {
unsafe {
ManuallyDrop::drop(&mut slot);
ManuallyDrop::drop(&mut slot);
}
continue;
}
Expand All @@ -20,6 +21,7 @@ fn evil_test(a: A) {
println!("{:?}", slot);
unsafe {
ManuallyDrop::drop(&mut slot);
ManuallyDrop::drop(&mut slot);
}
break;
}
Expand All @@ -30,4 +32,4 @@ fn main() {
let a2 = A::A2;
evil_test(a1);
evil_test(a2);
}
}

0 comments on commit 3f5dfeb

Please sign in to comment.