From fbee81e609e0cbc1bb8aff8ae649bca4e44ae465 Mon Sep 17 00:00:00 2001
From: 1zuna <1zuna@ccbluex.net>
Date: Sat, 2 Dec 2023 05:45:12 +0100
Subject: addd disconnect reason (#114)

---
 azalea-client/src/client.rs                        |  1 +
 azalea-client/src/disconnect.rs                    |  5 +++--
 azalea-client/src/events.rs                        | 13 ++++++++++++-
 azalea-client/src/packet_handling/configuration.rs |  1 +
 azalea-client/src/packet_handling/game.rs          |  1 +
 5 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 869ab0f3..95b7b747 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -422,6 +422,7 @@ impl Client {
     pub fn disconnect(&self) {
         self.ecs.lock().send_event(DisconnectEvent {
             entity: self.entity,
+            reason: None,
         });
     }
 
diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs
index 916ef1fb..ac663a66 100644
--- a/azalea-client/src/disconnect.rs
+++ b/azalea-client/src/disconnect.rs
@@ -33,6 +33,7 @@ impl Plugin for DisconnectPlugin {
 #[derive(Event)]
 pub struct DisconnectEvent {
     pub entity: Entity,
+    pub reason: Option<String>,
 }
 
 /// System that removes the [`JoinedClientBundle`] from the entity when it
@@ -41,7 +42,7 @@ pub fn remove_components_from_disconnected_players(
     mut commands: Commands,
     mut events: EventReader<DisconnectEvent>,
 ) {
-    for DisconnectEvent { entity } in events.read() {
+    for DisconnectEvent { entity, .. } in events.read() {
         commands.entity(*entity).remove::<JoinedClientBundle>();
     }
 }
@@ -64,7 +65,7 @@ fn disconnect_on_connection_dead(
 ) {
     for (entity, &is_connection_alive) in &query {
         if !*is_connection_alive {
-            disconnect_events.send(DisconnectEvent { entity });
+            disconnect_events.send(DisconnectEvent { entity, reason: None });
         }
     }
 }
diff --git a/azalea-client/src/events.rs b/azalea-client/src/events.rs
index c85f142e..0441ae98 100644
--- a/azalea-client/src/events.rs
+++ b/azalea-client/src/events.rs
@@ -24,7 +24,7 @@ use crate::{
         AddPlayerEvent, DeathEvent, KeepAliveEvent, PacketEvent, RemovePlayerEvent,
         UpdatePlayerEvent,
     },
-    PlayerInfo,
+    PlayerInfo, disconnect::DisconnectEvent,
 };
 
 // (for contributors):
@@ -93,6 +93,8 @@ pub enum Event {
     Death(Option<Arc<ClientboundPlayerCombatKillPacket>>),
     /// A `KeepAlive` packet was sent by the server.
     KeepAlive(u64),
+    /// The client disconnected from the server.
+    Disconnect(Option<String>),
 }
 
 /// A component that contains an event sender for events that are only
@@ -117,6 +119,7 @@ impl Plugin for EventPlugin {
                 remove_player_listener,
                 keepalive_listener,
                 death_listener,
+                disconnect_listener,
             ),
         )
         .add_systems(
@@ -229,3 +232,11 @@ fn keepalive_listener(query: Query<&LocalPlayerEvents>, mut events: EventReader<
             .unwrap();
     }
 }
+
+fn disconnect_listener(query: Query<&LocalPlayerEvents>, mut events: EventReader<DisconnectEvent>) {
+    for event in events.read() {
+        if let Ok(local_player_events) = query.get(event.entity) {
+            let _ = local_player_events.send(Event::Disconnect(event.reason.clone()));
+        }
+    }
+}
diff --git a/azalea-client/src/packet_handling/configuration.rs b/azalea-client/src/packet_handling/configuration.rs
index ab0ce089..82016fe6 100644
--- a/azalea-client/src/packet_handling/configuration.rs
+++ b/azalea-client/src/packet_handling/configuration.rs
@@ -104,6 +104,7 @@ pub fn process_packet_events(ecs: &mut World) {
                 let mut disconnect_events = system_state.get_mut(ecs);
                 disconnect_events.send(DisconnectEvent {
                     entity: player_entity,
+                    reason: Some(p.reason.to_ansi()),
                 });
             }
             ClientboundConfigurationPacket::FinishConfiguration(p) => {
diff --git a/azalea-client/src/packet_handling/game.rs b/azalea-client/src/packet_handling/game.rs
index 9f7c5e69..db8da69a 100644
--- a/azalea-client/src/packet_handling/game.rs
+++ b/azalea-client/src/packet_handling/game.rs
@@ -401,6 +401,7 @@ pub fn process_packet_events(ecs: &mut World) {
                 let mut disconnect_events = system_state.get_mut(ecs);
                 disconnect_events.send(DisconnectEvent {
                     entity: player_entity,
+                    reason: Some(p.reason.to_ansi()),
                 });
             }
             ClientboundGamePacket::UpdateRecipes(_p) => {
-- 
cgit v1.2.3