implement #[limit(i)] attribute for AzBuf derive macro

azalea-rs · Nov 27, 2024 · d047248 · d047248
1 parent f8cd7ed
commit d047248
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 4 deletions.
diff --git a/azalea-buf/azalea-buf-macros/src/lib.rs b/azalea-buf/azalea-buf-macros/src/lib.rs
@@ -19,7 +19,7 @@ pub fn derive_azaleawrite(input: TokenStream) -> TokenStream {
     write::create_impl_azaleawrite(&ident, &data).into()
 }
 
-#[proc_macro_derive(AzBuf, attributes(var))]
+#[proc_macro_derive(AzBuf, attributes(var, limit))]
 pub fn derive_azbuf(input: TokenStream) -> TokenStream {
     let DeriveInput { ident, data, .. } = parse_macro_input!(input);
 

diff --git a/azalea-buf/azalea-buf-macros/src/read.rs b/azalea-buf/azalea-buf-macros/src/read.rs
@@ -9,14 +9,35 @@ fn read_named_fields(
         .map(|f| {
             let field_name = &f.ident;
             let field_type = &f.ty;
+
+            let is_variable_length = f.attrs.iter().any(|a| a.path().is_ident("var"));
+            let limit = f
+                .attrs
+                .iter()
+                .find(|a| a.path().is_ident("limit"))
+                .map(|a| {
+                    a.parse_args::<syn::LitInt>()
+                        .unwrap()
+                        .base10_parse::<usize>()
+                        .unwrap()
+                });
+
+            if is_variable_length && limit.is_some() {
+                panic!("Fields cannot have both var and limit attributes");
+            }
+
             // do a different buf.write_* for each field depending on the type
             // if it's a string, use buf.write_string
             match field_type {
                 syn::Type::Path(_) | syn::Type::Array(_) => {
-                    if f.attrs.iter().any(|a| a.path().is_ident("var")) {
+                    if is_variable_length {
                         quote! {
                             let #field_name = azalea_buf::AzaleaReadVar::azalea_read_var(buf)?;
                         }
+                    } else if let Some(limit) = limit {
+                        quote! {
+                            let #field_name = azalea_buf::AzaleaReadLimited::azalea_read_limited(buf, #limit)?;
+                        }
                     } else {
                         quote! {
                             let #field_name = azalea_buf::AzaleaRead::azalea_read(buf)?;
@@ -108,10 +129,31 @@ pub fn create_impl_azalearead(ident: &Ident, data: &Data) -> proc_macro2::TokenS
                     syn::Fields::Unnamed(fields) => {
                         let mut reader_code = quote! {};
                         for f in &fields.unnamed {
-                            if f.attrs.iter().any(|attr| attr.path().is_ident("var")) {
+                            let is_variable_length =
+                                f.attrs.iter().any(|a| a.path().is_ident("var"));
+                            let limit =
+                                f.attrs
+                                    .iter()
+                                    .find(|a| a.path().is_ident("limit"))
+                                    .map(|a| {
+                                        a.parse_args::<syn::LitInt>()
+                                            .unwrap()
+                                            .base10_parse::<usize>()
+                                            .unwrap()
+                                    });
+
+                            if is_variable_length && limit.is_some() {
+                                panic!("Fields cannot have both var and limit attributes");
+                            }
+
+                            if is_variable_length {
                                 reader_code.extend(quote! {
                                     Self::#variant_name(azalea_buf::AzaleaReadVar::azalea_read_var(buf)?),
                                 });
+                            } else if let Some(limit) = limit {
+                                reader_code.extend(quote! {
+                                    Self::#variant_name(azalea_buf::AzaleaReadLimited::azalea_read_limited(buf, #limit)?),
+                                });
                             } else {
                                 reader_code.extend(quote! {
                                     Self::#variant_name(azalea_buf::AzaleaRead::azalea_read(buf)?),

diff --git a/azalea-buf/src/lib.rs b/azalea-buf/src/lib.rs
@@ -10,7 +10,7 @@ mod write;
 
 pub use azalea_buf_macros::*;
 pub use definitions::*;
-pub use read::{AzaleaRead, AzaleaReadVar, BufReadError};
+pub use read::{AzaleaRead, AzaleaReadLimited, AzaleaReadVar, BufReadError};
 pub use serializable_uuid::*;
 pub use write::{AzaleaWrite, AzaleaWriteVar};
 

diff --git a/azalea-buf/src/read.rs b/azalea-buf/src/read.rs
@@ -119,6 +119,16 @@ where
     fn azalea_read_var(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError>;
 }
 
+// note that there's no Write equivalent for this trait since we don't really
+// care if we're writing over the limit (and maybe we already know that the
+// server implementation accepts it)
+pub trait AzaleaReadLimited
+where
+    Self: Sized,
+{
+    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: usize) -> Result<Self, BufReadError>;
+}
+
 impl AzaleaRead for i32 {
     fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
         Ok(buf.read_i32::<BE>()?)
@@ -219,6 +229,11 @@ impl AzaleaRead for String {
         read_utf_with_len(buf, MAX_STRING_LENGTH.into())
     }
 }
+impl AzaleaReadLimited for String {
+    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: usize) -> Result<Self, BufReadError> {
+        read_utf_with_len(buf, limit as u32)
+    }
+}
 
 impl AzaleaRead for u32 {
     fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {

diff --git a/azalea-protocol/src/packets/login/s_hello.rs b/azalea-protocol/src/packets/login/s_hello.rs
@@ -4,6 +4,7 @@ use uuid::Uuid;
 
 #[derive(Clone, Debug, PartialEq, Eq, AzBuf, ServerboundLoginPacket)]
 pub struct ServerboundHello {
+    #[limit(16)]
     pub name: String,
     pub profile_id: Uuid,
 }