From b0f1e60e46c478ec1e855969c6d5e0fc081b5c92 Mon Sep 17 00:00:00 2001
From: Edvin Bryntesson <epost@edvinbryntesson.se>
Date: Mon, 9 Sep 2024 12:51:35 +0200
Subject: [PATCH] Add a `PumpkinError` trait, and implement it for
 `InventoryError` and `DeserializerError`

---
 pumpkin-inventory/src/error.rs | 14 --------
 pumpkin/src/entity/player.rs   | 28 +++++++++------
 pumpkin/src/error.rs           | 64 ++++++++++++++++++++++++++++++++++
 pumpkin/src/main.rs            |  1 +
 4 files changed, 82 insertions(+), 25 deletions(-)
 create mode 100644 pumpkin/src/error.rs

diff --git a/pumpkin-inventory/src/error.rs b/pumpkin-inventory/src/error.rs
index e74ece87e..008b22ff5 100644
--- a/pumpkin-inventory/src/error.rs
+++ b/pumpkin-inventory/src/error.rs
@@ -17,17 +17,3 @@ pub enum InventoryError {
     #[error("Player does not have enough permissions")]
     PermissionError,
 }
-
-impl InventoryError {
-    pub fn should_kick(&self) -> bool {
-        match self {
-            InventoryError::InvalidSlot
-            | InventoryError::ClosedContainerInteract(..)
-            | InventoryError::InvalidPacket
-            | InventoryError::PermissionError => true,
-            InventoryError::LockError
-            | InventoryError::OutOfOrderDragging
-            | InventoryError::MultiplePlayersDragging => false,
-        }
-    }
-}
diff --git a/pumpkin/src/entity/player.rs b/pumpkin/src/entity/player.rs
index 774d308ec..3b7ce24e2 100644
--- a/pumpkin/src/entity/player.rs
+++ b/pumpkin/src/entity/player.rs
@@ -10,7 +10,7 @@ use pumpkin_core::{
 use pumpkin_entity::{entity_type::EntityType, EntityId};
 use pumpkin_inventory::player::PlayerInventory;
 use pumpkin_protocol::{
-    bytebuf::{packet_id::Packet, DeserializerError},
+    bytebuf::packet_id::Packet,
     client::play::{
         CGameEvent, CPlayDisconnect, CPlayerAbilities, CPlayerInfoUpdate, CSyncPlayerPosition,
         CSystemChatMessage, PlayerAction,
@@ -27,14 +27,14 @@ use pumpkin_protocol::{
 use pumpkin_protocol::server::play::SCloseContainer;
 use pumpkin_world::item::ItemStack;
 
+use super::Entity;
+use crate::error::PumpkinError;
 use crate::{
     client::{authentication::GameProfile, Client, PlayerConfig},
     server::Server,
     world::World,
 };
 
-use super::Entity;
-
 pub struct PlayerAbilities {
     pub invulnerable: bool,
     pub flying: bool,
@@ -268,9 +268,17 @@ impl Player {
             match self.handle_play_packet(server, &mut packet).await {
                 Ok(_) => {}
                 Err(e) => {
-                    let text = format!("Error while reading incoming packet {}", e);
-                    log::error!("{}", text);
-                    self.kick(TextComponent::text(&text))
+                    if e.is_kick() {
+                        if let Some(kick_reason) = e.client_kick_reason() {
+                            self.kick(TextComponent::text(&kick_reason))
+                        } else {
+                            self.kick(TextComponent::text(&format!(
+                                "Error while reading incoming packet {}",
+                                e
+                            )));
+                        }
+                    }
+                    e.log();
                 }
             };
         }
@@ -280,7 +288,7 @@ impl Player {
         &mut self,
         server: &Arc<Server>,
         packet: &mut RawPacket,
-    ) -> Result<(), DeserializerError> {
+    ) -> Result<(), Box<dyn PumpkinError>> {
         let bytebuf = &mut packet.bytebuf;
         match packet.id.0 {
             SConfirmTeleport::PACKET_ID => {
@@ -353,8 +361,7 @@ impl Player {
                 Ok(())
             }
             SSetCreativeSlot::PACKET_ID => {
-                self.handle_set_creative_slot(server, SSetCreativeSlot::read(bytebuf)?)
-                    .unwrap();
+                self.handle_set_creative_slot(server, SSetCreativeSlot::read(bytebuf)?)?;
                 Ok(())
             }
             SPlayPingRequest::PACKET_ID => {
@@ -363,8 +370,7 @@ impl Player {
             }
             SClickContainer::PACKET_ID => {
                 self.handle_click_container(server, SClickContainer::read(bytebuf)?)
-                    .await
-                    .unwrap();
+                    .await?;
                 Ok(())
             }
             SCloseContainer::PACKET_ID => {
diff --git a/pumpkin/src/error.rs b/pumpkin/src/error.rs
new file mode 100644
index 000000000..caf9b98a6
--- /dev/null
+++ b/pumpkin/src/error.rs
@@ -0,0 +1,64 @@
+use log::log;
+use pumpkin_inventory::InventoryError;
+use pumpkin_protocol::bytebuf::DeserializerError;
+use std::fmt::Display;
+
+pub trait PumpkinError: std::error::Error + Display {
+    fn is_kick(&self) -> bool;
+
+    fn log(&self) {
+        log!(self.severity(), "{}", self.to_string());
+    }
+
+    fn severity(&self) -> log::Level;
+
+    fn client_kick_reason(&self) -> Option<String>;
+}
+
+impl<ErrorType: PumpkinError + 'static> From<ErrorType> for Box<dyn PumpkinError> {
+    fn from(error: ErrorType) -> Self {
+        Box::new(error)
+    }
+}
+impl PumpkinError for InventoryError {
+    fn is_kick(&self) -> bool {
+        match self {
+            InventoryError::InvalidSlot
+            | InventoryError::ClosedContainerInteract(..)
+            | InventoryError::InvalidPacket
+            | InventoryError::PermissionError => true,
+            InventoryError::LockError
+            | InventoryError::OutOfOrderDragging
+            | InventoryError::MultiplePlayersDragging => false,
+        }
+    }
+    fn severity(&self) -> log::Level {
+        match self {
+            InventoryError::LockError
+            | InventoryError::InvalidSlot
+            | InventoryError::ClosedContainerInteract(..)
+            | InventoryError::InvalidPacket
+            | InventoryError::PermissionError => log::Level::Error,
+            InventoryError::OutOfOrderDragging => log::Level::Info,
+            InventoryError::MultiplePlayersDragging => log::Level::Warn,
+        }
+    }
+
+    fn client_kick_reason(&self) -> Option<String> {
+        None
+    }
+}
+
+impl PumpkinError for DeserializerError {
+    fn is_kick(&self) -> bool {
+        true
+    }
+
+    fn severity(&self) -> log::Level {
+        log::Level::Error
+    }
+
+    fn client_kick_reason(&self) -> Option<String> {
+        None
+    }
+}
diff --git a/pumpkin/src/main.rs b/pumpkin/src/main.rs
index 4be8e54c6..c28cfe5e7 100644
--- a/pumpkin/src/main.rs
+++ b/pumpkin/src/main.rs
@@ -17,6 +17,7 @@ use server::Server;
 pub mod client;
 pub mod commands;
 pub mod entity;
+mod error;
 pub mod proxy;
 pub mod rcon;
 pub mod server;