path: root/src/unityfs.rs

use crate::helper::{AlignExt, ReadExt};
use anyhow::{Result, anyhow, bail};
use log::{debug, info, trace};
use std::io::{Cursor, ErrorKind, Read, Seek, SeekFrom, Take};

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>>,
    size: u64,
}

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)
        };

        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 & 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 {
            size: node.size,
            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];
        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)
    }
}
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 {
    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 => {
            //     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!(),
        }
    }
}