From eaca6155100b9071978a85f092b50eac19c7074c Mon Sep 17 00:00:00 2001
From: _ <_@_>
Date: Sun, 5 Dec 2021 19:58:18 +0000
Subject: [PATCH] working on TLV-formatted messages

---
 src/main.rs    |  60 ++++++++++++++++++++++------
 src/message.rs | 106 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/tlv.rs     |  30 +++++++++++---
 3 files changed, 178 insertions(+), 18 deletions(-)
 create mode 100644 src/message.rs

diff --git a/src/main.rs b/src/main.rs
index 09a9358..2cc8cbc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,6 +2,7 @@ use std::{
 	env,
 	net::{
 		Ipv4Addr,
+		SocketAddr,
 		SocketAddrV4,
 		UdpSocket,
 	},
@@ -10,12 +11,26 @@ use std::{
 use mac_address::get_mac_address;
 use thiserror::Error;
 
+mod message;
 mod tlv;
 
+use message::{
+	PACKET_SIZE,
+	Message,
+};
+
 #[derive (Debug, Error)]
 enum AppError {
 	#[error (transparent)]
 	CliArgs (#[from] CliArgError),
+	#[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),
 }
 
 #[derive (Debug, Error)]
@@ -58,10 +73,10 @@ fn main () -> Result <(), AppError> {
 	}
 	
 	match args.next ().as_ref ().map (|s| &s[..]) {
-		None => Err (CliArgError::MissingSubcommand)?,
+		None => return Err (CliArgError::MissingSubcommand.into ()),
 		Some ("client") => client ()?,
 		Some ("server") => server ()?,
-		Some (x) => Err (CliArgError::UnknownSubcommand (x.to_string ()))?,
+		Some (x) => return Err (CliArgError::UnknownSubcommand (x.to_string ()).into ()),
 	}
 	
 	Ok (())
@@ -73,30 +88,51 @@ fn client () -> Result <(), AppError> {
 	
 	socket.join_multicast_v4 (&params.multicast_addr, &([0u8, 0, 0, 0].into ())).unwrap ();
 	
-	socket.send_to ("hi there".as_bytes (), (params.multicast_addr, params.server_port)).unwrap ();
+	let msg = Message::Request (None).to_vec ()?;
+	
+	socket.send_to (&msg, (params.multicast_addr, params.server_port)).unwrap ();
+	
+	let (resp, remote_addr) = recv_msg_from (&socket)?;
 	
-	let mut buf = vec! [0u8; 4096];
-	let (bytes_recved, remote_addr) = socket.recv_from (&mut buf).unwrap ();
-	buf.truncate (bytes_recved);
-	let _buf = buf;
 	dbg! (remote_addr);
 	
 	Ok (())
 }
 
 fn server () -> Result <(), AppError> {
+	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 params = CommonParams::default ();
 	let socket = UdpSocket::bind (SocketAddrV4::new (Ipv4Addr::UNSPECIFIED, params.server_port)).unwrap ();
 	
 	socket.join_multicast_v4 (&params.multicast_addr, &([0u8, 0, 0, 0].into ())).unwrap ();
 	
-	let mut buf = vec! [0u8; 4096];
-	let (bytes_recved, remote_addr) = socket.recv_from (&mut buf).unwrap ();
-	buf.truncate (bytes_recved);
-	let _buf = buf;
+	let (req, remote_addr) = recv_msg_from (&socket)?;
 	dbg! (remote_addr);
 	
-	socket.send_to ("hi there".as_bytes (), remote_addr).unwrap ();
+	let resp = match req {
+		Message::Request (None) => {
+			Some (Message::Response (our_mac))
+		},
+		_ => None,
+	};
+	
+	if let Some (resp) = resp {
+		socket.send_to (&resp.to_vec ()?, remote_addr).unwrap ();
+	}
 	
 	Ok (())
 }
+
+fn recv_msg_from (socket: &UdpSocket) -> Result <(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)?;
+	
+	Ok ((msg, remote_addr))
+}
diff --git a/src/message.rs b/src/message.rs
new file mode 100644
index 0000000..99bf52d
--- /dev/null
+++ b/src/message.rs
@@ -0,0 +1,106 @@
+use std::{
+	io::Cursor,
+};
+
+use crate::tlv;
+
+use thiserror::Error;
+
+const MAGIC_NUMBER: [u8; 4] = [0x9a, 0x4a, 0x43, 0x81];
+pub const PACKET_SIZE: usize = 1024;
+
+pub enum Message {
+	Request (Option <[u8; 6]>),
+	Response (Option <[u8; 6]>),
+}
+
+#[derive (Debug, Error)]
+pub enum MessageError {
+	#[error (transparent)]
+	Io (#[from] std::io::Error),
+	#[error (transparent)]
+	Tlv (#[from] tlv::TlvError),
+	#[error ("Unknown type")]
+	UnknownType,
+}
+
+impl Message {
+	pub fn write <W: std::io::Write> (&self, w: &mut W) -> Result <(), std::io::Error> 
+	{
+		w.write_all (&MAGIC_NUMBER)?;
+		
+		match self {
+			Self::Request (mac) => {
+				w.write_all (&[1])?;
+				Self::write_mac_opt (w, *mac)?;
+			},
+			Self::Response (mac) => {
+				w.write_all (&[2])?;
+				Self::write_mac_opt (w, *mac)?;
+			},
+		}
+		
+		Ok (())
+	}
+	
+	fn write_mac_opt <W: std::io::Write> (w: &mut W, mac: Option <[u8; 6]>) -> Result <(), std::io::Error>
+	{
+		match mac {
+			Some (mac) => {
+				w.write_all (&[1])?;
+				w.write_all (&mac[..])?;
+			},
+			None => w.write_all (&[0])?,
+		}
+		Ok (())
+	}
+	
+	pub fn to_vec (&self) -> Result <Vec <u8>, tlv::TlvError> {
+		let mut cursor = Cursor::new (Vec::with_capacity (PACKET_SIZE));
+		self.write (&mut cursor)?;
+		Ok (cursor.into_inner ())
+	}
+	
+	pub fn read <R: std::io::Read> (r: &mut R) -> Result <Self, MessageError> {
+		tlv::Reader::expect (r, &MAGIC_NUMBER)?;
+		let t = tlv::Reader::u8 (r)?;
+		
+		Ok (match t {
+			1 => {
+				let mac = Self::read_mac_opt (r)?;
+				Self::Request (mac)
+			},
+			2 => {
+				let mac = Self::read_mac_opt (r)?;
+				Self::Response (mac)
+			},
+			_ => return Err (MessageError::UnknownType),
+		})
+	}
+	
+	fn read_mac_opt <R: std::io::Read> (r: &mut R) 
+	-> Result <Option <[u8; 6]>, std::io::Error> 
+	{
+		Ok (if tlv::Reader::u8 (r)? == 1 {
+			let mut mac = [0u8; 6];
+			r.read_exact (&mut mac)?;
+			Some (mac)
+		}
+		else {
+			None
+		})
+	}
+	
+	pub fn from_slice (buf: &[u8]) -> Result <Self, MessageError> {
+		let mut cursor = Cursor::new (buf);
+		Self::read (&mut cursor)
+	}
+}
+
+#[cfg (test)]
+mod test {
+	#[test]
+	fn test_1 () {
+		
+	}
+}
diff --git a/src/tlv.rs b/src/tlv.rs
index beac0a5..a083dd9 100644
--- a/src/tlv.rs
+++ b/src/tlv.rs
@@ -10,11 +10,13 @@ pub enum TlvError {
 	CallerBufferTooSmall,
 	#[error (transparent)]
 	Io (#[from] std::io::Error),
+	#[error ("Actual bytes didn't match expected bytes")]
+	NotExpected,
 	#[error (transparent)]
 	TryFromInt (#[from] std::num::TryFromIntError),
 }
 
-struct Writer <W> {
+pub struct Writer <W> {
 	_x: std::marker::PhantomData <W>,
 }
 
@@ -24,9 +26,9 @@ impl <W: std::io::Write> Writer <W> {
 		Ok (())
 	}
 	
-	fn lv_bytes (w: &mut W, b: &[u8]) -> Result <()> {
+	pub fn lv_bytes (w: &mut W, b: &[u8]) -> Result <()> {
 		if b.len () > 2_000_000_000 {
-			Err (TlvError::BufferTooBig)?;
+			return Err (TlvError::BufferTooBig);
 		}
 		
 		let l = u32::try_from (b.len ())?;
@@ -38,11 +40,20 @@ impl <W: std::io::Write> Writer <W> {
 	}
 }
 
-struct Reader <R> {
+pub struct Reader <R> {
 	_x: std::marker::PhantomData <R>,
 }
 
 impl <R: std::io::Read> Reader <R> {
+	pub fn expect (r: &mut R, expected: &[u8]) -> Result <()> {
+		let mut actual = vec! [0u8; expected.len ()];
+		r.read_exact (&mut actual)?;
+		if actual != expected {
+			return Err (TlvError::NotExpected);
+		}
+		Ok (())
+	}
+	
 	fn length (r: &mut R) -> Result <u32> {
 		let mut buf = [0; 4];
 		r.read_exact (&mut buf)?;
@@ -50,16 +61,23 @@ impl <R: std::io::Read> Reader <R> {
 		Ok (u32::from_le_bytes (buf))
 	}
 	
-	fn lv_bytes (r: &mut R, buf: &mut [u8]) -> Result <u32> {
+	pub fn lv_bytes (r: &mut R, buf: &mut [u8]) -> Result <u32> {
 		let l = Self::length (r)?;
 		if usize::try_from (l)? > buf.len () {
-			Err (TlvError::CallerBufferTooSmall)?;
+			return Err (TlvError::CallerBufferTooSmall);
 		}
 		
 		r.read_exact (&mut buf [0..usize::try_from (l)?])?;
 		
 		Ok (l)
 	}
+	
+	pub fn u8 (r: &mut R) -> std::io::Result <u8> {
+		let mut buf = [0];
+		r.read_exact (&mut buf)?;
+		
+		Ok (buf [0])
+	}
 }
 
 #[cfg (test)]