From d833582c7bd2b4dcbdc43222a2520105584401de Mon Sep 17 00:00:00 2001 From: Alexander Medvedev Date: Thu, 26 Dec 2024 14:09:18 +0100 Subject: [PATCH] Check byte size on custom payloads --- pumpkin-protocol/src/bytebuf/mod.rs | 23 +++++++++++++++++++ .../src/server/config/s_cookie_response.rs | 11 ++------- .../src/server/config/s_plugin_message.rs | 3 ++- .../src/server/login/s_cookie_response.rs | 11 ++------- .../src/server/login/s_plugin_response.rs | 5 +++- .../src/server/play/s_cookie_response.rs | 11 ++------- 6 files changed, 35 insertions(+), 29 deletions(-) diff --git a/pumpkin-protocol/src/bytebuf/mod.rs b/pumpkin-protocol/src/bytebuf/mod.rs index b1ee75ae..edfa15c7 100644 --- 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; + fn try_copy_to_bytes_len( + &mut self, + len: usize, + max_length: usize, + ) -> Result; + fn try_copy_to_slice(&mut self, dst: &mut [u8]) -> Result<(), ReadingError>; fn try_get_var_int(&mut self) -> Result; @@ -176,6 +182,23 @@ impl ByteBuf for T { } } + fn try_copy_to_bytes_len( + &mut self, + len: usize, + max_size: usize, + ) -> Result { + 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 index b9e3a0a9..9f683a5c 100644 --- 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, // 5120, } -const MAX_PAYLOAD_SIZE: i32 = 5120; +const MAX_COOKIE_LENGTH: usize = 5120; impl ServerPacket for SConfigCookieResponse { fn read(bytebuf: &mut impl Buf) -> Result { @@ -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 index d9e0abae..e5313443 100644 --- 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 { 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 index 5e498a1b..6a6d7383 100644 --- 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, // 5120, } -const MAX_PAYLOAD_SIZE: i32 = 5120; +const MAX_COOKIE_LENGTH: usize = 5120; impl ServerPacket for SLoginCookieResponse { fn read(bytebuf: &mut impl Buf) -> Result { @@ -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 index dd92a1bb..bffc72b9 100644 --- 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 { 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 index 7779d8ac..e8e5d4bb 100644 --- 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, // 5120, } -const MAX_PAYLOAD_SIZE: i32 = 5120; +const MAX_COOKIE_LENGTH: usize = 5120; impl ServerPacket for SCookieResponse { fn read(bytebuf: &mut impl Buf) -> Result { @@ -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,