From 021b3643ccb4237dfae6c8686b902f65cdb260b9 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Mon, 2 Apr 2018 11:45:40 -0400
Subject: [PATCH] Hand payment preimages into channel monitors when claimed

---
 src/ln/channel.rs        |  2 ++
 src/ln/channelmanager.rs | 33 ++++++++++++++++++++++-----------
 src/ln/channelmonitor.rs |  6 ++++++
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/ln/channel.rs b/src/ln/channel.rs
index b8a5bb66c..4b6cd08ff 100644
--- a/src/ln/channel.rs
+++ b/src/ln/channel.rs
@@ -874,6 +874,7 @@ impl Channel {
 		if htlc_amount_msat == 0 {
 			return Err(HandleError{err: "Unable to find a pending HTLC which matched the given payment preimage", msg: None});
 		}
+		self.channel_monitor.provide_payment_preimage(&payment_preimage);
 
 		//TODO: This is racy af, they may have pending messages in flight to us that will not have
 		//received this yet!
@@ -1220,6 +1221,7 @@ impl Channel {
 				self.value_to_self_msat -= htlc.amount_msat;
 			}
 		}
+		self.channel_monitor.provide_payment_preimage(&msg.payment_preimage);
 		Ok(())
 	}
 
diff --git a/src/ln/channelmanager.rs b/src/ln/channelmanager.rs
index d2854ca74..227581389 100644
--- a/src/ln/channelmanager.rs
+++ b/src/ln/channelmanager.rs
@@ -838,7 +838,7 @@ impl ChannelManager {
 				false
 			},
 			PendingOutboundHTLC::IntermediaryHopData { source_short_channel_id, .. } => {
-				let (node_id, fulfill_msg) = {
+				let (node_id, fulfill_msg, monitor) = {
 					let chan_id = match channel_state.short_to_id.get(&source_short_channel_id) {
 						Some(chan_id) => chan_id.clone(),
 						None => return false
@@ -846,7 +846,7 @@ impl ChannelManager {
 
 					let chan = channel_state.by_id.get_mut(&chan_id).unwrap();
 					match chan.get_update_fulfill_htlc(payment_preimage) {
-						Ok(msg) => (chan.get_their_node_id(), msg),
+						Ok(msg) => (chan.get_their_node_id(), msg, if from_user { Some(chan.channel_monitor()) } else { None }),
 						Err(_e) => {
 							//TODO: Do something with e?
 							return false;
@@ -854,14 +854,23 @@ impl ChannelManager {
 					}
 				};
 
-				mem::drop(channel_state);
-				let mut pending_events = self.pending_events.lock().unwrap();
-				pending_events.push(events::Event::SendFulfillHTLC {
-					node_id: node_id,
-					msg: fulfill_msg
-				});
+				{
+					mem::drop(channel_state);
+					let mut pending_events = self.pending_events.lock().unwrap();
+					pending_events.push(events::Event::SendFulfillHTLC {
+						node_id: node_id,
+						msg: fulfill_msg
+					});
+				}
 
-				true
+				//TODO: It may not be possible to handle add_update_monitor fails gracefully, maybe
+				//it should return no Err? Sadly, panic!()s instead doesn't help much :(
+				if from_user {
+					match self.monitor.add_update_monitor(monitor.as_ref().unwrap().get_funding_txo().unwrap(), monitor.unwrap()) {
+						Ok(()) => true,
+						Err(_) => true,
+					}
+				} else { true }
 			},
 		}
 	}
@@ -1319,7 +1328,7 @@ impl ChannelMessageHandler for ChannelManager {
 	}
 
 	fn handle_update_fulfill_htlc(&self, their_node_id: &PublicKey, msg: &msgs::UpdateFulfillHTLC) -> Result<(), HandleError> {
-		{
+		let monitor = {
 			let mut channel_state = self.channel_state.lock().unwrap();
 			match channel_state.by_id.get_mut(&msg.channel_id) {
 				Some(chan) => {
@@ -1327,12 +1336,14 @@ impl ChannelMessageHandler for ChannelManager {
 						return Err(HandleError{err: "Got a message for a channel from the wrong node!", msg: None})
 					}
 					chan.update_fulfill_htlc(&msg)?;
+					chan.channel_monitor()
 				},
 				None => return Err(HandleError{err: "Failed to find corresponding channel", msg: None})
 			}
-		}
+		};
 		//TODO: Delay the claimed_funds relaying just like we do outbound relay!
 		self.claim_funds_internal(msg.payment_preimage.clone(), false);
+		self.monitor.add_update_monitor(monitor.get_funding_txo().unwrap(), monitor)?;
 		Ok(())
 	}
 
diff --git a/src/ln/channelmonitor.rs b/src/ln/channelmonitor.rs
index 53440f436..27010d522 100644
--- a/src/ln/channelmonitor.rs
+++ b/src/ln/channelmonitor.rs
@@ -270,6 +270,12 @@ impl ChannelMonitor {
 		min
 	}
 
+	pub fn provide_payment_preimage(&mut self, payment_preimage: &[u8; 32]) {
+		//TODO: Some kind of timeout here or ability to mark all states containing this preimage
+		//revoked?
+		self.payment_preimages.push(payment_preimage.clone());
+	}
+
 	#[inline]
 	fn check_spend_transaction(&self, tx: &Transaction, height: u32) -> Vec<Transaction> {
 		// Most secp and related errors trying to create keys means we have no hope of constructing
-- 
2.39.5