From 9f020e4aa51c9d3bbff673fe7a117750762066b7 Mon Sep 17 00:00:00 2001 From: Sejong Kim Date: Sun, 1 Oct 2023 01:01:03 +0900 Subject: [PATCH] Add test cases --- network/Cargo.toml | 1 + network/src/nat_traversal/mod.rs | 211 +++++++++++++++++++------------ 2 files changed, 129 insertions(+), 83 deletions(-) diff --git a/network/Cargo.toml b/network/Cargo.toml index b5d15436..64380a1c 100644 --- a/network/Cargo.toml +++ b/network/Cargo.toml @@ -30,6 +30,7 @@ tokio-stream = { version = "0.1.11", features = ["fs"] } ip_rfc = "0.1.0" parking_lot = "0.12.1" stun = "0.4.4" +regex = "1.7.0" [dev-dependencies] rand = "0.8.5" diff --git a/network/src/nat_traversal/mod.rs b/network/src/nat_traversal/mod.rs index 9477e1a9..62350251 100644 --- a/network/src/nat_traversal/mod.rs +++ b/network/src/nat_traversal/mod.rs @@ -23,76 +23,16 @@ use std::io::{BufRead, BufReader}; use std::sync::Arc; use std::time::Duration; -/* -////////////////////////////////////////////////////////////// -Here is an example for the server behind NAT. -////////////////////////////////////////////////////////////// - -let bind_addr = SocketAddr::from(([0, 0, 0, 0], bind_port)); -let sock = UdpSocket::bind(bind_addr).await?; -let conn = Arc::new(sock); - -tokio::spawn(keep_connecting_to_available_stun( - conn.clone(), - stun_cands_file, - request_period, -)); - -tokio::spawn(keep_connecting_to_clients( - conn.clone(), - allowed_clients_file, - request_period, -)); - -println!("Listening on: {}", conn.clone().local_addr()?); -let mut msg = Message::new(); -let mut buf = [0; 1024]; -loop { - match conn.recv_from(&mut buf).await { - Ok((len, addr)) => { - if !is_stun_response(buf) { - // Handle the DMS message. - } - }, - Err(_) => { - println!("error!"); - } - } -} -*/ - -/* -////////////////////////////////////////////////////////////// -Here is an example for the client behind NAT. -////////////////////////////////////////////////////////////// - -let bind_addr = SocketAddr::from(([0, 0, 0, 0], bind_port)); -let sock = UdpSocket::bind(bind_addr).await?; -let conn = Arc::new(sock); - -tokio::spawn(keep_connecting_to_available_stun( - conn.clone(), - stun_cands_file, - request_period, -)); - -println!("Send a request to the peer: {}", peer_addr); -conn.send_to(b"DMS:CONSENSUS:PING", peer_addr.clone()) - .await?; - -match conn.recv_from(&mut buf).await { - Ok((len, addr)) => { - if !is_stun_response(buf) { - // Handle the DMS message. - } - }, - Err(_) => { - println!("error!"); - } +// send_stun_request sends a request to the given stun addr using the given conn. +async fn send_stun_request(conn: Arc, addr: String) -> Result<(), Error> { + let mut msg = Message::new(); + msg.build(&[Box::::default(), Box::new(BINDING_REQUEST)])?; + + conn.send_to(&msg.raw, addr).await?; + Ok(()) } -*/ -// keep_connecting_to_available_stun helps the given conn keep the public address +// keep_connecting_to_available_stun helps the given conn to maintain the public address // by periodically sending a request to the available STUN server. pub async fn keep_connecting_to_available_stun( conn: Arc, @@ -103,14 +43,9 @@ pub async fn keep_connecting_to_available_stun( let file = File::open(stun_cands_file_name.clone())?; let reader = BufReader::new(file); - let mut msg = Message::new(); - msg.build(&[Box::::default(), Box::new(BINDING_REQUEST)])?; - for line in reader.lines() { match line { - Ok(addr) => { - conn.send_to(&msg.raw, addr).await?; - } + Ok(addr) => send_stun_request(conn.clone(), addr).await?, Err(err) => { eprintln!("Error reading line: {}", err); } @@ -145,20 +80,130 @@ pub async fn keep_connecting_to_clients( } } -// is_stun_response returns whether the given data represents a response from the STUN server, and -// prints out the peer's current public address. -// You should broadcast your address to other peers somehow. -pub fn is_stun_response(data: &Vec) -> bool { +// decode_stun_response returns the public address by decoding the response from the stun server. +// Note(sejongk): You should broadcast your public address to other peers somehow for hole-punching. +pub fn decode_stun_response(data: &[u8]) -> Option { let mut msg = Message::new(); - match msg.unmarshal_binary(&data) { + match msg.unmarshal_binary(data) { Ok(_) => { let mut xor_addr = XorMappedAddress::default(); xor_addr.get_from(&msg).unwrap(); - println!("Your current public address: {xor_addr}"); - return true; + Some(xor_addr.to_string()) + } + Err(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use regex::Regex; + use std::net::SocketAddr; + + #[tokio::test] + async fn get_public_address_via_stun() { + let bind_port = 8080; + let bind_addr = SocketAddr::from(([0, 0, 0, 0], bind_port)); + let sock = UdpSocket::bind(bind_addr).await.unwrap(); + let conn = Arc::new(sock); + + let stun_addr = "stun.l.google.com:19302"; + send_stun_request(conn.clone(), stun_addr.to_string()) + .await + .unwrap(); + + let mut buf = [0; 1024]; + let (_len, _addr) = conn.recv_from(&mut buf).await.unwrap(); + let public_addr = decode_stun_response(&buf); + assert!(public_addr.is_some()); + + let re = Regex::new(r"^((25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.?\b){4}:\d{1,5}$").unwrap(); + assert!(re.is_match(&public_addr.unwrap())); + } + + #[ignore] + #[tokio::test] + async fn run_server_behind_nat() { + let bind_port = 8080; + let bind_addr = SocketAddr::from(([0, 0, 0, 0], bind_port)); + let sock = UdpSocket::bind(bind_addr).await.unwrap(); + let conn = Arc::new(sock); + + let request_period = Duration::from_secs(10); + let stun_cands_file = "stun_cands.txt".to_string(); + let allowed_clients_file = "allowed_clients.txt".to_string(); + + tokio::spawn(keep_connecting_to_available_stun( + conn.clone(), + stun_cands_file, + request_period, + )); + + tokio::spawn(keep_connecting_to_clients( + conn.clone(), + allowed_clients_file, + request_period, + )); + + println!("Listening on: {}", conn.clone().local_addr().unwrap()); + let mut buf = [0; 1024]; + loop { + match conn.recv_from(&mut buf).await { + Ok((len, addr)) => match decode_stun_response(&buf) { + Some(public_addr) => { + println!("Public addr: {}", public_addr) + } + None => { + // Handle the DMS message + println!("{:?} bytes received from the Peer ({:?})", len, addr); + let len = conn.send_to(&buf[..len], addr).await.unwrap(); + println!("{:?} bytes sent to the Peer ({:?})", len, addr); + } + }, + Err(_) => { + println!("error!"); + } + } } - Err(_) => { - return false; + } + + #[ignore] + #[tokio::test] + async fn run_client_behind_nat() { + let bind_port = 8080; + let bind_addr = SocketAddr::from(([0, 0, 0, 0], bind_port)); + let sock = UdpSocket::bind(bind_addr).await.unwrap(); + let conn = Arc::new(sock); + + let request_period = Duration::from_secs(10); + let stun_cands_file = "stun_cands.txt".to_string(); + + tokio::spawn(keep_connecting_to_available_stun( + conn.clone(), + stun_cands_file, + request_period, + )); + + let server_addr = "127.0.0.1:8888".to_string(); + println!("Send a request to the peer: {}", server_addr); + conn.send_to(b"DMS:CONSENSUS:PING", server_addr) + .await + .unwrap(); + + let mut buf = [0; 1024]; + match conn.recv_from(&mut buf).await { + Ok((len, addr)) => match decode_stun_response(&buf) { + Some(public_addr) => { + println!("Public addr: {}", public_addr) + } + None => { + // Handle the DMS message + println!("{:?} bytes received from the Server ({:?})", len, addr); + } + }, + Err(_) => { + println!("error!"); + } } } }