Skip to content

Commit

Permalink
Use evaluator to eliminate if and switch branches.
Browse files Browse the repository at this point in the history
  • Loading branch information
chriseth committed Nov 1, 2022
1 parent c318ba3 commit b57985c
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 53 deletions.
6 changes: 6 additions & 0 deletions src/encoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,9 @@ impl<InstructionsType: Instructions> Encoder<InstructionsType> {
fn encode_if(&mut self, expr: &yul::If) {
let cond = self.encode_expression(&expr.condition);
assert!(cond.len() == 1);
if self.evaluator.variable_known_equal(&cond[0], &"0".to_string()) {
return;
}
let prev_ssa = self.ssa_tracker.copy_current_ssa();

self.push_path_condition(smt::neq(cond[0].clone(), 0));
Expand Down Expand Up @@ -303,6 +306,9 @@ impl<InstructionsType: Instructions> Encoder<InstructionsType> {
} in &switch.cases
{
let is_default = literal.is_none();
if !is_default && self.evaluator.variable_known_unequal(&discriminator[0], &literal.as_ref().unwrap().literal) {
continue;
}

self.ssa_tracker.set_current_ssa(pre_switch_ssa.clone());

Expand Down
142 changes: 89 additions & 53 deletions src/evaluator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ enum Value {
struct ValueByte {
data: Value,
/// Byte of the data, 0 is most significant byte.
byte: u32
byte: u32,
}

impl Display for Value {
Expand Down Expand Up @@ -82,8 +82,29 @@ impl Evaluator {
/// TODO should also set the current address.
pub fn new_transaction(&mut self, calldata: Vec<u8>) {
self.calldata = Some(calldata);
self.memory.clear();
self.memory_slices.clear();
self.unknown_memory_is_zero = true;
}
/// Returns true if the evaluator can determine that the variable is equal
/// to the given literal value with the previously set restrictions.
pub fn variable_known_equal(&self, var: &SMTVariable, value: &String) -> bool {
if let Some(Value::Concrete(v)) = self.ssa_values.get(&var.name) {
if let Some(lit) = literal_value(value) {
return *v == lit;
}
}
false
}
pub fn variable_known_unequal(&self, var: &SMTVariable, value: &String) -> bool {
if let Some(Value::Concrete(v)) = self.ssa_values.get(&var.name) {
if let Some(lit) = literal_value(value) {
return *v != lit;
}
}
false
}

pub fn define_from_literal(&mut self, var: &SMTVariable, literal: &Literal) {
let value = &literal.literal;
//println!("{} := {literal}", var.name);
Expand Down Expand Up @@ -113,37 +134,6 @@ impl Evaluator {
arguments: &Vec<SMTVariable>,
return_vars: &[SMTVariable],
) {
if builtin.name == "create" {
println!(
"Create: {:?} {:?} {:?}",
arguments[0], arguments[1], arguments[2]
);

// self.ssa_values
// .insert(return_vars[0].name.clone(), BigUint::from(1234u64));
}
// if builtin.name == "sstore" {
// if let (Some(key), Some(value)) = (
// self.ssa_values.get(&arguments[0].name).cloned(),
// self.ssa_values.get(&arguments[1].name).cloned(),
// ) {
// self.storage.insert(key, value);
// }
// }
// if builtin.name == "sload" {
// if let Some(key) = self.ssa_values.get(&arguments[0].name).cloned() {
// if let Some(value) = self.storage.get(&key).cloned() {
// self.ssa_values.insert(return_vars[0].name.clone(), value);
// } else {
// // TODO assume unknown storage is some weird value - should use unknown bits later
// self.ssa_values.insert(
// return_vars[0].name.clone(),
// BigUint::from(0x1234567812345678u64),
// );
// }
// }
// }

let mask = (BigUint::from(1u64) << 256) - BigUint::from(1u64);
let two256 = BigUint::from(1u64) << 256;
let address_mask = (BigUint::from(1u64) << 160) - BigUint::from(1u64);
Expand Down Expand Up @@ -200,6 +190,18 @@ impl Evaluator {
self.memory_slices.insert(addr.clone(), Value::DataRef(offset.clone()));
None
}
("calldatasize", []) => {
self.calldata.as_ref().map(|cd| Value::Concrete(BigUint::from(cd.len())))
}
("calldataload", [Some(Value::Concrete(addr))]) => {
self.calldata.as_ref().map(|cd| {
let mut result = vec![0u8; 32];
for i in 0..32usize {
result[i] = cd.get(addr.to_usize().unwrap() + i).copied().unwrap_or_default();
}
Value::Concrete(BigUint::from_bytes_be(&result))
})
}
("create", [_ether_value, Some(Value::Concrete(offset)), _size /* TODO check size */]) => {
if let Some(Value::DataRef(name)) = self.memory_slices.get(offset) {
Some(Value::KnownContractAddress(name.clone(), 0)) // TODO increment coutner
Expand Down Expand Up @@ -233,14 +235,17 @@ impl Evaluator {
self.ssa_values.insert(return_vars[0].name.clone(), result);
}
if builtin.name == "call" {
for i in 0x80..0xc4 {
println!("{:x}: {:x?}", i,self.read_memory_byte(&BigUint::from(i as u32)));
}
// for i in 0x80..0xc4 {
// println!(
// "{:x}: {:x?}",
// i,
// self.read_memory_byte(&BigUint::from(i as u32))
// );
// }
// if let (Some(Value::Concrete(in_offset)), Some(Value::Concrete(in_len))) = (&arg_values[3], &arg_values[4]) {
// println!("{:x?}", &self.get_memory_slice(in_offset.clone(), in_len.clone()).unwrap());

// }

}
match builtin.name {
"create" | "sstore" | "sload" | "call" | "datacopy" | "mstore" => {
Expand Down Expand Up @@ -294,8 +299,17 @@ impl Evaluator {

fn read_memory(&self, offset: &BigUint) -> Option<Value> {
let candidate = self.memory.get(offset);
if let Some(ValueByte{data: candidate, ..}) = candidate {
if (0..32u32).all(|i| self.memory.get(&(offset.clone() + BigUint::from(i))) == Some(&ValueByte{data: candidate.clone(), byte: i}) ) {
if let Some(ValueByte {
data: candidate, ..
}) = candidate
{
if (0..32u32).all(|i| {
self.memory.get(&(offset.clone() + BigUint::from(i)))
== Some(&ValueByte {
data: candidate.clone(),
byte: i,
})
}) {
return Some(candidate.clone());
}
}
Expand All @@ -313,18 +327,30 @@ impl Evaluator {

fn read_memory_byte(&self, offset: &BigUint) -> Option<u8> {
match self.memory.get(offset) {
Some(ValueByte{data: Value::Concrete(v), byte: i}) => {
Some(v.to_bytes_le().get((31 - i) as usize).copied().unwrap_or_default())
}
_ => None
Some(ValueByte {
data: Value::Concrete(v),
byte: i,
}) => Some(
v.to_bytes_le()
.get((31 - i) as usize)
.copied()
.unwrap_or_default(),
),
_ => None,
}
}

fn write_memory(&mut self, offset: BigUint, value: Option<Value>) {
match value {
Some(value) => {
for i in 0..32u32 {
self.memory.insert(offset.clone() + BigUint::from(i), ValueByte{data: value.clone(), byte: i});
self.memory.insert(
offset.clone() + BigUint::from(i),
ValueByte {
data: value.clone(),
byte: i,
},
);
}
}
None => {
Expand All @@ -337,16 +363,16 @@ impl Evaluator {

fn get_memory_slice(&self, offset: BigUint, len: BigUint) -> Option<Vec<u8>> {
// TODO aliasing?
(0..(len.to_u64().unwrap())).map(|i| {
self.read_memory_byte(&(offset.clone() + BigUint::from(i)))
}).fold(Some(Vec::<u8>::new()), |acc, v| {
if let (Some(mut acc), Some(v)) = (acc, v) {
acc.push(v);
Some(acc)
} else {
None
}
})
(0..(len.to_u64().unwrap()))
.map(|i| self.read_memory_byte(&(offset.clone() + BigUint::from(i))))
.fold(Some(Vec::<u8>::new()), |acc, v| {
if let (Some(mut acc), Some(v)) = (acc, v) {
acc.push(v);
Some(acc)
} else {
None
}
})
}
}

Expand All @@ -355,3 +381,13 @@ fn wrap(mut x: BigUint) -> BigUint {
// TODO optimization: work directly on limbs
x & mask
}

fn literal_value(value: &String) -> Option<BigUint> {
if let Some(hex) = value.strip_prefix("0x") {
Some(BigUint::from_str_radix(hex, 16).unwrap())
} else if matches!(value.chars().next().unwrap(), '0'..='9') {
Some(BigUint::from_str_radix(value, 10).unwrap())
} else {
None
}
}
2 changes: 2 additions & 0 deletions tests/integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ fn recognize_known_contract() {
let content = fs::read_to_string(file).expect("I need to read this file.");
let main_ast = parse_and_resolve::<EVMInstructionsWithAssert>(&content, file);

println!("=========== SETUP ===================");
let mut evaluator = Evaluator::default();
evaluator.new_transaction(b"\x0a\x92\x54\xe4".to_vec());
// TODO also provide calldata to encoder?
Expand All @@ -71,6 +72,7 @@ fn recognize_known_contract() {
loop_unroll_default(&content),
evaluator,
);
println!("=========== CALL ===================");
evaluator.new_transaction(b"\x85\x63\x28\x95".to_vec());
encoder::encode_with_evaluator::<EVMInstructionsWithAssert>(
&main_ast,
Expand Down

0 comments on commit b57985c

Please sign in to comment.