Skip to content

Commit

Permalink
Update Ping Pong roles and introduce new test
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
shaavan committed May 23, 2024
1 parent 164095c commit 594dc0f
Showing 1 changed file with 92 additions and 8 deletions.
100 changes: 92 additions & 8 deletions lightning/src/onion_message/functional_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,16 +111,37 @@ impl Writeable for TestCustomMessage {
}

struct TestCustomMessageHandler {
expected_messages: Mutex<VecDeque<TestCustomMessage>>,
expected_messages: Mutex<VecDeque<Expectations>>,
}

struct Expectations {
message: TestCustomMessage,
add_reply_path: bool,
}

impl TestCustomMessageHandler {
fn new() -> Self {
Self { expected_messages: Mutex::new(VecDeque::new()) }
Self {
expected_messages: Mutex::new(VecDeque::new()),
}
}

fn expect_message(&self, message: TestCustomMessage) {
self.expected_messages.lock().unwrap().push_back(message);
self.expected_messages.lock().unwrap().push_back(
Expectations {
message,
add_reply_path: false,
}
);
}

fn expect_message_and_response(&self, message: TestCustomMessage) {
self.expected_messages.lock().unwrap().push_back(
Expectations {
message,
add_reply_path: true,
}
);
}
}

Expand All @@ -138,16 +159,23 @@ impl Drop for TestCustomMessageHandler {
impl CustomOnionMessageHandler for TestCustomMessageHandler {
type CustomMessage = TestCustomMessage;
fn handle_custom_message(&self, msg: Self::CustomMessage, responder: Option<Responder>) -> ResponseInstruction<Self::CustomMessage> {
match self.expected_messages.lock().unwrap().pop_front() {
Some(expected_msg) => assert_eq!(expected_msg, msg),
let add_reply_path = match self.expected_messages.lock().unwrap().pop_front() {
Some(Expectations { message: expected_msg, add_reply_path }) => {
assert_eq!(expected_msg, msg);
add_reply_path
},
None => panic!("Unexpected message: {:?}", msg),
}
};
let response_option = match msg {
TestCustomMessage::Ping => Some(TestCustomMessage::Pong),
TestCustomMessage::Pong => None,
TestCustomMessage::Pong => Some(TestCustomMessage::Ping),
};
if let (Some(response), Some(responder)) = (response_option, responder) {
responder.respond(response)
if add_reply_path {
responder.respond_with_reply_path(response)
} else {
responder.respond(response)
}
} else {
ResponseInstruction::NoResponse
}
Expand Down Expand Up @@ -401,6 +429,62 @@ fn async_response_over_one_blinded_hop() {
pass_along_path(&nodes);
}

fn do_test_async_response_with_reply_path_over_one_blinded_hop(reply_path_succeed: bool) {
// Simulate an asynchronous interaction between two nodes, Alice and Bob.

let mut nodes = create_nodes(2);
let alice = &nodes[0];
let bob = &nodes[1];

// Alice receives a message from Bob with a reply path
let message = TestCustomMessage::Ping;
let path_id = Some([2; 32]);

let secp_ctx = Secp256k1::new();
let reply_path = BlindedPath::new_for_message(&[bob.node_id], &*bob.entropy_source, &secp_ctx).unwrap();

if reply_path_succeed {
// Add a channel so that nodes are announced to each other.
// This will allow creating the reply path by Alice to include in the response.
add_channel_to_graph(alice, bob, &secp_ctx, 24);
}

let responder = Some(Responder::new(reply_path, path_id));
alice.custom_message_handler.expect_message_and_response(message.clone());

// Alice handles the message reponse, and creates the appropriate ResponseInstruction for it.
let response_instruction = alice.custom_message_handler.handle_custom_message(message, responder);

if !reply_path_succeed {
// Simulate Alice attempting to asynchronously respond back to Bob
// but failing to create a reply path.
assert_eq!(
alice.messenger.handle_onion_message_response(response_instruction),
Err(SendError::PathNotFound),
);
} else {
// Simulate Alice asynchronously responding back to Bob with a response.
assert_eq!(
alice.messenger.handle_onion_message_response(response_instruction),
Ok(Some(SendSuccess::Buffered)),
);

bob.custom_message_handler.expect_message(TestCustomMessage::Pong);
pass_along_path(&nodes);

// Simulate Bob responding back to Alice through the reply path created by her.
alice.custom_message_handler.expect_message(TestCustomMessage::Ping);
nodes.reverse();
pass_along_path(&nodes);
}
}

#[test]
fn async_response_with_reply_path_over_one_blinded_hop() {
do_test_async_response_with_reply_path_over_one_blinded_hop(true);
do_test_async_response_with_reply_path_over_one_blinded_hop(false);
}

#[test]
fn too_big_packet_error() {
// Make sure we error as expected if a packet is too big to send.
Expand Down

0 comments on commit 594dc0f

Please sign in to comment.