Skip to main content

pito_codecs/mkv/
reader.rs

1use std::collections::HashSet;
2use std::fs::File;
3use std::path::Path;
4use std::sync::OnceLock;
5
6use super::ebml::{self, Children, Lacing, Size, parse_lacing, read_id, read_size};
7use super::ids::*;
8use super::parse::{self, Cluster, CuePoint, Info, SeekEntry};
9use super::track::{Sample, Track};
10use crate::media::{Error, Packet, Result};
11use crate::mp4::{ByteSource, TrackKind};
12
13const WIN: usize = 64 * 1024;
14const MAX_WIN: usize = 4 << 20;
15const MAX_META: u64 = 256 << 20;
16const MAX_SAMPLES: usize = 1 << 26;
17const MAX_WARNINGS: usize = 256;
18const CUE_SPAN: u64 = 64 << 20;
19
20fn warn(warnings: &mut Vec<String>, line: impl FnOnce() -> String) {
21    match warnings.len().cmp(&MAX_WARNINGS) {
22        std::cmp::Ordering::Less => warnings.push(line()),
23        std::cmp::Ordering::Equal => warnings.push("further damage not listed".into()),
24        std::cmp::Ordering::Greater => {}
25    }
26}
27
28#[derive(Clone, Copy, Debug, PartialEq, Eq)]
29pub struct Options {
30    pub index: bool,
31}
32
33impl Default for Options {
34    fn default() -> Self {
35        Self { index: true }
36    }
37}
38
39#[derive(Default, Clone, Debug)]
40pub(crate) struct Cache {
41    buf: Vec<u8>,
42    start: u64,
43    win: usize,
44    scanned: Option<Scanned>,
45}
46
47#[derive(Clone, Copy, Debug)]
48struct Scanned {
49    from: u64,
50    to: u64,
51    limit: u64,
52    found: bool,
53}
54
55pub(crate) struct Io<'a, S: ByteSource + ?Sized> {
56    src: &'a S,
57    pub(crate) len: u64,
58    cache: &'a mut Cache,
59}
60
61#[derive(Clone, Copy, Debug)]
62pub(crate) struct Hdr {
63    pub(crate) id: u32,
64    pub(crate) pos: u64,
65    pub(crate) size: Size,
66    hlen: usize,
67}
68
69impl Hdr {
70    pub(crate) fn data_start(&self) -> u64 {
71        self.pos + self.hlen as u64
72    }
73
74    pub(crate) fn end(&self, limit: u64) -> u64 {
75        match self.size {
76            Size::Known(n) => self.data_start().saturating_add(n).min(limit),
77            Size::Unknown => limit,
78        }
79    }
80
81    pub(crate) fn overruns(&self, limit: u64) -> bool {
82        matches!(self.size, Size::Known(n) if self.data_start().saturating_add(n) > limit)
83    }
84}
85
86fn eof() -> Error {
87    Error::Invalid("truncated Matroska data".into())
88}
89
90impl<'a, S: ByteSource + ?Sized> Io<'a, S> {
91    pub(crate) fn new(src: &'a S, cache: &'a mut Cache) -> Self {
92        let len = src.len();
93        Self { src, len, cache }
94    }
95
96    pub(crate) fn with_len(src: &'a S, len: u64, cache: &'a mut Cache) -> Self {
97        Self { src, len, cache }
98    }
99
100    fn cached(&self, pos: u64, end: u64) -> Option<std::ops::Range<usize>> {
101        let c = &*self.cache;
102        let from = usize::try_from(pos.checked_sub(c.start)?).ok()?;
103        let to = usize::try_from(end.checked_sub(c.start)?).ok()?;
104        (from <= to && to <= c.buf.len()).then_some(from..to)
105    }
106
107    fn refill(&mut self, pos: u64) -> Result<()> {
108        let c = &*self.cache;
109        let end = c.start.saturating_add(c.buf.len() as u64);
110        let ahead = !c.buf.is_empty() && pos >= c.start && pos <= end.saturating_add(c.win as u64);
111        let win = if ahead {
112            c.win.saturating_mul(2).clamp(WIN, MAX_WIN)
113        } else {
114            WIN
115        };
116        let fill = self.len.saturating_sub(pos).min(win as u64) as usize;
117        let mut buf = std::mem::take(&mut self.cache.buf);
118        buf.resize(fill, 0);
119        let read = self.src.read_at(pos, &mut buf);
120        if read.is_err() {
121            buf.clear();
122        }
123        self.cache.buf = buf;
124        self.cache.start = pos;
125        self.cache.win = win;
126        Ok(read?)
127    }
128
129    pub(crate) fn peek(&mut self, pos: u64, n: usize) -> Result<&[u8]> {
130        let n = n.min(WIN);
131        if pos >= self.len {
132            return Ok(&[]);
133        }
134        let want_end = pos.saturating_add(n as u64).min(self.len);
135        let range = match self.cached(pos, want_end) {
136            Some(r) => r,
137            None => {
138                self.refill(pos)?;
139                self.cached(pos, want_end).ok_or_else(eof)?
140            }
141        };
142        self.cache.buf.get(range).ok_or_else(eof)
143    }
144
145    pub(crate) fn read_into(&mut self, pos: u64, n: u64, out: &mut Vec<u8>) -> Result<()> {
146        let end = pos
147            .checked_add(n)
148            .filter(|&e| e <= self.len)
149            .ok_or_else(eof)?;
150        if let Some(r) = self.cached(pos, end) {
151            out.extend_from_slice(self.cache.buf.get(r).unwrap_or_default());
152            return Ok(());
153        }
154        let n = usize::try_from(n).map_err(|_| eof())?;
155        if n <= WIN / 2 {
156            let b = self.peek(pos, n)?;
157            out.extend_from_slice(b);
158            return Ok(());
159        }
160        let at = out.len();
161        out.resize(at.checked_add(n).ok_or_else(eof)?, 0);
162        self.src.read_at(pos, &mut out[at..])?;
163        Ok(())
164    }
165
166    fn read_vec(&mut self, pos: u64, n: u64) -> Result<Vec<u8>> {
167        let mut v = Vec::new();
168        self.read_into(pos, n, &mut v)?;
169        Ok(v)
170    }
171
172    pub(crate) fn header(&mut self, pos: u64, limit: u64) -> Result<Option<Hdr>> {
173        if pos >= limit {
174            return Ok(None);
175        }
176        let b = self.peek(pos, 12)?;
177        let Some((id, il)) = read_id(b) else {
178            return Ok(None);
179        };
180        let Some((size, sl)) = read_size(&b[il..]) else {
181            return Ok(None);
182        };
183        let hlen = il + sl;
184        if pos + hlen as u64 > limit {
185            return Ok(None);
186        }
187        Ok(Some(Hdr {
188            id,
189            pos,
190            size,
191            hlen,
192        }))
193    }
194
195    pub(crate) fn uint(&mut self, h: &Hdr, end: u64) -> Result<u64> {
196        let n = end.saturating_sub(h.data_start()).min(8) as usize;
197        Ok(ebml::uint(self.peek(h.data_start(), n)?))
198    }
199
200    pub(crate) fn resync(&mut self, from: u64, limit: u64) -> Result<Option<u64>> {
201        let memo = self
202            .cache
203            .scanned
204            .filter(|m| m.limit == limit && m.to >= from);
205        let last = limit.saturating_sub(3);
206        let mut pos = from;
207        loop {
208            if let Some(m) = memo
209                && m.from <= pos
210            {
211                return Ok(self.remember(from.min(m.from), m.to, limit, m.found));
212            }
213            let stop = memo.map_or(last, |m| m.from.min(last));
214            if pos >= stop {
215                return Ok(self.remember(from, limit, limit, false));
216            }
217            let hit = {
218                let chunk = self.peek(pos, WIN)?;
219                let n = chunk.len().saturating_sub(3);
220                if n == 0 {
221                    return Ok(None);
222                }
223                let n = n.min(usize::try_from(stop - pos).unwrap_or(usize::MAX));
224                let found = chunk.windows(4).take(n).enumerate().find_map(|(i, w)| {
225                    let &[a, b, c, d] = w else {
226                        return None;
227                    };
228                    if a & 0xF0 != 0x10 {
229                        return None;
230                    }
231                    let id = u32::from_be_bytes([a, b, c, d]);
232                    (is_top_level(id) && id != EBML && id != SEGMENT).then_some((i, id))
233                });
234                match found {
235                    Some((i, id)) => Ok((pos + i as u64, id)),
236                    None => Err(pos + n as u64),
237                }
238            };
239            match hit {
240                Ok((p, id)) => {
241                    if self.plausible(p, id, limit)? {
242                        return Ok(self.remember(from, p, limit, true));
243                    }
244                    pos = p + 1;
245                }
246                Err(next) => pos = next,
247            }
248        }
249    }
250
251    fn remember(&mut self, from: u64, to: u64, limit: u64, found: bool) -> Option<u64> {
252        self.cache.scanned = Some(Scanned {
253            from,
254            to,
255            limit,
256            found,
257        });
258        found.then_some(to)
259    }
260
261    fn plausible(&mut self, p: u64, id: u32, limit: u64) -> Result<bool> {
262        let Some(h) = self.header(p, limit)? else {
263            return Ok(false);
264        };
265        if h.id != id || h.overruns(limit) {
266            return Ok(false);
267        }
268        if h.size == Size::Unknown && id != CLUSTER {
269            return Ok(false);
270        }
271        let mut cp = h.data_start();
272        let end = h.end(limit);
273        for _ in 0..3 {
274            let Some(c) = self.header(cp, end)? else {
275                return Ok(false);
276            };
277            if c.overruns(end) {
278                return Ok(false);
279            }
280            if id != CLUSTER {
281                return Ok(true);
282            }
283            match c.id {
284                TIMESTAMP => return Ok(true),
285                CRC32 | VOID => cp = c.end(end),
286                _ => return Ok(false),
287            }
288        }
289        Ok(false)
290    }
291
292    pub(crate) fn unknown_end(&mut self, data_start: u64, limit: u64) -> Result<u64> {
293        let mut pos = data_start;
294        while pos < limit {
295            let Some(h) = self.header(pos, limit)? else {
296                return Ok(pos);
297            };
298            if is_top_level(h.id) || h.size == Size::Unknown {
299                return Ok(pos);
300            }
301            pos = h.end(limit);
302        }
303        Ok(limit)
304    }
305}
306
307#[derive(Clone, Copy, Debug)]
308pub(crate) struct BlockRef {
309    pub(crate) track: usize,
310    pub(crate) offset: u64,
311    pub(crate) ts: i64,
312    pub(crate) duration: Option<u64>,
313    pub(crate) keyframe: bool,
314}
315
316#[derive(Clone, Copy, Debug)]
317pub(crate) struct FrameRef {
318    pub(crate) track: usize,
319    pub(crate) offset: u64,
320    pub(crate) size: u32,
321    pub(crate) pts: i64,
322    pub(crate) duration: u32,
323    pub(crate) keyframe: bool,
324    pub(crate) block: u64,
325    pub(crate) lace: u16,
326}
327
328pub(crate) fn track_index(tracks: &[Track], number: u64) -> Option<usize> {
329    tracks.iter().position(|t| t.number == number)
330}
331
332fn parse_block<S: ByteSource + ?Sized>(
333    io: &mut Io<'_, S>,
334    data_start: u64,
335    len: u64,
336    tracks: &[Track],
337    frames: &mut Vec<(u64, u32)>,
338) -> Result<Option<(usize, i16, u8)>> {
339    frames.clear();
340    let head = io.peek(data_start, 16.min(len as usize))?;
341    let Some((number, tl, _)) = ebml::read_vint(head) else {
342        return Ok(None);
343    };
344    if head.len() < tl + 3 {
345        return Ok(None);
346    }
347    let rel = i16::from_be_bytes([head[tl], head[tl + 1]]);
348    let flags = head[tl + 2];
349    let Some(track) = track_index(tracks, number) else {
350        return Ok(None);
351    };
352    let hdr = (tl + 3) as u64;
353    let lacing = Lacing::from_flags(flags);
354    if lacing == Lacing::None {
355        frames.push((data_start + hdr, (len - hdr) as u32));
356    } else {
357        let body = io.read_vec(data_start + hdr, len - hdr)?;
358        match parse_lacing(lacing, &body) {
359            Ok(f) => frames.extend(
360                f.into_iter()
361                    .map(|(o, l)| (data_start + hdr + o as u64, l as u32)),
362            ),
363            Err(_) => return Ok(None),
364        }
365    }
366    Ok(Some((track, rel, flags)))
367}
368
369pub(crate) fn parse_block_elem<S: ByteSource + ?Sized>(
370    io: &mut Io<'_, S>,
371    h: &Hdr,
372    end: u64,
373    cluster_ts: i64,
374    tracks: &[Track],
375    frames: &mut Vec<(u64, u32)>,
376) -> Result<Option<BlockRef>> {
377    let ds = h.data_start();
378    match h.id {
379        SIMPLE_BLOCK => {
380            let Some((track, rel, flags)) = parse_block(io, ds, end - ds, tracks, frames)? else {
381                return Ok(None);
382            };
383            Ok(Some(BlockRef {
384                track,
385                offset: h.pos,
386                ts: cluster_ts.saturating_add(rel as i64),
387                duration: None,
388                keyframe: flags & 0x80 != 0,
389            }))
390        }
391        BLOCK_GROUP => {
392            let mut pos = ds;
393            let mut block = None;
394            let (mut duration, mut has_ref) = (None, false);
395            while pos < end {
396                let Some(c) = io.header(pos, end)? else {
397                    break;
398                };
399                let ce = c.end(end);
400                match c.id {
401                    BLOCK => {
402                        block =
403                            parse_block(io, c.data_start(), ce - c.data_start(), tracks, frames)?
404                    }
405                    BLOCK_DURATION => duration = Some(io.uint(&c, ce)?),
406                    REFERENCE_BLOCK => has_ref = true,
407                    _ => {}
408                }
409                pos = ce;
410            }
411            let Some((track, rel, _)) = block else {
412                return Ok(None);
413            };
414            Ok(Some(BlockRef {
415                track,
416                offset: h.pos,
417                ts: cluster_ts.saturating_add(rel as i64),
418                duration,
419                keyframe: !has_ref,
420            }))
421        }
422        _ => Ok(None),
423    }
424}
425
426pub(crate) fn expand(
427    b: &BlockRef,
428    frames: &[(u64, u32)],
429    t: &Track,
430    scale: u64,
431    skip: u16,
432    mut out: impl FnMut(FrameRef),
433) {
434    let n = frames.len() as u64;
435    let block_ns: Option<u128> = match (b.duration, t.default_duration) {
436        (Some(d), _) => Some(d as u128 * scale as u128),
437        (None, Some(dd)) => Some(dd as u128 * n as u128),
438        _ => None,
439    };
440    let scale = scale as i128;
441    let ticks = |ns: i128| -> i64 {
442        (if ns >= 0 {
443            (ns + scale / 2) / scale
444        } else {
445            (ns - scale / 2) / scale
446        })
447        .clamp(i64::MIN as i128, i64::MAX as i128) as i64
448    };
449    let delay = ticks(t.codec_delay as i128);
450    let base = b.ts as i128 * scale;
451    for (i, &(offset, size)) in frames.iter().enumerate().skip(skip as usize) {
452        let (start, end) = match block_ns {
453            Some(total) => (
454                (total * i as u128 / n as u128) as i128,
455                Some((total * (i as u128 + 1) / n as u128) as i128),
456            ),
457            None => (0, None),
458        };
459        let duration = end.map_or(0, |e| {
460            ticks(e)
461                .saturating_sub(ticks(start))
462                .clamp(0, u32::MAX as i64)
463        });
464        out(FrameRef {
465            track: b.track,
466            offset,
467            size,
468            pts: ticks(base + start).saturating_sub(delay),
469            duration: duration as u32,
470            keyframe: b.keyframe,
471            block: b.offset,
472            lace: i as u16,
473        });
474    }
475}
476
477pub(crate) fn scan_cluster<S: ByteSource + ?Sized>(
478    io: &mut Io<'_, S>,
479    h: &Hdr,
480    limit: u64,
481    tracks: &[Track],
482    warnings: &mut Vec<String>,
483    frames: &mut Vec<(u64, u32)>,
484    mut each: impl FnMut(&BlockRef, &[(u64, u32)]),
485) -> Result<Cluster> {
486    let ds = h.data_start();
487    let known_end = match h.size {
488        Size::Known(_) => Some(h.end(limit)),
489        Size::Unknown => None,
490    };
491    let lim = known_end.unwrap_or(limit);
492    let mut ts: i64 = 0;
493    let mut pos = ds;
494    let end = loop {
495        if pos >= lim {
496            break lim;
497        }
498        let Some(c) = io.header(pos, lim)? else {
499            warn(warnings, || {
500                format!("damaged cluster at {}: undecodable element at {pos}", h.pos)
501            });
502            break io.resync(pos + 1, limit)?.unwrap_or(limit);
503        };
504        if known_end.is_none() && is_top_level(c.id) {
505            break pos;
506        }
507        if c.size == Size::Unknown || c.overruns(lim) {
508            warn(warnings, || {
509                format!(
510                    "damaged cluster at {}: element 0x{:X} at {pos} overruns",
511                    h.pos, c.id
512                )
513            });
514            if c.overruns(lim)
515                && matches!(c.id, SIMPLE_BLOCK | BLOCK_GROUP)
516                && io.len <= c.data_start() + c.size.known().unwrap_or(0)
517            {
518                break lim;
519            }
520            break io.resync(pos + 1, limit)?.unwrap_or(limit);
521        }
522        let ce = c.end(lim);
523        match c.id {
524            TIMESTAMP => ts = io.uint(&c, ce)? as i64,
525            SIMPLE_BLOCK | BLOCK_GROUP => {
526                if let Some(b) = parse_block_elem(io, &c, ce, ts, tracks, frames)? {
527                    each(&b, frames);
528                }
529            }
530            _ => {}
531        }
532        pos = ce;
533    };
534    Ok(Cluster {
535        offset: h.pos,
536        data_start: ds,
537        end,
538        timestamp: ts.max(0) as u64,
539    })
540}
541
542pub struct Reader<S: ByteSource = File> {
543    source: S,
544    len: u64,
545    pub doc_type: String,
546    pub info: Info,
547    pub tracks: Vec<Track>,
548    pub cues: Vec<CuePoint>,
549    pub clusters: Vec<Cluster>,
550    pub warnings: Vec<String>,
551    pub indexed: bool,
552    pub(crate) segment_start: u64,
553    pub(crate) segment_end: u64,
554    pub(crate) first_cluster: Option<u64>,
555    full: OnceLock<std::result::Result<Box<Reader<S>>, Error>>,
556    cue_keys: OnceLock<Option<Vec<i64>>>,
557}
558
559struct Head {
560    doc_type: String,
561    info: Info,
562    tracks: Vec<Track>,
563    cues: Vec<CuePoint>,
564    clusters: Vec<Cluster>,
565    warnings: Vec<String>,
566    segment_start: u64,
567    segment_end: u64,
568    first_cluster: Option<u64>,
569    seek_head: Vec<SeekEntry>,
570}
571
572impl Reader<File> {
573    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
574        let path = path.as_ref();
575        File::open(path)
576            .map_err(Error::from)
577            .and_then(Self::new)
578            .map_err(|e| e.in_file(path))
579    }
580
581    pub fn open_with(path: impl AsRef<Path>, options: Options) -> Result<Self> {
582        let path = path.as_ref();
583        File::open(path)
584            .map_err(Error::from)
585            .and_then(|f| Self::with_options(f, options))
586            .map_err(|e| e.in_file(path))
587    }
588}
589
590impl<S: ByteSource> Reader<S> {
591    pub fn new(source: S) -> Result<Self> {
592        Self::with_options(source, Options::default())
593    }
594
595    pub fn with_options(source: S, options: Options) -> Result<Self> {
596        Self::open_inner(source, options, false)
597    }
598
599    fn open_inner(source: S, options: Options, video_only: bool) -> Result<Self> {
600        let mut cache = Cache::default();
601        let mut io = Io::new(&source, &mut cache);
602        let len = io.len;
603        let start = {
604            let b = io.peek(0, WIN)?;
605            b.windows(4)
606                .position(|w| w == [0x1A, 0x45, 0xDF, 0xA3])
607                .ok_or_else(|| Error::Invalid("not a Matroska file: no EBML header".into()))?
608                as u64
609        };
610        let eh = io
611            .header(start, len)?
612            .ok_or_else(|| Error::Invalid("not a Matroska file: bad EBML header".into()))?;
613        let hdr_end = eh.end(len);
614        let hdr = io.read_vec(eh.data_start(), hdr_end - eh.data_start())?;
615        let mut doc_type = "matroska".to_string();
616        for e in Children::new(&hdr) {
617            match e.id {
618                DOC_TYPE => doc_type = ebml::string(e.data),
619                EBML_MAX_ID_LENGTH if ebml::uint(e.data) > 4 => {
620                    return Err(Error::Unsupported("EBMLMaxIDLength above 4".into()));
621                }
622                EBML_MAX_SIZE_LENGTH if ebml::uint(e.data) > 8 => {
623                    return Err(Error::Unsupported("EBMLMaxSizeLength above 8".into()));
624                }
625                _ => {}
626            }
627        }
628        if doc_type != "matroska" && doc_type != "webm" {
629            return Err(Error::Invalid(format!(
630                "not a Matroska file: DocType {doc_type:?}"
631            )));
632        }
633        let mut pos = hdr_end;
634        let seg = loop {
635            let h = io
636                .header(pos, len)?
637                .ok_or_else(|| Error::Invalid("no Matroska Segment".into()))?;
638            if h.id == SEGMENT {
639                break h;
640            }
641            if h.size == Size::Unknown {
642                return Err(Error::Invalid(
643                    "unknown-size element before the Segment".into(),
644                ));
645            }
646            pos = h.end(len);
647        };
648        let mut head = Head {
649            doc_type,
650            info: Info::default(),
651            tracks: Vec::new(),
652            cues: Vec::new(),
653            clusters: Vec::new(),
654            warnings: Vec::new(),
655            segment_start: seg.data_start(),
656            segment_end: seg.end(len),
657            first_cluster: None,
658            seek_head: Vec::new(),
659        };
660        if seg.overruns(len) {
661            head.warnings
662                .push("Segment extends past the end of the file".into());
663        }
664        let mut w = Walker {
665            index: options.index,
666            video_only,
667            parsed: HashSet::new(),
668            frames: Vec::new(),
669            room: len / 2,
670        };
671        w.walk(&mut io, &mut head)?;
672        for t in &mut head.tracks {
673            fill_durations(t);
674            t.finish_indexes();
675        }
676        Ok(Reader {
677            source,
678            len,
679            doc_type: head.doc_type,
680            info: head.info,
681            tracks: head.tracks,
682            cues: head.cues,
683            clusters: head.clusters,
684            warnings: head.warnings,
685            indexed: options.index,
686            segment_start: head.segment_start,
687            segment_end: head.segment_end,
688            first_cluster: head.first_cluster,
689            full: OnceLock::new(),
690            cue_keys: OnceLock::new(),
691        })
692    }
693
694    pub fn source(&self) -> &S {
695        &self.source
696    }
697
698    pub fn is_webm(&self) -> bool {
699        self.doc_type == "webm"
700    }
701
702    pub fn duration_seconds(&self) -> f64 {
703        if let Some(d) = self.info.duration.filter(|&d| d > 0.0) {
704            return d * self.info.timestamp_scale as f64 / 1e9;
705        }
706        self.tracks
707            .iter()
708            .map(Track::duration_seconds)
709            .fold(0.0, f64::max)
710    }
711
712    pub fn video(&self) -> Option<usize> {
713        let mut it = self
714            .tracks
715            .iter()
716            .enumerate()
717            .filter(|(_, t)| t.kind == TrackKind::Video);
718        let first = it.clone().next().map(|x| x.0);
719        it.find(|(_, t)| t.default).map(|x| x.0).or(first)
720    }
721
722    pub fn audio(&self) -> Vec<usize> {
723        self.tracks
724            .iter()
725            .enumerate()
726            .filter(|(_, t)| t.kind == TrackKind::Audio)
727            .map(|(i, _)| i)
728            .collect()
729    }
730
731    pub fn read(&self, track: usize, index: usize) -> Result<Vec<u8>> {
732        let mut data = Vec::new();
733        self.read_into(track, index, &mut data)?;
734        Ok(data)
735    }
736
737    pub fn read_into(&self, track: usize, index: usize, data: &mut Vec<u8>) -> Result<()> {
738        let t = self
739            .tracks
740            .get(track)
741            .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
742        let s = t
743            .samples
744            .get(index)
745            .ok_or_else(|| Error::Invalid(format!("no sample {index} in track {track}")))?;
746        read_frame(&self.source, self.len, t, s.offset, s.size, data)
747    }
748
749    pub fn packet(&self, track: usize, index: usize) -> Result<Packet> {
750        let data = self.read(track, index)?;
751        let s = self
752            .tracks
753            .get(track)
754            .and_then(|t| t.samples.get(index))
755            .copied()
756            .ok_or_else(|| Error::Invalid(format!("no sample {index} in track {track}")))?;
757        Ok(Packet {
758            data,
759            pts: s.pts,
760            dts: s.pts,
761            keyframe: s.keyframe,
762        })
763    }
764
765    pub(crate) fn len(&self) -> u64 {
766        self.len
767    }
768
769    pub(crate) fn cue_keys(&self) -> Option<&[i64]> {
770        self.cue_keys
771            .get_or_init(|| self.keys_from_cues())
772            .as_deref()
773    }
774
775    fn keys_from_cues(&self) -> Option<Vec<i64>> {
776        let v = self.video()?;
777        let t = self.tracks.get(v)?;
778        let mut cues = self
779            .cues
780            .iter()
781            .filter(|c| c.track == t.number)
782            .map(|c| {
783                Some((
784                    self.segment_start.checked_add(c.cluster)?,
785                    i64::try_from(c.time).ok()?,
786                ))
787            })
788            .collect::<Option<Vec<_>>>()?;
789        cues.sort_unstable();
790        cues.dedup();
791        let mut at: Vec<u64> = cues.iter().map(|c| c.0).collect();
792        at.dedup();
793        let last = at.len().checked_sub(1)?;
794        let first = self.first_cluster?;
795        let mut picks = vec![0, last / 2, last];
796        picks.dedup();
797        let mut cache = Cache::default();
798        let mut io = Io::with_len(&self.source, self.len, &mut cache);
799        for i in picks {
800            let from = if i == 0 { first.min(at[0]) } else { at[i] };
801            let (seen, mut found) = self.keys_in(&mut io, v, from, at.get(i + 1).copied())?;
802            let mut want: Vec<i64> = cues
803                .iter()
804                .filter(|c| seen.contains(&c.0))
805                .map(|c| c.1)
806                .collect();
807            want.sort_unstable();
808            want.dedup();
809            found.sort_unstable();
810            found.dedup();
811            if found != want {
812                return None;
813            }
814        }
815        let scale = self.info.timestamp_scale.max(1) as u128;
816        let delay = ((t.codec_delay as u128 + scale / 2) / scale).min(i64::MAX as u128) as i64;
817        let mut keys: Vec<i64> = cues.iter().map(|c| c.1.saturating_sub(delay)).collect();
818        keys.sort_unstable();
819        keys.dedup();
820        Some(keys)
821    }
822
823    fn keys_in(
824        &self,
825        io: &mut Io<'_, S>,
826        track: usize,
827        from: u64,
828        through: Option<u64>,
829    ) -> Option<(Vec<u64>, Vec<i64>)> {
830        let end = self.segment_end;
831        let (mut seen, mut keys) = (Vec::new(), Vec::new());
832        let (mut warnings, mut frames) = (Vec::new(), Vec::new());
833        let mut pos = from;
834        while pos < end {
835            if pos - from > CUE_SPAN {
836                return None;
837            }
838            let h = io.header(pos, end).ok()??;
839            if h.id != CLUSTER {
840                if through.is_some() || h.size == Size::Unknown || h.end(end) <= pos {
841                    return None;
842                }
843                pos = h.end(end);
844                continue;
845            }
846            let mut laced = false;
847            let ci = scan_cluster(
848                io,
849                &h,
850                end,
851                &self.tracks,
852                &mut warnings,
853                &mut frames,
854                |b, f| {
855                    if b.track == track && b.keyframe {
856                        laced |= f.len() != 1;
857                        keys.push(b.ts);
858                    }
859                },
860            )
861            .ok()?;
862            if laced || !warnings.is_empty() {
863                return None;
864            }
865            seen.push(h.pos);
866            match through {
867                Some(t) if h.pos == t => return Some((seen, keys)),
868                Some(t) if h.pos > t => return None,
869                _ => {}
870            }
871            if ci.end <= h.pos {
872                return None;
873            }
874            pos = ci.end;
875        }
876        through.is_none().then_some((seen, keys))
877    }
878}
879
880impl Reader<File> {
881    #[cfg_attr(
882        not(all(
883            feature = "h264",
884            feature = "hevc",
885            feature = "aac",
886            feature = "gpu",
887            feature = "filter",
888            feature = "vp9"
889        )),
890        allow(dead_code)
891    )]
892    pub(crate) fn full(&self) -> Result<&Reader<File>> {
893        if self.indexed {
894            return Ok(self);
895        }
896        let built = self.full.get_or_init(|| {
897            self.source
898                .try_clone()
899                .map_err(Error::from)
900                .and_then(|f| Self::open_inner(f, Options::default(), true))
901                .map(Box::new)
902        });
903        match built {
904            Ok(r) => Ok(r),
905            Err(e) => Err(again(e)),
906        }
907    }
908}
909
910fn again(e: &Error) -> Error {
911    match e {
912        Error::Unsupported(s) => Error::Unsupported(s.clone()),
913        Error::Invalid(s) => Error::Invalid(s.clone()),
914        Error::Backend(s) => Error::Backend(s.clone()),
915        Error::Io(e) => Error::Io(std::io::Error::new(e.kind(), e.to_string())),
916        Error::Cancelled => Error::Cancelled,
917    }
918}
919
920pub(crate) fn read_frame<S: ByteSource + ?Sized>(
921    src: &S,
922    len: u64,
923    t: &Track,
924    offset: u64,
925    size: u32,
926    data: &mut Vec<u8>,
927) -> Result<()> {
928    if !t.frames_readable() {
929        return Err(Error::Unsupported(format!(
930            "compressed or encrypted frames in track {}",
931            t.number
932        )));
933    }
934    if offset.checked_add(size as u64).is_none_or(|e| e > len) {
935        return Err(eof());
936    }
937    let prefix = t.stripped_header().unwrap_or(&[]);
938    data.clear();
939    data.reserve(prefix.len() + size as usize);
940    data.extend_from_slice(prefix);
941    data.resize(prefix.len() + size as usize, 0);
942    src.read_at(offset, &mut data[prefix.len()..])?;
943    Ok(())
944}
945
946pub(crate) fn read_frame_io<S: ByteSource + ?Sized>(
947    io: &mut Io<'_, S>,
948    t: &Track,
949    offset: u64,
950    size: u32,
951    data: &mut Vec<u8>,
952) -> Result<()> {
953    if !t.frames_readable() {
954        return Err(Error::Unsupported(format!(
955            "compressed or encrypted frames in track {}",
956            t.number
957        )));
958    }
959    data.clear();
960    data.extend_from_slice(t.stripped_header().unwrap_or(&[]));
961    io.read_into(offset, size as u64, data)
962}
963
964fn fill_durations(t: &mut Track) {
965    if t.samples.iter().all(|s| s.duration > 0) {
966        return;
967    }
968    let mut order: Vec<usize> = (0..t.samples.len()).collect();
969    order.sort_by_key(|&i| (t.samples[i].pts, i));
970    let mut last = 0u32;
971    for w in 0..order.len() {
972        let i = order[w];
973        if t.samples[i].duration > 0 {
974            last = t.samples[i].duration;
975            continue;
976        }
977        let d = match order.get(w + 1) {
978            Some(&j) => t.samples[j]
979                .pts
980                .saturating_sub(t.samples[i].pts)
981                .clamp(0, u32::MAX as i64) as u32,
982            None => last,
983        };
984        t.samples[i].duration = d;
985        if d > 0 {
986            last = d;
987        }
988    }
989}
990
991struct Walker {
992    index: bool,
993    video_only: bool,
994    parsed: HashSet<u64>,
995    frames: Vec<(u64, u32)>,
996    room: u64,
997}
998
999impl Walker {
1000    fn walk<S: ByteSource + ?Sized>(&mut self, io: &mut Io<'_, S>, r: &mut Head) -> Result<()> {
1001        let seg_end = r.segment_end;
1002        let mut pos = r.segment_start;
1003        while pos < seg_end {
1004            let Some(h) = io.header(pos, seg_end)? else {
1005                match io.resync(pos + 1, seg_end)? {
1006                    Some(p) => {
1007                        warn(&mut r.warnings, || {
1008                            format!("undecodable element at {pos}; resynced at {p}")
1009                        });
1010                        pos = p;
1011                        continue;
1012                    }
1013                    None => break,
1014                }
1015            };
1016            match h.id {
1017                CLUSTER => {
1018                    if r.first_cluster.is_none() {
1019                        r.first_cluster = Some(h.pos);
1020                        self.follow_seek_head(io, r, &[INFO, TRACKS])?;
1021                        if !self.index {
1022                            break;
1023                        }
1024                    }
1025                    let cluster = r.clusters.len() as u32;
1026                    let scale = r.info.timestamp_scale;
1027                    let mut samples: Vec<Vec<Sample>> = r
1028                        .tracks
1029                        .iter_mut()
1030                        .map(|t| std::mem::take(&mut t.samples))
1031                        .collect();
1032                    let tracks = &r.tracks;
1033                    let video_only = self.video_only;
1034                    let room = &mut self.room;
1035                    let mut full = false;
1036                    let ci = scan_cluster(
1037                        io,
1038                        &h,
1039                        seg_end,
1040                        tracks,
1041                        &mut r.warnings,
1042                        &mut self.frames,
1043                        |b, frames| {
1044                            if full || video_only && tracks[b.track].kind != TrackKind::Video {
1045                                return;
1046                            }
1047                            expand(b, frames, &tracks[b.track], scale, 0, |f| {
1048                                let Some(list) = samples.get_mut(f.track) else {
1049                                    return;
1050                                };
1051                                if *room == 0 || list.len() >= MAX_SAMPLES {
1052                                    full = true;
1053                                    return;
1054                                }
1055                                *room -= 1;
1056                                list.push(Sample {
1057                                    offset: f.offset,
1058                                    size: f.size,
1059                                    pts: f.pts,
1060                                    duration: f.duration,
1061                                    keyframe: f.keyframe,
1062                                    cluster,
1063                                    block: f.block,
1064                                    lace: f.lace,
1065                                })
1066                            })
1067                        },
1068                    )?;
1069                    for (t, s) in r.tracks.iter_mut().zip(samples) {
1070                        t.samples = s;
1071                    }
1072                    if full {
1073                        return Err(Error::Invalid(
1074                            "more frames than a Matroska file of this size can index".into(),
1075                        ));
1076                    }
1077                    r.clusters.push(ci);
1078                    pos = ci.end.max(h.data_start());
1079                }
1080                SEEK_HEAD | INFO | TRACKS | CUES | CHAPTERS | TAGS | ATTACHMENTS => {
1081                    pos = self.level1(io, r, &h, seg_end)?;
1082                }
1083                EBML | SEGMENT => break,
1084                _ => {
1085                    if h.size == Size::Unknown {
1086                        pos = io.unknown_end(h.data_start(), seg_end)?;
1087                    } else if h.overruns(seg_end) && h.id != VOID {
1088                        match io.resync(pos + 1, seg_end)? {
1089                            Some(p) => {
1090                                warn(&mut r.warnings, || {
1091                                    format!(
1092                                        "oversized element 0x{:X} at {pos}; resynced at {p}",
1093                                        h.id
1094                                    )
1095                                });
1096                                pos = p;
1097                            }
1098                            None => break,
1099                        }
1100                    } else {
1101                        pos = h.end(seg_end);
1102                    }
1103                }
1104            }
1105        }
1106        self.follow_seek_head(io, r, &[INFO, TRACKS, CUES, SEEK_HEAD])?;
1107        self.follow_seek_head(io, r, &[CUES])?;
1108        r.cues.sort_by_key(|c| c.time);
1109        Ok(())
1110    }
1111
1112    fn follow_seek_head<S: ByteSource + ?Sized>(
1113        &mut self,
1114        io: &mut Io<'_, S>,
1115        r: &mut Head,
1116        ids: &[u32],
1117    ) -> Result<()> {
1118        let entries = r.seek_head.clone();
1119        for e in entries {
1120            if !ids.contains(&e.id) {
1121                continue;
1122            }
1123            let Some(abs) = r.segment_start.checked_add(e.position) else {
1124                continue;
1125            };
1126            if self.parsed.contains(&abs) || abs >= r.segment_end {
1127                continue;
1128            }
1129            match io.header(abs, r.segment_end)? {
1130                Some(h) if h.id == e.id => {
1131                    self.level1(io, r, &h, r.segment_end)?;
1132                }
1133                _ => warn(&mut r.warnings, || {
1134                    format!(
1135                        "SeekHead entry 0x{:X} -> {abs} does not point at that element",
1136                        e.id
1137                    )
1138                }),
1139            }
1140        }
1141        Ok(())
1142    }
1143
1144    fn level1<S: ByteSource + ?Sized>(
1145        &mut self,
1146        io: &mut Io<'_, S>,
1147        r: &mut Head,
1148        h: &Hdr,
1149        limit: u64,
1150    ) -> Result<u64> {
1151        let end = if h.size == Size::Unknown {
1152            io.unknown_end(h.data_start(), limit)?
1153        } else {
1154            h.end(limit)
1155        };
1156        if h.overruns(limit) {
1157            warn(&mut r.warnings, || {
1158                format!("element 0x{:X} at {} truncated", h.id, h.pos)
1159            });
1160        }
1161        if !self.parsed.insert(h.pos) || !matches!(h.id, SEEK_HEAD | INFO | TRACKS | CUES) {
1162            return Ok(end);
1163        }
1164        let n = end - h.data_start();
1165        if n > MAX_META {
1166            warn(&mut r.warnings, || {
1167                format!(
1168                    "element 0x{:X} at {} too large ({n} bytes); skipped",
1169                    h.id, h.pos
1170                )
1171            });
1172            return Ok(end);
1173        }
1174        let data = io.read_vec(h.data_start(), n)?;
1175        match h.id {
1176            SEEK_HEAD => r.seek_head.extend(parse::seek_head(&data)),
1177            INFO => r.info = parse::info(&data),
1178            TRACKS if r.tracks.is_empty() => {
1179                r.tracks = parse::tracks(&data, r.info.timestamp_scale);
1180            }
1181            CUES => parse::cues(&data, &mut r.cues),
1182            _ => {}
1183        }
1184        Ok(end)
1185    }
1186}