From 9f9449ddd58eb07d20e14e7a75c7387c9cc17ebe Mon Sep 17 00:00:00 2001 From: metamuffin Date: Mon, 16 Jan 2023 10:57:40 +0100 Subject: write more code --- matroska/Cargo.toml | 10 ++ matroska/src/bin/experiment.rs | 32 ++++ matroska/src/bin/mkvdump.rs | 18 +++ matroska/src/lib.rs | 16 ++ matroska/src/matroska.rs | 330 +++++++++++++++++++++++++++++++++++++++++ matroska/src/read.rs | 227 ++++++++++++++++++++++++++++ matroska/src/size.rs | 21 +++ matroska/src/unflatten.rs | 74 +++++++++ matroska/src/write.rs | 169 +++++++++++++++++++++ 9 files changed, 897 insertions(+) create mode 100644 matroska/Cargo.toml create mode 100644 matroska/src/bin/experiment.rs create mode 100644 matroska/src/bin/mkvdump.rs create mode 100644 matroska/src/lib.rs create mode 100644 matroska/src/matroska.rs create mode 100644 matroska/src/read.rs create mode 100644 matroska/src/size.rs create mode 100644 matroska/src/unflatten.rs create mode 100644 matroska/src/write.rs (limited to 'matroska') diff --git a/matroska/Cargo.toml b/matroska/Cargo.toml new file mode 100644 index 0000000..0d6e492 --- /dev/null +++ b/matroska/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "jellymatroska" +version = "0.1.0" +edition = "2021" + +[dependencies] +ebml_derive = { path = "../ebml_derive" } +anyhow = "1.0.68" +log = "0.4.17" +env_logger = "0.10.0" diff --git a/matroska/src/bin/experiment.rs b/matroska/src/bin/experiment.rs new file mode 100644 index 0000000..caeae09 --- /dev/null +++ b/matroska/src/bin/experiment.rs @@ -0,0 +1,32 @@ +use jellymatroska::{matroska::MatroskaTag, read::EbmlReader, write::EbmlWriter}; +use std::{ + fs::File, + io::{stdout, BufReader, BufWriter}, +}; + +fn main() -> anyhow::Result<()> { + env_logger::init_from_env("LOG"); + let path = std::env::args().skip(1).next().unwrap(); + let mut r = EbmlReader::new(BufReader::new(File::open(path)?)); + let mut w = EbmlWriter::new(BufWriter::new(stdout()), 0); + + // r.seek( + // 631147167 + 52, + // ebml::matroska::MatroskaTag::Cues(Master::Start), + // ) + // .unwrap(); + + while let Some(tag) = r.next() { + let tag = tag?; + // println!("{} {tag:?}", r.position); + match tag { + tag => { + if MatroskaTag::is_master(tag.id())? { + eprintln!("{tag:?}"); + } + w.write_tag(&tag)?; + } + } + } + Ok(()) +} diff --git a/matroska/src/bin/mkvdump.rs b/matroska/src/bin/mkvdump.rs new file mode 100644 index 0000000..d5b6de4 --- /dev/null +++ b/matroska/src/bin/mkvdump.rs @@ -0,0 +1,18 @@ +use jellymatroska::{matroska::MatroskaTag, read::EbmlReader}; +use std::{fs::File, io::BufReader}; + +fn main() -> anyhow::Result<()> { + env_logger::init_from_env("LOG"); + let path = std::env::args().skip(1).next().unwrap(); + let mut r = EbmlReader::new(BufReader::new(File::open(path)?)); + + while let Some(tag) = r.next() { + let tag = tag?; + match tag { + MatroskaTag::SimpleBlock(_) => (), + MatroskaTag::Block(_) => (), + _ => eprintln!("{} {tag:?}", r.position), + } + } + Ok(()) +} diff --git a/matroska/src/lib.rs b/matroska/src/lib.rs new file mode 100644 index 0000000..bcba8fd --- /dev/null +++ b/matroska/src/lib.rs @@ -0,0 +1,16 @@ +pub mod matroska; +pub mod read; +pub mod size; +pub mod write; +pub mod unflatten; + +use matroska::MatroskaTag; +pub use read::ReadValue; +pub use write::WriteValue; + +#[derive(Debug, Clone, PartialEq)] +pub enum Master { + Collected(Vec), + Start, + End, +} diff --git a/matroska/src/matroska.rs b/matroska/src/matroska.rs new file mode 100644 index 0000000..b4078ab --- /dev/null +++ b/matroska/src/matroska.rs @@ -0,0 +1,330 @@ +use ebml_derive::define_ebml; + +define_ebml! { + global Crc32[0xbf]: Binary, + global Void[0xec]: Binary, + + Ebml[0x1a45dfa3]: { + EbmlVersion[0x4286]: Uint, + EbmlReadVersion[0x42f7]: Uint, + EbmlMaxIdLength[0x42f2]: Uint, + EbmlMaxSizeLength[0x42f3]: Uint, + DocType[0x4282]: Utf8, + DocTypeVersion[0x4287]: Uint, + DocTypeReadVersion[0x4285]: Uint, + DocTypeExtension[0x4281]: { + DocTypeExtensionName[0x4283]: Utf8, + DocTypeExtensionVersion[0x4284]: Uint, + }, + }, + + Segment[0x18538067]: { + Attachments[0x1941A469]: { + AttachedFile[0x61A7]: { + FileData[0x465C]: Binary, + FileDescription[0x467E]: Utf8, + FileMimeType[0x4660]: Utf8, + FileName[0x466E]: Utf8, + FileReferral[0x4675]: Binary, + FileUID[0x46AE]: Uint, + FileUsedEndTime[0x4662]: Uint, + FileUsedStartTime[0x4661]: Uint, + }, + }, + + Chapters[0x1043A770]: { + EditionEntry[0x45B9]: { + ChapterAtom[0xB6]: { + ChapProcess[0x6944]: { + ChapProcessCodecID[0x6955]: Uint, + ChapProcessCommand[0x6911]: { + ChapProcessData[0x6933]: Binary, + ChapProcessTime[0x6922]: Uint, + }, + ChapProcessPrivate[0x450D]: Binary, + }, + ChapterDisplay[0x80]: { + ChapCountry[0x437E]: Utf8, + ChapLanguage[0x437C]: Utf8, + ChapLanguageIETF[0x437D]: Utf8, + ChapString[0x85]: Utf8, + }, + ChapterFlagEnabled[0x4598]: Uint, + ChapterFlagHidden[0x98]: Uint, + ChapterPhysicalEquiv[0x63C3]: Uint, + ChapterSegmentEditionUID[0x6EBC]: Uint, + ChapterSegmentUID[0x6E67]: Binary, + ChapterStringUID[0x5654]: Utf8, + ChapterTimeEnd[0x92]: Uint, + ChapterTimeStart[0x91]: Uint, + ChapterUID[0x73C4]: Uint, + ChapterTrack[0x8F]: { + ChapterTrackUID[0x89]: Uint, + }, + }, + EditionFlagDefault[0x45DB]: Uint, + EditionFlagHidden[0x45BD]: Uint, + EditionFlagOrdered[0x45DD]: Uint, + EditionUID[0x45BC]: Uint, + }, + }, + + Cluster[0x1F43B675]: { + BlockGroup[0xA0]: { + Block[0xA1]: Binary, + BlockAdditions[0x75A1]: { + BlockMore[0xA6]: { + BlockAddID[0xEE]: Uint, + BlockAdditional[0xA5]: Binary, + }, + }, + BlockDuration[0x9B]: Uint, + BlockVirtual[0xA2]: Binary, + CodecState[0xA4]: Binary, + DiscardPadding[0x75A2]: Int, + ReferenceBlock[0xFB]: Int, + ReferenceFrame[0xC8]: { + ReferenceOffset[0xC9]: Uint, + ReferenceTimestamp[0xCA]: Uint, + }, + ReferencePriority[0xFA]: Uint, + ReferenceVirtual[0xFD]: Int, + Slices[0x8E]: { + TimeSlice[0xE8]: { + BlockAdditionID[0xCB]: Uint, + Delay[0xCE]: Uint, + FrameNumber[0xCD]: Uint, + LaceNumber[0xCC]: Uint, + SliceDuration[0xCF]: Uint, + }, + }, + }, + EncryptedBlock[0xAF]: Binary, + Position[0xA7]: Uint, + PrevSize[0xAB]: Uint, + SilentTracks[0x5854]: { + SilentTrackNumber[0x58D7]: Uint, + }, + SimpleBlock[0xA3]: Binary, + Timestamp[0xE7]: Uint, + }, + + Cues[0x1C53BB6B]: { + CuePoint[0xBB]: { + CueTime[0xB3]: Uint, + CueTrackPositions[0xB7]: { + CueBlockNumber[0x5378]: Uint, + CueClusterPosition[0xF1]: Uint, + CueCodecState[0xEA]: Uint, + CueDuration[0xB2]: Uint, + CueReference[0xDB]: { + CueRefCluster[0x97]: Uint, + CueRefCodecState[0xEB]: Uint, + CueRefNumber[0x535F]: Uint, + CueRefTime[0x96]: Uint, + }, + CueRelativePosition[0xF0]: Uint, + CueTrack[0xF7]: Uint, + }, + }, + }, + + Info[0x1549A966]: { + ChapterTranslate[0x6924]: { + ChapterTranslateCodec[0x69BF]: Uint, + ChapterTranslateEditionUID[0x69FC]: Uint, + ChapterTranslateID[0x69A5]: Binary, + }, + DateUTC[0x4461]: Int, + Duration[0x4489]: Float, + MuxingApp[0x4D80]: Utf8, + NextFilename[0x3E83BB]: Utf8, + NextUID[0x3EB923]: Binary, + PrevFilename[0x3C83AB]: Utf8, + PrevUID[0x3CB923]: Binary, + SegmentFamily[0x4444]: Binary, + SegmentFilename[0x7384]: Utf8, + SegmentUID[0x73A4]: Binary, + TimestampScale[0x2AD7B1]: Uint, + Title[0x7BA9]: Utf8, + WritingApp[0x5741]: Utf8, + }, + + SeekHead[0x114D9B74]: { + Seek[0x4DBB]: { + SeekID[0x53AB]: Binary, + SeekPosition[0x53AC]: Uint, + }, + }, + + Tags[0x1254C367]: { + Tag[0x7373]: { + SimpleTag[0x67C8]: { + TagBinary[0x4485]: Binary, + TagDefault[0x4484]: Uint, + TagDefaultBogus[0x44B4]: Uint, + TagLanguage[0x447A]: Utf8, + TagLanguageIETF[0x447B]: Utf8, + TagName[0x45A3]: Utf8, + TagString[0x4487]: Utf8, + }, + Targets[0x63C0]: { + TagAttachmentUID[0x63C6]: Uint, + TagChapterUID[0x63C4]: Uint, + TagEditionUID[0x63C9]: Uint, + TagTrackUID[0x63C5]: Uint, + TargetType[0x63CA]: Utf8, + TargetTypeValue[0x68CA]: Uint, + }, + }, + }, + + Tracks[0x1654AE6B]: { + TrackEntry[0xAE]: { + AttachmentLink[0x7446]: Uint, + Audio[0xE1]: { + BitDepth[0x6264]: Uint, + ChannelPositions[0x7D7B]: Binary, + Channels[0x9F]: Uint, + OutputSamplingFrequency[0x78B5]: Float, + SamplingFrequency[0xB5]: Float, + }, + BlockAdditionMapping[0x41E4]: { + BlockAddIDExtraData[0x41ED]: Binary, + BlockAddIDName[0x41A4]: Utf8, + BlockAddIDType[0x41E7]: Uint, + BlockAddIDValue[0x41F0]: Uint, + }, + CodecDecodeAll[0xAA]: Uint, + CodecDelay[0x56AA]: Uint, + CodecDownloadURL[0x26B240]: Utf8, + CodecID[0x86]: Utf8, + CodecInfoURL[0x3B4040]: Utf8, + CodecName[0x258688]: Utf8, + CodecPrivate[0x63A2]: Binary, + CodecSettings[0x3A9697]: Utf8, + ContentEncodings[0x6D80]: { + ContentEncoding[0x6240]: { + ContentCompression[0x5034]: { + ContentCompAlgo[0x4254]: Uint, + ContentCompSettings[0x4255]: Binary, + }, + ContentEncodingOrder[0x5031]: Uint, + ContentEncodingScope[0x5032]: Uint, + ContentEncodingType[0x5033]: Uint, + ContentEncryption[0x5035]: { + ContentEncAESSettings[0x47E7]: { + AESSettingsCipherMode[0x47E8]: Uint, + }, + ContentEncAlgo[0x47E1]: Uint, + ContentEncKeyID[0x47E2]: Binary, + ContentSigAlgo[0x47E5]: Uint, + ContentSigHashAlgo[0x47E6]: Uint, + ContentSigKeyID[0x47E4]: Binary, + ContentSignature[0x47E3]: Binary, + }, + }, + }, + DefaultDecodedFieldDuration[0x234E7A]: Uint, + DefaultDuration[0x23E383]: Uint, + FlagCommentary[0x55AF]: Uint, + FlagDefault[0x88]: Uint, + FlagEnabled[0xB9]: Uint, + FlagForced[0x55AA]: Uint, + FlagHearingImpaired[0x55AB]: Uint, + FlagLacing[0x9C]: Uint, + FlagOriginal[0x55AE]: Uint, + FlagTextDescriptions[0x55AD]: Uint, + FlagVisualImpaired[0x55AC]: Uint, + Language[0x22B59C]: Utf8, + LanguageIETF[0x22B59D]: Utf8, + MaxBlockAdditionID[0x55EE]: Uint, + MaxCache[0x6DF8]: Uint, + MinCache[0x6DE7]: Uint, + Name[0x536E]: Utf8, + SeekPreRoll[0x56BB]: Uint, + TrackNumber[0xD7]: Uint, + TrackOffset[0x537F]: Int, + TrackOperation[0xE2]: { + TrackCombinePlanes[0xE3]: { + TrackPlane[0xE4]: { + TrackPlaneType[0xE6]: Uint, + TrackPlaneUID[0xE5]: Uint, + }, + }, + TrackJoinBlocks[0xE9]: { + TrackJoinUID[0xED]: Uint, + }, + }, + TrackOverlay[0x6FAB]: Uint, + TrackTimestampScale[0x23314F]: Float, + TrackTranslate[0x6624]: { + TrackTranslateCodec[0x66BF]: Uint, + TrackTranslateEditionUID[0x66FC]: Uint, + TrackTranslateTrackID[0x66A5]: Binary, + }, + TrackType[0x83]: Uint, + TrackUID[0x73C5]: Uint, + TrickMasterTrackSegmentUID[0xC4]: Binary, + TrickMasterTrackUID[0xC7]: Uint, + TrickTrackFlag[0xC6]: Uint, + TrickTrackSegmentUID[0xC1]: Binary, + TrickTrackUID[0xC0]: Uint, + Video[0xE0]: { + AlphaMode[0x53C0]: Uint, + AspectRatioType[0x54B3]: Uint, + Colour[0x55B0]: { + BitsPerChannel[0x55B2]: Uint, + CbSubsamplingHorz[0x55B5]: Uint, + CbSubsamplingVert[0x55B6]: Uint, + ChromaSitingHorz[0x55B7]: Uint, + ChromaSitingVert[0x55B8]: Uint, + ChromaSubsamplingHorz[0x55B3]: Uint, + ChromaSubsamplingVert[0x55B4]: Uint, + MasteringMetadata[0x55D0]: { + LuminanceMax[0x55D9]: Float, + LuminanceMin[0x55DA]: Float, + PrimaryBChromaticityX[0x55D5]: Float, + PrimaryBChromaticityY[0x55D6]: Float, + PrimaryGChromaticityX[0x55D3]: Float, + PrimaryGChromaticityY[0x55D4]: Float, + PrimaryRChromaticityX[0x55D1]: Float, + PrimaryRChromaticityY[0x55D2]: Float, + WhitePointChromaticityX[0x55D7]: Float, + WhitePointChromaticityY[0x55D8]: Float, + }, + MatrixCoefficients[0x55B1]: Uint, + MaxCLL[0x55BC]: Uint, + MaxFALL[0x55BD]: Uint, + Primaries[0x55BB]: Uint, + Range[0x55B9]: Uint, + TransferCharacteristics[0x55BA]: Uint, + }, + DisplayHeight[0x54BA]: Uint, + DisplayUnit[0x54B2]: Uint, + DisplayWidth[0x54B0]: Uint, + FieldOrder[0x9D]: Uint, + FlagInterlaced[0x9A]: Uint, + FrameRate[0x2383E3]: Float, + GammaValue[0x2FB523]: Float, + OldStereoMode[0x53B9]: Uint, + PixelCropBottom[0x54AA]: Uint, + PixelCropLeft[0x54CC]: Uint, + PixelCropRight[0x54DD]: Uint, + PixelCropTop[0x54BB]: Uint, + PixelHeight[0xBA]: Uint, + PixelWidth[0xB0]: Uint, + Projection[0x7670]: { + ProjectionPosePitch[0x7674]: Float, + ProjectionPoseRoll[0x7675]: Float, + ProjectionPoseYaw[0x7673]: Float, + ProjectionPrivate[0x7672]: Binary, + ProjectionType[0x7671]: Uint, + }, + StereoMode[0x53B8]: Uint, + UncompressedFourCC[0x2EB524]: Binary, + }, + }, + }, + }, +} \ No newline at end of file diff --git a/matroska/src/read.rs b/matroska/src/read.rs new file mode 100644 index 0000000..95a98b5 --- /dev/null +++ b/matroska/src/read.rs @@ -0,0 +1,227 @@ +use crate::{matroska::MatroskaTag, size::EbmlSize, Master}; +use anyhow::{anyhow, bail, Result}; +use log::{debug, warn}; +use std::{ + collections::VecDeque, + io::{Read, Seek, SeekFrom}, +}; + +trait ReadAndSeek: Read + Seek {} +impl ReadAndSeek for T {} + +#[derive(Debug, Clone, Copy)] +pub struct StackTag { + end: Option, + id: u64, +} + +pub struct EbmlReader { + inner: Box, + stack: Vec, + queue: VecDeque, + pub position: usize, +} + +impl EbmlReader { + pub fn new(inner: T) -> Self { + Self { + queue: VecDeque::new(), + inner: Box::new(inner), + stack: vec![], + position: 0, + } + } + + pub fn read_byte(&mut self) -> Result { + let mut b = [0u8]; + self.inner.read_exact(&mut b)?; + self.position += 1; + Ok(b[0]) + } + pub fn read_buf(&mut self, size: impl Into) -> Result> { + let size = size.into(); + let mut b = vec![0u8; size]; + self.inner.read_exact(&mut b)?; + self.position += size; + Ok(b) + } + pub fn read_vint_len(&mut self) -> Result<(u64, usize)> { + let s = self.read_byte()?; + let len = s.leading_zeros() + 1; + if len > 8 { + bail!("varint too long"); + } + let mut value = s as u64; + value -= 1 << (8 - len); + for _ in 1..len { + value <<= 8; + value += self.read_byte()? as u64; + } + Ok((value, len as usize)) + } + pub fn read_vint(&mut self) -> Result { + Ok(self.read_vint_len()?.0) + } + pub fn read_utf8(&mut self, size: impl Into) -> Result { + let b = self.read_buf(size)?; + Ok(String::from_utf8(b)?) + } + pub fn read_tag_id(&mut self) -> Result { + let (value, len) = self.read_vint_len()?; + Ok(value + (1 << (7 * len))) + } + pub fn read_tag_size(&mut self) -> Result { + Ok(EbmlSize::from_vint(self.read_vint_len()?)) + } + pub fn read_stuff(&mut self) -> Result<()> { + while let Some(e) = self.stack.last().map(|e| *e) { + if let Some(end) = e.end { + if self.position >= end { + if self.position != end { + warn!("we missed the end by {} bytes", self.position - end) + } + self.stack.pop(); + self.queue + .push_back(MatroskaTag::construct_master(e.id, Master::End)?); + } else { + break; + } + } else { + break; + } + } + + let id = self.read_tag_id()?; + let size = self.read_tag_size()?; + let is_master = MatroskaTag::is_master(id)?; + let tag = if is_master { + MatroskaTag::construct_master(id, Master::Start)? + } else { + let data = self.read_buf(size.some().unwrap())?; + MatroskaTag::parse(id, &data)? + }; + + if let Some(path) = tag.path() { + // we have slightly different rules for closing tags implicitly + // this closes as many tags as needed to make the next tag a valid child + while let Some(tag @ StackTag { end: None, .. }) = self.stack.last() { + if path.last() == Some(&tag.id) { + break; + } else { + self.queue.push_back(MatroskaTag::construct_master( + self.stack.pop().unwrap().id, + Master::End, + )?); + } + } + } + + if is_master { + self.stack.push(StackTag { + end: size.some().map(|s| s + self.position), + id, + }); + } + self.queue.push_back(tag); + Ok(()) + } + + /// context should be the next expected tag, such that the stack can be derived from its path. + pub fn seek(&mut self, position: usize, context: MatroskaTag) -> Result<()> { + let path = context + .path() + .ok_or(anyhow!("global tags dont give context"))?; + debug!( + "seeking to {position} with a context restored from path {:x?}", + path + ); + self.queue.clear(); + self.position = position; + self.inner.seek(SeekFrom::Start(position as u64))?; + self.stack = path + .iter() + .map(|id| StackTag { id: *id, end: None }) + .collect(); + Ok(()) + } +} + +impl Iterator for EbmlReader { + type Item = Result; + + fn next(&mut self) -> Option { + if let Some(t) = self.queue.pop_front() { + // match t { + // MatroskaTag::SimpleBlock(_) | MatroskaTag::Block(_) => (), + // _ => debug!("reader yield: {t:?}"), + // }; + Some(Ok(t)) + } else { + match self.read_stuff() { + Ok(()) => self.next(), + Err(e) => Some(Err(e)), + } + } + } +} + +pub trait ReadValue: Sized { + fn from_buf(buf: &[u8]) -> anyhow::Result; +} + +impl ReadValue for u64 { + fn from_buf(buf: &[u8]) -> anyhow::Result { + if buf.len() > 8 { + bail!("u64 too big") + } + let mut val = 0u64; + for byte in buf { + val <<= 8; + val |= *byte as u64; + } + Ok(val) + } +} +impl ReadValue for i64 { + fn from_buf(buf: &[u8]) -> anyhow::Result { + if buf.len() > 8 { + bail!("i64 too big") + } + Ok(if buf[0] > 127 { + if buf.len() == 8 { + i64::from_be_bytes(buf.try_into().unwrap()) + } else { + -((1 << (buf.len() * 8)) - (u64::from_buf(buf)? as i64)) + } + } else { + u64::from_buf(buf)? as i64 + }) + } +} +impl ReadValue for f64 { + fn from_buf(buf: &[u8]) -> anyhow::Result { + Ok(if buf.len() == 4 { + f32::from_be_bytes(buf.try_into().unwrap()) as f64 + } else if buf.len() == 8 { + f64::from_be_bytes(buf.try_into().unwrap()) + } else { + bail!("float is not 4 or 8 bytes long"); + }) + } +} + +impl ReadValue for Vec { + fn from_buf(buf: &[u8]) -> anyhow::Result { + Ok(buf.to_vec()) + } +} +impl ReadValue for String { + fn from_buf(buf: &[u8]) -> anyhow::Result { + Ok(String::from_utf8(Vec::from(buf))?) + } +} +impl ReadValue for Master { + fn from_buf(_: &[u8]) -> anyhow::Result { + panic!("master shall not be read like this") + } +} diff --git a/matroska/src/size.rs b/matroska/src/size.rs new file mode 100644 index 0000000..e774f0a --- /dev/null +++ b/matroska/src/size.rs @@ -0,0 +1,21 @@ + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EbmlSize { + Exact(usize), + Unknown, +} +impl EbmlSize { + pub fn from_vint((value, len): (u64, usize)) -> EbmlSize { + if value == ((1 << (7 * len)) - 1) { + Self::Unknown + } else { + Self::Exact(value as usize) + } + } + pub fn some(self) -> Option { + match self { + EbmlSize::Exact(s) => Some(s), + EbmlSize::Unknown => None, + } + } +} diff --git a/matroska/src/unflatten.rs b/matroska/src/unflatten.rs new file mode 100644 index 0000000..03ca0f5 --- /dev/null +++ b/matroska/src/unflatten.rs @@ -0,0 +1,74 @@ +use crate::{matroska::MatroskaTag, Master}; +use anyhow::Result; + +pub struct Unflat<'a> { + pub item: MatroskaTag, + pub children: Option>, +} + +pub struct Unflatten<'a> { + inner: &'a mut dyn Iterator>, + stop: bool, + end: Option, +} + +impl<'a> Unflatten<'a> { + pub fn new(inner: &'a mut dyn Iterator>) -> Self { + Self { + inner, + stop: false, + end: None, + } + } + pub fn new_with_end( + inner: &'a mut dyn Iterator>, + start: MatroskaTag, + ) -> Self { + Self { + inner, + stop: false, + end: Some(MatroskaTag::construct_master(start.id(), Master::End).unwrap()), + } + } + + pub fn next(&mut self) -> Option> { + if self.stop { + return None; + } + match self.inner.next() { + None => None, + Some(Err(e)) => Some(Err(e)), + Some(Ok(item)) => { + let master = MatroskaTag::is_master(item.id()).unwrap(); + if Some(&item) == self.end.as_ref() { + self.stop = true; + None + } else { + Some(Ok(Unflat { + children: if master { + let end = + MatroskaTag::construct_master(item.id(), Master::End).unwrap(); + if end == item { + return None; + } + Some(Unflatten { + inner: self.inner, + stop: false, + end: Some(end), + }) + } else { + None + }, + item, + })) + } + } + } + } +} + +impl Drop for Unflatten<'_> { + fn drop(&mut self) { + while let Some(_) = self.next() {} + } +} diff --git a/matroska/src/write.rs b/matroska/src/write.rs new file mode 100644 index 0000000..fc12ffc --- /dev/null +++ b/matroska/src/write.rs @@ -0,0 +1,169 @@ +use anyhow::{bail, Result}; +use std::io::Write; + +use crate::{matroska::MatroskaTag, size::EbmlSize, Master}; + +pub struct EbmlWriter { + inner: Box, + position: usize, +} + +impl EbmlWriter { + pub fn new(inner: T, position: usize) -> Self { + Self { + inner: Box::new(inner), + position, + } + } + + pub fn write(&mut self, data: &[u8]) -> Result<()> { + self.inner.write_all(data)?; + self.position += data.len(); + Ok(()) + } + + pub fn write_tag(&mut self, tag: &MatroskaTag) -> Result<()> { + let mut buf = vec![]; + tag.write_full(&mut buf)?; + self.write(&buf)?; + Ok(()) + } + + pub fn write_vint(&mut self, i: u64) -> Result<()> { + if i > (1 << 56) - 1 { + bail!("vint does not fit"); + } + let mut len = 1; + while len <= 8 { + if i < (1 << ((7 * len) - 1)) { + break; + } + len += 1; + } + let mut bytes = i.to_be_bytes(); + let trunc = &mut bytes[(8 - len)..]; + trunc[0] |= 1 << (8 - len); + self.write(&trunc) + } +} + +impl MatroskaTag { + pub fn write_full(&self, w: &mut Vec) -> Result<()> { + let mut buf = vec![]; + buf.extend(self.id().to_be_bytes().iter().skip_while(|&v| *v == 0u8)); + // note: it is relevant here, to pass the buffer with the id, such that closing tags, can clear it + self.write(&mut buf)?; + w.extend_from_slice(&buf); + Ok(()) + } +} + +pub trait WriteValue { + fn write_to(&self, w: &mut Vec) -> Result<()>; +} + +impl WriteValue for i64 { + fn write_to(&self, w: &mut Vec) -> Result<()> { + Ok(match 64 - self.leading_zeros() { + x if x <= 8 => { + w.push(0x81); + w.extend_from_slice(&(*self as i8).to_be_bytes()); + } + x if x <= 16 => { + w.push(0x82); + w.extend_from_slice(&(*self as i16).to_be_bytes()); + } + x if x <= 32 => { + w.push(0x84); + w.extend_from_slice(&(*self as i32).to_be_bytes()); + } + _ => { + w.push(0x88); + w.extend_from_slice(&self.to_be_bytes()); + } + }) + } +} +impl WriteValue for u64 { + fn write_to(&self, w: &mut Vec) -> Result<()> { + Ok(match 64 - self.leading_zeros() { + x if x <= 8 => { + w.push(0x81); + w.extend_from_slice(&(*self as u8).to_be_bytes()); + } + x if x <= 16 => { + w.push(0x82); + w.extend_from_slice(&(*self as u16).to_be_bytes()); + } + x if x <= 32 => { + w.push(0x84); + w.extend_from_slice(&(*self as u32).to_be_bytes()); + } + _ => { + w.push(0x88); + w.extend_from_slice(&self.to_be_bytes()); + } + }) + } +} +impl WriteValue for f64 { + fn write_to(&self, w: &mut Vec) -> Result<(), anyhow::Error> { + w.push(0x88); + w.extend_from_slice(&self.to_be_bytes()); + Ok(()) + } +} +impl WriteValue for Vec { + fn write_to(&self, w: &mut Vec) -> Result<(), anyhow::Error> { + write_vint(w, self.len() as u64)?; + w.extend_from_slice(&self); + Ok(()) + } +} +impl WriteValue for String { + fn write_to(&self, w: &mut Vec) -> Result<(), anyhow::Error> { + let sl = self.as_bytes(); + write_vint(w, sl.len() as u64)?; + w.extend_from_slice(sl); + Ok(()) + } +} +impl WriteValue for EbmlSize { + fn write_to(&self, w: &mut Vec) -> Result<()> { + match self { + EbmlSize::Exact(s) => write_vint(w, *s as u64)?, + EbmlSize::Unknown => w.extend_from_slice(&(u64::MAX >> 7).to_be_bytes()), + } + Ok(()) + } +} + +impl WriteValue for Master { + fn write_to(&self, w: &mut Vec) -> Result<()> { + match self { + Master::Start => EbmlSize::Unknown.write_to(w), + Master::End => Ok(w.clear()), + Master::Collected(c) => { + let mut ib = vec![]; + for c in c { + c.write_full(&mut ib)?; + } + EbmlSize::Exact(ib.len()).write_to(w)?; + w.extend_from_slice(&ib); + Ok(()) + } + } + } +} + +pub fn write_vint(w: &mut Vec, i: u64) -> Result<()> { + if i > (1 << 56) - 1 { + bail!("vint does not fit"); + } + let len = (64 - i.leading_zeros() as usize) / 7 + 1; + let mut bytes = i.to_be_bytes(); + let trunc = &mut bytes[(8 - len)..]; + trunc[0] |= 1 << (8 - len); + w.extend_from_slice(&trunc); + Ok(()) +} -- cgit v1.2.3-70-g09d2