From aa75119f39d271f91c3510d072e24d2f4e443a17 Mon Sep 17 00:00:00 2001 From: _ <_@_> Date: Wed, 8 Dec 2021 02:10:58 +0000 Subject: [PATCH] :heavy_plus_sign: add nicknames --- src/main.rs | 96 +++++++++++--- src/message.rs | 351 ++++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 409 insertions(+), 38 deletions(-) diff --git a/src/main.rs b/src/main.rs index 9021d2e..d263230 100644 --- a/src/main.rs +++ b/src/main.rs @@ -40,10 +40,14 @@ enum AppError { #[derive (Debug, Error)] 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 { @@ -88,6 +92,11 @@ fn main () -> Result <(), AppError> { Ok (()) } +struct ServerResponse { + mac: Option <[u8; 6]>, + nickname: Option , +} + fn client () -> Result <(), AppError> { use rand::RngCore; @@ -100,7 +109,7 @@ fn client () -> Result <(), AppError> { let mut idem_id = [0u8; 8]; rand::thread_rng ().fill_bytes (&mut idem_id); - let msg = Message::Request { + let msg = Message::Request1 { idem_id, mac: None, }.to_vec ()?; @@ -115,38 +124,71 @@ fn client () -> Result <(), AppError> { let mut peers = HashMap::with_capacity (10); while Instant::now () < start_time + Duration::from_secs (2) { - let (resp, remote_addr) = match recv_msg_from (&socket) { + let (msgs, remote_addr) = match recv_msg_from (&socket) { Err (_) => continue, Ok (x) => x, }; - let peer_mac_addr = match resp { - Message::Response (mac) => mac, - _ => continue, + let mut resp = ServerResponse { + mac: None, + nickname: None, }; - peers.insert (remote_addr, peer_mac_addr); + 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); } let mut peers: Vec <_> = peers.into_iter ().collect (); - peers.sort (); + peers.sort_by_key (|(k, v)| v.mac); println! ("Found {} peers:", peers.len ()); - for (ip, mac) in &peers { - match mac { - Some (mac) => println! ("{} = {}", MacAddress::new (*mac), ip.ip ()), - None => println! (" = {}", ip), - } + 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 (()) } -fn server > (args: I) -> Result <(), AppError> +fn server > (mut args: I) -> Result <(), AppError> { let mut common_params = CommonParams::default (); let mut nickname = String::new (); + while let Some (arg) = args.next () { + match arg.as_str () { + "--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"); @@ -160,10 +202,18 @@ fn server > (args: I) -> Result <(), AppError> loop { println! ("Waiting for messages..."); - let (req, remote_addr) = recv_msg_from (&socket)?; + let (req_msgs, remote_addr) = recv_msg_from (&socket)?; + + 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::Request { + Message::Request1 { mac: None, idem_id, } => { @@ -174,24 +224,30 @@ fn server > (args: I) -> Result <(), AppError> else { recent_idem_ids.insert (0, idem_id); recent_idem_ids.truncate (30); - Some (Message::Response (our_mac)) + Some (vec! [ + Message::Response1 (our_mac), + Message::Response2 (message::Response2 { + idem_id, + nickname: nickname.clone (), + }), + ]) } }, _ => continue, }; if let Some (resp) = resp { - socket.send_to (&resp.to_vec ()?, remote_addr).unwrap (); + socket.send_to (&Message::many_to_vec (&resp)?, remote_addr).unwrap (); } } } -fn recv_msg_from (socket: &UdpSocket) -> Result <(Message, SocketAddr), AppError> +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)?; buf.truncate (bytes_recved); - let msg = Message::from_slice (&buf)?; + let msgs = Message::from_slice2 (&buf)?; - Ok ((msg, remote_addr)) + Ok ((msgs, remote_addr)) } diff --git a/src/message.rs b/src/message.rs index 5e1aea9..3acdeb2 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1,5 +1,8 @@ use std::{ - io::Cursor, + io::{ + Cursor, + Write, + }, }; use crate::tlv; @@ -9,32 +12,65 @@ use thiserror::Error; const MAGIC_NUMBER: [u8; 4] = [0x9a, 0x4a, 0x43, 0x81]; pub const PACKET_SIZE: usize = 1024; -#[derive (Debug)] +type Mac = [u8; 6]; + +#[derive (Debug, PartialEq)] pub enum Message { - Request { + // 1 + Request1 { idem_id: [u8; 8], mac: Option <[u8; 6]> }, - Response (Option <[u8; 6]>), + // 2 + Response1 (Option <[u8; 6]>), + // 3 + Response2 (Response2), +} + +#[derive (Debug, PartialEq)] +pub struct Response2 { + pub idem_id: [u8; 8], + pub nickname: String, } #[derive (Debug, Error)] pub enum MessageError { #[error (transparent)] Io (#[from] std::io::Error), + #[error ("Length prefix too long")] + LengthPrefixTooLong ((usize, usize)), #[error (transparent)] Tlv (#[from] tlv::TlvError), + #[error (transparent)] + TryFromInt (#[from] std::num::TryFromIntError), #[error ("Unknown type")] UnknownType, + #[error (transparent)] + FromUtf8 (#[from] std::string::FromUtf8Error), +} + +#[derive (Default)] +struct DummyWriter { + position: usize, +} + +impl Write for DummyWriter { + fn flush (&mut self) -> std::io::Result <()> { + Ok (()) + } + + fn write (&mut self, buf: &[u8]) -> std::io::Result { + self.position += buf.len (); + Ok (buf.len ()) + } } impl Message { - pub fn write (&self, w: &mut W) -> Result <(), std::io::Error> + pub fn write (&self, w: &mut Cursor ) -> Result <(), std::io::Error> + where Cursor : Write { - w.write_all (&MAGIC_NUMBER)?; - match self { - Self::Request { + Self::Request1 { idem_id, mac, }=> { @@ -42,16 +78,41 @@ impl Message { w.write_all (&idem_id[..])?; Self::write_mac_opt (w, *mac)?; }, - Self::Response (mac) => { + Self::Response1 (mac) => { w.write_all (&[2])?; Self::write_mac_opt (w, *mac)?; }, + Self::Response2 (x) => { + w.write_all (&[3])?; + // Measure length with dummy writes + // This is dumb, I'm just messing around to see if I can do + // this without allocating. + let mut dummy_writer = DummyWriter::default (); + + Self::write_response_2 (&mut dummy_writer, x)?; + + // Write length and real params to real output + let len = u32::try_from (dummy_writer.position).unwrap (); + w.write_all (&len.to_le_bytes ())?; + Self::write_response_2 (w, x)?; + }, } Ok (()) } - fn write_mac_opt (w: &mut W, mac: Option <[u8; 6]>) -> Result <(), std::io::Error> + fn write_response_2 (w: &mut W, params: &Response2) + -> Result <(), std::io::Error> + { + w.write_all (¶ms.idem_id)?; + let nickname = params.nickname.as_bytes (); + let nickname_len = u32::try_from (nickname.len ()).unwrap (); + w.write_all (&nickname_len.to_le_bytes ())?; + w.write_all (&nickname)?; + Ok (()) + } + + fn write_mac_opt (w: &mut W, mac: Option <[u8; 6]>) -> Result <(), std::io::Error> { match mac { Some (mac) => { @@ -65,12 +126,21 @@ impl Message { pub fn to_vec (&self) -> Result , tlv::TlvError> { let mut cursor = Cursor::new (Vec::with_capacity (PACKET_SIZE)); + cursor.write_all (&MAGIC_NUMBER)?; self.write (&mut cursor)?; Ok (cursor.into_inner ()) } - pub fn read (r: &mut R) -> Result { - tlv::Reader::expect (r, &MAGIC_NUMBER)?; + pub fn many_to_vec (msgs: &[Self]) -> Result , tlv::TlvError> { + let mut cursor = Cursor::new (Vec::with_capacity (PACKET_SIZE)); + cursor.write_all (&MAGIC_NUMBER)?; + for msg in msgs { + msg.write (&mut cursor)?; + } + Ok (cursor.into_inner ()) + } + + fn read1 (r: &mut R) -> Result { let t = tlv::Reader::u8 (r)?; Ok (match t { @@ -79,14 +149,62 @@ impl Message { r.read_exact (&mut idem_id)?; let mac = Self::read_mac_opt (r)?; - Self::Request { + Self::Request1 { idem_id, mac, } }, 2 => { let mac = Self::read_mac_opt (r)?; - Self::Response (mac) + Self::Response1 (mac) + }, + _ => return Err (MessageError::UnknownType), + }) + } + + fn read2 (r: &mut R) -> Result { + let t = tlv::Reader::u8 (r)?; + + Ok (match t { + 1 => { + let mut idem_id = [0u8; 8]; + r.read_exact (&mut idem_id)?; + + let mac = Self::read_mac_opt (r)?; + Self::Request1 { + idem_id, + mac, + } + }, + 2 => { + let mac = Self::read_mac_opt (r)?; + Self::Response1 (mac) + }, + 3 => { + let mut len = [0; 4]; + r.read_exact (&mut len)?; + let _len = len; + + let mut idem_id = [0; 8]; + r.read_exact (&mut idem_id)?; + + let mut nickname_len = [0; 4]; + r.read_exact (&mut nickname_len)?; + let nickname_len = u32::from_le_bytes (nickname_len); + let nickname_len = usize::try_from (nickname_len)?; + + if nickname_len > 64 { + return Err (MessageError::LengthPrefixTooLong ((64, nickname_len))); + } + + let mut nickname = vec! [0u8; nickname_len]; + r.read_exact (&mut nickname)?; + let nickname = String::from_utf8 (nickname)?; + + Self::Response2 (Response2 { + idem_id, + nickname, + }) }, _ => return Err (MessageError::UnknownType), }) @@ -105,9 +223,23 @@ impl Message { }) } - pub fn from_slice (buf: &[u8]) -> Result { + pub fn from_slice1 (buf: &[u8]) -> Result { let mut cursor = Cursor::new (buf); - Self::read (&mut cursor) + tlv::Reader::expect (&mut cursor, &MAGIC_NUMBER)?; + Self::read1 (&mut cursor) + } + + pub fn from_slice2 (buf: &[u8]) -> Result , MessageError> { + let mut cursor = Cursor::new (buf); + tlv::Reader::expect (&mut cursor, &MAGIC_NUMBER)?; + + let mut msgs = Vec::with_capacity (2); + + while cursor.position () < u64::try_from (buf.len ())? { + let msg = Self::read2 (&mut cursor)?; + msgs.push (msg); + } + Ok (msgs) } } @@ -116,10 +248,79 @@ mod test { use super::*; #[test] - fn test_1 () { + fn test_write_2 () { for (input, expected) in [ ( - Message::Response (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), + vec! [ + Message::Request1 { + idem_id: [1, 2, 3, 4, 5, 6, 7, 8,], + mac: None, + }, + ], + vec! [ + 154, 74, 67, 129, + // Request tag + 1, + // Idem ID + 1, 2, 3, 4, 5, 6, 7, 8, + // MAC is None + 0, + ], + ), + ( + vec! [ + Message::Response1 (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), + Message::Response2 (Response2 { + idem_id: [1, 2, 3, 4, 5, 6, 7, 8,], + nickname: ":V".to_string (), + }), + ], + vec! [ + // Magic number for LookAround packets + 154, 74, 67, 129, + // Response1 tag + 2, + // MAC is Some + 1, + // MAC + 17, 34, 51, 68, 85, 102, + // Response2 tag + 3, + // Length prefix + 14, 0, 0, 0, + // Idem ID + 1, 2, 3, 4, 5, 6, 7, 8, + // Length-prefixed string + 2, 0, 0, 0, + 58, 86, + ], + ), + ] { + let actual = Message::many_to_vec (&input).unwrap (); + assert_eq! (actual, expected, "{:?}", input); + } + } + + #[test] + fn test_write_1 () { + for (input, expected) in [ + ( + Message::Request1 { + idem_id: [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08], + mac: None, + }, + vec! [ + 154, 74, 67, 129, + // Request tag + 1, + // Idem ID + 1, 2, 3, 4, 5, 6, 7, 8, + // MAC is None + 0, + ], + ), + ( + Message::Response1 (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), vec! [ // Magic number for LookAround packets 154, 74, 67, 129, @@ -131,9 +332,123 @@ mod test { 17, 34, 51, 68, 85, 102, ], ), + ( + Message::Response1 (None), + vec! [ + // Magic number for LookAround packets + 154, 74, 67, 129, + // Response tag + 2, + // MAC is None + 0, + ], + ), ].into_iter () { let actual = input.to_vec ().unwrap (); assert_eq! (actual, expected, "{:?}", input); } } + + #[test] + fn test_read_2 () { + for input in [ + vec! [ + Message::Request1 { + idem_id: [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08], + mac: None, + }, + ], + vec! [ + Message::Response1 (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), + ], + vec! [ + Message::Response1 (None), + ], + vec! [ + Message::Response1 (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), + Message::Response2 (Response2 { + idem_id: [1, 2, 3, 4, 5, 6, 7, 8,], + nickname: ":V".to_string (), + }), + ], + ].into_iter () { + let encoded = Message::many_to_vec (&input).unwrap (); + let decoded = Message::from_slice2 (&encoded).unwrap (); + assert_eq! (input, decoded); + } + } + + #[test] + fn test_read_1 () { + for input in [ + Message::Request1 { + idem_id: [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08], + mac: None, + }, + Message::Response1 (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), + Message::Response1 (None), + ].into_iter () { + let encoded = input.to_vec ().unwrap (); + let decoded = Message::from_slice1 (&encoded).unwrap (); + assert_eq! (input, decoded); + } + + for (expected, input) in [ + ( + Message::Request1 { + idem_id: [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08], + mac: None, + }, + vec! [ + 154, 74, 67, 129, + // Request tag + 1, + // Idem ID + 1, 2, 3, 4, 5, 6, 7, 8, + // MAC is None + 0, + ], + ), + ( + Message::Response1 (Some ([0x11, 0x22, 0x33, 0x44, 0x55, 0x66])), + vec! [ + // Magic number for LookAround packets + 154, 74, 67, 129, + // Response tag + 2, + // MAC is Some + 1, + // MAC + 17, 34, 51, 68, 85, 102, + ], + ), + ( + Message::Response1 (None), + vec! [ + // Magic number for LookAround packets + 154, 74, 67, 129, + // Response tag + 2, + // MAC is None + 0, + ], + ), + ( + Message::Response1 (None), + vec! [ + // Magic number for LookAround packets + 154, 74, 67, 129, + // Response tag + 2, + // MAC is None + 0, + // New tag that older versions will just ignore + 255, + ], + ), + ].into_iter () { + let actual = Message::from_slice1 (&input).unwrap (); + assert_eq! (actual, expected, "{:?}", actual); + } + } }