Skip to content

Commit f5daf5f

Browse files
committed
Update Ping Pong roles and introduce new test
1. Update the Ping and Pong roles to make them each other's response. This allows unlimited back-and-forth communication in the test suite. 2. Introduce a new member variable for `TestCustomMessageHandler`, which signals to `handle_custom_message` which `ResponseInstruction` variant it should create. 3. Introduce a new test for `ResponseInstruction::WithReplyPath`, verifying both successful and unsuccessful `reply_path` creation.
1 parent 164095c commit f5daf5f

File tree

1 file changed

+79
-3
lines changed

1 file changed

+79
-3
lines changed

lightning/src/onion_message/functional_tests.rs

Lines changed: 79 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,16 +112,31 @@ impl Writeable for TestCustomMessage {
112112

113113
struct TestCustomMessageHandler {
114114
expected_messages: Mutex<VecDeque<TestCustomMessage>>,
115+
add_reply_path: Mutex<bool>,
115116
}
116117

117118
impl TestCustomMessageHandler {
118119
fn new() -> Self {
119-
Self { expected_messages: Mutex::new(VecDeque::new()) }
120+
Self {
121+
expected_messages: Mutex::new(VecDeque::new()),
122+
add_reply_path: Mutex::new(false),
123+
}
120124
}
121125

122126
fn expect_message(&self, message: TestCustomMessage) {
123127
self.expected_messages.lock().unwrap().push_back(message);
124128
}
129+
130+
fn include_reply_path(&self) {
131+
*self.add_reply_path.lock().unwrap() = true;
132+
}
133+
134+
fn read_and_reset_add_reply_path(&self) -> bool {
135+
let mut add_reply_path = self.add_reply_path.lock().unwrap();
136+
let was_set = *add_reply_path;
137+
*add_reply_path = false;
138+
was_set
139+
}
125140
}
126141

127142
impl Drop for TestCustomMessageHandler {
@@ -144,10 +159,14 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
144159
}
145160
let response_option = match msg {
146161
TestCustomMessage::Ping => Some(TestCustomMessage::Pong),
147-
TestCustomMessage::Pong => None,
162+
TestCustomMessage::Pong => Some(TestCustomMessage::Ping),
148163
};
149164
if let (Some(response), Some(responder)) = (response_option, responder) {
150-
responder.respond(response)
165+
if self.read_and_reset_add_reply_path() {
166+
responder.respond_with_reply_path(response)
167+
} else {
168+
responder.respond(response)
169+
}
151170
} else {
152171
ResponseInstruction::NoResponse
153172
}
@@ -401,6 +420,63 @@ fn async_response_over_one_blinded_hop() {
401420
pass_along_path(&nodes);
402421
}
403422

423+
fn do_test_async_response_with_reply_path_over_one_blinded_hop(reply_path_succeed: bool) {
424+
// Simulate an asynchronous interaction between two nodes, Alice and Bob.
425+
426+
let mut nodes = create_nodes(2);
427+
let alice = &nodes[0];
428+
let bob = &nodes[1];
429+
430+
// Alice receives a message from Bob with a reply path
431+
let message = TestCustomMessage::Ping;
432+
let path_id = Some([2; 32]);
433+
434+
let secp_ctx = Secp256k1::new();
435+
let reply_path = BlindedPath::new_for_message(&[bob.node_id], &*bob.entropy_source, &secp_ctx).unwrap();
436+
437+
if reply_path_succeed {
438+
// Add a channel so that nodes are announced to each other.
439+
// This will allow creating the reply path by Alice to include in the response.
440+
add_channel_to_graph(alice, bob, &secp_ctx, 24);
441+
}
442+
443+
let responder = Some(Responder::new(reply_path, path_id));
444+
alice.custom_message_handler.expect_message(message.clone());
445+
alice.custom_message_handler.include_reply_path();
446+
447+
// Alice handles the message reponse, and creates the appropriate ResponseInstruction for it.
448+
let response_instruction = alice.custom_message_handler.handle_custom_message(message, responder);
449+
450+
if !reply_path_succeed {
451+
// Simulate Alice attempting to asynchronously respond back to Bob
452+
// but failing to create a reply path.
453+
assert_eq!(
454+
alice.messenger.handle_onion_message_response(response_instruction),
455+
Err(SendError::PathNotFound),
456+
);
457+
} else {
458+
// Simulate Alice asynchronously responding back to Bob with a response.
459+
assert_eq!(
460+
alice.messenger.handle_onion_message_response(response_instruction),
461+
Ok(Some(SendSuccess::Buffered)),
462+
);
463+
464+
bob.custom_message_handler.expect_message(TestCustomMessage::Pong);
465+
pass_along_path(&nodes);
466+
467+
// Simulate Bob responding back to Alice through the reply path created by her.
468+
alice.custom_message_handler.expect_message(TestCustomMessage::Ping);
469+
nodes.reverse();
470+
pass_along_path(&nodes);
471+
}
472+
}
473+
474+
#[test]
475+
fn async_response_with_reply_path_over_one_blinded_hop() {
476+
do_test_async_response_with_reply_path_over_one_blinded_hop(true);
477+
do_test_async_response_with_reply_path_over_one_blinded_hop(false);
478+
}
479+
404480
#[test]
405481
fn too_big_packet_error() {
406482
// Make sure we error as expected if a packet is too big to send.

0 commit comments

Comments
 (0)