Rollback latest EVM block (#2002)

lib/ain-evm/src/block.rs

@@ -30,7 +30,7 @@ impl BlockHandler {
     }
 
     pub fn connect_block(&self, block: BlockAny) {
-        self.storage.put_latest_block(&block);
+        self.storage.put_latest_block(Some(&block));
         self.storage.put_block(&block);
     }
 }

lib/ain-evm/src/handler.rs

@@ -171,8 +171,8 @@ impl Handlers {
 
         if update_state {
             debug!(
-                "[finalize_block] Updating state with new state_root : {:#x}",
-                block.header.state_root
+                "[finalize_block] Finalizing block number {:#x}, state_root {:#x}",
+                block.header.number, block.header.state_root
             );
             self.block.connect_block(block.clone());
             self.receipt.put_receipts(receipts);

lib/ain-evm/src/receipt.rs

@@ -33,7 +33,7 @@ fn get_contract_address(sender: &H160, nonce: &U256) -> H160 {
     stream.append(sender);
     stream.append(nonce);
 
-    return H160::from(keccak(stream.as_raw()));
+    H160::from(keccak(stream.as_raw()))
 }
 
 impl ReceiptHandler {
@@ -81,8 +81,7 @@ impl ReceiptHandler {
                     from: signed_tx.sender,
                     to: signed_tx.to(),
                     tx_index: index,
-                    tx_type: EnvelopedEncodable::type_id(&signed_tx.transaction)
-                        .unwrap_or_default(),
+                    tx_type: signed_tx.transaction.type_id().unwrap_or_default(),
                     contract_address: signed_tx
                         .to()
                         .is_none()

lib/ain-evm/src/storage/cache.rs

@@ -5,7 +5,7 @@ use lru::LruCache;
 use primitive_types::{H256, U256};
 use std::borrow::ToOwned;
 
-use super::traits::{BlockStorage, TransactionStorage};
+use super::traits::{BlockStorage, Rollback, TransactionStorage};
 
 #[derive(Debug)]
 pub struct Cache {
@@ -71,9 +71,9 @@ impl BlockStorage for Cache {
             .map(ToOwned::to_owned)
     }
 
-    fn put_latest_block(&self, block: &BlockAny) {
+    fn put_latest_block(&self, block: Option<&BlockAny>) {
         let mut cache = self.latest_block.write().unwrap();
-        *cache = Some(block.clone());
+        *cache = block.cloned();
     }
 }
 
@@ -130,3 +130,19 @@ impl TransactionStorage for Cache {
             .put(transaction.hash(), transaction.clone());
     }
 }
+
+impl Rollback for Cache {
+    fn disconnect_latest_block(&self) {
+        if let Some(block) = self.get_latest_block() {
+            let mut transaction_cache = self.transactions.write().unwrap();
+            for tx in &block.transactions {
+                transaction_cache.pop(&tx.hash());
+            }
+
+            self.block_hashes.write().unwrap().pop(&block.header.hash());
+            self.blocks.write().unwrap().pop(&block.header.number);
+
+            self.put_latest_block(self.get_block_by_hash(&block.header.parent_hash).as_ref())
+        }
+    }
+}

lib/ain-evm/src/storage/code.rs

@@ -0,0 +1,43 @@
+use primitive_types::{H256, U256};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+use super::traits::PersistentState;
+
+/// `CodeHistory` maintains a history of accounts' codes.
+///
+/// It tracks the current state (`code_map`), as well as a history (`history`) of code hashes
+/// that should be removed if a specific block is rolled back. The correct account code_hash
+/// is tracked by the state trie.
+/// This structure is solely required for rolling back and preventing ghost entries.
+#[derive(Default, Debug, Serialize, Deserialize)]
+pub struct CodeHistory {
+    /// The current state of each code
+    code_map: HashMap<H256, Vec<u8>>,
+    /// A map from block number to a vector of code hashes to remove for that block.
+    history: HashMap<U256, Vec<H256>>,
+}
+
+impl PersistentState for CodeHistory {}
+
+impl CodeHistory {
+    pub fn insert(&mut self, block_number: U256, code_hash: H256, code: Vec<u8>) {
+        self.code_map.insert(code_hash, code.clone());
+        self.history
+            .entry(block_number)
+            .or_insert_with(Vec::new)
+            .push(code_hash);
+    }
+
+    pub fn get(&self, code_hash: &H256) -> Option<&Vec<u8>> {
+        self.code_map.get(code_hash)
+    }
+
+    pub fn rollback(&mut self, block_number: U256) {
+        if let Some(code_hashes) = self.history.remove(&block_number) {
+            for code_hash in &code_hashes {
+                self.code_map.remove(code_hash);
+            }
+        }
+    }
+}

lib/ain-evm/src/storage/data_handler.rs

@@ -6,30 +6,32 @@ use std::borrow::ToOwned;
 
 use crate::receipt::Receipt;
 
-use super::traits::{
-    BlockStorage, FlushableStorage, PersistentState, PersistentStateError, ReceiptStorage,
-    TransactionStorage,
+use super::{
+    code::CodeHistory,
+    traits::{
+        BlockStorage, FlushableStorage, PersistentState, PersistentStateError, ReceiptStorage,
+        Rollback, TransactionStorage,
+    },
 };
 
 pub static BLOCK_MAP_PATH: &str = "block_map.bin";
 pub static BLOCK_DATA_PATH: &str = "block_data.bin";
 pub static LATEST_BLOCK_DATA_PATH: &str = "latest_block_data.bin";
 pub static RECEIPT_MAP_PATH: &str = "receipt_map.bin";
 pub static CODE_MAP_PATH: &str = "code_map.bin";
-// pub static TRANSACTION_DATA_PATH: &str = "transaction_data.bin";
+pub static TRANSACTION_DATA_PATH: &str = "transaction_data.bin";
 
 type BlockHashtoBlock = HashMap<H256, U256>;
 type Blocks = HashMap<U256, BlockAny>;
 type TxHashToTx = HashMap<H256, TransactionV2>;
 type LatestBlockNumber = U256;
 type TransactionHashToReceipt = HashMap<H256, Receipt>;
-type CodeHashToCode = HashMap<H256, Vec<u8>>;
 
 impl PersistentState for BlockHashtoBlock {}
 impl PersistentState for Blocks {}
 impl PersistentState for LatestBlockNumber {}
 impl PersistentState for TransactionHashToReceipt {}
-impl PersistentState for CodeHashToCode {}
+impl PersistentState for TxHashToTx {}
 
 #[derive(Debug)]
 pub struct BlockchainDataHandler {
@@ -42,28 +44,32 @@ pub struct BlockchainDataHandler {
     blocks: RwLock<Blocks>,
     latest_block_number: RwLock<Option<LatestBlockNumber>>,
 
-    code_map: RwLock<CodeHashToCode>,
+    code_map: RwLock<CodeHistory>,
 }
 
 impl BlockchainDataHandler {
     pub fn new() -> Self {
-        let blocks = Blocks::load_from_disk(BLOCK_DATA_PATH).expect("Error loading blocks data");
         BlockchainDataHandler {
-            transactions: RwLock::new(HashMap::new()),
+            transactions: RwLock::new(
+                TxHashToTx::load_from_disk(TRANSACTION_DATA_PATH)
+                    .expect("Error loading blocks data"),
+            ),
             block_map: RwLock::new(
                 BlockHashtoBlock::load_from_disk(BLOCK_MAP_PATH)
                     .expect("Error loading block_map data"),
             ),
             latest_block_number: RwLock::new(
                 LatestBlockNumber::load_from_disk(LATEST_BLOCK_DATA_PATH).ok(),
             ),
-            blocks: RwLock::new(blocks),
+            blocks: RwLock::new(
+                Blocks::load_from_disk(BLOCK_DATA_PATH).expect("Error loading blocks data"),
+            ),
             receipts: RwLock::new(
                 TransactionHashToReceipt::load_from_disk(RECEIPT_MAP_PATH)
                     .expect("Error loading receipts data"),
             ),
             code_map: RwLock::new(
-                CodeHashToCode::load_from_disk(CODE_MAP_PATH).expect("Error loading code data"),
+                CodeHistory::load_from_disk(CODE_MAP_PATH).expect("Error loading code data"),
             ),
         }
     }
@@ -80,13 +86,12 @@ impl TransactionStorage for BlockchainDataHandler {
         }
     }
 
-    fn get_transaction_by_hash(&self, _hash: &H256) -> Option<TransactionV2> {
-        None
-        // self.transactions
-        //     .read()
-        //     .unwrap()
-        //     .get(hash)
-        //     .map(ToOwned::to_owned)
+    fn get_transaction_by_hash(&self, hash: &H256) -> Option<TransactionV2> {
+        self.transactions
+            .read()
+            .unwrap()
+            .get(hash)
+            .map(ToOwned::to_owned)
     }
 
     fn get_transaction_by_block_hash_and_index(
@@ -162,9 +167,9 @@ impl BlockStorage for BlockchainDataHandler {
             .and_then(|number| self.get_block_by_number(number))
     }
 
-    fn put_latest_block(&self, block: &BlockAny) {
+    fn put_latest_block(&self, block: Option<&BlockAny>) {
         let mut latest_block_number = self.latest_block_number.write().unwrap();
-        *latest_block_number = Some(block.header.number);
+        *latest_block_number = block.map(|b| b.header.number);
     }
 }
 
@@ -197,6 +202,10 @@ impl FlushableStorage for BlockchainDataHandler {
             .write()
             .unwrap()
             .save_to_disk(RECEIPT_MAP_PATH)?;
+        self.transactions
+            .write()
+            .unwrap()
+            .save_to_disk(TRANSACTION_DATA_PATH)?;
         self.code_map.write().unwrap().save_to_disk(CODE_MAP_PATH)
     }
 }
@@ -210,7 +219,36 @@ impl BlockchainDataHandler {
             .map(ToOwned::to_owned)
     }
 
-    pub fn put_code(&self, hash: &H256, code: &[u8]) -> Option<Vec<u8>> {
-        self.code_map.write().unwrap().insert(*hash, code.to_vec())
+    pub fn put_code(&self, hash: &H256, code: &[u8]) {
+        let block_number = self
+            .get_latest_block()
+            .map(|b| b.header.number)
+            .unwrap_or_default()
+            + 1;
+        self.code_map
+            .write()
+            .unwrap()
+            .insert(block_number, *hash, code.to_vec())
+    }
+}
+
+impl Rollback for BlockchainDataHandler {
+    fn disconnect_latest_block(&self) {
+        if let Some(block) = self.get_latest_block() {
+            println!("disconnecting block number : {:x?}", block.header.number);
+            let mut transactions = self.transactions.write().unwrap();
+            let mut receipts = self.receipts.write().unwrap();
+            for tx in &block.transactions {
+                let hash = &tx.hash();
+                transactions.remove(hash);
+                receipts.remove(hash);
+            }
+
+            self.block_map.write().unwrap().remove(&block.header.hash());
+            self.blocks.write().unwrap().remove(&block.header.number);
+            self.code_map.write().unwrap().rollback(block.header.number);
+
+            self.put_latest_block(self.get_block_by_hash(&block.header.parent_hash).as_ref())
+        }
     }
 }

lib/ain-evm/src/storage/mod.rs

@@ -1,4 +1,5 @@
 mod cache;
+mod code;
 mod data_handler;
 pub mod traits;
 
@@ -11,7 +12,8 @@ use self::{
     cache::Cache,
     data_handler::BlockchainDataHandler,
     traits::{
-        BlockStorage, FlushableStorage, PersistentStateError, ReceiptStorage, TransactionStorage,
+        BlockStorage, FlushableStorage, PersistentStateError, ReceiptStorage, Rollback,
+        TransactionStorage,
     },
 };
 
@@ -66,13 +68,13 @@ impl BlockStorage for Storage {
         self.cache.get_latest_block().or_else(|| {
             let latest_block = self.blockchain_data_handler.get_latest_block();
             if let Some(ref block) = latest_block {
-                self.cache.put_latest_block(block);
+                self.cache.put_latest_block(Some(block));
             }
             latest_block
         })
     }
 
-    fn put_latest_block(&self, block: &BlockAny) {
+    fn put_latest_block(&self, block: Option<&BlockAny>) {
         self.cache.put_latest_block(block);
         self.blockchain_data_handler.put_latest_block(block);
     }
@@ -160,7 +162,7 @@ impl Storage {
         self.blockchain_data_handler.get_code_by_hash(&hash)
     }
 
-    pub fn put_code(&self, hash: H256, code: Vec<u8>) -> Option<Vec<u8>> {
+    pub fn put_code(&self, hash: H256, code: Vec<u8>) {
         self.blockchain_data_handler.put_code(&hash, &code)
     }
 }
@@ -173,3 +175,10 @@ impl Storage {
         );
     }
 }
+
+impl Rollback for Storage {
+    fn disconnect_latest_block(&self) {
+        self.cache.disconnect_latest_block();
+        self.blockchain_data_handler.disconnect_latest_block();
+    }
+}

lib/ain-evm/src/storage/traits.rs

@@ -17,7 +17,7 @@ pub trait BlockStorage {
     fn get_block_by_hash(&self, block_hash: &H256) -> Option<BlockAny>;
     fn put_block(&self, block: &BlockAny);
     fn get_latest_block(&self) -> Option<BlockAny>;
-    fn put_latest_block(&self, block: &BlockAny);
+    fn put_latest_block(&self, block: Option<&BlockAny>);
 }
 
 pub trait TransactionStorage {
@@ -45,6 +45,10 @@ pub trait FlushableStorage {
     fn flush(&self) -> Result<(), PersistentStateError>;
 }
 
+pub trait Rollback {
+    fn disconnect_latest_block(&self);
+}
+
 pub trait PersistentState {
     fn save_to_disk(&self, file_path: &str) -> Result<(), PersistentStateError>
     where

lib/ain-grpc/src/rpc/eth.rs

@@ -506,7 +506,7 @@ impl MetachainRPCServer for MetachainRPCModule {
             Some(TransactionMessage::Legacy(mut m)) => {
                 m.nonce = nonce;
                 m.chain_id = Some(chain_id);
-                m.gas_limit = U256::from(1);
+                m.gas_limit = gas_limit;
                 if gas_price.is_none() {
                     m.gas_price = self.gas_price().unwrap();
                 }

lib/ain-rs-exports/src/lib.rs

@@ -1,4 +1,7 @@
-use ain_evm::transaction::{self, SignedTx};
+use ain_evm::{
+    storage::traits::Rollback,
+    transaction::{self, SignedTx},
+};
 use ain_grpc::{init_evm_runtime, start_servers, stop_evm_runtime};
 
 use ain_evm::runtime::RUNTIME;
@@ -64,6 +67,8 @@ pub mod ffi {
         fn stop_evm_runtime();
 
         fn create_and_sign_tx(ctx: CreateTransactionContext) -> Result<Vec<u8>>;
+
+        fn evm_disconnect_latest_block() -> Result<()>;
     }
 }
 
@@ -332,3 +337,8 @@ fn evm_finalize(
 pub fn preinit() {
     ain_grpc::preinit();
 }
+
+fn evm_disconnect_latest_block() -> Result<(), Box<dyn Error>> {
+    RUNTIME.handlers.storage.disconnect_latest_block();
+    Ok(())
+}