Merge "Early import of pending upstream commmit" am: 89c9c4f7cb am: 1f3bdbcffc

Original change: https://android-review.googlesource.com/c/platform/external/rust/crates/ciborium/+/2625934

Change-Id: I362fa18768cca6ed13e6cf050ddb479fe1fdfee2
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

2 files changed, 69 insertions, 1 deletions

diff --git a/src/de/mod.rs b/src/de/mod.rs
index f59a319..b68458d 100644
--- a/src/de/mod.rs
+++ b/src/de/mod.rs
@@ -831,3 +831,27 @@ where
 
     T::deserialize(&mut reader)
 }
+
+/// Deserializes as CBOR from a type with [`impl ciborium_io::Read`](ciborium_io::Read), with
+/// a specified maximum recursion limit.  Inputs that are nested beyond the specified limit
+/// will result in [`Error::RecursionLimitExceeded`] .
+///
+/// Set a high recursion limit at your own risk (of stack exhaustion)!
+#[inline]
+pub fn from_reader_with_recursion_limit<T: de::DeserializeOwned, R: Read>(
+    reader: R,
+    recurse_limit: usize,
+) -> Result<T, Error<R::Error>>
+where
+    R::Error: core::fmt::Debug,
+{
+    let mut scratch = [0; 4096];
+
+    let mut reader = Deserializer {
+        decoder: reader.into(),
+        scratch: &mut scratch,
+        recurse: recurse_limit,
+    };
+
+    T::deserialize(&mut reader)
+}
diff --git a/tests/recursion.rs b/tests/recursion.rs
index a340b2d..cda1ce2 100644
--- a/tests/recursion.rs
+++ b/tests/recursion.rs
@@ -7,7 +7,7 @@
 //! test each of these types here to ensure there is no stack overflow.
 
 use ciborium::{
-    de::{from_reader, Error},
+    de::{from_reader, from_reader_with_recursion_limit, Error},
     value::Value,
 };
 
@@ -46,3 +46,47 @@ fn text() {
         e => panic!("incorrect error: {:?}", e),
     }
 }
+
+#[test]
+fn array_limit() {
+    let bytes = [0x9f; 128 * 1024];
+    for limit in 16..256 {
+        match from_reader_with_recursion_limit::<Value, _>(&bytes[..], limit).unwrap_err() {
+            Error::RecursionLimitExceeded => (),
+            e => panic!("incorrect error with limit {}: {:?}", limit, e),
+        }
+        // Data that is nested beyond the limit should fail with `RecursionLimitExceeded`
+        match from_reader_with_recursion_limit::<Value, _>(&bytes[..limit + 1], limit).unwrap_err()
+        {
+            Error::RecursionLimitExceeded => (),
+            e => panic!("incorrect error with limit {}: {:?}", limit, e),
+        }
+        // Data that is nested within the limit fails with a different error.
+        match from_reader_with_recursion_limit::<Value, _>(&bytes[..limit], limit).unwrap_err() {
+            Error::Io(..) => (),
+            e => panic!("incorrect error with limit {}: {:?}", limit, e),
+        }
+    }
+}
+
+#[test]
+fn map_limit() {
+    let bytes = [0xbf; 128 * 1024];
+    for limit in 16..256 {
+        match from_reader_with_recursion_limit::<Value, _>(&bytes[..], limit).unwrap_err() {
+            Error::RecursionLimitExceeded => (),
+            e => panic!("incorrect error with limit {}: {:?}", limit, e),
+        }
+        // Data that is nested beyond the limit should fail with `RecursionLimitExceeded`
+        match from_reader_with_recursion_limit::<Value, _>(&bytes[..limit + 1], limit).unwrap_err()
+        {
+            Error::RecursionLimitExceeded => (),
+            e => panic!("incorrect error with limit {}: {:?}", limit, e),
+        }
+        // Data that is nested within the limit fails with a different error.
+        match from_reader_with_recursion_limit::<Value, _>(&bytes[..limit], limit).unwrap_err() {
+            Error::Io(..) => (),
+            e => panic!("incorrect error with limit {}: {:?}", limit, e),
+        }
+    }
+}