Begin fallible allocation; Fix WAV/AIFF OOM

src/error.rs

@@ -5,6 +5,7 @@
 
 use crate::types::file::FileType;
 
+use std::collections::TryReserveError;
 use std::fmt::{Debug, Display, Formatter};
 
 use ogg_pager::PageError;
@@ -25,7 +26,6 @@ pub enum ErrorKind {
 	// File data related errors
 	/// Attempting to read/write an abnormally large amount of data
 	TooMuchData,
-	// TODO: Fallible allocation
 	/// Errors that occur while decoding a file
 	FileDecoding(FileDecodingError),
 	/// Errors that occur while encoding a file
@@ -59,6 +59,8 @@ pub enum ErrorKind {
 	StrFromUtf8(std::str::Utf8Error),
 	/// Represents all cases of [`std::io::Error`].
 	Io(std::io::Error),
+	/// TODO
+	Alloc(TryReserveError),
 }
 
 #[derive(Debug, Clone)]
@@ -338,6 +340,14 @@ impl From<std::str::Utf8Error> for LoftyError {
 	}
 }
 
+impl From<std::collections::TryReserveError> for LoftyError {
+	fn from(input: TryReserveError) -> Self {
+		Self {
+			kind: ErrorKind::Alloc(input),
+		}
+	}
+}
+
 impl Display for LoftyError {
 	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
 		match self.kind {
@@ -346,6 +356,7 @@ impl Display for LoftyError {
 			ErrorKind::StringFromUtf8(ref err) => write!(f, "{}", err),
 			ErrorKind::StrFromUtf8(ref err) => write!(f, "{}", err),
 			ErrorKind::Io(ref err) => write!(f, "{}", err),
+			ErrorKind::Alloc(ref err) => write!(f, "{}", err),
 
 			ErrorKind::BadExtension(ref ext) => {
 				write!(f, "Found unknown file extension \"{}\"", ext)

src/id3/v2/write/chunk_file.rs

@@ -4,17 +4,19 @@ use crate::iff::chunk::Chunks;
 use std::fs::File;
 use std::io::{Read, Seek, SeekFrom, Write};
 
-use byteorder::{ByteOrder, WriteBytesExt};
+use byteorder::{ByteOrder, WriteBytesExt, ReadBytesExt};
 
 pub(in crate::id3::v2) fn write_to_chunk_file<B>(data: &mut File, tag: &[u8]) -> Result<()>
 where
 	B: ByteOrder,
 {
-	data.seek(SeekFrom::Current(12))?;
+	data.seek(SeekFrom::Current(4))?;
+	let file_size = data.read_u32::<B>()?;
+	data.seek(SeekFrom::Current(4))?;
 
 	let mut id3v2_chunk = (None, None);
 
-	let mut chunks = Chunks::<B>::new();
+	let mut chunks = Chunks::<B>::new(file_size);
 
 	while chunks.next(data).is_ok() {
 		if &chunks.fourcc == b"ID3 " || &chunks.fourcc == b"id3 " {

src/iff/aiff/read.rs

@@ -12,7 +12,7 @@ use std::io::{Read, Seek, SeekFrom};
 
 use byteorder::{BigEndian, ReadBytesExt};
 
-pub(in crate::iff) fn verify_aiff<R>(data: &mut R) -> Result<()>
+pub(in crate::iff) fn verify_aiff<R>(data: &mut R) -> Result<u32>
 where
 	R: Read + Seek,
 {
@@ -23,14 +23,14 @@ where
 		return Err(LoftyError::new(ErrorKind::UnknownFormat));
 	}
 
-	Ok(())
+	Ok(u32::from_be_bytes(id[4..8].try_into().unwrap()))
 }
 
 pub(crate) fn read_from<R>(data: &mut R, read_properties: bool) -> Result<AiffFile>
 where
 	R: Read + Seek,
 {
-	verify_aiff(data)?;
+	let file_size = verify_aiff(data)?;
 
 	let mut comm = None;
 	let mut stream_len = 0;
@@ -45,7 +45,7 @@ where
 	#[cfg(feature = "id3v2")]
 	let mut id3v2_tag: Option<Id3v2Tag> = None;
 
-	let mut chunks = Chunks::<BigEndian>::new();
+	let mut chunks = Chunks::<BigEndian>::new(file_size);
 
 	while chunks.next(data).is_ok() {
 		match &chunks.fourcc {
@@ -82,7 +82,7 @@ where
 					let marker_id = data.read_u16::<BigEndian>()?;
 					let size = data.read_u16::<BigEndian>()?;
 
-					let text = chunks.read_pstring(data, Some(size as usize))?;
+					let text = chunks.read_pstring(data, Some(u32::from(size)))?;
 
 					comments.push(Comment {
 						timestamp,

src/iff/aiff/tag.rs

@@ -355,13 +355,13 @@ where
 	}
 
 	fn write_to_inner(data: &mut File, mut tag: AiffTextChunksRef<T, AI>) -> Result<()> {
-		super::read::verify_aiff(data)?;
+		let file_size = super::read::verify_aiff(data)?;
 
 		let text_chunks = Self::create_text_chunks(&mut tag)?;
 
 		let mut chunks_remove = Vec::new();
 
-		let mut chunks = Chunks::<BigEndian>::new();
+		let mut chunks = Chunks::<BigEndian>::new(file_size);
 
 		while chunks.next(data).is_ok() {
 			match &chunks.fourcc {
@@ -468,7 +468,7 @@ mod tests {
 			.unwrap();
 
 		// Create a fake AIFF signature
-		let mut writer = vec![b'F', b'O', b'R', b'M', 0, 0, 0, 0, b'A', b'I', b'F', b'F'];
+		let mut writer = vec![b'F', b'O', b'R', b'M', 0, 0, 0, 0xC6, b'A', b'I', b'F', b'F'];
 		parsed_tag.dump_to(&mut writer).unwrap();
 
 		let temp_parsed_tag = super::super::read::read_from(&mut Cursor::new(writer), false)

src/iff/chunk.rs

@@ -1,9 +1,10 @@
-use crate::error::Result;
+use crate::error::{ErrorKind, LoftyError, Result};
 #[cfg(feature = "id3v2")]
 use crate::id3::v2::read::parse_id3v2;
 use crate::id3::v2::read_id3v2_header;
 #[cfg(feature = "id3v2")]
 use crate::id3::v2::tag::Id3v2Tag;
+use crate::macros::try_vec;
 
 use std::io::{Read, Seek, SeekFrom};
 use std::marker::PhantomData;
@@ -16,14 +17,16 @@ where
 {
 	pub fourcc: [u8; 4],
 	pub size: u32,
+	file_size: u32,
 	_phantom: PhantomData<B>,
 }
 
 impl<B: ByteOrder> Chunks<B> {
-	pub fn new() -> Self {
+	pub fn new(file_size: u32) -> Self {
 		Self {
 			fourcc: [0; 4],
 			size: 0,
+			file_size,
 			_phantom: PhantomData,
 		}
 	}
@@ -47,18 +50,15 @@ impl<B: ByteOrder> Chunks<B> {
 
 		let value_str = std::str::from_utf8(&cont)?;
 
-		Ok(value_str.trim_matches('\0').to_string())
+		Ok(value_str.trim_end_matches('\0').to_string())
 	}
 
-	pub fn read_pstring<R>(&mut self, data: &mut R, size: Option<usize>) -> Result<String>
+	pub fn read_pstring<R>(&mut self, data: &mut R, size: Option<u32>) -> Result<String>
 	where
 		R: Read + Seek,
 	{
 		let cont = if let Some(size) = size {
-			let mut v = vec![0; size];
+			self.read(data, size)?
-			data.read_exact(&mut v)?;
-
-			v
 		} else {
 			self.content(data)?
 		};
@@ -74,7 +74,18 @@ impl<B: ByteOrder> Chunks<B> {
 	where
 		R: Read,
 	{
-		let mut content = vec![0; self.size as usize];
+		self.read(data, self.size)
+	}
+
+	fn read<R>(&self, data: &mut R, size: u32) -> Result<Vec<u8>>
+	where
+		R: Read,
+	{
+		if size + 4 > self.file_size {
+			return Err(LoftyError::new(ErrorKind::TooMuchData));
+		}
+
+		let mut content = try_vec![0; size as usize];
 		data.read_exact(&mut content)?;
 
 		Ok(content)
@@ -103,28 +114,6 @@ impl<B: ByteOrder> Chunks<B> {
 		Ok(id3v2)
 	}
 
-	#[cfg(not(feature = "id3v2"))]
-	pub fn id3_chunk<R>(&mut self, data: &mut R) -> Result<()>
-	where
-		R: Read + Seek,
-	{
-		let mut value = vec![0; self.size as usize];
-		data.read_exact(&mut value)?;
-
-		let reader = &mut &*value;
-
-		let header = read_id3v2_header(reader)?;
-
-		// Skip over the footer
-		if header.flags.footer {
-			data.seek(SeekFrom::Current(10))?;
-		}
-
-		self.correct_position(data)?;
-
-		Ok(())
-	}
-
 	pub fn skip<R>(&mut self, data: &mut R) -> Result<()>
 	where
 		R: Read + Seek,

src/iff/wav/read.rs

@@ -12,7 +12,7 @@ use std::io::{Read, Seek, SeekFrom};
 
 use byteorder::{LittleEndian, ReadBytesExt};
 
-pub(in crate::iff) fn verify_wav<T>(data: &mut T) -> Result<()>
+pub(in crate::iff) fn verify_wav<T>(data: &mut T) -> Result<u32>
 where
 	T: Read + Seek,
 {
@@ -31,14 +31,14 @@ where
 		);
 	}
 
-	Ok(())
+	Ok(u32::from_le_bytes(id[4..8].try_into().unwrap()))
 }
 
 pub(crate) fn read_from<R>(data: &mut R, read_properties: bool) -> Result<WavFile>
 where
 	R: Read + Seek,
 {
-	verify_wav(data)?;
+	let file_size = verify_wav(data)?;
 
 	let mut stream_len = 0_u32;
 	let mut total_samples = 0_u32;
@@ -49,7 +49,7 @@ where
 	#[cfg(feature = "id3v2")]
 	let mut id3v2_tag: Option<Id3v2Tag> = None;
 
-	let mut chunks = Chunks::<LittleEndian>::new();
+	let mut chunks = Chunks::<LittleEndian>::new(file_size);
 
 	while chunks.next(data).is_ok() {
 		match &chunks.fourcc {
@@ -78,20 +78,20 @@ where
 				let mut list_type = [0; 4];
 				data.read_exact(&mut list_type)?;
 
-				#[cfg(feature = "riff_info_list")]
+				match &list_type {
-				if &list_type == b"INFO" {
+					#[cfg(feature = "riff_info_list")]
-					let end = data.seek(SeekFrom::Current(0))? + u64::from(chunks.size - 4);
+					b"INFO" => {
-					super::tag::read::parse_riff_info(data, end, &mut riff_info)?;
+						let end = data.seek(SeekFrom::Current(0))? + u64::from(chunks.size - 4);
-					chunks.correct_position(data)?;
+						super::tag::read::parse_riff_info(data, &mut chunks, end, &mut riff_info)?;
+					},
+					_ => {
+						data.seek(SeekFrom::Current(-4))?;
+						chunks.skip(data)?;
+					},
 				}
-
-				#[cfg(not(feature = "riff_info_list"))]
-				chunks.skip(data)?;
 			},
 			#[cfg(feature = "id3v2")]
 			b"ID3 " | b"id3 " => id3v2_tag = Some(chunks.id3_chunk(data)?),
-			#[cfg(not(feature = "id3v2"))]
-			b"ID3 " | b"id3 " => chunks.id3_chunk(data)?,
 			_ => chunks.skip(data)?,
 		}
 	}

src/iff/wav/tag/mod.rs

@@ -231,6 +231,8 @@ mod tests {
 	use crate::{Tag, TagIO, TagType};
 
 	use std::io::Cursor;
+	use byteorder::LittleEndian;
+	use crate::iff::chunk::Chunks;
 
 	#[test]
 	fn parse_riff_info() {
@@ -248,6 +250,7 @@ mod tests {
 
 		super::read::parse_riff_info(
 			&mut Cursor::new(&tag[..]),
+			&mut Chunks::<LittleEndian>::new(tag.len() as u32),
 			(tag.len() - 1) as u64,
 			&mut parsed_tag,
 		)
@@ -263,6 +266,7 @@ mod tests {
 
 		super::read::parse_riff_info(
 			&mut Cursor::new(&tag[..]),
+			&mut Chunks::<LittleEndian>::new(tag.len() as u32),
 			(tag.len() - 1) as u64,
 			&mut parsed_tag,
 		)
@@ -276,6 +280,7 @@ mod tests {
 		// Remove the LIST....INFO from the tag
 		super::read::parse_riff_info(
 			&mut Cursor::new(&writer[12..]),
+			&mut Chunks::<LittleEndian>::new(tag.len() as u32),
 			(tag.len() - 13) as u64,
 			&mut temp_parsed_tag,
 		)
@@ -291,7 +296,7 @@ mod tests {
 		let mut reader = std::io::Cursor::new(&tag_bytes[..]);
 		let mut riff_info = RiffInfoList::default();
 
-		super::read::parse_riff_info(&mut reader, (tag_bytes.len() - 1) as u64, &mut riff_info)
+		super::read::parse_riff_info(&mut reader, &mut Chunks::<LittleEndian>::new(tag_bytes.len() as u32), (tag_bytes.len() - 1) as u64, &mut riff_info)
 			.unwrap();
 
 		let tag: Tag = riff_info.into();

src/iff/wav/tag/read.rs

@@ -9,14 +9,13 @@ use byteorder::LittleEndian;
 
 pub(in crate::iff::wav) fn parse_riff_info<R>(
 	data: &mut R,
+	chunks: &mut Chunks<LittleEndian>,
 	end: u64,
 	tag: &mut RiffInfoList,
 ) -> Result<()>
 where
 	R: Read + Seek,
 {
-	let mut chunks = Chunks::<LittleEndian>::new();
-
 	while data.seek(SeekFrom::Current(0))? != end && chunks.next(data).is_ok() {
 		let key_str = String::from_utf8(chunks.fourcc.to_vec()).map_err(|_| {
 			FileDecodingError::new(FileType::WAV, "Non UTF-8 item key found in RIFF INFO")
@@ -34,7 +33,7 @@ where
 		}
 
 		tag.items.push((
-			key_str.to_string(),
+			key_str,
 			chunks.read_cstring(data).map_err(|_| {
 				FileDecodingError::new(FileType::WAV, "Failed to read RIFF INFO item value")
 			})?,

src/iff/wav/tag/write.rs

@@ -12,14 +12,12 @@ pub(in crate::iff::wav) fn write_riff_info(
 	data: &mut File,
 	tag: &mut RiffInfoListRef,
 ) -> Result<()> {
-	verify_wav(data)?;
+	let file_size = verify_wav(data)?;
 
 	let mut riff_info_bytes = Vec::new();
 	create_riff_info(&mut tag.items, &mut riff_info_bytes)?;
 
-	let (info_list, info_list_size) = find_info_list(data)?;
+	if let Some(info_list_size) = find_info_list(data, file_size)? {
-
-	if info_list {
 		let info_list_start = data.seek(SeekFrom::Current(-12))? as usize;
 		let info_list_end = info_list_start + 8 + info_list_size as usize;
 
@@ -50,13 +48,13 @@ pub(in crate::iff::wav) fn write_riff_info(
 	Ok(())
 }
 
-fn find_info_list<R>(data: &mut R) -> Result<(bool, u32)>
+fn find_info_list<R>(data: &mut R, file_size: u32) -> Result<Option<u32>>
 where
 	R: Read + Seek,
 {
-	let mut info = (false, 0);
+	let mut info = None;
 
-	let mut chunks = Chunks::<LittleEndian>::new();
+	let mut chunks = Chunks::<LittleEndian>::new(file_size);
 
 	while chunks.next(data).is_ok() {
 		if &chunks.fourcc == b"LIST" {
@@ -64,16 +62,14 @@ where
 			data.read_exact(&mut list_type)?;
 
 			if &list_type == b"INFO" {
-				info = (true, chunks.size);
+				info = Some(chunks.size);
 				break;
 			}
 
 			data.seek(SeekFrom::Current(-8))?;
 		}
 
-		data.seek(SeekFrom::Current(i64::from(chunks.size)))?;
+		chunks.skip(data)?;
-
-		chunks.correct_position(data)?;
 	}
 
 	Ok(info)

src/macros.rs

@@ -42,4 +42,14 @@ macro_rules! feature_locked {
 	}
 }
 
-pub(crate) use {feature_locked, tag_methods};
+macro_rules! try_vec {
+	($elem:expr; $size:expr) => {{
+		let mut v = Vec::new();
+		v.try_reserve($size)?;
+		v.resize($size, $elem);
+
+		v
+	}};
+}
+
+pub(crate) use {feature_locked, tag_methods, try_vec};