From 75f36c4a49a53dbd5a07e9e9aa9a968dc9be070c Mon Sep 17 00:00:00 2001 From: armaniferrante Date: Thu, 24 Sep 2020 23:43:49 -0700 Subject: [PATCH] common: Change unpack_unchecked api --- common/src/pack.rs | 84 ++++++++++++++++++++++++++++++---------------- 1 file changed, 56 insertions(+), 28 deletions(-) diff --git a/common/src/pack.rs b/common/src/pack.rs index dd907fe..aa26258 100644 --- a/common/src/pack.rs +++ b/common/src/pack.rs @@ -11,13 +11,16 @@ pub trait Pack<'a>: serde::Serialize + serde::Deserialize<'a> { fn pack(src: Self, dst: &mut [u8]) -> Result<(), ProgramError>; /// Deserializes `src` into Self. The deserialized object need not - /// use all the bytes in `src`. This is the case, for example, - /// when encoding variable length data, say, when one has a src - /// array of all zeroes, and a Self that has a Vec. + /// use all the bytes in `src` and should mutate the slice so that + /// it's new len() becomes the size of the bytes deserialized. + /// + /// This is the case, for example, when encoding variable length data, + /// say, when one has a src array of all zeroes, and a Self that has a + /// Vec. /// /// Use care when using this directly. If using fixed length structs /// always use the `unpack` method, instead. - fn unpack_unchecked(src: &[u8]) -> Result; + fn unpack_unchecked(src: &mut &[u8]) -> Result; /// Returns the size of the byte array required to serialize Self. fn size(&self) -> Result; @@ -25,8 +28,9 @@ pub trait Pack<'a>: serde::Serialize + serde::Deserialize<'a> { /// Analogue to pack, performing a check on the size of the given byte /// array. fn unpack(src: &[u8]) -> Result { - Pack::unpack_unchecked(src).and_then(|r: Self| { - if r.size()? != src.len() as u64 { + let mut src_mut = src.as_ref(); + Pack::unpack_unchecked(&mut src_mut).and_then(|r: Self| { + if src_mut.len() != 0 { return Err(ProgramError::InvalidAccountData); } Ok(r) @@ -49,7 +53,7 @@ pub trait Pack<'a>: serde::Serialize + serde::Deserialize<'a> { where F: FnMut(&mut Self) -> Result, { - let mut t = Self::unpack_unchecked(input)?; + let mut t = Self::unpack_unchecked(&mut input.as_ref())?; let u = f(&mut t)?; Self::pack(t, input)?; Ok(u) @@ -75,9 +79,11 @@ macro_rules! packable { } serum_common::pack::into_bytes(&src, dst) } - fn unpack_unchecked(src: &[u8]) -> Result<$my_struct, ProgramError> { - serum_common::pack::from_bytes::<$my_struct>(src) + + fn unpack_unchecked(src: &mut &[u8]) -> Result<$my_struct, ProgramError> { + serum_common::pack::from_reader(src) } + fn size(&self) -> Result { serum_common::pack::bytes_size(&self) } @@ -99,6 +105,7 @@ where let cursor = std::io::Cursor::new(dst); bincode::serialize_into(cursor, i).map_err(|_| ProgramError::InvalidAccountData) } + pub fn from_bytes<'a, T>(data: &'a [u8]) -> Result where T: serde::de::Deserialize<'a>, @@ -106,6 +113,14 @@ where bincode::deserialize(data).map_err(|_| ProgramError::InvalidAccountData) } +pub fn from_reader<'a, T, R>(rdr: R) -> Result +where + R: std::io::Read, + T: serde::de::DeserializeOwned, +{ + bincode::deserialize_from(rdr).map_err(|_| ProgramError::InvalidAccountData) +} + pub fn bytes_size(value: &T) -> Result where T: serde::Serialize, @@ -137,40 +152,53 @@ mod tests { #[test] fn unpack_too_small() { let data = vec![0; 8]; - let result = TestStruct::unpack(&data); - match result { - Ok(_) => panic!("expect error"), - Err(e) => assert_eq!(e, ProgramError::InvalidAccountData), - } + let r = TestStruct::unpack(&data); + assert_eq!(r.unwrap_err(), ProgramError::InvalidAccountData); } #[test] fn unpack_too_large() { let data = vec![0; 100]; - let result = TestStruct::unpack(&data); - match result { - Ok(_) => panic!("expect error"), - Err(e) => assert_eq!(e, ProgramError::InvalidAccountData), - } + let r = TestStruct::unpack(&data); + assert_eq!(r.unwrap_err(), ProgramError::InvalidAccountData); } #[test] fn pack_too_small() { let mut data = vec![0; 8]; - let result = TestStruct::pack(Default::default(), &mut data); - match result { - Ok(_) => panic!("expect error"), - Err(e) => assert_eq!(e, ProgramError::InvalidAccountData), - } + let r = TestStruct::pack(Default::default(), &mut data); + assert_eq!(r.unwrap_err(), ProgramError::InvalidAccountData); } #[test] fn pack_too_large() { let mut data = vec![0; 100]; - let result = TestStruct::pack(Default::default(), &mut data); - match result { - Ok(_) => panic!("expect error"), - Err(e) => assert_eq!(e, ProgramError::InvalidAccountData), + let r = TestStruct::pack(Default::default(), &mut data); + assert_eq!(r.unwrap_err(), ProgramError::InvalidAccountData); + } + + mod dyn_size { + use super::*; + #[derive(Clone, Debug, Default, PartialEq, serde::Serialize, serde::Deserialize)] + pub struct VarLenStruct { + a: u64, + v: Vec, } + packable!(VarLenStruct); + } + + #[test] + fn var_len_struct_unpack_unchecked() { + let mut data = [0; 100].as_ref(); + let r = dyn_size::VarLenStruct::unpack_unchecked(&mut data); + assert!(r.is_ok()); + assert_eq!(data.len(), 84); + } + + #[test] + fn var_len_struct_unpack_checked() { + let data = [0; 100].as_ref(); + let r = dyn_size::VarLenStruct::unpack(&data); + assert_eq!(r.unwrap_err(), ProgramError::InvalidAccountData); } }