alloc: Add fallible Vec::with_capacity

Author: Serial-ATA

lofty/src/id3/v2/frame/mod.rs

@@ -212,13 +212,13 @@ impl<'a> Frame<'a> {
 
 				attached_picture.as_bytes(version)?
 			},
-			Frame::Popularimeter(popularimeter) => popularimeter.as_bytes(),
+			Frame::Popularimeter(popularimeter) => popularimeter.as_bytes()?,
 			Frame::KeyValue(content) => content.as_bytes(is_id3v23),
 			Frame::RelativeVolumeAdjustment(frame) => frame.as_bytes(),
 			Frame::UniqueFileIdentifier(frame) => frame.as_bytes(),
 			Frame::Ownership(frame) => frame.as_bytes(is_id3v23)?,
 			Frame::EventTimingCodes(frame) => frame.as_bytes(),
-			Frame::Private(frame) => frame.as_bytes(),
+			Frame::Private(frame) => frame.as_bytes()?,
 			Frame::Timestamp(frame) => frame.as_bytes(is_id3v23)?,
 			Frame::Binary(frame) => frame.as_bytes(),
 		})

lofty/src/id3/v2/items/popularimeter.rs

@@ -1,5 +1,6 @@
 use crate::error::Result;
 use crate::id3::v2::{FrameFlags, FrameHeader, FrameId};
+use crate::util::alloc::VecFallibleCapacity;
 use crate::util::text::{decode_text, encode_text, TextDecodeOptions, TextEncoding};
 
 use std::borrow::Cow;
@@ -114,8 +115,12 @@ impl<'a> PopularimeterFrame<'a> {
 	/// Convert a [`PopularimeterFrame`] into an ID3v2 POPM frame byte Vec
 	///
 	/// NOTE: This does not include a frame header
-	pub fn as_bytes(&self) -> Vec<u8> {
-		let mut content = Vec::with_capacity(self.email.len() + 9);
+	///
+	/// # Errors
+	///
+	/// * The resulting [`Vec`] exceeds [`GlobalOptions::allocation_limit`](crate::config::GlobalOptions::allocation_limit)
+	pub fn as_bytes(&self) -> Result<Vec<u8>> {
+		let mut content = Vec::try_with_capacity_stable(self.email.len() + 9)?;
 		content.extend(encode_text(self.email.as_str(), TextEncoding::Latin1, true));
 		content.push(self.rating);
 
@@ -132,7 +137,7 @@ impl<'a> PopularimeterFrame<'a> {
 			content.extend(&counter_bytes[i..]);
 		}
 
-		content
+		Ok(content)
 	}
 }
 
@@ -146,7 +151,7 @@ mod tests {
 		let rating = popm.rating;
 		let counter = popm.counter;
 
-		let popm_bytes = popm.as_bytes();
+		let popm_bytes = popm.as_bytes().unwrap();
 		assert_eq!(&popm_bytes[..email.len()], email.as_bytes());
 		assert_eq!(popm_bytes[email.len()], 0);
 		assert_eq!(popm_bytes[email.len() + 1], rating);

lofty/src/id3/v2/items/private_frame.rs

@@ -1,5 +1,6 @@
 use crate::error::Result;
 use crate::id3::v2::{FrameFlags, FrameHeader, FrameId};
+use crate::util::alloc::VecFallibleCapacity;
 use crate::util::text::{decode_text, encode_text, TextDecodeOptions, TextEncoding};
 
 use std::borrow::Cow;
@@ -81,18 +82,22 @@ impl<'a> PrivateFrame<'a> {
 	}
 
 	/// Convert an [`PrivateFrame`] to a byte vec
-	pub fn as_bytes(&self) -> Vec<u8> {
+	///
+	/// # Errors
+	///
+	/// * The resulting [`Vec`] exceeds [`GlobalOptions::allocation_limit`](crate::config::GlobalOptions::allocation_limit)
+	pub fn as_bytes(&self) -> Result<Vec<u8>> {
 		let Self {
 			owner,
 			private_data,
 			..
 		} = self;
 
-		let mut content = Vec::with_capacity(owner.len() + private_data.len());
+		let mut content = Vec::try_with_capacity_stable(owner.len() + private_data.len())?;
 		content.extend(encode_text(owner.as_str(), TextEncoding::Latin1, true));
 		content.extend_from_slice(private_data);
 
-		content
+		Ok(content)
 	}
 }
 
@@ -121,7 +126,7 @@ mod tests {
 
 	#[test]
 	fn priv_encode() {
-		let encoded = expected().as_bytes();
+		let encoded = expected().as_bytes().unwrap();
 
 		let expected_bytes =
 			crate::tag::utils::test_utils::read_path("tests/tags/assets/id3v2/test.priv");

lofty/src/mp4/ilst/write.rs

@@ -9,6 +9,7 @@ use crate::mp4::read::{atom_tree, meta_is_full, nested_atom, verify_mp4, AtomRea
 use crate::mp4::write::{AtomWriter, AtomWriterCompanion, ContextualAtom};
 use crate::mp4::AtomData;
 use crate::picture::{MimeType, Picture};
+use crate::util::alloc::VecFallibleCapacity;
 use crate::util::io::{FileLike, Length, Truncate};
 
 use std::io::{Cursor, Seek, SeekFrom, Write};
@@ -540,7 +541,7 @@ fn create_udta(ilst: &[u8]) -> Result<Vec<u8>> {
 
 	// `udta` + `meta` + `hdlr` + `ilst`
 	let capacity = ATOM_HEADER_LEN + FULL_ATOM_SIZE + HDLR_SIZE + ilst.len() as u64;
-	let mut buf = Vec::with_capacity(capacity as usize);
+	let mut buf = Vec::try_with_capacity_stable(capacity as usize)?;
 
 	buf.write_all(&UDTA_HEADER)?;
 

lofty/src/util/alloc.rs

@@ -56,6 +56,29 @@ where
 	Vec::new().fallible_repeat(element, expected_size)
 }
 
+/// Provides the `try_with_capacity` method on `Vec`
+///
+/// This can be used directly.
+pub(crate) trait VecFallibleCapacity<T>: Sized {
+	/// Same as `Vec::with_capacity`, but takes `GlobalOptions::allocation_limit` into account.
+	///
+	/// Named `try_with_capacity_stable` to avoid conflicts with the nightly `Vec::try_with_capacity`.
+	fn try_with_capacity_stable(capacity: usize) -> Result<Self>;
+}
+
+impl<T> VecFallibleCapacity<T> for Vec<T> {
+	fn try_with_capacity_stable(capacity: usize) -> Result<Self> {
+		if capacity > unsafe { global_options().allocation_limit } {
+			err!(TooMuchData);
+		}
+
+		let mut v = Vec::new();
+		v.try_reserve(capacity)?;
+
+		Ok(v)
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use crate::util::alloc::fallible_vec_from_element;