use crate::secret;
use either::Either;
use haybale::{Error, Project, Result, ReturnValue, State};
use haybale::backend::BV;
use haybale::function_hooks::IsCall;
use llvm_ir::{Constant, Name, Operand, Type, Typed};
use log::info;
pub fn pitchfork_default_hook(
proj: &Project,
state: &mut State<secret::Backend>,
call: &dyn IsCall,
) -> Result<ReturnValue<secret::BV>> {
let called_funcname = match call.get_called_func() {
Either::Left(_) => panic!("invoked default hook for an inline assembly call"),
Either::Right(Operand::ConstantOperand(Constant::GlobalReference { name: Name::Name(name), .. })) => Some(name),
Either::Right(Operand::ConstantOperand(Constant::GlobalReference { name, .. })) => panic!("Function with a numbered name: {:?}", name),
Either::Right(_) => None,
};
let pretty_funcname = match called_funcname {
Some(funcname) => format!("a function named {:?}", state.demangle(funcname)),
None => "a function pointer".into(),
};
info!("Using Pitchfork default hook for {}", pretty_funcname);
for (i, arg) in call.get_arguments().iter().map(|(arg, _)| arg).enumerate() {
let arg_bv = state.operand_to_bv(arg)?;
match is_or_points_to_secret(proj, state, &arg_bv, &arg.get_type())? {
ArgumentKind::Secret => match called_funcname {
Some(funcname) => {
let demangled = state.demangle(funcname);
return Err(Error::OtherError(format!("Encountered a call of {}, but didn't find an LLVM definition or function hook for it; and its argument #{} (zero-indexed) may refer to secret data.\nTo fix this error, you can do one of these three options:\n (1) choose to simply ignore this function call by adding it to `config.function_hooks` with `haybale::function_hooks::generic_stub_hook` as the hook;\n (2) rerun with more bitcode files in the `Project` so that the symbolic execution can find an LLVM definition for {:?};\n (3) write your own custom hook for {:?}", pretty_funcname, i, demangled, demangled)));
},
None => {
return Err(Error::OtherError(format!("pitchfork_default_hook on a function pointer, and argument #{} (zero-indexed) may refer to secret data", i)));
},
},
ArgumentKind::Unknown => match called_funcname {
Some(funcname) => {
let demangled = state.demangle(funcname);
return Err(Error::OtherError(format!("Encountered a call of {}, but didn't find an LLVM definition or function hook for it; and its argument #{} (zero-indexed) involves an opaque struct type, so we're not sure if it may contain secret data.\nTo fix this error, you can do one of these three options:\n (1) choose to simply ignore this function call by adding it to `config.function_hooks` with `haybale::function_hooks::generic_stub_hook` as the hook;\n (2) rerun with more bitcode files in the `Project` so that the symbolic execution can find an LLVM definition for {:?};\n (3) write your own custom hook for {:?}", pretty_funcname, i, demangled, demangled)));
},
None => {
return Err(Error::OtherError(format!("pitchfork_default_hook on a function pointer, and argument #{} (zero-indexed) involves an opaque struct type, so we're not sure if it may contain secret data", i)));
},
},
ArgumentKind::Public => {},
}
}
haybale::function_hooks::generic_stub_hook(proj, state, call)
}
#[derive(Clone, Debug)]
pub(crate) enum ArgumentKind {
Public,
Secret,
Unknown,
}
pub(crate) fn is_or_points_to_secret(proj: &Project, state: &mut State<secret::Backend>, bv: &secret::BV, ty: &llvm_ir::Type) -> Result<ArgumentKind> {
if bv.is_secret() {
Ok(ArgumentKind::Secret)
} else {
match ty {
Type::PointerType { pointee_type, .. } => {
if let Type::FuncType { .. } = &**pointee_type {
return Ok(ArgumentKind::Public);
}
let pointee_size_bits = match haybale::layout::size_opaque_aware(&**pointee_type, proj) {
None => return Ok(ArgumentKind::Unknown),
Some(size) => size,
};
let mut need_pop = false;
if state.bvs_can_be_equal(&bv, &state.zero(bv.get_width()))? {
state.solver.push(1);
need_pop = true;
bv._ne(&state.zero(bv.get_width())).assert()?;
}
let pointee = match state.read(&bv, pointee_size_bits as u32) {
Ok(pointee) => pointee,
Err(e) => {
if need_pop {
state.solver.pop(1);
}
return Err(e);
},
};
let retval = is_or_points_to_secret(proj, state, &pointee, &**pointee_type);
if need_pop {
state.solver.pop(1);
}
retval
},
Type::VectorType { element_type, num_elements } | Type::ArrayType { element_type, num_elements } => {
let element_bits = match haybale::layout::size_opaque_aware(&**element_type, proj) {
None => return Ok(ArgumentKind::Unknown),
Some(size) => size as u32,
};
if element_bits == 0 {
Ok(ArgumentKind::Public)
} else {
let mut retval = ArgumentKind::Public;
for i in 0 .. *num_elements {
let i = i as u32;
let element = bv.slice((i+1) * element_bits - 1, i * element_bits);
match is_or_points_to_secret(proj, state, &element, &**element_type)? {
ArgumentKind::Secret => return Ok(ArgumentKind::Secret),
ArgumentKind::Unknown => retval = ArgumentKind::Unknown,
ArgumentKind::Public => {},
}
}
Ok(retval)
}
},
Type::StructType { element_types, .. } => {
let mut offset_bits = 0;
let mut retval = ArgumentKind::Public;
for element_ty in element_types {
let element_bits = match haybale::layout::size_opaque_aware(element_ty, proj) {
None => return Ok(ArgumentKind::Unknown),
Some(size) => size as u32,
};
if element_bits == 0 {
} else {
let element = bv.slice(offset_bits + element_bits - 1, offset_bits);
match is_or_points_to_secret(proj, state, &element, element_ty)? {
ArgumentKind::Secret => return Ok(ArgumentKind::Secret),
ArgumentKind::Unknown => retval = ArgumentKind::Unknown,
ArgumentKind::Public => {},
}
offset_bits += element_bits;
assert_eq!(offset_bits % 8, 0, "Struct offset of {} bits is not a multiple of 8 bits", offset_bits);
}
}
Ok(retval)
},
Type::NamedStructType { .. } => {
match proj.get_inner_struct_type_from_named(ty) {
None => Ok(ArgumentKind::Unknown),
Some(arc) => is_or_points_to_secret(proj, state, bv, &arc.read().unwrap()),
}
},
_ => Ok(ArgumentKind::Public),
}
}
}