Skip to content

Commit 27899b1

Browse files
committed
Introduce new parameter in handle_custom_message.
- Add an option to direct handle_custom_message to call respond_with_reply_path. This expands the test suite to test the respond_with_reply_path capabilities. - Introduce a new test for ResponseInstruction::WithReplyPath, testing both the success and failure of reply_path creation.
1 parent 903cc60 commit 27899b1

File tree

3 files changed

+76
-9
lines changed

3 files changed

+76
-9
lines changed

lightning/src/ln/peer_handler.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ impl OffersMessageHandler for IgnoringMessageHandler {
143143
}
144144
impl CustomOnionMessageHandler for IgnoringMessageHandler {
145145
type CustomMessage = Infallible;
146-
fn handle_custom_message(&self, _message: Self::CustomMessage, _responder: Option<Responder>) -> ResponseInstruction<Self::CustomMessage> {
146+
fn handle_custom_message(&self, _message: Self::CustomMessage, _responder: Option<Responder>, _add_reply_path: bool) -> ResponseInstruction<Self::CustomMessage> {
147147
// Since we always return `None` in the read the handle method should never be called.
148148
unreachable!();
149149
}

lightning/src/onion_message/functional_tests.rs

Lines changed: 73 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ impl Drop for TestCustomMessageHandler {
149149

150150
impl CustomOnionMessageHandler for TestCustomMessageHandler {
151151
type CustomMessage = TestCustomMessage;
152-
fn handle_custom_message(&self, msg: Self::CustomMessage, responder: Option<Responder>) -> ResponseInstruction<Self::CustomMessage> {
152+
fn handle_custom_message(&self, msg: Self::CustomMessage, responder: Option<Responder>, add_reply_path: bool) -> ResponseInstruction<Self::CustomMessage> {
153153
match self.expected_messages.lock().unwrap().pop_front() {
154154
Some(expected_msg) => assert_eq!(expected_msg, msg),
155155
None => panic!("Unexpected message: {:?}", msg),
@@ -160,12 +160,18 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
160160
TestCustomMessage::ResponseA => Some(TestCustomMessage::ResponseB),
161161
TestCustomMessage::ResponseB => Some(TestCustomMessage::ResponseA),
162162
};
163-
if let (Some(response), Some(responder)) = (response_option, responder) {
164-
responder.respond(response)
165-
} else {
166-
ResponseInstruction::NoResponse
163+
match (response_option, responder) {
164+
(Some(response), Some(responder)) => {
165+
if add_reply_path {
166+
responder.respond_with_reply_path(response)
167+
} else {
168+
responder.respond(response)
169+
}
170+
}
171+
_ => ResponseInstruction::NoResponse,
167172
}
168173
}
174+
169175
fn read_custom_message<R: io::Read>(&self, message_type: u64, buffer: &mut R) -> Result<Option<Self::CustomMessage>, DecodeError> where Self: Sized {
170176
match message_type {
171177
CUSTOM_REQUEST_MESSAGE_TYPE => {
@@ -412,7 +418,7 @@ fn async_response_over_one_blinded_hop() {
412418

413419
// 5. Expect Alice to receive the message and create a response instruction for it.
414420
alice.custom_message_handler.expect_message(message.clone());
415-
let response_instruction = nodes[0].custom_message_handler.handle_custom_message(message, responder);
421+
let response_instruction = nodes[0].custom_message_handler.handle_custom_message(message, responder, false);
416422

417423
// 6. Simulate Alice asynchronously responding back to Bob with a response.
418424
assert_eq!(
@@ -425,6 +431,67 @@ fn async_response_over_one_blinded_hop() {
425431
pass_along_path(&nodes);
426432
}
427433

434+
fn do_test_async_response_with_reply_path_over_one_blinded_hop(reply_path_succeed: bool) {
435+
// Simulate an asynchronous interaction between two nodes, Alice and Bob.
436+
437+
// 1. Set up the network with two nodes: Alice and Bob.
438+
let mut nodes = create_nodes(2);
439+
let alice = &nodes[0];
440+
let bob = &nodes[1];
441+
442+
// 2. Define the message sent from Bob to Alice.
443+
let message = TestCustomMessage::ResponseA;
444+
let path_id = Some([2; 32]);
445+
446+
// 3. Simulate the creation of a Blinded Reply path provided by Bob.
447+
let secp_ctx = Secp256k1::new();
448+
let reply_path = BlindedPath::new_for_message(&[bob.node_id], &*bob.entropy_source, &secp_ctx).unwrap();
449+
450+
if reply_path_succeed {
451+
// Add a channel so that nodes are announced to each other.
452+
// This will allow creating the reply path.
453+
add_channel_to_graph(alice, bob, &secp_ctx, 24);
454+
}
455+
456+
// 4. Create a responder using the reply path for Alice.
457+
let responder = Some(Responder::new(reply_path, path_id));
458+
459+
// 5. Expect Alice to receive the message and create a response instruction for it.
460+
alice.custom_message_handler.expect_message(message.clone());
461+
let response_instruction = alice.custom_message_handler.handle_custom_message(message, responder, true);
462+
463+
if !reply_path_succeed {
464+
// 6. Simulate Alice attempting to asynchronously respond back to Bob
465+
// but failing to create a reply path.
466+
assert_eq!(
467+
alice.messenger.handle_onion_message_response(response_instruction),
468+
Err(SendError::PathNotFound),
469+
);
470+
return;
471+
}
472+
473+
// 6. Simulate Alice asynchronously responding back to Bob with a response.
474+
assert_eq!(
475+
alice.messenger.handle_onion_message_response(response_instruction),
476+
Ok(Some(SendSuccess::Buffered)),
477+
);
478+
479+
// 7. Expect Bob to receive the response and handle it accordingly.
480+
bob.custom_message_handler.expect_message(TestCustomMessage::ResponseB);
481+
pass_along_path(&nodes);
482+
483+
// 8. Expect Alice to receive Bob's response through the created reply path and handle it.
484+
alice.custom_message_handler.expect_message(TestCustomMessage::ResponseA);
485+
nodes.reverse();
486+
pass_along_path(&nodes);
487+
}
488+
489+
#[test]
490+
fn async_response_with_reply_path_over_one_blinded_hop() {
491+
do_test_async_response_with_reply_path_over_one_blinded_hop(true);
492+
do_test_async_response_with_reply_path_over_one_blinded_hop(false);
493+
}
494+
428495
#[test]
429496
fn too_big_packet_error() {
430497
// Make sure we error as expected if a packet is too big to send.

lightning/src/onion_message/messenger.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -613,7 +613,7 @@ pub trait CustomOnionMessageHandler {
613613
/// Called with the custom message that was received, returning a response to send, if any.
614614
///
615615
/// The returned [`Self::CustomMessage`], if any, is enqueued to be sent by [`OnionMessenger`].
616-
fn handle_custom_message(&self, message: Self::CustomMessage, responder: Option<Responder>) -> ResponseInstruction<Self::CustomMessage>;
616+
fn handle_custom_message(&self, message: Self::CustomMessage, responder: Option<Responder>, add_reply_path: bool) -> ResponseInstruction<Self::CustomMessage>;
617617

618618
/// Read a custom message of type `message_type` from `buffer`, returning `Ok(None)` if the
619619
/// message type is unknown.
@@ -1227,7 +1227,7 @@ where
12271227
let responder = reply_path.map(
12281228
|reply_path| Responder::new(reply_path, path_id)
12291229
);
1230-
let response_instructions = self.custom_handler.handle_custom_message(msg, responder);
1230+
let response_instructions = self.custom_handler.handle_custom_message(msg, responder, false);
12311231
let _ = self.handle_onion_message_response(response_instructions);
12321232
},
12331233
}

0 commit comments

Comments
 (0)