Fix stack overflow when decoding big array newtypes wrapped in a Box (

#462)

* Fix stack overflow when decoding big array newtypes wrapped in a `Box`

* Bump version to 3.6.2

* Update changelog

* Fix indentation

* Add an extra assert to make sure the number of non-ZST fields is correct

* Add even more tests

* Add release date to changelog

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this crate are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this crate adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [3.6.2] - 2023-06-30
+
+### Fixed
+
+- Trying to deserialize a boxed newtype containing a big array won't overflow the stack anymore.
+
 ## [3.6.1] - 2023-06-19
 
 ### Fixed

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "parity-scale-codec"
 description = "SCALE - Simple Concatenating Aggregated Little Endians"
-version = "3.6.1"
+version = "3.6.2"
 authors = ["Parity Technologies <[email protected]>"]
 license = "Apache-2.0"
 repository = "https://github.com/paritytech/parity-scale-codec"
@@ -12,7 +12,7 @@ rust-version = "1.60.0"
 [dependencies]
 arrayvec = { version = "0.7", default-features = false }
 serde = { version = "1.0.164", optional = true }
-parity-scale-codec-derive = { path = "derive", version = ">= 3.6.1", default-features = false, optional = true }
+parity-scale-codec-derive = { path = "derive", version = ">= 3.6.2", default-features = false, optional = true }
 bitvec = { version = "1", default-features = false, features = [ "alloc" ], optional = true }
 bytes = { version = "1", default-features = false, optional = true }
 byte-slice-cast = { version = "1.2.2", default-features = false }

derive/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "parity-scale-codec-derive"
 description = "Serialization and deserialization derive macro for Parity SCALE Codec"
-version = "3.6.1"
+version = "3.6.2"
 authors = ["Parity Technologies <[email protected]>"]
 license = "Apache-2.0"
 edition = "2021"

derive/src/decode.rs

@@ -113,6 +113,91 @@ pub fn quote(
 	}
 }
 
