Support onion message replies in OnionMessenger

Modify onion message handlers to return an optional response message for
OnionMessenger to reply with.

Author: jkczyz

fuzz/src/onion_message.rs

@@ -68,7 +68,9 @@ impl MessageRouter for TestMessageRouter {
 struct TestOffersMessageHandler {}
 
 impl OffersMessageHandler for TestOffersMessageHandler {
-	fn handle_message(&self, _message: OffersMessage) {}
+	fn handle_message(&self, _message: OffersMessage) -> Option<OffersMessage> {
+		None
+	}
 }
 
 struct TestCustomMessage {}
@@ -92,7 +94,9 @@ struct TestCustomMessageHandler {}
 
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
-	fn handle_custom_message(&self, _msg: Self::CustomMessage) {}
+	fn handle_custom_message(&self, _msg: Self::CustomMessage) -> Option<Self::CustomMessage> {
+		None
+	}
 	fn read_custom_message<R: io::Read>(&self, _message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, msgs::DecodeError> {
 		let mut buf = Vec::new();
 		buffer.read_to_end(&mut buf)?;

lightning/src/ln/peer_handler.rs

@@ -119,11 +119,11 @@ impl OnionMessageHandler for IgnoringMessageHandler {
 	}
 }
 impl OffersMessageHandler for IgnoringMessageHandler {
-	fn handle_message(&self, _msg: OffersMessage) {}
+	fn handle_message(&self, _msg: OffersMessage) -> Option<OffersMessage> { None }
 }
 impl CustomOnionMessageHandler for IgnoringMessageHandler {
 	type CustomMessage = Infallible;
-	fn handle_custom_message(&self, _msg: Infallible) {
+	fn handle_custom_message(&self, _msg: Infallible) -> Option<Infallible> {
 		// Since we always return `None` in the read the handle method should never be called.
 		unreachable!();
 	}

lightning/src/onion_message/functional_tests.rs

@@ -58,8 +58,8 @@ impl MessageRouter for TestMessageRouter {
 struct TestOffersMessageHandler {}
 
 impl OffersMessageHandler for TestOffersMessageHandler {
-	fn handle_message(&self, _message: OffersMessage) {
-		todo!()
+	fn handle_message(&self, _message: OffersMessage) -> Option<OffersMessage> {
+		None
 	}
 }
 
@@ -104,8 +104,9 @@ impl Drop for TestCustomMessageHandler {
 
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
-	fn handle_custom_message(&self, _msg: Self::CustomMessage) {
+	fn handle_custom_message(&self, _msg: Self::CustomMessage) -> Option<Self::CustomMessage> {
 		self.num_messages_expected.fetch_sub(1, Ordering::SeqCst);
+		None
 	}
 	fn read_custom_message<R: io::Read>(&self, message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, DecodeError> where Self: Sized {
 		if message_type == CUSTOM_MESSAGE_TYPE {

lightning/src/onion_message/messenger.rs

@@ -224,8 +224,10 @@ pub trait CustomOnionMessageHandler {
 	/// The message known to the handler. To support multiple message types, you may want to make this
 	/// an enum with a variant for each supported message.
 	type CustomMessage: CustomOnionMessageContents;
-	/// Called with the custom message that was received.
-	fn handle_custom_message(&self, msg: Self::CustomMessage);
+
+	/// Called with the custom message that was received, returning a response to send, if any.
+	fn handle_custom_message(&self, msg: Self::CustomMessage) -> Option<Self::CustomMessage>;
+
 	/// Read a custom message of type `message_type` from `buffer`, returning `Ok(None)` if the
 	/// message type is unknown.
 	fn read_custom_message<R: io::Read>(&self, message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, msgs::DecodeError>;
@@ -320,6 +322,56 @@ where
 		}
 	}
 
+	fn respond_with_onion_message<T: CustomOnionMessageContents>(
+		&self, response: OnionMessageContents<T>, path_id: Option<[u8; 32]>,
+		reply_path: Option<BlindedPath>
+	) {
+		let sender = match self.node_signer.get_node_id(Recipient::Node) {
+			Ok(node_id) => node_id,
+			Err(_) => {
+				log_warn!(
+					self.logger, "Unable to retrieve node id when responding to onion message with \
+					path_id {:02x?}", path_id
+				);
+				return;
+			}
+		};
+
+		let peers = self.pending_messages.lock().unwrap().keys().copied().collect();
+
+		let destination = match reply_path {
+			Some(reply_path) => Destination::BlindedPath(reply_path),
+			None => {
+				log_trace!(
+					self.logger, "Missing reply path when responding to onion message with path_id \
+					{:02x?}", path_id
+				);
+				return;
+			},
+		};
+
+		let path = match self.message_router.find_path(sender, peers, destination) {
+			Ok(path) => path,
+			Err(()) => {
+				log_trace!(
+					self.logger, "Failed to find path when responding to onion message with \
+					path_id {:02x?}", path_id
+				);
+				return;
+			},
+		};
+
+		log_trace!(self.logger, "Responding to onion message with path_id {:02x?}", path_id);
+
+		if let Err(e) = self.send_onion_message(path, response, None) {
+			log_trace!(
+				self.logger, "Failed responding to onion message with path_id {:02x?}: {:?}",
+				path_id, e
+			);
+			return;
+		}
+	}
+
 	#[cfg(test)]
 	pub(super) fn release_pending_msgs(&self) -> HashMap<PublicKey, VecDeque<msgs::OnionMessage>> {
 		let mut pending_msgs = self.pending_messages.lock().unwrap();
@@ -403,9 +455,20 @@ where
 				log_info!(self.logger,
 					"Received an onion message with path_id {:02x?} and {} reply_path",
 						path_id, if reply_path.is_some() { "a" } else { "no" });
-				match message {
-					OnionMessageContents::Offers(msg) => self.offers_handler.handle_message(msg),
-					OnionMessageContents::Custom(msg) => self.custom_handler.handle_custom_message(msg),
+
+				let response = match message {
+					OnionMessageContents::Offers(msg) => {
+						self.offers_handler.handle_message(msg)
+							.map(|msg| OnionMessageContents::Offers(msg))
+					},
+					OnionMessageContents::Custom(msg) => {
+						self.custom_handler.handle_custom_message(msg)
+							.map(|msg| OnionMessageContents::Custom(msg))
+					},
+				};
+
+				if let Some(response) = response {
+					self.respond_with_onion_message(response, path_id, reply_path);
 				}
 			},
 			Ok((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {

lightning/src/onion_message/offers.rs

@@ -32,7 +32,7 @@ const INVOICE_ERROR_TLV_TYPE: u64 = 68;
 pub trait OffersMessageHandler {
 	/// Handles the given message by either responding with an [`Invoice`], sending a payment, or
 	/// replying with an error.
-	fn handle_message(&self, message: OffersMessage);
+	fn handle_message(&self, message: OffersMessage) -> Option<OffersMessage>;
 }
 
 /// Possible BOLT 12 Offers messages sent and received via an [`OnionMessage`].