➕ add nicknames

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 <String>,
+}
+
 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 {
+		let mut resp = ServerResponse {
-			Message::Response (mac) => mac,
+			mac: None,
-			_ => continue,
+			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 {
+	for (ip, resp) in peers.into_iter () {
-		match mac {
+		let mac = match resp.mac {
-			Some (mac) => println! ("{} = {}", MacAddress::new (*mac), ip.ip ()),
+			None => {
-			None => println! ("<Unknown> = {}", ip),
+				println! ("<Unknown> = {}", 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 <I: Iterator <Item=String>> (args: I) -> Result <(), AppError> 
+fn server <I: Iterator <Item=String>> (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 <I: Iterator <Item=String>> (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 <I: Iterator <Item=String>> (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 <Message>, 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))
 }

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 <usize> {
+		self.position += buf.len ();
+		Ok (buf.len ())
+	}
 }
 
 impl Message {
-	pub fn write <W: std::io::Write> (&self, w: &mut W) -> Result <(), std::io::Error> 
+	pub fn write <T> (&self, w: &mut Cursor <T>) -> Result <(), std::io::Error> 
+	where Cursor <T>: 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: std::io::Write> (w: &mut W, mac: Option <[u8; 6]>) -> Result <(), std::io::Error>
+	fn write_response_2 <W: Write> (w: &mut W, params: &Response2) 
+	-> Result <(), std::io::Error>
+	{
+		w.write_all (&params.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: Write> (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 <Vec <u8>, 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: std::io::Read> (r: &mut R) -> Result <Self, MessageError> {
+	pub fn many_to_vec (msgs: &[Self]) -> Result <Vec <u8>, tlv::TlvError> {
-		tlv::Reader::expect (r, &MAGIC_NUMBER)?;
+		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: std::io::Read> (r: &mut R) -> Result <Self, MessageError> {
 		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: std::io::Read> (r: &mut R) -> Result <Self, MessageError> {
+		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 <Self, MessageError> {
+	pub fn from_slice1 (buf: &[u8]) -> Result <Self, MessageError> {
 		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 <Vec <Self>, 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);
+		}
+	}
 }