Merge pull request #2912 from jkczyz/2024-02-prefer-more-connections
[rust-lightning] / lightning / src / onion_message / functional_tests.rs
index 82a08504fe4c7418ea7f9140f92f3bd329a381d8..271b56dc0c92a7a2de07e7d3b6c652777eb608d0 100644 (file)
 //! Onion message testing and test utilities live here.
 
 use crate::blinded_path::BlindedPath;
+use crate::events::{Event, EventsProvider};
 use crate::ln::features::InitFeatures;
-use crate::ln::msgs::{self, DecodeError, OnionMessageHandler};
+use crate::ln::msgs::{self, DecodeError, OnionMessageHandler, SocketAddress};
 use crate::sign::{NodeSigner, Recipient};
 use crate::util::ser::{FixedLengthReader, LengthReadable, Writeable, Writer};
 use crate::util::test_utils;
-use super::{CustomOnionMessageHandler, Destination, MessageRouter, OffersMessage, OffersMessageHandler, OnionMessageContents, OnionMessagePath, OnionMessenger, PendingOnionMessage, SendError};
+use super::messenger::{CustomOnionMessageHandler, Destination, MessageRouter, OnionMessagePath, OnionMessenger, PendingOnionMessage, SendError};
+use super::offers::{OffersMessage, OffersMessageHandler};
+use super::packet::{OnionMessageContents, Packet};
 
 use bitcoin::network::constants::Network;
 use bitcoin::hashes::hex::FromHex;
-use bitcoin::secp256k1::{PublicKey, Secp256k1, SecretKey};
+use bitcoin::secp256k1::{PublicKey, Secp256k1, SecretKey, self};
 
 use crate::io;
 use crate::io_extras::read_to_end;
@@ -50,9 +53,18 @@ impl MessageRouter for TestMessageRouter {
                Ok(OnionMessagePath {
                        intermediate_nodes: vec![],
                        destination,
-                       addresses: None,
+                       first_node_addresses:
+                               Some(vec![SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port: 1000 }]),
                })
        }
+
+       fn create_blinded_paths<
+               T: secp256k1::Signing + secp256k1::Verification
+       >(
+               &self, _recipient: PublicKey, _peers: Vec<PublicKey>, _secp_ctx: &Secp256k1<T>,
+       ) -> Result<Vec<BlindedPath>, ()> {
+               unreachable!()
+       }
 }
 
 struct TestOffersMessageHandler {}
@@ -180,15 +192,30 @@ fn create_nodes_using_secrets(secrets: Vec<SecretKey>) -> Vec<MessengerNode> {
                });
        }
        for i in 0..nodes.len() - 1 {
-               let mut features = InitFeatures::empty();
-               features.set_onion_messages_optional();
-               let init_msg = msgs::Init { features, networks: None, remote_network_address: None };
-               nodes[i].messenger.peer_connected(&nodes[i + 1].node_id, &init_msg.clone(), true).unwrap();
-               nodes[i + 1].messenger.peer_connected(&nodes[i].node_id, &init_msg.clone(), false).unwrap();
+               connect_peers(&nodes[i], &nodes[i + 1]);
        }
        nodes
 }
 
+fn connect_peers(node_a: &MessengerNode, node_b: &MessengerNode) {
+       let mut features = InitFeatures::empty();
+       features.set_onion_messages_optional();
+       let init_msg = msgs::Init { features, networks: None, remote_network_address: None };
+       node_a.messenger.peer_connected(&node_b.node_id, &init_msg.clone(), true).unwrap();
+       node_b.messenger.peer_connected(&node_a.node_id, &init_msg.clone(), false).unwrap();
+}
+
+fn disconnect_peers(node_a: &MessengerNode, node_b: &MessengerNode) {
+       node_a.messenger.peer_disconnected(&node_b.node_id);
+       node_b.messenger.peer_disconnected(&node_a.node_id);
+}
+
+fn release_events(node: &MessengerNode) -> Vec<Event> {
+       let events = core::cell::RefCell::new(Vec::new());
+       node.messenger.process_pending_events(&|e| events.borrow_mut().push(e));
+       events.into_inner()
+}
+
 fn pass_along_path(path: &Vec<MessengerNode>) {
        let mut prev_node = &path[0];
        for node in path.into_iter().skip(1) {
@@ -211,7 +238,7 @@ fn one_unblinded_hop() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::Node(nodes[1].node_id),
-               addresses: None,
+               first_node_addresses: None,
        };
        nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
        nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -226,7 +253,7 @@ fn two_unblinded_hops() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![nodes[1].node_id],
                destination: Destination::Node(nodes[2].node_id),
-               addresses: None,
+               first_node_addresses: None,
        };
        nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
        nodes[2].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -243,7 +270,7 @@ fn one_blinded_hop() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
        nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
        nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -260,7 +287,7 @@ fn two_unblinded_two_blinded() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![nodes[1].node_id, nodes[2].node_id],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
 
        nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
@@ -278,7 +305,7 @@ fn three_blinded_hops() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
 
        nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
@@ -297,7 +324,7 @@ fn too_big_packet_error() {
        let path = OnionMessagePath {
                intermediate_nodes: hops,
                destination: Destination::Node(hop_node_id),
-               addresses: None,
+               first_node_addresses: None,
        };
        let err = nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap_err();
        assert_eq!(err, SendError::TooBigPacket);
@@ -315,7 +342,7 @@ fn we_are_intro_node() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
 
        nodes[0].messenger.send_onion_message_using_path(path, test_msg.clone(), None).unwrap();
@@ -327,7 +354,7 @@ fn we_are_intro_node() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
        nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
        nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -348,7 +375,7 @@ fn invalid_blinded_path_error() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
        let err = nodes[0].messenger.send_onion_message_using_path(path, test_msg.clone(), None).unwrap_err();
        assert_eq!(err, SendError::TooFewBlindedHops);
@@ -364,7 +391,7 @@ fn reply_path() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![nodes[1].node_id, nodes[2].node_id],
                destination: Destination::Node(nodes[3].node_id),
-               addresses: None,
+               first_node_addresses: None,
        };
        let reply_path = BlindedPath::new_for_message(&[nodes[2].node_id, nodes[1].node_id, 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();
@@ -380,7 +407,7 @@ fn reply_path() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::BlindedPath(blinded_path),
-               addresses: None,
+               first_node_addresses: None,
        };
        let reply_path = BlindedPath::new_for_message(&[nodes[2].node_id, nodes[1].node_id, nodes[0].node_id], &*nodes[0].entropy_source, &secp_ctx).unwrap();
 
@@ -415,7 +442,7 @@ fn invalid_custom_message_type() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::Node(nodes[1].node_id),
-               addresses: None,
+               first_node_addresses: None,
        };
        let err = nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap_err();
        assert_eq!(err, SendError::InvalidMessage);
@@ -428,7 +455,7 @@ fn peer_buffer_full() {
        let path = OnionMessagePath {
                intermediate_nodes: vec![],
                destination: Destination::Node(nodes[1].node_id),
-               addresses: None,
+               first_node_addresses: None,
        };
        for _ in 0..188 { // Based on MAX_PER_PEER_BUFFER_SIZE in OnionMessenger
                nodes[0].messenger.send_onion_message_using_path(path.clone(), test_msg.clone(), None).unwrap();
@@ -453,13 +480,79 @@ fn many_hops() {
        let path = OnionMessagePath {
                intermediate_nodes,
                destination: Destination::Node(nodes[num_nodes-1].node_id),
-               addresses: None,
+               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);
        pass_along_path(&nodes);
 }
 
+#[test]
+fn requests_peer_connection_for_buffered_messages() {
+       let nodes = create_nodes(3);
+       let message = TestCustomMessage::Request;
+       let secp_ctx = Secp256k1::new();
+       let blinded_path = BlindedPath::new_for_message(
+               &[nodes[1].node_id, nodes[2].node_id], &*nodes[0].entropy_source, &secp_ctx
+       ).unwrap();
+       let destination = Destination::BlindedPath(blinded_path);
+
+       // Buffer an onion message for a connected peer
+       nodes[0].messenger.send_onion_message(message.clone(), destination.clone(), None).unwrap();
+       assert!(release_events(&nodes[0]).is_empty());
+       assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_some());
+       assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_none());
+
+       // Buffer an onion message for a disconnected peer
+       disconnect_peers(&nodes[0], &nodes[1]);
+       assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_none());
+       nodes[0].messenger.send_onion_message(message, destination, None).unwrap();
+
+       // Check that a ConnectionNeeded event for the peer is provided
+       let events = release_events(&nodes[0]);
+       assert_eq!(events.len(), 1);
+       match &events[0] {
+               Event::ConnectionNeeded { node_id, .. } => assert_eq!(*node_id, nodes[1].node_id),
+               e => panic!("Unexpected event: {:?}", e),
+       }
+
+       // Release the buffered onion message when reconnected
+       connect_peers(&nodes[0], &nodes[1]);
+       assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_some());
+       assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_none());
+}
+
+#[test]
+fn drops_buffered_messages_waiting_for_peer_connection() {
+       let nodes = create_nodes(3);
+       let message = TestCustomMessage::Request;
+       let secp_ctx = Secp256k1::new();
+       let blinded_path = BlindedPath::new_for_message(
+               &[nodes[1].node_id, nodes[2].node_id], &*nodes[0].entropy_source, &secp_ctx
+       ).unwrap();
+       let destination = Destination::BlindedPath(blinded_path);
+
+       // Buffer an onion message for a disconnected peer
+       disconnect_peers(&nodes[0], &nodes[1]);
+       nodes[0].messenger.send_onion_message(message, destination, None).unwrap();
+
+       // Release the event so the timer can start ticking
+       let events = release_events(&nodes[0]);
+       assert_eq!(events.len(), 1);
+       match &events[0] {
+               Event::ConnectionNeeded { node_id, .. } => assert_eq!(*node_id, nodes[1].node_id),
+               e => panic!("Unexpected event: {:?}", e),
+       }
+
+       // Drop buffered messages for a disconnected peer after some timer ticks
+       use crate::onion_message::messenger::MAX_TIMER_TICKS;
+       for _ in 0..=MAX_TIMER_TICKS {
+               nodes[0].messenger.timer_tick_occurred();
+       }
+       connect_peers(&nodes[0], &nodes[1]);
+       assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_none());
+}
+
 #[test]
 fn spec_test_vector() {
        let secret_keys = [
@@ -479,8 +572,8 @@ fn spec_test_vector() {
        let sender_to_alice_packet_bytes_len = sender_to_alice_packet_bytes.len() as u64;
        let mut reader = io::Cursor::new(sender_to_alice_packet_bytes);
        let mut packet_reader = FixedLengthReader::new(&mut reader, sender_to_alice_packet_bytes_len);
-       let sender_to_alice_packet: super::Packet =
-               <super::Packet as LengthReadable>::read(&mut packet_reader).unwrap();
+       let sender_to_alice_packet: Packet =
+               <Packet as LengthReadable>::read(&mut packet_reader).unwrap();
        let secp_ctx = Secp256k1::new();
        let sender_to_alice_om = msgs::OnionMessage {
                blinding_point: PublicKey::from_secret_key(&secp_ctx, &SecretKey::from_slice(&<Vec<u8>>::from_hex("6363636363636363636363636363636363636363636363636363636363636363").unwrap()).unwrap()),