Improve packable implementations

fuzz/.gitignore

@@ -0,0 +1,3 @@
+target
+corpus
+artifacts

fuzz/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "msgpacker-fuzz"
+version = "0.0.0"
+authors = ["Automatically generated"]
+publish = false
+edition = "2018"
+
+[package.metadata]
+cargo-fuzz = true
+
+[dependencies]
+arbitrary = { version = "1.1", features = ["derive"] }
+libfuzzer-sys = "0.4"
+msgpacker = { path = "../msgpacker" }
+
+[workspace]
+members = ["."]
+
+[[bin]]
+name = "messages"
+path = "fuzz_targets/messages.rs"
+test = false
+doc = false

fuzz/fuzz_targets/messages.rs

@@ -0,0 +1,44 @@
+#![no_main]
+use libfuzzer_sys::fuzz_target;
+
+use arbitrary::Arbitrary;
+use msgpacker::prelude::*;
+
+#[derive(Arbitrary, MsgPacker, Debug, Clone, PartialEq, Eq)]
+pub struct AllMessages {
+    pub uint64: Option<u64>,
+    pub int64: Option<i64>,
+    pub b: Option<bool>,
+    pub f_32: Option<f32>,
+    pub f_64: Option<f64>,
+    pub string: Option<String>,
+    pub bin: Option<Vec<u8>>,
+}
+
+fuzz_target!(|m: AllMessages| {
+    let mut m = m;
+
+    let buf: Vec<u8> = vec![];
+    let mut packer = CursorPacker::new(buf);
+
+    packer.pack(m.clone()).expect("failed to pack message");
+    packer.set_position(0);
+
+    let mut p: AllMessages = packer.unpack().expect("failed to unpack message");
+
+    // NaN equality doesn't hold
+    if m.f_32.filter(|f| f.is_nan()).is_some() {
+        m.f_32.take();
+    }
+    if m.f_64.filter(|f| f.is_nan()).is_some() {
+        m.f_64.take();
+    }
+    if p.f_32.filter(|f| f.is_nan()).is_some() {
+        p.f_32.take();
+    }
+    if p.f_64.filter(|f| f.is_nan()).is_some() {
+        p.f_64.take();
+    }
+
+    assert_eq!(m, p);
+});

msgpacker-derive/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "msgpacker-derive"
-version = "0.1.0"
+version = "0.1.2"
 authors = ["Victor Lopez <victor@codx.io>"]
 categories = ["compression", "encoding", "parser-implementations"]
 edition = "2021"

msgpacker-derive/src/lib.rs

