1 files changed, 54 insertions, 5 deletions

diff --git a/server/src/state.rs b/server/src/state.rs
index 97261fab..43ca29bd 100644
--- a/server/src/state.rs
+++ b/server/src/state.rs
@@ -15,18 +15,22 @@
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 */
-use crate::{data::DataIndex, game::Game};
+use crate::{data::DataIndex, game::Game, ConnectionID};
 use anyhow::{anyhow, bail, Result};
 use clap::{Parser, ValueEnum};
 use hurrycurry_protocol::{Message, PacketC, PacketS, PlayerID};
 use log::{debug, trace};
-use std::{collections::VecDeque, time::Duration};
+use std::{
+    collections::{HashMap, HashSet, VecDeque},
+    time::Duration,
+};
 use tokio::sync::broadcast::Sender;
 
 pub struct State {
     index: DataIndex,
     packet_out: VecDeque<PacketC>,
     tx: Sender<PacketC>,
+    connections: HashMap<ConnectionID, HashSet<PlayerID>>,
     pub game: Game,
 }
 
@@ -89,6 +93,7 @@ impl State {
             index,
             tx,
             packet_out,
+            connections: HashMap::new(),
         })
     }
 
@@ -110,14 +115,20 @@ impl State {
         }
         Ok(())
     }
-    pub async fn packet_in(&mut self, player: PlayerID, packet: PacketS) -> Result<Vec<PacketC>> {
+    pub async fn packet_in(&mut self, conn: ConnectionID, packet: PacketS) -> Result<Vec<PacketC>> {
+        if let Some(p) = get_packet_player(&packet) {
+            if !self.connections.entry(conn).or_default().contains(&p) {
+                bail!("Packet sent to player that is not owned by this connection.");
+            }
+        }
         let mut replies = Vec::new();
         match &packet {
             PacketS::Communicate {
                 message: Some(Message::Text(text)),
                 persist: false,
+                player,
             } if let Some(command) = text.strip_prefix("/") => {
-                match self.handle_command_parse(player, command).await {
+                match self.handle_command_parse(*player, command).await {
                     Ok(()) => return Ok(vec![]),
                     Err(e) => {
                         return Ok(vec![PacketC::ServerMessage {
@@ -126,10 +137,27 @@ impl State {
                     }
                 }
             }
+            PacketS::Leave { player } => {
+                self.connections.entry(conn).or_default().remove(player);
+            }
+            PacketS::Join { .. } => {
+                if self.connections.entry(conn).or_default().len() > 8 {
+                    bail!("Players per connection limit exceeded.")
+                }
+            }
             _ => (),
         }
         self.game
-            .packet_in(player, packet, &mut replies, &mut self.packet_out)?;
+            .packet_in(packet, &mut replies, &mut self.packet_out)?;
+
+        for p in &replies {
+            match p {
+                PacketC::Joined { id } => {
+                    self.connections.entry(conn).or_default().insert(*id);
+                }
+                _ => (),
+            }
+        }
 
         if self.game.count_chefs() <= 0 && !self.game.lobby {
             self.tx
@@ -146,6 +174,15 @@ impl State {
         Ok(replies)
     }
 
+    pub async fn disconnect(&mut self, conn: ConnectionID) {
+        if let Some(players) = self.connections.get(&conn) {
+            for player in players.to_owned() {
+                let _ = self.packet_in(conn, PacketS::Leave { player }).await;
+            }
+        }
+        self.connections.remove(&conn);
+    }
+
     async fn handle_command_parse(&mut self, player: PlayerID, command: &str) -> Result<()> {
         self.handle_command(
             player,
@@ -239,3 +276,15 @@ impl State {
         Ok(())
     }
 }
+
+fn get_packet_player(packet: &PacketS) -> Option<PlayerID> {
+    match packet {
+        PacketS::Join { .. } => None,
+        PacketS::Leave { player } => Some(*player),
+        PacketS::Movement { player, .. } => Some(*player),
+        PacketS::Interact { player, .. } => Some(*player),
+        PacketS::Communicate { player, .. } => Some(*player),
+        PacketS::ReplaceHand { player, .. } => Some(*player),
+        PacketS::ReplayTick { .. } => None,
+    }
+}