From ff8e3f2d9e4752cf3ecf8ed80041ab352cceb870 Mon Sep 17 00:00:00 2001 From: mat Date: Wed, 15 Dec 2021 19:10:44 -0600 Subject: [PATCH] reading packets is now friendlier --- minecraft-protocol/src/mc_buf.rs | 245 +++++++++--------- .../packets/login/clientbound_hello_packet.rs | 8 +- .../clientbound_status_response_packet.rs | 4 +- minecraft-protocol/src/read.rs | 6 +- 4 files changed, 133 insertions(+), 130 deletions(-) diff --git a/minecraft-protocol/src/mc_buf.rs b/minecraft-protocol/src/mc_buf.rs index 3490a66f..b42a33bb 100644 --- a/minecraft-protocol/src/mc_buf.rs +++ b/minecraft-protocol/src/mc_buf.rs @@ -2,61 +2,22 @@ use std::io::Write; +use async_trait::async_trait; use byteorder::{BigEndian, WriteBytesExt}; -use tokio::io::{AsyncRead, AsyncReadExt, BufReader}; +use tokio::io::{AsyncRead, AsyncReadExt}; // const DEFAULT_NBT_QUOTA: u32 = 2097152; const MAX_STRING_LENGTH: u16 = 32767; // const MAX_COMPONENT_STRING_LENGTH: u32 = 262144; -/// Read a single byte from the reader -pub async fn read_byte( - buf: &mut BufReader, -) -> Result { - match AsyncReadExt::read_u8(buf).await { - Ok(r) => Ok(r), - Err(_) => Err("Error reading byte".to_string()), - } -} - pub fn write_byte(buf: &mut Vec, n: u8) { WriteBytesExt::write_u8(buf, n).unwrap(); } -pub async fn read_bytes( - buf: &mut BufReader, - n: usize, -) -> Result, String> { - let mut bytes = vec![0; n]; - match AsyncReadExt::read_exact(buf, &mut bytes).await { - Ok(_) => Ok(bytes), - Err(_) => Err("Error reading bytes".to_string()), - } -} - pub fn write_bytes(buf: &mut Vec, bytes: &[u8]) { buf.extend_from_slice(bytes); } -// fast varints stolen from https://github.com/luojia65/mc-varint/blob/master/src/lib.rs#L67 -/// Read a single varint from the reader and return the value, along with the number of bytes read -pub async fn read_varint( - buf: &mut BufReader, -) -> Result<(i32, u8), String> { - let mut buffer = [0]; - let mut ans = 0; - for i in 0..4 { - buf.read_exact(&mut buffer) - .await - .map_err(|_| "Invalid VarInt".to_string())?; - ans |= ((buffer[0] & 0b0111_1111) as i32) << (7 * i); - if buffer[0] & 0b1000_0000 == 0 { - return Ok((ans, i + 1)); - } - } - Ok((ans, 5)) -} - pub fn write_varint(buf: &mut Vec, mut value: i32) { let mut buffer = [0]; if value == 0 { @@ -72,10 +33,128 @@ pub fn write_varint(buf: &mut Vec, mut value: i32) { } } +pub fn write_utf_with_len(buf: &mut Vec, string: &str, len: usize) { + if string.len() > len { + panic!( + "String too big (was {} bytes encoded, max {})", + string.len(), + len + ); + } + write_varint(buf, string.len() as i32); + write_bytes(buf, string.as_bytes()); +} +pub fn write_utf(buf: &mut Vec, string: &str) { + write_utf_with_len(buf, string, MAX_STRING_LENGTH.into()); +} + +pub fn write_short(buf: &mut Vec, n: u16) { + WriteBytesExt::write_u16::(buf, n).unwrap(); +} + +pub fn write_byte_array(buf: &mut Vec, bytes: &[u8]) { + write_varint(buf, bytes.len() as i32); + write_bytes(buf, bytes); +} + +#[async_trait] +pub trait Readable { + async fn read_varint(&mut self) -> Result<(i32, u8), String>; + async fn read_byte_array(&mut self) -> Result, String>; + async fn read_bytes(&mut self, n: usize) -> Result, String>; + async fn read_utf(&mut self) -> Result; + async fn read_utf_with_len(&mut self, max_length: u32) -> Result; + async fn read_byte(&mut self) -> Result; +} + +#[async_trait] +impl Readable for R +where + R: AsyncRead + std::marker::Unpin + std::marker::Send, +{ + // fast varints stolen from https://github.com/luojia65/mc-varint/blob/master/src/lib.rs#L67 + /// Read a single varint from the reader and return the value, along with the number of bytes read + async fn read_varint(&mut self) -> Result<(i32, u8), String> { + let mut buffer = [0]; + let mut ans = 0; + for i in 0..4 { + self.read_exact(&mut buffer) + .await + .map_err(|_| "Invalid VarInt".to_string())?; + ans |= ((buffer[0] & 0b0111_1111) as i32) << (7 * i); + if buffer[0] & 0b1000_0000 == 0 { + return Ok((ans, i + 1)); + } + } + Ok((ans, 5)) + } + + async fn read_byte_array(&mut self) -> Result, String> { + let length = self.read_varint().await?.0 as usize; + Ok(self.read_bytes(length).await?) + } + + async fn read_bytes(&mut self, n: usize) -> Result, String> { + let mut bytes = vec![0; n]; + match AsyncReadExt::read_exact(self, &mut bytes).await { + Ok(_) => Ok(bytes), + Err(_) => Err("Error reading bytes".to_string()), + } + } + + async fn read_utf(&mut self) -> Result { + self.read_utf_with_len(MAX_STRING_LENGTH.into()).await + } + + async fn read_utf_with_len(&mut self, max_length: u32) -> Result { + let (length, _length_varint_length) = self.read_varint().await?; + // i don't know why it's multiplied by 4 but it's like that in mojang's code so + if length < 0 { + return Err( + "The received encoded string buffer length is less than zero! Weird string!" + .to_string(), + ); + } + if length as u32 > max_length * 4 { + return Err(format!( + "The received encoded string buffer length is longer than maximum allowed ({} > {})", + length, + max_length * 4 + )); + } + + // this is probably quite inefficient, idk how to do it better + let mut string = String::new(); + let mut buffer = vec![0; length as usize]; + self.read_exact(&mut buffer) + .await + .map_err(|_| "Invalid UTF-8".to_string())?; + + string.push_str(std::str::from_utf8(&buffer).unwrap()); + if string.len() > length as usize { + return Err(format!( + "The received string length is longer than maximum allowed ({} > {})", + length, max_length + )); + } + + Ok(string) + } + + /// Read a single byte from the reader + async fn read_byte(&mut self) -> Result { + match AsyncReadExt::read_u8(self).await { + Ok(r) => Ok(r), + Err(_) => Err("Error reading byte".to_string()), + } + } +} + #[cfg(test)] mod tests { use super::*; use std::io::Cursor; + use tokio::io::BufReader; #[test] fn test_write_varint() { @@ -91,94 +170,18 @@ mod tests { #[tokio::test] async fn test_read_varint() { let mut buf = BufReader::new(Cursor::new(vec![192, 196, 7])); - assert_eq!(read_varint(&mut buf).await.unwrap(), (123456, 3)); + assert_eq!(buf.read_varint().await.unwrap(), (123456, 3)); let mut buf = BufReader::new(Cursor::new(vec![0])); - assert_eq!(read_varint(&mut buf).await.unwrap(), (0, 1)); + assert_eq!(buf.read_varint().await.unwrap(), (0, 1)); let mut buf = BufReader::new(Cursor::new(vec![1])); - assert_eq!(read_varint(&mut buf).await.unwrap(), (1, 1)); + assert_eq!(buf.read_varint().await.unwrap(), (1, 1)); } #[tokio::test] async fn test_read_varint_longer() { let mut buf = BufReader::new(Cursor::new(vec![138, 56, 0, 135, 56, 123])); - assert_eq!(read_varint(&mut buf).await.unwrap(), (7178, 2)); + assert_eq!(buf.read_varint().await.unwrap(), (7178, 2)); } } - -pub async fn read_utf_with_len( - buf: &mut BufReader, - max_length: u32, -) -> Result { - let (length, _length_varint_length) = read_varint(buf).await?; - // i don't know why it's multiplied by 4 but it's like that in mojang's code so - if length < 0 { - return Err( - "The received encoded string buffer length is less than zero! Weird string!" - .to_string(), - ); - } - if length as u32 > max_length * 4 { - return Err(format!( - "The received encoded string buffer length is longer than maximum allowed ({} > {})", - length, - max_length * 4 - )); - } - - // this is probably quite inefficient, idk how to do it better - let mut string = String::new(); - let mut buffer = vec![0; length as usize]; - buf.read_exact(&mut buffer) - .await - .map_err(|_| "Invalid UTF-8".to_string())?; - - string.push_str(std::str::from_utf8(&buffer).unwrap()); - if string.len() > length as usize { - return Err(format!( - "The received string length is longer than maximum allowed ({} > {})", - length, max_length - )); - } - - Ok(string) -} - -pub fn write_utf_with_len(buf: &mut Vec, string: &str, len: usize) { - if string.len() > len { - panic!( - "String too big (was {} bytes encoded, max {})", - string.len(), - len - ); - } - write_varint(buf, string.len() as i32); - write_bytes(buf, string.as_bytes()); -} - -pub async fn read_utf( - buf: &mut BufReader, -) -> Result { - read_utf_with_len(buf, MAX_STRING_LENGTH.into()).await -} - -pub fn write_utf(buf: &mut Vec, string: &str) { - write_utf_with_len(buf, string, MAX_STRING_LENGTH.into()); -} - -pub fn write_short(buf: &mut Vec, n: u16) { - WriteBytesExt::write_u16::(buf, n).unwrap(); -} - -pub async fn read_byte_array( - buf: &mut BufReader, -) -> Result, String> { - let length = read_varint(buf).await?.0 as usize; - Ok(read_bytes(buf, length).await?) -} - -pub fn write_byte_array(buf: &mut Vec, bytes: &[u8]) { - write_varint(buf, bytes.len() as i32); - write_bytes(buf, bytes); -} diff --git a/minecraft-protocol/src/packets/login/clientbound_hello_packet.rs b/minecraft-protocol/src/packets/login/clientbound_hello_packet.rs index 2041497a..36a48706 100644 --- a/minecraft-protocol/src/packets/login/clientbound_hello_packet.rs +++ b/minecraft-protocol/src/packets/login/clientbound_hello_packet.rs @@ -1,7 +1,7 @@ use std::hash::Hash; use tokio::io::BufReader; -use crate::mc_buf; +use crate::mc_buf::Readable; use super::LoginPacket; @@ -24,9 +24,9 @@ impl ClientboundHelloPacket { pub async fn read( buf: &mut BufReader, ) -> Result { - let server_id = mc_buf::read_utf_with_len(buf, 20).await?; - let public_key = mc_buf::read_byte_array(buf).await?; - let nonce = mc_buf::read_byte_array(buf).await?; + let server_id = buf.read_utf_with_len(20).await?; + let public_key = buf.read_byte_array().await?; + let nonce = buf.read_byte_array().await?; Ok(ClientboundHelloPacket { server_id, diff --git a/minecraft-protocol/src/packets/status/clientbound_status_response_packet.rs b/minecraft-protocol/src/packets/status/clientbound_status_response_packet.rs index fc5d8a2a..920e3484 100644 --- a/minecraft-protocol/src/packets/status/clientbound_status_response_packet.rs +++ b/minecraft-protocol/src/packets/status/clientbound_status_response_packet.rs @@ -3,7 +3,7 @@ use serde::Deserialize; use serde_json::Value; use tokio::io::BufReader; -use crate::mc_buf; +use crate::mc_buf::Readable; use super::StatusPacket; @@ -45,7 +45,7 @@ impl ClientboundStatusResponsePacket { pub async fn read( buf: &mut BufReader, ) -> Result { - let status_string = mc_buf::read_utf(buf).await?; + let status_string = buf.read_utf().await?; let status_json: Value = serde_json::from_str(status_string.as_str()).expect("Server status isn't valid JSON"); diff --git a/minecraft-protocol/src/read.rs b/minecraft-protocol/src/read.rs index f0c82641..6f242e8b 100644 --- a/minecraft-protocol/src/read.rs +++ b/minecraft-protocol/src/read.rs @@ -1,6 +1,6 @@ use tokio::{io::BufReader, net::TcpStream}; -use crate::{connect::PacketFlow, mc_buf, packets::ProtocolPacket}; +use crate::{connect::PacketFlow, mc_buf::Readable, packets::ProtocolPacket}; pub async fn read_packet( flow: &PacketFlow, @@ -15,10 +15,10 @@ pub async fn read_packet( // the first thing minecraft sends us is the length as a varint, which can be up to 5 bytes long let mut buf = BufReader::with_capacity(4 * 1024 * 1024, stream); - let (_packet_size, _packet_size_varint_size) = mc_buf::read_varint(&mut buf).await?; + let (_packet_size, _packet_size_varint_size) = buf.read_varint().await?; // then, minecraft tells us the packet id as a varint - let (packet_id, _packet_id_size) = mc_buf::read_varint(&mut buf).await?; + let (packet_id, _packet_id_size) = buf.read_varint().await?; // if we recognize the packet id, parse it