From ca5b6b51d16662d968bb7b49a28ae5c65ed8493f Mon Sep 17 00:00:00 2001 From: shaavan Date: Mon, 15 Apr 2024 15:22:25 +0530 Subject: [PATCH 1/5] Convert handle_onion_message_response to a public function and add test coverage This commit modifies handle_onion_message_response to be accessible publicly as handle_onion_message, enabling users to respond asynchronously. Additionally, a new test is introduced to validate this functionality. --- .../src/onion_message/functional_tests.rs | 31 +++++++++++++++++++ lightning/src/onion_message/messenger.rs | 13 ++++++-- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs index 029038a9fe7..a777dc41981 100644 --- a/lightning/src/onion_message/functional_tests.rs +++ b/lightning/src/onion_message/functional_tests.rs @@ -372,6 +372,37 @@ fn three_blinded_hops() { pass_along_path(&nodes); } +#[test] +fn async_response_over_one_blinded_hop() { + // Simulate an asynchronous interaction between two nodes, Alice and Bob. + + // 1. Set up the network with two nodes: Alice and Bob. + let nodes = create_nodes(2); + let alice = &nodes[0]; + let bob = &nodes[1]; + + // 2. Define the message sent from Bob to Alice. + let message = TestCustomMessage::Request; + let path_id = Some([2; 32]); + + // 3. Simulate the creation of a Blinded Reply path provided by Bob. + let secp_ctx = Secp256k1::new(); + let reply_path = BlindedPath::new_for_message(&[], nodes[1].node_id, &*nodes[1].entropy_source, &secp_ctx).unwrap(); + + // 4. Create a responder using the reply path for Alice. + let responder = Some(Responder::new(reply_path, path_id)); + + // 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, responder); + + // 6. Simulate Alice asynchronously responding back to Bob with a response. + nodes[0].messenger.handle_onion_message_response(response_instruction); + bob.custom_message_handler.expect_message(TestCustomMessage::Response); + + pass_along_path(&nodes); +} + #[test] fn too_big_packet_error() { // Make sure we error as expected if a packet is too big to send. diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs index fc3eead1e4f..8ab92ba647d 100644 --- a/lightning/src/onion_message/messenger.rs +++ b/lightning/src/onion_message/messenger.rs @@ -264,7 +264,7 @@ pub struct Responder { impl Responder { /// Creates a new [`Responder`] instance with the provided reply path. - fn new(reply_path: BlindedPath, path_id: Option<[u8; 32]>) -> Self { + pub(super) fn new(reply_path: BlindedPath, path_id: Option<[u8; 32]>) -> Self { Responder { reply_path, path_id, @@ -1053,7 +1053,16 @@ where ) } - fn handle_onion_message_response( + /// Handles the response to an [`OnionMessage`] based on its [`ResponseInstruction`], + /// enqueueing any response for sending. + /// + /// This function is useful for asynchronous handling of [`OnionMessage`]s. + /// Handlers have the option to return [`ResponseInstruction::NoResponse`], indicating that + /// no immediate response should be sent. Then, they can transfer the associated [`Responder`] + /// to another task responsible for generating the response asynchronously. Subsequently, when + /// the response is prepared and ready for sending, that task can invoke this method to enqueue + /// the response for delivery. + pub fn handle_onion_message_response( &self, response: ResponseInstruction ) { if let ResponseInstruction::WithoutReplyPath(response) = response { From 6904786ed4edd1fe5d85a5b886aa4be86052f338 Mon Sep 17 00:00:00 2001 From: shaavan Date: Mon, 15 Apr 2024 16:25:03 +0530 Subject: [PATCH 2/5] Introduce ResponseInstruction::WithReplyPath variant. And expand handle_onion_message_response return Type 1. Introduce a new function in OnionMessenger to create blinded paths. 2. Use it in handle_onion_message_response to create a reply_path for the right variant and use it in onion_message. 3. Expand the return type of handle_onion_message_response to handle three cases: 1. Ok(None) in case of no response to be sent. 2. Ok(Some(SendSuccess) and Err(SendError) in case of successful and unsuccessful queueing up of response messages respectively. This allows the user to get access to the Success/Failure status of the sending of response and handle it accordingly. --- .../src/onion_message/functional_tests.rs | 8 +- lightning/src/onion_message/messenger.rs | 92 +++++++++++++++---- 2 files changed, 82 insertions(+), 18 deletions(-) diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs index a777dc41981..03356b434a3 100644 --- a/lightning/src/onion_message/functional_tests.rs +++ b/lightning/src/onion_message/functional_tests.rs @@ -19,7 +19,7 @@ use crate::routing::test_utils::{add_channel, add_or_update_node}; use crate::sign::{NodeSigner, Recipient}; use crate::util::ser::{FixedLengthReader, LengthReadable, Writeable, Writer}; use crate::util::test_utils; -use super::messenger::{CustomOnionMessageHandler, DefaultMessageRouter, Destination, OnionMessagePath, OnionMessenger, PendingOnionMessage, Responder, ResponseInstruction, SendError}; +use super::messenger::{CustomOnionMessageHandler, DefaultMessageRouter, Destination, OnionMessagePath, OnionMessenger, PendingOnionMessage, Responder, ResponseInstruction, SendError, SendSuccess}; use super::offers::{OffersMessage, OffersMessageHandler}; use super::packet::{OnionMessageContents, Packet}; @@ -397,7 +397,11 @@ fn async_response_over_one_blinded_hop() { let response_instruction = nodes[0].custom_message_handler.handle_custom_message(message, responder); // 6. Simulate Alice asynchronously responding back to Bob with a response. - nodes[0].messenger.handle_onion_message_response(response_instruction); + assert_eq!( + nodes[0].messenger.handle_onion_message_response(response_instruction), + Ok(Some(SendSuccess::Buffered)), + ); + bob.custom_message_handler.expect_message(TestCustomMessage::Response); pass_along_path(&nodes); diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs index 8ab92ba647d..aa5f6589492 100644 --- a/lightning/src/onion_message/messenger.rs +++ b/lightning/src/onion_message/messenger.rs @@ -271,7 +271,9 @@ impl Responder { } } - /// Creates the appropriate [`ResponseInstruction`] for a given response. + /// Creates a [`ResponseInstruction::WithoutReplyPath`] for a given response. + /// + /// Use when the recipient doesn't need to send back a reply to us. pub fn respond(self, response: T) -> ResponseInstruction { ResponseInstruction::WithoutReplyPath(OnionMessageResponse { message: response, @@ -279,6 +281,17 @@ impl Responder { path_id: self.path_id, }) } + + /// Creates a [`ResponseInstruction::WithReplyPath`] for a given response. + /// + /// Use when the recipient needs to send back a reply to us. + pub fn respond_with_reply_path(self, response: T) -> ResponseInstruction { + ResponseInstruction::WithReplyPath(OnionMessageResponse { + message: response, + reply_path: self.reply_path, + path_id: self.path_id, + }) + } } /// This struct contains the information needed to reply to a received message. @@ -290,6 +303,9 @@ pub struct OnionMessageResponse { /// `ResponseInstruction` represents instructions for responding to received messages. pub enum ResponseInstruction { + /// Indicates that a response should be sent including a reply path for + /// the recipient to respond back. + WithReplyPath(OnionMessageResponse), /// Indicates that a response should be sent without including a reply path /// for the recipient to respond back. WithoutReplyPath(OnionMessageResponse), @@ -564,7 +580,11 @@ pub enum SendError { TooFewBlindedHops, /// The first hop is not a peer and doesn't have a known [`SocketAddress`]. InvalidFirstHop(PublicKey), - /// A path from the sender to the destination could not be found by the [`MessageRouter`]. + /// Indicates that a path could not be found by the [`MessageRouter`]. + /// + /// This occurs when either: + /// - No path from the sender to the destination was found to send the onion message + /// - No reply path to the sender could be created when responding to an onion message PathNotFound, /// Onion message contents must have a TLV type >= 64. InvalidMessage, @@ -981,6 +1001,27 @@ where .map_err(|_| SendError::PathNotFound) } + fn create_blinded_path(&self) -> Result { + let recipient = self.node_signer + .get_node_id(Recipient::Node) + .map_err(|_| SendError::GetNodeIdFailed)?; + let secp_ctx = &self.secp_ctx; + + let peers = self.message_recipients.lock().unwrap() + .iter() + .filter(|(_, peer)| matches!(peer, OnionMessageRecipient::ConnectedPeer(_))) + .map(|(node_id, _ )| ForwardNode { + node_id: *node_id, + short_channel_id: None, + }) + .collect::>(); + + self.message_router + .create_blinded_paths(recipient, peers, secp_ctx) + .and_then(|paths| paths.into_iter().next().ok_or(())) + .map_err(|_| SendError::PathNotFound) + } + fn enqueue_onion_message( &self, path: OnionMessagePath, contents: T, reply_path: Option, log_suffix: fmt::Arguments @@ -1064,18 +1105,37 @@ where /// the response for delivery. pub fn handle_onion_message_response( &self, response: ResponseInstruction - ) { - if let ResponseInstruction::WithoutReplyPath(response) = response { - let message_type = response.message.msg_type(); - let _ = self.find_path_and_enqueue_onion_message( - response.message, Destination::BlindedPath(response.reply_path), None, - format_args!( - "when responding with {} to an onion message with path_id {:02x?}", - message_type, - response.path_id - ) - ); - } + ) -> Result, SendError> { + let (response, create_reply_path) = match response { + ResponseInstruction::WithReplyPath(response) => (response, true), + ResponseInstruction::WithoutReplyPath(response) => (response, false), + ResponseInstruction::NoResponse => return Ok(None), + }; + + let message_type = response.message.msg_type(); + let reply_path = if create_reply_path { + match self.create_blinded_path() { + Ok(reply_path) => Some(reply_path), + Err(err) => { + log_trace!( + self.logger, + "Failed to create reply path when responding with {} to an onion message \ + with path_id {:02x?}: {:?}", + message_type, response.path_id, err + ); + return Err(err); + } + } + } else { None }; + + self.find_path_and_enqueue_onion_message( + response.message, Destination::BlindedPath(response.reply_path), reply_path, + format_args!( + "when responding with {} to an onion message with path_id {:02x?}", + message_type, + response.path_id + ) + ).map(|result| Some(result)) } #[cfg(test)] @@ -1181,14 +1241,14 @@ where |reply_path| Responder::new(reply_path, path_id) ); let response_instructions = self.offers_handler.handle_message(msg, responder); - self.handle_onion_message_response(response_instructions); + let _ = self.handle_onion_message_response(response_instructions); }, ParsedOnionMessageContents::Custom(msg) => { let responder = reply_path.map( |reply_path| Responder::new(reply_path, path_id) ); let response_instructions = self.custom_handler.handle_custom_message(msg, responder); - self.handle_onion_message_response(response_instructions); + let _ = self.handle_onion_message_response(response_instructions); }, } }, From 7d5dd6b2994752d9f2ae57538a0b25af475ee90d Mon Sep 17 00:00:00 2001 From: shaavan Date: Thu, 23 May 2024 11:33:34 +0530 Subject: [PATCH 3/5] Refactor: Rename Request & Response to Ping & Pong 1. These two variants will be modified in an upcoming commit to be each other's response. 2. The names are updated to better fit their new roles. --- .../src/onion_message/functional_tests.rs | 94 +++++++++---------- 1 file changed, 47 insertions(+), 47 deletions(-) diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs index 03356b434a3..fe2b3799e72 100644 --- a/lightning/src/onion_message/functional_tests.rs +++ b/lightning/src/onion_message/functional_tests.rs @@ -81,20 +81,20 @@ impl OffersMessageHandler for TestOffersMessageHandler { #[derive(Clone, Debug, PartialEq)] enum TestCustomMessage { - Request, - Response, + Ping, + Pong, } -const CUSTOM_REQUEST_MESSAGE_TYPE: u64 = 4242; -const CUSTOM_RESPONSE_MESSAGE_TYPE: u64 = 4343; -const CUSTOM_REQUEST_MESSAGE_CONTENTS: [u8; 32] = [42; 32]; -const CUSTOM_RESPONSE_MESSAGE_CONTENTS: [u8; 32] = [43; 32]; +const CUSTOM_PING_MESSAGE_TYPE: u64 = 4242; +const CUSTOM_PONG_MESSAGE_TYPE: u64 = 4343; +const CUSTOM_PING_MESSAGE_CONTENTS: [u8; 32] = [42; 32]; +const CUSTOM_PONG_MESSAGE_CONTENTS: [u8; 32] = [43; 32]; impl OnionMessageContents for TestCustomMessage { fn tlv_type(&self) -> u64 { match self { - TestCustomMessage::Request => CUSTOM_REQUEST_MESSAGE_TYPE, - TestCustomMessage::Response => CUSTOM_RESPONSE_MESSAGE_TYPE, + TestCustomMessage::Ping => CUSTOM_PING_MESSAGE_TYPE, + TestCustomMessage::Pong => CUSTOM_PONG_MESSAGE_TYPE, } } fn msg_type(&self) -> &'static str { @@ -105,8 +105,8 @@ impl OnionMessageContents for TestCustomMessage { impl Writeable for TestCustomMessage { fn write(&self, w: &mut W) -> Result<(), io::Error> { match self { - TestCustomMessage::Request => Ok(CUSTOM_REQUEST_MESSAGE_CONTENTS.write(w)?), - TestCustomMessage::Response => Ok(CUSTOM_RESPONSE_MESSAGE_CONTENTS.write(w)?), + TestCustomMessage::Ping => Ok(CUSTOM_PING_MESSAGE_CONTENTS.write(w)?), + TestCustomMessage::Pong => Ok(CUSTOM_PONG_MESSAGE_CONTENTS.write(w)?), } } } @@ -144,8 +144,8 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler { None => panic!("Unexpected message: {:?}", msg), } let response_option = match msg { - TestCustomMessage::Request => Some(TestCustomMessage::Response), - TestCustomMessage::Response => None, + TestCustomMessage::Ping => Some(TestCustomMessage::Pong), + TestCustomMessage::Pong => None, }; if let (Some(response), Some(responder)) = (response_option, responder) { responder.respond(response) @@ -155,15 +155,15 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler { } fn read_custom_message(&self, message_type: u64, buffer: &mut R) -> Result, DecodeError> where Self: Sized { match message_type { - CUSTOM_REQUEST_MESSAGE_TYPE => { + CUSTOM_PING_MESSAGE_TYPE => { let buf = read_to_end(buffer)?; - assert_eq!(buf, CUSTOM_REQUEST_MESSAGE_CONTENTS); - Ok(Some(TestCustomMessage::Request)) + assert_eq!(buf, CUSTOM_PING_MESSAGE_CONTENTS); + Ok(Some(TestCustomMessage::Ping)) }, - CUSTOM_RESPONSE_MESSAGE_TYPE => { + CUSTOM_PONG_MESSAGE_TYPE => { let buf = read_to_end(buffer)?; - assert_eq!(buf, CUSTOM_RESPONSE_MESSAGE_CONTENTS); - Ok(Some(TestCustomMessage::Response)) + assert_eq!(buf, CUSTOM_PONG_MESSAGE_CONTENTS); + Ok(Some(TestCustomMessage::Pong)) }, _ => Ok(None), } @@ -298,18 +298,18 @@ fn pass_along_path(path: &Vec) { #[test] fn one_unblinded_hop() { let nodes = create_nodes(2); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let destination = Destination::Node(nodes[1].node_id); nodes[0].messenger.send_onion_message(test_msg, destination, None).unwrap(); - nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[1].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } #[test] fn two_unblinded_hops() { let nodes = create_nodes(3); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let path = OnionMessagePath { intermediate_nodes: vec![nodes[1].node_id], @@ -318,27 +318,27 @@ fn two_unblinded_hops() { }; nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap(); - nodes[2].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[2].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } #[test] fn one_blinded_hop() { let nodes = create_nodes(2); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let secp_ctx = Secp256k1::new(); let blinded_path = BlindedPath::new_for_message(&[], nodes[1].node_id, &*nodes[1].entropy_source, &secp_ctx).unwrap(); let destination = Destination::BlindedPath(blinded_path); nodes[0].messenger.send_onion_message(test_msg, destination, None).unwrap(); - nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[1].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } #[test] fn two_unblinded_two_blinded() { let nodes = create_nodes(5); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let secp_ctx = Secp256k1::new(); let intermediate_nodes = [ForwardNode { node_id: nodes[3].node_id, short_channel_id: None }]; @@ -350,14 +350,14 @@ fn two_unblinded_two_blinded() { }; nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap(); - nodes[4].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[4].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } #[test] fn three_blinded_hops() { let nodes = create_nodes(4); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let secp_ctx = Secp256k1::new(); let intermediate_nodes = [ @@ -368,7 +368,7 @@ fn three_blinded_hops() { let destination = Destination::BlindedPath(blinded_path); nodes[0].messenger.send_onion_message(test_msg, destination, None).unwrap(); - nodes[3].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[3].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } @@ -382,7 +382,7 @@ fn async_response_over_one_blinded_hop() { let bob = &nodes[1]; // 2. Define the message sent from Bob to Alice. - let message = TestCustomMessage::Request; + let message = TestCustomMessage::Ping; let path_id = Some([2; 32]); // 3. Simulate the creation of a Blinded Reply path provided by Bob. @@ -402,7 +402,7 @@ fn async_response_over_one_blinded_hop() { Ok(Some(SendSuccess::Buffered)), ); - bob.custom_message_handler.expect_message(TestCustomMessage::Response); + bob.custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } @@ -411,7 +411,7 @@ fn async_response_over_one_blinded_hop() { fn too_big_packet_error() { // Make sure we error as expected if a packet is too big to send. let nodes = create_nodes(2); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let hop_node_id = nodes[1].node_id; let hops = vec![hop_node_id; 400]; @@ -429,7 +429,7 @@ fn we_are_intro_node() { // If we are sending straight to a blinded path and we are the introduction node, we need to // advance the blinded path by 1 hop so the second hop is the new introduction node. let mut nodes = create_nodes(3); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let secp_ctx = Secp256k1::new(); let intermediate_nodes = [ @@ -440,7 +440,7 @@ fn we_are_intro_node() { let destination = Destination::BlindedPath(blinded_path); nodes[0].messenger.send_onion_message(test_msg.clone(), destination, None).unwrap(); - nodes[2].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[2].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); // Try with a two-hop blinded path where we are the introduction node. @@ -448,7 +448,7 @@ fn we_are_intro_node() { let blinded_path = BlindedPath::new_for_message(&intermediate_nodes, nodes[1].node_id, &*nodes[1].entropy_source, &secp_ctx).unwrap(); let destination = Destination::BlindedPath(blinded_path); nodes[0].messenger.send_onion_message(test_msg, destination, None).unwrap(); - nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[1].custom_message_handler.expect_message(TestCustomMessage::Pong); nodes.remove(2); pass_along_path(&nodes); } @@ -457,7 +457,7 @@ fn we_are_intro_node() { fn invalid_blinded_path_error() { // Make sure we error as expected if a provided blinded path has 0 hops. let nodes = create_nodes(3); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let secp_ctx = Secp256k1::new(); let intermediate_nodes = [ForwardNode { node_id: nodes[1].node_id, short_channel_id: None }]; @@ -471,7 +471,7 @@ fn invalid_blinded_path_error() { #[test] fn reply_path() { let mut nodes = create_nodes(4); - let test_msg = TestCustomMessage::Request; + let test_msg = TestCustomMessage::Ping; let secp_ctx = Secp256k1::new(); // Destination::Node @@ -486,10 +486,10 @@ fn reply_path() { ]; let reply_path = BlindedPath::new_for_message(&intermediate_nodes, nodes[0].node_id, &*nodes[0].entropy_source, &secp_ctx).unwrap(); nodes[0].messenger.send_onion_message_using_path(path, test_msg.clone(), Some(reply_path)).unwrap(); - nodes[3].custom_message_handler.expect_message(TestCustomMessage::Request); + nodes[3].custom_message_handler.expect_message(TestCustomMessage::Ping); pass_along_path(&nodes); // Make sure the last node successfully decoded the reply path. - nodes[0].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[0].custom_message_handler.expect_message(TestCustomMessage::Pong); nodes.reverse(); pass_along_path(&nodes); @@ -507,11 +507,11 @@ fn reply_path() { let reply_path = BlindedPath::new_for_message(&intermediate_nodes, nodes[0].node_id, &*nodes[0].entropy_source, &secp_ctx).unwrap(); nodes[0].messenger.send_onion_message(test_msg, destination, Some(reply_path)).unwrap(); - nodes[3].custom_message_handler.expect_message(TestCustomMessage::Request); + nodes[3].custom_message_handler.expect_message(TestCustomMessage::Ping); pass_along_path(&nodes); // Make sure the last node successfully decoded the reply path. - nodes[0].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[0].custom_message_handler.expect_message(TestCustomMessage::Pong); nodes.reverse(); pass_along_path(&nodes); } @@ -545,7 +545,7 @@ fn invalid_custom_message_type() { #[test] fn peer_buffer_full() { let nodes = create_nodes(2); - let test_msg = TestCustomMessage::Request; + let test_msg = TestCustomMessage::Ping; let destination = Destination::Node(nodes[1].node_id); for _ in 0..188 { // Based on MAX_PER_PEER_BUFFER_SIZE in OnionMessenger nodes[0].messenger.send_onion_message(test_msg.clone(), destination.clone(), None).unwrap(); @@ -560,7 +560,7 @@ fn many_hops() { // of size [`crate::onion_message::packet::BIG_PACKET_HOP_DATA_LEN`]. let num_nodes: usize = 25; let nodes = create_nodes(num_nodes as u8); - let test_msg = TestCustomMessage::Response; + let test_msg = TestCustomMessage::Pong; let mut intermediate_nodes = vec![]; for i in 1..(num_nodes-1) { @@ -573,14 +573,14 @@ fn many_hops() { first_node_addresses: None, }; nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap(); - nodes[num_nodes-1].custom_message_handler.expect_message(TestCustomMessage::Response); + nodes[num_nodes-1].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&nodes); } #[test] fn requests_peer_connection_for_buffered_messages() { let nodes = create_nodes(3); - let message = TestCustomMessage::Request; + let message = TestCustomMessage::Ping; let secp_ctx = Secp256k1::new(); add_channel_to_graph(&nodes[0], &nodes[1], &secp_ctx, 42); @@ -618,7 +618,7 @@ fn requests_peer_connection_for_buffered_messages() { #[test] fn drops_buffered_messages_waiting_for_peer_connection() { let nodes = create_nodes(3); - let message = TestCustomMessage::Request; + let message = TestCustomMessage::Ping; let secp_ctx = Secp256k1::new(); add_channel_to_graph(&nodes[0], &nodes[1], &secp_ctx, 42); @@ -670,7 +670,7 @@ fn intercept_offline_peer_oms() { } } - let message = TestCustomMessage::Response; + let message = TestCustomMessage::Pong; let secp_ctx = Secp256k1::new(); let intermediate_nodes = [ForwardNode { node_id: nodes[1].node_id, short_channel_id: None }]; let blinded_path = BlindedPath::new_for_message( @@ -710,7 +710,7 @@ fn intercept_offline_peer_oms() { } nodes[1].messenger.forward_onion_message(onion_message, &final_node_vec[0].node_id).unwrap(); - final_node_vec[0].custom_message_handler.expect_message(TestCustomMessage::Response); + final_node_vec[0].custom_message_handler.expect_message(TestCustomMessage::Pong); pass_along_path(&vec![nodes.remove(1), final_node_vec.remove(0)]); } From b95cfed7bb230d6eea26a404f58070f3affd911f Mon Sep 17 00:00:00 2001 From: shaavan Date: Thu, 23 May 2024 11:40:12 +0530 Subject: [PATCH 4/5] Refactor TestCustomMessageHandler - Introduce a new struct for keeping expectations organized. - Add a boolean field to track whether a response is expected, and hence whether a `reply_path` should be included with the response. - Update Ping and Pong roles for bidirectional communication. - Introduce panic for when there is no responder and we were expecting to include a `reply_path`. - Refactor `handle_custom_message` code. --- .../src/onion_message/functional_tests.rs | 60 ++++++++++++++----- 1 file changed, 45 insertions(+), 15 deletions(-) diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs index fe2b3799e72..0c2342ca154 100644 --- a/lightning/src/onion_message/functional_tests.rs +++ b/lightning/src/onion_message/functional_tests.rs @@ -112,16 +112,39 @@ impl Writeable for TestCustomMessage { } struct TestCustomMessageHandler { - expected_messages: Mutex>, + expectations: Mutex>, +} + +struct OnHandleCustomMessage { + expect: TestCustomMessage, + include_reply_path: bool, } impl TestCustomMessageHandler { fn new() -> Self { - Self { expected_messages: Mutex::new(VecDeque::new()) } + Self { expectations: Mutex::new(VecDeque::new()) } } fn expect_message(&self, message: TestCustomMessage) { - self.expected_messages.lock().unwrap().push_back(message); + self.expectations.lock().unwrap().push_back( + OnHandleCustomMessage { + expect: message, + include_reply_path: false, + } + ); + } + + fn expect_message_and_response(&self, message: TestCustomMessage) { + self.expectations.lock().unwrap().push_back( + OnHandleCustomMessage { + expect: message, + include_reply_path: true, + } + ); + } + + fn get_next_expectation(&self) -> OnHandleCustomMessage { + self.expectations.lock().unwrap().pop_front().expect("No expectations remaining") } } @@ -132,25 +155,32 @@ impl Drop for TestCustomMessageHandler { return; } } - assert!(self.expected_messages.lock().unwrap().is_empty()); + assert!(self.expectations.lock().unwrap().is_empty()); } } impl CustomOnionMessageHandler for TestCustomMessageHandler { type CustomMessage = TestCustomMessage; fn handle_custom_message(&self, msg: Self::CustomMessage, responder: Option) -> ResponseInstruction { - match self.expected_messages.lock().unwrap().pop_front() { - Some(expected_msg) => assert_eq!(expected_msg, msg), - None => panic!("Unexpected message: {:?}", msg), - } - let response_option = match msg { - TestCustomMessage::Ping => Some(TestCustomMessage::Pong), - TestCustomMessage::Pong => None, + let expectation = self.get_next_expectation(); + assert_eq!(msg, expectation.expect); + + let response = match msg { + TestCustomMessage::Ping => TestCustomMessage::Pong, + TestCustomMessage::Pong => TestCustomMessage::Ping, }; - if let (Some(response), Some(responder)) = (response_option, responder) { - responder.respond(response) - } else { - ResponseInstruction::NoResponse + + // Sanity check: expecting to include reply path when responder is absent should panic. + if expectation.include_reply_path && responder.is_none() { + panic!("Expected to include a reply_path, but the responder was absent.") + } + + match responder { + Some(responder) if expectation.include_reply_path => { + responder.respond_with_reply_path(response) + }, + Some(responder) => responder.respond(response), + None => ResponseInstruction::NoResponse, } } fn read_custom_message(&self, message_type: u64, buffer: &mut R) -> Result, DecodeError> where Self: Sized { From d8adbb0a5b7dcc39a41596274129ff3d69ec08de Mon Sep 17 00:00:00 2001 From: shaavan Date: Fri, 24 May 2024 14:30:34 +0530 Subject: [PATCH 5/5] Introduce new tests for `ResponseInstruction::WithReplyPath` - The new tests check both successful and unsuccessful `reply_path` creation conditions. --- .../src/onion_message/functional_tests.rs | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs index 0c2342ca154..0c3c67d4483 100644 --- a/lightning/src/onion_message/functional_tests.rs +++ b/lightning/src/onion_message/functional_tests.rs @@ -437,6 +437,73 @@ fn async_response_over_one_blinded_hop() { pass_along_path(&nodes); } +#[test] +fn async_response_with_reply_path_succeeds() { + // Simulate an asynchronous interaction between two nodes, Alice and Bob. + // Create a channel between the two nodes to establish them as announced nodes, + // which allows the creation of the reply_path for successful communication. + + let mut nodes = create_nodes(2); + let alice = &nodes[0]; + let bob = &nodes[1]; + let secp_ctx = Secp256k1::new(); + + add_channel_to_graph(alice, bob, &secp_ctx, 24); + + // Alice receives a message from Bob with an added reply_path for responding back. + let message = TestCustomMessage::Ping; + let path_id = Some([2; 32]); + let reply_path = BlindedPath::new_for_message(&[], bob.node_id, &*bob.entropy_source, &secp_ctx).unwrap(); + + // Alice asynchronously responds to Bob, expecting a response back from him. + let responder = Responder::new(reply_path, path_id); + alice.custom_message_handler.expect_message_and_response(message.clone()); + let response_instruction = alice.custom_message_handler.handle_custom_message(message, Some(responder)); + + assert_eq!( + alice.messenger.handle_onion_message_response(response_instruction), + Ok(Some(SendSuccess::Buffered)), + ); + + // Set Bob's expectation and pass the Onion Message along the path. + bob.custom_message_handler.expect_message(TestCustomMessage::Pong); + pass_along_path(&nodes); + + // Bob responds back to Alice using the reply_path she included with the OnionMessage. + // Set Alice's expectation and reverse the path for the response. + alice.custom_message_handler.expect_message(TestCustomMessage::Ping); + nodes.reverse(); + pass_along_path(&nodes); +} + +#[test] +fn async_response_with_reply_path_fails() { + // Simulate an asynchronous interaction between two unannounced nodes, Alice and Bob. + // Since the nodes are unannounced, attempting to respond using a reply_path + // will fail, leading to an expected failure in communication. + + let nodes = create_nodes(2); + let alice = &nodes[0]; + let bob = &nodes[1]; + let secp_ctx = Secp256k1::new(); + + // Alice receives a message from Bob with an added reply_path for responding back. + let message = TestCustomMessage::Ping; + let path_id = Some([2; 32]); + let reply_path = BlindedPath::new_for_message(&[], bob.node_id, &*bob.entropy_source, &secp_ctx).unwrap(); + + // Alice tries to asynchronously respond to Bob, but fails because the nodes are unannounced. + // Therefore, the reply_path cannot be used for the response. + let responder = Responder::new(reply_path, path_id); + alice.custom_message_handler.expect_message_and_response(message.clone()); + let response_instruction = alice.custom_message_handler.handle_custom_message(message, Some(responder)); + + assert_eq!( + alice.messenger.handle_onion_message_response(response_instruction), + Err(SendError::PathNotFound), + ); +} + #[test] fn too_big_packet_error() { // Make sure we error as expected if a packet is too big to send.