Fix for UniBTC hack (#542)

* fix for unibtc

* fix

---------

Co-authored-by: dev4 <dev4@fuzzland>

onchain_scripts/debug_file_generator.py

@@ -119,6 +119,8 @@ def generate_debug_file(target, data):
             contract = w3.eth.contract(address=i["target"], abi=abi)
             abi_encoded = contract.encodeABI(fn_name=i["name"], args=i["args"]).replace("0x", "")
             current["direct_data"] = abi_encoded
+        elif "direct_data" in i:
+            current["direct_data"] = i["direct_data"]
         current["input_type"] = "ABI" if ("ty" not in i or i["ty"] == "abi") else "Borrow"
         current["caller"] = i["caller"]
         current["contract"] = i["target"]
@@ -218,4 +220,25 @@ def generate_debug_file(target, data):
 ]
 
 
-generate_debug_file("bsc", AES_DEFLATE)
+BTC_MINTER = Web3.to_checksum_address("0x047d41f2544b7f63a8e991af2068a363d210d6da")
+
+UNIBTC = [
+    {
+        "ty": "abi",
+        "caller": ATTACKER,
+        "target": BTC_MINTER,
+        "direct_data": "1249c58b",
+        "value": int(1e18),
+        "liquidation_percent": 10,  # sell 100% of the tokens
+    }, 
+    {
+        "ty": "abi",
+        "caller": ATTACKER,
+        "target": BTC_MINTER,
+        "direct_data": "1249c58b",
+        "value": int(1e18),
+        "liquidation_percent": 10,  # sell 100% of the tokens
+    }
+]
+
+generate_debug_file("eth", UNIBTC)

src/evm/corpus_initializer.rs

@@ -116,11 +116,11 @@ impl ABIMap {
 
 #[macro_export]
 macro_rules! handle_contract_insertion {
-    ($state: expr, $host: expr, $deployed_address: expr, $abi: expr) => {
+    ($state: expr, $host: expr, $deployed_address: expr, $impl_address: expr, $abi: expr) => {
         let (is_erc20, is_pair) = match $host.flashloan_middleware {
             Some(ref middleware) => {
                 let mut mid = middleware.deref().borrow_mut();
-                mid.on_contract_insertion(&$deployed_address, &$abi, $state)
+                mid.on_contract_insertion(&$deployed_address, &$impl_address, &$abi, $state)
             }
             None => (false, false),
         };
@@ -387,6 +387,7 @@ where
                     self.state,
                     self.executor.host,
                     contract.deployed_address,
+                    contract.deployed_address,
                     contract.abi.clone()
                 );
             }

src/evm/host.rs

@@ -1409,7 +1409,7 @@ where
                         parsed_abi = abis;
                     }
                     // notify flashloan and blacklisting flashloan addresses
-                    handle_contract_insertion!(state, self, r_addr, parsed_abi);
+                    handle_contract_insertion!(state, self, r_addr, r_addr, parsed_abi);
 
                     parsed_abi.iter().filter(|v| !v.is_constructor).for_each(|abi| {
                         #[cfg(not(feature = "fuzz_static"))]

src/evm/input.rs

@@ -46,7 +46,7 @@ pub enum EVMInputTy {
     Liquidate,
 }
 
-const CALL_VALUE_MAX_BYTES: usize = 21; // 309M ether
+const CALL_VALUE_MAX_BYTES: usize = 23; // 309M ether
 
 /// EVM Input Trait
 pub trait EVMInputT {
@@ -884,7 +884,7 @@ impl EVMInput {
         }
         add_mutator!(caller);
         add_mutator!(balance, !ap.balance.is_empty());
-        if self.get_txn_value().is_some() {
+        if ap.call_value || self.get_txn_value().is_some() {
             mutators.push(&EVMInput::call_value as &dyn Fn(&mut EVMInput, &mut S) -> MutationResult);
         }
         add_mutator!(gas_price);

src/evm/mutator.rs

@@ -95,6 +95,7 @@ impl AccessPattern {
             0x45 => self.gas_limit = true,
             0x46 => self.chain_id = true,
             0x48 => self.basefee = true,
+            0x34 => self.call_value = true,
             _ => {}
         }
     }

src/evm/onchain/flashloan.rs

@@ -140,14 +140,15 @@ impl Flashloan {
     pub fn on_contract_insertion(
         &mut self,
         addr: &EVMAddress,
+        impl_addr: &EVMAddress,
         abi: &[ABIConfig],
         _state: &mut EVMFuzzState,
     ) -> (bool, bool) {
         // should not happen, just sanity check
-        if self.known_addresses.contains(addr) {
+        if self.known_addresses.contains(impl_addr) {
             return (false, false);
         }
-        self.known_addresses.insert(*addr);
+        self.known_addresses.insert(*impl_addr);
 
         // balanceOf(address) - 70a08231
         // allowance(address,address) - dd62ed3e

src/evm/onchain/mod.rs

@@ -457,6 +457,7 @@ impl OnChain {
                 state,
                 host,
                 target,
+                address_h160,
                 parsed_abi
                     .iter()
                     .filter(|x| abi_hashes_to_add.contains(&x.function))

src/evm/vm.rs

@@ -677,15 +677,18 @@ where
             if tx_value > EVMU256::ZERO {
                 let caller_balance = *vm_state.get_balance(&input.get_caller()).unwrap_or(&EVMU256::ZERO);
                 let contract_balance = *vm_state.get_balance(&input.get_contract()).unwrap_or(&EVMU256::ZERO);
-                if caller_balance < tx_value {
-                    return ExecutionResult {
-                        output: vec![],
-                        reverted: true,
-                        new_state: StagedVMState::new_uninitialized(),
-                        additional_info: None,
-                    };
+                if !state.has_caller(&input.get_caller()) {
+                    if caller_balance < tx_value {
+                        return ExecutionResult {
+                            output: vec![],
+                            reverted: true,
+                            new_state: StagedVMState::new_uninitialized(),
+                            additional_info: None,
+                        };
+                    }
+                    vm_state.set_balance(input.get_caller(), caller_balance - tx_value);
                 }
-                vm_state.set_balance(input.get_caller(), caller_balance - tx_value);
+
                 vm_state.set_balance(input.get_contract(), contract_balance + tx_value);
             }
         }