Skip to content

Commit

Permalink
EBML: Support locking at multiple depths
Browse files Browse the repository at this point in the history
  • Loading branch information
Serial-ATA committed Aug 31, 2024
1 parent e641666 commit 66c0544
Show file tree
Hide file tree
Showing 8 changed files with 219 additions and 141 deletions.
102 changes: 82 additions & 20 deletions lofty/src/ebml/element_reader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 },
Expand All @@ -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: {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
}

Expand All @@ -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,
}
}
Expand Down Expand Up @@ -285,6 +317,7 @@ impl ElementReaderYield {
}
}

/// An EBML element reader.
pub struct ElementReader<R> {
reader: R,
ctx: ElementReaderContext,
Expand All @@ -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)
}
}
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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")
}

Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
}
}

Expand Down
107 changes: 56 additions & 51 deletions lofty/src/ebml/read/segment.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand All @@ -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:?}"),
}
}

Expand Down
10 changes: 5 additions & 5 deletions lofty/src/ebml/read/segment_attachments.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions lofty/src/ebml/read/segment_chapters.rs
Original file line number Diff line number Diff line change
@@ -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<()>
Expand Down
4 changes: 2 additions & 2 deletions lofty/src/ebml/read/segment_cluster.rs
Original file line number Diff line number Diff line change
@@ -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<()>
Expand Down
Loading

0 comments on commit 66c0544

Please sign in to comment.