1use std::{
2 cmp::{max, min},
3 fmt::{Debug, Display},
4 fs::File,
5 io::{self, Read, Seek},
6 marker::PhantomData,
7 net::TcpStream,
8 ops::Range,
9 path::Path,
10};
11
12use crate::{
13 buffer::Buffers,
14 error::{ParsedError, ParsingError, ParsingErrorState},
15 exif::{parse_exif_iter, TiffHeader},
16 file::Mime,
17 partial_vec::PartialVec,
18 skip::Skip,
19 video::parse_track_info,
20 ExifIter, Seekable, TrackInfo, Unseekable,
21};
22
23pub struct MediaSource<R, S = Seekable> {
45 pub(crate) reader: R,
46 pub(crate) buf: Vec<u8>,
47 pub(crate) mime: Mime,
48 phantom: PhantomData<S>,
49}
50
51impl<R, S: Skip<R>> Debug for MediaSource<R, S> {
52 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53 f.debug_struct("MediaSource")
54 .field("mime", &self.mime)
56 .field("seekable", &S::debug())
57 .finish_non_exhaustive()
58 }
59}
60
61const HEADER_PARSE_BUF_SIZE: usize = 128;
63
64impl<R: Read, S: Skip<R>> MediaSource<R, S> {
65 #[tracing::instrument(skip(reader))]
66 fn build(mut reader: R) -> crate::Result<Self> {
67 let mut buf = Vec::with_capacity(HEADER_PARSE_BUF_SIZE);
69 reader
70 .by_ref()
71 .take(HEADER_PARSE_BUF_SIZE as u64)
72 .read_to_end(&mut buf)?;
73 let mime: Mime = buf.as_slice().try_into()?;
74 tracing::debug!(?mime);
75 Ok(Self {
76 reader,
77 buf,
78 mime,
79 phantom: PhantomData,
80 })
81 }
82
83 pub fn has_track(&self) -> bool {
84 match self.mime {
85 Mime::Image(_) => false,
86 Mime::Video(_) => true,
87 }
88 }
89
90 pub fn has_exif(&self) -> bool {
91 match self.mime {
92 Mime::Image(_) => true,
93 Mime::Video(_) => false,
94 }
95 }
96}
97
98impl<R: Read + Seek> MediaSource<R, Seekable> {
99 pub fn seekable(reader: R) -> crate::Result<Self> {
104 Self::build(reader)
105 }
106}
107
108impl<R: Read> MediaSource<R, Unseekable> {
109 pub fn unseekable(reader: R) -> crate::Result<Self> {
115 Self::build(reader)
116 }
117}
118
119impl MediaSource<File, Seekable> {
120 pub fn file_path<P: AsRef<Path>>(path: P) -> crate::Result<Self> {
121 Self::seekable(File::open(path)?)
122 }
123
124 pub fn file(file: File) -> crate::Result<Self> {
125 Self::seekable(file)
126 }
127}
128
129impl MediaSource<TcpStream, Unseekable> {
130 pub fn tcp_stream(stream: TcpStream) -> crate::Result<Self> {
131 Self::unseekable(stream)
132 }
133}
134
135pub(crate) const INIT_BUF_SIZE: usize = 4096;
137pub(crate) const MIN_GROW_SIZE: usize = 4096;
138pub(crate) const MAX_ALLOC_SIZE: usize = 1024 * 1024 * 1024;
142
143pub(crate) trait Buf {
144 fn buffer(&self) -> &[u8];
145 fn clear(&mut self);
146
147 fn set_position(&mut self, pos: usize);
148 #[allow(unused)]
149 fn position(&self) -> usize;
150}
151
152#[derive(Debug, Clone)]
153pub(crate) enum ParsingState {
154 TiffHeader(TiffHeader),
155 HeifExifSize(usize),
156 Cr3ExifSize(usize),
157}
158
159impl Display for ParsingState {
160 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161 match self {
162 ParsingState::TiffHeader(h) => Display::fmt(&format!("ParsingState: {h:?})"), f),
163 ParsingState::HeifExifSize(n) => Display::fmt(&format!("ParsingState: {n}"), f),
164 ParsingState::Cr3ExifSize(n) => Display::fmt(&format!("ParsingState: {n}"), f),
165 }
166 }
167}
168
169pub(crate) trait BufParser: Buf + Debug {
170 fn fill_buf<R: Read>(&mut self, reader: &mut R, size: usize) -> io::Result<usize>;
171 fn load_and_parse<R: Read, S: Skip<R>, P, O>(
172 &mut self,
173 reader: &mut R,
174 mut parse: P,
175 ) -> Result<O, ParsedError>
176 where
177 P: FnMut(&[u8], Option<ParsingState>) -> Result<O, ParsingErrorState>,
178 {
179 self.load_and_parse_with_offset::<R, S, _, _>(
180 reader,
181 |data, _, state| parse(data, state),
182 0,
183 )
184 }
185
186 #[tracing::instrument(skip_all)]
187 fn load_and_parse_with_offset<R: Read, S: Skip<R>, P, O>(
188 &mut self,
189 reader: &mut R,
190 mut parse: P,
191 offset: usize,
192 ) -> Result<O, ParsedError>
193 where
194 P: FnMut(&[u8], usize, Option<ParsingState>) -> Result<O, ParsingErrorState>,
195 {
196 if offset >= self.buffer().len() {
197 self.fill_buf(reader, MIN_GROW_SIZE)?;
198 }
199
200 let mut parsing_state: Option<ParsingState> = None;
201 loop {
202 let res = parse(self.buffer(), offset, parsing_state.take());
203 match res {
204 Ok(o) => return Ok(o),
205 Err(es) => {
206 tracing::debug!(?es);
207 parsing_state = es.state;
208
209 match es.err {
210 ParsingError::ClearAndSkip(n) => {
211 self.clear_and_skip::<R, S>(reader, n)?;
212 }
213 ParsingError::Need(i) => {
214 tracing::debug!(need = i, "need more bytes");
215 let to_read = max(i, MIN_GROW_SIZE);
216 let n = self.fill_buf(reader, to_read)?;
219 if n == 0 {
220 return Err(ParsedError::NoEnoughBytes);
221 }
222 tracing::debug!(n, "actual read");
223 }
224 ParsingError::Failed(s) => return Err(ParsedError::Failed(s)),
225 }
226 }
227 }
228 }
229 }
230
231 #[tracing::instrument(skip(reader))]
232 fn clear_and_skip<R: Read, S: Skip<R>>(
233 &mut self,
234 reader: &mut R,
235 n: usize,
236 ) -> Result<(), ParsedError> {
237 tracing::debug!("ClearAndSkip");
238 if n <= self.buffer().len() {
239 tracing::debug!(n, "skip by set_position");
240 self.set_position(n);
241 return Ok(());
242 }
243
244 let skip_n = n - self.buffer().len();
245 tracing::debug!(skip_n, "clear and skip bytes");
246 self.clear();
247
248 let done = S::skip_by_seek(
249 reader,
250 skip_n
251 .try_into()
252 .map_err(|_| ParsedError::Failed("skip too many bytes".into()))?,
253 )?;
254 if !done {
255 tracing::debug!(skip_n, "skip by using our buffer");
256 let mut skipped = 0;
257 while skipped < skip_n {
258 let mut to_skip = skip_n - skipped;
259 to_skip = min(to_skip, MAX_ALLOC_SIZE);
260 let n = self.fill_buf(reader, to_skip)?;
261 skipped += n;
262 if skipped <= skip_n {
263 self.clear();
264 } else {
265 let remain = skipped - skip_n;
266 self.set_position(self.buffer().len() - remain);
267 break;
268 }
269 }
270 } else {
271 tracing::debug!(skip_n, "skip with seek");
272 }
273
274 if self.buffer().is_empty() {
275 self.fill_buf(reader, MIN_GROW_SIZE)?;
276 }
277 Ok(())
278 }
279}
280
281impl BufParser for MediaParser {
282 #[tracing::instrument(skip(self, reader), fields(buf_len=self.buf().len()))]
283 fn fill_buf<R: Read>(&mut self, reader: &mut R, size: usize) -> io::Result<usize> {
284 if size.saturating_add(self.buf().len()) > MAX_ALLOC_SIZE {
285 tracing::error!(?size, "the requested buffer size is too big");
286 return Err(io::ErrorKind::Unsupported.into());
287 }
288 self.buf_mut().reserve_exact(size);
289
290 let n = reader.take(size as u64).read_to_end(self.buf_mut())?;
291 if n == 0 {
292 tracing::error!(buf_len = self.buf().len(), "fill_buf: EOF");
293 return Err(std::io::ErrorKind::UnexpectedEof.into());
294 }
295
296 tracing::debug!(
297 ?size,
298 ?n,
299 buf_len = self.buf().len(),
300 "fill_buf: read bytes"
301 );
302
303 Ok(n)
304 }
305}
306
307impl Buf for MediaParser {
308 fn buffer(&self) -> &[u8] {
309 &self.buf()[self.position..]
310 }
311
312 fn clear(&mut self) {
313 self.buf_mut().clear();
314 }
315
316 fn set_position(&mut self, pos: usize) {
317 self.position = pos;
318 }
319
320 fn position(&self) -> usize {
321 self.position
322 }
323}
324
325pub trait ParseOutput<R, S>: Sized {
326 fn parse(parser: &mut MediaParser, ms: MediaSource<R, S>) -> crate::Result<Self>;
327}
328
329impl<R: Read, S: Skip<R>> ParseOutput<R, S> for ExifIter {
330 fn parse(parser: &mut MediaParser, mut ms: MediaSource<R, S>) -> crate::Result<Self> {
331 if !ms.has_exif() {
332 return Err(crate::Error::ParseFailed("no Exif data here".into()));
333 }
334 parse_exif_iter::<R, S>(parser, ms.mime.unwrap_image(), &mut ms.reader)
335 }
336}
337
338impl<R: Read, S: Skip<R>> ParseOutput<R, S> for TrackInfo {
339 fn parse(parser: &mut MediaParser, mut ms: MediaSource<R, S>) -> crate::Result<Self> {
340 if !ms.has_track() {
341 return Err(crate::Error::ParseFailed("no track info here".into()));
342 }
343 let out = parser.load_and_parse::<R, S, _, _>(ms.reader.by_ref(), |data, _| {
344 parse_track_info(data, ms.mime.unwrap_video())
345 .map_err(|e| ParsingErrorState::new(e, None))
346 })?;
347 Ok(out)
348 }
349}
350
351pub struct MediaParser {
402 bb: Buffers,
403 buf: Option<Vec<u8>>,
404 position: usize,
405}
406
407impl Debug for MediaParser {
408 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
409 f.debug_struct("MediaParser")
410 .field("buffers", &self.bb)
411 .field("buf len", &self.buf.as_ref().map(|x| x.len()))
412 .field("position", &self.position)
413 .finish_non_exhaustive()
414 }
415}
416
417impl Default for MediaParser {
418 fn default() -> Self {
419 Self {
420 bb: Buffers::new(),
421 buf: None,
422 position: 0,
423 }
424 }
425}
426
427pub(crate) trait ShareBuf {
428 fn share_buf(&mut self, range: Range<usize>) -> PartialVec;
429}
430
431impl ShareBuf for MediaParser {
432 fn share_buf(&mut self, mut range: Range<usize>) -> PartialVec {
433 let buf = self.buf.take().unwrap();
434 let vec = self.bb.release_to_share(buf);
435 range.start += self.position;
436 range.end += self.position;
437 PartialVec::new(vec, range)
438 }
439}
440
441impl MediaParser {
442 pub fn new() -> Self {
443 Self::default()
444 }
445
446 pub fn parse<R: Read, S, O: ParseOutput<R, S>>(
466 &mut self,
467 mut ms: MediaSource<R, S>,
468 ) -> crate::Result<O> {
469 self.reset();
470 self.acquire_buf();
471
472 self.buf_mut().append(&mut ms.buf);
473 let res = self.do_parse(ms);
474
475 self.reset();
476 res
477 }
478
479 fn do_parse<R: Read, S, O: ParseOutput<R, S>>(
480 &mut self,
481 mut ms: MediaSource<R, S>,
482 ) -> Result<O, crate::Error> {
483 self.fill_buf(&mut ms.reader, INIT_BUF_SIZE)?;
484 let res = ParseOutput::parse(self, ms)?;
485 Ok(res)
486 }
487
488 fn reset(&mut self) {
489 if let Some(buf) = self.buf.take() {
491 self.bb.release(buf);
492 }
493
494 self.set_position(0);
496 }
497
498 pub(crate) fn buf(&self) -> &Vec<u8> {
499 match self.buf.as_ref() {
500 Some(b) => b,
501 None => panic!("no buf here"),
502 }
503 }
504
505 fn buf_mut(&mut self) -> &mut Vec<u8> {
506 match self.buf.as_mut() {
507 Some(b) => b,
508 None => panic!("no buf here"),
509 }
510 }
511
512 fn acquire_buf(&mut self) {
513 assert!(self.buf.is_none());
514 self.buf = Some(self.bb.acquire());
515 }
516}
517
518#[cfg(test)]
519mod tests {
520 use std::sync::{LazyLock, Mutex, MutexGuard};
521
522 use super::*;
523 use test_case::case;
524
525 enum TrackExif {
526 Track,
527 Exif,
528 NoData,
529 Invalid,
530 }
531 use TrackExif::*;
532
533 static PARSER: LazyLock<Mutex<MediaParser>> = LazyLock::new(|| Mutex::new(MediaParser::new()));
534 fn parser() -> MutexGuard<'static, MediaParser> {
535 PARSER.lock().unwrap()
536 }
537
538 #[case("3gp_640x360.3gp", Track)]
539 #[case("broken.jpg", Exif)]
540 #[case("compatible-brands-fail.heic", Invalid)]
541 #[case("compatible-brands-fail.mov", Invalid)]
542 #[case("compatible-brands.heic", NoData)]
543 #[case("compatible-brands.mov", NoData)]
544 #[case("embedded-in-heic.mov", Track)]
545 #[case("exif.heic", Exif)]
546 #[case("exif.jpg", Exif)]
547 #[case("exif-no-tz.jpg", Exif)]
548 #[case("fujifilm_x_t1_01.raf.meta", Exif)]
549 #[case("meta.mov", Track)]
550 #[case("meta.mp4", Track)]
551 #[case("mka.mka", Track)]
552 #[case("mkv_640x360.mkv", Track)]
553 #[case("exif-one-entry.heic", Exif)]
554 #[case("no-exif.jpg", NoData)]
555 #[case("tif.tif", Exif)]
556 #[case("ramdisk.img", Invalid)]
557 #[case("webm_480.webm", Track)]
558 fn parse_media(path: &str, te: TrackExif) {
559 let mut parser = parser();
560 let ms = MediaSource::file_path(Path::new("testdata").join(path));
561 match te {
562 Track => {
563 let ms = ms.unwrap();
564 assert!(ms.has_track());
566 let _: TrackInfo = parser.parse(ms).unwrap();
567 }
568 Exif => {
569 let ms = ms.unwrap();
570 assert!(ms.has_exif());
572 let mut it: ExifIter = parser.parse(ms).unwrap();
573 let _ = it.parse_gps_info();
574
575 if path.contains("one-entry") {
576 assert!(it.next().is_some());
577 assert!(it.next().is_none());
578
579 let exif: crate::Exif = it.clone_and_rewind().into();
580 assert!(exif.get(ExifTag::Orientation).is_some());
581 } else {
582 let _: crate::Exif = it.clone_and_rewind().into();
583 }
584 }
585 NoData => {
586 let ms = ms.unwrap();
587 if ms.has_exif() {
589 let res: Result<ExifIter, _> = parser.parse(ms);
590 res.unwrap_err();
591 } else if ms.has_track() {
592 let res: Result<TrackInfo, _> = parser.parse(ms);
593 res.unwrap_err();
594 }
595 }
596 Invalid => {
597 ms.unwrap_err();
598 }
599 }
600 }
601
602 use crate::testkit::open_sample;
603 use crate::{EntryValue, Exif, ExifTag, TrackInfoTag};
604 use chrono::{DateTime, FixedOffset, NaiveDateTime};
605 use test_case::test_case;
606
607 #[test_case("exif.jpg", ExifTag::DateTimeOriginal, DateTime::parse_from_str("2023-07-09T20:36:33+08:00", "%+").unwrap().into())]
608 #[test_case("exif.heic", ExifTag::DateTimeOriginal, DateTime::parse_from_str("2022-07-22T21:26:32+08:00", "%+").unwrap().into())]
609 #[test_case("exif.jpg", ExifTag::DateTimeOriginal,
610 (NaiveDateTime::parse_from_str("2023-07-09T20:36:33", "%Y-%m-%dT%H:%M:%S").unwrap(),
611 Some(FixedOffset::east_opt(8*3600).unwrap())).into())]
612 #[test_case("exif-no-tz.jpg", ExifTag::DateTimeOriginal,
613 (NaiveDateTime::parse_from_str("2023-07-09T20:36:33", "%Y-%m-%dT%H:%M:%S").unwrap(), None).into())]
614 fn parse_exif(path: &str, tag: ExifTag, v: EntryValue) {
615 let mut parser = parser();
616
617 let mf = MediaSource::seekable(open_sample(path).unwrap()).unwrap();
618 assert!(mf.has_exif());
619 let iter: ExifIter = parser.parse(mf).unwrap();
620 let exif: Exif = iter.into();
621 assert_eq!(exif.get(tag).unwrap(), &v);
622
623 let mf = MediaSource::unseekable(open_sample(path).unwrap()).unwrap();
624 assert!(mf.has_exif());
625 let iter: ExifIter = parser.parse(mf).unwrap();
626 let exif: Exif = iter.into();
627 assert_eq!(exif.get(tag).unwrap(), &v);
628 }
629
630 use crate::video::TrackInfoTag::*;
631
632 #[test_case("mkv_640x360.mkv", ImageWidth, 640_u32.into())]
633 #[test_case("mkv_640x360.mkv", ImageHeight, 360_u32.into())]
634 #[test_case("mkv_640x360.mkv", DurationMs, 13346_u64.into())]
635 #[test_case("mkv_640x360.mkv", CreateDate, DateTime::parse_from_str("2008-08-08T08:08:08Z", "%+").unwrap().into())]
636 #[test_case("meta.mov", Make, "Apple".into())]
637 #[test_case("meta.mov", Model, "iPhone X".into())]
638 #[test_case("meta.mov", GpsIso6709, "+27.1281+100.2508+000.000/".into())]
639 #[test_case("meta.mov", CreateDate, DateTime::parse_from_str("2019-02-12T15:27:12+08:00", "%+").unwrap().into())]
640 #[test_case("meta.mp4", ImageWidth, 1920_u32.into())]
641 #[test_case("meta.mp4", ImageHeight, 1080_u32.into())]
642 #[test_case("meta.mp4", DurationMs, 1063_u64.into())]
643 #[test_case("meta.mp4", GpsIso6709, "+27.2939+112.6932/".into())]
644 #[test_case("meta.mp4", CreateDate, DateTime::parse_from_str("2024-02-03T07:05:38Z", "%+").unwrap().into())]
645 #[test_case("udta.auth.mp4", Author, "ReplayKitRecording".into(); "udta author")]
646 #[test_case("auth.mov", Author, "ReplayKitRecording".into(); "mov author")]
647 fn parse_track_info(path: &str, tag: TrackInfoTag, v: EntryValue) {
648 let mut parser = parser();
649
650 let mf = MediaSource::seekable(open_sample(path).unwrap()).unwrap();
651 let info: TrackInfo = parser.parse(mf).unwrap();
652 assert_eq!(info.get(tag).unwrap(), &v);
653
654 let mf = MediaSource::unseekable(open_sample(path).unwrap()).unwrap();
655 let info: TrackInfo = parser.parse(mf).unwrap();
656 assert_eq!(info.get(tag).unwrap(), &v);
657 }
658
659 #[test_case("crash_moov-trak")]
660 #[test_case("crash_skip_large")]
661 #[test_case("crash_add_large")]
662 fn parse_track_crash(path: &str) {
663 let mut parser = parser();
664
665 let mf = MediaSource::file(open_sample(path).unwrap()).unwrap();
666 let _: TrackInfo = parser.parse(mf).unwrap_or_default();
667
668 let mf = MediaSource::unseekable(open_sample(path).unwrap()).unwrap();
669 let _: TrackInfo = parser.parse(mf).unwrap_or_default();
670 }
671}