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

@@ -191,7 +191,8 @@ 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);
@@ -209,7 +210,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())),
@@ -507,7 +508,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();
@@ -546,7 +547,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!(
@@ -590,7 +591,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,19 @@ 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 +442,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 +471,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.
@@ -886,7 +888,8 @@ 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
@@ -1332,10 +1335,14 @@ 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!(
@@ -1349,7 +1356,7 @@ where
 			},
 			MessageSendInstructions::WithoutReplyPath { destination }
 			| MessageSendInstructions::ForReply {
-				instructions: ResponseInstruction { destination, context: None },
+				instructions: ResponseInstruction { destination, reply_data: (None, _) },
 			} => (destination, None),
 		};
 
@@ -1407,7 +1414,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
@@ -1424,7 +1431,7 @@ where
 				.collect::<Vec<_>>()
 		};
 
-		let recipient_tlvs = ReceiveTlvs { context: Some(context), custom_data: None };
+		let recipient_tlvs = ReceiveTlvs { context: Some(context), custom_data };
 
 		self.message_router
 			.create_blinded_paths(recipient, recipient_tlvs, peers, secp_ctx)
@@ -1932,8 +1939,12 @@ where
 								return;
 							},
 						};
-						let response_instructions =
-							self.custom_handler.handle_custom_message(msg, context, responder);
+						let response_instructions = 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);
 						}