Просмотр исходного кода

serial: Add deserialization method with max_len enforcing

x 7 месяцев назад
Родитель
Сommit
2b9e894214
2 измененных файлов с 72 добавлено и 3 удалено
  1. 36 0
      src/serial/src/async_lib.rs
  2. 36 3
      src/serial/src/lib.rs

+ 36 - 0
src/serial/src/async_lib.rs

@@ -52,6 +52,24 @@ pub async fn serialize_async<T: AsyncEncodable + ?Sized>(data: &T) -> Vec<u8> {
     encoder
 }
 
+/// Asynchronously deserialize a variable-length object from a vector, enforce
+/// a maximum size limit, but do not error if the entire vector is not consumed.
+///
+/// Expects objects with VarInt length encoding.
+pub async fn deserialize_async_limited_partial<T: AsyncDecodable>(
+    data: &[u8],
+    max_len: usize,
+) -> Result<(T, usize)> {
+    let (len, _) = deserialize_async_partial::<VarInt>(data).await?;
+    let obj_len = len.0 as usize;
+
+    if obj_len > max_len {
+        return Err(Error::other("Data size exceeds set `max_len`"))
+    }
+
+    deserialize_async_partial(data).await
+}
+
 /// Asynchronously deserialize an object from a vector, but do not error if the
 /// entire vector is not consumed.
 pub async fn deserialize_async_partial<T: AsyncDecodable>(data: &[u8]) -> Result<(T, usize)> {
@@ -62,6 +80,24 @@ pub async fn deserialize_async_partial<T: AsyncDecodable>(data: &[u8]) -> Result
     Ok((rv, consumed))
 }
 
+/// Asynchronously deserialize a variable-length object from a vector, but
+/// enforce a maximum size limit.
+///
+/// Expects objects with VarInt length encoding.
+pub async fn deserialize_async_limited<T: AsyncDecodable>(
+    data: &[u8],
+    max_len: usize,
+) -> Result<T> {
+    let (len, _) = deserialize_async_partial::<VarInt>(data).await?;
+    let obj_len = len.0 as usize;
+
+    if obj_len > max_len {
+        return Err(Error::other("Data size exceeds set `max_len`"))
+    }
+
+    deserialize_async(data).await
+}
+
 /// Asynchronously deserialize an object from a vector.
 /// Will error if said deserialization doesn't consume the entire vector.
 pub async fn deserialize_async<T: AsyncDecodable>(data: &[u8]) -> Result<T> {

+ 36 - 3
src/serial/src/lib.rs

@@ -28,9 +28,9 @@ pub use darkfi_derive::{SerialDecodable, SerialEncodable};
 mod async_lib;
 #[cfg(feature = "async")]
 pub use async_lib::{
-    async_trait, deserialize_async, deserialize_async_partial, serialize_async, AsyncDecodable,
-    AsyncEncodable, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, FutAsyncReadExt,
-    FutAsyncWriteExt,
+    async_trait, deserialize_async, deserialize_async_limited, deserialize_async_limited_partial,
+    deserialize_async_partial, serialize_async, AsyncDecodable, AsyncEncodable, AsyncRead,
+    AsyncReadExt, AsyncWrite, AsyncWriteExt, FutAsyncReadExt, FutAsyncWriteExt,
 };
 
 mod endian;
@@ -57,6 +57,24 @@ pub fn serialize<T: Encodable + ?Sized>(data: &T) -> Vec<u8> {
     encoder
 }
 
+/// Deserialize a variable-length object from a vector, enforce a maximum
+/// size limit, but do not error if the entire vector is not consumed.
+///
+/// Expects objects with VarInt length encoding.
+pub fn deserialize_limited_partial<T: Decodable>(
+    data: &[u8],
+    max_len: usize,
+) -> Result<(T, usize), Error> {
+    let (len, _) = deserialize_partial::<VarInt>(data)?;
+    let obj_len = len.0 as usize;
+
+    if obj_len > max_len {
+        return Err(Error::other("Data size exceeds set `max_len`"))
+    }
+
+    deserialize_partial(data)
+}
+
 /// Deserialize an object from a vector, but do not error if the entire
 /// vector is not consumed.
 pub fn deserialize_partial<T: Decodable>(data: &[u8]) -> Result<(T, usize), Error> {
@@ -67,6 +85,21 @@ pub fn deserialize_partial<T: Decodable>(data: &[u8]) -> Result<(T, usize), Erro
     Ok((rv, consumed))
 }
 
+/// Deserialize a variable-length object from a vector, but enforce a
+/// maximum size limit.
+///
+/// Expects objects with VarInt length encoding.
+pub fn deserialize_limited<T: Decodable>(data: &[u8], max_len: usize) -> Result<T, Error> {
+    let (len, _) = deserialize_partial::<VarInt>(data)?;
+    let obj_len = len.0 as usize;
+
+    if obj_len > max_len {
+        return Err(Error::other("Data size exceeds set `max_len`"))
+    }
+
+    deserialize(data)
+}
+
 /// Deserialize an object from a vector.
 /// Will error if said deserialization doesn't consume the entire vector.
 pub fn deserialize<T: Decodable>(data: &[u8]) -> Result<T, Error> {