MutinyWallet · TonyGiorgio · Nov 1, 2023 · Oct 31, 2023
diff --git a/mutiny-core/src/node.rs b/mutiny-core/src/node.rs
@@ -1,7 +1,11 @@
 use crate::labels::LabelStorage;
 use crate::ldkstorage::{persist_monitor, ChannelOpenParams};
+use crate::multiesplora::MultiEsploraClient;
 use crate::nodemanager::ChannelClosure;
+use crate::peermanager::LspMessageRouter;
+use crate::scb::message_handler::SCBMessageHandler;
 use crate::scb::StaticChannelBackup;
+use crate::utils::get_monitor_version;
 use crate::{
     background::process_events_async,
     chain::MutinyChain,
@@ -18,33 +22,24 @@ use crate::{
     peermanager::{GossipMessageHandler, PeerManager, PeerManagerImpl},
     utils::{self, sleep},
 };
-use crate::{keymanager::PhantomKeysManager, scorer::HubPreferentialScorer};
-
-use crate::scb::message_handler::SCBMessageHandler;
 use crate::{fees::P2WSH_OUTPUT_SIZE, peermanager::connect_peer_if_necessary};
+use crate::{keymanager::PhantomKeysManager, scorer::HubPreferentialScorer};
 use crate::{lspclient::FeeRequest, storage::MutinyStorage};
 use anyhow::{anyhow, Context};
 use bdk::FeeRate;
 use bitcoin::hashes::{hex::ToHex, sha256::Hash as Sha256};
 use bitcoin::secp256k1::rand;
+use bitcoin::util::bip32::ExtendedPrivKey;
 use bitcoin::{hashes::Hash, secp256k1::PublicKey, BlockHash, Network, OutPoint};
 use core::time::Duration;
-use lightning::chain::channelmonitor::ChannelMonitor;
-use lightning::util::ser::{ReadableArgs, Writeable};
-use lightning::{
-    ln::channelmanager::{RecipientOnionFields, RetryableSendFailure},
-    routing::scoring::ProbabilisticScoringFeeParameters,
-    util::config::ChannelConfig,
-};
-
-use crate::multiesplora::MultiEsploraClient;
-use crate::utils::get_monitor_version;
-use bitcoin::util::bip32::ExtendedPrivKey;
 use futures_util::lock::Mutex;
+use lightning::chain::channelmonitor::ChannelMonitor;
 use lightning::events::bump_transaction::{BumpTransactionEventHandler, Wallet};
 use lightning::ln::PaymentSecret;
+use lightning::onion_message::OnionMessenger as LdkOnionMessenger;
 use lightning::sign::{EntropySource, InMemorySigner, NodeSigner, Recipient};
 use lightning::util::config::MaxDustHTLCExposure;
+use lightning::util::ser::{ReadableArgs, Writeable};
 use lightning::{
     chain::{chainmonitor, Filter, Watch},
     ln::{
@@ -63,6 +58,11 @@ use lightning::{
         logger::Logger,
     },
 };
+use lightning::{
+    ln::channelmanager::{RecipientOnionFields, RetryableSendFailure},
+    routing::scoring::ProbabilisticScoringFeeParameters,
+    util::config::ChannelConfig,
+};
 use lightning_invoice::payment::PaymentError;
 use lightning_invoice::{
     payment::{pay_invoice, pay_zero_value_invoice},
@@ -94,10 +94,19 @@ pub(crate) type RapidGossipSync =
 
 pub(crate) type NetworkGraph = gossip::NetworkGraph<Arc<MutinyLogger>>;
 
+pub(crate) type OnionMessenger<S: MutinyStorage> = LdkOnionMessenger<
+    Arc<PhantomKeysManager<S>>,
+    Arc<PhantomKeysManager<S>>,
+    Arc<MutinyLogger>,
+    Arc<LspMessageRouter>,
+    Arc<PhantomChannelManager<S>>,
+    IgnoringMessageHandler,
+>;
+
 pub(crate) type MessageHandler<S: MutinyStorage> = LdkMessageHandler<
     Arc<PhantomChannelManager<S>>,
     Arc<GossipMessageHandler<S>>,
-    Arc<IgnoringMessageHandler>,
+    Arc<OnionMessenger<S>>,
     Arc<SCBMessageHandler>,
 >;
 
@@ -310,21 +319,6 @@ impl<S: MutinyStorage> Node<S> {
             }
         }
 
-        let route_handler = Arc::new(GossipMessageHandler {
-            storage: persister.storage.clone(),
-            network_graph: gossip_sync.network_graph().clone(),
-            logger: logger.clone(),
-        });
-
-        // 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: Arc::new(IgnoringMessageHandler {}),
-            custom_message_handler: scb_message_handler.clone(),
-        };
-
         log_info!(logger, "creating lsp client");
         let lsp_client: Option<LspClient> = match node_index.lsp {
             None => {
@@ -341,6 +335,31 @@ impl<S: MutinyStorage> Node<S> {
             }
             Some(ref lsp) => lsp_clients.iter().find(|c| &c.url == lsp).cloned(),
         };
+        let lsp_client_pubkey = lsp_client.clone().map(|lsp| lsp.pubkey);
+        let message_router = Arc::new(LspMessageRouter::new(lsp_client_pubkey));
+        let onion_message_handler = Arc::new(OnionMessenger::new(
+            keys_manager.clone(),
+            keys_manager.clone(),
+            logger.clone(),
+            message_router,
+            channel_manager.clone(),
+            IgnoringMessageHandler {},
+        ));
+
+        let route_handler = Arc::new(GossipMessageHandler {
+            storage: persister.storage.clone(),
+            network_graph: gossip_sync.network_graph().clone(),
+            logger: logger.clone(),
+        });
+
+        // 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),
@@ -349,8 +368,6 @@ impl<S: MutinyStorage> Node<S> {
             Arc::clone(&logger),
         ));
 
-        let lsp_client_pubkey = lsp_client.clone().map(|lsp| lsp.pubkey);
-
         // init event handler
         let event_handler = EventHandler::new(
             channel_manager.clone(),

diff --git a/mutiny-core/src/peermanager.rs b/mutiny-core/src/peermanager.rs
@@ -1,28 +1,28 @@
-use crate::node::NetworkGraph;
+use crate::networking::socket::{schedule_descriptor_read, MutinySocketDescriptor};
+use crate::node::{NetworkGraph, OnionMessenger};
+use crate::scb::message_handler::SCBMessageHandler;
 use crate::storage::MutinyStorage;
 use crate::{error::MutinyError, fees::MutinyFeeEstimator};
 use crate::{gossip, ldkstorage::PhantomChannelManager, logging::MutinyLogger};
 use crate::{gossip::read_peer_info, node::PubkeyConnectionInfo};
 use crate::{keymanager::PhantomKeysManager, node::ConnectionType};
 use bitcoin::secp256k1::PublicKey;
-use lightning::{
-    ln::{msgs::SocketAddress, peer_handler::SocketDescriptor as LdkSocketDescriptor},
-    log_debug, log_trace,
-};
-use std::{net::SocketAddr, sync::atomic::AtomicBool};
-
-use crate::networking::socket::{schedule_descriptor_read, MutinySocketDescriptor};
-use crate::scb::message_handler::SCBMessageHandler;
 use lightning::events::{MessageSendEvent, MessageSendEventsProvider};
 use lightning::ln::features::{InitFeatures, NodeFeatures};
 use lightning::ln::msgs;
 use lightning::ln::msgs::{LightningError, RoutingMessageHandler};
 use lightning::ln::peer_handler::PeerHandleError;
-use lightning::ln::peer_handler::{IgnoringMessageHandler, PeerManager as LdkPeerManager};
+use lightning::ln::peer_handler::PeerManager as LdkPeerManager;
 use lightning::log_warn;
+use lightning::onion_message::{Destination, MessageRouter, OnionMessagePath};
 use lightning::routing::gossip::NodeId;
 use lightning::util::logger::Logger;
+use lightning::{
+    ln::{msgs::SocketAddress, peer_handler::SocketDescriptor as LdkSocketDescriptor},
+    log_debug, log_trace,
+};
 use std::sync::Arc;
+use std::{net::SocketAddr, sync::atomic::AtomicBool};
 
 #[cfg(target_arch = "wasm32")]
 use crate::networking::ws_socket::WsTcpSocketDescriptor;
@@ -91,7 +91,7 @@ pub(crate) type PeerManagerImpl<S: MutinyStorage> = LdkPeerManager<
     MutinySocketDescriptor,
     Arc<PhantomChannelManager<S>>,
     Arc<GossipMessageHandler<S>>,
-    Arc<IgnoringMessageHandler>,
+    Arc<OnionMessenger<S>>,
     Arc<MutinyLogger>,
     Arc<SCBMessageHandler>,
     Arc<PhantomKeysManager<S>>,
@@ -294,6 +294,49 @@ impl<S: MutinyStorage> RoutingMessageHandler for GossipMessageHandler<S> {
     }
 }
 
+/// LDK currently can't route onion messages, so we need to do it ourselves
+/// We just assume they are connected to us or the LSP.
+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,