misc: refactor types

src/main.rs

@@ -1,19 +1,22 @@
 mod types;
 
-use crate::types::{Color, MatchMaker, Role, AI};
+use crate::types::{Color, Match};
 use futures_util::{SinkExt, StreamExt};
 use std::collections::HashMap;
 use std::net::SocketAddr;
-use std::num::ParseIntError;
-use std::str::FromStr;
 use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
+use tokio::sync::mpsc::error::SendError;
+use tokio::sync::mpsc::UnboundedSender;
 use tokio::sync::RwLock;
 use tokio_tungstenite::{accept_async, tungstenite::Message};
 use tracing::{error, info, warn};
 use types::Client;
 
 type Clients = Arc<RwLock<HashMap<SocketAddr, Client<'static>>>>;
+type Observers = Arc<RwLock<HashMap<SocketAddr, UnboundedSender<Message>>>>;
+type Matches = Arc<RwLock<HashMap<u32, Match<'static>>>>;
+
 
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
@@ -24,15 +27,17 @@ async fn main() -> Result<(), anyhow::Error> {
     let listener = TcpListener::bind(&addr).await?;
     info!("WebSocket server listening on: {}", addr);
 
-    let mut clients: Clients = Arc::new(RwLock::new(HashMap::new()));
+    let clients: Clients = Arc::new(RwLock::new(HashMap::new()));
-    let mut match_maker: Arc<RwLock<MatchMaker>> = Arc::new(RwLock::new(MatchMaker::new()));
+	let observers: Observers = Arc::new(RwLock::new(HashMap::new()));
+    let matches: Matches = Arc::new(RwLock::new(HashMap::new()));
 
     while let Ok((stream, addr)) = listener.accept().await {
         tokio::spawn(handle_connection(
             stream,
             addr,
             clients.clone(),
-            match_maker.clone(),
+			observers.clone(),
+            matches.clone(),
         ));
     }
 
@@ -43,7 +48,8 @@ async fn handle_connection(
     stream: TcpStream,
     addr: SocketAddr,
     clients: Clients,
-    match_maker: Arc<RwLock<MatchMaker>>,
+	observers: Observers,
+    matches: Matches,
 ) -> Result<(), anyhow::Error> {
     info!("New WebSocket connection from: {}", addr);
 
@@ -53,10 +59,10 @@ async fn handle_connection(
 
     // Store the client
     {
-        clients
+        observers
             .write()
             .await
-            .insert(addr, Client::new(String::new(), Role::Observer, tx));
+            .insert(addr, tx.clone());
     }
 
     // Spawn task to handle outgoing messages
@@ -78,141 +84,83 @@ async fn handle_connection(
                     let requested_username = text.split(":").collect::<Vec<&str>>()[1].to_string();
                     let mut is_taken = false;
                     for client in clients.read().await.values() {
-                        if let Some(username) = &client.username {
+                        if requested_username == client.username {
-                            if *username == requested_username {
+							let _ = send(&tx, &format!("ERROR:INVALID:ID:{}", requested_username));
-								let _ = clients
-									.read()
-									.await
-									.get(&addr)
-									.unwrap()
-									.send(&format!("ERROR:INVALID:ID:{}", requested_username));
 							is_taken = true;
 							break;
                         }
                     }
-                    }
 
                     if is_taken { continue; }
 
                     // not taken
-                    clients.write().await.get_mut(&addr).unwrap().role = Role::Player;
+					observers.write().await.remove(&addr);
-                    clients.write().await.get_mut(&addr).unwrap().username = Some(requested_username);
+                    clients.write().await.insert(addr.to_string().parse()?, Client::new(requested_username, tx.clone()));
 
-					let _ = clients
+					let _ = send(&tx, "CONNECT:ACK");
-						.read()
-						.await
-						.get(&addr)
-						.unwrap()
-						.send("CONNECT:ACK");
                 }
 				else if text.starts_with("READY") {
-                    if clients.read().await.get(&addr).unwrap().role != Role::Player {
+                    if clients.read().await.get(&addr).is_none() {
-                        let _ = clients
+                        let _ = send(&tx, "ERROR:INVALID");
-                            .read()
-                            .await
-                            .get(&addr)
-                            .unwrap()
-                            .send("ERROR:INVALID");
                         continue;
                     }
 
-					let mut already_ready = false;
+					if clients.read().await.get(&addr).unwrap().ready {
-					for ready_player in &match_maker.read().await.ready_players {
+						let _ = send(&tx, "ERROR:INVALID");
-						if ready_player.username.eq(clients.read().await.get(&addr).unwrap().username.as_ref().unwrap()) {
+						continue;
-							let _ = clients
-								.read()
-								.await
-								.get(&addr)
-								.unwrap()
-								.send("ERROR:INVALID");
-							already_ready = true;
-							break;
-						}
 					}
 
-					if already_ready { continue; }
+					clients.write().await.get_mut(&addr).unwrap().ready = true;
 
-                    match_maker.write().await.ready_players.push(AI::new(
+					let _ = send(&tx,"READY:ACK");
-                        clients
-                            .read()
-                            .await
-                            .get(&addr)
-                            .unwrap()
-                            .username
-                            .as_ref()
-                            .unwrap(),
-                        Color::None,
-						addr.to_string()
-                    ));
-
-					let _ = clients
-						.read()
-						.await
-						.get(&addr)
-						.unwrap()
-						.send("READY:ACK");
                 }
 				else if text.starts_with("PLAY") {
 					let read = clients.read().await;
-					let client = read.get(&addr).unwrap();
+					let client = read.get(&addr);
 
 					// Check if client is valid
-					if client.role != Role::Player || client.current_match.is_none() {
+					if client.is_none() || client.unwrap().current_match.is_none() {
-						let _ = clients
+						let _ = send(&tx, "ERROR:INVALID:MOVE");
-							.read()
-							.await
-							.get(&addr)
-							.unwrap()
-							.send("ERROR:INVALID:MOVE");
 						continue;
 					}
 
-					let current_match = client.current_match.unwrap();
+					let current_match = client.unwrap().current_match.unwrap();
-					let ai = if *client.username.as_ref().unwrap() == current_match.player1.username
+					let current_player = if client.unwrap().username == current_match.player1.username
-					{ &current_match.player1 } else { &current_match.player2 };
+					{ current_match.player1 } else { current_match.player2 };
-					let opponent = if *client.username.as_ref().unwrap() == current_match.player1.username
+					let opponent = if client.unwrap().username == current_match.player1.username
-					{ &current_match.player2 } else { &current_match.player1 };
+					{ current_match.player2 } else { current_match.player1 };
 
 					// Check if it's their move
 					if (current_match.ledger.is_empty() &&
-						current_match.first.username != *client.username.as_ref().unwrap()) ||
+						current_match.first.username != client.unwrap().username) ||
-						(client.current_match.unwrap().ledger.last().unwrap().0.username == *client.username.as_ref().unwrap())
+						(client.unwrap().current_match.unwrap().ledger.last().unwrap().0.username == client.unwrap().username)
 					{
-						let _ = clients
+						let _ = send(&tx, "ERROR:INVALID:MOVE");
-							.read()
-							.await
-							.get(&addr)
-							.unwrap()
-							.send("ERROR:INVALID:MOVE");
 						continue;
 					}
 
+					let column_parse = text.split(":").collect::<Vec<&str>>()[1].parse::<usize>();
+
 					// Check if valid move
-					if let Ok(column) = text.split(":").collect::<Vec<&str>>()[1].parse::<usize>() {
+					if let Ok(column) = column_parse {
 						if current_match.board[column][4] != Color::None {
-							let _ = clients
+							let _ = send(&tx, "ERROR:INVALID:MOVE");
-								.read()
-								.await
-								.get(&addr)
-								.unwrap()
-								.send("ERROR:INVALID:MOVE");
 							continue;
 						}
 
 						// Place it
 						for i in 0..current_match.board[column].len() {
 							if current_match.board[column][i] == Color::None {
-								match_maker.write().await.matches.get_mut(&current_match.id).unwrap().board[column][i] = ai.color.clone();
+								matches
+									.write()
+									.await
+									.get_mut(&current_match.id).unwrap()
+									.board[column][i] = current_player.color.clone();
 								break;
 							}
 						}
 					} else {
-						let _ = clients
+						let _ = send(&tx, "ERROR:INVALID:MOVE");
-							.read()
-							.await
-							.get(&addr)
-							.unwrap()
-							.send("ERROR:INVALID:MOVE");
 						continue;
 					}
 
@@ -256,70 +204,29 @@ async fn handle_connection(
 					};
 
 					if winner != Color::None {
-						if winner == ai.color {
+						if winner == current_player.color {
-							let _ = clients
+							let _ = send(&tx, "GAME:WINS");
-								.read()
+							let _ = send(&opponent.connection, "GAME:LOSS");
-								.await
-								.get(&addr)
-								.unwrap()
-								.send("GAME:WINS");
 
-							let _ = clients
-								.read()
-								.await
-								.get(&SocketAddr::from_str(&opponent.addr)?)
-								.unwrap()
-								.send("GAME:LOSS");
 						} else {
-							let _ = clients
+							let _ = send(&tx, "GAME:LOSS");
-								.read()
+							let _ = send(&opponent.connection, "GAME:WINS");
-								.await
-								.get(&addr)
-								.unwrap()
-								.send("GAME:LOSS");
 
-							let _ = clients
-								.read()
-								.await
-								.get(&SocketAddr::from_str(&opponent.addr)?)
-								.unwrap()
-								.send("GAME:WINS");
 						}
 					}
 					else if filled {
-						let _ = clients
+						let _ = send(&tx, "GAME:DRAW");
-							.read()
+						let _ = send(&opponent.connection, "GAME:DRAW");
-							.await
-							.get(&addr)
-							.unwrap()
-							.send("GAME:DRAW");
-
-						let _ = clients
-							.read()
-							.await
-							.get(&SocketAddr::from_str(&opponent.addr)?)
-							.unwrap()
-							.send("GAME:DRAW");
 					}
 					else {
-						let _ = clients
+						let _ = send(&opponent.connection, &format!("OPPONENT:{}", column_parse?));
-							.read()
-							.await
-							.get(&SocketAddr::from_str(&opponent.addr)?)
-							.unwrap()
-							.send(&format!("OPPONENT:{}", text.split(":").collect::<Vec<&str>>()[1].parse::<usize>()?));
 					}
 
 					// TODO: remove match from matchmaker
 					// TODO: broadcast moves to viewers
 				}
 				else {
-                    let _ = clients
+					let _ = send(&tx, "GAME:UNKNOWN");
-                        .read()
-                        .await
-                        .get(&addr)
-                        .unwrap()
-                        .send("ERROR:UNKNOWN");
 				}
             }
             Ok(Message::Close(_)) => {
@@ -327,12 +234,7 @@ async fn handle_connection(
                 break;
             }
             Ok(_) => {
-                let _ = clients
+				let _ = send(&tx, "GAME:UNKNOWN");
-                    .read()
-                    .await
-                    .get(&addr)
-                    .unwrap()
-                    .send("ERROR:UNKNOWN");
 			}
             Err(e) => {
                 error!("WebSocket error for {}: {}", addr, e);
@@ -353,11 +255,30 @@ async fn handle_connection(
     Ok(())
 }
 
-async fn broadcast_message(clients: &Clients, msg: Message) {
+async fn broadcast_message(observers: &Observers, msg: &str) {
-    let clients = clients.read().await;
+    let observers = observers.read().await;
-    for (_, client) in clients.iter() {
+    for (_, tx) in observers.iter() {
-        if client.role == Role::Admin || client.role == Role::Observer {
+		let _ = send(tx, msg);
-            client.connection.send(msg.clone()).ok();
 	}
 }
+
+async fn watch(matches: &Matches, new_match_id: u32, addr: SocketAddr) {
+	for a_match in &mut matches.write().await.values_mut() {
+		let mut found = false;
+		for i in 0..a_match.viewers.len() {
+			if a_match.viewers[i] == addr {
+				a_match.viewers.remove(i);
+				found = true;
+				break;
+			}
+		}
+
+		if found { break; }
+	}
+
+	matches.write().await.get_mut(&new_match_id).unwrap().viewers.push(addr);
+}
+
+fn send(tx: &UnboundedSender<Message>, text: &str) -> Result<(), SendError<Message>> {
+	tx.send(Message::text(text))
 }

src/types.rs

@@ -6,13 +6,6 @@ use tokio::sync::mpsc::error::SendError;
 use tokio::sync::mpsc::UnboundedSender;
 use tokio_tungstenite::tungstenite::Message;
 
-#[derive(PartialEq)]
-pub enum Role {
-    Admin,
-    Observer,
-    Player,
-}
-
 #[derive(PartialEq, Clone)]
 pub enum Color {
 	Red,
@@ -21,20 +14,20 @@ pub enum Color {
 }
 
 pub struct Client<'a> {
-    pub client_id: String,
+	pub username: String,
-    pub role: Role,
     pub connection: UnboundedSender<Message>,
-	pub username: Option<String>,
+	pub ready: bool,
+	pub color: Color,
 	pub current_match: Option<&'a Match<'static>>,
 }
 
 impl Client<'static> {
-	pub fn new(client_id: String, role: Role, connection: UnboundedSender<Message>) -> Client<'static> {
+	pub fn new(username: String, connection: UnboundedSender<Message>) -> Client<'static> {
 		Client {
-			client_id,
+			username,
-			role,
 			connection,
-			username: None,
+			ready: false,
+			color: Color::None,
 			current_match: None,
 		}
 	}
@@ -44,61 +37,19 @@ impl Client<'static> {
 	}
 }
 
-#[derive(Clone)]
-pub struct AI {
-	pub username: String,
-	pub color: Color,
-	pub ready: bool,
-	pub addr: String,
-}
-
-impl AI {
-	pub fn new(username: &str, color: Color, addr: String) -> AI {
-		AI { username: username.to_string(), color, ready: false, addr }
-	}
-}
-
 pub struct Match<'a> {
 	pub id: u32,
 	pub board: Vec<Vec<Color>>,
 	pub viewers: Vec<SocketAddr>,
-	pub ledger: Vec<(&'a AI, usize)>,
+	pub ledger: Vec<(&'a Client<'a>, usize)>,
-	pub first: AI,
+	pub first: &'a Client<'a>,
-	pub player1: AI,
+	pub player1: &'a Client<'a>,
-	pub player2: AI,
+	pub player2: &'a Client<'a>,
 }
 
-impl Match<'static> {
+impl<'a> Match<'a> {
-	pub fn new(id: u32, player1: AI, player2: AI) -> Match<'static> {
+	pub fn new(id: u32, player1: &'a Client, player2: &'a Client) -> Match<'a> {
-		let first = if rand::rng().random_range(0..=1) == 0 { player1.clone() } else { player2.clone() };
+		let first = if rand::rng().random_range(0..=1) == 0 { player1 } else { player2 };
 		Match { id, board: vec![vec![Color::None; 5]; 6], viewers: Vec::new(), ledger: Vec::new(), first, player1, player2 }
 	}
 }
-
-pub struct MatchMaker {
-	pub matches: HashMap<u32, Match<'static>>,
-	pub ready_players: Vec<AI>,
-}
-
-impl MatchMaker {
-	pub fn new() -> MatchMaker {
-		MatchMaker { matches: HashMap::new(), ready_players: Vec::new() }
-	}
-
-	pub fn watch(&mut self, viewer: SocketAddr, match_id: u32) {
-		for aMatch in &mut self.matches.values_mut() {
-			let mut found = false;
-			for i in 0..aMatch.viewers.len() {
-				if aMatch.viewers[i] == viewer {
-					aMatch.viewers.remove(i);
-					found = true;
-					break;
-				}
-			}
-
-			if found { break; }
-		}
-
-		self.matches.get_mut(&match_id).unwrap().viewers.push(viewer);
-	}
-}