From 3bc12e192feb25f1bd91dbf3d14bc541fa78cfe5 Mon Sep 17 00:00:00 2001
From: Carla Kirk-Cohen <kirkcohenc@gmail.com>
Date: Mon, 27 Mar 2023 13:20:47 -0400
Subject: [PATCH] main: add incoming message events to consumer

---
 Cargo.toml  |  3 ++-
 src/main.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index d99557b..38a204f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,9 +21,10 @@ tonic = "0.8.3"
 tonic_lnd = "0.5.1"
 hex = "0.4.3"
 configure_me = "0.4.0"
+bytes = "1.4.0"
 
 [dev-dependencies]
 mockall = "0.11.3"
 
 [build-dependencies]
-configure_me_codegen = "0.4.4"
\ No newline at end of file
+configure_me_codegen = "0.4.4"
diff --git a/src/main.rs b/src/main.rs
index 0ca2532..883742b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -20,7 +20,7 @@ use futures::executor::block_on;
 use internal::*;
 use lightning::chain::keysinterface::{EntropySource, KeyMaterial, NodeSigner, Recipient};
 use lightning::ln::features::InitFeatures;
-use lightning::ln::msgs::{Init, OnionMessageHandler, UnsignedGossipMessage};
+use lightning::ln::msgs::{Init, OnionMessage, OnionMessageHandler, UnsignedGossipMessage};
 use lightning::ln::peer_handler::IgnoringMessageHandler;
 use lightning::onion_message::{CustomOnionMessageHandler, OnionMessenger};
 use lightning::util::logger::{Level, Logger, Record};
@@ -362,6 +362,7 @@ impl fmt::Display for ConsumerError {
 enum MessengerEvents {
     PeerConnected(PublicKey, bool),
     PeerDisconnected(PublicKey),
+    IncomingMessage(PublicKey, OnionMessage),
     ProducerExit(ConsumerError),
 }
 
@@ -375,6 +376,9 @@ impl fmt::Display for MessengerEvents {
                 )
             }
             MessengerEvents::PeerDisconnected(p) => write!(f, "messenger event: {p} disconnected"),
+            MessengerEvents::IncomingMessage(p, _) => {
+                write!(f, "messenger event: {p} incoming onion message")
+            }
             MessengerEvents::ProducerExit(s) => write!(f, "messenger event: {s} exited"),
         }
     }
@@ -409,6 +413,9 @@ async fn consume_messenger_events(
                     .map_err(|_| ConsumerError::OnionMessengerFailure)?
             }
             MessengerEvents::PeerDisconnected(pubkey) => onion_messenger.peer_disconnected(&pubkey),
+            MessengerEvents::IncomingMessage(pubkey, onion_message) => {
+                onion_messenger.handle_onion_message(&pubkey, &onion_message)
+            }
             MessengerEvents::ProducerExit(e) => {
                 return Err(e);
             }
@@ -688,12 +695,15 @@ async fn set_feature_bit(
 mod tests {
     use super::*;
     use bitcoin::secp256k1::PublicKey;
+    use bytes::BufMut;
     use hex;
     use lightning::ln::features::{InitFeatures, NodeFeatures};
     use lightning::ln::msgs::{OnionMessage, OnionMessageHandler};
     use lightning::util::events::OnionMessageProvider;
+    use lightning::util::ser::Readable;
     use mockall::mock;
     use std::collections::HashMap;
+    use std::io::Cursor;
     use tokio::sync::mpsc::channel;
 
     mock! {
@@ -842,6 +852,33 @@ mod tests {
         .unwrap()
     }
 
+    // Produces an OnionMessage that can be used for tests. We need to manually write individual bytes because onion
+    // messages in LDK can only be created using read/write impls that deal with raw bytes (since some other fields
+    // are not public).
+    fn onion_message() -> OnionMessage {
+        let mut w = vec![];
+        let pubkey_bytes = pubkey().serialize();
+
+        // Blinding point for the onion message.
+        w.put_slice(&pubkey_bytes);
+
+        // Write the length of the onion packet:
+        // Version: 1
+        // Ephemeral Key: 33
+        // Hop Payloads: 1300
+        // HMAC: 32.
+        w.put_u16(1 + 33 + 1300 + 32);
+
+        // Write meaningless contents for the actual values.
+        w.put_u8(0);
+        w.put_slice(&pubkey_bytes);
+        w.put_bytes(1, 1300);
+        w.put_bytes(2, 32);
+
+        let mut readable = Cursor::new(w);
+        OnionMessage::read(&mut readable).unwrap()
+    }
+
     mock! {
             OnionHandler{}
 
@@ -870,7 +907,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_consume_messenger_events() {
-        let (sender, receiver) = channel(4);
+        let (sender, receiver) = channel(5);
 
         let pk = pubkey();
         let mut mock = MockOnionHandler::new();
@@ -902,6 +939,14 @@ mod tests {
             .unwrap();
         mock.expect_peer_disconnected().return_once(|_| ());
 
+        // Cover incoming onion messages.
+        let onion_message = onion_message();
+        sender
+            .send(MessengerEvents::IncomingMessage(pk, onion_message))
+            .await
+            .unwrap();
+        mock.expect_handle_onion_message().return_once(|_, _| ());
+
         // Finally, send a producer exit event to test exit.
         sender
             .send(MessengerEvents::ProducerExit(