DisconnectEvent

5 files changed, 76 insertions, 18 deletions

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 13b0522a..8e1a7938 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -1,12 +1,13 @@
 use crate::{
     chat::ChatPlugin,
+    disconnect::{DisconnectEvent, DisconnectPlugin},
     events::{Event, EventPlugin, LocalPlayerEvents},
     local_player::{
         death_event, handle_send_packet_event, update_in_loaded_chunk, GameProfileComponent,
         LocalPlayer, PhysicsState, SendPacketEvent,
     },
     movement::{local_player_ai_step, send_position, sprint_listener, walk_listener},
-    packet_handling::{self, PacketHandlerPlugin},
+    packet_handling::{self, PacketHandlerPlugin, PacketReceiver},
     player::retroactively_add_game_profile_component,
     task_pool::TaskPoolPlugin,
     Account, PlayerInfo, StartSprintEvent, StartWalkEvent,
@@ -16,6 +17,7 @@ use azalea_auth::{game_profile::GameProfile, sessionserver::ClientSessionServerE
 use azalea_chat::FormattedText;
 use azalea_ecs::{
     app::{App, Plugin, PluginGroup, PluginGroupBuilder},
+    bundle::Bundle,
     component::Component,
     entity::Entity,
     schedule::{IntoSystemDescriptor, Schedule, Stage, SystemSet},
@@ -225,14 +227,14 @@ impl Client {
         local_player.tasks.push(read_packets_task);
         local_player.tasks.push(write_packets_task);
 
-        ecs.entity_mut(entity).insert((
+        ecs.entity_mut(entity).insert(JoinedClientBundle {
             local_player,
             packet_receiver,
-            GameProfileComponent(game_profile),
-            PhysicsState::default(),
-            Local,
-            LocalPlayerEvents(tx),
-        ));
+            game_profile: GameProfileComponent(game_profile),
+            physics_state: PhysicsState::default(),
+            local_player_events: LocalPlayerEvents(tx),
+            _local: Local,
+        });
 
         Ok((client, rx))
     }
@@ -371,7 +373,9 @@ impl Client {
     /// The OwnedReadHalf for the TCP connection is in one of the tasks, so it
     /// automatically closes the connection when that's dropped.
     pub fn disconnect(&self) {
-        self.local_player_mut(&mut self.ecs.lock()).disconnect();
+        self.ecs.lock().send_event(DisconnectEvent {
+            entity: self.entity,
+        });
     }
 
     pub fn local_player<'a>(&'a self, ecs: &'a mut Ecs) -> &'a LocalPlayer {
@@ -468,6 +472,18 @@ impl Client {
     }
 }
 
+/// A bundle for the components that are present on a local player that received
+/// a login packet. If you want to filter for this, just use [`Local`].
+#[derive(Bundle)]
+pub struct JoinedClientBundle {
+    pub local_player: LocalPlayer,
+    pub packet_receiver: PacketReceiver,
+    pub game_profile: GameProfileComponent,
+    pub physics_state: PhysicsState,
+    pub local_player_events: LocalPlayerEvents,
+    pub _local: Local,
+}
+
 pub struct AzaleaPlugin;
 impl Plugin for AzaleaPlugin {
     fn build(&self, app: &mut App) {
@@ -596,5 +612,6 @@ impl PluginGroup for DefaultPlugins {
             .add(EventPlugin)
             .add(TaskPoolPlugin::default())
             .add(ChatPlugin)
+            .add(DisconnectPlugin)
     }
 }
diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs
new file mode 100644
index 00000000..100618c3
--- /dev/null
+++ b/azalea-client/src/disconnect.rs
@@ -0,0 +1,31 @@
+//! Disconnect a client from the server.
+
+use azalea_ecs::{
+    app::{App, CoreStage, Plugin},
+    entity::Entity,
+    event::EventReader,
+    system::Commands,
+};
+
+use crate::client::JoinedClientBundle;
+
+pub struct DisconnectPlugin;
+impl Plugin for DisconnectPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_event::<DisconnectEvent>()
+            .add_system_to_stage(CoreStage::PostUpdate, handle_disconnect);
+    }
+}
+
+/// An event sent when a client is getting disconnected.
+pub struct DisconnectEvent {
+    pub entity: Entity,
+}
+
+/// System that removes the [`JoinedClientBundle`] from the entity when it
+/// receives a [`DisconnectEvent`].
+pub fn handle_disconnect(mut commands: Commands, mut events: EventReader<DisconnectEvent>) {
+    for DisconnectEvent { entity } in events.iter() {
+        commands.entity(*entity).remove::<JoinedClientBundle>();
+    }
+}
diff --git a/azalea-client/src/lib.rs b/azalea-client/src/lib.rs
index a6782fa1..8e119c0e 100644
--- a/azalea-client/src/lib.rs
+++ b/azalea-client/src/lib.rs
@@ -14,6 +14,7 @@
 mod account;
 pub mod chat;
 mod client;
+pub mod disconnect;
 mod entity_query;
 mod events;
 mod get_mc_dir;
diff --git a/azalea-client/src/local_player.rs b/azalea-client/src/local_player.rs
index bf3f18f1..8298cf0a 100644
--- a/azalea-client/src/local_player.rs
+++ b/azalea-client/src/local_player.rs
@@ -112,12 +112,11 @@ impl LocalPlayer {
             .send(packet)
             .expect("write_packet shouldn't be able to be called if the connection is closed");
     }
+}
 
-    /// Disconnect this client from the server by ending all tasks.
-    ///
-    /// The OwnedReadHalf for the TCP connection is in one of the tasks, so it
-    /// automatically closes the connection when that's dropped.
-    pub fn disconnect(&self) {
+impl Drop for LocalPlayer {
+    /// Stop every active task when the `LocalPlayer` is dropped.
+    fn drop(&mut self) {
         for task in &self.tasks {
             task.abort();
         }
diff --git a/azalea-client/src/packet_handling.rs b/azalea-client/src/packet_handling.rs
index 1dcf3c06..516c6784 100644
--- a/azalea-client/src/packet_handling.rs
+++ b/azalea-client/src/packet_handling.rs
@@ -38,6 +38,7 @@ use tokio::sync::mpsc;
 
 use crate::{
     chat::{ChatPacket, ChatReceivedEvent},
+    disconnect::DisconnectEvent,
     local_player::{GameProfileComponent, LocalPlayer},
     ClientInformation, PlayerInfo,
 };
@@ -286,10 +287,14 @@ fn handle_packets(ecs: &mut Ecs) {
                 }
                 ClientboundGamePacket::Disconnect(p) => {
                     debug!("Got disconnect packet {:?}", p);
-                    let mut system_state: SystemState<Query<&LocalPlayer>> = SystemState::new(ecs);
-                    let query = system_state.get(ecs);
-                    let local_player = query.get(player_entity).unwrap();
-                    local_player.disconnect();
+                    let mut system_state: SystemState<EventWriter<DisconnectEvent>> =
+                        SystemState::new(ecs);
+                    let mut disconnect_events = system_state.get_mut(ecs);
+                    disconnect_events.send(DisconnectEvent {
+                        entity: player_entity,
+                    });
+                    // bye
+                    return;
                 }
                 ClientboundGamePacket::UpdateRecipes(_p) => {
                     debug!("Got update recipes packet");
@@ -950,10 +955,14 @@ impl PacketReceiver {
                     if !matches!(*error, ReadPacketError::ConnectionClosed) {
                         error!("Error reading packet from Client: {error:?}");
                     }
-                    return;
+                    break;
                 }
             }
         }
+        // TODO: it should send a DisconnectEvent here somehow
+        // maybe use a tokio::sync::oneshot that tells it to close and have the
+        // receiver in localplayer and have a system that watches that or
+        // something?
     }
 
     /// Consume the [`ServerboundGamePacket`] queue and actually write the
@@ -970,6 +979,7 @@ impl PacketReceiver {
                 break;
             };
         }
+        println!("Write task finished");
         // receiver is automatically closed when it's dropped
     }
 }