Skip to content

Commit

Permalink
Cross-function union instrumentation (#3465)
Browse files Browse the repository at this point in the history 
This PR introduces support for memory initialization checks for unions
passed across the function boundary.

Whenever a union is passed as an argument, we need to make sure that its
initialization state is preserved. Unlike pointers, unions do not have a
stable memory address which could identify them in shadow memory. Hence,
we need to pass extra information across function boundary since unions
are passed “by value”.

We introduce a global variable to store the previous address of unions
passed as function arguments, which allows us to effectively tie the
initialization state of unions passed between functions. This struct is
written to by the caller and read from by the callee.

For more information about planned functionality, see
#3300

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.
  • Loading branch information
@artemagvanian
artemagvanian authored Sep 4, 2024
1 parent edded56 commit 93a29af
Show file tree
Hide file tree
Showing 9 changed files with 516 additions and 120 deletions.
68 changes: 66 additions & 2 deletions kani-compiler/src/kani_middle/transform/check_uninit/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ const KANI_IS_STR_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniIsStrPtrInitialized";
const KANI_SET_STR_PTR_INITIALIZED_DIAGNOSTIC: &str = "KaniSetStrPtrInitialized";
const KANI_COPY_INIT_STATE_DIAGNOSTIC: &str = "KaniCopyInitState";
const KANI_COPY_INIT_STATE_SINGLE_DIAGNOSTIC: &str = "KaniCopyInitStateSingle";
const KANI_LOAD_ARGUMENT_DIAGNOSTIC: &str = "KaniLoadArgument";
const KANI_STORE_ARGUMENT_DIAGNOSTIC: &str = "KaniStoreArgument";
const KANI_RESET_ARGUMENT_BUFFER_DIAGNOSTIC: &str = "KaniResetArgumentBuffer";

// Function bodies of those functions will not be instrumented as not to cause infinite recursion.
const SKIPPED_DIAGNOSTIC_ITEMS: &[&str] = &[
Expand All @@ -58,6 +61,9 @@ const SKIPPED_DIAGNOSTIC_ITEMS: &[&str] = &[
KANI_SET_STR_PTR_INITIALIZED_DIAGNOSTIC,
KANI_COPY_INIT_STATE_DIAGNOSTIC,
KANI_COPY_INIT_STATE_SINGLE_DIAGNOSTIC,
KANI_LOAD_ARGUMENT_DIAGNOSTIC,
KANI_STORE_ARGUMENT_DIAGNOSTIC,
KANI_RESET_ARGUMENT_BUFFER_DIAGNOSTIC,
];

/// Instruments the code with checks for uninitialized memory, agnostic to the source of targets.
Expand Down Expand Up @@ -159,9 +165,9 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
// Calculate pointee layout for byte-by-byte memory initialization checks.
match PointeeInfo::from_ty(pointee_ty) {
Ok(type_info) => type_info,
Err(_) => {
Err(reason) => {
let reason = format!(
"Kani currently doesn't support checking memory initialization for pointers to `{pointee_ty}.",
"Kani currently doesn't support checking memory initialization for pointers to `{pointee_ty}. {reason}",
);
self.inject_assert_false(self.tcx, body, source, operation.position(), &reason);
return;
Expand All @@ -185,6 +191,9 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
MemoryInitOp::AssignUnion { .. } => {
self.build_assign_union(body, source, operation, pointee_info)
}
MemoryInitOp::StoreArgument { .. } | MemoryInitOp::LoadArgument { .. } => {
self.build_argument_operation(body, source, operation, pointee_info)
}
MemoryInitOp::Unsupported { .. } | MemoryInitOp::TriviallyUnsafe { .. } => {
unreachable!()
}
Expand Down Expand Up @@ -532,6 +541,61 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
);
}

/// Instrument the code to pass information about arguments containing unions. Whenever a
/// function is called and some of the arguments contain unions, we store the information. And
/// when we enter the callee, we load the information.
fn build_argument_operation(
&mut self,
body: &mut MutableBody,
source: &mut SourceInstruction,
operation: MemoryInitOp,
pointee_info: PointeeInfo,
) {
let ret_place = Place {
local: body.new_local(Ty::new_tuple(&[]), source.span(body.blocks()), Mutability::Not),
projection: vec![],
};
let mut statements = vec![];
let layout_size = pointee_info.layout().maybe_size().unwrap();
let diagnostic = match operation {
MemoryInitOp::LoadArgument { .. } => KANI_LOAD_ARGUMENT_DIAGNOSTIC,
MemoryInitOp::StoreArgument { .. } => KANI_STORE_ARGUMENT_DIAGNOSTIC,
_ => unreachable!(),
};
let argument_operation_instance = resolve_mem_init_fn(
get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
layout_size,
*pointee_info.ty(),
);
let operand = operation.mk_operand(body, &mut statements, source);
let argument_no = operation.expect_argument_no();
let terminator = Terminator {
kind: TerminatorKind::Call {
func: Operand::Copy(Place::from(body.new_local(
argument_operation_instance.ty(),
source.span(body.blocks()),
Mutability::Not,
))),
args: vec![
operand,
Operand::Constant(ConstOperand {
span: source.span(body.blocks()),
user_ty: None,
const_: MirConst::try_from_uint(argument_no as u128, UintTy::Usize)
.unwrap(),
}),
],
destination: ret_place.clone(),
target: Some(0), // this will be overriden in add_bb
unwind: UnwindAction::Terminate,
},
span: source.span(body.blocks()),
};

// Construct the basic block and insert it into the body.
body.insert_bb(BasicBlock { statements, terminator }, source, operation.position());
}

/// Copy memory initialization state from one union variable to another.
fn build_assign_union(
&mut self,
Expand Down
132 changes: 96 additions & 36 deletions kani-compiler/src/kani_middle/transform/check_uninit/ptr_uninit/uninit_visitor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@ use crate::{
body::{InsertPosition, MutableBody, SourceInstruction},
check_uninit::{
relevant_instruction::{InitRelevantInstruction, MemoryInitOp},
ty_layout::tys_layout_compatible_to_size,
TargetFinder,
ty_layout::{tys_layout_compatible_to_size, LayoutComputationError},
PointeeInfo, TargetFinder,
},
},
};
Expand Down Expand Up @@ -38,11 +38,38 @@ impl TargetFinder for CheckUninitVisitor {
fn find_all(mut self, body: &MutableBody) -> Vec<InitRelevantInstruction> {
self.locals = body.locals().to_vec();
for (bb_idx, bb) in body.blocks().iter().enumerate() {
self.current_target = InitRelevantInstruction {
source: SourceInstruction::Statement { idx: 0, bb: bb_idx },
before_instruction: vec![],
after_instruction: vec![],
// Set the first target to start iterating from.
self.current_target = if !bb.statements.is_empty() {
InitRelevantInstruction {
source: SourceInstruction::Statement { idx: 0, bb: bb_idx },
before_instruction: vec![],
after_instruction: vec![],
}
} else {
InitRelevantInstruction {
source: SourceInstruction::Terminator { bb: bb_idx },
before_instruction: vec![],
after_instruction: vec![],
}
};
if bb_idx == 0 {
let union_args: Vec<_> = body
.locals()
.iter()
.enumerate()
.skip(1)
.take(body.arg_count())
.filter(|(_, local)| local.ty.kind().is_union())
.collect();
if !union_args.is_empty() {
for (idx, _) in union_args {
self.push_target(MemoryInitOp::LoadArgument {
operand: Operand::Copy(Place { local: idx, projection: vec![] }),
argument_no: idx,
})
}
}
}
self.visit_basic_block(bb);
}
self.targets
Expand Down Expand Up @@ -84,6 +111,22 @@ impl MirVisitor for CheckUninitVisitor {
StatementKind::Assign(place, rvalue) => {
// First check rvalue.
self.visit_rvalue(rvalue, location);

// Preemptively check if we do not support computing the layout of a type because of
// inner union fields. This allows to inject `assert!(false)` early.
if let Err(reason @ LayoutComputationError::UnionAsField(_)) =
PointeeInfo::from_ty(place.ty(&self.locals).unwrap())
{
self.push_target(MemoryInitOp::Unsupported {
reason: format!(
"Checking memory initialization of type {} is not supported. {}",
place.ty(&self.locals).unwrap(),
reason
),
});
return;
}

// Check whether we are assigning into a dereference (*ptr = _).
if let Some(place_without_deref) = try_remove_topmost_deref(place) {
// First, check that we are not dereferencing extra pointers along the way
Expand Down Expand Up @@ -127,9 +170,10 @@ impl MirVisitor for CheckUninitVisitor {
// if a union as a subfield is detected, `assert!(false)` will be injected from
// the type layout code.
let is_inside_union = {
let mut contains_union = false;
let mut place_to_add_projections =
Place { local: place.local, projection: vec![] };
let mut contains_union =
place_to_add_projections.ty(&self.locals).unwrap().kind().is_union();
for projection_elem in place.projection.iter() {
if place_to_add_projections.ty(&self.locals).unwrap().kind().is_union() {
contains_union = true;
Expand All @@ -143,35 +187,37 @@ impl MirVisitor for CheckUninitVisitor {
// Need to copy some information about union initialization, since lvalue is
// either a union or a field inside a union.
if is_inside_union {
if let Rvalue::Use(operand) = rvalue {
// This is a union-to-union assignment, so we need to copy the
// initialization state.
if place.ty(&self.locals).unwrap().kind().is_union() {
self.push_target(MemoryInitOp::AssignUnion {
lvalue: place.clone(),
rvalue: operand.clone(),
});
} else {
// This is assignment to a field of a union.
self.push_target(MemoryInitOp::SetRef {
operand: Operand::Copy(place.clone()),
value: true,
position: InsertPosition::After,
});
match rvalue {
Rvalue::Use(operand) => {
// This is a union-to-union assignment, so we need to copy the
// initialization state.
if place.ty(&self.locals).unwrap().kind().is_union() {
self.push_target(MemoryInitOp::AssignUnion {
lvalue: place.clone(),
rvalue: operand.clone(),
});
} else {
// This is assignment to a field of a union.
self.push_target(MemoryInitOp::SetRef {
operand: Operand::Copy(place.clone()),
value: true,
position: InsertPosition::After,
});
}
}
}
}

// Create a union from scratch as an aggregate. We handle it here because we
// need to know which field is getting assigned.
if let Rvalue::Aggregate(AggregateKind::Adt(adt_def, _, _, _, union_field), _) =
rvalue
{
if adt_def.kind() == AdtKind::Union {
self.push_target(MemoryInitOp::CreateUnion {
operand: Operand::Copy(place.clone()),
field: union_field.unwrap(), // Safe to unwrap because we know this is a union.
});
Rvalue::Aggregate(AggregateKind::Adt(adt_def, _, _, _, union_field), _) => {
// Create a union from scratch as an aggregate. We handle it here because we
// need to know which field is getting assigned.
if adt_def.kind() == AdtKind::Union {
self.push_target(MemoryInitOp::CreateUnion {
operand: Operand::Copy(place.clone()),
field: union_field.unwrap(), // Safe to unwrap because we know this is a union.
});
}
}
// TODO: add support for Rvalue::Cast, etc.
_ => self
.push_target(MemoryInitOp::Unsupported { reason: "Performing a union assignment with a non-supported construct as an Rvalue".to_string() }),
}
}
}
Expand Down Expand Up @@ -218,7 +264,7 @@ impl MirVisitor for CheckUninitVisitor {
before_instruction: vec![],
};
} else {
unreachable!()
// The only instruction in this basic block is the terminator, which was already set.
}
// Leave it as an exhaustive match to be notified when a new kind is added.
match &term.kind {
Expand Down Expand Up @@ -340,6 +386,20 @@ impl MirVisitor for CheckUninitVisitor {
}
_ => {}
}
} else {
let union_args: Vec<_> = args
.iter()
.enumerate()
.filter(|(_, arg)| arg.ty(&self.locals).unwrap().kind().is_union())
.collect();
if !union_args.is_empty() {
for (idx, operand) in union_args {
self.push_target(MemoryInitOp::StoreArgument {
operand: operand.clone(),
argument_no: idx + 1, // since arguments are 1-indexed
})
}
}
}
}
_ => {}
Expand Down
Loading

0 comments on commit 93a29af

Please sign in to comment.