Implement WatchtowerPersister to mock forming justice tx from monitor updates

Author: alecchendev

lightning/src/util/test_utils.rs

@@ -32,6 +32,7 @@ use crate::util::enforcing_trait_impls::{EnforcingSigner, EnforcementState};
 use crate::util::logger::{Logger, Level, Record};
 use crate::util::ser::{Readable, ReadableArgs, Writer, Writeable};
 
+use bitcoin::{TxIn, Witness, EcdsaSighashType};
 use bitcoin::blockdata::constants::genesis_block;
 use bitcoin::blockdata::transaction::{Transaction, TxOut};
 use bitcoin::blockdata::script::{Builder, Script};
@@ -264,6 +265,157 @@ impl<'a> chain::Watch<EnforcingSigner> for TestChainMonitor<'a> {
 	}
 }
 
+pub enum WatchtowerState {
+	/// Upon a new commitment signed, we'll get a
+	/// ChannelMonitorUpdateStep::LatestCounterpartyCommitmentTxInfo. We'll store the commitment txid and
+	/// revokeable output index to use to form the justice tx once we get a revoke_and_ack with the
+	/// commitment secret.
+	CounterpartyCommitmentTxSeen {
+		commitment_txid: Txid,
+		output_idx: u32,
+		value: u64,
+	},
+	/// After receiving a revoke_and_ack for a commitment number, we'll form and store the justice
+	/// tx which would be used to provide the watchtower with the data it needs.
+	JusticeTxFormed(Transaction),
+}
+
+pub struct WatchtowerPersister {
+	pub persister: TestPersister,
+	pub destination_script: Script,
+	pub channel_watchtower_state: Mutex<HashMap<OutPoint, HashMap<u64, WatchtowerState>>>,
+}
+
+impl WatchtowerPersister {
+	pub(crate) fn new(destination_script: Script) -> Self {
+		WatchtowerPersister {
+			persister: TestPersister::new(),
+			destination_script,
+			channel_watchtower_state: Mutex::new(HashMap::new()),
+		}
+	}
+
+	fn build_justice_tx(&self, commitment_txid: Txid, output_idx: u32, value: u64) -> Transaction {
+		let mut justice_tx = Transaction {
+			version: 2,
+			lock_time: bitcoin::PackedLockTime::ZERO,
+			input: vec![TxIn {
+				previous_output: bitcoin::OutPoint {
+					txid: commitment_txid,
+					vout: output_idx,
+				},
+				script_sig: Script::new(),
+				sequence: Sequence::ENABLE_RBF_NO_LOCKTIME,
+				witness: Witness::new(),
+			}],
+			output: vec![TxOut {
+				script_pubkey: self.destination_script.clone(),
+				value,
+			}],
+		};
+		// Naive fee calculation to pass tests
+		let min_fee = (justice_tx.weight() as u64 + 3) / 4;
+		justice_tx.output[0].value -= min_fee * 2;
+		justice_tx
+	}
+
+	fn sign_justice_tx<Signer: sign::WriteableEcdsaChannelSigner>(&self, mut justice_tx: Transaction, secret: &[u8; 32], value: u64, data: &channelmonitor::ChannelMonitor<Signer>) -> Result<Transaction, ()> {
+		let secp_ctx = Secp256k1::new();
+		let per_commitment_key = SecretKey::from_slice(secret).unwrap();
+		let per_commitment_point = PublicKey::from_secret_key(&secp_ctx, &per_commitment_key);
+		let revokeable_redeemscript = data.get_revokeable_redeemscript(per_commitment_point);
+
+		let input_idx = 0;
+		let sig = match data.sign_justice_revoked_output(&justice_tx, input_idx, value, &per_commitment_key) {
+			Ok(sig) => sig,
+			Err(_) => return Err(()),
+		};
+
+		let mut ser_sig = sig.serialize_der().to_vec();
+		ser_sig.push(EcdsaSighashType::All as u8);
+		justice_tx.input[0].witness.push(ser_sig);
+		justice_tx.input[0].witness.push(vec!(1));
+		justice_tx.input[0].witness.push(revokeable_redeemscript.clone().into_bytes());
+
+		Ok(justice_tx)
+	}
+}
+
+impl<Signer: sign::WriteableEcdsaChannelSigner> chainmonitor::Persist<Signer> for WatchtowerPersister {
+	fn persist_new_channel(&self, funding_txo: OutPoint, data: &channelmonitor::ChannelMonitor<Signer>, id: MonitorUpdateId) -> chain::ChannelMonitorUpdateStatus {
+		assert!(self.channel_watchtower_state.lock().unwrap()
+			.insert(funding_txo, HashMap::new()).is_none());
+
+		// No ChannelMonitorUpdateStep::LatestCounterpartyCommitmentTxInfo is persisted on channel creation
+		// so we get the same info and update the watchtower
+		let commitment_number = (1u64 << 48) - 1;
+		let output_idx = 0;
+		let commitment_txid = data.get_current_counterparty_commitment_txid().expect("Should have at least one counterparty commitment tx");
+		let value = data.initial_to_counterparty_value();
+		if value != 0 {
+			assert!(self.channel_watchtower_state.lock().unwrap().get_mut(&funding_txo).unwrap()
+				.insert(commitment_number, WatchtowerState::CounterpartyCommitmentTxSeen {
+					commitment_txid,
+					output_idx,
+					value,
+				}).is_none());
+		}
+
+		self.persister.persist_new_channel(funding_txo, data, id)
+	}
+
+	fn update_persisted_channel(&self, funding_txo: OutPoint, update: Option<&channelmonitor::ChannelMonitorUpdate>, data: &channelmonitor::ChannelMonitor<Signer>, update_id: MonitorUpdateId) -> chain::ChannelMonitorUpdateStatus {
+		if let Some(up) = update {
+			for step in up.updates.iter() {
+				match step {
+					channelmonitor::ChannelMonitorUpdateStep::LatestCounterpartyCommitmentTXInfo { commitment_txid, htlc_outputs: _, commitment_number, their_per_commitment_point, non_htlc_outputs } => {
+						// Find the revokeable output and store the data needed to build the
+						// justice tx later
+						let revokeable_redeemscript = data.get_revokeable_redeemscript(*their_per_commitment_point);
+						let revokeable_p2wsh = revokeable_redeemscript.to_v0_p2wsh();
+						let (output_idx, output) = match non_htlc_outputs.iter().enumerate()
+							.find(|(_, output)| output.script_pubkey == revokeable_p2wsh) {
+								Some((idx, output)) => (idx, output),
+								None => {
+									println!("No revokeable output found, revokeable output must be below dust limit.");
+									continue;
+								},
+							};
+
+						assert!(self.channel_watchtower_state.lock().unwrap().get_mut(&funding_txo).unwrap()
+							.insert(*commitment_number, WatchtowerState::CounterpartyCommitmentTxSeen {
+								commitment_txid: *commitment_txid,
+								output_idx: output_idx as u32,
+								value: output.value,
+							}).is_none());
+					},
+					channelmonitor::ChannelMonitorUpdateStep::CommitmentSecret { idx, secret } => {
+						// Build the justice tx and store it
+						let mut channel_watchtower_state = self.channel_watchtower_state.lock().unwrap();
+						let justice_tx = match channel_watchtower_state.get(&funding_txo).unwrap().get(idx) {
+							Some(WatchtowerState::CounterpartyCommitmentTxSeen { commitment_txid, output_idx, value }) => {
+								let justice_tx = self.build_justice_tx(*commitment_txid, *output_idx, *value);
+								self.sign_justice_tx(justice_tx, secret, *value, data).expect("Should be able to sign justice tx")
+							},
+							_ => {
+								println!("No CounterpartyCommitmentTxSeen found, revokeable output must have been below dust limit.");
+								continue;
+							},
+						};
+
+						channel_watchtower_state.get_mut(&funding_txo).unwrap()
+							.insert(*idx, WatchtowerState::JusticeTxFormed(justice_tx))
+							.expect("Should only happen after previous update");
+					},
+					_ => {},
+				}
+			}
+		}
+
+		self.persister.update_persisted_channel(funding_txo, update, data, update_id)
+	}
+}
+
 pub struct TestPersister {
 	/// The queue of update statuses we'll return. If none are queued, ::Completed will always be
 	/// returned.