Introduce custom data in handle_custom_message

Author: shaavan

fuzz/src/onion_message.rs

@@ -161,7 +161,7 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
 	fn handle_custom_message(
 		&self, message: Self::CustomMessage, _context: Option<Vec<u8>>,
-		responder: Option<Responder>,
+		_custom_data: Option<Vec<u8>>, responder: Option<Responder>,
 	) -> Option<(Self::CustomMessage, ResponseInstruction)> {
 		match responder {
 			Some(responder) => Some((message, responder.respond())),

lightning/src/ln/channelmanager.rs

@@ -12296,7 +12296,7 @@ where
 						let nonce = Nonce::from_entropy_source(&*self.entropy_source);
 						let hmac = payment_hash.hmac_for_offer_payment(nonce, expanded_key);
 						let context = MessageContext::Offers(OffersContext::InboundPayment { payment_hash, nonce, hmac });
-						Some((OffersMessage::Invoice(invoice), responder.respond_with_reply_path(context)))
+						Some((OffersMessage::Invoice(invoice), responder.respond_with_reply_path(context, None)))
 					},
 					Err(error) => Some((OffersMessage::InvoiceError(error.into()), responder.respond())),
 				}

lightning/src/ln/peer_handler.rs

@@ -169,7 +169,7 @@ impl DNSResolverMessageHandler for IgnoringMessageHandler {
 }
 impl CustomOnionMessageHandler for IgnoringMessageHandler {
 	type CustomMessage = Infallible;
-	fn handle_custom_message(&self, _message: Infallible, _context: Option<Vec<u8>>, _responder: Option<Responder>) -> Option<(Infallible, ResponseInstruction)> {
+	fn handle_custom_message(&self, _message: Infallible, _context: Option<Vec<u8>>, _custom_data: Option<Vec<u8>>, _responder: Option<Responder>) -> Option<(Infallible, ResponseInstruction)> {
 		// Since we always return `None` in the read the handle method should never be called.
 		unreachable!();
 	}

lightning/src/onion_message/functional_tests.rs

@@ -188,7 +188,7 @@ impl Drop for TestCustomMessageHandler {
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
 	fn handle_custom_message(
-		&self, msg: Self::CustomMessage, context: Option<Vec<u8>>, responder: Option<Responder>,
+		&self, msg: Self::CustomMessage, context: Option<Vec<u8>>, custom_data: Option<Vec<u8>>, responder: Option<Responder>,
 	) -> Option<(Self::CustomMessage, ResponseInstruction)> {
 		let expectation = self.get_next_expectation();
 		assert_eq!(msg, expectation.expect);
@@ -206,7 +206,7 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
 		match responder {
 			Some(responder) if expectation.include_reply_path => {
 				let context = MessageContext::Custom(context.unwrap_or_else(Vec::new));
-				let reply = responder.respond_with_reply_path(context);
+				let reply = responder.respond_with_reply_path(context, custom_data);
 				Some((response, reply))
 			},
 			Some(responder) => Some((response, responder.respond())),
@@ -502,7 +502,7 @@ fn async_response_over_one_blinded_hop() {
 	// 5. Expect Alice to receive the message and create a response instruction for it.
 	alice.custom_message_handler.expect_message(message.clone());
 	let response_instruction =
-		nodes[0].custom_message_handler.handle_custom_message(message, None, responder);
+		nodes[0].custom_message_handler.handle_custom_message(message, None, None, responder);
 
 	// 6. Simulate Alice asynchronously responding back to Bob with a response.
 	let (msg, instructions) = response_instruction.unwrap();
@@ -543,7 +543,7 @@ fn async_response_with_reply_path_succeeds() {
 	let responder = Responder::new(reply_path);
 	alice.custom_message_handler.expect_message_and_response(message.clone());
 	let response_instruction =
-		alice.custom_message_handler.handle_custom_message(message, None, Some(responder));
+		alice.custom_message_handler.handle_custom_message(message, None, None, Some(responder));
 
 	let (msg, instructions) = response_instruction.unwrap();
 	assert_eq!(
@@ -589,7 +589,7 @@ fn async_response_with_reply_path_fails() {
 	let responder = Responder::new(reply_path);
 	alice.custom_message_handler.expect_message_and_response(message.clone());
 	let response_instruction =
-		alice.custom_message_handler.handle_custom_message(message, None, Some(responder));
+		alice.custom_message_handler.handle_custom_message(message, None, None, Some(responder));
 
 	let (msg, instructions) = response_instruction.unwrap();
 	assert_eq!(

lightning/src/onion_message/messenger.rs

@@ -418,17 +418,17 @@ impl Responder {
 	pub fn respond(self) -> ResponseInstruction {
 		ResponseInstruction {
 			destination: Destination::BlindedPath(self.reply_path),
-			context: None,
+			reply_data: (None, None)
 		}
 	}
 
 	/// Creates a [`ResponseInstruction`] for responding including a reply path.
 	///
 	/// Use when the recipient needs to send back a reply to us.
-	pub fn respond_with_reply_path(self, context: MessageContext) -> ResponseInstruction {
+	pub fn respond_with_reply_path(self, context: MessageContext, custom_data: Option<Vec<u8>>) -> ResponseInstruction {
 		ResponseInstruction {
 			destination: Destination::BlindedPath(self.reply_path),
-			context: Some(context),
+			reply_data: (Some(context), custom_data),
 		}
 	}
 }
@@ -440,7 +440,7 @@ pub struct ResponseInstruction {
 	/// [`Destination`] rather than an explicit [`BlindedMessagePath`] simplifies the logic in
 	/// [`OnionMessenger::send_onion_message_internal`] somewhat.
 	destination: Destination,
-	context: Option<MessageContext>,
+	reply_data: (Option<MessageContext>, Option<Vec<u8>>)
 }
 
 impl ResponseInstruction {
@@ -469,7 +469,7 @@ pub enum MessageSendInstructions {
 		destination: Destination,
 		/// The context to include in the reply path we'll give the recipient so they can respond
 		/// to us.
-		context: MessageContext,
+		reply_data: (MessageContext, Option<Vec<u8>>),
 	},
 	/// Indicates that a message should be sent without including a reply path, preventing the
 	/// recipient from responding.
@@ -874,7 +874,7 @@ pub trait CustomOnionMessageHandler {
 	///
 	/// The returned [`Self::CustomMessage`], if any, is enqueued to be sent by [`OnionMessenger`].
 	fn handle_custom_message(
-		&self, message: Self::CustomMessage, context: Option<Vec<u8>>, responder: Option<Responder>,
+		&self, message: Self::CustomMessage, context: Option<Vec<u8>>, custom_data: Option<Vec<u8>>, responder: Option<Responder>,
 	) -> Option<(Self::CustomMessage, ResponseInstruction)>;
 
 	/// Read a custom message of type `message_type` from `buffer`, returning `Ok(None)` if the
@@ -1315,10 +1315,10 @@ where
 			MessageSendInstructions::WithSpecifiedReplyPath { destination, reply_path } => {
 				(destination, Some(reply_path))
 			},
-			MessageSendInstructions::WithReplyPath { destination, context }
+			MessageSendInstructions::WithReplyPath { destination, reply_data: (context, custom_data) }
 			| MessageSendInstructions::ForReply {
-				instructions: ResponseInstruction { destination, context: Some(context) },
-			} => match self.create_blinded_path(context) {
+				instructions: ResponseInstruction { destination, reply_data: (Some(context), custom_data) },
+			} => match self.create_blinded_path(context, custom_data) {
 				Ok(reply_path) => (destination, Some(reply_path)),
 				Err(err) => {
 					log_trace!(
@@ -1332,7 +1332,7 @@ where
 			},
 			MessageSendInstructions::WithoutReplyPath { destination }
 			| MessageSendInstructions::ForReply {
-				instructions: ResponseInstruction { destination, context: None },
+				instructions: ResponseInstruction { destination, reply_data: (None, _) },
 			} => (destination, None),
 		};
 
@@ -1390,7 +1390,7 @@ where
 	}
 
 	fn create_blinded_path(
-		&self, context: MessageContext,
+		&self, context: MessageContext, custom_data: Option<Vec<u8>>,
 	) -> Result<BlindedMessagePath, SendError> {
 		let recipient = self
 			.node_signer
@@ -1409,7 +1409,7 @@ where
 
 		let recipient_tlvs = ReceiveTlvs {
 			context: Some(context),
-			custom_data: None,
+			custom_data,
 		};
 
 		self.message_router
@@ -1915,7 +1915,7 @@ where
 							},
 						};
 						let response_instructions =
-							self.custom_handler.handle_custom_message(msg, context, responder);
+							self.custom_handler.handle_custom_message(msg, context, custom_data, responder);
 						if let Some((msg, instructions)) = response_instructions {
 							let _ = self.handle_onion_message_response(msg, instructions);
 						}