From 6a9a584c43db85b311af6c800c2e5ab4f004f914 Mon Sep 17 00:00:00 2001
From: Ignotus Peverell <igno.peverell@protonmail.com>
Date: Mon, 8 Jan 2018 01:23:23 +0000
Subject: [PATCH] Fix pruning of last PMMR leaf, additional tests

Due to the construction of PMMRs the last element, when its a leaf,
can never be pruned as it has no parent yet and it will be needed
to calculate that hash. To work around this, we now insert coinbase
outputs first to add at least one output of padding.

Also changed the `set_sumtree_root` function on chain a bit to allow
setting the roots on a fork. Mostly useful for tests.

Added new test case to handle both the issue above and spending
transactions within a fork.
---
 chain/src/chain.rs                    |   8 +-
 chain/src/pipe.rs                     |  10 +-
 chain/src/sumtree.rs                  | 161 +++++++++++++++-----------
 chain/tests/mine_simple_chain.rs      | 134 +++++++++++++++++----
 chain/tests/test_coinbase_maturity.rs |   8 +-
 grin/src/miner.rs                     |   2 +-
 6 files changed, 224 insertions(+), 99 deletions(-)

diff --git a/chain/src/chain.rs b/chain/src/chain.rs
index 59e2efcb5..52c7bb8d9 100644
--- a/chain/src/chain.rs
+++ b/chain/src/chain.rs
@@ -358,11 +358,15 @@ impl Chain {
 
 	/// Sets the sumtree roots on a brand new block by applying the block on the
 	/// current sumtree state.
-	pub fn set_sumtree_roots(&self, b: &mut Block) -> Result<(), Error> {
+	pub fn set_sumtree_roots(&self, b: &mut Block, is_fork: bool) -> Result<(), Error> {
 		let mut sumtrees = self.sumtrees.write().unwrap();
+		let store = self.store.clone();
 
 		let roots = sumtree::extending(&mut sumtrees, |extension| {
 			// apply the block on the sumtrees and check the resulting root
+			if is_fork {
+				pipe::rewind_and_apply_fork(b, store, extension)?;
+			}
 			extension.apply_block(b)?;
 			extension.force_rollback();
 			Ok(extension.roots())
@@ -374,7 +378,7 @@ impl Chain {
 		Ok(())
 	}
 
-	/// returs sumtree roots
+	/// Returns current sumtree roots
 	pub fn get_sumtree_roots(
 		&self,
 	) -> (
diff --git a/chain/src/pipe.rs b/chain/src/pipe.rs
index 0ba3e9771..1c6804b1b 100644
--- a/chain/src/pipe.rs
+++ b/chain/src/pipe.rs
@@ -393,11 +393,11 @@ fn update_header_head(bh: &BlockHeader, ctx: &mut BlockContext) -> Result<Option
 	}
 }
 
-// Utility function to handle forks. From the forked block, jump backward
-// to find to fork root. Rewind the sumtrees to the root and apply all the
-// forked blocks prior to the one being processed to set the sumtrees in
-// the expected state.
-fn rewind_and_apply_fork(
+/// Utility function to handle forks. From the forked block, jump backward
+/// to find to fork root. Rewind the sumtrees to the root and apply all the
+/// forked blocks prior to the one being processed to set the sumtrees in
+/// the expected state.
+pub fn rewind_and_apply_fork(
 	b: &Block,
 	store: Arc<ChainStore>,
 	ext: &mut sumtree::Extension,
diff --git a/chain/src/sumtree.rs b/chain/src/sumtree.rs
index fa16fb5c4..6e96c9197 100644
--- a/chain/src/sumtree.rs
+++ b/chain/src/sumtree.rs
@@ -22,7 +22,7 @@ use std::sync::Arc;
 
 use util::secp::pedersen::{RangeProof, Commitment};
 
-use core::core::{Block, SumCommit, TxKernel};
+use core::core::{Block, SumCommit, Input, Output, TxKernel, COINBASE_OUTPUT};
 use core::core::pmmr::{HashSum, NoSum, Summable, PMMR};
 use core::core::hash::Hashed;
 use grin_store;
@@ -244,79 +244,32 @@ impl<'a> Extension<'a> {
 	/// prune MMR data.
 	pub fn apply_block(&mut self, b: &Block) -> Result<(), Error> {
 
-		// doing inputs first guarantees an input can't spend an output in the
+		// first applying coinbase outputs. due to the construction of PMMRs the
+		// last element, when its a leaf, can never be pruned as it has no parent
+		// yet and it will be needed to calculate that hash. to work around this,
+		// we insert coinbase outputs first to add at least one output of padding
+		for out in &b.outputs {
+			if out.features.contains(COINBASE_OUTPUT) {
+				self.apply_output(out)?;
+			}
+		}
+	
+		// then doing inputsm guarantees an input can't spend an output in the
 		// same block, enforcing block cut-through
 		for input in &b.inputs {
-			let commit = input.commitment();
-			let pos_res = self.get_output_pos(&commit);
-			if let Ok(pos) = pos_res {
-				match self.output_pmmr.prune(pos, b.header.height as u32) {
-					Ok(true) => {
-						self.rproof_pmmr
-							.prune(pos, b.header.height as u32)
-							.map_err(|s| Error::SumTreeErr(s))?;
-					}
-					Ok(false) => return Err(Error::AlreadySpent(commit)),
-					Err(s) => return Err(Error::SumTreeErr(s)),
-				}
-			} else {
-				return Err(Error::AlreadySpent(commit));
-			}
+			self.apply_input(input, b.header.height)?;
 		}
 
+		// now all regular, non coinbase outputs
 		for out in &b.outputs {
-			let commit = out.commitment();
-			let switch_commit_hash = out.switch_commit_hash();
-			let sum_commit = SumCommit {
-				commit,
-				switch_commit_hash,
-			};
-
-			if let Ok(pos) = self.get_output_pos(&commit) {
-				// we need to check whether the commitment is in the current MMR view
-				// as well as the index doesn't support rewind and is non-authoritative
-				// (non-historical node will have a much smaller one)
-				// note that this doesn't show the commitment *never* existed, just
-				// that this is not an existing unspent commitment right now
-				if let Some(c) = self.output_pmmr.get(pos) {
-					let hashsum = HashSum::from_summable(pos, &sum_commit);
-
-					// processing a new fork so we may get a position on the old
-					// fork that exists but matches a different node
-					// filtering that case out
-					if c.hash == hashsum.hash {
-						return Err(Error::DuplicateCommitment(commit));
-					}
-				}
+			if !out.features.contains(COINBASE_OUTPUT) {
+				self.apply_output(out)?;
 			}
-			// push new outputs commitments in their MMR and save them in the index
-			let pos = self.output_pmmr
-				.push(sum_commit)
-				.map_err(&Error::SumTreeErr)?;
-
-			self.new_output_commits.insert(out.commitment(), pos);
-
-			// push range proofs in their MMR
-			self.rproof_pmmr
-				.push(NoSum(out.proof))
-				.map_err(&Error::SumTreeErr)?;
 		}
-
+	
+		// finally, applying all kernels
 		for kernel in &b.kernels {
-			if let Ok(pos) = self.get_kernel_pos(&kernel.excess) {
-				// same as outputs
-				if let Some(k) = self.kernel_pmmr.get(pos) {
-					let hashsum = HashSum::from_summable(pos, &NoSum(kernel));
-					if k.hash == hashsum.hash {
-						return Err(Error::DuplicateKernel(kernel.excess.clone()));
-					}
-				}
-			}
-			// push kernels in their MMR
-			let pos = self.kernel_pmmr
-				.push(NoSum(kernel.clone()))
-				.map_err(&Error::SumTreeErr)?;
-			self.new_kernel_excesses.insert(kernel.excess, pos);
+			self.apply_kernel(kernel)?;
 		}
 		Ok(())
 	}
@@ -347,6 +300,82 @@ impl<'a> Extension<'a> {
 		Ok(())
 	}
 
+	fn apply_input(&mut self, input: &Input, height: u64) -> Result<(), Error> {
+		let commit = input.commitment();
+		let pos_res = self.get_output_pos(&commit);
+		if let Ok(pos) = pos_res {
+			match self.output_pmmr.prune(pos, height as u32) {
+				Ok(true) => {
+					self.rproof_pmmr
+						.prune(pos, height as u32)
+						.map_err(|s| Error::SumTreeErr(s))?;
+				}
+				Ok(false) => return Err(Error::AlreadySpent(commit)),
+				Err(s) => return Err(Error::SumTreeErr(s)),
+			}
+		} else {
+			return Err(Error::AlreadySpent(commit));
+		}
+		Ok(())
+	}
+
+	fn apply_output(&mut self, out: &Output) -> Result<(), Error> {
+		let commit = out.commitment();
+		let switch_commit_hash = out.switch_commit_hash();
+		let sum_commit = SumCommit {
+			commit,
+			switch_commit_hash,
+		};
+
+		if let Ok(pos) = self.get_output_pos(&commit) {
+			// we need to check whether the commitment is in the current MMR view
+			// as well as the index doesn't support rewind and is non-authoritative
+			// (non-historical node will have a much smaller one)
+			// note that this doesn't show the commitment *never* existed, just
+			// that this is not an existing unspent commitment right now
+			if let Some(c) = self.output_pmmr.get(pos) {
+				let hashsum = HashSum::from_summable(pos, &sum_commit);
+
+				// processing a new fork so we may get a position on the old
+				// fork that exists but matches a different node
+				// filtering that case out
+				if c.hash == hashsum.hash {
+					return Err(Error::DuplicateCommitment(commit));
+				}
+			}
+		}
+		// push new outputs commitments in their MMR and save them in the index
+		let pos = self.output_pmmr
+			.push(sum_commit)
+			.map_err(&Error::SumTreeErr)?;
+
+		self.new_output_commits.insert(out.commitment(), pos);
+
+		// push range proofs in their MMR
+		self.rproof_pmmr
+			.push(NoSum(out.proof))
+			.map_err(&Error::SumTreeErr)?;
+		Ok(())
+	}
+
+	fn apply_kernel(&mut self, kernel: &TxKernel) -> Result<(), Error> {
+		if let Ok(pos) = self.get_kernel_pos(&kernel.excess) {
+			// same as outputs
+			if let Some(k) = self.kernel_pmmr.get(pos) {
+				let hashsum = HashSum::from_summable(pos, &NoSum(kernel));
+				if k.hash == hashsum.hash {
+					return Err(Error::DuplicateKernel(kernel.excess.clone()));
+				}
+			}
+		}
+		// push kernels in their MMR
+		let pos = self.kernel_pmmr
+			.push(NoSum(kernel.clone()))
+			.map_err(&Error::SumTreeErr)?;
+		self.new_kernel_excesses.insert(kernel.excess, pos);
+		Ok(())
+	}
+
 	/// Rewinds the MMRs to the provided position, given the last output and
 	/// last kernel of the block we want to rewind to.
 	pub fn rewind(&mut self, block: &Block) -> Result<(), Error> {
diff --git a/chain/tests/mine_simple_chain.rs b/chain/tests/mine_simple_chain.rs
index 499aa407a..448bc6c49 100644
--- a/chain/tests/mine_simple_chain.rs
+++ b/chain/tests/mine_simple_chain.rs
@@ -17,6 +17,7 @@ extern crate grin_chain as chain;
 extern crate grin_core as core;
 extern crate grin_keychain as keychain;
 extern crate grin_pow as pow;
+extern crate grin_util as util;
 extern crate rand;
 extern crate time;
 
@@ -25,7 +26,7 @@ use std::sync::Arc;
 
 use chain::Chain;
 use chain::types::*;
-use core::core::{Block, BlockHeader};
+use core::core::{Block, BlockHeader, Transaction, build};
 use core::core::hash::Hashed;
 use core::core::target::Difficulty;
 use core::consensus;
@@ -79,7 +80,7 @@ fn mine_empty_chain() {
 
 		let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
 		b.header.difficulty = difficulty.clone();
-		chain.set_sumtree_roots(&mut b).unwrap();
+		chain.set_sumtree_roots(&mut b, false).unwrap();
 
 		pow::pow_size(
 			&mut cuckoo_miner,
@@ -123,10 +124,11 @@ fn mine_empty_chain() {
 #[test]
 fn mine_forks() {
 	let chain = setup(".grin2");
+	let kc = Keychain::from_random_seed().unwrap();
 
 	// add a first block to not fork genesis
 	let prev = chain.head_header().unwrap();
-	let b = prepare_block(&prev, &chain, 2);
+	let b = prepare_block(&kc, &prev, &chain, 2);
 	chain.process_block(b, chain::SKIP_POW).unwrap();
 
 	// mine and add a few blocks
@@ -134,10 +136,10 @@ fn mine_forks() {
 	for n in 1..4 {
 		// first block for one branch
 		let prev = chain.head_header().unwrap();
-		let b1 = prepare_block(&prev, &chain, 3 * n);
+		let b1 = prepare_block(&kc, &prev, &chain, 3 * n);
 
 		// 2nd block with higher difficulty for other branch
-		let b2 = prepare_block(&prev, &chain, 3 * n + 1);
+		let b2 = prepare_block(&kc, &prev, &chain, 3 * n + 1);
 
 		// process the first block to extend the chain
 		let bhash = b1.hash();
@@ -163,24 +165,25 @@ fn mine_forks() {
 
 #[test]
 fn mine_losing_fork() {
+	let kc = Keychain::from_random_seed().unwrap();
 	let chain = setup(".grin3");
 
 	// add a first block we'll be forking from
 	let prev = chain.head_header().unwrap();
-	let b1 = prepare_block(&prev, &chain, 2);
+	let b1 = prepare_block(&kc, &prev, &chain, 2);
 	let b1head = b1.header.clone();
 	chain.process_block(b1, chain::SKIP_POW).unwrap();
 
 	// prepare the 2 successor, sibling blocks, one with lower diff
-	let b2 = prepare_block(&b1head, &chain, 4);
+	let b2 = prepare_block(&kc, &b1head, &chain, 4);
 	let b2head = b2.header.clone();
-	let bfork = prepare_block(&b1head, &chain, 3);
+	let bfork = prepare_block(&kc, &b1head, &chain, 3);
 
 	// add higher difficulty first, prepare its successor, then fork
  // with lower diff
 	chain.process_block(b2, chain::SKIP_POW).unwrap();
 	assert_eq!(chain.head_header().unwrap().hash(), b2head.hash());
-	let b3 = prepare_block(&b2head, &chain, 5);
+	let b3 = prepare_block(&kc, &b2head, &chain, 5);
 	chain.process_block(bfork, chain::SKIP_POW).unwrap();
 
 	// adding the successor
@@ -191,17 +194,18 @@ fn mine_losing_fork() {
 
 #[test]
 fn longer_fork() {
+	let kc = Keychain::from_random_seed().unwrap();
 	// to make it easier to compute the sumtree roots in the test, we
- // prepare 2 chains, the 2nd will be have the forked blocks we can
- // then send back on the 1st
+	// prepare 2 chains, the 2nd will be have the forked blocks we can
+	// then send back on the 1st
 	let chain = setup(".grin4");
 	let chain_fork = setup(".grin5");
 
 	// add blocks to both chains, 20 on the main one, only the first 5
- // for the forked chain
+	// for the forked chain
 	let mut prev = chain.head_header().unwrap();
 	for n in 0..10 {
-		let b = prepare_block(&prev, &chain, n + 2);
+		let b = prepare_block(&kc, &prev, &chain, 2*n + 2);
 		let bh = b.header.clone();
 
 		if n < 5 {
@@ -222,7 +226,7 @@ fn longer_fork() {
 
 	let mut prev_fork = head_fork.clone();
 	for n in 0..7 {
-		let b_fork = prepare_block(&prev_fork, &chain_fork, n + 7);
+		let b_fork = prepare_block(&kc, &prev_fork, &chain_fork, 2*n + 11);
 		let bh_fork = b_fork.header.clone();
 
 		let b = b_fork.clone();
@@ -233,17 +237,105 @@ fn longer_fork() {
 	}
 }
 
-fn prepare_block(prev: &BlockHeader, chain: &Chain, diff: u64) -> Block {
-	let mut b = prepare_block_nosum(prev, diff);
-	chain.set_sumtree_roots(&mut b).unwrap();
+#[test]
+fn spend_in_fork() {
+	util::init_test_logger();
+	let chain = setup(".grin6");
+	let prev = chain.head_header().unwrap();
+	let kc = Keychain::from_random_seed().unwrap();
+
+	// mine 4 blocks, the 4th will be the root of the fork
+	let mut fork_head = prev;
+	for n in 2..6 {
+		let b = prepare_block(&kc, &fork_head, &chain, n);
+		fork_head = b.header.clone();
+		chain.process_block(b, chain::SKIP_POW).unwrap();
+	}
+
+	let (tx1, _) = build::transaction(
+		vec![
+			build::input(consensus::REWARD, kc.derive_key_id(2).unwrap()),
+			build::output(consensus::REWARD - 20000, kc.derive_key_id(30).unwrap()),
+			build::with_fee(20000),
+		],
+		&kc,
+	).unwrap();
+
+	let next = prepare_block_tx(&kc, &fork_head, &chain, 7, vec![&tx1]);
+	let prev_main = next.header.clone();
+	chain.process_block(next, chain::SKIP_POW).unwrap();
+
+	let (tx2, _) = build::transaction(
+		vec![
+			build::input(consensus::REWARD - 20000, kc.derive_key_id(30).unwrap()),
+			build::output(consensus::REWARD - 40000, kc.derive_key_id(31).unwrap()),
+			build::with_fee(20000),
+		],
+		&kc,
+	).unwrap();
+
+	let next = prepare_block_tx(&kc, &prev_main, &chain, 9, vec![&tx2]);
+	let prev_main = next.header.clone();
+	chain.process_block(next, chain::SKIP_POW).unwrap();
+
+	// mine 2 forked blocks from the first
+	let fork = prepare_fork_block_tx(&kc, &fork_head, &chain, 6, vec![&tx1]);
+	let prev_fork = fork.header.clone();
+	chain.process_block(fork, chain::SKIP_POW).unwrap();
+
+	let fork_next = prepare_fork_block_tx(&kc, &prev_fork, &chain, 8, vec![&tx2]);
+	let prev_fork = fork_next.header.clone();
+	chain.process_block(fork_next, chain::SKIP_POW).unwrap();
+
+	// check state
+	let head = chain.head_header().unwrap();
+	assert_eq!(head.height, 6);
+	assert_eq!(head.hash(), prev_main.hash());
+	assert!(chain.is_unspent(&tx2.outputs[0].commitment()).unwrap());
+	let res = chain.is_unspent(&tx1.outputs[0].commitment());
+	assert!(!res.unwrap());
+
+	// make the fork win
+	let fork_next = prepare_fork_block(&kc, &prev_fork, &chain, 10);
+	let prev_fork = fork_next.header.clone();
+	chain.process_block(fork_next, chain::SKIP_POW).unwrap();
+
+	// check state
+	let head = chain.head_header().unwrap();
+	assert_eq!(head.height, 7);
+	assert_eq!(head.hash(), prev_fork.hash());
+	assert!(chain.is_unspent(&tx2.outputs[0].commitment()).unwrap());
+	assert!(!chain.is_unspent(&tx1.outputs[0].commitment()).unwrap());
+}
+
+fn prepare_block(kc: &Keychain, prev: &BlockHeader, chain: &Chain, diff: u64) -> Block {
+	let mut b = prepare_block_nosum(kc, prev, diff, vec![]);
+	chain.set_sumtree_roots(&mut b, false).unwrap();
 	b
 }
 
-fn prepare_block_nosum(prev: &BlockHeader, diff: u64) -> Block {
-	let keychain = Keychain::from_random_seed().unwrap();
-	let key_id = keychain.derive_key_id(1).unwrap();
+fn prepare_block_tx(kc: &Keychain, prev: &BlockHeader, chain: &Chain, diff: u64, txs: Vec<&Transaction>) -> Block {
+	let mut b = prepare_block_nosum(kc, prev, diff, txs);
+	chain.set_sumtree_roots(&mut b, false).unwrap();
+	b
+}
 
-	let mut b = core::core::Block::new(prev, vec![], &keychain, &key_id).unwrap();
+fn prepare_fork_block(kc: &Keychain, prev: &BlockHeader, chain: &Chain, diff: u64) -> Block {
+	let mut b = prepare_block_nosum(kc, prev, diff, vec![]);
+	chain.set_sumtree_roots(&mut b, true).unwrap();
+	b
+}
+
+fn prepare_fork_block_tx(kc: &Keychain, prev: &BlockHeader, chain: &Chain, diff: u64, txs: Vec<&Transaction>) -> Block {
+	let mut b = prepare_block_nosum(kc, prev, diff, txs);
+	chain.set_sumtree_roots(&mut b, true).unwrap();
+	b
+}
+
+fn prepare_block_nosum(kc: &Keychain, prev: &BlockHeader, diff: u64, txs: Vec<&Transaction>) -> Block {
+	let key_id = kc.derive_key_id(diff as u32).unwrap();
+
+	let mut b = core::core::Block::new(prev, txs, kc, &key_id).unwrap();
 	b.header.timestamp = prev.timestamp + time::Duration::seconds(60);
 	b.header.total_difficulty = Difficulty::from_num(diff);
 	b
diff --git a/chain/tests/test_coinbase_maturity.rs b/chain/tests/test_coinbase_maturity.rs
index 60c1bb085..dab4bea5f 100644
--- a/chain/tests/test_coinbase_maturity.rs
+++ b/chain/tests/test_coinbase_maturity.rs
@@ -79,7 +79,7 @@ fn test_coinbase_maturity() {
 
 	let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
 	block.header.difficulty = difficulty.clone();
-	chain.set_sumtree_roots(&mut block).unwrap();
+	chain.set_sumtree_roots(&mut block, false).unwrap();
 
 	pow::pow_size(
 		&mut cuckoo_miner,
@@ -115,7 +115,7 @@ fn test_coinbase_maturity() {
 
 	let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
 	block.header.difficulty = difficulty.clone();
-	chain.set_sumtree_roots(&mut block).unwrap();
+	chain.set_sumtree_roots(&mut block, false).unwrap();
 
 	pow::pow_size(
 		&mut cuckoo_miner,
@@ -143,7 +143,7 @@ fn test_coinbase_maturity() {
 
 		let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
 		block.header.difficulty = difficulty.clone();
-		chain.set_sumtree_roots(&mut block).unwrap();
+		chain.set_sumtree_roots(&mut block, false).unwrap();
 
 		pow::pow_size(
 			&mut cuckoo_miner,
@@ -164,7 +164,7 @@ fn test_coinbase_maturity() {
 
 	let difficulty = consensus::next_difficulty(chain.difficulty_iter()).unwrap();
 	block.header.difficulty = difficulty.clone();
-	chain.set_sumtree_roots(&mut block).unwrap();
+	chain.set_sumtree_roots(&mut block, false).unwrap();
 
 	pow::pow_size(
 		&mut cuckoo_miner,
diff --git a/grin/src/miner.rs b/grin/src/miner.rs
index e68a93239..73fa692dd 100644
--- a/grin/src/miner.rs
+++ b/grin/src/miner.rs
@@ -591,7 +591,7 @@ impl Miner {
 		b.header.timestamp = time::at_utc(time::Timespec::new(now_sec, 0));
 		trace!(LOGGER, "Block: {:?}", b);
 	
-		let roots_result = self.chain.set_sumtree_roots(&mut b);
+		let roots_result = self.chain.set_sumtree_roots(&mut b, false);
 		match roots_result {
 			Ok(_) => Ok((b, block_fees)),