diff options
Diffstat (limited to 'src/unityfs.rs')
| -rw-r--r-- | src/unityfs.rs | 278 |
1 files changed, 278 insertions, 0 deletions
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!(), + } + } +} |