1 files changed, 28 insertions, 8 deletions
diff --git a/src/unityfs.rs b/src/unityfs.rs
index d7659b8..bd0c7bb 100644
--- a/src/unityfs.rs
+++ b/src/unityfs.rs
@@ -1,7 +1,7 @@
 use crate::helper::{AlignExt, ReadExt};
 use anyhow::{Result, anyhow, bail};
 use log::{debug, info, trace};
-use std::io::{Cursor, ErrorKind, Read, Seek, SeekFrom, Take};
+use std::io::{Cursor, Error, ErrorKind, Read, Seek, SeekFrom};
 
 pub struct UnityFS<T> {
     nodes: Vec<NodeInfo>,
@@ -23,7 +23,9 @@ struct BlockInfo {
 }
 
 pub struct NodeReader<'a, T> {
-    inner: Take<&'a mut BlocksReader<T>>,
+    inner: &'a mut BlocksReader<T>,
+    position: u64,
+    offset: u64,
     size: u64,
 }
 
@@ -131,7 +133,9 @@ impl<T: Read + Seek> UnityFS<T> {
         self.reader.seek(SeekFrom::Start(node.offset))?;
         Ok(NodeReader {
             size: node.size,
-            inner: (&mut self.reader).take(node.size),
+            offset: node.offset,
+            position: 0,
+            inner: &mut self.reader,
         })
     }
 }
@@ -213,10 +217,14 @@ impl<T: Seek + Read> Seek for BlocksReader<T> {
 
         let block_off = pos - decomp_off;
         debug!("target is block={i} offset={block_off}");
-        debug!("seek comp to {comp_off}");
         self.inner.seek(SeekFrom::Start(comp_off))?;
-        self.nblock_index = i;
-        self.load_next_block()?;
+        if self.nblock_index == i + 1 {
+            debug!("intra-block seek")
+        } else {
+            debug!("seek comp to {comp_off}");
+            self.nblock_index = i;
+            self.load_next_block()?;
+        }
         self.cblock_off = block_off as usize;
 
         Ok(pos)
@@ -225,7 +233,11 @@ impl<T: Seek + Read> Seek for BlocksReader<T> {
 
 impl<T: Read> Read for NodeReader<'_, T> {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-        self.inner.read(buf)
+        let bytes_left = self.size - self.position;
+        let end = buf.len().min(bytes_left as usize);
+        let size = self.inner.read(&mut buf[..end])?;
+        self.position += size as u64;
+        Ok(size)
     }
 }
 impl<T: Seek + Read> Seek for NodeReader<'_, T> {
@@ -237,11 +249,19 @@ impl<T: Seek + Read> Seek for NodeReader<'_, T> {
                 }
                 Ok(self.stream_position()?)
             }
+            SeekFrom::Start(n) => {
+                debug!("seek node to {n} (off={})", self.offset);
+                if n > self.size {
+                    return Err(Error::new(ErrorKind::NotSeekable, "seek out of bounds"));
+                }
+                self.position = n;
+                self.inner.seek(SeekFrom::Start(self.offset + n))
+            }
             _ => unimplemented!(),
         }
     }
     fn stream_position(&mut self) -> std::io::Result<u64> {
-        Ok(self.size - self.inner.limit())
+        Ok(self.position)
     }
 }