Unverified Commit e1aa580d authored by thiolliere's avatar thiolliere Committed by GitHub
Browse files

Fix preallocation (#108)

parent cebf0575
Pipeline #41385 passed with stages
in 1 minute and 12 seconds
......@@ -15,7 +15,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
[[package]]
name = "parity-codec"
version = "3.5.1"
version = "3.5.2"
dependencies = [
"arrayvec 0.4.10 (registry+https://github.com/rust-lang/crates.io-index)",
"parity-codec-derive 3.3.0",
......
[package]
name = "parity-codec"
description = "Lightweight, efficient, binary serialization and deserialization codec"
version = "3.5.1"
version = "3.5.2"
authors = ["Parity Technologies <admin@parity.io>"]
license = "Apache-2.0"
repository = "https://github.com/paritytech/parity-codec"
......
......@@ -18,6 +18,10 @@ use crate::alloc::vec::Vec;
use crate::alloc::boxed::Box;
use crate::alloc::collections::btree_map::BTreeMap;
/// When decoding input we don't allocate too much in advance so a small crafted input can't
/// trigger a big allocation.
const MAX_PREALLOCATION_SIZE: usize = 4 * 1024;
#[cfg(any(feature = "std", feature = "full"))]
use crate::alloc::{
string::String,
......@@ -57,8 +61,8 @@ impl<'a> Input for &'a [u8] {
#[cfg(feature = "std")]
impl<R: std::io::Read> Input for R {
fn read(&mut self, into: &mut [u8]) -> usize {
match (self as &mut dyn std::io::Read).read_exact(into) {
Ok(()) => into.len(),
match (self as &mut dyn std::io::Read).read(into) {
Ok(len) => len,
Err(_) => 0,
}
}
......@@ -781,10 +785,32 @@ impl Decode for Vec<u8> {
fn decode<I: Input>(input: &mut I) -> Option<Self> {
<Compact<u32>>::decode(input).and_then(move |Compact(len)| {
let len = len as usize;
let mut vec = vec![0; len];
if input.read(&mut vec[..len]) != len {
None
if len < MAX_PREALLOCATION_SIZE {
let mut vec = vec![0; len]; // len is ok here
if input.read(&mut vec[..]) != len {
None
} else {
Some(vec)
}
} else {
// if len is considered too much for preallocation then use dynamic allocation
let mut vec = Vec::new();
let mut remains = len;
let buffer_len = MAX_PREALLOCATION_SIZE;
let mut buffer = vec![0; buffer_len];
while remains != 0 {
let len_read = input.read(&mut buffer[..buffer_len.min(remains)]);
if len_read == 0 {
return None
}
remains -= len_read;
vec.extend_from_slice(&buffer[..len_read]);
}
Some(vec)
}
})
......@@ -864,7 +890,10 @@ impl<T: Encode> Encode for Vec<T> {
impl<T: Decode> Decode for Vec<T> {
fn decode<I: Input>(input: &mut I) -> Option<Self> {
<Compact<u32>>::decode(input).and_then(move |Compact(len)| {
let mut r = Vec::with_capacity(len as usize);
let max_pre_allocated_len = MAX_PREALLOCATION_SIZE / mem::size_of::<T>();
let pre_allocated_len = (len as usize).min(max_pre_allocated_len);
let mut r = Vec::with_capacity(pre_allocated_len);
for _ in 0..len {
r.push(T::decode(input)?);
}
......@@ -1472,4 +1501,12 @@ mod tests {
CompactRef(&std::u64::MAX).using_encoded(|_| {});
CompactRef(&std::u128::MAX).using_encoded(|_| {});
}
#[test]
fn encode_for_very_large_vec_works() {
let vec_u8: Vec<u8> = (0..MAX_PREALLOCATION_SIZE*3).map(|i| i as u8).collect();
let vec_u16: Vec<u16> = (0..MAX_PREALLOCATION_SIZE*3).map(|i| i as u16).collect();
assert_eq!(Vec::<u8>::decode(&mut &vec_u8.encode()[..][..]).unwrap(), vec_u8);
assert_eq!(Vec::<u16>::decode(&mut &vec_u16.encode()[..][..]).unwrap(), vec_u16);
}
}
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment