diff --git a/src/handler/mod.rs b/src/handler/mod.rs index cd3c3bc7..d092a897 100644 --- a/src/handler/mod.rs +++ b/src/handler/mod.rs @@ -162,6 +162,8 @@ pub struct WhoAreYouRef(pub NodeAddress, MessageNonce); #[derive(Debug)] /// A Challenge (WHOAREYOU) object used to handle and send WHOAREYOU requests. pub struct Challenge { + /// The WHOAREYOU packet sent to the remote. + packet: Packet, /// The challenge data received from the node. data: ChallengeData, /// The remote's ENR if we know it. We can receive a challenge from an unknown node. @@ -584,8 +586,13 @@ impl Handler { let node_address = wru_ref.0; let message_nonce = wru_ref.1; - if self.active_challenges.get(&node_address).is_some() { - warn!(%node_address, "WHOAREYOU already sent."); + if let Some(packet) = self + .active_challenges + .get(&node_address) + .map(|challenge| challenge.packet.clone()) + { + debug!(%node_address, "Resending WHOAREYOU"); + self.send(node_address, packet).await; return; } @@ -608,10 +615,11 @@ impl Handler { .expect("Must be the correct challenge size"); debug!(%node_address, "Sending WHOAREYOU"); self.add_expected_response(node_address.socket_addr); - self.send(node_address.clone(), packet).await; + self.send(node_address.clone(), packet.clone()).await; self.active_challenges.insert( node_address, Challenge { + packet, data: challenge_data, remote_enr, }, diff --git a/src/handler/tests.rs b/src/handler/tests.rs index 79d2b2e5..fa1a6156 100644 --- a/src/handler/tests.rs +++ b/src/handler/tests.rs @@ -178,6 +178,108 @@ async fn simple_session_message() { } } +#[tokio::test] +async fn handshake_resends_active_challenge() { + init(); + + let receiver_socket = Arc::new( + tokio::net::UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .unwrap(), + ); + let receiver_addr = receiver_socket.local_addr().unwrap(); + let sender_socket = tokio::net::UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .unwrap(); + let sender_addr = sender_socket.local_addr().unwrap(); + + let sender_key = CombinedKey::generate_secp256k1(); + let receiver_key = CombinedKey::generate_secp256k1(); + let sender_enr = Enr::builder() + .ip4(Ipv4Addr::LOCALHOST) + .udp4(sender_addr.port()) + .build(&sender_key) + .unwrap(); + let receiver_enr = Enr::builder() + .ip4(Ipv4Addr::LOCALHOST) + .udp4(receiver_addr.port()) + .build(&receiver_key) + .unwrap(); + + let config = ConfigBuilder::new(ListenConfig::FromSockets { + ipv4: Some(receiver_socket), + ipv6: None, + }) + .build(); + let (exit, receiver_send, mut receiver_recv) = + Handler::spawn(arc_rw!(receiver_enr.clone()), arc_rw!(receiver_key), config) + .await + .unwrap(); + + let first_nonce = [1; crate::packet::MESSAGE_NONCE_LENGTH]; + let mut first_challenge = None; + let mut first_challenge_data = None; + + for (request_id, message_nonce) in [ + (1, first_nonce), + (2, [2; crate::packet::MESSAGE_NONCE_LENGTH]), + ] { + let ping = Packet::new_message( + sender_enr.node_id(), + message_nonce, + ProtocolIdentity::default(), + Request { + id: RequestId(vec![request_id]), + body: RequestBody::Ping { enr_seq: 1 }, + } + .encode(), + ) + .encode(&receiver_enr.node_id()); + sender_socket.send_to(&ping, receiver_addr).await.unwrap(); + + let whoareyou_ref = match tokio::time::timeout(Duration::from_secs(2), receiver_recv.recv()) + .await + .unwrap() + .unwrap() + { + HandlerOut::WhoAreYou(whoareyou_ref) => whoareyou_ref, + message => panic!("expected WHOAREYOU request, got {:?}", message), + }; + assert_eq!(whoareyou_ref.0.socket_addr, sender_addr); + assert_eq!(whoareyou_ref.1, message_nonce); + receiver_send + .send(HandlerIn::WhoAreYou(whoareyou_ref, None)) + .unwrap(); + + let mut response = [0; crate::packet::MAX_PACKET_SIZE]; + let (response_len, _) = tokio::time::timeout( + Duration::from_secs(2), + sender_socket.recv_from(&mut response), + ) + .await + .unwrap() + .unwrap(); + let response = &response[..response_len]; + let (challenge, challenge_data) = + Packet::decode(&sender_enr.node_id(), ProtocolIdentity::default(), response).unwrap(); + assert!(matches!( + challenge.header.kind, + PacketKind::WhoAreYou { .. } + )); + assert_eq!(challenge.header.message_nonce, first_nonce); + + if let Some(ref expected_challenge) = first_challenge { + assert_eq!(response, expected_challenge); + assert_eq!(Some(challenge_data), first_challenge_data); + } else { + first_challenge = Some(response.to_vec()); + first_challenge_data = Some(challenge_data); + } + } + + exit.send(()).unwrap(); +} + #[tokio::test] // Tests sending multiple messages on an encrypted session async fn multiple_messages() {