f: Address comments and feedbacks

1. Corrected implementation for onion_message.rs handle_message
2. Renamed received_onion_message -> message
3. Removed redundant MessageWithID struct
4. Seperated message from Responder.
5. Add docs

Author: shaavan

fuzz/src/onion_message.rs

@@ -96,9 +96,7 @@ impl MessageRouter for TestMessageRouter {
 struct TestOffersMessageHandler {}
 
 impl OffersMessageHandler for TestOffersMessageHandler {
-	fn handle_message(&self, _message: OffersMessage) -> Option<OffersMessage> {
-		None
-	}
+	fn handle_message<OMH: OnionMessageHandler>(&self, _message: &ReceivedOnionMessage<OMH, OffersMessage>) {}
 }
 
 #[derive(Debug)]
@@ -123,8 +121,8 @@ struct TestCustomMessageHandler {}
 
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
-	fn handle_custom_message<OMH: OnionMessageHandler, T: OnionMessageContents>(&self, received_onion_message: &ReceivedOnionMessage<OMH, T>) {
-		if let ReceivedOnionMessage::WithReplyPath(responder) = received_onion_message {
+	fn handle_custom_message<OMH: OnionMessageHandler, T: OnionMessageContents>(&self, message: &ReceivedOnionMessage<OMH, T>) {
+		if let ReceivedOnionMessage::WithReplyPath({_, responder}) = message {
 			responder.respond(TestCustomMessage {})
 		}
 	}

lightning/src/ln/channelmanager.rs

@@ -9275,12 +9275,12 @@ where
 	R::Target: Router,
 	L::Target: Logger,
 {
-	fn handle_message<OMH: OnionMessageHandler>(&self, received_onion_message: &ReceivedOnionMessage<OMH, OffersMessage>) {
+	fn handle_message<OMH: OnionMessageHandler>(&self, message: &ReceivedOnionMessage<OMH, OffersMessage>) {
 		let secp_ctx = &self.secp_ctx;
 		let expanded_key = &self.inbound_payment_key;
 
-		if let ReceivedOnionMessage::WithReplyPath(responder) = received_onion_message {
-			let response_option = match &responder.message {
+		if let ReceivedOnionMessage::WithReplyPath{message, responder} = message {
+			let response_option = match &message {
 				OffersMessage::InvoiceRequest(invoice_request) => {
 					let amount_msats = match InvoiceBuilder::<DerivedSigningPubkey>::amount_msats(
 						&invoice_request

lightning/src/ln/peer_handler.rs

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

lightning/src/onion_message/functional_tests.rs

@@ -70,7 +70,7 @@ impl MessageRouter for TestMessageRouter {
 struct TestOffersMessageHandler {}
 
 impl OffersMessageHandler for TestOffersMessageHandler {
-	fn handle_message<OMH: OnionMessageHandler>(&self, _received_onion_message: &ReceivedOnionMessage<OMH, OffersMessage>) {}
+	fn handle_message<OMH: OnionMessageHandler>(&self, _message: &ReceivedOnionMessage<OMH, OffersMessage>) {}
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -132,14 +132,14 @@ impl Drop for TestCustomMessageHandler {
 
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
-	fn handle_custom_message<OMH: OnionMessageHandler>(&self, received_onion_message: &ReceivedOnionMessage<OMH, Self::CustomMessage>) {
-		if let ReceivedOnionMessage::WithReplyPath(responder) = received_onion_message {
+	fn handle_custom_message<OMH: OnionMessageHandler>(&self, message: &ReceivedOnionMessage<OMH, Self::CustomMessage>) {
+		if let ReceivedOnionMessage::WithReplyPath{message, responder} = message {
 			match self.expected_messages.lock().unwrap().pop_front() {
-				Some(expected_msg) => assert_eq!(expected_msg, responder.message),
-				None => panic!("Unexpected message: {:?}", responder.message),
+				Some(expected_msg) => assert_eq!(expected_msg, *message),
+				None => panic!("Unexpected message: {:?}", message),
 			}
 
-			let response_option = match responder.message {
+			let response_option = match message {
 				TestCustomMessage::Request => Some(TestCustomMessage::Response),
 				TestCustomMessage::Response => None,
 			};

lightning/src/onion_message/messenger.rs

@@ -243,50 +243,62 @@ impl OnionMessageRecipient {
 	}
 }
 
+/// A struct handling response to an [`OnionMessage`]
 pub struct Responder<'a, OMH: OnionMessageHandler, T: OnionMessageContents> {
 	messenger: &'a OMH,
-	pub message: T,
 	pub reply_path: BlindedPath,
-	pub path_id: Option<[u8; 32]>
+	pub path_id: Option<[u8; 32]>,
+	// This phantom Data is used to ensure that we use T in the struct definition
+	// This allow us to ensure at compile time that the received message type and response type will be same
+	_phantom: std::marker::PhantomData<T>
 }
 
 impl<'a, OMH: OnionMessageHandler, T: OnionMessageContents> Responder<'a, OMH, T> {
     pub fn respond(&self, response: T) {
+		// Utilising the fact that we ensure at compile time that
+		// received message type, and response type will be same
+		let message_type = T::msg_type(&response);
         self.messenger.handle_onion_message_response(
 			response, self.reply_path.clone(), format_args!(
 				"when responding to {} onion message with path_id {:02x?}",
-				self.message.msg_type(),
+				message_type,
 				self.path_id.clone()
 			)
 		);
     }
 }
 
-pub struct MessageWithId<T: OnionMessageContents> {
-    pub message: T,
-    pub path_id: Option<[u8; 32]>,
-}
-
+/// A enum to handle received [`OnionMessage`]
 pub enum ReceivedOnionMessage<'a, OMH: OnionMessageHandler, T: OnionMessageContents> {
-	WithReplyPath(Responder<'a, OMH, T>),
-	WithoutReplyPath(MessageWithId<T>),
+	WithReplyPath {
+		message: T,
+		responder: Responder<'a, OMH, T>,
+	},
+	WithoutReplyPath {
+		message: T,
+		path_id: Option<[u8; 32]>,
+	},
 }
 
 impl<'a, OMH: OnionMessageHandler, T: OnionMessageContents> ReceivedOnionMessage<'a, OMH, T> {
 	fn new(messenger: &'a OMH, message: T, reply_path_option: Option<BlindedPath>, path_id: Option<[u8; 32]> ) -> Self {
 		match reply_path_option {
 			Some(reply_path) => {
-				ReceivedOnionMessage::WithReplyPath(Responder {
+				let responder = Responder {
 					messenger,
-					message,
 					reply_path,
-					path_id
-				})
+					path_id,
+					_phantom: std::marker::PhantomData
+				};
+				ReceivedOnionMessage::WithReplyPath {
+					message,
+					responder
+				}
 			}
-			None => ReceivedOnionMessage::WithoutReplyPath(MessageWithId {
+			None => ReceivedOnionMessage::WithoutReplyPath {
 				message,
 				path_id
-			})
+			}
 		}
     }
 }
@@ -548,7 +560,7 @@ pub trait CustomOnionMessageHandler {
 	/// Called with the custom message that was received, returning a response to send, if any.
 	///
 	/// The returned [`Self::CustomMessage`], if any, is enqueued to be sent by [`OnionMessenger`].
-	fn handle_custom_message<OMH: OnionMessageHandler>(&self, received_onion_message: &ReceivedOnionMessage<OMH, Self::CustomMessage>);
+	fn handle_custom_message<OMH: OnionMessageHandler>(&self, message: &ReceivedOnionMessage<OMH, Self::CustomMessage>);
 
 	/// Read a custom message of type `message_type` from `buffer`, returning `Ok(None)` if the
 	/// message type is unknown.
@@ -953,12 +965,12 @@ where
 
 				match message {
 					ParsedOnionMessageContents::Offers(msg) => {
-						let received_onion_message = ReceivedOnionMessage::new(self, msg, reply_path, path_id);
-						self.offers_handler.handle_message(&received_onion_message);
+						let message = ReceivedOnionMessage::new(self, msg, reply_path, path_id);
+						self.offers_handler.handle_message(&message);
 					},
 					ParsedOnionMessageContents::Custom(msg) => {
-						let received_onion_message = ReceivedOnionMessage::new(self, msg, reply_path, path_id);
-						self.custom_handler.handle_custom_message(&received_onion_message);
+						let message = ReceivedOnionMessage::new(self, msg, reply_path, path_id);
+						self.custom_handler.handle_custom_message(&message);
 					},
 				}
 			},

lightning/src/onion_message/offers.rs

@@ -42,7 +42,7 @@ pub trait OffersMessageHandler {
 	/// The returned [`OffersMessage`], if any, is enqueued to be sent by [`OnionMessenger`].
 	///
 	/// [`OnionMessenger`]: crate::onion_message::messenger::OnionMessenger
-	fn handle_message<OMH: OnionMessageHandler>(&self, received_onion_message: &ReceivedOnionMessage<OMH, OffersMessage>);
+	fn handle_message<OMH: OnionMessageHandler>(&self, message: &ReceivedOnionMessage<OMH, OffersMessage>);
 
 	/// Releases any [`OffersMessage`]s that need to be sent.
 	///