From bb57273f48294355d7ac863c291d80878f711c16 Mon Sep 17 00:00:00 2001 From: mat Date: Sun, 2 Jan 2022 17:03:34 -0600 Subject: [PATCH] start improving packet macros --- azalea-protocol/packet-macros/src/lib.rs | 42 +- azalea-protocol/src/mc_buf.rs | 496 ----------------------- azalea-protocol/src/mc_buf/mod.rs | 163 ++++++++ azalea-protocol/src/mc_buf/read.rs | 230 +++++++++++ azalea-protocol/src/mc_buf/write.rs | 180 ++++++++ 5 files changed, 592 insertions(+), 519 deletions(-) delete mode 100644 azalea-protocol/src/mc_buf.rs create mode 100644 azalea-protocol/src/mc_buf/mod.rs create mode 100644 azalea-protocol/src/mc_buf/read.rs create mode 100644 azalea-protocol/src/mc_buf/write.rs diff --git a/azalea-protocol/packet-macros/src/lib.rs b/azalea-protocol/packet-macros/src/lib.rs index 11878abf..22c39ea9 100644 --- a/azalea-protocol/packet-macros/src/lib.rs +++ b/azalea-protocol/packet-macros/src/lib.rs @@ -32,13 +32,6 @@ fn as_packet_derive( // i don't know how to do this in a way that isn't terrible } else if path.to_token_stream().to_string() == "Vec < u8 >" { quote! { buf.write_bytes(&self.#field_name)?; } - } else if path.is_ident("i32") { - // only treat it as a varint if it has the varint attribute - if f.attrs.iter().any(|attr| attr.path.is_ident("varint")) { - quote! { buf.write_varint(self.#field_name)?; } - } else { - quote! { buf.write_i32(self.#field_name)?; } - } } else if path.is_ident("u32") { if f.attrs.iter().any(|attr| attr.path.is_ident("varint")) { quote! { buf.write_varint(self.#field_name as i32)?; } @@ -50,10 +43,16 @@ fn as_packet_derive( } else if path.is_ident("ConnectionProtocol") { quote! { buf.write_varint(self.#field_name.clone() as i32)?; } } else { - panic!( - "#[derive(*Packet)] doesn't know how to write {}", - path.to_token_stream() - ); + if f.attrs.iter().any(|attr| attr.path.is_ident("varint")) { + quote! { + crate::mc_buf::McBufVarintWritable::varint_write_into(&self.#field_name, buf)?; + } + } else { + quote! { + crate::mc_buf::McBufVarintWritable::write_into(&self.#field_name, buf)?; + } + } + } } _ => panic!( @@ -78,16 +77,8 @@ fn as_packet_derive( quote! { let #field_name = buf.read_utf().await?; } } else if path.is_ident("ResourceLocation") { quote! { let #field_name = buf.read_resource_location().await?; } - // i don't know how to do this in a way that isn't terrible } else if path.to_token_stream().to_string() == "Vec < u8 >" { quote! { let #field_name = buf.read_bytes().await?; } - } else if path.is_ident("i32") { - // only treat it as a varint if it has the varint attribute - if f.attrs.iter().any(|a| a.path.is_ident("varint")) { - quote! { let #field_name = buf.read_varint().await?; } - } else { - quote! { let #field_name = buf.read_i32().await?; } - } } else if path.is_ident("u32") { if f.attrs.iter().any(|a| a.path.is_ident("varint")) { quote! { let #field_name = buf.read_varint().await? as u32; } @@ -102,10 +93,15 @@ fn as_packet_derive( .ok_or_else(|| "Invalid intention".to_string())?; } } else { - panic!( - "#[derive(*Packet)] doesn't know how to read {}", - path.to_token_stream() - ); + if f.attrs.iter().any(|a| a.path.is_ident("varint")) { + quote! { + let #field_name = crate::mc_buf::McBufVarintReadable::varint_read_into(buf).await?; + } + } else { + quote! { + let #field_name = crate::mc_buf::McBufReadable::read_into(buf).await?; + } + } } } _ => panic!( diff --git a/azalea-protocol/src/mc_buf.rs b/azalea-protocol/src/mc_buf.rs deleted file mode 100644 index 72583d5a..00000000 --- a/azalea-protocol/src/mc_buf.rs +++ /dev/null @@ -1,496 +0,0 @@ -//! Utilities for reading and writing for the Minecraft protocol - -use std::io::Write; - -use async_trait::async_trait; -use azalea_core::resource_location::ResourceLocation; -use byteorder::{BigEndian, WriteBytesExt}; -use tokio::io::{AsyncRead, AsyncReadExt}; - -// const DEFAULT_NBT_QUOTA: u32 = 2097152; -const MAX_STRING_LENGTH: u16 = 32767; -// const MAX_COMPONENT_STRING_LENGTH: u32 = 262144; - -#[async_trait] -pub trait Writable { - fn write_list(&mut self, list: &[T], writer: F) -> Result<(), std::io::Error> - where - F: FnOnce(&mut Self, &T) -> Result<(), std::io::Error> + Copy, - T: Sized, - Self: Sized; - fn write_int_id_list(&mut self, list: Vec) -> Result<(), std::io::Error>; - fn write_map( - &mut self, - map: Vec<(KT, VT)>, - key_writer: KF, - value_writer: VF, - ) -> Result<(), std::io::Error> - where - KF: Fn(&mut Self, KT) -> Result<(), std::io::Error> + Copy, - VF: Fn(&mut Self, VT) -> Result<(), std::io::Error> + Copy, - Self: Sized; - - fn write_byte(&mut self, n: u8) -> Result<(), std::io::Error>; - fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), std::io::Error>; - fn write_varint(&mut self, value: i32) -> Result<(), std::io::Error>; - fn write_utf_with_len(&mut self, string: &str, len: usize) -> Result<(), std::io::Error>; - fn write_utf(&mut self, string: &str) -> Result<(), std::io::Error>; - fn write_short(&mut self, n: i16) -> Result<(), std::io::Error>; - fn write_byte_array(&mut self, bytes: &[u8]) -> Result<(), std::io::Error>; - fn write_int(&mut self, n: i32) -> Result<(), std::io::Error>; - fn write_boolean(&mut self, b: bool) -> Result<(), std::io::Error>; - fn write_nbt(&mut self, nbt: &azalea_nbt::Tag) -> Result<(), std::io::Error>; - fn write_long(&mut self, n: i64) -> Result<(), std::io::Error>; - fn write_resource_location( - &mut self, - location: &ResourceLocation, - ) -> Result<(), std::io::Error>; -} - -#[async_trait] -impl Writable for Vec { - fn write_list(&mut self, list: &[T], writer: F) -> Result<(), std::io::Error> - where - F: FnOnce(&mut Self, &T) -> Result<(), std::io::Error> + Copy, - Self: Sized, - { - self.write_varint(list.len() as i32)?; - for item in list { - writer(self, item)?; - } - Ok(()) - } - - fn write_int_id_list(&mut self, list: Vec) -> Result<(), std::io::Error> { - self.write_list(&list, |buf, n| buf.write_varint(*n)) - } - - fn write_map( - &mut self, - map: Vec<(KT, VT)>, - key_writer: KF, - value_writer: VF, - ) -> Result<(), std::io::Error> - where - KF: Fn(&mut Self, KT) -> Result<(), std::io::Error> + Copy, - VF: Fn(&mut Self, VT) -> Result<(), std::io::Error> + Copy, - Self: Sized, - { - self.write_varint(map.len() as i32)?; - for (key, value) in map { - key_writer(self, key)?; - value_writer(self, value)?; - } - Ok(()) - } - - fn write_byte(&mut self, n: u8) -> Result<(), std::io::Error> { - WriteBytesExt::write_u8(self, n) - } - - fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), std::io::Error> { - self.extend_from_slice(bytes); - Ok(()) - } - - fn write_varint(&mut self, mut value: i32) -> Result<(), std::io::Error> { - let mut buffer = [0]; - if value == 0 { - self.write_all(&buffer).unwrap(); - } - while value != 0 { - buffer[0] = (value & 0b0111_1111) as u8; - value = (value >> 7) & (i32::max_value() >> 6); - if value != 0 { - buffer[0] |= 0b1000_0000; - } - self.write_all(&buffer)?; - } - Ok(()) - } - - fn write_utf_with_len(&mut self, string: &str, len: usize) -> Result<(), std::io::Error> { - if string.len() > len { - panic!( - "String too big (was {} bytes encoded, max {})", - string.len(), - len - ); - } - self.write_varint(string.len() as i32)?; - self.write_bytes(string.as_bytes()) - } - - fn write_utf(&mut self, string: &str) -> Result<(), std::io::Error> { - self.write_utf_with_len(string, MAX_STRING_LENGTH.into()) - } - - fn write_short(&mut self, n: i16) -> Result<(), std::io::Error> { - WriteBytesExt::write_i16::(self, n) - } - - fn write_byte_array(&mut self, bytes: &[u8]) -> Result<(), std::io::Error> { - self.write_varint(bytes.len() as i32)?; - self.write_bytes(bytes) - } - - fn write_int(&mut self, n: i32) -> Result<(), std::io::Error> { - WriteBytesExt::write_i32::(self, n) - } - - fn write_boolean(&mut self, b: bool) -> Result<(), std::io::Error> { - self.write_byte(if b { 1 } else { 0 }) - } - - fn write_nbt(&mut self, nbt: &azalea_nbt::Tag) -> Result<(), std::io::Error> { - nbt.write(self) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string())) - } - - fn write_long(&mut self, n: i64) -> Result<(), std::io::Error> { - WriteBytesExt::write_i64::(self, n) - } - - fn write_resource_location( - &mut self, - location: &ResourceLocation, - ) -> Result<(), std::io::Error> { - self.write_utf(&location.to_string()) - } -} - -#[async_trait] -pub trait Readable { - async fn read_int_id_list(&mut self) -> Result, String>; - async fn read_varint(&mut self) -> Result; - fn get_varint_size(&mut self, value: i32) -> u8; - fn get_varlong_size(&mut self, value: i32) -> u8; - async fn read_byte_array(&mut self) -> Result, String>; - async fn read_bytes_with_len(&mut self, n: usize) -> Result, String>; - async fn read_bytes(&mut self) -> 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 fn read_int(&mut self) -> Result; - async fn read_boolean(&mut self) -> Result; - async fn read_nbt(&mut self) -> Result; - async fn read_long(&mut self) -> Result; - async fn read_resource_location(&mut self) -> Result; - async fn read_short(&mut self) -> Result; -} - -#[async_trait] -impl Readable for R -where - R: AsyncRead + std::marker::Unpin + std::marker::Send, -{ - async fn read_int_id_list(&mut self) -> Result, String> { - let len = self.read_varint().await?; - let mut list = Vec::with_capacity(len as usize); - for _ in 0..len { - list.push(self.read_varint().await?); - } - Ok(list) - } - - // 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 { - 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); - } - } - Ok(ans) - } - - fn get_varint_size(&mut self, value: i32) -> u8 { - for i in 1..5 { - if (value & -1 << (i * 7)) != 0 { - continue; - } - return i; - } - 5 - } - - fn get_varlong_size(&mut self, value: i32) -> u8 { - for i in 1..10 { - if (value & -1 << (i * 7)) != 0 { - continue; - } - return i; - } - 10 - } - - async fn read_byte_array(&mut self) -> Result, String> { - let length = self.read_varint().await? as usize; - Ok(self.read_bytes_with_len(length).await?) - } - - async fn read_bytes_with_len(&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_bytes(&mut self) -> Result, String> { - // read to end of the buffer - let mut bytes = vec![]; - AsyncReadExt::read_to_end(self, &mut bytes) - .await - .map_err(|_| "Error reading bytes".to_string())?; - Ok(bytes) - } - - 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 = 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()), - } - } - - async fn read_int(&mut self) -> Result { - match AsyncReadExt::read_i32(self).await { - Ok(r) => Ok(r), - Err(_) => Err("Error reading int".to_string()), - } - } - - async fn read_boolean(&mut self) -> Result { - match self.read_byte().await { - Ok(0) => Ok(false), - Ok(1) => Ok(true), - _ => Err("Error reading boolean".to_string()), - } - } - - async fn read_nbt(&mut self) -> Result { - match azalea_nbt::Tag::read(self).await { - Ok(r) => Ok(r), - // Err(e) => Err(e.to_string()), - Err(e) => Err(e.to_string()).unwrap(), - } - } - - async fn read_long(&mut self) -> Result { - match AsyncReadExt::read_i64(self).await { - Ok(r) => Ok(r), - Err(_) => Err("Error reading long".to_string()), - } - } - - async fn read_resource_location(&mut self) -> Result { - // get the resource location from the string - let location_string = self.read_utf().await?; - let location = ResourceLocation::new(&location_string)?; - Ok(location) - } - - async fn read_short(&mut self) -> Result { - match AsyncReadExt::read_i16(self).await { - Ok(r) => Ok(r), - Err(_) => Err("Error reading short".to_string()), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::{collections::HashMap, io::Cursor}; - use tokio::io::BufReader; - - #[test] - fn test_write_varint() { - let mut buf = Vec::new(); - buf.write_varint(123456).unwrap(); - assert_eq!(buf, vec![192, 196, 7]); - - let mut buf = Vec::new(); - buf.write_varint(0).unwrap(); - assert_eq!(buf, vec![0]); - } - - #[tokio::test] - async fn test_read_varint() { - let mut buf = BufReader::new(Cursor::new(vec![192, 196, 7])); - assert_eq!(buf.read_varint().await.unwrap(), 123456); - assert_eq!(buf.get_varint_size(123456), 3); - - let mut buf = BufReader::new(Cursor::new(vec![0])); - assert_eq!(buf.read_varint().await.unwrap(), 0); - assert_eq!(buf.get_varint_size(0), 1); - - let mut buf = BufReader::new(Cursor::new(vec![1])); - assert_eq!(buf.read_varint().await.unwrap(), 1); - assert_eq!(buf.get_varint_size(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!(buf.read_varint().await.unwrap(), 7178); - } - - #[tokio::test] - async fn test_list() { - let mut buf = Vec::new(); - buf.write_list(&vec!["a", "bc", "def"], |buf, s| buf.write_utf(s)) - .unwrap(); - - // there's no read_list because idk how to do it in rust - let mut buf = BufReader::new(Cursor::new(buf)); - - let mut result = Vec::new(); - let length = buf.read_varint().await.unwrap(); - for _ in 0..length { - result.push(buf.read_utf().await.unwrap()); - } - - assert_eq!(result, vec!["a", "bc", "def"]); - } - - #[tokio::test] - async fn test_int_id_list() { - let mut buf = Vec::new(); - buf.write_list(&vec![1, 2, 3], |buf, i| buf.write_varint(*i)) - .unwrap(); - - let mut buf = BufReader::new(Cursor::new(buf)); - - let result = buf.read_int_id_list().await.unwrap(); - assert_eq!(result, vec![1, 2, 3]); - } - - #[tokio::test] - async fn test_map() { - let mut buf = Vec::new(); - buf.write_map( - vec![("a", 1), ("bc", 23), ("def", 456)], - Vec::write_utf, - Vec::write_varint, - ) - .unwrap(); - - let mut buf = BufReader::new(Cursor::new(buf)); - - let mut result = Vec::new(); - let length = buf.read_varint().await.unwrap(); - for _ in 0..length { - result.push(( - buf.read_utf().await.unwrap(), - buf.read_varint().await.unwrap(), - )); - } - - assert_eq!( - result, - vec![ - ("a".to_string(), 1), - ("bc".to_string(), 23), - ("def".to_string(), 456) - ] - ); - } - - #[tokio::test] - async fn test_nbt() { - let mut buf = Vec::new(); - buf.write_nbt(&azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( - "hello world".to_string(), - azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( - "name".to_string(), - azalea_nbt::Tag::String("Bananrama".to_string()), - )])), - )]))) - .unwrap(); - - let mut buf = BufReader::new(Cursor::new(buf)); - - let result = buf.read_nbt().await.unwrap(); - assert_eq!( - result, - azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( - "hello world".to_string(), - azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( - "name".to_string(), - azalea_nbt::Tag::String("Bananrama".to_string()), - )])), - )])) - ); - } - - #[tokio::test] - async fn test_long() { - let mut buf = Vec::new(); - buf.write_long(123456).unwrap(); - - let mut buf = BufReader::new(Cursor::new(buf)); - - assert_eq!(buf.read_long().await.unwrap(), 123456); - } - - #[tokio::test] - async fn test_resource_location() { - let mut buf = Vec::new(); - buf.write_resource_location(&ResourceLocation::new("minecraft:dirt").unwrap()) - .unwrap(); - - let mut buf = BufReader::new(Cursor::new(buf)); - - assert_eq!( - buf.read_resource_location().await.unwrap(), - ResourceLocation::new("minecraft:dirt").unwrap() - ); - } -} diff --git a/azalea-protocol/src/mc_buf/mod.rs b/azalea-protocol/src/mc_buf/mod.rs new file mode 100644 index 00000000..a924431e --- /dev/null +++ b/azalea-protocol/src/mc_buf/mod.rs @@ -0,0 +1,163 @@ +//! Utilities for reading and writing for the Minecraft protocol + +mod read; +mod write; + +pub use read::{McBufReadable, McBufVarintReadable, Readable}; +pub use write::{McBufVarintWritable, McBufWritable, Writable}; + +// const DEFAULT_NBT_QUOTA: u32 = 2097152; +const MAX_STRING_LENGTH: u16 = 32767; +// const MAX_COMPONENT_STRING_LENGTH: u32 = 262144; + +#[cfg(test)] +mod tests { + use super::*; + use azalea_core::resource_location::ResourceLocation; + use std::{collections::HashMap, io::Cursor}; + use tokio::io::BufReader; + + #[test] + fn test_write_varint() { + let mut buf = Vec::new(); + buf.write_varint(123456).unwrap(); + assert_eq!(buf, vec![192, 196, 7]); + + let mut buf = Vec::new(); + buf.write_varint(0).unwrap(); + assert_eq!(buf, vec![0]); + } + + #[tokio::test] + async fn test_read_varint() { + let mut buf = BufReader::new(Cursor::new(vec![192, 196, 7])); + assert_eq!(buf.read_varint().await.unwrap(), 123456); + assert_eq!(buf.get_varint_size(123456), 3); + + let mut buf = BufReader::new(Cursor::new(vec![0])); + assert_eq!(buf.read_varint().await.unwrap(), 0); + assert_eq!(buf.get_varint_size(0), 1); + + let mut buf = BufReader::new(Cursor::new(vec![1])); + assert_eq!(buf.read_varint().await.unwrap(), 1); + assert_eq!(buf.get_varint_size(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!(buf.read_varint().await.unwrap(), 7178); + } + + #[tokio::test] + async fn test_list() { + let mut buf = Vec::new(); + buf.write_list(&vec!["a", "bc", "def"], |buf, s| buf.write_utf(s)) + .unwrap(); + + // there's no read_list because idk how to do it in rust + let mut buf = BufReader::new(Cursor::new(buf)); + + let mut result = Vec::new(); + let length = buf.read_varint().await.unwrap(); + for _ in 0..length { + result.push(buf.read_utf().await.unwrap()); + } + + assert_eq!(result, vec!["a", "bc", "def"]); + } + + #[tokio::test] + async fn test_int_id_list() { + let mut buf = Vec::new(); + buf.write_list(&vec![1, 2, 3], |buf, i| buf.write_varint(*i)) + .unwrap(); + + let mut buf = BufReader::new(Cursor::new(buf)); + + let result = buf.read_int_id_list().await.unwrap(); + assert_eq!(result, vec![1, 2, 3]); + } + + #[tokio::test] + async fn test_map() { + let mut buf = Vec::new(); + buf.write_map( + vec![("a", 1), ("bc", 23), ("def", 456)], + Vec::write_utf, + Vec::write_varint, + ) + .unwrap(); + + let mut buf = BufReader::new(Cursor::new(buf)); + + let mut result = Vec::new(); + let length = buf.read_varint().await.unwrap(); + for _ in 0..length { + result.push(( + buf.read_utf().await.unwrap(), + buf.read_varint().await.unwrap(), + )); + } + + assert_eq!( + result, + vec![ + ("a".to_string(), 1), + ("bc".to_string(), 23), + ("def".to_string(), 456) + ] + ); + } + + #[tokio::test] + async fn test_nbt() { + let mut buf = Vec::new(); + buf.write_nbt(&azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( + "hello world".to_string(), + azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( + "name".to_string(), + azalea_nbt::Tag::String("Bananrama".to_string()), + )])), + )]))) + .unwrap(); + + let mut buf = BufReader::new(Cursor::new(buf)); + + let result = buf.read_nbt().await.unwrap(); + assert_eq!( + result, + azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( + "hello world".to_string(), + azalea_nbt::Tag::Compound(HashMap::from_iter(vec![( + "name".to_string(), + azalea_nbt::Tag::String("Bananrama".to_string()), + )])), + )])) + ); + } + + #[tokio::test] + async fn test_long() { + let mut buf = Vec::new(); + buf.write_long(123456).unwrap(); + + let mut buf = BufReader::new(Cursor::new(buf)); + + assert_eq!(buf.read_long().await.unwrap(), 123456); + } + + #[tokio::test] + async fn test_resource_location() { + let mut buf = Vec::new(); + buf.write_resource_location(&ResourceLocation::new("minecraft:dirt").unwrap()) + .unwrap(); + + let mut buf = BufReader::new(Cursor::new(buf)); + + assert_eq!( + buf.read_resource_location().await.unwrap(), + ResourceLocation::new("minecraft:dirt").unwrap() + ); + } +} diff --git a/azalea-protocol/src/mc_buf/read.rs b/azalea-protocol/src/mc_buf/read.rs new file mode 100644 index 00000000..c429eb7f --- /dev/null +++ b/azalea-protocol/src/mc_buf/read.rs @@ -0,0 +1,230 @@ +use async_trait::async_trait; +use azalea_core::resource_location::ResourceLocation; +use tokio::io::{AsyncRead, AsyncReadExt}; + +use super::MAX_STRING_LENGTH; + +#[async_trait] +pub trait Readable { + async fn read_int_id_list(&mut self) -> Result, String>; + async fn read_varint(&mut self) -> Result; + fn get_varint_size(&mut self, value: i32) -> u8; + fn get_varlong_size(&mut self, value: i32) -> u8; + async fn read_byte_array(&mut self) -> Result, String>; + async fn read_bytes_with_len(&mut self, n: usize) -> Result, String>; + async fn read_bytes(&mut self) -> 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 fn read_int(&mut self) -> Result; + async fn read_boolean(&mut self) -> Result; + async fn read_nbt(&mut self) -> Result; + async fn read_long(&mut self) -> Result; + async fn read_resource_location(&mut self) -> Result; + async fn read_short(&mut self) -> Result; +} + +#[async_trait] +impl Readable for R +where + R: AsyncRead + std::marker::Unpin + std::marker::Send, +{ + async fn read_int_id_list(&mut self) -> Result, String> { + let len = self.read_varint().await?; + let mut list = Vec::with_capacity(len as usize); + for _ in 0..len { + list.push(self.read_varint().await?); + } + Ok(list) + } + + // 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 { + 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); + } + } + Ok(ans) + } + + fn get_varint_size(&mut self, value: i32) -> u8 { + for i in 1..5 { + if (value & -1 << (i * 7)) != 0 { + continue; + } + return i; + } + 5 + } + + fn get_varlong_size(&mut self, value: i32) -> u8 { + for i in 1..10 { + if (value & -1 << (i * 7)) != 0 { + continue; + } + return i; + } + 10 + } + + async fn read_byte_array(&mut self) -> Result, String> { + let length = self.read_varint().await? as usize; + Ok(self.read_bytes_with_len(length).await?) + } + + async fn read_bytes_with_len(&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_bytes(&mut self) -> Result, String> { + // read to end of the buffer + let mut bytes = vec![]; + AsyncReadExt::read_to_end(self, &mut bytes) + .await + .map_err(|_| "Error reading bytes".to_string())?; + Ok(bytes) + } + + 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 = 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()), + } + } + + async fn read_int(&mut self) -> Result { + match AsyncReadExt::read_i32(self).await { + Ok(r) => Ok(r), + Err(_) => Err("Error reading int".to_string()), + } + } + + async fn read_boolean(&mut self) -> Result { + match self.read_byte().await { + Ok(0) => Ok(false), + Ok(1) => Ok(true), + _ => Err("Error reading boolean".to_string()), + } + } + + async fn read_nbt(&mut self) -> Result { + match azalea_nbt::Tag::read(self).await { + Ok(r) => Ok(r), + // Err(e) => Err(e.to_string()), + Err(e) => Err(e.to_string()).unwrap(), + } + } + + async fn read_long(&mut self) -> Result { + match AsyncReadExt::read_i64(self).await { + Ok(r) => Ok(r), + Err(_) => Err("Error reading long".to_string()), + } + } + + async fn read_resource_location(&mut self) -> Result { + // get the resource location from the string + let location_string = self.read_utf().await?; + let location = ResourceLocation::new(&location_string)?; + Ok(location) + } + + async fn read_short(&mut self) -> Result { + match AsyncReadExt::read_i16(self).await { + Ok(r) => Ok(r), + Err(_) => Err("Error reading short".to_string()), + } + } +} + +#[async_trait] +pub trait McBufReadable +where + Self: Sized, +{ + async fn read_into(buf: &mut R) -> Result + where + R: AsyncRead + std::marker::Unpin + std::marker::Send; +} + +#[async_trait] +pub trait McBufVarintReadable +where + Self: Sized, +{ + async fn varint_read_into(buf: &mut R) -> Result + where + R: AsyncRead + std::marker::Unpin + std::marker::Send; +} + +#[async_trait] +impl McBufReadable for i32 { + async fn read_into(buf: &mut R) -> Result + where + R: AsyncRead + std::marker::Unpin + std::marker::Send, + { + buf.read_int().await + } +} + +#[async_trait] +impl McBufVarintReadable for i32 { + async fn varint_read_into(buf: &mut R) -> Result + where + R: AsyncRead + std::marker::Unpin + std::marker::Send, + { + buf.read_varint().await + } +} diff --git a/azalea-protocol/src/mc_buf/write.rs b/azalea-protocol/src/mc_buf/write.rs new file mode 100644 index 00000000..3c3375b9 --- /dev/null +++ b/azalea-protocol/src/mc_buf/write.rs @@ -0,0 +1,180 @@ +use async_trait::async_trait; +use azalea_core::resource_location::ResourceLocation; +use byteorder::{BigEndian, WriteBytesExt}; +use std::io::Write; + +use super::MAX_STRING_LENGTH; + +#[async_trait] +pub trait Writable { + fn write_list(&mut self, list: &[T], writer: F) -> Result<(), std::io::Error> + where + F: FnOnce(&mut Self, &T) -> Result<(), std::io::Error> + Copy, + T: Sized, + Self: Sized; + fn write_int_id_list(&mut self, list: Vec) -> Result<(), std::io::Error>; + fn write_map( + &mut self, + map: Vec<(KT, VT)>, + key_writer: KF, + value_writer: VF, + ) -> Result<(), std::io::Error> + where + KF: Fn(&mut Self, KT) -> Result<(), std::io::Error> + Copy, + VF: Fn(&mut Self, VT) -> Result<(), std::io::Error> + Copy, + Self: Sized; + + fn write_byte(&mut self, n: u8) -> Result<(), std::io::Error>; + fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), std::io::Error>; + fn write_varint(&mut self, value: i32) -> Result<(), std::io::Error>; + fn write_utf_with_len(&mut self, string: &str, len: usize) -> Result<(), std::io::Error>; + fn write_utf(&mut self, string: &str) -> Result<(), std::io::Error>; + fn write_short(&mut self, n: i16) -> Result<(), std::io::Error>; + fn write_byte_array(&mut self, bytes: &[u8]) -> Result<(), std::io::Error>; + fn write_int(&mut self, n: i32) -> Result<(), std::io::Error>; + fn write_boolean(&mut self, b: bool) -> Result<(), std::io::Error>; + fn write_nbt(&mut self, nbt: &azalea_nbt::Tag) -> Result<(), std::io::Error>; + fn write_long(&mut self, n: i64) -> Result<(), std::io::Error>; + fn write_resource_location( + &mut self, + location: &ResourceLocation, + ) -> Result<(), std::io::Error>; +} + +#[async_trait] +impl Writable for Vec { + fn write_list(&mut self, list: &[T], writer: F) -> Result<(), std::io::Error> + where + F: FnOnce(&mut Self, &T) -> Result<(), std::io::Error> + Copy, + Self: Sized, + { + self.write_varint(list.len() as i32)?; + for item in list { + writer(self, item)?; + } + Ok(()) + } + + fn write_int_id_list(&mut self, list: Vec) -> Result<(), std::io::Error> { + self.write_list(&list, |buf, n| buf.write_varint(*n)) + } + + fn write_map( + &mut self, + map: Vec<(KT, VT)>, + key_writer: KF, + value_writer: VF, + ) -> Result<(), std::io::Error> + where + KF: Fn(&mut Self, KT) -> Result<(), std::io::Error> + Copy, + VF: Fn(&mut Self, VT) -> Result<(), std::io::Error> + Copy, + Self: Sized, + { + self.write_varint(map.len() as i32)?; + for (key, value) in map { + key_writer(self, key)?; + value_writer(self, value)?; + } + Ok(()) + } + + fn write_byte(&mut self, n: u8) -> Result<(), std::io::Error> { + WriteBytesExt::write_u8(self, n) + } + + fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), std::io::Error> { + self.extend_from_slice(bytes); + Ok(()) + } + + fn write_varint(&mut self, mut value: i32) -> Result<(), std::io::Error> { + let mut buffer = [0]; + if value == 0 { + self.write_all(&buffer).unwrap(); + } + while value != 0 { + buffer[0] = (value & 0b0111_1111) as u8; + value = (value >> 7) & (i32::max_value() >> 6); + if value != 0 { + buffer[0] |= 0b1000_0000; + } + self.write_all(&buffer)?; + } + Ok(()) + } + + fn write_utf_with_len(&mut self, string: &str, len: usize) -> Result<(), std::io::Error> { + if string.len() > len { + panic!( + "String too big (was {} bytes encoded, max {})", + string.len(), + len + ); + } + self.write_varint(string.len() as i32)?; + self.write_bytes(string.as_bytes()) + } + + fn write_utf(&mut self, string: &str) -> Result<(), std::io::Error> { + self.write_utf_with_len(string, MAX_STRING_LENGTH.into()) + } + + fn write_short(&mut self, n: i16) -> Result<(), std::io::Error> { + WriteBytesExt::write_i16::(self, n) + } + + fn write_byte_array(&mut self, bytes: &[u8]) -> Result<(), std::io::Error> { + self.write_varint(bytes.len() as i32)?; + self.write_bytes(bytes) + } + + fn write_int(&mut self, n: i32) -> Result<(), std::io::Error> { + WriteBytesExt::write_i32::(self, n) + } + + fn write_boolean(&mut self, b: bool) -> Result<(), std::io::Error> { + self.write_byte(if b { 1 } else { 0 }) + } + + fn write_nbt(&mut self, nbt: &azalea_nbt::Tag) -> Result<(), std::io::Error> { + nbt.write(self) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e.to_string())) + } + + fn write_long(&mut self, n: i64) -> Result<(), std::io::Error> { + WriteBytesExt::write_i64::(self, n) + } + + fn write_resource_location( + &mut self, + location: &ResourceLocation, + ) -> Result<(), std::io::Error> { + self.write_utf(&location.to_string()) + } +} + +pub trait McBufWritable +where + Self: Sized, +{ + fn write_into(&self, buf: &mut Vec) -> Result<(), std::io::Error>; +} + +pub trait McBufVarintWritable +where + Self: Sized, +{ + fn varint_write_into(&self, buf: &mut Vec) -> Result<(), std::io::Error>; +} + +impl McBufWritable for i32 { + fn write_into(&self, buf: &mut Vec) -> Result<(), std::io::Error> { + Writable::write_int(buf, *self) + } +} + +impl McBufVarintWritable for i32 { + fn varint_write_into(&self, buf: &mut Vec) -> Result<(), std::io::Error> { + buf.write_varint(*self) + } +}