@@ -16,7 +16,7 @@ pub fn msg_packer(input: TokenStream) -> TokenStream {
     let data = input.data;
 
     let mut values: Punctuated<FieldValue, Token![,]> = Punctuated::new();
-    let block: Block = match data {
+    let (block, block_size): (Block, Block) = match data {
         Data::Struct(syn::DataStruct {
             struct_token: _,
             fields: Fields::Named(f),
@@ -25,30 +25,54 @@ pub fn msg_packer(input: TokenStream) -> TokenStream {
             .named
             .into_pairs()
             .map(|p| p.into_value())
-            .fold(syn::parse_str("{}").unwrap(), |mut block, field| {
+            .fold((syn::parse_str("{}").unwrap(), syn::parse_str("{}").unwrap()), |(mut block, mut block_size), field| {
                 let ident = field.ident.as_ref().cloned().unwrap();
                 let ty = field.ty;
 
+                block_size.stmts.push(parse_quote! {
+                    n += <#ty as msgpacker::prelude::SizeableMessage>::packed_len(&self.#ident);
+                });
+
                 block.stmts.push(parse_quote! {
                     n += <#ty as msgpacker::prelude::Packable>::pack(&self.#ident, packer.by_ref())?;
                 });
 
                 let fv = FieldValue {
                     attrs: vec![],
-                    member: Member::Named(ident.clone()),
+                    member: Member::Named(ident),
                     colon_token: Some(<Token![:]>::default()),
                     expr: parse_quote! {
                         <#ty as msgpacker::prelude::Unpackable>::unpack(unpacker.by_ref())?
                     },
                 };
                 values.push(fv);
 
-                block
+                (block, block_size)
             }),
         _ => todo!(),
     };
 
     let expanded = quote! {
+        impl msgpacker::prelude::SizeableMessage for #name {
+            fn packed_len(&self) -> usize {
+                let mut n = 0;
+
+                #block_size
+
+                n
+            }
+        }
+
+        impl<'a> msgpacker::prelude::SizeableMessage for &'a #name {
+            fn packed_len(&self) -> usize {
+                let mut n = 0;
+
+                #block_size
+
+                n
+            }
+        }
+
         impl msgpacker::prelude::Packable for #name {
             fn pack<W>(&self, mut packer: W) -> std::io::Result<usize>
             where
@@ -62,6 +86,19 @@ pub fn msg_packer(input: TokenStream) -> TokenStream {
             }
         }
 
+        impl<'a> msgpacker::prelude::Packable for &'a #name {
+            fn pack<W>(&self, mut packer: W) -> std::io::Result<usize>
+            where
+                W: std::io::Write
+            {
+                let mut n = 0;
+
+                #block
+
+                Ok(n)
+            }
+        }
+
         impl msgpacker::prelude::Unpackable for #name {
             fn unpack<R>(mut unpacker: R) -> std::io::Result<Self>
             where

msgpacker/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "msgpacker"
-version = "0.2.1"
+version = "0.2.2"
 authors = ["Victor Lopez <victor@codx.io>"]
 categories = ["compression", "encoding", "parser-implementations"]
 edition = "2021"

msgpacker/src/buffer.rs

@@ -59,3 +59,18 @@ where
 pub const unsafe fn cast_fixed_array<const M: usize, const N: usize>(array: [u8; M]) -> [u8; N] {
     *mem::transmute::<&[u8; M], &[u8; N]>(&array)
 }
+
+#[test]
+fn take_buf_wont_panic_for_small_buf() {
+    use std::io::Read;
+
+    const LEN: usize = 10;
+
+    let mut cursor = io::Cursor::new([0u8; LEN]);
+
+    let err = unsafe { take_buf(cursor.by_ref(), LEN + 1) }
+        .err()
+        .expect("buffer isn't big enough");
+
+    assert_eq!(io::ErrorKind::UnexpectedEof, err.kind());
+}

msgpacker/src/consts.rs

@@ -0,0 +1,11 @@
+/// Packable representation for variant `None` of [`Option`]
+pub const OPTION_NONE: isize = 0x00;
+
+/// Packable representation for variant `Some` of [`Option`]
+pub const OPTION_SOME: isize = 0x01;
+
+/// Packable representation for variant `Ok` of [`Result`]
+pub const RESULT_OK: isize = 0x00;
+
+/// Packable representation for variant `Err` of [`Result`]
+pub const RESULT_ERR: isize = 0x01;

msgpacker/src/float.rs

@@ -1,5 +1,6 @@
 use crate::buffer;
 use crate::format::MessageFormat;
+use crate::packer::SizeableMessage;
 
 use std::io;
 
@@ -57,6 +58,18 @@ impl Float {
             }
         }
 
+        debug_assert_eq!(n, self.packed_len());
+
         Ok(n)
     }
 }
+
+impl SizeableMessage for Float {
+    fn packed_len(&self) -> usize {
+        match self {
+            Self::F32(_) => 5,
+
+            Self::F64(_) => 9,
+        }
+    }
+}

msgpacker/src/format.rs

@@ -1,3 +1,5 @@
+use crate::packer::SizeableMessage;
+
 /// [specs](https://github.com/msgpack/msgpack/blob/master/spec.md#formats)
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum MessageFormat {
@@ -77,6 +79,12 @@ pub enum MessageFormat {
     NegativeFixInt(i8),
 }
 
+impl SizeableMessage for MessageFormat {
+    fn packed_len(&self) -> usize {
+        1
+    }
+}
+
 impl From<u8> for MessageFormat {
     fn from(b: u8) -> Self {
         use MessageFormat::*;

msgpacker/src/integer.rs

@@ -1,5 +1,6 @@
 use crate::buffer;
 use crate::format::MessageFormat;
+use crate::packer::SizeableMessage;
 
 use std::io;
 
@@ -158,10 +159,48 @@ impl Integer {
             }
         }
 
+        debug_assert_eq!(n, self.packed_len());
+
         Ok(n)
     }
 }
 
+impl SizeableMessage for Integer {
+    fn packed_len(&self) -> usize {
+        match self {
+            Self::Int64(i) if *i <= i32::MIN as i64 => 9,
+
+            Self::Int64(i) if *i <= i16::MIN as i64 => 5,
+
+            Self::Int64(i) if *i <= i8::MIN as i64 => 3,
+
+            Self::Int64(i) if *i <= -33 => 2,
+
+            Self::Int64(i) if *i <= -1 => 1,
+
+            Self::Int64(i) if *i <= 127 => 1,
+
+            Self::Uint64(i) if *i <= 127 => 1,
+
+            Self::Int64(i) if *i <= u8::MAX as i64 => 2,
+
+            Self::Uint64(i) if *i <= u8::MAX as u64 => 2,
+
+            Self::Int64(i) if *i <= u16::MAX as i64 => 3,
+
+            Self::Uint64(i) if *i <= u16::MAX as u64 => 3,
+
+            Self::Int64(i) if *i <= u32::MAX as i64 => 5,
+
+            Self::Uint64(i) if *i <= u32::MAX as u64 => 5,
+
+            Self::Int64(_) => 9,
+
+            Self::Uint64(_) => 9,
+        }
+    }
+}
+
 impl From<Integer> for i64 {
     fn from(i: Integer) -> i64 {
         match i {

msgpacker/src/lib.rs

@@ -11,6 +11,9 @@ mod message;
 mod message_ref;
 mod packer;
 
+/// Preset constants
+pub mod consts;
+
 pub use message::Message;
 pub use message_ref::MessageRef;
 
@@ -27,7 +30,10 @@ pub mod types {
 pub mod prelude {
     pub use crate::message::Message;
     pub use crate::message_ref::MessageRef;
-    pub use crate::packer::{CursorPacker, MessagePacker, MessageUnpacker, Packable, Unpackable};
+    pub use crate::packer::{
+        BufferedUnpacker, CursorPacker, MessagePacker, MessageUnpacker, Packable, SizeableMessage,
+        Unpackable, UnpackableIter,
+    };
     pub use crate::types::*;
 
     #[cfg(feature = "derive")]

msgpacker/src/message.rs

@@ -5,6 +5,7 @@ use crate::format::MessageFormat;
 use crate::integer::Integer;
 use crate::map::MapEntry;
 use crate::message_ref::MessageRef;
+use crate::packer::SizeableMessage;
 
 use std::io;
 use std::ops::Index;
@@ -112,10 +113,7 @@ impl Message {
 
     /// Return `true` if the message is nil
     pub const fn is_nil(&self) -> bool {
-        match self {
-            Self::Nil => true,
-            _ => false,
-        }
+        matches!(self, Self::Nil)
     }
 
     /// Create a new unsigned interger
@@ -215,7 +213,7 @@ impl Message {
 
             MessageFormat::Int8 => {
                 buffer::take(reader, &mut buf, 1)?;
-                let n = Integer::Int64(buf[0] as i64);
+                let n = Integer::Int64((buf[0] as i8) as i64);
 
                 Ok(Self::Integer(n))
             }
@@ -759,6 +757,8 @@ impl Message {
             )),
         }
 
+        debug_assert_eq!(n, self.packed_len());
+
         Ok(n)
     }
 }
@@ -770,11 +770,16 @@ impl<M: Into<Message>> Index<M> for Message {
         let i = i.into();
 
         self.as_map()
-            .map(|m| {
+            .and_then(|m| {
                 m.iter()
                     .find_map(|m| if m.key() == &i { Some(m.val()) } else { None })
             })
-            .flatten()
             .unwrap_or(&Message::Nil)
     }
 }
+
+impl SizeableMessage for Message {
+    fn packed_len(&self) -> usize {
+        self.to_ref().packed_len()
+    }
+}