EBML: Support locking at multiple depths

Author: Serial-ATA

lofty/src/ebml/element_reader.rs

@@ -107,7 +107,7 @@ ebml_master_elements! {
 	Segment: {
 		id: 0x1853_8067,
 		children: [
-			SeekHead: { 0x114D_9B74, Master },
+			// SeekHead: { 0x114D_9B74, Master },
 			Info: { 0x1549_A966, Master },
 			Cluster: { 0x1F43_B675, Master },
 			Tracks: { 0x1654_AE6B, Master },
@@ -118,12 +118,12 @@ ebml_master_elements! {
 	},
 
 	// segment.seekHead
-	SeekHead: {
-		id: 0x114D_9B74,
-		children: [
-			Seek: { 0x4DBB, Master },
-		],
-	},
+	// SeekHead: {
+	// 	id: 0x114D_9B74,
+	// 	children: [
+	// 		Seek: { 0x4DBB, Master },
+	// 	],
+	// },
 
 	// segment.info
 	Info: {
@@ -190,6 +190,33 @@ ebml_master_elements! {
 		],
 	},
 
+	// segment.tags.tag.targets
+	Targets: {
+		id: 0x63C0,
+		children: [
+			TargetTypeValue: { 0x68CA, UnsignedInt },
+			TargetType: { 0x63CA, String },
+			TagTrackUID: { 0x63C5, UnsignedInt },
+			TagEditionUID: { 0x63C9, UnsignedInt },
+			TagChapterUID: { 0x63C4, UnsignedInt },
+			TagAttachmentUID: { 0x63C6, UnsignedInt },
+		],
+	},
+
+	// segment.tags.tag.simpleTag
+	SimpleTag: {
+		id: 0x67C8,
+		children: [
+			TagName: { 0x45A3, Utf8 },
+			TagLanguage: { 0x447A, String },
+			TagLanguageBCP47: { 0x447B, String },
+			TagDefault: { 0x4484, UnsignedInt },
+			TagDefaultBogus: { 0x44B4, UnsignedInt },
+			TagString: { 0x4487, Utf8 },
+			TagBinary: { 0x4485, Binary },
+		],
+	},
+
 	// segment.attachments
 	Attachments: {
 		id: 0x1941_A469,
@@ -235,8 +262,13 @@ struct ElementReaderContext {
 	/// This is set with [`ElementReader::lock`], and is used to prevent
 	/// the reader from reading past the end of the current master element.
 	locked: bool,
-	/// The depth at which we are locked to
-	lock_depth: u8,
+	/// The depths at which we are locked
+	///
+	/// When we reach the end of one lock and unlock the reader, we need
+	/// to know which depth to lock the reader at again (if any).
+	///
+	/// This will **always** be sorted, so the current lock will be at the end.
+	lock_depths: Vec<u8>,
 	lock_len: VInt,
 }
 
@@ -250,7 +282,7 @@ impl Default for ElementReaderContext {
 			// https://www.rfc-editor.org/rfc/rfc8794.html#name-ebmlmaxsizelength-element
 			max_size_length: 8,
 			locked: false,
-			lock_depth: 0,
+			lock_depths: Vec::with_capacity(MAX_DEPTH as usize),
 			lock_len: VInt::ZERO,
 		}
 	}
@@ -285,6 +317,7 @@ impl ElementReaderYield {
 	}
 }
 
+/// An EBML element reader.
 pub struct ElementReader<R> {
 	reader: R,
 	ctx: ElementReaderContext,
@@ -308,6 +341,11 @@ where
 		let ret = self.reader.read(buf)?;
 		let len = self.current_master_length();
 		self.set_current_master_length(len.saturating_sub(ret as u64));
+
+		if self.ctx.locked {
+			self.ctx.lock_len = self.ctx.lock_len.saturating_sub(ret as u64);
+		}
+
 		Ok(ret)
 	}
 }
@@ -355,10 +393,6 @@ where
 			return;
 		}
 
-		if self.ctx.locked {
-			self.ctx.lock_len = length;
-		}
-
 		self.ctx.masters[(self.ctx.depth - 1) as usize].remaining_length = length;
 	}
 
