1 files changed, 120 insertions, 16 deletions

diff --git a/src/lobby.rs b/src/lobby.rs
index ac6ab305414a..496f98bd7b6f 100644
--- a/src/lobby.rs
+++ b/src/lobby.rs
@@ -3,7 +3,7 @@
 use crate::{
     data::{Color, ColorPalette},
     users::MetaUser,
-    wire::{Lobby, LobbyId, LobbyUser},
+    wire::{Lobby, LobbyErr, LobbyId, LobbyUpdate, LobbyUser, User, UserId},
 };
 use async_std::sync::{Arc, RwLock};
 use std::{
@@ -14,15 +14,60 @@ use std::{
 /// A list of all the lobbies on the server
 pub struct LobbyList {
     max: AtomicUsize,
-    users: RwLock<BTreeMap<LobbyId, Arc<MetaLobby>>>,
+    lobbies: RwLock<BTreeMap<LobbyId, MetaLobby>>,
+}
+
+impl LobbyList {
+    pub fn new() -> Self {
+        Self {
+            max: 0.into(),
+            lobbies: Default::default(),
+        }
+    }
+
+    /// Create a new lobby
+    pub async fn create(&self, map: String) -> LobbyId {
+        let id = self.max.fetch_add(1, Ordering::Relaxed);
+        self.lobbies
+            .write()
+            .await
+            .insert(id, MetaLobby::create(id, map));
+        id
+    }
+
+    /// Remove a lobby by ID
+    pub async fn destroy(&self, id: LobbyId) -> Result<(), LobbyErr> {
+        self.consume(id).await.map(|_| ())
+    }
+
+    /// Remove and return the lobby
+    pub async fn consume(&self, id: LobbyId) -> Result<MetaLobby, LobbyErr> {
+        self.lobbies
+            .write()
+            .await
+            .remove(&id)
+            .map_or(Err(LobbyErr::NoSuchRoom), |l| Ok(l))
+    }
+
+    /// Get mutable access to a lobby
+    pub async fn get_mut<F, T>(&self, id: LobbyId, cb: F) -> Result<T, LobbyErr>
+    where
+        F: Fn(&mut MetaLobby) -> T,
+    {
+        self.lobbies
+            .write()
+            .await
+            .get_mut(&id)
+            .map_or(Err(LobbyErr::OtherError), |ref mut l| Ok(cb(l)))
+    }
 }
 
 /// Additional state held by the server
 ///
 /// The meta lobby will also sync updates to all connected users, when updates are made to the lobby
 pub struct MetaLobby {
-    palette: Vec<Color>,
-    inner: Lobby,
+    pub palette: Vec<Color>,
+    pub inner: Lobby,
 }
 
 impl MetaLobby {
@@ -38,21 +83,19 @@ impl MetaLobby {
         }
     }
 
-    pub fn join(&mut self, user: MetaUser) {
+    pub fn join(&mut self, user: &MetaUser) -> Lobby {
         let color = if &user.name == "spacekookie" {
             let color = Color::blue();
+            self.palette.without(&color);
 
-            let num = self.inner.players.len();
-            self.palette = Vec::without(&color);
-            let in_use = self.palette.split_off(num);
-
-            self.inner
+            if let Some(user) = self
+                .inner
                 .players
                 .iter_mut()
-                .zip(in_use.into_iter())
-                .for_each(|(user, color)| {
-                    user.color = color;
-                });
+                .find(|u| u.color == Color::blue())
+            {
+                user.color = self.palette.remove(0);
+            }
 
             color
         } else {
@@ -60,14 +103,17 @@ impl MetaLobby {
         };
 
         self.inner.players.push(LobbyUser {
+            admin: false,
             id: user.id,
-            name: user.name,
+            name: user.name.clone(),
             ready: false,
             color,
         });
+
+        self.inner.clone()
     }
 
-    pub fn leave(&mut self, user: MetaUser) {
+    pub fn leave(&mut self, user: &MetaUser) {
         let (pos, user) = self
             .inner
             .players
@@ -84,4 +130,62 @@ impl MetaLobby {
         self.palette.remix(user.color);
         self.inner.players.remove(pos);
     }
+
+    /// Check if a user is even present in a lobby
+    ///
+    /// Perform this prerequisite check before making other user-specific changes to the lobby
+    pub fn in_lobby(&self, user: UserId) -> bool {
+        self.inner
+            .players
+            .iter()
+            .find(|u| u.id == user)
+            .map(|_| true)
+            .unwrap_or(false)
+    }
+
+    /// Set the ready state for a user
+    pub fn ready(&mut self, user: User, ready: bool) -> LobbyUpdate {
+        if let Some(user) = self
+            .inner
+            .players
+            .iter_mut()
+            .find(|u| u.id == user.id)
+            .as_mut()
+        {
+            user.ready = ready;
+        }
+
+        LobbyUpdate::Ready(
+            self.inner
+                .players
+                .iter()
+                .filter_map(|u| if u.ready { Some(u.id) } else { None })
+                .collect(),
+        )
+    }
+
+    /// Try to start a game, if the user can and everybody is ready
+    pub fn start(&mut self, user: UserId) -> Result<(), LobbyErr> {
+        if let Some(_) = self
+            .inner
+            .players
+            .iter()
+            .filter(|u| u.admin)
+            .find(|u| u.id == user)
+        {
+            return Err(LobbyErr::NotAuthorized);
+        };
+
+        match self
+            .inner
+            .players
+            .iter()
+            .filter(|u| !u.ready)
+            .collect::<Vec<_>>()
+            .len()
+        {
+            0 => Err(LobbyErr::NotAllReady),
+            _ => Ok(()),
+        }
+    }
 }