Skip to content

Commit 594dc0f

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 for unlimited back-and-forth communication in the test suite. 2. Introduce new structs for the expectation fields. This allows for keeping the expectations ordered in a neat way. 3. Introduce a new boolean field in Expectations. This keeps track of whether a response is expected or not and creates the appropriate ResponseInstruction accordingly in handle_custom_message. 4. Introduce a new test for `ResponseInstruction::WithReplyPath`, verifying both successful and unsuccessful `reply_path` creation.
1 parent 164095c commit 594dc0f

File tree

1 file changed

+92
-8
lines changed

1 file changed

+92
-8
lines changed

lightning/src/onion_message/functional_tests.rs

Lines changed: 92 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -111,16 +111,37 @@ impl Writeable for TestCustomMessage {
111111
}
112112

113113
struct TestCustomMessageHandler {
114-
expected_messages: Mutex<VecDeque<TestCustomMessage>>,
114+
expected_messages: Mutex<VecDeque<Expectations>>,
115+
}
116+
117+
struct Expectations {
118+
message: TestCustomMessage,
119+
add_reply_path: bool,
115120
}
116121

117122
impl TestCustomMessageHandler {
118123
fn new() -> Self {
119-
Self { expected_messages: Mutex::new(VecDeque::new()) }
124+
Self {
125+
expected_messages: Mutex::new(VecDeque::new()),
126+
}
120127
}
121128

122129
fn expect_message(&self, message: TestCustomMessage) {
123-
self.expected_messages.lock().unwrap().push_back(message);
130+
self.expected_messages.lock().unwrap().push_back(
131+
Expectations {
132+
message,
133+
add_reply_path: false,
134+
}
135+
);
136+
}
137+
138+
fn expect_message_and_response(&self, message: TestCustomMessage) {
139+
self.expected_messages.lock().unwrap().push_back(
140+
Expectations {
141+
message,
142+
add_reply_path: true,
143+
}
144+
);
124145
}
125146
}
126147

@@ -138,16 +159,23 @@ impl Drop for TestCustomMessageHandler {
138159
impl CustomOnionMessageHandler for TestCustomMessageHandler {
139160
type CustomMessage = TestCustomMessage;
140161
fn handle_custom_message(&self, msg: Self::CustomMessage, responder: Option<Responder>) -> ResponseInstruction<Self::CustomMessage> {
141-
match self.expected_messages.lock().unwrap().pop_front() {
142-
Some(expected_msg) => assert_eq!(expected_msg, msg),
162+
let add_reply_path = match self.expected_messages.lock().unwrap().pop_front() {
163+
Some(Expectations { message: expected_msg, add_reply_path }) => {
164+
assert_eq!(expected_msg, msg);
165+
add_reply_path
166+
},
143167
None => panic!("Unexpected message: {:?}", msg),
144-
}
168+
};
145169
let response_option = match msg {
146170
TestCustomMessage::Ping => Some(TestCustomMessage::Pong),
147-
TestCustomMessage::Pong => None,
171+
TestCustomMessage::Pong => Some(TestCustomMessage::Ping),
148172
};
149173
if let (Some(response), Some(responder)) = (response_option, responder) {
150-
responder.respond(response)
174+
if add_reply_path {
175+
responder.respond_with_reply_path(response)
176+
} else {
177+
responder.respond(response)
178+
}
151179
} else {
152180
ResponseInstruction::NoResponse
153181
}
@@ -401,6 +429,62 @@ fn async_response_over_one_blinded_hop() {
401429
pass_along_path(&nodes);
402430
}
403431

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

0 commit comments

Comments
 (0)