Add mkv duration extractor (#140)

2 files changed, 64 insertions, 1 deletions

diff --git a/src/util/duration/mkv.rs b/src/util/duration/mkv.rs
new file mode 100644
index 0000000..297f6db
--- /dev/null
+++ b/src/util/duration/mkv.rs
@@ -0,0 +1,60 @@
+use mp3_metadata::MP3Metadata;
+use std::fs::File;
+use std::io;
+use std::path::Path;
+
+use crate::util::duration::DurationExtractor;
+use crate::util::Duration;
+
+use matroska::MatroskaError;
+
+pub struct MkvDurationExtractor;
+
+impl DurationExtractor for MkvDurationExtractor {
+    fn supports_ext(&self, ext_lowercase: &str) -> bool {
+        "mkv" == ext_lowercase || "webm" == ext_lowercase
+    }
+
+    fn try_read_duration(
+        &self,
+        path: &Path,
+        _: &Option<MP3Metadata>,
+    ) -> io::Result<Option<Duration>> {
+        let fd = File::open(path)?;
+        let matroska = matroska::Matroska::open(fd).map_err(|err| match err {
+            MatroskaError::Io(io) => io,
+            MatroskaError::UTF8(utf8) => io::Error::new(io::ErrorKind::InvalidData, utf8),
+            e => io::Error::new(io::ErrorKind::InvalidData, e),
+        })?;
+
+        match matroska.info.duration {
+            Some(duration) => {
+                return Ok(Some(Duration {
+                    length: duration.as_secs() as usize,
+                }))
+            }
+            None => return Ok(None),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::MkvDurationExtractor;
+    use crate::util::duration::DurationExtractor;
+    use crate::util::Duration;
+    use std::error::Error;
+    use std::path::PathBuf;
+
+    #[test]
+    fn test_success() -> Result<(), Box<dyn Error>> {
+        let path_string =
+            std::env::var("CARGO_MANIFEST_DIR")? + "/resources/test/" + "video/rust-logo-blk.mkv";
+        let path = PathBuf::from(path_string);
+        assert_eq!(
+            MkvDurationExtractor.try_read_duration(&path, &None)?,
+            Some(Duration { length: 1 }),
+        );
+        Ok(())
+    }
+}
diff --git a/src/util/duration/mod.rs b/src/util/duration/mod.rs
index 0a4e7a3..8353e49 100644
--- a/src/util/duration/mod.rs
+++ b/src/util/duration/mod.rs
@@ -1,5 +1,6 @@
 use std::io;
 
+mod mkv;
 mod mp3;
 mod mp4;
 
@@ -7,6 +8,7 @@ use std::path::Path;
 
 use mp3_metadata::MP3Metadata;
 
+use mkv::MkvDurationExtractor;
 use mp3::Mp3DurationExtractor;
 use mp4::Mp4DurationExtractor;
 
@@ -20,9 +22,10 @@ pub trait DurationExtractor {
     fn try_read_duration(&self, path: &Path, mp3_metadata: &Option<MP3Metadata>) -> io::Result<Option<Duration>>;
 }
 
-const EXTRACTORS: [&dyn DurationExtractor; 2] = [
+const EXTRACTORS: [&dyn DurationExtractor; 3] = [
     &Mp3DurationExtractor,
     &Mp4DurationExtractor,
+    &MkvDurationExtractor,
 ];
 
 pub fn get_duration<T: AsRef<Path>>(path: T, mp3_metadata: &Option<MP3Metadata>) -> Option<Duration> {