+pub fn quote_decode_into(
+	data: &Data,
+	crate_path: &syn::Path,
+	input: &TokenStream,
+	attrs: &[syn::Attribute]
+) -> Option<TokenStream> {
+	// Make sure the type is `#[repr(transparent)]`, as this guarantees that
+	// there can be only one field that is not zero-sized.
+	if !crate::utils::is_transparent(attrs) {
+		return None;
+	}
+
+	let fields = match data {
+		Data::Struct(
+			syn::DataStruct {
+				fields: Fields::Named(syn::FieldsNamed { named: fields, .. }) |
+				        Fields::Unnamed(syn::FieldsUnnamed { unnamed: fields, .. }),
+				..
+			}
+		) => {
+			fields
+		},
+		_ => return None
+	};
+
+	if fields.is_empty() {
+		return None;
+	}
+
+	// Bail if there are any extra attributes which could influence how the type is decoded.
+	if fields.iter().any(|field|
+		utils::get_encoded_as_type(field).is_some() ||
+		utils::is_compact(field) ||
+		utils::should_skip(&field.attrs)
+	) {
+		return None;
+	}
+
+	// Go through each field and call `decode_into` on it.
+	//
+	// Normally if there's more than one field in the struct this would be incorrect,
+	// however since the struct's marked as `#[repr(transparent)]` we're guaranteed that
+	// there's at most one non zero-sized field, so only one of these `decode_into` calls
+	// should actually do something, and the rest should just be dummy calls that do nothing.
+	let mut decode_fields = Vec::new();
+	let mut sizes = Vec::new();
+	let mut non_zst_field_count = Vec::new();
+	for field in fields {
+		let field_type = &field.ty;
+		decode_fields.push(quote! {{
+			let dst_: &mut ::core::mem::MaybeUninit<Self> = dst_; // To make sure the type is what we expect.
+
+			// Here we cast `&mut MaybeUninit<Self>` into a `&mut MaybeUninit<#field_type>`.
+			//
+			// SAFETY: The struct is marked as `#[repr(transparent)]` so the address of every field will
+			//         be the same as the address of the struct itself.
+			let dst_: &mut ::core::mem::MaybeUninit<#field_type> = unsafe {
+				&mut *dst_.as_mut_ptr().cast::<::core::mem::MaybeUninit<#field_type>>()
+			};
+			<#field_type as #crate_path::Decode>::decode_into(#input, dst_)?;
+		}});
+
+		if !sizes.is_empty() {
+			sizes.push(quote! { + });
+		}
+		sizes.push(quote! { ::core::mem::size_of::<#field_type>() });
+
+		if !non_zst_field_count.is_empty() {
+			non_zst_field_count.push(quote! { + });
+		}
+		non_zst_field_count.push(quote! { if ::core::mem::size_of::<#field_type>() > 0 { 1 } else { 0 } });
+	}
+
+	Some(quote!{
+		// Just a sanity check. These should always be true and will be optimized-out.
+		assert_eq!(#(#sizes)*, ::core::mem::size_of::<Self>());
+		assert!(#(#non_zst_field_count)* <= 1);
+
+		#(#decode_fields)*
+
+		// SAFETY: We've successfully called `decode_into` for all of the fields.
+		unsafe { Ok(#crate_path::DecodeFinished::assert_decoding_finished()) }
+	})
+}
+
 fn create_decode_expr(field: &Field, name: &str, input: &TokenStream, crate_path: &syn::Path) -> TokenStream {
 	let encoded_as = utils::get_encoded_as_type(field);
 	let compact = utils::is_compact(field);

derive/src/lib.rs

@@ -208,6 +208,26 @@ pub fn decode_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream
 	let decoding =
 		decode::quote(&input.data, name, &quote!(#ty_gen_turbofish), &input_, &crate_path);
 
+	let decode_into_body = decode::quote_decode_into(
+		&input.data,
+		&crate_path,
+		&input_,
+		&input.attrs
+	);
+
+	let impl_decode_into = if let Some(body) = decode_into_body {
+		quote! {
+			fn decode_into<__CodecInputEdqy: #crate_path::Input>(
+				#input_: &mut __CodecInputEdqy,
+				dst_: &mut ::core::mem::MaybeUninit<Self>,
+			) -> ::core::result::Result<#crate_path::DecodeFinished, #crate_path::Error> {
+				#body
+			}
+		}
+	} else {
+		quote! {}
+	};
+
 	let impl_block = quote! {
 		#[automatically_derived]
 		impl #impl_generics #crate_path::Decode for #name #ty_generics #where_clause {
@@ -216,6 +236,8 @@ pub fn decode_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream
 			) -> ::core::result::Result<Self, #crate_path::Error> {
 				#decoding
 			}
+
+			#impl_decode_into
 		}
 	};
 

derive/src/utils.rs

@@ -20,7 +20,7 @@
 use std::str::FromStr;
 
 use proc_macro2::TokenStream;
-use quote::quote;
+use quote::{ToTokens, quote};
 use syn::{
 	parse::Parse, punctuated::Punctuated, spanned::Spanned, token, Attribute, Data, DeriveInput,
 	Field, Fields, FieldsNamed, FieldsUnnamed, Lit, Meta, MetaNameValue, NestedMeta, Path, Variant,
@@ -431,3 +431,23 @@ fn check_top_attribute(attr: &Attribute) -> syn::Result<()> {
 		Ok(())
 	}
 }
+
+fn check_repr(attrs: &[syn::Attribute], value: &str) -> bool {
+	let mut result = false;
+	for raw_attr in attrs {
+		let path = raw_attr.path.clone().into_token_stream().to_string();
+		if path != "repr" {
+			continue;
+		}
+
+		result = raw_attr.tokens.clone().into_token_stream().to_string() == value;
+	}
+
+	result
+}
+
+/// Checks whether the given attributes contain a `#[repr(transparent)]`.
+pub fn is_transparent(attrs: &[syn::Attribute]) -> bool {
+	// TODO: When migrating to syn 2 the `"(transparent)"` needs to be changed into `"transparent"`.
+	check_repr(attrs, "(transparent)")
+}

tests/mod.rs

@@ -662,6 +662,88 @@ fn decoding_a_huge_array_inside_of_arc_does_not_overflow_the_stack() {
 	let _ = std::sync::Arc::<[u8; 100 * 1024 * 1024]>::decode(&mut data.as_slice());
 }
 
+#[test]
+fn decoding_a_huge_boxed_newtype_array_does_not_overflow_the_stack() {
+	#[derive(DeriveDecode)]
+	#[repr(transparent)]
+	struct HugeArrayNewtype([u8; 100 * 1024 * 1024]);
+
+	#[derive(DeriveDecode)]
+	struct HugeArrayNewtypeBox(Box<HugeArrayNewtype>);
+
+	let data = &[];
+	assert!(HugeArrayNewtypeBox::decode(&mut data.as_slice()).is_err());
+}
+
+#[test]
+fn decoding_two_indirectly_boxed_arrays_works() {
+	// This test will fail if the check for `#[repr(transparent)]` in the derive crate
+	// doesn't work when implementing `Decode::decode_into`.
+	#[derive(DeriveDecode)]
+	#[derive(PartialEq, Eq, Debug)]
+	struct SmallArrays([u8; 2], [u8; 2]);
+
+	#[derive(DeriveDecode)]
+	struct SmallArraysBox(Box<SmallArrays>);
+
+	let data = &[1, 2, 3, 4];
+	assert_eq!(*SmallArraysBox::decode(&mut data.as_slice()).unwrap().0, SmallArrays([1, 2], [3, 4]));
+}
+
+#[test]
+fn zero_sized_types_are_properly_decoded_in_a_transparent_boxed_struct() {
+	#[derive(DeriveDecode)]
+	#[repr(transparent)]
+	struct ZstTransparent;
+
+	#[derive(DeriveDecode)]
+	struct ZstNonTransparent;
+
+	struct ConsumeByte;
+
+	#[derive(DeriveDecode)]
+	#[repr(transparent)]
+	struct NewtypeWithZst {
+		_zst_1: ConsumeByte,
+		_zst_2: ZstTransparent,
+		_zst_3: ZstNonTransparent,
+		field: [u8; 1],
+		_zst_4: ConsumeByte
+	}
+
+	#[derive(DeriveDecode)]
+	struct NewtypeWithZstBox(Box<NewtypeWithZst>);
+
+	impl Decode for ConsumeByte {
+		fn decode<I: parity_scale_codec::Input>(input: &mut I) -> Result<Self, Error> {
+			let mut buffer = [0; 1];
+			input.read(&mut buffer).unwrap();
+			Ok(Self)
+		}
+	}
+
+	let data = &[1, 2, 3];
+	assert_eq!(NewtypeWithZst::decode(&mut data.as_slice()).unwrap().field, [2]);
+}
+
+#[test]
+fn boxed_zero_sized_newtype_with_everything_being_transparent_is_decoded_correctly() {
+	#[derive(DeriveDecode)]
+	#[repr(transparent)]
+	struct Zst;
+
+	#[derive(DeriveDecode)]
+	#[repr(transparent)]
+	struct NewtypeWithZst(Zst);
+
+	#[derive(DeriveDecode)]
+	#[repr(transparent)]
+	struct NewtypeWithZstBox(Box<NewtypeWithZst>);
+
+	let data = &[];
+	assert!(NewtypeWithZst::decode(&mut data.as_slice()).is_ok());
+}
+
 #[test]
 fn decoding_an_array_of_boxed_zero_sized_types_works() {
 	#[cfg(not(miri))]