@@ -405,7 +439,7 @@ where
 	}
 
 	fn goto_previous_master(&mut self) -> Result<()> {
-		if self.ctx.depth == 0 || self.ctx.depth == self.ctx.lock_depth {
+		if self.ctx.depth == 0 || self.ctx.lock_depths.last() == Some(&self.ctx.depth) {
 			decode_err!(@BAIL Ebml, "Cannot go to previous master element, already at root")
 		}
 
@@ -468,10 +502,22 @@ where
 	pub(crate) fn lock(&mut self) {
 		self.ctx.locked = true;
 		self.ctx.lock_len = self.current_master_length();
+		self.ctx.lock_depths.push(self.ctx.depth);
 	}
 
 	pub(crate) fn unlock(&mut self) {
-		self.ctx.locked = false;
+		let _ = self.ctx.lock_depths.pop();
+
+		let [.., last] = &*self.ctx.lock_depths else {
+			// We can only ever *truly* unlock if we are at the root level.
+			log::trace!("Lock freed");
+
+			self.ctx.locked = false;
+			return;
+		};
+
+		log::trace!("Moving lock to depth: {}", last);
+		self.ctx.lock_len = self.ctx.masters[(*last - 1) as usize].remaining_length;
 	}
 
 	pub(crate) fn children(&mut self) -> ElementChildIterator<'_, R> {
@@ -595,6 +641,17 @@ where
 	}
 }
 
+/// An iterator over the children of an EBML master element.
+///
+/// This is created by calling [`ElementReader::children`].
+///
+/// This is essentially a fancy wrapper around `ElementReader` that:
+///
+/// * Automatically skips unknown elements ([`ElementReaderYield::Unknown`]).
+/// * [`Deref`]s to `ElementReader` so you can access the reader's methods.
+/// * Unlocks the reader when dropped.
+///   * If the reader is locked at multiple depths (meaning [`ElementReader::children`] was called
+///     multiple times), it will move the lock to the previously locked depth.
 pub(crate) struct ElementChildIterator<'a, R>
 where
 	R: Read,
@@ -622,10 +679,15 @@ where
 	}
 
 	pub(crate) fn master_exhausted(&self) -> bool {
-		let lock_depth = self.reader.ctx.lock_depth;
-		assert!(lock_depth < self.reader.ctx.depth);
-
-		self.reader.ctx.masters[lock_depth as usize].remaining_length == 0
+		let lock_depth = *self
+			.reader
+			.ctx
+			.lock_depths
+			.last()
+			.expect("a child iterator should always have a lock depth");
+		assert!(lock_depth <= self.reader.ctx.depth);
+
+		self.reader.ctx.masters[(lock_depth - 1) as usize].remaining_length == 0
 	}
 }
 

lofty/src/ebml/read/segment.rs

@@ -1,14 +1,10 @@
-use super::{
-	segment_attachments, segment_chapters, segment_cluster, segment_info, segment_tags,
-	segment_tracks,
-};
+use super::{segment_attachments, segment_cluster, segment_info, segment_tags, segment_tracks};
 use crate::config::ParseOptions;
 use crate::ebml::element_reader::{ElementHeader, ElementIdent, ElementReader, ElementReaderYield};
 use crate::ebml::properties::EbmlProperties;
 use crate::ebml::tag::EbmlTag;
 use crate::ebml::VInt;
 use crate::error::Result;
-use crate::macros::decode_err;
 
 use std::io::{Read, Seek};
 
