refactoring

src/code_generator.rs

@@ -0,0 +1,42 @@
+use rand::RngCore;
+use sha2::{Sha256, Digest};
+
+pub struct CodeGenerator {
+    counter: u64,
+    salt: [u8; 32],
+}
+
+impl CodeGenerator {
+    pub fn generate(&mut self) -> String {
+        let count = self.counter;
+        self.counter += 1;
+
+        let mut hasher = Sha256::new();
+        hasher.update(self.salt);
+        hasher.update(count.to_be_bytes());
+
+        format!("{:x}", hasher.finalize())[..6].to_string()
+    }
+}
+
+impl Default for CodeGenerator {
+    fn default() -> Self {
+        let mut salt = [0; 32];
+        rand::thread_rng().fill_bytes(&mut salt);
+        
+        CodeGenerator {
+            counter: 0,
+            salt,
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    #[test]
+    fn test_generate() {
+        let code = CodeGenerator::default().generate();
+        assert_eq!(code.len(), 6);
+    }
+}

src/main.rs

@@ -5,12 +5,10 @@ use std::collections::HashMap;
 
 mod message;
 use message::Message;
-
-use rand::RngCore;
-
-use tungstenite::protocol::Message as WsMessage;
-
-use sha2::{Sha256, Digest};
+mod code_generator;
+use code_generator::CodeGenerator;
+mod websocket;
+use websocket::WebsocketWrapper;
 
 fn main() {
     let code_generator = Arc::new(Mutex::new(CodeGenerator::default()));
@@ -24,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!(),
                 }
@@ -63,33 +52,3 @@ struct Room {
 struct Player {
     name: String,
 }
-
-struct CodeGenerator {
-    counter: u64,
-    salt: [u8; 32],
-}
-
-impl CodeGenerator {
-    fn generate(&mut self) -> String {
-        let count = self.counter;
-        self.counter += 1;
-
-        let mut hasher = Sha256::new();
-        hasher.update(self.salt);
-        hasher.update(count.to_be_bytes());
-
-        format!("{:x}", hasher.finalize())[..6].to_string()
-    }
-}
-
-impl Default for CodeGenerator {
-    fn default() -> Self {
-        let mut salt = [0; 32];
-        rand::thread_rng().fill_bytes(&mut salt);
-        
-        CodeGenerator {
-            counter: 0,
-            salt,
-        }
-    }
-}

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::*;
@@ -45,7 +77,14 @@ mod test {
 	    command: "COMMAND",
 	    args: vec!["arg1", "arg2"],
 	}, 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());
+    }
 }

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()
+            }
+        }
+    }
+}
+        
+