From 90f59060822fceea42841a8a03fed01f2c39eeaa Mon Sep 17 00:00:00 2001
From: Valentine Wallace <vwallace@protonmail.com>
Date: Fri, 9 Sep 2022 12:29:13 -0400
Subject: [PATCH] OR InitFeatures and NodeFeatures from onion message handler

Similar to how we OR our InitFeaures and NodeFeatures across both our channel
and routing message handlers, we also want to OR the features of our onion
message handler.
---
 lightning/src/ln/msgs.rs                 | 13 +++++++++++++
 lightning/src/ln/peer_handler.rs         | 13 ++++++++++---
 lightning/src/onion_message/messenger.rs | 13 +++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/lightning/src/ln/msgs.rs b/lightning/src/ln/msgs.rs
index a6d096bff..a43156965 100644
--- a/lightning/src/ln/msgs.rs
+++ b/lightning/src/ln/msgs.rs
@@ -980,6 +980,19 @@ pub trait OnionMessageHandler : OnionMessageProvider {
 	/// Indicates a connection to the peer failed/an existing connection was lost. Allows handlers to
 	/// drop and refuse to forward onion messages to this peer.
 	fn peer_disconnected(&self, their_node_id: &PublicKey, no_connection_possible: bool);
+
+	// Handler information:
+	/// Gets the node feature flags which this handler itself supports. All available handlers are
+	/// queried similarly and their feature flags are OR'd together to form the [`NodeFeatures`]
+	/// which are broadcasted in our [`NodeAnnouncement`] message.
+	fn provided_node_features(&self) -> NodeFeatures;
+
+	/// Gets the init feature flags which should be sent to the given peer. All available handlers
+	/// are queried similarly and their feature flags are OR'd together to form the [`InitFeatures`]
+	/// which are sent in our [`Init`] message.
+	///
+	/// Note that this method is called before [`Self::peer_connected`].
+	fn provided_init_features(&self, their_node_id: &PublicKey) -> InitFeatures;
 }
 
 mod fuzzy_internal_msgs {
diff --git a/lightning/src/ln/peer_handler.rs b/lightning/src/ln/peer_handler.rs
index 5e6938993..838927cb2 100644
--- a/lightning/src/ln/peer_handler.rs
+++ b/lightning/src/ln/peer_handler.rs
@@ -89,6 +89,10 @@ impl OnionMessageHandler for IgnoringMessageHandler {
 	fn handle_onion_message(&self, _their_node_id: &PublicKey, _msg: &msgs::OnionMessage) {}
 	fn peer_connected(&self, _their_node_id: &PublicKey, _init: &msgs::Init) {}
 	fn peer_disconnected(&self, _their_node_id: &PublicKey, _no_connection_possible: bool) {}
+	fn provided_node_features(&self) -> NodeFeatures { NodeFeatures::empty() }
+	fn provided_init_features(&self, _their_node_id: &PublicKey) -> InitFeatures {
+		InitFeatures::empty()
+	}
 }
 impl Deref for IgnoringMessageHandler {
 	type Target = IgnoringMessageHandler;
@@ -1062,7 +1066,8 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 								peer.their_node_id = Some(their_node_id);
 								insert_node_id!();
 								let features = self.message_handler.chan_handler.provided_init_features(&their_node_id)
-									.or(self.message_handler.route_handler.provided_init_features(&their_node_id));
+									.or(self.message_handler.route_handler.provided_init_features(&their_node_id))
+									.or(self.message_handler.onion_message_handler.provided_init_features(&their_node_id));
 								let resp = msgs::Init { features, remote_network_address: filter_addresses(peer.their_net_address.clone()) };
 								self.enqueue_message(peer, &resp);
 								peer.awaiting_pong_timer_tick_intervals = 0;
@@ -1075,7 +1080,8 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 								peer.their_node_id = Some(their_node_id);
 								insert_node_id!();
 								let features = self.message_handler.chan_handler.provided_init_features(&their_node_id)
-									.or(self.message_handler.route_handler.provided_init_features(&their_node_id));
+									.or(self.message_handler.route_handler.provided_init_features(&their_node_id))
+									.or(self.message_handler.onion_message_handler.provided_init_features(&their_node_id));
 								let resp = msgs::Init { features, remote_network_address: filter_addresses(peer.their_net_address.clone()) };
 								self.enqueue_message(peer, &resp);
 								peer.awaiting_pong_timer_tick_intervals = 0;
@@ -1971,7 +1977,8 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 		addresses.sort_by_key(|addr| addr.get_id());
 
 		let features = self.message_handler.chan_handler.provided_node_features()
-			.or(self.message_handler.route_handler.provided_node_features());
+			.or(self.message_handler.route_handler.provided_node_features())
+			.or(self.message_handler.onion_message_handler.provided_node_features());
 		let announcement = msgs::UnsignedNodeAnnouncement {
 			features,
 			timestamp: self.last_node_announcement_serial.fetch_add(1, Ordering::AcqRel) as u32,
diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
index e38a0d10f..2df01c2ba 100644
--- a/lightning/src/onion_message/messenger.rs
+++ b/lightning/src/onion_message/messenger.rs
@@ -16,6 +16,7 @@ use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::secp256k1::{self, PublicKey, Scalar, Secp256k1, SecretKey};
 
 use chain::keysinterface::{InMemorySigner, KeysInterface, KeysManager, Recipient, Sign};
+use ln::features::{InitFeatures, NodeFeatures};
 use ln::msgs::{self, OnionMessageHandler};
 use ln::onion_utils;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
@@ -345,6 +346,18 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessageHandler for OnionMessenger<Si
 		let mut pending_msgs = self.pending_messages.lock().unwrap();
 		pending_msgs.remove(their_node_id);
 	}
+
+	fn provided_node_features(&self) -> NodeFeatures {
+		let mut features = NodeFeatures::empty();
+		features.set_onion_messages_optional();
+		features
+	}
+
+	fn provided_init_features(&self, _their_node_id: &PublicKey) -> InitFeatures {
+		let mut features = InitFeatures::empty();
+		features.set_onion_messages_optional();
+		features
+	}
 }
 
 impl<Signer: Sign, K: Deref, L: Deref> OnionMessageProvider for OnionMessenger<Signer, K, L>
-- 
2.39.5