Add pruning of preimages no longer needed + tests

Author: Antoine Riard

src/ln/channelmonitor.rs

@@ -286,12 +286,27 @@ impl ChannelMonitor {
 				}
 			}
 		}
-		// TODO: Prune payment_preimages no longer needed by the revocation (just have to check
-		// that non-revoked remote commitment tx(n) do not need it, and our latest local commitment
-		// tx does not need it.
+
+		let mut pruned_payment_preimages = HashMap::with_capacity(self.payment_preimages.capacity());
+		for (k, v) in self.payment_preimages.drain() {
+			for (htlc, _s1, _s2) in &self.current_local_signed_commitment_tx.as_ref().expect("Channel need at least an initial commitment tx !").htlc_outputs {
+				if k == htlc.payment_hash {
+					pruned_payment_preimages.insert(k, v); 
+				}
+			}
+			for (_tx, htlcs) in &self.remote_claimable_outpoints {
+				for htlc in htlcs {
+					if k == htlc.payment_hash {
+						pruned_payment_preimages.insert(k, v);
+					}
+				}
+			}
+		}
+		self.payment_preimages = pruned_payment_preimages;
+
 		Ok(())
 	}
-
+	
 	/// Informs this monitor of the latest remote (ie non-broadcastable) commitment transaction.
 	/// The monitor watches for it to be broadcasted and then uses the HTLC information (and
 	/// possibly future revocation/preimage information) to claim outputs where possible.
