Skip to content

Commit

Permalink
Limit pre-allocation from serde size hints
Browse files Browse the repository at this point in the history
  • Loading branch information
smoelius authored and cuviper committed Aug 22, 2023
1 parent 2cea7f4 commit 44c87c1
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 4 deletions.
1 change: 1 addition & 0 deletions ci/big_serde/Cargo.toml
Expand Up @@ -6,6 +6,7 @@ edition = "2018"

[dependencies]
num-traits = "0.2.11"
serde = "1.0"
serde_test = "1.0"

[dependencies.num-bigint]
Expand Down
35 changes: 34 additions & 1 deletion ci/big_serde/src/lib.rs
Expand Up @@ -12,7 +12,9 @@

use num_bigint::{BigInt, BigUint};
use num_traits::{One, Zero};
use serde_test::{assert_tokens, Token};
use serde::{de::DeserializeOwned, Serialize};
use serde_test::{assert_de_tokens, assert_ser_tokens, assert_tokens, Token};
use std::{fmt::Debug, panic::catch_unwind};

#[test]
fn biguint_zero() {
Expand Down Expand Up @@ -128,3 +130,34 @@ fn big_digits() {
assert_tokens(&-n, &tokens);
}
}

#[test]
fn bad_size_hint_int() {
bad_size_hint::<BigInt>(&[Token::Tuple { len: 2 }, Token::I8(1)], &[Token::TupleEnd]);
}

#[test]
fn bad_size_hint_uint() {
bad_size_hint::<BigUint>(&[], &[]);
}

fn bad_size_hint<T: Debug + DeserializeOwned + One + PartialEq + Serialize>(
prefix: &[Token],
suffix: &[Token],
) {
let mut tokens = [
prefix,
&[Token::Seq { len: Some(1) }, Token::U32(1), Token::SeqEnd],
suffix,
]
.concat();

assert_tokens(&T::one(), &tokens);

tokens[prefix.len()] = Token::Seq {
len: Some(usize::max_value()),
};

catch_unwind(|| assert_ser_tokens(&T::one(), &tokens)).unwrap_err();
assert_de_tokens(&T::one(), &tokens);
}
17 changes: 14 additions & 3 deletions src/biguint/serde.rs
Expand Up @@ -2,10 +2,21 @@ use super::{biguint_from_vec, BigUint};

use crate::std_alloc::Vec;

use core::fmt;
use core::{cmp, fmt, mem};
use serde::de::{SeqAccess, Visitor};
use serde::{Deserialize, Deserializer, Serialize, Serializer};

// `cautious` is based on the function of the same name in `serde`, but specialized to `u32`:
// https://github.com/dtolnay/serde/blob/399ef081ecc36d2f165ff1f6debdcbf6a1dc7efb/serde/src/private/size_hint.rs#L11-L22
fn cautious(hint: Option<usize>) -> usize {
const MAX_PREALLOC_BYTES: usize = 1024 * 1024;

cmp::min(
hint.unwrap_or(0),
MAX_PREALLOC_BYTES / mem::size_of::<u32>(),
)
}

impl Serialize for BigUint {
#[cfg(not(u64_digit))]
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
Expand Down Expand Up @@ -70,7 +81,7 @@ impl<'de> Visitor<'de> for U32Visitor {
where
S: SeqAccess<'de>,
{
let len = seq.size_hint().unwrap_or(0);
let len = cautious(seq.size_hint());
let mut data = Vec::with_capacity(len);

while let Some(value) = seq.next_element::<u32>()? {
Expand All @@ -88,7 +99,7 @@ impl<'de> Visitor<'de> for U32Visitor {
use crate::big_digit::BigDigit;
use num_integer::Integer;

let u32_len = seq.size_hint().unwrap_or(0);
let u32_len = cautious(seq.size_hint());
let len = Integer::div_ceil(&u32_len, &2);
let mut data = Vec::with_capacity(len);

Expand Down

0 comments on commit 44c87c1

Please sign in to comment.