fix: malformed packet crash (#889)

When a client sends a malformed packet, the server will crash because it
used entity.destruct to disconnect players for malformed packets which
leads to an abort because, when the server receives PlayerDisconnect
from the proxy, the server will try to set PendingRemove on the
already-destructed entity.

This PR fixes these errors on player disconnects.

Author: TestingPlant

crates/hyperion-proto/src/server_to_proxy.rs

@@ -47,6 +47,14 @@ pub struct Unicast<'a> {
 #[rkyv(derive(Debug))]
 pub struct Flush;
 
+/// The server must be prepared to handle other additional packets with this stream from the proxy after the server
+/// sends [`Shutdown`] until the server receives [`crate::PlayerDisconnect`] because proxy to server packets may
+/// already be in transit.
+#[derive(Archive, Deserialize, Serialize, Clone, Copy, PartialEq, Debug)]
+pub struct Shutdown {
+    pub stream: u64,
+}
+
 #[derive(Archive, Deserialize, Serialize, Clone, PartialEq)]
 pub enum ServerToProxyMessage<'a> {
     UpdatePlayerChunkPositions(UpdatePlayerChunkPositions),
@@ -55,4 +63,5 @@ pub enum ServerToProxyMessage<'a> {
     Unicast(Unicast<'a>),
     SetReceiveBroadcasts(SetReceiveBroadcasts),
     Flush(Flush),
+    Shutdown(Shutdown),
 }

crates/hyperion-proxy/src/cache.rs

@@ -323,6 +323,9 @@ impl BufferedEgress {
                     egress.handle_broadcast_local(instruction);
                 });
             }
+            ArchivedServerToProxyMessage::Shutdown(pkt) => {
+                self.egress.handle_shutdown(pkt);
+            }
         }
     }
 

crates/hyperion-proxy/src/data.rs

@@ -40,21 +40,11 @@ impl OrderedBytes {
         data: Bytes::from_static(b""),
         exclusions: None,
     };
-    pub const SHUTDOWN: Self = Self {
-        order: u32::MAX - 1,
-        offset: 0,
-        data: Bytes::from_static(b""),
-        exclusions: None,
-    };
 
     pub const fn is_flush(&self) -> bool {
         self.order == u32::MAX
     }
 
