Hand payment preimages into channel monitors when claimed

diff --git a/src/ln/channel.rs b/src/ln/channel.rs
index b8a5bb66c5d39279f2e322d37ea7e4dfe4861703..4b6cd08ff4ada2fc3b122c6332bb8b6fac6691b7 100644 (file)
--- 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 d2854ca741145eaeaefc11ab673cbe53c71d3f2c..227581389dc9b5a90909323dae68d42959337e78 100644 (file)
--- 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 53440f4361341c7af9d5d441ebca3559030698e4..27010d522dfe7037485c4bfac6f18237417df829 100644 (file)
--- 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