Check byte size on custom payloads

Pumpkin-MC · Dec 26, 2024 · d833582 · d833582
1 parent febbfd8
commit d833582
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 29 deletions.
diff --git a/pumpkin-protocol/src/bytebuf/mod.rs b/pumpkin-protocol/src/bytebuf/mod.rs
@@ -54,6 +54,12 @@ pub trait ByteBuf: Buf {
 
     fn try_copy_to_bytes(&mut self, len: usize) -> Result<bytes::Bytes, ReadingError>;
 
+    fn try_copy_to_bytes_len(
+        &mut self,
+        len: usize,
+        max_length: usize,
+    ) -> Result<bytes::Bytes, ReadingError>;
+
     fn try_copy_to_slice(&mut self, dst: &mut [u8]) -> Result<(), ReadingError>;
 
     fn try_get_var_int(&mut self) -> Result<VarInt, ReadingError>;
@@ -176,6 +182,23 @@ impl<T: Buf> ByteBuf for T {
         }
     }
 
+    fn try_copy_to_bytes_len(
+        &mut self,
+        len: usize,
+        max_size: usize,
+    ) -> Result<bytes::Bytes, ReadingError> {
+        if len > max_size {
+            return Err(ReadingError::Message(
+                "Tried to copy bytes but length exceeds maximum length".to_string(),
+            ));
+        }
+        if self.remaining() >= len {
+            Ok(self.copy_to_bytes(len))
+        } else {
+            Err(ReadingError::Message("Unable to copy bytes".to_string()))
+        }
+    }
+
     fn try_copy_to_slice(&mut self, dst: &mut [u8]) -> Result<(), ReadingError> {
         if self.remaining() >= dst.len() {
             self.copy_to_slice(dst);

diff --git a/pumpkin-protocol/src/server/config/s_cookie_response.rs b/pumpkin-protocol/src/server/config/s_cookie_response.rs
@@ -1,6 +1,5 @@
 use bytes::Buf;
 use pumpkin_macros::server_packet;
-use serde::de;
 
 use crate::{
     bytebuf::{ByteBuf, ReadingError},
@@ -18,7 +17,7 @@ pub struct SConfigCookieResponse {
     pub payload: Option<bytes::Bytes>, // 5120,
 }
 
-const MAX_PAYLOAD_SIZE: i32 = 5120;
+const MAX_COOKIE_LENGTH: usize = 5120;
 
 impl ServerPacket for SConfigCookieResponse {
     fn read(bytebuf: &mut impl Buf) -> Result<Self, ReadingError> {
@@ -37,13 +36,7 @@ impl ServerPacket for SConfigCookieResponse {
         let payload_length = bytebuf.try_get_var_int()?;
         let length = payload_length.0;
 
-        if length > MAX_PAYLOAD_SIZE {
-            return Err(de::Error::custom(
-                "Payload exceeds the maximum allowed size (5120 bytes)",
-            ));
-        }
-
-        let payload = bytebuf.try_copy_to_bytes(length as usize)?;
+        let payload = bytebuf.try_copy_to_bytes_len(length as usize, MAX_COOKIE_LENGTH)?;
 
         Ok(Self {
             key,

diff --git a/pumpkin-protocol/src/server/config/s_plugin_message.rs b/pumpkin-protocol/src/server/config/s_plugin_message.rs
@@ -6,6 +6,7 @@ use crate::{
     codec::identifier::Identifier,
     ServerPacket,
 };
+const MAX_PAYLOAD_SIZE: usize = 1048576;
 
 #[server_packet("config:custom_payload")]
 pub struct SPluginMessage {
@@ -17,7 +18,7 @@ impl ServerPacket for SPluginMessage {
     fn read(bytebuf: &mut impl Buf) -> Result<Self, ReadingError> {
         Ok(Self {
             channel: bytebuf.try_get_identifer()?,
-            data: bytebuf.try_copy_to_bytes(bytebuf.remaining())?,
+            data: bytebuf.try_copy_to_bytes_len(bytebuf.remaining(), MAX_PAYLOAD_SIZE)?,
         })
     }
 }
diff --git a/pumpkin-protocol/src/server/login/s_cookie_response.rs b/pumpkin-protocol/src/server/login/s_cookie_response.rs
@@ -5,7 +5,6 @@ use crate::{
 };
 use bytes::Buf;
 use pumpkin_macros::server_packet;
-use serde::de;
 
 #[server_packet("login:cookie_response")]
 /// Response to a Cookie Request (login) from the server.
@@ -17,7 +16,7 @@ pub struct SLoginCookieResponse {
     pub payload: Option<bytes::Bytes>, // 5120,
 }
 
-const MAX_PAYLOAD_SIZE: i32 = 5120;
+const MAX_COOKIE_LENGTH: usize = 5120;
 
 impl ServerPacket for SLoginCookieResponse {
     fn read(bytebuf: &mut impl Buf) -> Result<Self, ReadingError> {
@@ -36,13 +35,7 @@ impl ServerPacket for SLoginCookieResponse {
         let payload_length = bytebuf.try_get_var_int()?;
         let length = payload_length.0;
 
-        if length > MAX_PAYLOAD_SIZE {
-            return Err(de::Error::custom(
-                "Payload exceeds the maximum allowed size (5120 bytes)",
-            ));
-        }
-
-        let payload = bytebuf.try_copy_to_bytes(length as usize)?;
+        let payload = bytebuf.try_copy_to_bytes_len(length as usize, MAX_COOKIE_LENGTH)?;
 
         Ok(Self {
             key,

diff --git a/pumpkin-protocol/src/server/login/s_plugin_response.rs b/pumpkin-protocol/src/server/login/s_plugin_response.rs
@@ -5,6 +5,8 @@ use crate::{
 use bytes::{Buf, Bytes};
 use pumpkin_macros::server_packet;
 
+const MAX_PAYLOAD_SIZE: usize = 1048576;
+
 #[server_packet("login:custom_query_answer")]
 pub struct SLoginPluginResponse {
     pub message_id: VarInt,
@@ -15,7 +17,8 @@ impl ServerPacket for SLoginPluginResponse {
     fn read(bytebuf: &mut impl Buf) -> Result<Self, ReadingError> {
         Ok(Self {
             message_id: bytebuf.try_get_var_int()?,
-            data: bytebuf.try_get_option(|v| v.try_copy_to_bytes(v.remaining()))?,
+            data: bytebuf
+                .try_get_option(|v| v.try_copy_to_bytes_len(v.remaining(), MAX_PAYLOAD_SIZE))?,
         })
     }
 }
diff --git a/pumpkin-protocol/src/server/play/s_cookie_response.rs b/pumpkin-protocol/src/server/play/s_cookie_response.rs
@@ -5,7 +5,6 @@ use crate::{
 };
 use bytes::Buf;
 use pumpkin_macros::server_packet;
-use serde::de;
 
 #[server_packet("play:cookie_response")]
 /// Response to a Cookie Request (play) from the server.
@@ -17,7 +16,7 @@ pub struct SCookieResponse {
     pub payload: Option<bytes::Bytes>, // 5120,
 }
 
-const MAX_PAYLOAD_SIZE: i32 = 5120;
+const MAX_COOKIE_LENGTH: usize = 5120;
 
 impl ServerPacket for SCookieResponse {
     fn read(bytebuf: &mut impl Buf) -> Result<Self, ReadingError> {
@@ -36,13 +35,7 @@ impl ServerPacket for SCookieResponse {
         let payload_length = bytebuf.try_get_var_int()?;
         let length = payload_length.0;
 
-        if length > MAX_PAYLOAD_SIZE {
-            return Err(de::Error::custom(
-                "Payload exceeds the maximum allowed size (5120 bytes)",
-            ));
-        }
-
-        let payload = bytebuf.try_copy_to_bytes(length as usize)?;
+        let payload = bytebuf.try_copy_to_bytes_len(length as usize, MAX_COOKIE_LENGTH)?;
 
         Ok(Self {
             key,