From 44c87c1d17d7651079aaecda72fcf22589947c3a Mon Sep 17 00:00:00 2001
From: Samuel Moelius <sam.moelius@trailofbits.com>
Date: Sun, 13 Aug 2023 15:13:23 -0400
Subject: [PATCH] Limit pre-allocation from serde size hints

---
 ci/big_serde/Cargo.toml |  1 +
 ci/big_serde/src/lib.rs | 35 ++++++++++++++++++++++++++++++++++-
 src/biguint/serde.rs    | 17 ++++++++++++++---
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/ci/big_serde/Cargo.toml b/ci/big_serde/Cargo.toml
index d8fed73f..f18413a2 100644
--- a/ci/big_serde/Cargo.toml
+++ b/ci/big_serde/Cargo.toml
@@ -6,6 +6,7 @@ edition = "2018"
 
 [dependencies]
 num-traits = "0.2.11"
+serde = "1.0"
 serde_test = "1.0"
 
 [dependencies.num-bigint]
diff --git a/ci/big_serde/src/lib.rs b/ci/big_serde/src/lib.rs
index f1fa513d..68c2551b 100644
--- a/ci/big_serde/src/lib.rs
+++ b/ci/big_serde/src/lib.rs
@@ -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() {
@@ -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);
+}
diff --git a/src/biguint/serde.rs b/src/biguint/serde.rs
index ed663c6d..3240f093 100644
--- a/src/biguint/serde.rs
+++ b/src/biguint/serde.rs
@@ -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>
@@ -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>()? {
@@ -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);