Merge pull request #159 from ariard/channel_monitor

Add registration of commitment tx's outputs from check_spend_remote_transaction

Author: TheBlueMatt

fuzz/fuzz_targets/router_target.rs

@@ -76,7 +76,7 @@ struct DummyChainWatcher {
 }
 
 impl ChainWatchInterface for DummyChainWatcher {
-	fn install_watch_script(&self, _script_pub_key: &Script) { }
+	fn install_watch_tx(&self, _txid: &Sha256dHash, _script_pub_key: &Script) { }
 	fn install_watch_outpoint(&self, _outpoint: (Sha256dHash, u32), _out_script: &Script) { }
 	fn watch_all_txn(&self) { }
 	fn register_listener(&self, _listener: Weak<ChainListener>) { }

src/chain/chaininterface.rs

@@ -5,9 +5,12 @@ use bitcoin::blockdata::constants::genesis_block;
 use bitcoin::util::hash::Sha256dHash;
 use bitcoin::network::constants::Network;
 use bitcoin::network::serialize::BitcoinHash;
+
 use util::logger::Logger;
+
 use std::sync::{Mutex,Weak,MutexGuard,Arc};
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::collections::HashSet;
 
 /// Used to give chain error details upstream
 pub enum ChainError {
@@ -25,8 +28,8 @@ pub enum ChainError {
 /// called from inside the library in response to ChainListener events, P2P events, or timer
 /// events).
 pub trait ChainWatchInterface: Sync + Send {
-	/// Provides a scriptPubKey which much be watched for.
-	fn install_watch_script(&self, script_pub_key: &Script);
+	/// Provides a txid/random-scriptPubKey-in-the-tx which much be watched for.
+	fn install_watch_tx(&self, txid: &Sha256dHash, script_pub_key: &Script);
 
 	/// Provides an outpoint which must be watched for, providing any transactions which spend the
 	/// given outpoint.
@@ -54,9 +57,11 @@ pub trait BroadcasterInterface: Sync + Send {
 /// A trait indicating a desire to listen for events from the chain
 pub trait ChainListener: Sync + Send {
 	/// Notifies a listener that a block was connected.
-	/// Note that if a new script/transaction is watched during a block_connected call, the block
-	/// *must* be re-scanned with the new script/transaction and block_connected should be called
-	/// again with the same header and (at least) the new transactions.
+	/// Note that if a new transaction/outpoint is watched during a block_connected call, the block
+	/// *must* be re-scanned with the new transaction/outpoints and block_connected should be
+	/// called again with the same header and (at least) the new transactions.
+	/// Note that if non-new transaction/outpoints may be registered during a call, a second call
+	/// *must not* happen.
 	/// This also means those counting confirmations using block_connected callbacks should watch
 	/// for duplicate headers and not count them towards confirmations!
 	fn block_connected(&self, header: &BlockHeader, height: u32, txn_matched: &[&Transaction], indexes_of_txn_matched: &[u32]);
@@ -85,34 +90,124 @@ pub trait FeeEstimator: Sync + Send {
 	fn get_est_sat_per_1000_weight(&self, confirmation_target: ConfirmationTarget) -> u64;
 }
 
+/// Utility for tracking registered txn/outpoints and checking for matches
+pub struct ChainWatchedUtil {
+	watch_all: bool,
+
+	// We are more conservative in matching during testing to ensure everything matches *exactly*,
+	// even though during normal runtime we take more optimized match approaches...
+	#[cfg(test)]
+	watched_txn: HashSet<(Sha256dHash, Script)>,
+	#[cfg(not(test))]
+	watched_txn: HashSet<Script>,
+
+	watched_outpoints: HashSet<(Sha256dHash, u32)>,
+}
+
+impl ChainWatchedUtil {
+	/// Constructs an empty (watches nothing) ChainWatchedUtil
+	pub fn new() -> Self {
+		Self {
+			watch_all: false,
+			watched_txn: HashSet::new(),
+			watched_outpoints: HashSet::new(),
+		}
+	}
+
+	/// Registers a tx for monitoring, returning true if it was a new tx and false if we'd already
+	/// been watching for it.
+	pub fn register_tx(&mut self, txid: &Sha256dHash, script_pub_key: &Script) -> bool {
+		if self.watch_all { return false; }
+		#[cfg(test)]
+		{
+			self.watched_txn.insert((txid.clone(), script_pub_key.clone()))
+		}
+		#[cfg(not(test))]
+		{
+			let _tx_unused = txid; // Its used in cfg(test), though
+			self.watched_txn.insert(script_pub_key.clone())
+		}
+	}
+
+	/// Registers an outpoint for monitoring, returning true if it was a new outpoint and false if
+	/// we'd already been watching for it
+	pub fn register_outpoint(&mut self, outpoint: (Sha256dHash, u32), _script_pub_key: &Script) -> bool {
+		if self.watch_all { return false; }
+		self.watched_outpoints.insert(outpoint)
+	}
+
+	/// Sets us to match all transactions, returning true if this is a new setting anf false if
+	/// we'd already been set to match everything.
+	pub fn watch_all(&mut self) -> bool {
+		if self.watch_all { return false; }
+		self.watch_all = true;
+		true
+	}
+
+	/// Checks if a given transaction matches the current filter.
+	pub fn does_match_tx(&self, tx: &Transaction) -> bool {
+		if self.watch_all {
+			return true;
+		}
+		for out in tx.output.iter() {
+			#[cfg(test)]
+			for &(ref txid, ref script) in self.watched_txn.iter() {
+				if *script == out.script_pubkey {
+					if tx.txid() == *txid {
+						return true;
+					}
+				}
+			}
+			#[cfg(not(test))]
+			for script in self.watched_txn.iter() {
+				if *script == out.script_pubkey {
+					return true;
+				}
+			}
+		}
+		for input in tx.input.iter() {
+			for outpoint in self.watched_outpoints.iter() {
+				let &(outpoint_hash, outpoint_index) = outpoint;
+				if outpoint_hash == input.previous_output.txid && outpoint_index == input.previous_output.vout {
+					return true;
+				}
+			}
+		}
+		false
+	}
+}
+
 /// Utility to capture some common parts of ChainWatchInterface implementors.
 /// Keeping a local copy of this in a ChainWatchInterface implementor is likely useful.
 pub struct ChainWatchInterfaceUtil {
 	network: Network,
-	watched: Mutex<(Vec<Script>, Vec<(Sha256dHash, u32)>, bool)>, //TODO: Something clever to optimize this
+	watched: Mutex<ChainWatchedUtil>,
 	listeners: Mutex<Vec<Weak<ChainListener>>>,
 	reentered: AtomicUsize,
 	logger: Arc<Logger>,
 }
 
 /// Register listener
 impl ChainWatchInterface for ChainWatchInterfaceUtil {
-	fn install_watch_script(&self, script_pub_key: &Script) {
+	fn install_watch_tx(&self, txid: &Sha256dHash, script_pub_key: &Script) {
 		let mut watched = self.watched.lock().unwrap();
-		watched.0.push(script_pub_key.clone());
-		self.reentered.fetch_add(1, Ordering::Relaxed);
+		if watched.register_tx(txid, script_pub_key) {
+			self.reentered.fetch_add(1, Ordering::Relaxed);
+		}
 	}
 
-	fn install_watch_outpoint(&self, outpoint: (Sha256dHash, u32), _out_script: &Script) {
+	fn install_watch_outpoint(&self, outpoint: (Sha256dHash, u32), out_script: &Script) {
 		let mut watched = self.watched.lock().unwrap();
-		watched.1.push(outpoint);
-		self.reentered.fetch_add(1, Ordering::Relaxed);
+		if watched.register_outpoint(outpoint, out_script) {
+			self.reentered.fetch_add(1, Ordering::Relaxed);
+		}
 	}
 
 	fn watch_all_txn(&self) {
 		let mut watched = self.watched.lock().unwrap();
-		watched.2 = true;
-		self.reentered.fetch_add(1, Ordering::Relaxed);
+		if watched.watch_all() {
+			self.reentered.fetch_add(1, Ordering::Relaxed);
+		}
 	}
 
 	fn register_listener(&self, listener: Weak<ChainListener>) {
@@ -132,7 +227,7 @@ impl ChainWatchInterfaceUtil {
 	pub fn new(network: Network, logger: Arc<Logger>) -> ChainWatchInterfaceUtil {
 		ChainWatchInterfaceUtil {
 			network: network,
-			watched: Mutex::new((Vec::new(), Vec::new(), false)),
+			watched: Mutex::new(ChainWatchedUtil::new()),
 			listeners: Mutex::new(Vec::new()),
 			reentered: AtomicUsize::new(1),
 			logger: logger,
@@ -195,25 +290,7 @@ impl ChainWatchInterfaceUtil {
 		self.does_match_tx_unguarded (tx, &watched)
 	}
 
-	fn does_match_tx_unguarded(&self, tx: &Transaction, watched: &MutexGuard<(Vec<Script>, Vec<(Sha256dHash, u32)>, bool)>) -> bool {
-		if watched.2 {
-			return true;
-		}
-		for out in tx.output.iter() {
-			for script in watched.0.iter() {
-				if script[..] == out.script_pubkey[..] {
-					return true;
-				}
-			}
-		}
-		for input in tx.input.iter() {
-			for outpoint in watched.1.iter() {
-				let &(outpoint_hash, outpoint_index) = outpoint;
-				if outpoint_hash == input.previous_output.txid && outpoint_index == input.previous_output.vout {
-					return true;
-				}
-			}
-		}
-		false
+	fn does_match_tx_unguarded(&self, tx: &Transaction, watched: &MutexGuard<ChainWatchedUtil>) -> bool {
+		watched.does_match_tx(tx)
 	}
 }

src/ln/channelmanager.rs

@@ -3337,7 +3337,8 @@ mod tests {
 			nodes[1].chain_monitor.block_connected_with_filtering(&Block { header, txdata: vec![revoked_local_txn[0].clone()] }, 1);
 			{
 				let mut node_txn = nodes[1].tx_broadcaster.txn_broadcasted.lock().unwrap();
-				assert_eq!(node_txn.len(), 2);
+				assert_eq!(node_txn.len(), 3);
+				assert_eq!(node_txn.pop().unwrap(), node_txn[0]); // An outpoint registration will result in a 2nd block_connected
 				assert_eq!(node_txn[0].input.len(), 1);
 
 				let mut funding_tx_map = HashMap::new();

src/ln/channelmonitor.rs

@@ -46,6 +46,8 @@ pub enum ChannelMonitorUpdateErr {
 /// channel's monitor everywhere (including remote watchtowers) *before* this function returns. If
 /// an update occurs and a remote watchtower is left with old state, it may broadcast transactions
 /// which we have revoked, allowing our counterparty to claim all funds in the channel!
+/// A call to add_update_monitor is needed to register outpoint and its txid with ChainWatchInterface 
+/// after setting funding_txo in a ChannelMonitor
 pub trait ManyChannelMonitor: Send + Sync {
 	/// Adds or updates a monitor for the given `funding_txo`.
 	fn add_update_monitor(&self, funding_txo: OutPoint, monitor: ChannelMonitor) -> Result<(), ChannelMonitorUpdateErr>;
@@ -69,7 +71,12 @@ impl<Key : Send + cmp::Eq + hash::Hash> ChainListener for SimpleManyChannelMonit
 	fn block_connected(&self, _header: &BlockHeader, height: u32, txn_matched: &[&Transaction], _indexes_of_txn_matched: &[u32]) {
 		let monitors = self.monitors.lock().unwrap();
 		for monitor in monitors.values() {
-			monitor.block_connected(txn_matched, height, &*self.broadcaster);
+			let txn_outputs = monitor.block_connected(txn_matched, height, &*self.broadcaster);
+			for (ref txid, ref outputs) in txn_outputs {
+				for (idx, output) in outputs.iter().enumerate() {
+					self.chain_monitor.install_watch_outpoint((txid.clone(), idx as u32), &output.script_pubkey);
+				}
+			}
 		}
 	}
 
@@ -97,7 +104,7 @@ impl<Key : Send + cmp::Eq + hash::Hash + 'static> SimpleManyChannelMonitor<Key>
 		match &monitor.funding_txo {
 			&None => self.chain_monitor.watch_all_txn(),
 			&Some((ref outpoint, ref script)) => {
-				self.chain_monitor.install_watch_script(script);
+				self.chain_monitor.install_watch_tx(&outpoint.txid, script);
 				self.chain_monitor.install_watch_outpoint((outpoint.txid, outpoint.index as u32), script);
 			},
 		}
@@ -464,8 +471,9 @@ impl ChannelMonitor {
 	/// optional, without it this monitor cannot be used in an SPV client, but you may wish to
 	/// avoid this (or call unset_funding_info) on a monitor you wish to send to a watchtower as it
 	/// provides slightly better privacy.
+	/// It's the responsability of the caller to register outpoint and script with passing the former
+	/// value as key to add_update_monitor.
 	pub(super) fn set_funding_info(&mut self, funding_info: (OutPoint, Script)) {
-		//TODO: Need to register the given script here with a chain_monitor
 		self.funding_txo = Some(funding_info);
 	}
 
@@ -908,22 +916,24 @@ impl ChannelMonitor {
 	/// height > height + CLTV_SHARED_CLAIM_BUFFER. In any case, will install monitoring for
 	/// HTLC-Success/HTLC-Timeout transactions, and claim them using the revocation key (if
 	/// applicable) as well.
-	fn check_spend_remote_transaction(&self, tx: &Transaction, height: u32) -> Vec<Transaction> {
+	fn check_spend_remote_transaction(&self, tx: &Transaction, height: u32) -> (Vec<Transaction>, (Sha256dHash, Vec<TxOut>)) {
 		// Most secp and related errors trying to create keys means we have no hope of constructing
 		// a spend transaction...so we return no transactions to broadcast
 		let mut txn_to_broadcast = Vec::new();
+		let mut watch_outputs = Vec::new();
+
+		let commitment_txid = tx.txid(); //TODO: This is gonna be a performance bottleneck for watchtowers!
+		let per_commitment_option = self.remote_claimable_outpoints.get(&commitment_txid);
+
 		macro_rules! ignore_error {
 			( $thing : expr ) => {
 				match $thing {
 					Ok(a) => a,
-					Err(_) => return txn_to_broadcast
+					Err(_) => return (txn_to_broadcast, (commitment_txid, watch_outputs))
 				}
 			};
 		}
 
-		let commitment_txid = tx.txid(); //TODO: This is gonna be a performance bottleneck for watchtowers!
-		let per_commitment_option = self.remote_claimable_outpoints.get(&commitment_txid);
-
 		let commitment_number = 0xffffffffffff - ((((tx.input[0].sequence as u64 & 0xffffff) << 3*8) | (tx.lock_time as u64 & 0xffffff)) ^ self.commitment_transaction_number_obscure_factor);
 		if commitment_number >= self.get_min_seen_secret() {
 			let secret = self.get_secret(commitment_number).unwrap();
@@ -942,7 +952,7 @@ impl ChannelMonitor {
 			};
 			let delayed_key = ignore_error!(chan_utils::derive_public_key(&self.secp_ctx, &PublicKey::from_secret_key(&self.secp_ctx, &per_commitment_key), &self.delayed_payment_base_key));
 			let a_htlc_key = match self.their_htlc_base_key {
-				None => return txn_to_broadcast,
+				None => return (txn_to_broadcast, (commitment_txid, watch_outputs)),
 				Some(their_htlc_base_key) => ignore_error!(chan_utils::derive_public_key(&self.secp_ctx, &PublicKey::from_secret_key(&self.secp_ctx, &per_commitment_key), &their_htlc_base_key)),
 			};
 
@@ -1009,7 +1019,7 @@ impl ChannelMonitor {
 					if htlc.transaction_output_index as usize >= tx.output.len() ||
 							tx.output[htlc.transaction_output_index as usize].value != htlc.amount_msat / 1000 ||
 							tx.output[htlc.transaction_output_index as usize].script_pubkey != expected_script.to_v0_p2wsh() {
-						return txn_to_broadcast; // Corrupted per_commitment_data, fuck this user
+						return (txn_to_broadcast, (commitment_txid, watch_outputs)); // Corrupted per_commitment_data, fuck this user
 					}
 					let input = TxIn {
 						previous_output: BitcoinOutPoint {
@@ -1044,10 +1054,10 @@ impl ChannelMonitor {
 
 			if !inputs.is_empty() || !txn_to_broadcast.is_empty() { // ie we're confident this is actually ours
 				// We're definitely a remote commitment transaction!
-				// TODO: Register all outputs in commitment_tx with the ChainWatchInterface!
+				watch_outputs.append(&mut tx.output.clone());
 				self.remote_commitment_txn_on_chain.lock().unwrap().insert(commitment_txid, commitment_number);
 			}
-			if inputs.is_empty() { return txn_to_broadcast; } // Nothing to be done...probably a false positive/local tx
+			if inputs.is_empty() { return (txn_to_broadcast, (commitment_txid, watch_outputs)); } // Nothing to be done...probably a false positive/local tx
 
 			let outputs = vec!(TxOut {
 				script_pubkey: self.destination_script.clone(),
@@ -1077,7 +1087,7 @@ impl ChannelMonitor {
 			// already processed the block, resulting in the remote_commitment_txn_on_chain entry
 			// not being generated by the above conditional. Thus, to be safe, we go ahead and
 			// insert it here.
-			// TODO: Register all outputs in commitment_tx with the ChainWatchInterface!
+			watch_outputs.append(&mut tx.output.clone());
 			self.remote_commitment_txn_on_chain.lock().unwrap().insert(commitment_txid, commitment_number);
 
 			if let Some(revocation_points) = self.their_cur_revocation_points {
@@ -1098,7 +1108,7 @@ impl ChannelMonitor {
 						},
 					};
 					let a_htlc_key = match self.their_htlc_base_key {
-						None => return txn_to_broadcast,
+						None => return (txn_to_broadcast, (commitment_txid, watch_outputs)),
 						Some(their_htlc_base_key) => ignore_error!(chan_utils::derive_public_key(&self.secp_ctx, revocation_point, &their_htlc_base_key)),
 					};
 
@@ -1161,7 +1171,7 @@ impl ChannelMonitor {
 						}
 					}
 
-					if inputs.is_empty() { return txn_to_broadcast; } // Nothing to be done...probably a false positive/local tx
+					if inputs.is_empty() { return (txn_to_broadcast, (commitment_txid, watch_outputs)); } // Nothing to be done...probably a false positive/local tx
 
 					let outputs = vec!(TxOut {
 						script_pubkey: self.destination_script.clone(),
@@ -1189,7 +1199,7 @@ impl ChannelMonitor {
 			//TODO: For each input check if its in our remote_commitment_txn_on_chain map!
 		}
 
-		txn_to_broadcast
+		(txn_to_broadcast, (commitment_txid, watch_outputs))
 	}
 
 	fn broadcast_by_local_state(&self, local_tx: &LocalSignedTx) -> Vec<Transaction> {
@@ -1250,11 +1260,15 @@ impl ChannelMonitor {
 		Vec::new()
 	}
 
-	fn block_connected(&self, txn_matched: &[&Transaction], height: u32, broadcaster: &BroadcasterInterface) {
+	fn block_connected(&self, txn_matched: &[&Transaction], height: u32, broadcaster: &BroadcasterInterface)-> Vec<(Sha256dHash, Vec<TxOut>)> {
+		let mut watch_outputs = Vec::new();
 		for tx in txn_matched {
 			for txin in tx.input.iter() {
 				if self.funding_txo.is_none() || (txin.previous_output.txid == self.funding_txo.as_ref().unwrap().0.txid && txin.previous_output.vout == self.funding_txo.as_ref().unwrap().0.index as u32) {
-					let mut txn = self.check_spend_remote_transaction(tx, height);
+					let (mut txn, new_outputs) = self.check_spend_remote_transaction(tx, height);
+					if !new_outputs.1.is_empty() {
+						watch_outputs.push(new_outputs);
+					}
 					if txn.is_empty() {
 						txn = self.check_spend_local_transaction(tx, height);
 					}
@@ -1281,6 +1295,7 @@ impl ChannelMonitor {
 				}
 			}
 		}
+		watch_outputs
 	}
 
 	pub fn would_broadcast_at_height(&self, height: u32) -> bool {