Merge pull request #2787 from jkczyz/2023-12-onion-messenger-assert

[rust-lightning] / lightning / src / onion_message / messenger.rs

diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
index a6f11bf6a7d75e07be52442569c42ed87d610870..8a44eb2a5beebf3128582bdbce23cc8edfbda381 100644 (file)
--- a/lightning/src/onion_message/messenger.rs
+++ b/lightning/src/onion_message/messenger.rs
@@ -223,6 +223,13 @@ impl OnionMessageRecipient {
                        *self = OnionMessageRecipient::ConnectedPeer(new_pending_messages);
                }
        }
+
+       fn is_connected(&self) -> bool {
+               match self {
+                       OnionMessageRecipient::ConnectedPeer(..) => true,
+                       OnionMessageRecipient::PendingConnection(..) => false,
+               }
+       }
 }
 
 /// An [`OnionMessage`] for [`OnionMessenger`] to send.
@@ -729,7 +736,11 @@ where
                        },
                        hash_map::Entry::Occupied(mut e) => {
                                e.get_mut().enqueue_message(onion_message);
-                               Ok(SendSuccess::Buffered)
+                               if e.get().is_connected() {
+                                       Ok(SendSuccess::Buffered)
+                               } else {
+                                       Ok(SendSuccess::BufferedAwaitingConnection(first_node_id))
+                               }
                        },
                }
        }
@@ -903,7 +914,8 @@ where
        fn peer_disconnected(&self, their_node_id: &PublicKey) {
                match self.message_recipients.lock().unwrap().remove(their_node_id) {
                        Some(OnionMessageRecipient::ConnectedPeer(..)) => {},
-                       _ => debug_assert!(false),
+                       Some(_) => debug_assert!(false),
+                       None => {},
                }
        }