@@ -796,9 +811,15 @@ impl ChannelMonitor {
 mod tests {
 	use bitcoin::util::misc::hex_bytes;
 	use bitcoin::blockdata::script::Script;
+	use bitcoin::util::hash::{Hash160,Sha256dHash};
+	use bitcoin::blockdata::transaction::Transaction;
 	use ln::channelmonitor::ChannelMonitor;
+	use ln::channelmonitor::LocalSignedTx;
+	use ln::chan_utils::HTLCOutputInCommitment;
 	use secp256k1::key::{SecretKey,PublicKey};
-	use secp256k1::Secp256k1;
+	use secp256k1::{Secp256k1, Signature};
+	use rand::{thread_rng,Rng};
+	use std::collections::HashMap;
 
 	#[test]
 	fn test_per_commitment_storage() {
@@ -1154,5 +1175,146 @@ mod tests {
 		}
 	}
 
+	macro_rules! gen_local_tx {
+		($hex : expr, $monitor : expr, $htlcs : expr, $rng : expr, $preimage : expr, $hash : expr) => {
+			{
+
+				let mut htlcs = Vec::new();
+
+				for _i in 0..$htlcs {
+					$rng.fill_bytes(&mut $preimage);
+					$hash[0..20].clone_from_slice(&Hash160::from_data(&$preimage)[0..20]);
+					$monitor.provide_payment_preimage(&$hash, &$preimage);
+					htlcs.push((HTLCOutputInCommitment {
+						offered : true,
+						amount_msat : 0,
+						cltv_expiry : 0,
+						payment_hash : $hash.clone(),
+						transaction_output_index : 0,
+					}, Signature::from_der(&Secp256k1::new(), $hex).unwrap(),
+					Signature::from_der(&Secp256k1::new(), $hex).unwrap()))
+				}
+  
+				Some(LocalSignedTx {
+					txid: Sha256dHash::from_data(&[]),
+					tx: Transaction {
+						version: 0,
+						lock_time: 0,
+						input: Vec::new(),
+						output: Vec::new(),
+					},
+					revocation_key: PublicKey::new(),
+					a_htlc_key: PublicKey::new(),
+					b_htlc_key: PublicKey::new(),
+					delayed_payment_key: PublicKey::new(),
+					feerate_per_kw: 0,
+					htlc_outputs: htlcs,
+				})
+			}
+		}
+	}
+
+	macro_rules! gen_remote_htlcs {
+		($monitor : expr, $tx : expr, $htlcs : expr, $rng : expr, $preimage : expr, $hash: expr) => {
+			{
+				let mut remote_outpoints = HashMap::new();
+
+				for i in 0..$tx {
+
+					let tx_zero = Transaction {
+						version : 0,
+						lock_time : i,
+						input : Vec::new(),
+						output: Vec::new(),
+					};
+
+					let mut htlcs = Vec::new();
+
+					for _i in 0..$htlcs {
+						$rng.fill_bytes(&mut $preimage);
+						$hash[0..20].clone_from_slice(&Hash160::from_data(&$preimage)[0..20]);
+						$monitor.provide_payment_preimage(&$hash, &$preimage);
+						htlcs.push(HTLCOutputInCommitment {
+							offered : true,
+							amount_msat : 0,
+							cltv_expiry : 0,
+							payment_hash : $hash.clone(),
+							transaction_output_index : 0,
+						});
+					}
+					remote_outpoints.insert(tx_zero.txid(), htlcs);
+				}
+				remote_outpoints
+			}
+		}
+	}
+
+	#[test]
+	fn test_prune_preimages() {
+
+		let mut monitor : ChannelMonitor;
+		let secp_ctx = Secp256k1::new();
+		let mut preimage : [u8;32] = [0;32];
+		let mut hash : [u8;32] = [0;32];
+		let mut rng  = thread_rng();
+
+		{
+			monitor = ChannelMonitor::new(&SecretKey::from_slice(&secp_ctx, &[42; 32]).unwrap(), &PublicKey::new(), &SecretKey::from_slice(&secp_ctx, &[43; 32]).unwrap(), 0, Script::new());
+
+			for _i in 0..30 {
+				rng.fill_bytes(&mut preimage);
+				hash[0..20].clone_from_slice(&Hash160::from_data(&preimage)[0..20]);
+				monitor.provide_payment_preimage(&hash, &preimage);
+			}
+			monitor.current_local_signed_commitment_tx = gen_local_tx!(&hex_bytes("3045022100fa86fa9a36a8cd6a7bb8f06a541787d51371d067951a9461d5404de6b928782e02201c8b7c334c10aed8976a3a465be9a28abff4cb23acbf00022295b378ce1fa3cd").unwrap()[..], monitor, 10, rng, preimage, hash);
+			monitor.remote_claimable_outpoints = gen_remote_htlcs!(monitor, 1, 10, rng, preimage, hash);
+			monitor.provide_secret(0, [0;32], None);
+			assert_eq!(monitor.payment_preimages.len(), 20);
+		}
+
+
+		{
+			monitor = ChannelMonitor::new(&SecretKey::from_slice(&secp_ctx, &[42; 32]).unwrap(), &PublicKey::new(), &SecretKey::from_slice(&secp_ctx, &[43; 32]).unwrap(), 0, Script::new());
+
+			for _i in 0..30 {
+				rng.fill_bytes(&mut preimage);
+				hash[0..20].clone_from_slice(&Hash160::from_data(&preimage)[0..20]);
+				monitor.provide_payment_preimage(&hash, &preimage);
+			}
+			monitor.current_local_signed_commitment_tx = gen_local_tx!(&hex_bytes("3045022100fa86fa9a36a8cd6a7bb8f06a541787d51371d067951a9461d5404de6b928782e02201c8b7c334c10aed8976a3a465be9a28abff4cb23acbf00022295b378ce1fa3cd").unwrap()[..], monitor, 0, rng, preimage, hash);
+			monitor.remote_claimable_outpoints = gen_remote_htlcs!(monitor, 0, 0, rng, preimage, hash);
+			monitor.provide_secret(0, [0;32], None);
+			assert_eq!(monitor.payment_preimages.len(), 0);
+		}
+
+		{
+			monitor = ChannelMonitor::new(&SecretKey::from_slice(&secp_ctx, &[42; 32]).unwrap(), &PublicKey::new(), &SecretKey::from_slice(&secp_ctx, &[43; 32]).unwrap(), 0, Script::new());
+
+			for _i in 0..30 {
+				rng.fill_bytes(&mut preimage);
+				hash[0..20].clone_from_slice(&Hash160::from_data(&preimage)[0..20]);
+				monitor.provide_payment_preimage(&hash, &preimage);
+			}
+			monitor.current_local_signed_commitment_tx = gen_local_tx!(&hex_bytes("3045022100fa86fa9a36a8cd6a7bb8f06a541787d51371d067951a9461d5404de6b928782e02201c8b7c334c10aed8976a3a465be9a28abff4cb23acbf00022295b378ce1fa3cd").unwrap()[..], monitor, 0, rng, preimage, hash);
+			monitor.remote_claimable_outpoints = gen_remote_htlcs!(monitor, 5, 5, rng, preimage, hash);
+			monitor.provide_secret(0, [0;32], None);
+			assert_eq!(monitor.payment_preimages.len(), 25);
+		}
+
+		{
+			monitor = ChannelMonitor::new(&SecretKey::from_slice(&secp_ctx, &[42; 32]).unwrap(), &PublicKey::new(), &SecretKey::from_slice(&secp_ctx, &[43; 32]).unwrap(), 0, Script::new());
+
+			for _i in 0..30 {
+				rng.fill_bytes(&mut preimage);
+				hash[0..20].clone_from_slice(&Hash160::from_data(&preimage)[0..20]);
+				monitor.provide_payment_preimage(&hash, &preimage);
+			}
+			monitor.current_local_signed_commitment_tx = gen_local_tx!(&hex_bytes("3045022100fa86fa9a36a8cd6a7bb8f06a541787d51371d067951a9461d5404de6b928782e02201c8b7c334c10aed8976a3a465be9a28abff4cb23acbf00022295b378ce1fa3cd").unwrap()[..], monitor, 25, rng, preimage, hash);
+			monitor.remote_claimable_outpoints = gen_remote_htlcs!(monitor, 0, 0, rng, preimage, hash);
+			monitor.provide_secret(0, [0;32], None);
+			assert_eq!(monitor.payment_preimages.len(), 25);
+		}
+	}
+ 
 	// Further testing is done in the ChannelManager integration tests.
 }