refactor burn key into key_overrides on keychain (#178)

* refactor burn key into key_overrides on keychain
* introduce UnconfirmedChange output status, we can potentially spend these with zero confirmations
* pass in burn_key_id for the burn enabled keychain, spend *all* coins when spending from a wallet, spend UnconfirmedChange coins also
* add comment about simplifying wallet_data.select logic
* replace UnconfirmedChange output status with a more flexible zero_ok, flag on the output data

keychain/src/keychain.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use rand::{thread_rng, Rng};
+use std::collections::HashMap;
 
 use secp;
 use secp::{Message, Secp256k1, Signature};
@@ -42,15 +43,11 @@ impl From<extkey::Error> for Error {
 	}
 }
 
-
 #[derive(Clone, Debug)]
 pub struct Keychain {
 	secp: Secp256k1,
 	extkey: extkey::ExtendedKey,
-
-	/// for tests and burn only, associate the zero fingerprint to a known
-	/// dummy private key
-	pub enable_burn_key: bool,
+	key_overrides: HashMap<Identifier, SecretKey>,
 }
 
 impl Keychain {
@@ -58,13 +55,27 @@ impl Keychain {
 		self.extkey.root_key_id.clone()
 	}
 
+	// For tests and burn only, associate a key identifier with a known secret key.
+	//
+	pub fn burn_enabled(keychain: &Keychain, burn_key_id: &Identifier) -> Keychain {
+		let mut key_overrides = HashMap::new();
+		key_overrides.insert(
+			burn_key_id.clone(),
+			SecretKey::from_slice(&keychain.secp, &[1; 32]).unwrap(),
+		);
+		Keychain {
+			key_overrides: key_overrides,
+			..keychain.clone()
+		}
+	}
+
 	pub fn from_seed(seed: &[u8]) -> Result<Keychain, Error> {
 		let secp = secp::Secp256k1::with_caps(secp::ContextFlag::Commit);
 		let extkey = extkey::ExtendedKey::from_seed(&secp, seed)?;
 		let keychain = Keychain {
 			secp: secp,
 			extkey: extkey,
-			enable_burn_key: false,
+			key_overrides: HashMap::new(),
 		};
 		Ok(keychain)
 	}
@@ -82,36 +93,20 @@ impl Keychain {
 		Ok(key_id)
 	}
 
-	// TODO - this is a work in progress
-	// TODO - smarter lookups - can we cache key_id/fingerprint -> derivation
-	// number somehow?
 	fn derived_key(&self, key_id: &Identifier) -> Result<SecretKey, Error> {
-		if self.enable_burn_key {
-			// for tests and burn only, associate the zero fingerprint to a known
-			// dummy private key
-			if *key_id == Identifier::zero() {
-				return Ok(SecretKey::from_slice(&self.secp, &[1; 32])?);
-			}
+		if let Some(key) = self.key_overrides.get(key_id) {
+			return Ok(*key);
 		}
+
 		for i in 1..10000 {
 			let extkey = self.extkey.derive(&self.secp, i)?;
 			if extkey.identifier(&self.secp)? == *key_id {
 				return Ok(extkey.key);
 			}
 		}
-		Err(Error::KeyDerivation(format!("cannot find extkey for {:?}", key_id)))
-	}
-
-	// TODO - clean this and derived_key up, rename them?
-	// TODO - maybe wallet deals exclusively with key_ids and not derivations - this leaks?
-	pub fn derivation_from_key_id(&self, key_id: &Identifier) -> Result<u32, Error> {
-		for i in 1..10000 {
-			let extkey = self.extkey.derive(&self.secp, i)?;
-			if extkey.identifier(&self.secp)? == *key_id {
-				return Ok(extkey.n_child);
-			}
-		}
-		Err(Error::KeyDerivation(format!("cannot find extkey for {:?}", key_id)))
+		Err(Error::KeyDerivation(
+			format!("cannot find extkey for {:?}", key_id),
+		))
 	}
 
 	pub fn commit(&self, amount: u64, key_id: &Identifier) -> Result<Commitment, Error> {
@@ -246,19 +241,25 @@ mod test {
 		let key_id2 = keychain.derive_key_id(2).unwrap();
 
 		// cannot rewind with a different nonce
-		let proof_info = keychain.rewind_range_proof(&key_id2, commit, proof).unwrap();
+		let proof_info = keychain
+			.rewind_range_proof(&key_id2, commit, proof)
+			.unwrap();
 		assert_eq!(proof_info.success, false);
 		assert_eq!(proof_info.value, 0);
 
 		// cannot rewind with a commitment to the same value using a different key
 		let commit2 = keychain.commit(5, &key_id2).unwrap();
-		let proof_info = keychain.rewind_range_proof(&key_id, commit2, proof).unwrap();
+		let proof_info = keychain
+			.rewind_range_proof(&key_id, commit2, proof)
+			.unwrap();
 		assert_eq!(proof_info.success, false);
 		assert_eq!(proof_info.value, 0);
 
 		// cannot rewind with a commitment to a different value
 		let commit3 = keychain.commit(4, &key_id).unwrap();
-		let proof_info = keychain.rewind_range_proof(&key_id, commit3, proof).unwrap();
+		let proof_info = keychain
+			.rewind_range_proof(&key_id, commit3, proof)
+			.unwrap();
 		assert_eq!(proof_info.success, false);
 		assert_eq!(proof_info.value, 0);
 	}

src/bin/grin.rs

@@ -325,7 +325,7 @@ fn wallet_command(wallet_args: &ArgMatches) {
 
 	// TODO do something closer to BIP39, eazy solution right now
 	let seed = blake2::blake2b::blake2b(32, &[], hd_seed.as_bytes());
-	let mut keychain = Keychain::from_seed(seed.as_bytes()).expect(
+	let keychain = Keychain::from_seed(seed.as_bytes()).expect(
 		"Failed to initialize keychain from the provided seed.",
 	);
 
@@ -390,7 +390,6 @@ fn wallet_command(wallet_args: &ArgMatches) {
 				.expect("Amount to burn required")
 				.parse()
 				.expect("Could not parse amount as a whole number.");
-			keychain.enable_burn_key = true;
 			wallet::issue_burn_tx(&wallet_config, &keychain, amount).unwrap();
 		}
 		("info", Some(_)) => {

wallet/src/info.rs

@@ -24,7 +24,7 @@ pub fn show_info(config: &WalletConfig, keychain: &Keychain) {
 	let _ = WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
 
 		println!("Outputs - ");
-		println!("key_id, height, lock_height, status, value");
+		println!("key_id, height, lock_height, status, zero_ok, value");
 		println!("----------------------------------");
 
 		let mut outputs = wallet_data
@@ -35,11 +35,12 @@ pub fn show_info(config: &WalletConfig, keychain: &Keychain) {
 		outputs.sort_by_key(|out| out.n_child);
 		for out in outputs {
 			println!(
-				"{}, {}, {}, {:?}, {}",
+				"{}, {}, {}, {:?}, {}, {}",
 				out.key_id,
 				out.height,
 				out.lock_height,
 				out.status,
+				out.zero_ok,
 				out.value
 			);
 		}

wallet/src/receiver.rs

@@ -183,8 +183,11 @@ fn receive_coinbase(config: &WalletConfig,
 		let key_id = block_fees.key_id();
 		let (key_id, derivation) = match key_id {
 			Some(key_id) => {
-				let derivation = keychain.derivation_from_key_id(&key_id)?;
-				(key_id.clone(), derivation)
+				if let Some(existing) = wallet_data.get_output(&key_id) {
+					(existing.key_id.clone(), existing.n_child)
+				} else {
+					panic!("should never happen");
+				}
 			},
 			None => {
 				let derivation = wallet_data.next_child(root_key_id.clone());
@@ -202,6 +205,7 @@ fn receive_coinbase(config: &WalletConfig,
 			status: OutputStatus::Unconfirmed,
 			height: 0,
 			lock_height: 0,
+			zero_ok: false,
 		});
 
 		debug!(
@@ -276,6 +280,7 @@ fn receive_transaction(
 			status: OutputStatus::Unconfirmed,
 			height: 0,
 			lock_height: 0,
+			zero_ok: false,
 		});
 		debug!(
 			LOGGER,

wallet/src/sender.rs

@@ -16,7 +16,7 @@ use api;
 use checker;
 use core::core::{Transaction, build};
 use core::ser;
-use keychain::{BlindingFactor, Keychain, Identifier, IDENTIFIER_SIZE};
+use keychain::{BlindingFactor, Keychain, Identifier};
 use receiver::TxWrapper;
 use types::*;
 use util::LOGGER;
@@ -72,12 +72,10 @@ fn build_send_tx(
 	WalletData::with_wallet(&config.data_file_dir, |wallet_data| {
 
 		// select some suitable outputs to spend from our local wallet
-		let (coins, change) = wallet_data.select(key_id.clone(), amount);
-		if change < 0 {
-			return Err(Error::NotEnoughFunds((-change) as u64));
-		}
+		let (coins, _) = wallet_data.select(key_id.clone(), u64::max_value());
 
 		// build transaction skeleton with inputs and change
+		// TODO - should probably also check we are sending enough to cover the fees + non-zero output
 		let mut parts = inputs_and_change(&coins, keychain, key_id, wallet_data, amount)?;
 
 		// This is more proof of concept than anything but here we set a
@@ -92,8 +90,11 @@ fn build_send_tx(
 }
 
 pub fn issue_burn_tx(config: &WalletConfig, keychain: &Keychain, amount: u64) -> Result<(), Error> {
+	let keychain = &Keychain::burn_enabled(keychain, &Identifier::zero());
+
 	let _ = checker::refresh_outputs(config, keychain);
-	let key_id = keychain.clone().root_key_id();
+
+	let key_id = keychain.root_key_id();
 
 	// operate within a lock on wallet data
 	WalletData::with_wallet(&config.data_file_dir, |mut wallet_data| {
@@ -105,10 +106,8 @@ pub fn issue_burn_tx(config: &WalletConfig, keychain: &Keychain, amount: u64) ->
 		let mut parts = inputs_and_change(&coins, keychain, key_id, &mut wallet_data, amount)?;
 
 		// add burn output and fees
-		parts.push(build::output(
-			amount,
-			Identifier::from_bytes(&[0; IDENTIFIER_SIZE]),
-		));
+		let fee = tx_fee(coins.len(), 2, None);
+		parts.push(build::output(amount - fee, Identifier::zero()));
 
 		// finalize the burn transaction and send
 		let (tx_burn, _) = build::transaction(parts, &keychain)?;
@@ -162,8 +161,7 @@ fn inputs_and_change(
 	let change_key = keychain.derive_key_id(change_derivation)?;
 	parts.push(build::output(change, change_key.clone()));
 
-	// we got that far, time to start tracking the new output
-	// and lock the outputs used
+	// we got that far, time to start tracking the output representing our change
 	wallet_data.add_output(OutputData {
 		root_key_id: root_key_id.clone(),
 		key_id: change_key.clone(),
@@ -172,9 +170,10 @@ fn inputs_and_change(
 		status: OutputStatus::Unconfirmed,
 		height: 0,
 		lock_height: 0,
+		zero_ok: true,
 	});
 
-	// lock the ouputs we're spending
+	// now lock the ouputs we're spending so we avoid accidental double spend attempt
 	for coin in coins {
 		wallet_data.lock_output(coin);
 	}

wallet/src/types.rs

@@ -168,6 +168,8 @@ pub struct OutputData {
 	pub height: u64,
 	/// Height we are locked until
 	pub lock_height: u64,
+	/// Can we spend with zero confirmations? (Did it originate from us, change output etc.)
+	pub zero_ok: bool,
 }
 
 impl OutputData {
@@ -307,16 +309,22 @@ impl WalletData {
 		}
 	}
 
+	pub fn get_output(&self, key_id: &keychain::Identifier) -> Option<&OutputData> {
+		self.outputs.get(&key_id.to_hex())
+	}
+
 	/// Select a subset of unspent outputs to spend in a transaction
 	/// transferring the provided amount.
 	pub fn select(&self, root_key_id: keychain::Identifier, amount: u64) -> (Vec<OutputData>, i64) {
 		let mut to_spend = vec![];
 		let mut input_total = 0;
 
-		// TODO very naive impl for now - definitely better coin selection
-		// algos available
 		for out in self.outputs.values() {
-			if out.status == OutputStatus::Unspent && out.root_key_id == root_key_id {
+			if out.root_key_id == root_key_id
+				&& (out.status == OutputStatus::Unspent)
+					// the following will let us spend zero confirmation change outputs
+					// || (out.status == OutputStatus::Unconfirmed && out.zero_ok))
+			{
 				to_spend.push(out.clone());
 				input_total += out.value;
 				if input_total >= amount {