Check byte size on custom payloads

This commit is contained in:
Alexander Medvedev 2024-12-26 14:09:18 +01:00
parent febbfd89e6
commit d833582c7b
6 changed files with 35 additions and 29 deletions

View File

@ -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);

View File

@ -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,

View File

@ -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)?,
})
}
}

View File

@ -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,

View File

@ -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))?,
})
}
}

View File

@ -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,