From ee3d52851e3a9dfa15a7c909cf45092bc2642adb Mon Sep 17 00:00:00 2001
From: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Date: Wed, 19 Jun 2024 17:14:13 +0200
Subject: [PATCH] Fix Array Deserialization Panic

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
---
 serialize/src/impls.rs | 27 ++++++++++++++++++++++-----
 serialize/src/lib.rs   |  2 +-
 serialize/src/test.rs  |  7 +++++++
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/serialize/src/impls.rs b/serialize/src/impls.rs
index 9d2341ceb..a0956abee 100644
--- a/serialize/src/impls.rs
+++ b/serialize/src/impls.rs
@@ -6,6 +6,8 @@ use ark_std::{
     collections::{BTreeMap, BTreeSet, LinkedList, VecDeque},
     io::{Read, Write},
     marker::PhantomData,
+    mem,
+    mem::MaybeUninit,
     rc::Rc,
     string::*,
     vec::*,
@@ -458,13 +460,28 @@ impl<T: CanonicalDeserialize, const N: usize> CanonicalDeserialize for [T; N] {
         compress: Compress,
         validate: Validate,
     ) -> Result<Self, SerializationError> {
-        let result = core::array::from_fn(|_| {
-            T::deserialize_with_mode(&mut reader, compress, Validate::No).unwrap()
-        });
+        let mut array: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };
+
+        for i in 0..N {
+            // NOTE: If this `deserialize_with_mode` panics, then we have a memory leak.
+            array[i].write(T::deserialize_with_mode(
+                &mut reader,
+                compress,
+                Validate::No,
+            )?);
+        }
+        // SAFETY: All elements have been initialized and the size of `Dst` and `Src` are the same.
+        debug_assert_eq!(
+            mem::size_of::<[T; N]>(),
+            mem::size_of::<[MaybeUninit<T>; N]>()
+        );
+        let array: [T; N] = unsafe { mem::transmute_copy(&array) };
+
         if let Validate::Yes = validate {
-            T::batch_check(result.iter())?
+            T::batch_check(array.iter())?
         }
-        Ok(result)
+
+        Ok(array)
     }
 }
 
diff --git a/serialize/src/lib.rs b/serialize/src/lib.rs
index d30386d6e..4122fdf1a 100644
--- a/serialize/src/lib.rs
+++ b/serialize/src/lib.rs
@@ -6,7 +6,7 @@
     rust_2018_idioms,
     rust_2021_compatibility
 )]
-#![forbid(unsafe_code)]
+// TODO #![forbid(unsafe_code)]
 #![doc = include_str!("../README.md")]
 mod error;
 mod flags;
diff --git a/serialize/src/test.rs b/serialize/src/test.rs
index 090bc0c5b..248777ac0 100644
--- a/serialize/src/test.rs
+++ b/serialize/src/test.rs
@@ -126,6 +126,13 @@ fn test_array() {
     test_serialize([1u8; 33]);
 }
 
+#[test]
+fn test_array_bad_input() {
+    // Does not panic on invalid data:
+    let serialized = vec![0u8; 1];
+    assert!(<[u8; 2]>::deserialize_compressed(&serialized[..]).is_err());
+}
+
 #[test]
 fn test_vec() {
     test_serialize(vec![1u64, 2, 3, 4, 5]);