From 533357ceaaefbf78ccb579d5d61a9c35ecf4515f Mon Sep 17 00:00:00 2001
From: Snowiiii <lilalexmed@proton.me>
Date: Sun, 13 Oct 2024 15:17:09 +0200
Subject: [PATCH] Fix: kick in player state when client fails

---
 pumpkin/src/client/mod.rs    | 8 +++++++-
 pumpkin/src/entity/player.rs | 3 +--
 pumpkin/src/main.rs          | 3 +--
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/pumpkin/src/client/mod.rs b/pumpkin/src/client/mod.rs
index 252c5b9e2..2fcd008a9 100644
--- a/pumpkin/src/client/mod.rs
+++ b/pumpkin/src/client/mod.rs
@@ -17,9 +17,10 @@ use crossbeam::atomic::AtomicCell;
 use mio::{event::Event, net::TcpStream, Token};
 use parking_lot::Mutex;
 use pumpkin_config::compression::CompressionInfo;
+use pumpkin_core::text::TextComponent;
 use pumpkin_protocol::{
     bytebuf::{packet_id::Packet, DeserializerError},
-    client::{config::CConfigDisconnect, login::CLoginDisconnect},
+    client::{config::CConfigDisconnect, login::CLoginDisconnect, play::CPlayDisconnect},
     packet_decoder::PacketDecoder,
     packet_encoder::PacketEncoder,
     server::{
@@ -416,6 +417,11 @@ impl Client {
                 self.try_send_packet(&CConfigDisconnect::new(reason))
                     .unwrap_or_else(|_| self.close());
             }
+            // This way players get kicked when players using client functions (e.g. poll, send_packet)
+            ConnectionState::Play => {
+                self.try_send_packet(&CPlayDisconnect::new(&TextComponent::text(reason)))
+                    .unwrap_or_else(|_| self.close());
+            }
             _ => {
                 log::warn!("Can't kick in {:?} State", self.connection_state)
             }
diff --git a/pumpkin/src/entity/player.rs b/pumpkin/src/entity/player.rs
index fc39cdf2b..496cedfe0 100644
--- a/pumpkin/src/entity/player.rs
+++ b/pumpkin/src/entity/player.rs
@@ -26,7 +26,7 @@ use pumpkin_protocol::{
         SPlayerPositionRotation, SPlayerRotation, SSetCreativeSlot, SSetHeldItem, SSetPlayerGround,
         SSwingArm, SUseItem, SUseItemOn,
     },
-    ConnectionState, RawPacket, ServerPacket, VarInt,
+    RawPacket, ServerPacket, VarInt,
 };
 
 use pumpkin_protocol::server::play::{SCloseContainer, SKeepAlive};
@@ -215,7 +215,6 @@ impl Player {
 
     /// Kicks the Client with a reason depending on the connection state
     pub fn kick(&self, reason: TextComponent) {
-        assert!(self.client.connection_state.load() == ConnectionState::Play);
         assert!(!self
             .client
             .closed
diff --git a/pumpkin/src/main.rs b/pumpkin/src/main.rs
index 976910003..8cf3c0153 100644
--- a/pumpkin/src/main.rs
+++ b/pumpkin/src/main.rs
@@ -234,7 +234,7 @@ fn main() -> io::Result<()> {
                         }
                         clients.insert(token, client);
                     },
-
+                    // Maybe received an event for a TCP connection.
                     token => {
                         // Poll Players
                         if let Some(player) = players.get_mut(&token) {
@@ -256,7 +256,6 @@ fn main() -> io::Result<()> {
                         };
 
                         // Poll current Clients (non players)
-                        // Maybe received an event for a TCP connection.
                         let (done, make_player) = if let Some(client) = clients.get_mut(&token) {
                             client.poll(event).await;
                             let closed = client.closed.load(std::sync::atomic::Ordering::Relaxed);