2 files changed, 51 insertions, 9 deletions

diff --git a/server/src/main.rs b/server/src/main.rs
index 09e6186..38238b3 100644
--- a/server/src/main.rs
+++ b/server/src/main.rs
@@ -18,15 +18,15 @@ pub mod network;
 
 use anyhow::Result;
 use clap::Parser;
-use log::{debug, warn};
+use log::{debug, info, warn};
 use network::ServerNetwork;
 use std::{
-    collections::HashSet,
+    collections::{BTreeSet, HashMap, HashSet},
     marker::PhantomData,
     net::{IpAddr, SocketAddr},
 };
 use weareshared::{
-    packets::{Data, Packet, Resource},
+    packets::{Data, Object, Packet, Resource},
     resources::{Prefab, PrefabIndex},
     store::ResourceStore,
     tree::SceneTree,
@@ -46,6 +46,7 @@ struct State {
     prefab_index: PrefabIndex,
     prefab_index_pending: HashSet<Resource<Prefab>>,
     prefab_index_have: HashSet<Resource<Prefab>>,
+    conn_object_ownership: HashMap<u128, BTreeSet<Object>>,
 }
 
 fn main() -> Result<()> {
@@ -72,6 +73,7 @@ impl State {
             prefab_index: PrefabIndex::default(),
             prefab_index_pending: HashSet::new(),
             prefab_index_have: HashSet::new(),
+            conn_object_ownership: HashMap::new(),
         })
     }
     pub fn prime_client(&self, conn: u128, net: &ServerNetwork) -> Result<()> {
@@ -85,12 +87,18 @@ impl State {
         Ok(())
     }
     pub fn handle_packet(&mut self, conn: u128, packet: Packet, net: &ServerNetwork) -> Result<()> {
-        self.tree.packet(&packet);
         match packet {
             Packet::Connect(_) => {
                 self.prime_client(conn, net)?;
             }
-            Packet::Disconnect => {}
+            Packet::Disconnect => {
+                if let Some(owned) = self.conn_object_ownership.remove(&conn) {
+                    for o in owned {
+                        debug!("removing conn owned object {o}");
+                        self.tree.remove_recursive(o, &mut net.br());
+                    }
+                }
+            }
             Packet::RequestResource(resource) => {
                 if let Some(r) = self.store.get_raw(resource)? {
                     debug!("resource is cached");
@@ -121,18 +129,31 @@ impl State {
                 }
             }
             Packet::Add(object, resource) => {
+                self.conn_object_ownership
+                    .entry(conn)
+                    .or_default()
+                    .insert(object);
+                self.tree.add(object, resource.clone());
                 net.broadcast(Packet::Add(object, resource), true);
             }
             Packet::Remove(object) => {
+                self.conn_object_ownership
+                    .entry(conn)
+                    .or_default()
+                    .remove(&object);
+                self.tree.remove_reparent(object, &mut net.br());
                 net.broadcast(Packet::Remove(object), true);
             }
             Packet::Position(object, pos, rot) => {
+                self.tree.packet(&Packet::Position(object, pos, rot));
                 net.broadcast(Packet::Position(object, pos, rot), false);
             }
             Packet::Pose(object, vec) => {
+                self.tree.packet(&Packet::Pose(object, vec.clone()));
                 net.broadcast(Packet::Pose(object, vec), false);
             }
             Packet::Parent(parent, child) => {
+                self.tree.reparent(parent, child);
                 net.broadcast(Packet::Parent(parent, child), true);
             }
             Packet::Sound(object, vec) => {
@@ -150,6 +171,10 @@ impl State {
                     );
                 }
             }
+            Packet::Chat(object, message) => {
+                info!("[CHAT] {}", message.0);
+                net.broadcast(Packet::Chat(object, message), true);
+            }
         }
         Ok(())
     }
diff --git a/server/src/network.rs b/server/src/network.rs
index 5c6f173..e632e97 100644
--- a/server/src/network.rs
+++ b/server/src/network.rs
@@ -27,7 +27,7 @@ use std::{
     thread::spawn,
     time::Instant,
 };
-use weareshared::{helper::ReadWrite, packets::Packet};
+use weareshared::{helper::ReadWrite, packets::Packet, tree::PacketSink};
 
 pub struct ServerNetwork {
     conns: Arc<Mutex<HashMap<u128, (SyncSender<Arc<Vec<u8>>>, Option<SocketAddr>)>>>,
@@ -93,9 +93,10 @@ impl ServerNetwork {
                             warn!("client outbound error: {e}");
                         }
                     });
-                    if let Err(e) = handle_conn_read(conn, sock, recv_tx) {
+                    if let Err(e) = handle_conn_read(conn, sock, &recv_tx) {
                         warn!("client inbound error: {e}");
                     }
+                    recv_tx.send((conn, Packet::Disconnect)).unwrap();
                     conns.lock().unwrap().remove(&conn);
                 });
             }
@@ -192,12 +193,16 @@ impl ServerNetwork {
     }
 }
 
-fn handle_conn_read(conn: u128, sock: TcpStream, tx: Sender<(u128, Packet)>) -> Result<()> {
+fn handle_conn_read(conn: u128, sock: TcpStream, tx: &Sender<(u128, Packet)>) -> Result<()> {
     let mut sock = BufReader::new(sock);
     loop {
         let packet = Packet::read(&mut sock)?;
         debug!("{conn} <- {packet:?}");
-        tx.send((conn, packet)).unwrap();
+        if matches!(packet, Packet::Disconnect) {
+            return Ok(()); // not sending because generic disconnect sends it too
+        } else {
+            tx.send((conn, packet)).unwrap();
+        }
     }
 }
 fn handle_conn_write(sock: TcpStream, rx: Receiver<Arc<Vec<u8>>>) -> Result<()> {
@@ -208,3 +213,15 @@ fn handle_conn_write(sock: TcpStream, rx: Receiver<Arc<Vec<u8>>>) -> Result<()>
     }
     Ok(())
 }
+
+impl ServerNetwork {
+    pub fn br<'a>(&'a self) -> Broadcaster<'a> {
+        Broadcaster(self)
+    }
+}
+pub struct Broadcaster<'a>(&'a ServerNetwork);
+impl<'a> PacketSink for Broadcaster<'a> {
+    fn push(&mut self, p: Packet) {
+        self.0.broadcast(p, true);
+    }
+}