-    pub const fn is_shutdown(&self) -> bool {
-        self.order == u32::MAX - 1
-    }
-
     pub const fn no_order(data: Bytes) -> Self {
         Self {
             order: 0,
@@ -120,8 +110,8 @@ impl PlayerHandle {
     }
 
     pub fn shutdown(&self) {
-        let _ = self.writer.try_send(OrderedBytes::SHUTDOWN);
-        self.writer.close().unwrap();
+        // Ignore error for if the channel is already closed
+        let _ = self.writer.close();
     }
 
     pub fn enable_receive_broadcasts(&self) {
@@ -143,7 +133,7 @@ impl PlayerHandle {
                 bail!("failed to send packet to player, channel is full: {is_full}");
             }
             Err(e) => {
-                self.writer.close().unwrap();
+                self.shutdown();
                 bail!("failed to send packet to player: {e}");
             }
         }

crates/hyperion-proxy/src/egress.rs

@@ -4,8 +4,8 @@ use bvh::{Aabb, Bvh};
 use bytes::Bytes;
 use glam::I16Vec2;
 use hyperion_proto::{
-    ArchivedSetReceiveBroadcasts, ArchivedUnicast, ArchivedUpdatePlayerChunkPositions,
-    ChunkPosition,
+    ArchivedSetReceiveBroadcasts, ArchivedShutdown, ArchivedUnicast,
+    ArchivedUpdatePlayerChunkPositions, ChunkPosition,
 };
 use rustc_hash::FxBuildHasher;
 use tracing::{Instrument, debug, error, info_span, instrument, warn};
@@ -19,8 +19,6 @@ use crate::{
 pub struct Egress {
     // todo: can we do some type of EntityId and SlotMap
     player_registry: &'static papaya::HashMap<u64, PlayerHandle, FxBuildHasher>,
-
-    // todo: remove positions when player leaves
     positions: &'static papaya::HashMap<u64, ChunkPosition, FxBuildHasher>,
 }
 
@@ -81,7 +79,7 @@ impl Egress {
 
                 // imo it makes sense to read once... it is a fast loop
                 #[allow(clippy::significant_drop_in_scrutinee)]
-                for (player_id, player) in &players {
+                for (_, player) in &players {
                     if !player.can_receive_broadcasts() {
                         continue;
                     }
@@ -91,9 +89,7 @@ impl Egress {
 
                     if let Err(e) = player.send(to_send) {
                         warn!("Failed to send data to player: {:?}", e);
-                        if let Some(result) = players.remove(player_id) {
-                            result.shutdown();
-                        }
+                        player.shutdown();
                     }
                 }
             }
@@ -107,12 +103,10 @@ impl Egress {
 
         tokio::spawn(
             async move {
-                for (id, player) in &players {
+                for (_, player) in &players {
                     if let Err(e) = player.send(OrderedBytes::FLUSH) {
                         warn!("Failed to send data to player: {:?}", e);
-                        if let Some(result) = players.remove(id) {
-                            result.shutdown();
-                        }
+                        player.shutdown();
                     }
                 }
             }
@@ -171,9 +165,7 @@ impl Egress {
 
                         if let Err(e) = player.send(to_send) {
                             warn!("Failed to send data to player: {:?}", e);
-                            if let Some(result) = players.remove(id) {
-                                result.shutdown();
-                            }
+                            player.shutdown();
                         }
                     }
                 }
@@ -202,16 +194,14 @@ impl Egress {
 
         let Some(player) = players.get(&id) else {
             // expected to still happen infrequently
-            debug!("Player not found for id {id:?}");
+            warn!("Player not found for id {id:?}");
             return;
         };
 
         // todo: handle error; kick player if cannot send (buffer full)
         if let Err(e) = player.send(ordered) {
             warn!("Failed to send data to player: {:?}", e);
-            if let Some(result) = players.remove(&id) {
-                result.shutdown();
-            }
+            player.shutdown();
         }
 
         drop(players);
@@ -230,4 +220,17 @@ impl Egress {
 
         player.enable_receive_broadcasts();
     }
+
+    #[instrument(skip_all)]
+    pub fn handle_shutdown(&self, pkt: &ArchivedShutdown) {
+        let player_registry = self.player_registry;
+        let players = player_registry.pin();
+        let Ok(stream) = rkyv::deserialize::<u64, !>(&pkt.stream);
+
+        if let Some(result) = players.get(&stream) {
+            result.shutdown();
+        } else {
+            error!("Player not found for stream {stream:?}");
+        }
+    }
 }

crates/hyperion-proxy/src/player.rs

@@ -115,10 +115,6 @@ pub fn initiate_player_connection(
         let mut packet_writer = PlayerPacketWriter::new(socket_writer, player_id);
 
         while let Ok(outgoing_packet) = incoming_packet_receiver.recv().await {
-            if outgoing_packet.is_shutdown() {
-                return;
-            }
-
             if outgoing_packet.is_flush() {
                 let time_start = std::time::Instant::now();
                 if let Err(e) = packet_writer.flush_pending_packets().await {
@@ -175,15 +171,14 @@ pub fn initiate_player_connection(
                 if let Err(e) = server_sender.send(disconnect).await {
                     warn!("failed to send player disconnect to server: {e}");
                 }
-
-                let map_ref = player_registry.pin();
-                map_ref.remove(&player_id);
-
-                let map_ref = player_positions.pin();
-                map_ref.remove(&player_id);
-
             }
         }
+
+        let map_ref = player_registry.pin();
+        map_ref.remove(&player_id);
+
+        let map_ref = player_positions.pin();
+        map_ref.remove(&player_id);
     })
 }
 

crates/hyperion/src/egress/player_join/mod.rs

@@ -6,7 +6,7 @@ use glam::DVec3;
 use hyperion_crafting::{Action, CraftingRegistry, RecipeBookState};
 use hyperion_utils::EntityExt;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
-use tracing::{info, instrument};
+use tracing::{info, instrument, warn};
 use valence_protocol::{
     ByteAngle, GameMode, Ident, PacketEncoder, RawBytes, VarInt, Velocity,
     game_mode::OptGameMode,
@@ -29,7 +29,6 @@ pub use list::*;
 use crate::{
     config::Config,
     egress::metadata::show_all,
-    ingress::PendingRemove,
     net::{Compose, ConnectionId, DataBundle},
     simulation::{
         Comms, Name, Position, Uuid, Yaw,
@@ -598,7 +597,8 @@ impl Module for PlayerJoinModule {
                                 crafting_registry,
                                 config,
                             ) {
-                                entity.set(PendingRemove::new(e.to_string()));
+                                warn!("player_join_world error: {e:?}");
+                                compose.io_buf().shutdown(stream_id, world);
                             }
                         },
                     );

crates/hyperion/src/ingress/mod.rs

@@ -38,19 +38,10 @@ use crate::{
     util::{TracingExt, mojang::MojangClient},
 };
 
+/// This marks players who have already been disconnected and about to be destructed. This component should not be
+/// added to an entity to disconnect a player. Use [`crate::net::IoBuf::shutdown`] instead.
 #[derive(Component, Debug)]
-pub struct PendingRemove {
-    pub reason: String,
-}
-
-impl PendingRemove {
-    #[must_use]
-    pub fn new(reason: impl Into<String>) -> Self {
-        Self {
-            reason: reason.into(),
-        }
-    }
-}
+pub struct PendingRemove;
 
 fn process_handshake(
     login_state: &mut PacketState,
@@ -335,13 +326,11 @@ impl Module for IngressModule {
             for disconnect in recv.player_disconnect.lock().drain(..) {
                 // will initiate the removal of entity
                 info!("queue pending remove");
-                let Some(id) = lookup.get(&disconnect).copied() else {
+                let Some(entity_id) = lookup.get(&disconnect).copied() else {
                     error!("failed to get id for disconnect stream {disconnect:?}");
                     continue;
                 };
-                world
-                    .entity_from_id(*id)
-                    .set(PendingRemove::new("disconnected"));
+                world.entity_from_id(*entity_id).add(id::<PendingRemove>());
             }
         });
 
@@ -375,11 +364,10 @@ impl Module for IngressModule {
             world,
             &Uuid,
             &Compose($),
-            &ConnectionId,
-            &PendingRemove,
         )
         .kind(id::<flecs::pipeline::PostLoad>())
-        .each_iter(move |it, row, (uuid, compose, io, pending_remove)| {
+        .with(id::<PendingRemove>())
+        .each_iter(move |it, row, (uuid, compose)| {
             let system = it.system();
             let entity = it.entity(row).expect("row must be in bounds");
             let uuids = &[uuid.0];
@@ -402,16 +390,6 @@ impl Module for IngressModule {
             if let Err(e) = compose.broadcast(&pkt, system).send() {
                 error!("failed to send player remove packet: {e}");
             }
-
-            if !pending_remove.reason.is_empty() {
-                let pkt = play::DisconnectS2c {
-                    reason: pending_remove.reason.clone().into_cow_text(),
-                };
-
-                if let Err(e) = compose.unicast_no_compression(&pkt, *io, system) {
-                    error!("failed to send disconnect packet: {e}");
-                }
-            }
         });
 
         world
@@ -485,7 +463,7 @@ impl Module for IngressModule {
                         Ok(frame) => frame,
                         Err(e) => {
                             error!("failed to decode packet: {e}");
-                            entity.destruct();
+                            compose.io_buf().shutdown(io_ref, &world);
                             break;
                         }
                     };
@@ -498,9 +476,7 @@ impl Module for IngressModule {
                         PacketState::Handshake => {
                             if process_handshake(login_state, &frame).is_err() {
                                 error!("failed to process handshake");
-
-                                entity.destruct();
-
+                                compose.io_buf().shutdown(io_ref, &world);
                                 break;
                             }
                         }
@@ -509,7 +485,7 @@ impl Module for IngressModule {
                                 process_status(login_state, system, &frame, io_ref, compose)
                             {
                                 error!("failed to process status packet: {e}");
-                                entity.destruct();
+                                compose.io_buf().shutdown(io_ref, &world);
                                 break;
                             }
                         }
@@ -549,7 +525,7 @@ impl Module for IngressModule {
                                     error!("failed to send login disconnect packet: {e}");
                                 }
 
-                                entity.destruct();
+                                compose.io_buf().shutdown(io_ref, &world);
                                 break;
                             }
                         }

crates/hyperion/src/net/mod.rs

@@ -636,4 +636,24 @@ impl IoBuf {
         let packet_len = u64::try_from(new_len - len - size_of::<u64>()).unwrap();
         buffer[len..(len + 8)].copy_from_slice(&packet_len.to_be_bytes());
     }
+
+    pub fn shutdown(&self, stream: ConnectionId, world: &World) {
+        let buffer = self.buffer.get(world);
+        let buffer = &mut *buffer.borrow_mut();
+
+        let to_send = hyperion_proto::Shutdown {
+            stream: stream.stream_id,
+        };
+
+        let to_send = ServerToProxyMessage::Shutdown(to_send);
+
+        let len = buffer.len();
+        buffer.write_u64::<byteorder::BigEndian>(0x00).unwrap();
+
+        rkyv::api::high::to_bytes_in::<_, rkyv::rancor::Error>(&to_send, &mut *buffer).unwrap();
+
+        let new_len = buffer.len();
+        let packet_len = u64::try_from(new_len - len - size_of::<u64>()).unwrap();
+        buffer[len..(len + 8)].copy_from_slice(&packet_len.to_be_bytes());
+    }
 }