Implement Suggestions

Author: GStudiosX2

Diff for: src/bin/src/systems/keep_alive_system.rs

@@ -5,7 +5,7 @@ use ferrumc_net::connection::{ConnectionState, StreamWriter};
 use ferrumc_net::packets::incoming::keep_alive::IncomingKeepAlivePacket;
 use ferrumc_net::packets::outgoing::keep_alive::OutgoingKeepAlivePacket;
 use ferrumc_net::utils::broadcast::{BroadcastOptions, BroadcastToAll};
-use ferrumc_net::utils::state::terminate_connection;
+use ferrumc_net::utils::state::TerminateConnectionPlayerExt;
 use ferrumc_state::GlobalState;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
@@ -76,9 +76,8 @@ impl System for KeepAliveSystem {
 
                     if (current_time - keep_alive.timestamp) >= 30000 {
                         // two iterations missed
-                        if let Err(e) = terminate_connection(
+                        if let Err(e) = entity.terminate_connection(
                             state.clone(),
-                            *entity,
                             "Keep alive timeout".to_string(),
                         )
                         .await

Diff for: src/lib/net/crates/codec/src/net_types/length_prefixed_vec.rs

@@ -1,6 +1,7 @@
 use crate::decode::{NetDecode, NetDecodeOpts, NetDecodeResult};
 use crate::encode::{NetEncode, NetEncodeOpts, NetEncodeResult};
 use crate::net_types::var_int::VarInt;
+use std::ops::{Deref, DerefMut};
 use std::io::{Read, Write};
 use tokio::io::AsyncWrite;
 
@@ -13,17 +14,19 @@ impl<T> LengthPrefixedVec<T> {
     pub fn new(data: Vec<T>) -> Self {
         Self { data }
     }
+}
 
-    pub fn push(&mut self, value: T) {
-        self.data.push(value);
-    }
+impl<T> Deref for LengthPrefixedVec<T> {
+    type Target = Vec<T>;
 
-    pub fn pop(&mut self) -> Option<T> {
-        self.data.pop()
+    fn deref(&self) -> &Self::Target {
+        &self.data
     }
+}
 
-    pub fn length(&self) -> usize {
-        self.data.len()
+impl<T> DerefMut for LengthPrefixedVec<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.data
     }
 }
 
@@ -32,7 +35,7 @@ where
     T: NetEncode,
 {
     fn encode<W: Write>(&self, writer: &mut W, opts: &NetEncodeOpts) -> NetEncodeResult<()> {
-        VarInt::from(self.length()).encode(writer, opts)?;
+        VarInt::from(self.len()).encode(writer, opts)?;
 
         for item in &self.data {
             item.encode(writer, opts)?;
@@ -46,7 +49,7 @@ where
         writer: &mut W,
         opts: &NetEncodeOpts,
     ) -> NetEncodeResult<()> {
-        VarInt::from(self.length()).encode_async(writer, opts).await?;
+        VarInt::from(self.len()).encode_async(writer, opts).await?;
 
         for item in &self.data {
             item.encode_async(writer, opts).await?;

Diff for: src/lib/net/src/connection.rs

@@ -1,5 +1,5 @@
 use crate::packets::incoming::packet_skeleton::PacketSkeleton;
-use crate::utils::state::terminate_connection;
+use crate::utils::state::TerminateConnectionPlayerExt;
 use crate::{handle_packet, NetResult, packets::outgoing::disconnect::DISCONNECT_STRING};
 use crate::errors::NetError;
 use ferrumc_net_codec::encode::NetEncode;
@@ -164,7 +164,7 @@ pub async fn handle_connection(state: Arc<ServerState>, tcp_stream: TcpStream) -
         {
             match e {
                 NetError::Kick(msg) => {
-                    terminate_connection(state.clone(), entity, msg.clone())
+                    entity.terminate_connection(state.clone(), msg.clone())
                         .await?;
                 },
                 _ => {
@@ -174,7 +174,7 @@ pub async fn handle_connection(state: Arc<ServerState>, tcp_stream: TcpStream) -
                         packet_skele.id,
                         conn_state.as_str()
                     );
-                    terminate_connection(state.clone(), entity, DISCONNECT_STRING.to_string())
+                    entity.terminate_connection(state.clone(), DISCONNECT_STRING.to_string())
                         .await?;
                 }
             }

Diff for: src/lib/net/src/packets/incoming/keep_alive.rs

@@ -1,6 +1,6 @@
 use crate::packets::outgoing::keep_alive::OutgoingKeepAlivePacket;
 use crate::packets::IncomingPacket;
-use crate::utils::state::terminate_connection;
+use crate::utils::state::TerminateConnectionPlayerExt;
 use crate::NetResult;
 use ferrumc_macros::{packet, NetDecode};
 use ferrumc_state::ServerState;
@@ -22,7 +22,7 @@ impl IncomingPacket for IncomingKeepAlivePacket {
                 conn_id, self.timestamp, last_sent_keep_alive.timestamp
             );
             if let Err(e) =
-                terminate_connection(state, conn_id, "Invalid keep alive packet".to_string()).await
+                conn_id.terminate_connection(state, "Invalid keep alive packet".to_string()).await
             {
                 debug!("Error terminating connection: {:?}", e);
             }

Diff for: src/lib/net/src/utils/state.rs

@@ -10,61 +10,72 @@ use tracing::{trace, warn};
 
 use super::ecs_helpers::EntityExt;
 
-// used codium for this function comment, very useful
+pub trait TerminateConnectionPlayerExt {
+    #[allow(async_fn_in_trait)]
+    async fn terminate_connection(
+        &self,
+        state: GlobalState,
+        reason: impl Into<ferrumc_text::TextComponent>,
+    ) -> NetResult<()>;
+}
 
-/// Terminates the connection of an entity with the given `conn_id`.
-///
-/// Sends a disconnect packet with the given `reason` to the client, and marks the connection as
-/// terminated. This will cause the connection to be dropped on the next tick of the
-/// `ConnectionSystem`.
-///
-/// # Errors
-///
-/// Returns an error if the stream writer or connection control component cannot be accessed for
-/// the given `conn_id`.
-pub async fn terminate_connection(
-    state: GlobalState,
-    conn_id: usize,
-    reason: impl Into<ferrumc_text::TextComponent>,
-) -> NetResult<()> {
-    let mut writer = match conn_id.get_mut::<StreamWriter>(&state.clone()) {
-        Ok(writer) => writer,
-        Err(e) => {
-            warn!("Failed to get stream writer for entity {}: {}", conn_id, e);
-            return Err(NetError::ECSError(e));
-        }
-    };
+impl TerminateConnectionPlayerExt for usize {
+    // used codium for this function comment, very useful
 
-    let conn_state = conn_id.get::<ConnectionState>(&state.clone())?;
+    /// Terminates the connection of an entity with the given `conn_id`.
+    ///
+    /// Sends a disconnect packet with the given `reason` to the client, and marks the connection as
+    /// terminated. This will cause the connection to be dropped on the next tick of the
+    /// `ConnectionSystem`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the stream writer or connection control component cannot be accessed for
+    /// the given `conn_id`.
+    async fn terminate_connection(
+        &self,
+        state: GlobalState,
+        reason: impl Into<ferrumc_text::TextComponent>,
+    ) -> NetResult<()> {
+        let mut writer = match self.get_mut::<StreamWriter>(&state.clone()) {
+            Ok(writer) => writer,
+            Err(e) => {
+                warn!("Failed to get stream writer for entity {}: {}", self, e);
+                return Err(NetError::ECSError(e));
+            }
+        };
 
-    if let Err(e) = writer
-        .send_packet(
-            &DisconnectPacket::from(&conn_state, reason)?,
-            &NetEncodeOpts::WithLength,
-        )
-        .await
-    {
-        warn!(
-            "Failed to send disconnect packet to entity {}: {}",
-            conn_id, e
-        );
-        return Err(e);
-    }
+        let conn_state = self.get::<ConnectionState>(&state.clone())?;
 
-    match conn_id.get_mut::<ConnectionControl>(&state.clone()) {
-        Ok(mut control) => {
-            control.should_disconnect = true;
-
-            trace!("Set should_disconnect to true for entity {}", conn_id);
-        }
-        Err(e) => {
+        if let Err(e) = writer
+            .send_packet(
+                &DisconnectPacket::from(&conn_state, reason)?,
+                &NetEncodeOpts::WithLength,
+            )
+            .await
+        {
             warn!(
-                "Failed to get connection control for entity {}: {}",
-                conn_id, e
+                "Failed to send disconnect packet to entity {}: {}",
+                self, e
             );
-            return Err(NetError::ECSError(e));
+            return Err(e);
         }
-    }
 
-    Ok(())
+        match self.get_mut::<ConnectionControl>(&state.clone()) {
+            Ok(mut control) => {
+                control.should_disconnect = true;
+
+                trace!("Set should_disconnect to true for entity {}", self);
+            }
+            Err(e) => {
+                warn!(
+                    "Failed to get connection control for entity {}: {}",
+                    self, e
+                );
+                return Err(NetError::ECSError(e));
+            }
+        }
+
+        Ok(())
+    }
 }