From c21c81abe3ac317f90b8bf482de37dc87bc96a20 Mon Sep 17 00:00:00 2001
From: Svyatoslav Nikolsky <svyatonik@gmail.com>
Date: Wed, 23 Aug 2017 12:24:09 +0300
Subject: [PATCH] rollback command implemented

---
 db/src/block_chain.rs     |  3 +++
 db/src/block_chain_db.rs  | 28 ++++++++++++++++++++++++++++
 pbtc/cli.yml              |  6 ++++++
 pbtc/commands/mod.rs      |  2 ++
 pbtc/commands/rollback.rs | 39 +++++++++++++++++++++++++++++++++++++++
 pbtc/main.rs              |  1 +
 6 files changed, 79 insertions(+)
 create mode 100644 pbtc/commands/rollback.rs

diff --git a/db/src/block_chain.rs b/db/src/block_chain.rs
index 8a7e5f21..3f589913 100644
--- a/db/src/block_chain.rs
+++ b/db/src/block_chain.rs
@@ -15,6 +15,9 @@ pub trait BlockChain {
 	/// Inserts new block into blockchain
 	fn insert(&self, block: IndexedBlock) -> Result<(), Error>;
 
+	/// Rollbacks single best block. Returns new best block hash
+	fn rollback_best(&self) -> Result<H256, Error>;
+
 	/// Canonizes block with given hash
 	fn canonize(&self, block_hash: &H256) -> Result<(), Error>;
 
diff --git a/db/src/block_chain_db.rs b/db/src/block_chain_db.rs
index 7f58277e..7803637b 100644
--- a/db/src/block_chain_db.rs
+++ b/db/src/block_chain_db.rs
@@ -221,6 +221,30 @@ impl<T> BlockChainDatabase<T> where T: KeyValueDatabase {
 		self.db.write(update).map_err(Error::DatabaseError)
 	}
 
+	/// Rollbacks single best block
+	fn rollback_best(&self) -> Result<H256, Error> {
+		let decanonized = match self.block(self.best_block.read().hash.clone().into()) {
+			Some(block) => block,
+			None => return Ok(H256::default()),
+		};
+		let decanonized_hash = self.decanonize()?;
+		debug_assert_eq!(decanonized.hash(), decanonized_hash);
+
+		// and now remove decanonized block from database
+		// all code currently works in assumption that origin of all blocks is one of:
+		// {CanonChain, SideChain, SideChainBecomingCanonChain}
+		let mut update = DBTransaction::new();
+		update.delete(Key::BlockHeader(decanonized_hash.clone()));
+		update.delete(Key::BlockTransactions(decanonized_hash.clone()));
+		for tx in decanonized.transactions.into_iter() {
+			update.delete(Key::Transaction(tx.hash()));
+		}
+
+		self.db.write(update).map_err(Error::DatabaseError)?;
+
+		Ok(self.best_block().hash)
+	}
+
 	/// Marks block as a new best block.
 	/// Block must be already inserted into db, and it's parent must be current best block.
 	/// Updates meta data.
@@ -489,6 +513,10 @@ impl<T> BlockChain for BlockChainDatabase<T> where T: KeyValueDatabase {
 		BlockChainDatabase::insert(self, block)
 	}
 
+	fn rollback_best(&self) -> Result<H256, Error> {
+		BlockChainDatabase::rollback_best(self)
+	}
+
 	fn canonize(&self, block_hash: &H256) -> Result<(), Error> {
 		BlockChainDatabase::canonize(self, block_hash)
 	}
diff --git a/pbtc/cli.yml b/pbtc/cli.yml
index add1feb1..b62ce6fd 100644
--- a/pbtc/cli.yml
+++ b/pbtc/cli.yml
@@ -101,3 +101,9 @@ subcommands:
             - PATH:
                 required: true
                 help: Path of the bitcoin core database
+    - rollback:
+        about: Rollback database to given canon-chain block
+        args:
+            - BLOCK:
+                required: true
+                help: Either block hash, or block number 
diff --git a/pbtc/commands/mod.rs b/pbtc/commands/mod.rs
index 3713456a..c5b2e0b7 100644
--- a/pbtc/commands/mod.rs
+++ b/pbtc/commands/mod.rs
@@ -1,5 +1,7 @@
 mod import;
 mod start;
+mod rollback;
 
 pub use self::import::import;
 pub use self::start::start;
+pub use self::rollback::rollback;
\ No newline at end of file
diff --git a/pbtc/commands/rollback.rs b/pbtc/commands/rollback.rs
new file mode 100644
index 00000000..fc13a2c8
--- /dev/null
+++ b/pbtc/commands/rollback.rs
@@ -0,0 +1,39 @@
+use clap::ArgMatches;
+use db::BlockRef;
+use config::Config;
+use primitives::hash::H256;
+use util::{open_db, init_db};
+
+pub fn rollback(cfg: Config, matches: &ArgMatches) -> Result<(), String> {
+	let db = open_db(&cfg);
+	try!(init_db(&cfg, &db));
+
+	let block_ref = matches.value_of("BLOCK").expect("BLOCK is required in cli.yml; qed");
+	let block_ref = if block_ref.len() == 64 {
+		BlockRef::Hash({
+			let hash: H256 = block_ref.parse().map_err(|e| format!("Invalid block number: {}", e))?;
+			hash.reversed()
+		})
+	} else {
+		BlockRef::Number(block_ref.parse().map_err(|e| format!("Invalid block hash: {}", e))?)
+	};
+
+	let required_block_hash = db.block_header(block_ref.clone()).ok_or(format!("Block {:?} is unknown", block_ref))?.hash();
+	let genesis_hash = cfg.magic.genesis_block().hash();
+
+	let mut best_block_hash = db.best_block().hash;
+	debug_assert!(best_block_hash != H256::default()); // genesis inserted in init_db
+
+	loop {
+		if best_block_hash == required_block_hash {
+			info!("Reverted to block {:?}", block_ref);
+			return Ok(());
+		} 
+
+		if best_block_hash == genesis_hash {
+			return Err(format!("Failed to revert to block {:?}. Reverted to genesis", block_ref));
+		}
+
+		best_block_hash = db.rollback_best().map_err(|e| format!("{:?}", e))?;
+	}
+}
diff --git a/pbtc/main.rs b/pbtc/main.rs
index 4c0f76be..39613c31 100644
--- a/pbtc/main.rs
+++ b/pbtc/main.rs
@@ -64,6 +64,7 @@ fn run() -> Result<(), String> {
 
 	match matches.subcommand() {
 		("import", Some(import_matches)) => commands::import(cfg, import_matches),
+		("rollback", Some(rollback_matches)) => commands::rollback(cfg, rollback_matches),
 		_ => commands::start(cfg),
 	}
 }