From cb8f4f17b018f6a091388352327d577d011f2f38 Mon Sep 17 00:00:00 2001
From: benthecarman <benthecarman@live.com>
Date: Thu, 26 Oct 2023 18:32:55 -0500
Subject: [PATCH] Custom message router

---
 mutiny-core/src/node.rs        | 43 ++++++++++++++++++----------------
 mutiny-core/src/peermanager.rs | 42 +++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/mutiny-core/src/node.rs b/mutiny-core/src/node.rs
index 63d4179f2..bb1f604f1 100644
--- a/mutiny-core/src/node.rs
+++ b/mutiny-core/src/node.rs
@@ -38,12 +38,13 @@ use lightning::{
 };
 
 use crate::multiesplora::MultiEsploraClient;
+use crate::peermanager::LspMessageRouter;
 use crate::utils::get_monitor_version;
 use bitcoin::util::bip32::ExtendedPrivKey;
 use lightning::events::bump_transaction::{BumpTransactionEventHandler, Wallet};
 use lightning::ln::PaymentSecret;
 use lightning::offers::offer::Offer;
-use lightning::onion_message::{DefaultMessageRouter, OnionMessenger as LdkOnionMessenger};
+use lightning::onion_message::OnionMessenger as LdkOnionMessenger;
 use lightning::sign::{EntropySource, InMemorySigner, NodeSigner, Recipient};
 use lightning::util::config::MaxDustHTLCExposure;
 use lightning::{
@@ -99,7 +100,7 @@ pub(crate) type OnionMessenger<S: MutinyStorage> = LdkOnionMessenger<
     Arc<PhantomKeysManager<S>>,
     Arc<PhantomKeysManager<S>>,
     Arc<MutinyLogger>,
-    Arc<DefaultMessageRouter>,
+    Arc<LspMessageRouter>,
     Arc<PhantomChannelManager<S>>,
     IgnoringMessageHandler,
 >;
@@ -325,24 +326,6 @@ impl<S: MutinyStorage> Node<S> {
             logger: logger.clone(),
         });
 
-        let onion_message_handler = Arc::new(OnionMessenger::new(
-            keys_manager.clone(),
-            keys_manager.clone(),
-            logger.clone(),
-            Arc::new(DefaultMessageRouter {}),
-            channel_manager.clone(),
-            IgnoringMessageHandler {},
-        ));
-
-        // init peer manager
-        let scb_message_handler = Arc::new(SCBMessageHandler::new());
-        let ln_msg_handler = MessageHandler {
-            chan_handler: channel_manager.clone(),
-            route_handler,
-            onion_message_handler,
-            custom_message_handler: scb_message_handler.clone(),
-        };
-
         log_info!(logger, "creating lsp client");
         let lsp_client: Option<LspClient> = match node_index.lsp {
             None => {
@@ -360,6 +343,26 @@ impl<S: MutinyStorage> Node<S> {
             Some(ref lsp) => lsp_clients.iter().find(|c| &c.url == lsp).cloned(),
         };
 
+        let message_router = LspMessageRouter::new(lsp_client.as_ref().map(|l| l.pubkey));
+
+        let onion_message_handler = Arc::new(OnionMessenger::new(
+            keys_manager.clone(),
+            keys_manager.clone(),
+            logger.clone(),
+            Arc::new(message_router),
+            channel_manager.clone(),
+            IgnoringMessageHandler {},
+        ));
+
+        // init peer manager
+        let scb_message_handler = Arc::new(SCBMessageHandler::new());
+        let ln_msg_handler = MessageHandler {
+            chan_handler: channel_manager.clone(),
+            route_handler,
+            onion_message_handler,
+            custom_message_handler: scb_message_handler.clone(),
+        };
+
         let bump_tx_event_handler = Arc::new(BumpTransactionEventHandler::new(
             Arc::clone(&chain),
             Arc::new(Wallet::new(Arc::clone(&wallet), Arc::clone(&logger))),
diff --git a/mutiny-core/src/peermanager.rs b/mutiny-core/src/peermanager.rs
index 61c2f2180..fc0e1fe9a 100644
--- a/mutiny-core/src/peermanager.rs
+++ b/mutiny-core/src/peermanager.rs
@@ -33,6 +33,7 @@ use crate::networking::proxy::WsProxy;
 #[cfg(not(target_arch = "wasm32"))]
 use tokio::time;
 
+use lightning::onion_message::{Destination, MessageRouter, OnionMessagePath};
 #[cfg(not(target_arch = "wasm32"))]
 use std::time::Duration;
 
@@ -294,6 +295,47 @@ impl<S: MutinyStorage> RoutingMessageHandler for GossipMessageHandler<S> {
     }
 }
 
+pub struct LspMessageRouter {
+    intermediate_nodes: Vec<PublicKey>,
+}
+
+impl LspMessageRouter {
+    pub fn new(lsp_pubkey: Option<PublicKey>) -> Self {
+        let intermediate_nodes = match lsp_pubkey {
+            Some(pubkey) => vec![pubkey],
+            None => vec![],
+        };
+
+        Self { intermediate_nodes }
+    }
+}
+
+impl MessageRouter for LspMessageRouter {
+    fn find_path(
+        &self,
+        _sender: PublicKey,
+        peers: Vec<PublicKey>,
+        destination: Destination,
+    ) -> Result<OnionMessagePath, ()> {
+        let first_node = match &destination {
+            Destination::Node(node_id) => *node_id,
+            Destination::BlindedPath(path) => path.introduction_node_id,
+        };
+
+        if peers.contains(&first_node) {
+            Ok(OnionMessagePath {
+                intermediate_nodes: vec![],
+                destination,
+            })
+        } else {
+            Ok(OnionMessagePath {
+                intermediate_nodes: self.intermediate_nodes.clone(),
+                destination,
+            })
+        }
+    }
+}
+
 pub(crate) async fn connect_peer_if_necessary<S: MutinyStorage>(
     #[cfg(target_arch = "wasm32")] websocket_proxy_addr: &str,
     peer_connection_info: &PubkeyConnectionInfo,