Add WAD support to VFS & `wad` commands

This commit is contained in:
Luke Street 2024-11-07 08:43:20 -07:00
parent 146c4d2f8c
commit 1cc38ad621
12 changed files with 770 additions and 9 deletions

4
Cargo.lock generated
View File

@ -348,14 +348,16 @@ dependencies = [
[[package]]
name = "decomp-toolkit"
version = "1.2.0"
version = "1.3.0"
dependencies = [
"aes",
"anyhow",
"ar",
"argp",
"base16ct",
"base64",
"byteorder",
"cbc",
"crossterm",
"cwdemangle",
"cwextab",

View File

@ -3,7 +3,7 @@ name = "decomp-toolkit"
description = "Yet another GameCube/Wii decompilation toolkit."
authors = ["Luke Street <luke@street.dev>"]
license = "MIT OR Apache-2.0"
version = "1.2.0"
version = "1.3.0"
edition = "2021"
publish = false
repository = "https://github.com/encounter/decomp-toolkit"
@ -25,6 +25,7 @@ strip = "debuginfo"
codegen-units = 1
[dependencies]
aes = "0.8"
anyhow = { version = "1.0", features = ["backtrace"] }
ar = { git = "https://github.com/bjorn3/rust-ar.git", branch = "write_symbol_table" }
argp = "0.3"
@ -32,6 +33,7 @@ base16ct = "0.2"
base64 = "0.22"
byteorder = "1.5"
typed-path = "0.9"
cbc = "0.1"
crossterm = "0.28"
cwdemangle = "1.0"
cwextab = "1.0"

View File

@ -52,6 +52,9 @@ project structure and build system that uses decomp-toolkit under the hood.
- [yay0 compress](#yay0-compress)
- [yaz0 decompress](#yaz0-decompress)
- [yaz0 compress](#yaz0-compress)
- [wad info](#wad-info)
- [wad extract](#wad-extract)
- [wad verify](#wad-verify)
## Goals
@ -474,6 +477,7 @@ Supported containers:
- Disc images (see [disc info](#disc-info) for supported formats)
- RARC archives (older .arc)
- U8 archives (newer .arc)
- WAD files (Wii VC)
Supported compression formats are handled transparently:
- Yay0 (SZP) / Yaz0 (SZS)
@ -562,3 +566,31 @@ $ dtk yaz0 compress input.bin -o output.bin.yaz0
# or, for batch processing
$ dtk yaz0 compress rels/* -o rels
```
### wad info
Prints information about a WAD file.
```shell
$ dtk wad info input.wad
```
### wad extract
> [!NOTE]
> [vfs cp](#vfs-cp) is more flexible and supports WAD files.
> This command is now equivalent to `dtk vfs cp input.wad: output_dir`
Extracts the contents of a WAD file.
```shell
$ dtk wad extract input.wad -o output_dir
```
### wad verify
Verifies the contents of a WAD file.
```shell
$ dtk wad verify input.wad
```

View File

@ -14,5 +14,6 @@ pub mod rso;
pub mod shasum;
pub mod u8_arc;
pub mod vfs;
pub mod wad;
pub mod yay0;
pub mod yaz0;

108
src/cmd/wad.rs Normal file
View File

@ -0,0 +1,108 @@
use anyhow::Result;
use argp::FromArgs;
use size::Size;
use typed_path::Utf8NativePathBuf;
use crate::{
cmd::vfs,
util::{
path::native_path,
wad::{process_wad, verify_wad},
},
vfs::open_file,
};
#[derive(FromArgs, PartialEq, Debug)]
/// Commands for processing Wii WAD files.
#[argp(subcommand, name = "wad")]
pub struct Args {
#[argp(subcommand)]
command: SubCommand,
}
#[derive(FromArgs, PartialEq, Debug)]
#[argp(subcommand)]
enum SubCommand {
Extract(ExtractArgs),
Info(InfoArgs),
Verify(VerifyArgs),
}
#[derive(FromArgs, PartialEq, Eq, Debug)]
/// Extracts WAD file contents.
#[argp(subcommand, name = "extract")]
pub struct ExtractArgs {
#[argp(positional, from_str_fn(native_path))]
/// WAD file
file: Utf8NativePathBuf,
#[argp(option, short = 'o', from_str_fn(native_path))]
/// output directory
output: Option<Utf8NativePathBuf>,
#[argp(switch)]
/// Do not decompress files when copying.
no_decompress: bool,
#[argp(switch, short = 'q')]
/// Quiet output. Don't print anything except errors.
quiet: bool,
}
#[derive(FromArgs, PartialEq, Eq, Debug)]
/// Views WAD file information.
#[argp(subcommand, name = "info")]
pub struct InfoArgs {
#[argp(positional, from_str_fn(native_path))]
/// WAD file
file: Utf8NativePathBuf,
}
#[derive(FromArgs, PartialEq, Eq, Debug)]
/// Verifies WAD file integrity.
#[argp(subcommand, name = "verify")]
pub struct VerifyArgs {
#[argp(positional, from_str_fn(native_path))]
/// WAD file
file: Utf8NativePathBuf,
}
pub fn run(args: Args) -> Result<()> {
match args.command {
SubCommand::Info(c_args) => info(c_args),
SubCommand::Verify(c_args) => verify(c_args),
SubCommand::Extract(c_args) => extract(c_args),
}
}
fn info(args: InfoArgs) -> Result<()> {
let mut file = open_file(&args.file, true)?;
let wad = process_wad(file.as_mut())?;
println!("Title ID: {}", hex::encode(wad.ticket().title_id));
println!("Title key: {}", hex::encode(wad.title_key));
println!("Fake signed: {}", wad.fake_signed);
for content in wad.contents() {
println!(
"Content {:08x}: Offset {:#X}, size {}",
content.content_index.get(),
wad.content_offset(content.content_index.get()),
Size::from_bytes(content.size.get())
);
}
Ok(())
}
fn verify(args: VerifyArgs) -> Result<()> {
let mut file = open_file(&args.file, true)?;
let wad = process_wad(file.as_mut())?;
verify_wad(&wad, file.as_mut())?;
println!("Verification successful");
Ok(())
}
fn extract(args: ExtractArgs) -> Result<()> {
let path = Utf8NativePathBuf::from(format!("{}:", args.file));
let output = args.output.unwrap_or_else(|| Utf8NativePathBuf::from("."));
vfs::cp(vfs::CpArgs {
paths: vec![path, output],
no_decompress: args.no_decompress,
quiet: args.quiet,
})
}

View File

@ -106,6 +106,7 @@ enum SubCommand {
Vfs(cmd::vfs::Args),
Yay0(cmd::yay0::Args),
Yaz0(cmd::yaz0::Args),
Wad(cmd::wad::Args),
}
// Duplicated from supports-color so we can check early.
@ -181,6 +182,7 @@ fn main() {
SubCommand::Vfs(c_args) => cmd::vfs::run(c_args),
SubCommand::Yay0(c_args) => cmd::yay0::run(c_args),
SubCommand::Yaz0(c_args) => cmd::yaz0::run(c_args),
SubCommand::Wad(c_args) => cmd::wad::run(c_args),
});
if let Err(e) = result {
eprintln!("Failed: {e:?}");

View File

@ -18,6 +18,7 @@ pub mod nested;
pub mod nlzss;
pub mod path;
pub mod rarc;
pub mod read;
pub mod reader;
pub mod rel;
pub mod rso;
@ -25,6 +26,7 @@ pub mod signatures;
pub mod split;
pub mod take_seek;
pub mod u8_arc;
pub mod wad;
#[inline]
pub const fn align_up(value: u32, align: u32) -> u32 { (value + (align - 1)) & !(align - 1) }

26
src/util/read.rs Normal file
View File

@ -0,0 +1,26 @@
use std::{io, io::Read};
use zerocopy::{FromBytes, FromZeros, IntoBytes};
#[inline(always)]
pub fn read_from<T, R>(reader: &mut R) -> io::Result<T>
where
T: FromBytes + IntoBytes,
R: Read + ?Sized,
{
let mut ret = <T>::new_zeroed();
reader.read_exact(ret.as_mut_bytes())?;
Ok(ret)
}
#[inline(always)]
pub fn read_box_slice<T, R>(reader: &mut R, count: usize) -> io::Result<Box<[T]>>
where
T: FromBytes + IntoBytes,
R: Read + ?Sized,
{
let mut ret = <[T]>::new_box_zeroed_with_elems(count)
.map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
reader.read_exact(ret.as_mut().as_mut_bytes())?;
Ok(ret)
}

220
src/util/wad.rs Normal file
View File

@ -0,0 +1,220 @@
use std::{
io,
io::{BufRead, Read, Seek},
};
use aes::cipher::{BlockDecryptMut, KeyIvInit};
use anyhow::{bail, Result};
use nodtool::nod::{Ticket, TmdHeader};
use sha1::{Digest, Sha1};
use size::Size;
use zerocopy::{big_endian::*, FromBytes, FromZeros, Immutable, IntoBytes, KnownLayout};
use crate::{
array_ref_mut, static_assert,
util::read::{read_box_slice, read_from},
};
// TODO: other WAD types?
pub const WAD_MAGIC: [u8; 8] = [0x00, 0x00, 0x00, 0x20, 0x49, 0x73, 0x00, 0x00];
#[derive(Clone, Debug, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)]
#[repr(C, align(4))]
pub struct WadHeader {
pub header_size: U32,
pub wad_type: [u8; 0x2],
pub wad_version: U16,
pub cert_chain_size: U32,
pub _reserved1: [u8; 0x4],
pub ticket_size: U32,
pub tmd_size: U32,
pub data_size: U32,
pub footer_size: U32,
}
static_assert!(size_of::<WadHeader>() == 0x20);
#[derive(Clone, Debug, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)]
#[repr(C, align(4))]
pub struct ContentMetadata {
pub content_id: U32,
pub content_index: U16,
pub content_type: U16,
pub size: U64,
pub hash: HashBytes,
}
static_assert!(size_of::<ContentMetadata>() == 0x24);
impl ContentMetadata {
#[inline]
pub fn iv(&self) -> [u8; 0x10] {
let mut iv = [0u8; 0x10];
*array_ref_mut!(iv, 0, 2) = self.content_index.get().to_be_bytes();
iv
}
}
const ALIGNMENT: usize = 0x40;
#[inline(always)]
pub fn align_up(value: u64, alignment: u64) -> u64 { (value + alignment - 1) & !(alignment - 1) }
pub type HashBytes = [u8; 20];
pub type KeyBytes = [u8; 16];
type Aes128Cbc = cbc::Decryptor<aes::Aes128>;
#[derive(Debug, Clone)]
pub struct WadFile {
pub header: WadHeader,
pub title_key: KeyBytes,
pub fake_signed: bool,
pub raw_cert_chain: Box<[u8]>,
pub raw_ticket: Box<[u8]>,
pub raw_tmd: Box<[u8]>,
pub content_offset: u64,
}
impl WadFile {
pub fn ticket(&self) -> &Ticket {
Ticket::ref_from_bytes(&self.raw_ticket).expect("Invalid ticket alignment")
}
pub fn tmd(&self) -> &TmdHeader {
TmdHeader::ref_from_prefix(&self.raw_tmd).expect("Invalid TMD alignment").0
}
pub fn contents(&self) -> &[ContentMetadata] {
let (_, cmd_data) =
TmdHeader::ref_from_prefix(&self.raw_tmd).expect("Invalid TMD alignment");
<[ContentMetadata]>::ref_from_bytes(cmd_data).expect("Invalid CMD alignment")
}
pub fn content_offset(&self, content_index: u16) -> u64 {
let contents = self.contents();
let mut offset = self.content_offset;
for content in contents.iter().take(content_index as usize) {
offset = align_up(offset + content.size.get(), ALIGNMENT as u64);
}
offset
}
pub fn trailer_offset(&self) -> u64 {
let contents = self.contents();
let mut offset = self.content_offset;
for content in contents.iter() {
offset = align_up(offset + content.size.get(), ALIGNMENT as u64);
}
offset
}
}
pub fn process_wad<R>(reader: &mut R) -> Result<WadFile>
where R: BufRead + Seek + ?Sized {
let header: WadHeader = read_from(reader)?;
let mut offset = align_up(header.header_size.get() as u64, ALIGNMENT as u64);
reader.seek(io::SeekFrom::Start(offset))?;
let raw_cert_chain: Box<[u8]> = read_box_slice(reader, header.cert_chain_size.get() as usize)?;
offset = align_up(offset + header.cert_chain_size.get() as u64, ALIGNMENT as u64);
reader.seek(io::SeekFrom::Start(offset))?;
let raw_ticket: Box<[u8]> = read_box_slice(reader, header.ticket_size.get() as usize)?;
offset = align_up(offset + header.ticket_size.get() as u64, ALIGNMENT as u64);
reader.seek(io::SeekFrom::Start(offset))?;
let raw_tmd: Box<[u8]> = read_box_slice(reader, header.tmd_size.get() as usize)?;
offset = align_up(offset + header.tmd_size.get() as u64, ALIGNMENT as u64);
let content_offset = offset;
let mut file = WadFile {
header,
title_key: [0; 16],
fake_signed: false,
raw_cert_chain,
raw_ticket,
raw_tmd,
content_offset,
};
let mut title_key_found = false;
if file.ticket().header.sig.iter().all(|&x| x == 0) {
// Fake signed, try to determine common key index
file.fake_signed = true;
let contents = file.contents();
if let Some(smallest_content) = contents.iter().min_by_key(|x| x.size.get()) {
let mut ticket = file.ticket().clone();
for i in 0..2 {
ticket.common_key_idx = i;
let title_key = ticket.decrypt_title_key()?;
let offset = file.content_offset(smallest_content.content_index.get());
reader.seek(io::SeekFrom::Start(offset))?;
if verify_content(reader, smallest_content, &title_key)? {
file.title_key = title_key;
title_key_found = true;
break;
}
}
}
if !title_key_found {
bail!("Failed to determine title key for fake signed WAD");
}
}
if !title_key_found {
let title_key = file.ticket().decrypt_title_key()?;
file.title_key = title_key;
}
Ok(file)
}
pub fn verify_wad<R>(file: &WadFile, reader: &mut R) -> Result<()>
where R: Read + Seek + ?Sized {
for content in file.contents() {
let content_index = content.content_index.get();
println!(
"Verifying content {:08x} (size {})",
content_index,
Size::from_bytes(content.size.get())
);
let offset = file.content_offset(content_index);
reader.seek(io::SeekFrom::Start(offset))?;
if !verify_content(reader, content, &file.title_key)? {
bail!("Content {:08x} hash mismatch", content_index);
}
}
Ok(())
}
fn verify_content<R>(
reader: &mut R,
content: &ContentMetadata,
title_key: &KeyBytes,
) -> Result<bool>
where
R: Read + ?Sized,
{
let mut buf = <[[u8; 0x10]]>::new_box_zeroed_with_elems(0x200)
.map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
// Read full padded size for decryption
let read_size = align_up(content.size.get(), 0x40);
let mut decryptor = Aes128Cbc::new(title_key.into(), (&content.iv()).into());
let mut digest = Sha1::default();
let mut read = 0;
while read < read_size {
let len = buf.len().min(usize::try_from(read_size - read).unwrap_or(usize::MAX));
debug_assert_eq!(len % 0x10, 0);
reader.read_exact(&mut buf.as_mut_bytes()[..len])?;
for block in buf.iter_mut().take(len / 0x10) {
decryptor.decrypt_block_mut(block.into());
}
// Only hash up to content size
let hash_len = (read + len as u64).min(content.size.get()).saturating_sub(read) as usize;
if hash_len > 0 {
digest.update(&buf.as_bytes()[..hash_len]);
}
read += len as u64;
}
Ok(HashBytes::from(digest.finalize()) == content.hash)
}

View File

@ -112,7 +112,7 @@ impl Seek for WindowedFile {
}
#[inline]
fn stream_position(&mut self) -> io::Result<u64> { Ok(self.pos) }
fn stream_position(&mut self) -> io::Result<u64> { Ok(self.pos - self.begin) }
}
impl VfsFile for WindowedFile {

View File

@ -3,6 +3,7 @@ mod disc;
mod rarc;
mod std_fs;
mod u8_arc;
mod wad;
use std::{
error::Error,
@ -22,12 +23,14 @@ use rarc::RarcFs;
pub use std_fs::StdFs;
use typed_path::{Utf8NativePath, Utf8UnixPath, Utf8UnixPathBuf};
use u8_arc::U8Fs;
use wad::WadFs;
use crate::util::{
ncompress::{YAY0_MAGIC, YAZ0_MAGIC},
nlzss,
rarc::RARC_MAGIC,
u8_arc::U8_MAGIC,
wad::WAD_MAGIC,
};
pub trait Vfs: DynClone + Send + Sync {
@ -154,6 +157,7 @@ pub enum ArchiveKind {
Rarc,
U8,
Disc(nod::Format),
Wad,
}
impl Display for ArchiveKind {
@ -162,6 +166,7 @@ impl Display for ArchiveKind {
ArchiveKind::Rarc => write!(f, "RARC"),
ArchiveKind::U8 => write!(f, "U8"),
ArchiveKind::Disc(format) => write!(f, "Disc ({})", format),
ArchiveKind::Wad => write!(f, "WAD"),
}
}
}
@ -169,18 +174,19 @@ impl Display for ArchiveKind {
pub fn detect<R>(file: &mut R) -> io::Result<FileFormat>
where R: Read + Seek + ?Sized {
file.seek(SeekFrom::Start(0))?;
let mut magic = [0u8; 4];
let mut magic = [0u8; 8];
match file.read_exact(&mut magic) {
Ok(_) => {}
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(FileFormat::Regular),
Err(e) => return Err(e),
}
file.seek_relative(-4)?;
file.seek_relative(-8)?;
match magic {
YAY0_MAGIC => Ok(FileFormat::Compressed(CompressionKind::Yay0)),
YAZ0_MAGIC => Ok(FileFormat::Compressed(CompressionKind::Yaz0)),
RARC_MAGIC => Ok(FileFormat::Archive(ArchiveKind::Rarc)),
U8_MAGIC => Ok(FileFormat::Archive(ArchiveKind::U8)),
_ if magic.starts_with(&YAY0_MAGIC) => Ok(FileFormat::Compressed(CompressionKind::Yay0)),
_ if magic.starts_with(&YAZ0_MAGIC) => Ok(FileFormat::Compressed(CompressionKind::Yaz0)),
_ if magic.starts_with(&RARC_MAGIC) => Ok(FileFormat::Archive(ArchiveKind::Rarc)),
_ if magic.starts_with(&U8_MAGIC) => Ok(FileFormat::Archive(ArchiveKind::U8)),
WAD_MAGIC => Ok(FileFormat::Archive(ArchiveKind::Wad)),
_ => {
let format = nod::Disc::detect(file)?;
file.seek(SeekFrom::Start(0))?;
@ -332,6 +338,7 @@ pub fn open_fs(mut file: Box<dyn VfsFile>, kind: ArchiveKind) -> io::Result<Box<
disc.open_partition_kind(nod::PartitionKind::Data).map_err(nod_to_io_error)?;
Ok(Box::new(DiscFs::new(disc, partition, metadata.mtime)?))
}
ArchiveKind::Wad => Ok(Box::new(WadFs::new(file)?)),
}
}

359
src/vfs/wad.rs Normal file
View File

@ -0,0 +1,359 @@
use std::{
io,
io::{BufRead, Cursor, Read, Seek, SeekFrom},
sync::Arc,
};
use aes::cipher::{BlockDecryptMut, KeyIvInit};
use filetime::FileTime;
use nodtool::nod::DiscStream;
use typed_path::Utf8UnixPath;
use zerocopy::FromZeros;
use crate::{
array_ref,
util::wad::{align_up, process_wad, ContentMetadata, WadFile},
vfs::{
common::{StaticFile, WindowedFile},
Vfs, VfsError, VfsFile, VfsFileType, VfsMetadata, VfsResult,
},
};
#[derive(Clone)]
pub struct WadFs {
file: Box<dyn VfsFile>,
wad: WadFile,
}
enum WadFindResult<'a> {
Root,
Static(&'a [u8]),
Content(u16, &'a ContentMetadata),
Window(u64, u64),
}
impl WadFs {
pub fn new(mut file: Box<dyn VfsFile>) -> io::Result<Self> {
let wad = process_wad(file.as_mut())
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
Ok(Self { file, wad })
}
fn find(&self, path: &str) -> Option<WadFindResult> {
let filename = path.trim_start_matches('/');
if filename.contains('/') {
return None;
}
if filename.is_empty() {
return Some(WadFindResult::Root);
}
let filename = filename.to_ascii_lowercase();
if let Some(id) = filename.strip_suffix(".app") {
if let Ok(content_index) = u16::from_str_radix(id, 16) {
if let Some(content) = self.wad.contents().get(content_index as usize) {
return Some(WadFindResult::Content(content_index, content));
}
}
return None;
}
let title_id = hex::encode(self.wad.ticket().title_id);
match filename.strip_prefix(&title_id) {
Some(".tik") => Some(WadFindResult::Static(&self.wad.raw_ticket)),
Some(".tmd") => Some(WadFindResult::Static(&self.wad.raw_tmd)),
Some(".cert") => Some(WadFindResult::Static(&self.wad.raw_cert_chain)),
Some(".trailer") => {
if self.wad.header.footer_size.get() == 0 {
return None;
}
Some(WadFindResult::Window(
self.wad.trailer_offset(),
self.wad.header.footer_size.get() as u64,
))
}
_ => None,
}
}
}
impl Vfs for WadFs {
fn open(&mut self, path: &Utf8UnixPath) -> VfsResult<Box<dyn VfsFile>> {
if let Some(result) = self.find(path.as_str()) {
match result {
WadFindResult::Root => Err(VfsError::IsADirectory),
WadFindResult::Static(data) => {
Ok(Box::new(StaticFile::new(Arc::from(data), self.file.metadata()?.mtime)))
}
WadFindResult::Content(content_index, content) => {
let offset = self.wad.content_offset(content_index);
Ok(Box::new(WadContent::new(
AesCbcStream::new(
self.file.clone(),
offset,
content.size.get(),
&self.wad.title_key,
&content.iv(),
),
self.file.metadata()?.mtime,
)))
}
WadFindResult::Window(offset, len) => {
Ok(Box::new(WindowedFile::new(self.file.clone(), offset, len)?))
}
}
} else {
Err(VfsError::NotFound)
}
}
fn exists(&mut self, path: &Utf8UnixPath) -> VfsResult<bool> {
Ok(self.find(path.as_str()).is_some())
}
fn read_dir(&mut self, path: &Utf8UnixPath) -> VfsResult<Vec<String>> {
let path = path.as_str().trim_start_matches('/');
if !path.is_empty() {
return Err(VfsError::NotFound);
}
let title_id = hex::encode(self.wad.ticket().title_id);
let mut entries = Vec::new();
entries.push(format!("{}.tik", title_id));
entries.push(format!("{}.tmd", title_id));
entries.push(format!("{}.cert", title_id));
if self.wad.header.footer_size.get() > 0 {
entries.push(format!("{}.trailer", title_id));
}
for content in self.wad.contents() {
entries.push(format!("{:08x}.app", content.content_index.get()));
}
Ok(entries)
}
fn metadata(&mut self, path: &Utf8UnixPath) -> VfsResult<VfsMetadata> {
let mtime = self.file.metadata()?.mtime;
if let Some(result) = self.find(path.as_str()) {
match result {
WadFindResult::Root => {
Ok(VfsMetadata { file_type: VfsFileType::Directory, len: 0, mtime })
}
WadFindResult::Static(data) => {
Ok(VfsMetadata { file_type: VfsFileType::File, len: data.len() as u64, mtime })
}
WadFindResult::Content(_, content) => {
Ok(VfsMetadata { file_type: VfsFileType::File, len: content.size.get(), mtime })
}
WadFindResult::Window(_, len) => {
Ok(VfsMetadata { file_type: VfsFileType::File, len, mtime })
}
}
} else {
Err(VfsError::NotFound)
}
}
}
#[derive(Clone)]
enum WadContentInner {
Stream(AesCbcStream),
Mapped(Cursor<Arc<[u8]>>),
}
#[derive(Clone)]
struct WadContent {
inner: WadContentInner,
mtime: Option<FileTime>,
}
impl WadContent {
fn new(inner: AesCbcStream, mtime: Option<FileTime>) -> Self {
Self { inner: WadContentInner::Stream(inner), mtime }
}
fn convert_to_mapped(&mut self) {
match &mut self.inner {
WadContentInner::Stream(stream) => {
let pos = stream.stream_position().unwrap();
stream.seek(SeekFrom::Start(0)).unwrap();
let mut data = vec![0u8; stream.len() as usize];
stream.read_exact(&mut data).unwrap();
let mut cursor = Cursor::new(Arc::from(data.as_slice()));
cursor.set_position(pos);
self.inner = WadContentInner::Mapped(cursor);
}
WadContentInner::Mapped(_) => {}
};
}
}
impl BufRead for WadContent {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
match &mut self.inner {
WadContentInner::Stream(stream) => stream.fill_buf(),
WadContentInner::Mapped(data) => data.fill_buf(),
}
}
fn consume(&mut self, amt: usize) {
match &mut self.inner {
WadContentInner::Stream(stream) => stream.consume(amt),
WadContentInner::Mapped(data) => data.consume(amt),
}
}
}
impl Read for WadContent {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match &mut self.inner {
WadContentInner::Stream(stream) => stream.read(buf),
WadContentInner::Mapped(data) => data.read(buf),
}
}
}
impl Seek for WadContent {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
match &mut self.inner {
WadContentInner::Stream(stream) => stream.seek(pos),
WadContentInner::Mapped(data) => data.seek(pos),
}
}
}
impl VfsFile for WadContent {
fn map(&mut self) -> io::Result<&[u8]> {
self.convert_to_mapped();
match &mut self.inner {
WadContentInner::Stream(_) => unreachable!(),
WadContentInner::Mapped(data) => Ok(data.get_ref()),
}
}
fn metadata(&mut self) -> io::Result<VfsMetadata> {
match &mut self.inner {
WadContentInner::Stream(stream) => Ok(VfsMetadata {
file_type: VfsFileType::File,
len: stream.len(),
mtime: self.mtime,
}),
WadContentInner::Mapped(data) => Ok(VfsMetadata {
file_type: VfsFileType::File,
len: data.get_ref().len() as u64,
mtime: self.mtime,
}),
}
}
fn into_disc_stream(self: Box<Self>) -> Box<dyn DiscStream> { self }
}
#[derive(Clone)]
struct AesCbcStream {
inner: Box<dyn VfsFile>,
position: u64,
content_offset: u64,
content_size: u64,
key: [u8; 0x10],
init_iv: [u8; 0x10],
last_iv: [u8; 0x10],
block_idx: u64,
block: Box<[u8; 0x200]>,
}
impl AesCbcStream {
fn new(
inner: Box<dyn VfsFile>,
content_offset: u64,
content_size: u64,
key: &[u8; 0x10],
iv: &[u8; 0x10],
) -> Self {
let block = <[u8; 0x200]>::new_box_zeroed().unwrap();
Self {
inner,
position: 0,
content_offset,
content_size,
key: *key,
init_iv: *iv,
last_iv: [0u8; 0x10],
block_idx: u64::MAX,
block,
}
}
#[inline]
fn len(&self) -> u64 { self.content_size }
#[inline]
fn remaining(&self) -> u64 { self.content_size.saturating_sub(self.position) }
}
impl Read for AesCbcStream {
fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
let mut total = 0;
while !buf.is_empty() {
let block = self.fill_buf()?;
if block.is_empty() {
break;
}
let len = buf.len().min(block.len());
buf[..len].copy_from_slice(&block[..len]);
buf = &mut buf[len..];
self.consume(len);
total += len;
}
Ok(total)
}
}
impl BufRead for AesCbcStream {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
if self.position >= self.content_size {
return Ok(&[]);
}
let block_size = self.block.len();
let current_block = self.position / block_size as u64;
if current_block != self.block_idx {
let block_offset = current_block * block_size as u64;
let mut iv = [0u8; 0x10];
if current_block == 0 {
// Use the initial IV for the first block
self.inner.seek(SeekFrom::Start(self.content_offset))?;
iv = self.init_iv;
} else if self.block_idx.checked_add(1) == Some(current_block) {
// Shortcut to avoid seeking when reading sequentially
iv = self.last_iv;
} else {
// Read the IV from the previous block
self.inner.seek(SeekFrom::Start(self.content_offset + block_offset - 0x10))?;
self.inner.read_exact(&mut iv)?;
}
let aligned_size = align_up(self.content_size, 0x10);
let remaining = aligned_size.saturating_sub(block_offset);
let read = remaining.min(block_size as u64) as usize;
self.inner.read_exact(&mut self.block[..read])?;
self.last_iv = *array_ref!(self.block, read - 0x10, 0x10);
let mut decryptor =
cbc::Decryptor::<aes::Aes128>::new((&self.key).into(), (&iv).into());
for aes_block in self.block[..read].chunks_exact_mut(0x10) {
decryptor.decrypt_block_mut(aes_block.into());
}
self.block_idx = current_block;
}
let offset = (self.position % block_size as u64) as usize;
let len = self.remaining().min((block_size - offset) as u64) as usize;
Ok(&self.block[offset..offset + len])
}
fn consume(&mut self, amt: usize) { self.position = self.position.saturating_add(amt as u64); }
}
impl Seek for AesCbcStream {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
self.position = match pos {
SeekFrom::Start(p) => p,
SeekFrom::End(p) => self.content_size.saturating_add_signed(p),
SeekFrom::Current(p) => self.position.saturating_add_signed(p),
};
Ok(self.position)
}
}