From 85e890be4909361b7163fa09fa4bb271463af271 Mon Sep 17 00:00:00 2001
From: Valentine Wallace <vwallace@protonmail.com>
Date: Wed, 14 Apr 2021 14:58:27 -0400
Subject: [PATCH] Check payment secret before succeeding payment

---
 src/main.rs | 45 +++++++++++++++++++++++++++------------------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 7001129..927284a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -149,34 +149,43 @@ async fn handle_ldk_events(
 						.funding_transaction_generated(&temporary_channel_id, final_tx)
 						.unwrap();
 				}
-				Event::PaymentReceived { payment_hash, .. } => {
+				Event::PaymentReceived { payment_hash, payment_secret, amt } => {
 					let mut payments = inbound_payments.lock().unwrap();
 					if let Some(payment) = payments.get_mut(&payment_hash) {
-						assert!(loop_channel_manager.claim_funds(
-							payment.preimage.unwrap().clone(),
-							&payment.secret,
-							payment.amt_msat.0.unwrap(),
-						));
-						println!(
-							"\nEVENT: received payment from payment_hash {} of {} millisatoshis",
-							hex_utils::hex_str(&payment_hash.0),
-							payment.amt_msat
-						);
-						print!("> ");
-						io::stdout().flush().unwrap();
-						payment.status = HTLCStatus::Succeeded;
+						if payment.secret == payment_secret {
+							assert!(loop_channel_manager.claim_funds(
+								payment.preimage.unwrap().clone(),
+								&payment.secret,
+								payment.amt_msat.0.unwrap(),
+							));
+							println!(
+							        "\nEVENT: received payment from payment hash {} of {} millisatoshis",
+							        hex_utils::hex_str(&payment_hash.0),
+							        payment.amt_msat
+						      );
+							print!("> ");
+							io::stdout().flush().unwrap();
+							payment.status = HTLCStatus::Succeeded;
+						} else {
+							loop_channel_manager
+								.fail_htlc_backwards(&payment_hash, &payment.secret);
+							println!("\nERROR: we received a payment from payment hash {} but the payment secret didn't match", hex_utils::hex_str(&payment_hash.0));
+							print!("> ");
+							io::stdout().flush().unwrap();
+							payment.status = HTLCStatus::Failed;
+						}
 					} else {
-						println!("\nERROR: we received a payment but didn't know the preimage");
+						loop_channel_manager.fail_htlc_backwards(&payment_hash, &payment_secret);
+						println!("\nERROR: we received a payment for payment hash {} but didn't know the preimage", hex_utils::hex_str(&payment_hash.0));
 						print!("> ");
 						io::stdout().flush().unwrap();
-						loop_channel_manager.fail_htlc_backwards(&payment_hash, &None);
 						payments.insert(
 							payment_hash,
 							PaymentInfo {
 								preimage: None,
-								secret: None,
+								secret: payment_secret,
 								status: HTLCStatus::Failed,
-								amt_msat: MillisatAmount(None),
+								amt_msat: MillisatAmount(Some(amt)),
 							},
 						);
 					}
-- 
2.30.2