From f28c9fee834c4b13217e6263dd5bc8f73e24fa5f Mon Sep 17 00:00:00 2001
From: Dane Johnson <daneallenjohnson@protonmail.com>
Date: Mon, 3 Oct 2022 16:49:32 -0500
Subject: [PATCH] refactoring

---
 src/main.rs      | 44 +++++++++++++++----------------------
 src/message.rs   | 56 +++++++++++++++++++++++++++++++++++++++++-------
 src/websocket.rs | 36 +++++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 35 deletions(-)
 create mode 100644 src/websocket.rs

diff --git a/src/main.rs b/src/main.rs
index 2054433..ad36f4e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,11 +5,10 @@ use std::collections::HashMap;
 
 mod message;
 use message::Message;
-
 mod code_generator;
 use code_generator::CodeGenerator;
-
-use tungstenite::protocol::Message as WsMessage;
+mod websocket;
+use websocket::WebsocketWrapper;
 
 fn main() {
     let code_generator = Arc::new(Mutex::new(CodeGenerator::default()));
@@ -23,30 +22,21 @@ fn main() {
         thread::spawn (move || {
             let mut ws = tungstenite::accept(stream.unwrap()).unwrap();
             println!("New client!");
-            ws.write_message(WsMessage::Text("HOSTJOIN:".to_string())).unwrap();
-            loop {
-                let message = ws.read_message();
-                println!("{:?}", message);
-                match message {
-                    Err(_) => break,
-                    Ok(WsMessage::Close(_)) => break,
-                    Ok(WsMessage::Text(msg)) => {
-                        let msg = Message::parse(msg.as_str().trim());
-                        match msg {
-                            Ok(msg) =>
-                                match msg.command {
-                                    "HOST" => {
-                                        let code = code_generator.lock().unwrap().generate();
-                                        let mut room = Room::default();
-                                        let player = Player { name: "Guest".to_string() };
-                                        room.players.push(player);
-                                        rooms.lock().unwrap().insert(code.clone(), room);
-                                        ws.write_message(WsMessage::Text(code)).unwrap();
-                                    }
-                                    _ => todo!(),
-                                }
-                            Err(_) => todo!(),
-                        }                       
+            let mut ws = WebsocketWrapper::new(ws);
+            ws.send(msg!("HOSTJOIN"));
+
+            let msg = ws.recv();
+            match msg {                
+                None => (),
+
+                Some(msg) => match msg.command.as_str() {
+                    "HOST" => {
+                        let code = code_generator.lock().unwrap().generate();
+                        let mut room = Room::default();
+                        let player = Player { name: "Guest".to_string() };
+                        room.players.push(player);
+                        rooms.lock().unwrap().insert(code.clone(), room);
+                        ws.send(msg!("CODE", code));
                     }
                     _ => unimplemented!(),
                 }
diff --git a/src/message.rs b/src/message.rs
index 87570e4..7b8075d 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -1,7 +1,7 @@
 #[derive(PartialEq, Debug)]
-pub struct Message<'a> {
-    pub command: &'a str,
-    pub args: Vec<&'a str>,
+pub struct Message {
+    pub command: String,
+    pub args: Vec<String>,
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -13,18 +13,21 @@ pub enum Error {
     UnknownCommand,
 }
 
-impl<'a> Message<'a> {
-    pub fn parse(text: &'a str) -> Result<Message<'a>> {
+impl Message {
+    pub fn new(command: String, args: Vec<String>) -> Self {
+	Message { command, args }
+    }
+    pub fn parse(text: String) -> Result<Message> {
 	let re = regex::Regex::new(r"^([A-Z_]+):\s*(.*)").unwrap();
-	match re.captures(text) {
+	match re.captures(text.as_str()) {
             Some(captures) => {
 		if captures.len() < 3 {
 		    Err(Error::BadParse)
 		} else {
-		    let command = captures.get(1).unwrap().as_str();
+		    let command = captures.get(1).unwrap().as_str().to_string();
 		    let args = captures.get(2).unwrap().as_str()
 			.split(',')
-			.map(|s| s.trim())
+			.map(|s| s.trim().to_string())
 			.collect();
                     Ok(Message { command, args })
                 }
@@ -34,6 +37,35 @@ impl<'a> Message<'a> {
     }
 }
 
+#[macro_export]
+macro_rules! msg {
+    ( $command:expr) => {
+	{
+	    let command = $command.to_string();
+	    let args = Vec::new();
+	    Message { command, args }
+	}
+    };
+    ( $command:expr, $( $arg:expr ),*) => {
+	{
+	    let command = $command.to_string();
+	    let mut args = Vec::new();
+	    $(
+		args.push($arg.to_string());
+	    )*
+	    Message { command, args }
+	}
+    };
+}
+	    
+
+impl std::fmt::Display for Message {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
+	write!(f, "{}: {}", self.command, self.args.as_slice().join(", "))
+    }
+}
+
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -47,4 +79,12 @@ mod test {
 	}, msg);
 	Ok(())
     }
+    #[test]
+    fn test_to_string() {
+	let msg = Message {
+	    command: "COMMAND",
+	    args: vec!["arg1", "arg2"],
+	};
+	assert_eq!(msg.to_string(), "COMMAND: arg1, arg2".to_string());
+    }
 }
diff --git a/src/websocket.rs b/src/websocket.rs
new file mode 100644
index 0000000..68d3c6d
--- /dev/null
+++ b/src/websocket.rs
@@ -0,0 +1,36 @@
+use std::net::TcpStream;
+
+use tungstenite::protocol::{ WebSocket, Message as WsMessage };
+
+use crate::message::Message;
+
+pub struct WebsocketWrapper {
+    websocket: WebSocket<TcpStream>,
+}
+
+impl WebsocketWrapper {
+    pub fn new(websocket: WebSocket<TcpStream>) -> Self {
+        WebsocketWrapper { websocket }
+    }
+    pub fn send(&mut self, msg: Message) {
+        self.websocket.write_message(WsMessage::Text(msg.to_string())).unwrap();
+    }
+    pub fn recv(&mut self) -> Option<Message> {
+        match self.websocket.read_message() {
+            Err(_) | Ok(WsMessage::Close(_)) => None,
+            Ok(WsMessage::Text(text)) => match Message::parse(text) {
+                Ok(msg) => Some(msg),
+                Err(_) => {
+                    self.websocket.write_message(WsMessage::Text("ERROR: bad_format".to_string()));
+                    self.recv()
+                }
+            }
+            _ => {
+                self.websocket.write_message(WsMessage::Text("ERROR: bad_command".to_string()));
+                self.recv()
+            }
+        }
+    }
+}
+        
+