use crate::error::Error;
use byteorder::{WriteBytesExt, LE};
use snafu::Backtrace;
use snafu::GenerateBacktrace;
use std::cell::RefCell;
use std::fs::{self, File};
use std::io::{self, Cursor, Write};
use std::mem;
use std::path::{Path, PathBuf};
use std::rc::{Rc, Weak};
#[derive(Debug)]
struct RomFsDirEntCtx {
system_path: PathBuf,
name: String,
entry_offset: u32,
parent: Weak<RefCell<RomFsDirEntCtx>>,
child: Vec<Rc<RefCell<RomFsDirEntCtx>>>,
file: Vec<Rc<RefCell<RomFsFileEntCtx>>>,
}
impl RomFsDirEntCtx {
fn internal_path(&self) -> String {
let mut path = self.name.clone();
let mut cur = self.parent.upgrade().unwrap();
while cur.borrow().name != "" {
path = cur.borrow().name.clone() + "/" + &path;
let new_cur = cur.borrow().parent.upgrade().unwrap();
cur = new_cur;
}
path
}
}
#[derive(Debug)]
struct RomFsFileEntCtx {
system_path: PathBuf,
name: String,
entry_offset: u32,
offset: u64,
size: u64,
parent: Weak<RefCell<RomFsDirEntCtx>>,
}
impl RomFsFileEntCtx {
fn internal_path(&self) -> String {
let parent = self.parent.upgrade().unwrap();
let parent_borrow = parent.borrow();
parent_borrow.internal_path() + "/" + &self.name
}
}
#[repr(C)]
#[derive(Debug)]
struct RomFsDirEntryHdr {
parent: u32,
sibling: u32,
child: u32,
file: u32,
hash: u32,
name_size: u32,
}
#[repr(C)]
#[derive(Debug)]
struct RomFsFileEntryHdr {
parent: u32,
sibling: u32,
offset: u64,
size: u64,
hash: u32,
name_size: u32,
}
impl RomFsDirEntCtx {
#[allow(clippy::new_ret_no_self)]
fn new(parent: Weak<RefCell<RomFsDirEntCtx>>, path: PathBuf) -> Rc<RefCell<RomFsDirEntCtx>> {
let filename = path
.file_name()
.expect("Path to terminate properly")
.to_str()
.expect("Path to contain non-unicode chars")
.into();
Rc::new(RefCell::new(RomFsDirEntCtx {
system_path: path,
name: filename,
entry_offset: 0,
parent,
child: vec![],
file: vec![],
}))
}
fn new_root() -> Rc<RefCell<RomFsDirEntCtx>> {
let root = Rc::new(RefCell::new(RomFsDirEntCtx {
system_path: PathBuf::from(""),
name: String::from(""),
entry_offset: 0,
parent: Weak::new(),
child: vec![],
file: vec![],
}));
let weak = Rc::downgrade(&root);
root.borrow_mut().parent = weak;
root
}
}
fn romfs_get_hash_table_count(mut num_entries: usize) -> usize {
if num_entries < 3 {
3
} else if num_entries < 19 {
num_entries | 1
} else {
while num_entries % 2 == 0
|| num_entries % 3 == 0
|| num_entries % 5 == 0
|| num_entries % 7 == 0
|| num_entries % 11 == 0
|| num_entries % 13 == 0
|| num_entries % 17 == 0
{
num_entries += 1;
}
num_entries
}
}
fn align32(offset: u32, align: u32) -> u32 {
let mask = !(align - 1);
(offset + (align - 1)) & mask
}
fn align64(offset: u64, align: u64) -> u64 {
let mask = !(align - 1);
(offset + (align - 1)) & mask
}
fn calc_path_hash(parent: u32, path: &str) -> u32 {
let mut hash = parent ^ 123_456_789;
for c in path.as_bytes() {
hash = (hash >> 5) | (hash << 27);
hash ^= u32::from(*c);
}
hash
}
const ROMFS_FILEPARTITION_OFS: u64 = 0x200;
#[derive(Debug)]
pub struct RomFs {
dirs: Vec<Rc<RefCell<RomFsDirEntCtx>>>,
files: Vec<Rc<RefCell<RomFsFileEntCtx>>>,
dir_table_size: u64,
file_table_size: u64,
file_partition_size: u64,
}
#[allow(clippy::len_without_is_empty)]
impl RomFs {
pub fn push_file(&mut self, file_path: &Path, internal_path: &str) -> io::Result<()> {
let mut parent = self.dirs[0].clone();
let mut components = internal_path.split('/').peekable();
while let Some(component) = components.next() {
if components.peek().is_none() {
let metadata = file_path.metadata()?;
let file_to_add = Rc::new(RefCell::new(RomFsFileEntCtx {
system_path: PathBuf::from(file_path),
name: String::from(component),
entry_offset: 0,
offset: 0,
size: metadata.len(),
parent: Rc::downgrade(&parent),
}));
self.files.push(file_to_add.clone());
parent.borrow_mut().file.push(file_to_add.clone());
parent
.borrow_mut()
.file
.sort_by_key(|v| v.borrow().name.clone());
self.file_table_size += mem::size_of::<RomFsFileEntryHdr>() as u64
+ align64(file_to_add.borrow().name.len() as u64, 4);
} else {
if component == "" {
continue;
}
let new_parent = if let Some(child) = parent
.borrow()
.child
.iter()
.find(|v| v.borrow().name == component)
{
child.clone()
} else {
let child = Rc::new(RefCell::new(RomFsDirEntCtx {
system_path: PathBuf::from(""),
name: String::from(component),
entry_offset: 0,
parent: Rc::downgrade(&parent),
child: vec![],
file: vec![],
}));
self.dirs.push(child.clone());
parent.borrow_mut().child.push(child.clone());
parent
.borrow_mut()
.child
.sort_by_key(|v| v.borrow().name.clone());
self.dir_table_size += mem::size_of::<RomFsDirEntryHdr>() as u64
+ align64(child.borrow().name.len() as u64, 4);
child
};
parent = new_parent;
}
}
self.files.sort_by_key(|v| v.borrow().internal_path());
self.dirs.sort_by_key(|v| v.borrow().internal_path());
self.calculate_offsets();
Ok(())
}
pub fn empty() -> RomFs {
let root_folder = RomFsDirEntCtx::new_root();
let mut ctx = RomFs {
dirs: vec![root_folder],
files: vec![],
dir_table_size: mem::size_of::<RomFsDirEntryHdr>() as u64, file_table_size: 0,
file_partition_size: 0,
};
ctx.calculate_offsets();
ctx
}
pub fn from_directory(path: &Path) -> Result<RomFs, Error> {
let mut dirs = vec![];
let mut ctx = RomFs::empty();
ctx.dirs[0].borrow_mut().system_path = PathBuf::from(path);
dirs.push(ctx.dirs[0].clone());
while let Some(parent_dir) = dirs.pop() {
let path = parent_dir.borrow().system_path.clone();
for entry in fs::read_dir(&path).map_err(|err| (err, &path))? {
let entry = entry.map_err(|err| (err, &path))?;
let file_type = entry.file_type().map_err(|err| (err, entry.path()))?;
if file_type.is_dir() {
let new_dir = RomFsDirEntCtx::new(Rc::downgrade(&parent_dir), entry.path());
ctx.dirs.push(new_dir.clone());
dirs.push(new_dir.clone());
parent_dir.borrow_mut().child.push(new_dir.clone());
ctx.dir_table_size += mem::size_of::<RomFsDirEntryHdr>() as u64
+ align64(new_dir.borrow().name.len() as u64, 4);
} else if file_type.is_file() {
let file = Rc::new(RefCell::new(RomFsFileEntCtx {
system_path: entry.path(),
name: entry
.path()
.file_name()
.expect("Path to terminate properly")
.to_str()
.expect("Path to contain non-unicode chars")
.into(),
entry_offset: 0,
offset: 0,
size: entry.metadata().map_err(|err| (err, entry.path()))?.len(),
parent: Rc::downgrade(&parent_dir),
}));
ctx.files.push(file.clone());
parent_dir.borrow_mut().file.push(file.clone());
ctx.file_table_size += mem::size_of::<RomFsFileEntryHdr>() as u64
+ align64(file.borrow().name.len() as u64, 4);
} else if file_type.is_symlink() {
return Err(Error::RomFsSymlink {
error: entry.path(),
backtrace: Backtrace::generate(),
});
} else {
return Err(Error::RomFsFiletype {
error: entry.path(),
backtrace: Backtrace::generate(),
});
}
}
parent_dir
.borrow_mut()
.child
.sort_by_key(|v| v.borrow().name.clone());
parent_dir
.borrow_mut()
.file
.sort_by_key(|v| v.borrow().name.clone());
}
ctx.files.sort_by_key(|v| v.borrow().internal_path());
ctx.dirs.sort_by_key(|v| v.borrow().internal_path());
ctx.calculate_offsets();
Ok(ctx)
}
pub fn len(&self) -> usize {
(align64(ROMFS_FILEPARTITION_OFS + self.file_partition_size, 4)
+ romfs_get_hash_table_count(self.dirs.len() * mem::size_of::<u32>()) as u64
+ self.dir_table_size
+ romfs_get_hash_table_count(self.files.len() * mem::size_of::<u32>()) as u64
+ self.file_table_size) as usize
}
fn calculate_offsets(&mut self) {
let mut entry_offset = 0;
self.file_partition_size = 0;
for file in self.files.iter_mut() {
self.file_partition_size = align64(self.file_partition_size, 0x10);
file.borrow_mut().offset = self.file_partition_size;
self.file_partition_size += file.borrow().size;
file.borrow_mut().entry_offset = entry_offset;
entry_offset += mem::size_of::<RomFsFileEntryHdr>() as u32
+ align32(file.borrow().name.len() as u32, 4);
}
let mut entry_offset = 0;
for dir in self.dirs.iter_mut() {
dir.borrow_mut().entry_offset = entry_offset;
entry_offset += mem::size_of::<RomFsDirEntryHdr>() as u32
+ align32(dir.borrow().name.len() as u32, 4);
}
}
pub fn write(&self, to: &mut dyn Write) -> io::Result<()> {
const ROMFS_ENTRY_EMPTY: u32 = 0xFF_FF_FF_FF;
let mut dir_hash_table =
vec![ROMFS_ENTRY_EMPTY; romfs_get_hash_table_count(self.dirs.len())];
let mut file_hash_table =
vec![ROMFS_ENTRY_EMPTY; romfs_get_hash_table_count(self.files.len())];
let mut dir_table = vec![0u8; self.dir_table_size as usize];
let mut file_table = vec![0u8; self.file_table_size as usize];
for file in self.files.iter() {
let orig_file = file;
let file = file.borrow();
let parent = file.parent.upgrade().unwrap();
let parent = parent.borrow();
let sibling = parent
.file
.windows(2)
.find(|window| Rc::ptr_eq(&window[0], orig_file))
.map(|window| window[1].borrow().entry_offset);
let hash = calc_path_hash(parent.entry_offset, &file.name);
let mut cursor = Cursor::new(&mut file_table[file.entry_offset as usize..]);
cursor.write_u32::<LE>(parent.entry_offset)?;
cursor.write_u32::<LE>(sibling.unwrap_or(ROMFS_ENTRY_EMPTY))?;
cursor.write_u64::<LE>(file.offset)?;
cursor.write_u64::<LE>(file.size)?;
cursor.write_u32::<LE>(file_hash_table[hash as usize % file_hash_table.len()])?;
cursor.write_u32::<LE>(file.name.len() as u32)?;
cursor.write_all(file.name.as_bytes())?;
let cur_len = file_hash_table.len();
file_hash_table[hash as usize % cur_len] = file.entry_offset;
}
for dir in self.dirs.iter() {
let dir = dir.borrow();
let parent = dir.parent.upgrade().unwrap();
let parent = parent.borrow();
let sibling = parent
.child
.windows(2)
.find(|window| window[0].borrow().internal_path() == dir.internal_path())
.map(|window| window[1].borrow().entry_offset);
let hash = calc_path_hash(parent.entry_offset, &dir.name);
let mut cursor = Cursor::new(&mut dir_table[dir.entry_offset as usize..]);
cursor.write_u32::<LE>(parent.entry_offset)?;
cursor.write_u32::<LE>(sibling.unwrap_or(ROMFS_ENTRY_EMPTY))?;
cursor.write_u32::<LE>(
dir.child
.first()
.map(|v| v.borrow().entry_offset)
.unwrap_or(ROMFS_ENTRY_EMPTY),
)?;
cursor.write_u32::<LE>(
dir.file
.first()
.map(|v| v.borrow().entry_offset)
.unwrap_or(ROMFS_ENTRY_EMPTY),
)?;
cursor.write_u32::<LE>(dir_hash_table[hash as usize % dir_hash_table.len()])?;
cursor.write_u32::<LE>(dir.name.len() as u32)?;
cursor.write_all(dir.name.as_bytes())?;
let cur_len = dir_hash_table.len();
dir_hash_table[hash as usize % cur_len] = dir.entry_offset;
}
to.write_u64::<LE>(80)?;
let cur_ofs = align64(ROMFS_FILEPARTITION_OFS + self.file_partition_size, 4);
to.write_u64::<LE>(cur_ofs)?; to.write_u64::<LE>((dir_hash_table.len() * mem::size_of::<u32>()) as u64)?;
let cur_ofs = cur_ofs + (dir_hash_table.len() * mem::size_of::<u32>()) as u64;
to.write_u64::<LE>(cur_ofs)?; to.write_u64::<LE>(dir_table.len() as u64)?;
let cur_ofs = cur_ofs + dir_table.len() as u64;
to.write_u64::<LE>(cur_ofs)?; to.write_u64::<LE>((file_hash_table.len() * mem::size_of::<u32>()) as u64)?;
let cur_ofs = cur_ofs + (file_hash_table.len() * mem::size_of::<u32>()) as u64;
to.write_u64::<LE>(cur_ofs)?; to.write_u64::<LE>(file_table.len() as u64)?;
to.write_u64::<LE>(ROMFS_FILEPARTITION_OFS)?;
to.write_all(&[0; 0x1B0])?;
let mut cur_ofs = 0x200;
for file in self.files.iter() {
let new_cur_ofs = align64(cur_ofs, 0x10);
to.write_all(&vec![0; (new_cur_ofs - cur_ofs) as usize])?;
cur_ofs = new_cur_ofs;
println!(
"Writing {} to RomFS image...",
file.borrow().system_path.to_string_lossy()
);
assert_eq!(file.borrow().offset, cur_ofs - 0x200, "Wrong offset");
let len = io::copy(&mut File::open(&file.borrow().system_path)?, to)?;
assert_eq!(len, file.borrow().size, "File changed while building romfs");
cur_ofs += file.borrow().size;
}
let new_cur_ofs = align64(cur_ofs, 4);
to.write_all(&vec![0; (new_cur_ofs - cur_ofs) as usize])?;
let cur_ofs = new_cur_ofs;
assert_eq!(
cur_ofs,
align64(ROMFS_FILEPARTITION_OFS + self.file_partition_size, 4)
);
for hash in dir_hash_table {
to.write_u32::<LE>(hash)?;
}
to.write_all(&dir_table)?;
for hash in file_hash_table {
to.write_u32::<LE>(hash)?;
}
to.write_all(&file_table)?;
Ok(())
}
}