reparenting checks and conn owned objects

5 files changed, 133 insertions, 24 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);
+    }
+}
diff --git a/shared/src/helper.rs b/shared/src/helper.rs
index 14fb2bd..3ae7c37 100644
--- a/shared/src/helper.rs
+++ b/shared/src/helper.rs
@@ -14,7 +14,7 @@
     You should have received a copy of the GNU Affero General Public License
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
-use crate::packets::{Data, Object, Resource};
+use crate::packets::{Data, Message, Object, Resource};
 use anyhow::Result;
 use glam::{Affine3A, Vec2, Vec3, Vec3A};
 use std::{
@@ -304,6 +304,21 @@ impl ReadWrite for Data {
         Ok(Self(buf))
     }
 }
+impl ReadWrite for Message {
+    fn write(&self, w: &mut dyn Write) -> Result<()> {
+        w.write_all(&(self.0.len() as u32).to_be_bytes())?;
+        w.write_all(&self.0.as_bytes())?;
+        Ok(())
+    }
+    fn read(r: &mut dyn Read) -> Result<Self> {
+        let mut size = [0; { size_of::<u32>() }];
+        r.read_exact(&mut size)?;
+        let size = u32::from_be_bytes(size);
+        let mut buf = vec![0; size as usize];
+        r.read_exact(&mut buf)?;
+        Ok(Self(String::from_utf8_lossy_owned(buf)))
+    }
+}
 impl<T> ReadWrite for Resource<T> {
     fn write(&self, w: &mut dyn Write) -> Result<()> {
         w.write_all(&self.0)?;
diff --git a/shared/src/lib.rs b/shared/src/lib.rs
index 2c26486..e5cf666 100644
--- a/shared/src/lib.rs
+++ b/shared/src/lib.rs
@@ -14,7 +14,12 @@
     You should have received a copy of the GNU Affero General Public License
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
-#![feature(iter_array_chunks, array_try_map, debug_closure_helpers)]
+#![feature(
+    iter_array_chunks,
+    array_try_map,
+    debug_closure_helpers,
+    string_from_utf8_lossy_owned
+)]
 
 pub mod helper;
 pub mod packets;
diff --git a/shared/src/tree.rs b/shared/src/tree.rs
index 2de02d0..a0bbf43 100644
--- a/shared/src/tree.rs
+++ b/shared/src/tree.rs
@@ -19,6 +19,7 @@ use crate::{
     resources::Prefab,
 };
 use glam::Vec3A;
+use log::warn;
 use std::collections::{BTreeSet, HashMap};
 
 pub struct SceneTree {
@@ -43,21 +44,10 @@ impl SceneTree {
     pub fn packet(&mut self, p: &Packet) {
         match p {
             Packet::Add(object, res) => {
-                self.objects.insert(*object, ObjectData {
-                    parent: Object(0),
-                    pos: Vec3A::ZERO,
-                    rot: Vec3A::ZERO,
-                    pose: Vec::new(),
-                    res: res.clone(),
-                    children: BTreeSet::new(),
-                });
+                self.add(*object, res.to_owned());
             }
             Packet::Remove(object) => {
-                if let Some(o) = self.objects.remove(&object) {
-                    for c in o.children {
-                        self.reparent(o.parent, c);
-                    }
-                }
+                self.remove_reparent(*object, &mut ());
             }
             Packet::Position(object, pos, rot) => {
                 if let Some(o) = self.objects.get_mut(&object) {
@@ -77,8 +67,36 @@ impl SceneTree {
         }
     }
 
+    pub fn add(&mut self, object: Object, res: Resource<Prefab>) {
+        self.objects.insert(object, ObjectData {
+            parent: Object(0),
+            pos: Vec3A::ZERO,
+            rot: Vec3A::ZERO,
+            pose: Vec::new(),
+            res: res.clone(),
+            children: BTreeSet::new(),
+        });
+    }
+
     pub fn reparent(&mut self, parent: Object, child: Object) {
+        fn check_parent_loop(tree: &SceneTree, o: Object, test: Object) -> bool {
+            if o == test {
+                true
+            } else {
+                let parent = tree.objects[&o].parent;
+                if parent == o {
+                    false
+                } else {
+                    check_parent_loop(tree, parent, test)
+                }
+            }
+        }
         if !self.objects.contains_key(&parent) || !self.objects.contains_key(&child) {
+            warn!("reparent of missing objects");
+            return;
+        }
+        if check_parent_loop(&self, parent, child) {
+            warn!("cyclic parenting prevented");
             return;
         }
         if let Some(co) = self.objects.get(&child) {
@@ -94,6 +112,23 @@ impl SceneTree {
             co.parent = parent;
         }
     }
+    pub fn remove_reparent(&mut self, object: Object, ps: &mut dyn PacketSink) {
+        if let Some(o) = self.objects.remove(&object) {
+            for c in o.children {
+                self.reparent(o.parent, c);
+                ps.push(Packet::Parent(o.parent, c));
+            }
+            ps.push(Packet::Remove(object));
+        }
+    }
+    pub fn remove_recursive(&mut self, object: Object, ps: &mut dyn PacketSink) {
+        if let Some(o) = self.objects.remove(&object) {
+            for c in o.children {
+                self.remove_recursive(c, ps);
+            }
+            ps.push(Packet::Remove(object));
+        }
+    }
 
     pub fn prime_client(&self) -> impl Iterator<Item = Packet> {
         self.objects
@@ -114,3 +149,15 @@ impl SceneTree {
             .flatten()
     }
 }
+
+pub trait PacketSink {
+    fn push(&mut self, p: Packet);
+}
+impl PacketSink for Vec<Packet> {
+    fn push(&mut self, p: Packet) {
+        Vec::push(self, p);
+    }
+}
+impl PacketSink for () {
+    fn push(&mut self, _: Packet) {}
+}