From 0a7549ce8e9bbc90203bce224de01b74dea98929 Mon Sep 17 00:00:00 2001
From: Orbital <orbitalturtle@protonmail.com>
Date: Wed, 24 Jul 2024 21:22:59 -0500
Subject: [PATCH] messenger: override LDK's DefaultMessageRouter

---
 src/lib.rs             |  6 +++--
 src/onion_messenger.rs | 54 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 5a97d43a..6b1a57ea 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -16,7 +16,7 @@ use crate::lnd::{
     MIN_LND_PRE_RELEASE_VER,
 };
 use crate::lndk_offers::{OfferError, SendPaymentParams};
-use crate::onion_messenger::{LndkNodeIdLookUp, MessengerUtilities};
+use crate::onion_messenger::{LndkMessageRouter, LndkNodeIdLookUp, MessengerUtilities};
 use bitcoin::network::constants::Network;
 use bitcoin::secp256k1::{PublicKey, Secp256k1};
 use home::home_dir;
@@ -216,7 +216,9 @@ impl LndkOnionMessenger {
         let node_signer = LndNodeSigner::new(pubkey, &mut node_client);
         let messenger_utils = MessengerUtilities::new();
         let network_graph = &NetworkGraph::new(network, &messenger_utils);
-        let message_router = &DefaultMessageRouter::new(network_graph, &messenger_utils);
+        let default_message_router = DefaultMessageRouter::new(network_graph, &messenger_utils);
+        let message_router =
+            &LndkMessageRouter(default_message_router, client.clone().lightning_read_only());
         let node_id_lookup = LndkNodeIdLookUp::new(client.clone(), pubkey);
         let onion_messenger = OnionMessenger::new(
             &messenger_utils,
diff --git a/src/onion_messenger.rs b/src/onion_messenger.rs
index fa02151f..b049cab6 100644
--- a/src/onion_messenger.rs
+++ b/src/onion_messenger.rs
@@ -5,16 +5,18 @@ use crate::{LifecycleSignals, LndkOnionMessenger, LDK_LOGGER_NAME};
 use async_trait::async_trait;
 use bitcoin::blockdata::constants::ChainHash;
 use bitcoin::network::constants::Network;
-use bitcoin::secp256k1::PublicKey;
+use bitcoin::secp256k1::{PublicKey, Secp256k1, Signing, Verification};
 use core::ops::Deref;
 use futures::executor::block_on;
-use lightning::blinded_path::NodeIdLookUp;
+use lightning::blinded_path::{BlindedPath, NodeIdLookUp};
 use lightning::ln::features::InitFeatures;
 use lightning::ln::msgs::{Init, OnionMessage, OnionMessageHandler};
 use lightning::onion_message::messenger::{
-    CustomOnionMessageHandler, MessageRouter, OnionMessenger,
+    CustomOnionMessageHandler, DefaultMessageRouter, Destination, MessageRouter, OnionMessagePath,
+    OnionMessenger,
 };
 use lightning::onion_message::offers::OffersMessageHandler;
+use lightning::routing::gossip::NetworkGraph;
 use lightning::sign::EntropySource;
 use lightning::sign::NodeSigner;
 use lightning::util::logger::{Level, Logger, Record};
@@ -52,6 +54,52 @@ const DEFAULT_CALL_COUNT: u8 = 10;
 /// DEFAULT_CALL_FREQUENCY is the default period over which peers are rate limited.
 const DEFAULT_CALL_FREQUENCY: Duration = Duration::from_secs(1);
 
+/// LndkMessageRouter temporarily helps us to avoid a race condition we sometimes experience when
+/// making payments. What happens is, we auto-connect to the introduction node of a blinded path
+/// provided by a offer. But because the OnionMessenger hasn't yet processed the PeerConnected
+/// event, it thinks we aren't connected to the introduction node, and it throws an error that
+/// it can't find a route. So, here in find_path we add a quick check in ListPeers and add that
+/// peer to the list.
+///
+/// Once we've fully implemented a NetworkGraph we should remove this.
+#[allow(dead_code)]
+pub(crate) struct LndkMessageRouter<G: Deref<Target = NetworkGraph<L>>, L: Deref, ES: Deref>(
+    pub(crate) DefaultMessageRouter<G, L, ES>,
+    pub(crate) LightningClient,
+)
+where
+    L::Target: Logger,
+    ES::Target: EntropySource;
+
+impl<G: Deref<Target = NetworkGraph<L>>, L: Deref, ES: Deref> MessageRouter
+    for LndkMessageRouter<G, L, ES>
+where
+    L::Target: Logger,
+    ES::Target: EntropySource,
+{
+    fn find_path(
+        &self,
+        sender: PublicKey,
+        peers: Vec<PublicKey>,
+        destination: Destination,
+    ) -> Result<OnionMessagePath, ()> {
+        // TODO: Instead of simply returning an error here, we should somehow pass the error to
+        // OfferHandler. Perhaps by doing something like...  passing a channel to LndkMessageRouter
+        // when it's instantiated that OfferHandler also has access to. The OfferHandler will
+        // need access to this so it can return the error to the user properly.
+        self.0.find_path(sender, peers, destination)
+    }
+
+    fn create_blinded_paths<T: Signing + Verification>(
+        &self,
+        recipient: PublicKey,
+        peers: Vec<PublicKey>,
+        secp_ctx: &Secp256k1<T>,
+    ) -> Result<Vec<BlindedPath>, ()> {
+        self.0.create_blinded_paths(recipient, peers, secp_ctx)
+    }
+}
+
 /// Node Id LookUp is a utility struct implementing NodeIdLookUp trait for LDK's OnionMessenger.
 pub struct LndkNodeIdLookUp {
     client: Client,