unityfs library usage

3 files changed, 289 insertions, 144 deletions

diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..eed5b02
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,2 @@
+pub mod unityfs;
+pub mod helper;
diff --git a/src/main.rs b/src/main.rs
index 5ba3e3a..69ff519 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,156 +1,21 @@
-pub mod helper;
-
-use anyhow::{Result, anyhow};
-use helper::ReadExt;
-use log::{debug, info};
 use std::{
     env::args,
     fs::File,
-    io::{BufReader, Cursor, Read, Seek, SeekFrom, Write, copy},
+    io::{BufReader, copy},
 };
 
+use assetdebundle::unityfs::UnityFS;
+
 fn main() -> anyhow::Result<()> {
     env_logger::init_from_env("LOG");
-    let mut file = BufReader::new(File::open(args().nth(1).unwrap())?);
-
-    let signature = file.read_cstr()?;
-    let file_version = file.read_u32_be()?;
-    let player_version = file.read_cstr()?;
-    let unity_version = file.read_cstr()?;
-    let size = file.read_u64_be()?;
-    let blockindex_comp_size = file.read_u32_be()?;
-    let blockindex_decomp_size = file.read_u32_be()?;
-    let flags = file.read_u32_be()?;
-
-    let meta_comp_scheme = CompressionScheme::from_flag_num((flags & 0x3f) as u8)
-        .ok_or(anyhow!("unknown block compression"))?;
-    let blockindex_eof = flags & 0x80 != 0;
-    let blockindex_has_directory = flags & 0x40 != 0;
-    let blockindex_need_padding = flags & 0x200 != 0;
-
-    info!("signature={signature:?}");
-    info!("file_version={file_version:?}");
-    info!("player_version={player_version:?}");
-    info!("unity_version={unity_version:?}");
-    debug!("size={size:?}");
-    debug!("meta_comp_size={blockindex_comp_size:?}");
-    debug!("meta_decomp_size={blockindex_decomp_size:?}");
-    debug!("flags={flags:?}");
-    debug!("meta_comp_scheme={meta_comp_scheme:?}");
-    debug!("blockindex_eof={blockindex_eof:?}");
-    debug!("blockindex_has_directory={blockindex_has_directory:?}");
-    debug!("blockindex_need_padding={blockindex_need_padding:?}");
-
-    let mut blockindex = {
-        let restore_position = if blockindex_eof {
-            let pos = file.stream_position()?;
-            file.seek(SeekFrom::End(-(blockindex_comp_size as i64)))?;
-            Some(pos)
-        } else {
-            None
-        };
+    let file = BufReader::new(File::open(args().nth(1).unwrap())?);
+    let mut fs = UnityFS::open(file)?;
 
-        let mut blockindex = vec![0u8; blockindex_comp_size as usize];
-        file.read_exact(&mut blockindex)?;
-
-        if let Some(pos) = restore_position {
-            file.seek(SeekFrom::Start(pos))?;
-        }
-        let blockindex =
-            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)?;
-        }
-    }
-
-    blockindex.read_u128_be()?;
-
-    let num_blocks = blockindex.read_u32_be()?;
-    debug!("num_blocks={num_blocks:?}");
-    let mut blocks = Vec::new();
-    for _ in 0..num_blocks {
-        let block_decomp_size = blockindex.read_u32_be()?;
-        let block_comp_size = blockindex.read_u32_be()?;
-        let block_flags = blockindex.read_u16_be()?;
-        let block_comp_scheme = CompressionScheme::from_flag_num((block_flags & 0x3f) as u8)
-            .ok_or(anyhow!("unknown block compression"))?;
-        blocks.push((block_comp_size, block_decomp_size, block_comp_scheme))
-    }
-
-    let num_nodes = blockindex.read_u32_be()?;
-    debug!("num_nodes={num_nodes:?}");
-    let mut nodes = Vec::new();
-    for _ in 0..num_nodes {
-        let offset = blockindex.read_u64_be()?;
-        let size = blockindex.read_u64_be()?;
-        let status = blockindex.read_u32_be()?;
-        let name = blockindex.read_cstr()?;
-        info!("found node {name:?} (offset={offset}, size={size}, status={status})");
-        nodes.push((offset, size, status, name))
-    }
-
-    let mut storage = Vec::new();
-    for (comp_size, decomp_size, comp_scheme) in blocks {
-        let mut comp_buf = vec![0u8; comp_size as usize];
-        file.read_exact(&mut comp_buf)?;
-        let decomp_buf = comp_scheme.decompress(comp_buf, decomp_size as usize)?;
-        assert_eq!(decomp_size, decomp_buf.len() as u32);
-        storage.extend_from_slice(&decomp_buf);
-    }
-    let mut storage = Cursor::new(storage);
-
-    for (offset, size, _status, name) in nodes {
-        storage.seek(SeekFrom::Start(offset))?;
-        info!("extracting {name:?}");
-        let mut file = File::create(format!("/tmp/{}", name))?;
-        copy(&mut (&mut storage).take(size), &mut file)?;
+    for node in fs.nodes().to_vec() {
+        let mut reader = fs.read(&node)?;
+        let mut writer = File::create(format!("/tmp/{}", node.name))?;
+        copy(&mut reader, &mut writer)?;
     }
 
     Ok(())
 }