@@ -21,60 +17,69 @@ where
 	R: Read + Seek,
 {
 	let mut tags = None;
+	let mut children_reader = element_reader.children();
 
-	element_reader.lock();
-
-	loop {
-		let child = element_reader.next()?;
-
+	while let Some(child) = children_reader.next()? {
 		match child {
-			ElementReaderYield::Master((id, size)) => match id {
-				ElementIdent::Info if parse_options.read_properties => {
-					segment_info::read_from(element_reader, parse_options, properties)?
-				},
-				ElementIdent::Cluster if parse_options.read_properties => {
-					segment_cluster::read_from(element_reader, parse_options, properties)?
-				},
-				ElementIdent::Tracks if parse_options.read_properties => {
-					segment_tracks::read_from(element_reader, parse_options, properties)?
-				},
-				// TODO: ElementIdent::Chapters
-				ElementIdent::Tags if parse_options.read_tags => {
-					let mut tag = tags.unwrap_or_default();
+			ElementReaderYield::Master((id, size)) => {
+				match id {
+					ElementIdent::Info if parse_options.read_properties => {
+						segment_info::read_from(
+							&mut children_reader.children(),
+							parse_options,
+							properties,
+						)?;
+					},
+					ElementIdent::Cluster if parse_options.read_properties => {
+						segment_cluster::read_from(
+							&mut children_reader.children(),
+							parse_options,
+							properties,
+						)?;
+					},
+					ElementIdent::Tracks if parse_options.read_properties => {
+						segment_tracks::read_from(
+							&mut children_reader.children(),
+							parse_options,
+							properties,
+						)?;
+					},
+					// TODO: ElementIdent::Chapters
+					ElementIdent::Tags if parse_options.read_tags => {
+						let mut tag = tags.unwrap_or_default();
 
-					segment_tags::read_from(element_reader, parse_options, &mut tag)?;
+						segment_tags::read_from(
+							&mut children_reader.children(),
+							parse_options,
+							&mut tag,
+						)?;
 
-					tags = Some(tag);
-				},
-				ElementIdent::Attachments if parse_options.read_cover_art => {
-					let mut tag = tags.unwrap_or_default();
+						tags = Some(tag);
+					},
+					ElementIdent::Attachments if parse_options.read_cover_art => {
+						let mut tag = tags.unwrap_or_default();
 
-					segment_attachments::read_from(element_reader, parse_options, &mut tag)?;
+						segment_attachments::read_from(
+							&mut children_reader.children(),
+							parse_options,
+							&mut tag,
+						)?;
 
-					tags = Some(tag);
-				},
-				_ => {
-					// We do not end up using information from all of the segment
-					// elements, so we can just skip any useless ones.
+						tags = Some(tag);
+					},
+					_ => {
+						// We do not end up using information from all of the segment
+						// elements, so we can just skip any useless ones.
 
-					element_reader.skip_element(ElementHeader {
-						id: VInt(id as u64),
-						size,
-					})?;
-					continue;
-				},
-			},
-			ElementReaderYield::Unknown(header) => {
-				element_reader.skip_element(header)?;
-				continue;
-			},
-			ElementReaderYield::Child(_) => {
-				decode_err!(@BAIL Ebml, "Segment element should only contain master elements")
-			},
-			ElementReaderYield::Eof => {
-				element_reader.unlock();
-				break;
+						children_reader.skip_element(ElementHeader {
+							id: VInt(id as u64),
+							size,
+						})?;
+					},
+				}
 			},
+			ElementReaderYield::Eof => break,
+			_ => unreachable!("Unhandled child element in \\Ebml\\Segment: {child:?}"),
 		}
 	}
 

lofty/src/ebml/read/segment_attachments.rs

@@ -1,5 +1,7 @@
 use crate::config::ParseOptions;
-use crate::ebml::element_reader::{ElementIdent, ElementReader, ElementReaderYield};
+use crate::ebml::element_reader::{
+	ElementChildIterator, ElementIdent, ElementReader, ElementReaderYield,
+};
 use crate::ebml::{AttachedFile, EbmlTag};
 use crate::error::Result;
 use crate::macros::decode_err;
@@ -8,19 +10,17 @@ use crate::picture::MimeType;
 use std::io::{Read, Seek};
 
 pub(super) fn read_from<R>(
-	element_reader: &mut ElementReader<R>,
+	children_reader: &mut ElementChildIterator<'_, R>,
 	_parse_options: ParseOptions,
 	tag: &mut EbmlTag,
 ) -> Result<()>
 where
 	R: Read + Seek,
 {
-	let mut children_reader = element_reader.children();
-
 	while let Some(child) = children_reader.next()? {
 		match child {
 			ElementReaderYield::Master((ElementIdent::AttachedFile, size)) => {
-				let attached_file = read_attachment(&mut children_reader)?;
+				let attached_file = read_attachment(children_reader)?;
 				tag.attached_files.push(attached_file);
 			},
 			ElementReaderYield::Eof => break,

lofty/src/ebml/read/segment_chapters.rs

@@ -1,13 +1,13 @@
 use crate::config::ParseOptions;
-use crate::ebml::element_reader::ElementReader;
+use crate::ebml::element_reader::ElementChildIterator;
 use crate::ebml::EbmlTag;
 use crate::error::Result;
 
 use std::io::{Read, Seek};
 
 #[allow(dead_code)]
 pub(super) fn read_from<R>(
-	_element_reader: &mut ElementReader<R>,
+	_children_reader: &mut ElementChildIterator<'_, R>,
 	_parse_options: ParseOptions,
 	_tag: &mut EbmlTag,
 ) -> Result<()>

lofty/src/ebml/read/segment_cluster.rs

@@ -1,12 +1,12 @@
 use crate::config::ParseOptions;
-use crate::ebml::element_reader::ElementReader;
+use crate::ebml::element_reader::ElementChildIterator;
 use crate::ebml::properties::EbmlProperties;
 use crate::error::Result;
 
 use std::io::{Read, Seek};
 
 pub(super) fn read_from<R>(
-	_element_reader: &mut ElementReader<R>,
+	_children_reader: &mut ElementChildIterator<'_, R>,
 	_parse_options: ParseOptions,
 	_properties: &mut EbmlProperties,
 ) -> Result<()>