From c7681ce9f5f0e559647cde84303b52a2ae7f0d95 Mon Sep 17 00:00:00 2001 From: _ <_@_> Date: Thu, 9 Dec 2021 01:28:29 +0000 Subject: [PATCH] :recycle: refactor --- src/app_common.rs | 58 ++++++++++ src/client.rs | 99 +++++++++++++++++ src/ip.rs | 12 +- src/main.rs | 273 ++-------------------------------------------- src/prelude.rs | 37 +++++++ src/server.rs | 78 +++++++++++++ 6 files changed, 289 insertions(+), 268 deletions(-) create mode 100644 src/app_common.rs create mode 100644 src/client.rs create mode 100644 src/prelude.rs create mode 100644 src/server.rs diff --git a/src/app_common.rs b/src/app_common.rs new file mode 100644 index 0000000..4e7a8b5 --- /dev/null +++ b/src/app_common.rs @@ -0,0 +1,58 @@ +use crate::prelude::*; + +#[derive (Debug, thiserror::Error)] +pub enum AppError { + #[error (transparent)] + AddrParse (#[from] std::net::AddrParseError), + #[error (transparent)] + CliArgs (#[from] CliArgError), + #[error (transparent)] + Io (#[from] std::io::Error), + #[error (transparent)] + Ip (#[from] crate::ip::IpError), + #[error (transparent)] + MacAddr (#[from] mac_address::MacAddressError), + #[error (transparent)] + Message (#[from] crate::message::MessageError), + #[error (transparent)] + Tlv (#[from] crate::tlv::TlvError), +} + +#[derive (Debug, thiserror::Error)] +pub enum CliArgError { + #[error ("Missing value for argument `{0}`")] + MissingArgumentValue (String), + #[error ("First argument should be a subcommand")] + MissingSubcommand, + #[error ("Unknown subcommand `{0}`")] + UnknownSubcommand (String), + #[error ("Unrecognized argument `{0}`")] + UnrecognizedArgument (String), +} + +pub async fn recv_msg_from (socket: &UdpSocket) -> Result <(Vec , SocketAddr), AppError> +{ + let mut buf = vec! [0u8; PACKET_SIZE]; + let (bytes_recved, remote_addr) = socket.recv_from (&mut buf).await?; + buf.truncate (bytes_recved); + let msgs = Message::from_slice2 (&buf)?; + + Ok ((msgs, remote_addr)) +} + +pub struct Params { + // Servers bind on this port, clients must send to the port + pub server_port: u16, + + // Clients and servers will all join the same multicast addr + pub multicast_addr: Ipv4Addr, +} + +impl Default for Params { + fn default () -> Self { + Self { + server_port: 9040, + multicast_addr: Ipv4Addr::new (225, 100, 99, 98), + } + } +} diff --git a/src/client.rs b/src/client.rs new file mode 100644 index 0000000..6e0c39b --- /dev/null +++ b/src/client.rs @@ -0,0 +1,99 @@ +use crate::prelude::*; + +struct ServerResponse { + mac: Option <[u8; 6]>, + nickname: Option , +} + +pub async fn client > (mut args: I) -> Result <(), AppError> { + use rand::RngCore; + + let common_params = app_common::Params::default (); + let mut bind_addr = "0.0.0.0".to_string (); + + while let Some (arg) = args.next () { + match arg.as_str () { + "--bind-addr" => { + bind_addr = match args.next () { + None => return Err (CliArgError::MissingArgumentValue (arg).into ()), + Some (x) => x + }; + }, + _ => return Err (CliArgError::UnrecognizedArgument (arg).into ()), + } + } + + let socket = UdpSocket::bind (&format! ("{}:0", bind_addr)).await?; + + socket.join_multicast_v4 (common_params.multicast_addr, Ipv4Addr::from_str (&bind_addr)?)?; + + let mut idem_id = [0u8; 8]; + rand::thread_rng ().fill_bytes (&mut idem_id); + + let msg = Message::Request1 { + idem_id, + mac: None, + }.to_vec ()?; + + for _ in 0..10 { + socket.send_to (&msg, (common_params.multicast_addr, common_params.server_port)).await?; + sleep (Duration::from_millis (100)).await; + } + + let mut peers = HashMap::with_capacity (10); + + timeout (Duration::from_secs (2), listen_for_responses (&socket, &mut peers)).await.ok (); + + let mut peers: Vec <_> = peers.into_iter ().collect (); + peers.sort_by_key (|(_, v)| v.mac); + + println! ("Found {} peers:", peers.len ()); + for (ip, resp) in peers.into_iter () { + let mac = match resp.mac { + None => { + println! (" = {}", ip); + continue; + }, + Some (x) => x, + }; + + let nickname = match resp.nickname { + None => { + println! ("{} = {}", MacAddress::new (mac), ip.ip ()); + continue; + }, + Some (x) => x, + }; + + println! ("{} = {} `{}`", MacAddress::new (mac), ip.ip (), nickname); + } + + Ok (()) +} + +async fn listen_for_responses ( + socket: &UdpSocket, + peers: &mut HashMap +) { + loop { + let (msgs, remote_addr) = match recv_msg_from (socket).await { + Err (_) => continue, + Ok (x) => x, + }; + + let mut resp = ServerResponse { + mac: None, + nickname: None, + }; + + for msg in msgs.into_iter () { + match msg { + Message::Response1 (x) => resp.mac = x, + Message::Response2 (x) => resp.nickname = Some (x.nickname), + _ => (), + } + } + + peers.insert (remote_addr, resp); + } +} diff --git a/src/ip.rs b/src/ip.rs index 9fe9dca..8c5f3df 100644 --- a/src/ip.rs +++ b/src/ip.rs @@ -4,13 +4,19 @@ use std::{ str::FromStr, }; -use crate::AppError; +#[derive (Debug, thiserror::Error)] +pub enum IpError { + #[error (transparent)] + Io (#[from] std::io::Error), + #[error (transparent)] + FromUtf8 (#[from] std::string::FromUtf8Error), +} #[cfg(target_os = "linux")] pub mod linux { use super::*; - pub fn get_ip_addr_output () -> Result { + pub fn get_ip_addr_output () -> Result { let output = Command::new ("ip") .arg ("addr") .output ()?; @@ -35,7 +41,7 @@ pub mod linux { pub mod windows { use super::*; - pub fn get_ip_config_output () -> Result { + pub fn get_ip_config_output () -> Result { let output = Command::new ("ipconfig") .output ()?; let output = output.stdout.as_slice (); diff --git a/src/main.rs b/src/main.rs index dd3266f..f7930e5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,86 +1,13 @@ -use std::{ - collections::HashMap, - env, - net::{ - Ipv4Addr, - SocketAddr, - SocketAddrV4, - }, - str::FromStr, - time::{Duration}, -}; - -use mac_address::{ - MacAddress, - get_mac_address, -}; -use thiserror::Error; -use tokio::{ - net::UdpSocket, - time::{ - sleep, - timeout, - }, -}; +use prelude::*; +pub mod app_common; +mod client; mod ip; -mod message; +pub mod message; +mod prelude; +mod server; mod tlv; -use message::{ - PACKET_SIZE, - Message, -}; - -#[derive (Debug, Error)] -pub enum AppError { - #[error (transparent)] - AddrParse (#[from] std::net::AddrParseError), - #[error (transparent)] - CliArgs (#[from] CliArgError), - #[error (transparent)] - FromUtf8 (#[from] std::string::FromUtf8Error), - #[error (transparent)] - Io (#[from] std::io::Error), - #[error (transparent)] - MacAddr (#[from] mac_address::MacAddressError), - #[error (transparent)] - Message (#[from] message::MessageError), - #[error (transparent)] - Tlv (#[from] tlv::TlvError), - #[error (transparent)] - Utf8 (#[from] std::str::Utf8Error), -} - -#[derive (Debug, Error)] -pub enum CliArgError { - #[error ("Missing value for argument `{0}`")] - MissingArgumentValue (String), - #[error ("First argument should be a subcommand")] - MissingSubcommand, - #[error ("Unknown subcommand `{0}`")] - UnknownSubcommand (String), - #[error ("Unrecognized argument `{0}`")] - UnrecognizedArgument (String), -} - -struct CommonParams { - // Servers bind on this port, clients must send to the port - server_port: u16, - - // Clients and servers will all join the same multicast addr - multicast_addr: Ipv4Addr, -} - -impl Default for CommonParams { - fn default () -> Self { - Self { - server_port: 9040, - multicast_addr: Ipv4Addr::new (225, 100, 99, 98), - } - } -} - fn main () -> Result <(), AppError> { let rt = tokio::runtime::Builder::new_current_thread () .enable_io () @@ -109,9 +36,9 @@ async fn async_main () -> Result <(), AppError> { match subcommand.as_ref ().map (|x| &x[..]) { None => return Err (CliArgError::MissingSubcommand.into ()), - Some ("client") => client (args).await?, + Some ("client") => client::client (args).await?, Some ("my-ips") => my_ips ()?, - Some ("server") => server (args).await?, + Some ("server") => server::server (args).await?, Some (x) => return Err (CliArgError::UnknownSubcommand (x.to_string ()).into ()), } @@ -149,187 +76,3 @@ fn my_ips () -> Result <(), AppError> { Ok (()) } -struct ServerResponse { - mac: Option <[u8; 6]>, - nickname: Option , -} - -async fn client > (mut args: I) -> Result <(), AppError> { - use rand::RngCore; - - let common_params = CommonParams::default (); - let mut bind_addr = "0.0.0.0".to_string (); - - while let Some (arg) = args.next () { - match arg.as_str () { - "--bind-addr" => { - bind_addr = match args.next () { - None => return Err (CliArgError::MissingArgumentValue (arg).into ()), - Some (x) => x - }; - }, - _ => return Err (CliArgError::UnrecognizedArgument (arg).into ()), - } - } - - let socket = UdpSocket::bind (&format! ("{}:0", bind_addr)).await?; - - socket.join_multicast_v4 (common_params.multicast_addr, Ipv4Addr::from_str (&bind_addr)?)?; - - let mut idem_id = [0u8; 8]; - rand::thread_rng ().fill_bytes (&mut idem_id); - - let msg = Message::Request1 { - idem_id, - mac: None, - }.to_vec ()?; - - for _ in 0..10 { - socket.send_to (&msg, (common_params.multicast_addr, common_params.server_port)).await?; - sleep (Duration::from_millis (100)).await; - } - - let mut peers = HashMap::with_capacity (10); - - timeout (Duration::from_secs (2), listen_for_responses (&socket, &mut peers)).await.ok (); - - let mut peers: Vec <_> = peers.into_iter ().collect (); - peers.sort_by_key (|(_, v)| v.mac); - - println! ("Found {} peers:", peers.len ()); - for (ip, resp) in peers.into_iter () { - let mac = match resp.mac { - None => { - println! (" = {}", ip); - continue; - }, - Some (x) => x, - }; - - let nickname = match resp.nickname { - None => { - println! ("{} = {}", MacAddress::new (mac), ip.ip ()); - continue; - }, - Some (x) => x, - }; - - println! ("{} = {} `{}`", MacAddress::new (mac), ip.ip (), nickname); - } - - Ok (()) -} - -async fn listen_for_responses ( - socket: &UdpSocket, - peers: &mut HashMap -) { - loop { - let (msgs, remote_addr) = match recv_msg_from (socket).await { - Err (_) => continue, - Ok (x) => x, - }; - - let mut resp = ServerResponse { - mac: None, - nickname: None, - }; - - for msg in msgs.into_iter () { - match msg { - Message::Response1 (x) => resp.mac = x, - Message::Response2 (x) => resp.nickname = Some (x.nickname), - _ => (), - } - } - - peers.insert (remote_addr, resp); - } -} - -async fn server > (mut args: I) -> Result <(), AppError> -{ - let common_params = CommonParams::default (); - let mut bind_addr = "0.0.0.0".to_string (); - let mut nickname = String::new (); - - while let Some (arg) = args.next () { - match arg.as_str () { - "--bind-addr" => { - bind_addr = match args.next () { - None => return Err (CliArgError::MissingArgumentValue (arg).into ()), - Some (x) => x - }; - }, - "--nickname" => { - nickname = match args.next () { - None => return Err (CliArgError::MissingArgumentValue (arg).into ()), - Some (x) => x - }; - }, - _ => return Err (CliArgError::UnrecognizedArgument (arg).into ()), - } - } - - let our_mac = get_mac_address ()?.map (|x| x.bytes ()); - if our_mac.is_none () { - println! ("Warning: Can't find our own MAC address. We won't be able to respond to MAC-specific lookaround requests"); - } - - let socket = UdpSocket::bind (SocketAddrV4::new (Ipv4Addr::from_str (&bind_addr)?, common_params.server_port)).await.unwrap (); - - socket.join_multicast_v4 (common_params.multicast_addr, [0u8, 0, 0, 0].into ()).unwrap (); - - let mut recent_idem_ids = Vec::with_capacity (32); - - loop { - println! ("Waiting for messages..."); - let (req_msgs, remote_addr) = recv_msg_from (&socket).await?; - - let req = match req_msgs.into_iter ().next () { - Some (x) => x, - _ => { - println! ("Don't know how to handle this message, ignoring"); - continue; - }, - }; - - let resp = match req { - Message::Request1 { - mac: None, - idem_id, - } => { - if recent_idem_ids.contains (&idem_id) { - println! ("Ignoring request we already processed"); - None - } - else { - recent_idem_ids.insert (0, idem_id); - recent_idem_ids.truncate (30); - Some (vec! [ - Message::Response1 (our_mac), - Message::Response2 (message::Response2 { - idem_id, - nickname: nickname.clone (), - }), - ]) - } - }, - _ => continue, - }; - - if let Some (resp) = resp { - socket.send_to (&Message::many_to_vec (&resp)?, remote_addr).await.unwrap (); - } - } -} - -async fn recv_msg_from (socket: &UdpSocket) -> Result <(Vec , SocketAddr), AppError> -{ - let mut buf = vec! [0u8; PACKET_SIZE]; - let (bytes_recved, remote_addr) = socket.recv_from (&mut buf).await?; - buf.truncate (bytes_recved); - let msgs = Message::from_slice2 (&buf)?; - - Ok ((msgs, remote_addr)) -} diff --git a/src/prelude.rs b/src/prelude.rs new file mode 100644 index 0000000..e83b23b --- /dev/null +++ b/src/prelude.rs @@ -0,0 +1,37 @@ +pub use std::{ + collections::HashMap, + env, + net::{ + Ipv4Addr, + SocketAddr, + SocketAddrV4, + }, + str::FromStr, + time::Duration, +}; + +pub use mac_address::{ + MacAddress, + get_mac_address, +}; +pub use tokio::{ + net::UdpSocket, + time::{ + sleep, + timeout, + }, +}; + +pub use crate::{ + app_common::{ + self, + AppError, + CliArgError, + recv_msg_from, + }, + message::{ + self, + PACKET_SIZE, + Message, + }, +}; diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..eb865fe --- /dev/null +++ b/src/server.rs @@ -0,0 +1,78 @@ +use crate::prelude::*; + +pub async fn server > (mut args: I) -> Result <(), AppError> +{ + let common_params = app_common::Params::default (); + let mut bind_addr = "0.0.0.0".to_string (); + let mut nickname = String::new (); + + while let Some (arg) = args.next () { + match arg.as_str () { + "--bind-addr" => { + bind_addr = match args.next () { + None => return Err (CliArgError::MissingArgumentValue (arg).into ()), + Some (x) => x + }; + }, + "--nickname" => { + nickname = match args.next () { + None => return Err (CliArgError::MissingArgumentValue (arg).into ()), + Some (x) => x + }; + }, + _ => return Err (CliArgError::UnrecognizedArgument (arg).into ()), + } + } + + let our_mac = get_mac_address ()?.map (|x| x.bytes ()); + if our_mac.is_none () { + println! ("Warning: Can't find our own MAC address. We won't be able to respond to MAC-specific lookaround requests"); + } + + let socket = UdpSocket::bind (SocketAddrV4::new (Ipv4Addr::from_str (&bind_addr)?, common_params.server_port)).await.unwrap (); + + socket.join_multicast_v4 (common_params.multicast_addr, [0u8, 0, 0, 0].into ()).unwrap (); + + let mut recent_idem_ids = Vec::with_capacity (32); + + loop { + println! ("Waiting for messages..."); + let (req_msgs, remote_addr) = recv_msg_from (&socket).await?; + + let req = match req_msgs.into_iter ().next () { + Some (x) => x, + _ => { + println! ("Don't know how to handle this message, ignoring"); + continue; + }, + }; + + let resp = match req { + Message::Request1 { + mac: None, + idem_id, + } => { + if recent_idem_ids.contains (&idem_id) { + println! ("Ignoring request we already processed"); + None + } + else { + recent_idem_ids.insert (0, idem_id); + recent_idem_ids.truncate (30); + Some (vec! [ + Message::Response1 (our_mac), + Message::Response2 (message::Response2 { + idem_id, + nickname: nickname.clone (), + }), + ]) + } + }, + _ => continue, + }; + + if let Some (resp) = resp { + socket.send_to (&Message::many_to_vec (&resp)?, remote_addr).await.unwrap (); + } + } +}