-
-#[derive(Debug, Clone, Copy, PartialEq)]
-enum CompressionScheme {
-    None,
-    LZMA,
-    LZ4,
-    LZ4HC,
-    LZHAM,
-}
-impl CompressionScheme {
-    pub fn from_flag_num(n: u8) -> Option<CompressionScheme> {
-        Some(match n & 0x3f {
-            0 => CompressionScheme::None,
-            1 => CompressionScheme::LZMA,
-            2 => CompressionScheme::LZ4,
-            3 => CompressionScheme::LZ4HC,
-            4 => CompressionScheme::LZHAM,
-            _ => return None,
-        })
-    }
-    pub fn decompress(&self, block: Vec<u8>, decomp_size: usize) -> Result<Vec<u8>> {
-        match self {
-            CompressionScheme::None => Ok(block),
-            CompressionScheme::LZMA => {
-                let mut r = lzma::Reader::from(Cursor::new(block))?;
-                let mut buf = Vec::new();
-                r.read_to_end(&mut buf)?;
-                Ok(buf)
-            }
-            // CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
-            //     Ok(lz4_flex::block::decompress(&block, decomp_size).context("lz4 decomp")?)
-            // }
-            CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
-                File::create("/tmp/a")?.write_all(&block)?;
-                Ok(lz4::block::decompress(&block, Some(decomp_size as i32))?)
-                // Ok(lz4::block::decompress(&block, None)?)
-            }
-            CompressionScheme::LZHAM => todo!(),
-        }
-    }
-}
diff --git a/src/unityfs.rs b/src/unityfs.rs
new file mode 100644
index 0000000..0d5668c
--- /dev/null
+++ b/src/unityfs.rs
@@ -0,0 +1,278 @@
+use crate::helper::ReadExt;
+use anyhow::{Result, anyhow, bail};
+use log::{debug, info, trace};
+use std::{
+    fs::File,
+    io::{Cursor, ErrorKind, Read, Seek, SeekFrom, Take, Write},
+};
+
+pub struct UnityFS<T> {
+    nodes: Vec<NodeInfo>,
+    reader: BlocksReader<T>,
+}
+
+#[derive(Debug, Clone)]
+pub struct NodeInfo {
+    pub name: String,
+    pub size: u64,
+    offset: u64,
+    _status: u32,
+}
+
+struct BlockInfo {
+    comp_size: u32,
+    decomp_size: u32,
+    comp_scheme: CompressionScheme,
+}
+
+pub struct NodeReader<'a, T> {
+    inner: Take<&'a mut BlocksReader<T>>,
+}
+
+impl<T: Read + Seek> UnityFS<T> {
+    pub fn open(mut file: T) -> Result<Self> {
+        let signature = file.read_cstr()?;
+        let file_version = file.read_u32_be()?;
+        let player_version = file.read_cstr()?;
+        let unity_version = file.read_cstr()?;
+        let size = file.read_u64_be()?;
+        let blockindex_comp_size = file.read_u32_be()?;
+        let blockindex_decomp_size = file.read_u32_be()?;
+        let flags = file.read_u32_be()?;
+
+        let meta_comp_scheme = CompressionScheme::from_flag_num((flags & 0x3f) as u8)
+            .ok_or(anyhow!("unknown block compression"))?;
+        let blockindex_eof = flags & 0x80 != 0;
+        let blockindex_has_directory = flags & 0x40 != 0;
+        let blockindex_need_padding = flags & 0x200 != 0;
+
+        if signature.as_str() != "UnityFS" {
+            bail!("unknown signature {signature:?}")
+        }
+        info!("File Version: {file_version:?}");
+        info!("Player Version: {player_version:?}");
+        info!("Unity Version: {unity_version:?}");
+        debug!("size={size:?}");
+        debug!("meta_comp_size={blockindex_comp_size:?}");
+        debug!("meta_decomp_size={blockindex_decomp_size:?}");
+        debug!("flags={flags:?}");
+        debug!("meta_comp_scheme={meta_comp_scheme:?}");
+        debug!("blockindex_eof={blockindex_eof:?}");
+        debug!("blockindex_has_directory={blockindex_has_directory:?}");
+        debug!("blockindex_need_padding={blockindex_need_padding:?}");
+
+        let mut blockindex = {
+            let restore_position = if blockindex_eof {
+                let pos = file.stream_position()?;
+                file.seek(SeekFrom::End(-(blockindex_comp_size as i64)))?;
+                Some(pos)
+            } else {
+                None
+            };
+
+            let mut blockindex = vec![0u8; blockindex_comp_size as usize];
+            file.read_exact(&mut blockindex)?;
+
+            if let Some(pos) = restore_position {
+                file.seek(SeekFrom::Start(pos))?;
+            }
+            let blockindex =
+                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)?;
+            }
+        }
+
+        blockindex.read_u128_be()?;
+
+        let num_blocks = blockindex.read_u32_be()?;
+        info!("File has {num_blocks} blocks");
+        let mut blocks = Vec::new();
+        for _ in 0..num_blocks {
+            let decomp_size = blockindex.read_u32_be()?;
+            let comp_size = blockindex.read_u32_be()?;
+            let flags = blockindex.read_u16_be()?;
+            let comp_scheme = CompressionScheme::from_flag_num((flags & 0x3f) as u8)
+                .ok_or(anyhow!("unknown block compression"))?;
+            blocks.push(BlockInfo {
+                comp_size,
+                decomp_size,
+                comp_scheme,
+            })
+        }
+
+        let num_nodes = blockindex.read_u32_be()?;
+        debug!("num_nodes={num_nodes:?}");
+        let mut nodes = Vec::new();
+        for _ in 0..num_nodes {
+            let offset = blockindex.read_u64_be()?;
+            let size = blockindex.read_u64_be()?;
+            let status = blockindex.read_u32_be()?;
+            let name = blockindex.read_cstr()?;
+            info!("found node {name:?} (offset={offset}, size={size}, status={status})");
+            nodes.push(NodeInfo {
+                offset,
+                size,
+                _status: status,
+                name,
+            })
+        }
+
+        let position = file.stream_position()?;
+
+        Ok(Self {
+            nodes,
+            reader: BlocksReader::new(blocks, file, position),
+        })
+    }
+
+    pub fn nodes(&self) -> &[NodeInfo] {
+        &self.nodes
+    }
+    pub fn read<'a>(&'a mut self, node: &NodeInfo) -> std::io::Result<NodeReader<'a, T>> {
+        self.reader.seek(SeekFrom::Start(node.offset))?;
+        Ok(NodeReader {
+            inner: (&mut self.reader).take(node.size),
+        })
+    }
+}
+
+struct BlocksReader<T> {
+    blocks: Vec<BlockInfo>,
+    inner: T,
+    inner_seek_offset: u64,
+    nblock_index: usize,
+    cblock_data: Vec<u8>,
+    cblock_off: usize,
+}
+impl<T: Read> BlocksReader<T> {
+    pub fn new(blocks: Vec<BlockInfo>, inner: T, inner_seek_offset: u64) -> Self {
+        Self {
+            blocks,
+            inner,
+            inner_seek_offset,
+            nblock_index: 0,
+            cblock_data: Vec::new(),
+            cblock_off: 0,
+        }
+    }
+    pub fn load_next_block(&mut self) -> std::io::Result<()> {
+        trace!("loading block {}", self.nblock_index);
+        let block = &self.blocks[self.nblock_index];
+        let mut comp_buf = vec![0; block.comp_size as usize];
+        File::create("/tmp/a")?.write_all(&comp_buf)?;
+        self.inner.read_exact(&mut comp_buf)?;
+        let decomp_buf = block
+            .comp_scheme
+            .decompress(comp_buf, block.decomp_size as usize)
+            .map_err(|e| {
+                std::io::Error::new(
+                    ErrorKind::InvalidData,
+                    format!("decompression failure: {e}"),
+                )
+            })?;
+        self.nblock_index += 1;
+        self.cblock_data = decomp_buf;
+        self.cblock_off = 0;
+        Ok(())
+    }
+}
+impl<T: Read> Read for BlocksReader<T> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        if self.cblock_off >= self.cblock_data.len() {
+            self.load_next_block()?;
+        }
+        let size = (self.cblock_data.len() - self.cblock_off).min(buf.len());
+        buf[..size].copy_from_slice(&self.cblock_data[self.cblock_off..self.cblock_off + size]);
+        self.cblock_off += size;
+        Ok(size)
+    }
+}
+impl<T: Seek + Read> Seek for BlocksReader<T> {
+    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+        let SeekFrom::Start(pos) = pos else {
+            unimplemented!()
+        };
+        debug!("seek decomp to {pos}");
+        let mut comp_off = self.inner_seek_offset;
+        let mut decomp_off = 0;
+        let mut target_block = None;
+        for (i, b) in self.blocks.iter().enumerate() {
+            if pos <= decomp_off + b.decomp_size as u64 {
+                target_block = Some(i);
+                break;
+            }
+            decomp_off += b.decomp_size as u64;
+            comp_off += b.comp_size as u64;
+        }
+
+        let Some(i) = target_block else {
+            return Err(std::io::Error::new(
+                ErrorKind::UnexpectedEof,
+                "seek out of bounds",
+            ));
+        };
+
+        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()?;
+        self.cblock_off = block_off as usize;
+
+        Ok(pos)
+    }
+}
+
+impl<T: Read> Read for NodeReader<'_, T> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.inner.read(buf)
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+enum CompressionScheme {
+    None,
+    LZMA,
+    LZ4,
+    LZ4HC,
+    LZHAM,
+}
+impl CompressionScheme {
+    pub fn from_flag_num(n: u8) -> Option<CompressionScheme> {
+        Some(match n & 0x3f {
+            0 => CompressionScheme::None,
+            1 => CompressionScheme::LZMA,
+            2 => CompressionScheme::LZ4,
+            3 => CompressionScheme::LZ4HC,
+            4 => CompressionScheme::LZHAM,
+            _ => return None,
+        })
+    }
+    pub fn decompress(&self, block: Vec<u8>, decomp_size: usize) -> Result<Vec<u8>> {
+        match self {
+            CompressionScheme::None => Ok(block),
+            CompressionScheme::LZMA => {
+                let mut r = lzma::Reader::from(Cursor::new(block))?;
+                let mut buf = Vec::new();
+                r.read_to_end(&mut buf)?;
+                Ok(buf)
+            }
+            // CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
+            //     Ok(lz4_flex::block::decompress(&block, decomp_size).context("lz4 decomp")?)
+            // }
+            CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
+                File::create("/tmp/a")?.write_all(&block)?;
+                Ok(lz4::block::decompress(&block, Some(decomp_size as i32))?)
+                // Ok(lz4::block::decompress(&block, None)?)
+            }
+            CompressionScheme::LZHAM => todo!(),
+        }
+    }
+}