1 files changed, 55 insertions, 2 deletions

diff --git a/matroska/src/read.rs b/matroska/src/read.rs
index 95a98b5..6849f11 100644
--- a/matroska/src/read.rs
+++ b/matroska/src/read.rs
@@ -1,6 +1,6 @@
 use crate::{matroska::MatroskaTag, size::EbmlSize, Master};
 use anyhow::{anyhow, bail, Result};
-use log::{debug, warn};
+use log::{debug, error, warn};
 use std::{
     collections::VecDeque,
     io::{Read, Seek, SeekFrom},
@@ -22,6 +22,14 @@ pub struct EbmlReader {
     pub position: usize,
 }
 
+impl Read for EbmlReader {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let r = self.inner.read(buf)?;
+        self.position += r;
+        Ok(r)
+    }
+}
+
 impl EbmlReader {
     pub fn new<T: Seek + Read + 'static>(inner: T) -> Self {
         Self {
@@ -159,7 +167,22 @@ impl Iterator for EbmlReader {
         } else {
             match self.read_stuff() {
                 Ok(()) => self.next(),
-                Err(e) => Some(Err(e)),
+                // in case we reached the end (error: failed to fill whole buffer),
+                // return the rest in the queue and pop all items of the stack
+                Err(e) => {
+                    // TODO this is horrible, should use a custom error enum instead
+                    if format!("{e}").as_str() == "failed to fill whole buffer" {
+                        match self.queue.pop_front() {
+                            Some(q) => Some(Ok(q)),
+                            None => match self.stack.pop() {
+                                Some(q) => Some(MatroskaTag::construct_master(q.id, Master::End)),
+                                None => Some(Err(e)),
+                            },
+                        }
+                    } else {
+                        Some(Err(e))
+                    }
+                }
             }
         }
     }
@@ -225,3 +248,33 @@ impl ReadValue for Master {
         panic!("master shall not be read like this")
     }
 }
+
+pub trait ReadExt: Read {
+    fn read_byte(&mut self) -> Result<u8>;
+    fn read_vint_len(&mut self) -> Result<(u64, usize)>;
+    fn read_vint(&mut self) -> Result<u64>;
+}
+impl<T: Read> ReadExt for T {
+    fn read_byte(&mut self) -> Result<u8> {
+        let mut b = [0u8];
+        self.read_exact(&mut b)?;
+        Ok(b[0])
+    }
+    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))
+    }
+    fn read_vint(&mut self) -> Result<u64> {
+        Ok(self.read_vint_len()?.0)
+    }
+}