From 09c5fe3d8b57d9e305f14045d776dd10b8f6ff51 Mon Sep 17 00:00:00 2001
From: Dane Johnson <daneallenjohnson@protonmail.com>
Date: Thu, 29 Sep 2022 15:28:38 -0500
Subject: [PATCH] Generate room codes, python test script

---
 Cargo.lock   | 12 ++++++++++++
 Cargo.toml   |  3 ++-
 src/main.rs  | 42 ++++++++++++++++++++++++++++++++++++++----
 test/test.py | 14 ++++++++++++++
 4 files changed, 66 insertions(+), 5 deletions(-)
 create mode 100644 test/test.py

diff --git a/Cargo.lock b/Cargo.lock
index 2152ce0..7730e9b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -162,6 +162,7 @@ dependencies = [
 name = "hexland-server"
 version = "0.1.0"
 dependencies = [
+ "sha2",
  "tungstenite",
 ]
 
@@ -425,6 +426,17 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "sha2"
+version = "0.10.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "82e6b795fe2e3b1e845bafcb27aa35405c4d47cdfc92af5fc8d3002f76cebdc0"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
 [[package]]
 name = "syn"
 version = "1.0.100"
diff --git a/Cargo.toml b/Cargo.toml
index eb676c0..2a437ba 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,4 +6,5 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-tungstenite = { version = "0.17.3", features = ["native-tls"] }
\ No newline at end of file
+tungstenite = { version = "0.17.3", features = ["native-tls"] }
+sha2 = "0.10.6"
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index d1684ef..aa94e4f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,23 +1,39 @@
 use std::net::TcpListener;
 use std::thread;
+use std::sync::{Arc, Mutex};
 
 use tungstenite::protocol::Message;
 
+use sha2::{Sha256, Digest};
+
 fn main() {
+    let code_generator = Arc::new(Mutex::new(CodeGenerator {
+        counter: 0,
+        salt: [0; 32], // TODO have this be randomized on startup
+    }));
     let server = TcpListener::bind("127.0.0.1:8080").unwrap();
     for stream in server.incoming() {
-        thread::spawn ( move || {
+        thread::spawn ( { let code_generator = Arc::clone(&code_generator); move || {
             let mut ws = tungstenite::accept(stream.unwrap()).unwrap();
             println!("New client!");
             ws.write_message(Message::Text("HOSTJOIN:".to_string())).unwrap();
             loop {
                 let message = ws.read_message();
                 println!("{:?}", message);
-                if message.is_err() {
-                    break
+                match message {
+                    Err(_) => break,
+                    Ok(Message::Text(msg)) =>
+                        match msg.as_str().trim() {
+                            "HOST:" => {
+                                ws.write_message(Message::Text(code_generator.lock().unwrap().generate())).unwrap();
+                            },
+                            _ => unimplemented!(),
+                        },
+                    _ => unimplemented!(),
+                    
                 }
             }
-        });
+        }});
     }
 }
 
@@ -28,3 +44,21 @@ 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()
+    }
+}
diff --git a/test/test.py b/test/test.py
new file mode 100644
index 0000000..5777945
--- /dev/null
+++ b/test/test.py
@@ -0,0 +1,14 @@
+import asyncio
+import websockets
+import re
+
+async def test():
+    async with websockets.connect("ws://localhost:8080") as ws:
+        response = await ws.recv()
+        assert(response == "HOSTJOIN:")
+        await ws.send("HOST:")
+        response = await ws.recv()
+        assert(re.match(r"(\d|[a-f]){6}", response))
+        
+asyncio.get_event_loop().run_until_complete(test())
+print("All tests passed")