♻️ refactor: refactor the client a lot so I can reuse its code for new subcommands

main
_ 2021-12-09 16:21:14 +00:00
parent 5665f484a2
commit b261d7ba4a
4 changed files with 95 additions and 68 deletions

View File

@ -5,66 +5,22 @@ struct ServerResponse {
nickname: Option <String>, nickname: Option <String>,
} }
pub async fn client <I : Iterator <Item=String>> (mut args: I) -> Result <(), AppError> { struct ClientParams {
use rand::RngCore; common: app_common::Params,
bind_addrs: Vec <Ipv4Addr>,
timeout_ms: u64,
}
pub async fn client <I: Iterator <Item=String>> (args: I) -> Result <(), AppError> {
let params = configure_client (args)?;
let socket = make_socket (&params).await?;
let msg = Message::new_request1 ().to_vec ()?;
let common_params = app_common::Params::default (); tokio::spawn (send_requests (Arc::clone (&socket), params.common, msg));
let mut bind_addrs = vec! [];
let mut timeout_ms = 500;
while let Some (arg) = args.next () {
match arg.as_str () {
"--bind-addr" => {
bind_addrs.push (match args.next () {
None => return Err (CliArgError::MissingArgumentValue (arg).into ()),
Some (x) => Ipv4Addr::from_str (&x)?,
});
},
"--timeout-ms" => {
timeout_ms = match args.next () {
None => return Err (CliArgError::MissingArgumentValue (arg).into ()),
Some (x) => u64::from_str (&x)?,
};
},
_ => return Err (CliArgError::UnrecognizedArgument (arg).into ()),
}
}
if bind_addrs.is_empty () {
bind_addrs = get_ips ()?;
}
let socket = UdpSocket::bind (SocketAddrV4::new (Ipv4Addr::UNSPECIFIED, 0)).await?;
for bind_addr in bind_addrs {
if let Err (e) = socket.join_multicast_v4 (common_params.multicast_addr, bind_addr) {
println! ("Error joining multicast group with iface {}: {:?}", bind_addr, e);
}
}
let mut idem_id = [0u8; 8];
rand::thread_rng ().fill_bytes (&mut idem_id);
let msg = Message::Request1 {
idem_id,
mac: None,
}.to_vec ()?;
let socket = Arc::new (socket);
let socket2 = Arc::clone (&socket);
tokio::spawn (async move {
for _ in 0..10 {
socket2.send_to (&msg, (common_params.multicast_addr, common_params.server_port)).await?;
sleep (Duration::from_millis (100)).await;
}
Ok::<_, AppError> (())
});
let mut peers = HashMap::with_capacity (10); let mut peers = HashMap::with_capacity (10);
timeout (Duration::from_millis (timeout_ms), listen_for_responses (&*socket, &mut peers)).await.ok (); timeout (Duration::from_millis (params.timeout_ms), listen_for_responses (&*socket, &mut peers)).await.ok ();
let mut peers: Vec <_> = peers.into_iter ().collect (); let mut peers: Vec <_> = peers.into_iter ().collect ();
peers.sort_by_key (|(_, v)| v.mac); peers.sort_by_key (|(_, v)| v.mac);
@ -93,6 +49,68 @@ pub async fn client <I : Iterator <Item=String>> (mut args: I) -> Result <(), Ap
Ok (()) Ok (())
} }
fn configure_client <I: Iterator <Item=String>> (mut args: I)
-> Result <ClientParams, AppError>
{
let mut bind_addrs = vec! [];
let mut timeout_ms = 500;
while let Some (arg) = args.next () {
match arg.as_str () {
"--bind-addr" => {
bind_addrs.push (match args.next () {
None => return Err (CliArgError::MissingArgumentValue (arg).into ()),
Some (x) => Ipv4Addr::from_str (&x)?,
});
},
"--timeout-ms" => {
timeout_ms = match args.next () {
None => return Err (CliArgError::MissingArgumentValue (arg).into ()),
Some (x) => u64::from_str (&x)?,
};
},
_ => return Err (CliArgError::UnrecognizedArgument (arg).into ()),
}
}
if bind_addrs.is_empty () {
bind_addrs = get_ips ()?;
}
Ok (ClientParams {
common: Default::default (),
bind_addrs,
timeout_ms,
})
}
async fn make_socket (params: &ClientParams) -> Result <Arc <UdpSocket>, AppError> {
let socket = UdpSocket::bind (SocketAddrV4::new (Ipv4Addr::UNSPECIFIED, 0)).await?;
for bind_addr in &params.bind_addrs {
if let Err (e) = socket.join_multicast_v4 (params.common.multicast_addr, *bind_addr) {
println! ("Error joining multicast group with iface {}: {:?}", bind_addr, e);
}
}
Ok (Arc::new (socket))
}
async fn send_requests (
socket: Arc <UdpSocket>,
params: app_common::Params,
msg: Vec <u8>,
)
-> Result <(), AppError>
{
for _ in 0..10 {
socket.send_to (&msg, (params.multicast_addr, params.server_port)).await?;
sleep (Duration::from_millis (100)).await;
}
Ok::<_, AppError> (())
}
async fn listen_for_responses ( async fn listen_for_responses (
socket: &UdpSocket, socket: &UdpSocket,
peers: &mut HashMap <SocketAddr, ServerResponse> peers: &mut HashMap <SocketAddr, ServerResponse>

View File

@ -6,7 +6,7 @@ mod ip;
pub mod message; pub mod message;
mod prelude; mod prelude;
mod server; mod server;
mod tlv; pub mod tlv;
fn main () -> Result <(), AppError> { fn main () -> Result <(), AppError> {
let rt = tokio::runtime::Builder::new_current_thread () let rt = tokio::runtime::Builder::new_current_thread ()

View File

@ -1,13 +1,4 @@
use std::{ use crate::prelude::*;
io::{
Cursor,
Write,
},
};
use crate::tlv;
use thiserror::Error;
const MAGIC_NUMBER: [u8; 4] = [0x9a, 0x4a, 0x43, 0x81]; const MAGIC_NUMBER: [u8; 4] = [0x9a, 0x4a, 0x43, 0x81];
pub const PACKET_SIZE: usize = 1024; pub const PACKET_SIZE: usize = 1024;
@ -27,13 +18,25 @@ pub enum Message {
Response2 (Response2), Response2 (Response2),
} }
impl Message {
pub fn new_request1 () -> Message {
let mut idem_id = [0u8; 8];
rand::thread_rng ().fill_bytes (&mut idem_id);
Message::Request1 {
idem_id,
mac: None,
}
}
}
#[derive (Debug, PartialEq)] #[derive (Debug, PartialEq)]
pub struct Response2 { pub struct Response2 {
pub idem_id: [u8; 8], pub idem_id: [u8; 8],
pub nickname: String, pub nickname: String,
} }
#[derive (Debug, Error)] #[derive (Debug, thiserror::Error)]
pub enum MessageError { pub enum MessageError {
#[error (transparent)] #[error (transparent)]
Io (#[from] std::io::Error), Io (#[from] std::io::Error),

View File

@ -1,6 +1,10 @@
pub use std::{ pub use std::{
collections::HashMap, collections::HashMap,
env, env,
io::{
Cursor,
Write,
},
net::{ net::{
Ipv4Addr, Ipv4Addr,
SocketAddr, SocketAddr,
@ -18,6 +22,7 @@ pub use mac_address::{
MacAddress, MacAddress,
get_mac_address, get_mac_address,
}; };
pub use rand::RngCore;
pub use tokio::{ pub use tokio::{
net::UdpSocket, net::UdpSocket,
time::{ time::{
@ -39,4 +44,5 @@ pub use crate::{
PACKET_SIZE, PACKET_SIZE,
Message, Message,
}, },
tlv,
}; };