1 files changed, 21 insertions, 8 deletions
diff --git a/src/unityfs.rs b/src/unityfs.rs
index d106ce6..f9a8d34 100644
--- a/src/unityfs.rs
+++ b/src/unityfs.rs
@@ -1,4 +1,4 @@
-use crate::helper::ReadExt;
+use crate::helper::{AlignExt, ReadExt};
 use anyhow::{Result, anyhow, bail};
 use log::{debug, info, trace};
 use std::io::{Cursor, ErrorKind, Read, Seek, SeekFrom, Take};
@@ -24,6 +24,7 @@ struct BlockInfo {
 
 pub struct NodeReader<'a, T> {
     inner: Take<&'a mut BlocksReader<T>>,
+    size: u64,
 }
 
 impl<T: Read + Seek> UnityFS<T> {
@@ -77,13 +78,8 @@ impl<T: Read + Seek> UnityFS<T> {
                 meta_comp_scheme.decompress(blockindex, blockindex_decomp_size as usize)?;
             Cursor::new(blockindex)
         };
-        {
-            // align stream
-            let off = file.stream_position()? % 16;
-            if off != 0 {
-                file.seek_relative(16 - off as i64)?;
-            }
-        }
+
+        file.align(16)?;
 
         blockindex.read_u128_be()?;
 
@@ -134,6 +130,7 @@ impl<T: Read + Seek> UnityFS<T> {
     pub fn read<'a>(&'a mut self, node: &NodeInfo) -> std::io::Result<NodeReader<'a, T>> {
         self.reader.seek(SeekFrom::Start(node.offset))?;
         Ok(NodeReader {
+            size: node.size,
             inner: (&mut self.reader).take(node.size),
         })
     }
@@ -231,6 +228,22 @@ impl<T: Read> Read for NodeReader<'_, T> {
         self.inner.read(buf)
     }
 }
+impl<T: Seek + Read> Seek for NodeReader<'_, T> {
+    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+        match pos {
+            SeekFrom::Current(n) if n >= 0 => {
+                for _ in 0..n {
+                    self.read_exact(&mut [0u8])?;
+                }
+                Ok(self.stream_position()?)
+            }
+            _ => unimplemented!(),
+        }
+    }
+    fn stream_position(&mut self) -> std::io::Result<u64> {
+        Ok(self.size - self.inner.limit())
+    }
+}
 
 #[derive(Debug, Clone, Copy, PartialEq)]
 enum CompressionScheme {