more abstraction around unityfs to read multiple files from a single reader

1 files changed, 175 insertions, 0 deletions

diff --git a/src/unityfs/header.rs b/src/unityfs/header.rs
new file mode 100644
index 0000000..d4fc89f
--- /dev/null
+++ b/src/unityfs/header.rs
@@ -0,0 +1,175 @@
+use crate::helper::{AlignExt, ReadExt};
+use anyhow::{Result, anyhow, bail};
+use humansize::DECIMAL;
+use log::{debug, info};
+use std::io::{Cursor, Read, Seek, SeekFrom};
+
+#[derive(Debug, Clone)]
+pub struct NodeInfo {
+    pub name: String,
+    pub size: u64,
+    pub(super) offset: u64,
+    _status: u32,
+}
+
+pub struct BlockInfo {
+    pub comp_size: u32,
+    pub decomp_size: u32,
+    pub comp_scheme: CompressionScheme,
+}
+
+pub struct UnityFSHeader {
+    pub(crate) nodes: Vec<NodeInfo>,
+    pub file_version: u32,
+    pub player_version: String,
+    pub unity_version: String,
+}
+
+impl UnityFSHeader {
+    pub fn read(mut file: impl Read + Seek) -> Result<(Self, Vec<BlockInfo>)> {
+        let signature = file.read_cstr()?;
+        if signature.as_str() != "UnityFS" {
+            bail!("unknown signature {signature:?}")
+        }
+
+        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 as u8).ok_or(anyhow!(
+            "unknown block compression 0x{:02x}",
+            (flags & 0x3f) as u8
+        ))?;
+        let blockindex_eof = flags & 0x80 != 0;
+        let blockindex_has_directory = flags & 0x40 != 0;
+        let blockindex_need_padding = flags & 0x200 != 0;
+
+        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)
+        };
+
+        file.align(16)?;
+
+        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 as u8)
+                .ok_or(anyhow!("unknown block compression 0x{:02x}", flags & 0x3f))?;
+            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:?} (size={}, status={status})",
+                humansize::format_size(size, DECIMAL)
+            );
+            nodes.push(NodeInfo {
+                offset,
+                size,
+                _status: status,
+                name,
+            })
+        }
+
+        Ok((
+            Self {
+                file_version,
+                player_version,
+                unity_version,
+                nodes,
+            },
+            blocks,
+        ))
+    }
+    pub fn nodes(&self) -> &[NodeInfo] {
+        &self.nodes
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub 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)?)
+            }
+            // CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
+            //     Ok(lz4::block::decompress(&block, Some(decomp_size as i32))?)
+            // }
+            CompressionScheme::Lzham => todo!(),
+        }
+    }
+}