1 files changed, 24 insertions, 15 deletions
diff --git a/ebml/src/read.rs b/ebml/src/read.rs
index 7431683..23b4e21 100644
--- a/ebml/src/read.rs
+++ b/ebml/src/read.rs
@@ -1,5 +1,7 @@
+use anyhow::{bail, Result};
 use std::io::{Read, Seek};
-use anyhow::Result;
+
+use crate::matroska::MatroskaTag;
 
 trait ReadAndSeek: Read + Seek {}
 impl<T: Read + Seek> ReadAndSeek for T {}
@@ -37,9 +39,12 @@ impl EbmlReader {
     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 0..(len - 1) {
+        for _ in 1..len {
             value <<= 8;
             value += self.read_byte()? as u64;
         }
@@ -52,11 +57,23 @@ impl EbmlReader {
         let b = self.read_buf(size)?;
         Ok(String::from_utf8(b)?)
     }
-
     pub fn read_tag_id(&mut self) -> Result<u64> {
         let (value, len) = self.read_vint_len()?;
         Ok(value + (1 << (7 * len)))
     }
+    pub fn read_tag_size(&mut self) -> Result<EbmlSize> {
+        Ok(EbmlSize::from_vint(self.read_vint_len()?))
+    }
+    pub fn read_tag(&mut self) -> Result<MatroskaTag> {
+        let id = self.read_tag_id()?;
+        let size = self.read_tag_size()?;
+        if MatroskaTag::is_master(id)? {
+            Ok(MatroskaTag::parse(id, &[])?)
+        } else {
+            let data = self.read_buf(size)?;
+            Ok(MatroskaTag::parse(id, &data)?)
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -65,20 +82,12 @@ pub enum EbmlSize {
     Unknown,
 }
 impl EbmlSize {
-    #[rustfmt::skip]
     pub fn from_vint((value, len): (u64, usize)) -> EbmlSize {
-        match len {
-            1 => if value == ((1 << (7))     - 1) { return Self::Unknown; },
-            2 => if value == ((1 << (7 * 2)) - 1) { return Self::Unknown; },
-            3 => if value == ((1 << (7 * 3)) - 1) { return Self::Unknown; },
-            4 => if value == ((1 << (7 * 4)) - 1) { return Self::Unknown; },
-            5 => if value == ((1 << (7 * 5)) - 1) { return Self::Unknown; },
-            6 => if value == ((1 << (7 * 6)) - 1) { return Self::Unknown; },
-            7 => if value == ((1 << (7 * 7)) - 1) { return Self::Unknown; },
-            8 => if value == ((1 << (7 * 8)) - 1) { return Self::Unknown; },
-            _ => {},
+        if value == ((1 << (7 * len)) - 1) {
+            Self::Unknown
+        } else {
+            Self::Exact(value as usize)
         }
-        Self::Exact(len)
     }
 }
 impl Into<usize> for EbmlSize {