src/code_generator.rs

@@ -1,42 +0,0 @@
-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,10 +5,12 @@ use std::collections::HashMap;
 
 mod message;
 use message::Message;
-mod code_generator;
+
-use code_generator::CodeGenerator;
+use rand::RngCore;
-mod websocket;
+
-use websocket::WebsocketWrapper;
+use tungstenite::protocol::Message as WsMessage;
+
+use sha2::{Sha256, Digest};
 
 fn main() {
     let code_generator = Arc::new(Mutex::new(CodeGenerator::default()));
@@ -22,21 +24,30 @@ fn main() {
         thread::spawn (move || {
             let mut ws = tungstenite::accept(stream.unwrap()).unwrap();
             println!("New client!");
-            let mut ws = WebsocketWrapper::new(ws);
+            ws.write_message(WsMessage::Text("HOSTJOIN:".to_string())).unwrap();
-            ws.send(msg!("HOSTJOIN"));
+            loop {
-
+                let message = ws.read_message();
-            let msg = ws.recv();
+                println!("{:?}", message);
+                match message {
+                    Err(_) => break,
+                    Ok(WsMessage::Close(_)) => break,
+                    Ok(WsMessage::Text(msg)) => {
+                        let msg = Message::parse(msg.as_str().trim());
                         match msg {
-                None => (),
+                            Ok(msg) =>
-
+                                match msg.command {
-                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));
+                                        ws.write_message(WsMessage::Text(code)).unwrap();
+                                    }
+                                    _ => todo!(),
+                                }
+                            Err(_) => todo!(),
+                        }                       
                     }
                     _ => unimplemented!(),
                 }
@@ -52,3 +63,33 @@ 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 {
+pub struct Message<'a> {
-    pub command: String,
+    pub command: &'a str,
-    pub args: Vec<String>,
+    pub args: Vec<&'a str>,
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -13,21 +13,18 @@ pub enum Error {
     UnknownCommand,
 }
 
-impl Message {
+impl<'a> Message<'a> {
-    pub fn new(command: String, args: Vec<String>) -> Self {
+    pub fn parse(text: &'a str) -> Result<Message<'a>> {
-	Message { command, args }
-    }
-    pub fn parse(text: String) -> Result<Message> {
 	let re = regex::Regex::new(r"^([A-Z_]+):\s*(.*)").unwrap();
-	match re.captures(text.as_str()) {
+	match re.captures(text) {
             Some(captures) => {
 		if captures.len() < 3 {
 		    Err(Error::BadParse)
 		} else {
-		    let command = captures.get(1).unwrap().as_str().to_string();
+		    let command = captures.get(1).unwrap().as_str();
 		    let args = captures.get(2).unwrap().as_str()
 			.split(',')
-			.map(|s| s.trim().to_string())
+			.map(|s| s.trim())
 			.collect();
                     Ok(Message { command, args })
                 }
@@ -37,35 +34,6 @@ impl Message {
     }
 }
 
-#[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::*;
@@ -77,14 +45,7 @@ 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

@@ -1,36 +0,0 @@
-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()
-            }
-        }
-    }
-}
-        
-