Clean up VFS error handling

Allows renaming, for example, local statics from `test$1234`
to `test` for inclusion in the source function.

Cargo.lock

@@ -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",

Cargo.toml

@@ -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"

README.md

@@ -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
+```

src/cmd/dol.rs

@@ -296,6 +296,8 @@ pub struct ModuleConfig {
 pub struct ExtractConfig {
     /// The name of the symbol to extract.
     pub symbol: String,
+    /// Optionally rename the output symbol. (e.g. symbol$1234 -> symbol)
+    pub rename: Option<String>,
     /// If specified, the symbol's data will be extracted to the given file.
     /// Path is relative to `out_dir/bin`.
     #[serde(with = "unix_path_serde_option", default, skip_serializing_if = "Option::is_none")]
@@ -385,6 +387,7 @@ pub struct OutputModule {
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 pub struct OutputExtract {
     pub symbol: String,
+    pub rename: Option<String>,
     #[serde(with = "unix_path_serde_option")]
     pub binary: Option<Utf8UnixPathBuf>,
     #[serde(with = "unix_path_serde_option")]
@@ -1014,7 +1017,8 @@ fn split_write_obj(
 
         if header_kind != HeaderKind::None {
             if let Some(header) = &extract.header {
-                let header_string = bin2c(symbol, section, data, header_kind);
+                let header_string =
+                    bin2c(symbol, section, data, header_kind, extract.rename.as_deref());
                 let out_path = base_dir.join("include").join(header.with_encoding());
                 if let Some(parent) = out_path.parent() {
                     DirBuilder::new().recursive(true).create(parent)?;
@@ -1026,6 +1030,7 @@ fn split_write_obj(
         // Copy to output config
         out_config.extract.push(OutputExtract {
             symbol: symbol.name.clone(),
+            rename: extract.rename.clone(),
             binary: extract.binary.clone(),
             header: extract.header.clone(),
             header_type: header_kind.to_string(),

src/cmd/mod.rs

@@ -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;

src/cmd/wad.rs

@@ -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,
+    })
+}

src/main.rs

@@ -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:?}");

src/util/bin2c.rs

@@ -50,15 +50,26 @@ impl fmt::Display for HeaderKind {
 }
 
 /// Converts a binary blob into a C array.
-pub fn bin2c(symbol: &ObjSymbol, section: &ObjSection, data: &[u8], kind: HeaderKind) -> String {
+pub fn bin2c(
+    symbol: &ObjSymbol,
+    section: &ObjSection,
+    data: &[u8],
+    kind: HeaderKind,
+    rename: Option<&str>,
+) -> String {
     match kind {
         HeaderKind::None => String::new(),
-        HeaderKind::Symbol => bin2c_symbol(symbol, section, data),
+        HeaderKind::Symbol => bin2c_symbol(symbol, section, data, rename),
         HeaderKind::Raw => bin2c_raw(data),
     }
 }
 
-fn bin2c_symbol(symbol: &ObjSymbol, section: &ObjSection, data: &[u8]) -> String {
+fn bin2c_symbol(
+    symbol: &ObjSymbol,
+    section: &ObjSection,
+    data: &[u8],
+    rename: Option<&str>,
+) -> String {
     let mut output = String::new();
     output.push_str(PROLOGUE);
     output.push_str(&format!(
@@ -72,7 +83,11 @@ fn bin2c_symbol(symbol: &ObjSymbol, section: &ObjSection, data: &[u8]) -> String
         output.push_str("const ");
     }
     output.push_str("unsigned char ");
-    output.push_str(symbol.demangled_name.as_deref().unwrap_or(symbol.name.as_str()));
+    if let Some(rename) = rename {
+        output.push_str(rename);
+    } else {
+        output.push_str(symbol.demangled_name.as_deref().unwrap_or(symbol.name.as_str()));
+    }
     output.push_str(&format!("[] ATTRIBUTE_ALIGN({}) = {{", symbol.align.unwrap_or(4)));
     for (i, byte) in data.iter().enumerate() {
         if i % 16 == 0 {

src/util/mod.rs

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

src/util/read.rs

@@ -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)
+}

src/util/wad.rs

@@ -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)
+}

src/vfs/common.rs

@@ -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 {

src/vfs/disc.rs

@@ -10,7 +10,7 @@ use nodtool::{
     nod::{Disc, DiscStream, Fst, NodeKind, OwnedFileStream, PartitionBase, PartitionMeta},
 };
 use typed_path::Utf8UnixPath;
-use zerocopy::IntoBytes;
+use zerocopy::{FromZeros, IntoBytes};
 
 use super::{
     next_non_empty, StaticFile, Vfs, VfsError, VfsFile, VfsFileType, VfsMetadata, VfsResult,
@@ -252,19 +252,21 @@ impl DiscFile {
         Self { inner: DiscFileInner::Stream(file), mtime }
     }
 
-    fn convert_to_mapped(&mut self) {
+    fn convert_to_mapped(&mut self) -> io::Result<()> {
         match &mut self.inner {
             DiscFileInner::Stream(stream) => {
-                let pos = stream.stream_position().unwrap();
+                let pos = stream.stream_position()?;
-                stream.seek(SeekFrom::Start(0)).unwrap();
+                stream.seek(SeekFrom::Start(0))?;
-                let mut data = vec![0u8; stream.len() as usize];
+                let mut data = <[u8]>::new_box_zeroed_with_elems(stream.len() as usize)
-                stream.read_exact(&mut data).unwrap();
+                    .map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
-                let mut cursor = Cursor::new(Arc::from(data.as_slice()));
+                stream.read_exact(&mut data)?;
+                let mut cursor = Cursor::new(Arc::from(data));
                 cursor.set_position(pos);
                 self.inner = DiscFileInner::Mapped(cursor);
             }
             DiscFileInner::Mapped(_) => {}
         };
+        Ok(())
     }
 }
 
@@ -304,7 +306,7 @@ impl Seek for DiscFile {
 
 impl VfsFile for DiscFile {
     fn map(&mut self) -> io::Result<&[u8]> {
-        self.convert_to_mapped();
+        self.convert_to_mapped()?;
         match &mut self.inner {
             DiscFileInner::Stream(_) => unreachable!(),
             DiscFileInner::Mapped(data) => Ok(data.get_ref()),

src/vfs/mod.rs

@@ -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)),
+        _ if magic.starts_with(&YAY0_MAGIC) => Ok(FileFormat::Compressed(CompressionKind::Yay0)),
-        YAZ0_MAGIC => Ok(FileFormat::Compressed(CompressionKind::Yaz0)),
+        _ if magic.starts_with(&YAZ0_MAGIC) => Ok(FileFormat::Compressed(CompressionKind::Yaz0)),
-        RARC_MAGIC => Ok(FileFormat::Archive(ArchiveKind::Rarc)),
+        _ if magic.starts_with(&RARC_MAGIC) => Ok(FileFormat::Archive(ArchiveKind::Rarc)),
-        U8_MAGIC => Ok(FileFormat::Archive(ArchiveKind::U8)),
+        _ 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)?)),
     }
 }
 

src/vfs/std_fs.rs

@@ -55,10 +55,10 @@ impl StdFile {
     pub fn new(path: Utf8NativePathBuf) -> Self { StdFile { path, file: None, mmap: None } }
 
     pub fn file(&mut self) -> io::Result<&mut BufReader<fs::File>> {
-        if self.file.is_none() {
+        Ok(match self.file {
-            self.file = Some(BufReader::new(fs::File::open(&self.path)?));
+            Some(ref mut file) => file,
-        }
+            None => self.file.insert(BufReader::new(fs::File::open(&self.path)?)),
-        Ok(self.file.as_mut().unwrap())
+        })
     }
 }
 
@@ -86,13 +86,15 @@ impl Seek for StdFile {
 
 impl VfsFile for StdFile {
     fn map(&mut self) -> io::Result<&[u8]> {
-        if self.file.is_none() {
+        let file = match self.file {
-            self.file = Some(BufReader::new(fs::File::open(&self.path)?));
+            Some(ref mut file) => file,
-        }
+            None => self.file.insert(BufReader::new(fs::File::open(&self.path)?)),
-        if self.mmap.is_none() {
+        };
-            self.mmap = Some(unsafe { memmap2::Mmap::map(self.file.as_ref().unwrap().get_ref())? });
+        let mmap = match self.mmap {
-        }
+            Some(ref mmap) => mmap,
-        Ok(self.mmap.as_ref().unwrap())
+            None => self.mmap.insert(unsafe { memmap2::Mmap::map(file.get_ref())? }),
+        };
+        Ok(mmap)
     }
 
     fn metadata(&mut self) -> io::Result<VfsMetadata> {

src/vfs/wad.rs

@@ -0,0 +1,366 @@
+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,
+    mtime: Option<FileTime>,
+}
+
+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 mtime = file.metadata()?.mtime;
+        let wad = process_wad(file.as_mut())
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+        Ok(Self { file, wad, mtime })
+    }
+
+    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.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.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> {
+        if let Some(result) = self.find(path.as_str()) {
+            match result {
+                WadFindResult::Root => {
+                    Ok(VfsMetadata { file_type: VfsFileType::Directory, len: 0, mtime: self.mtime })
+                }
+                WadFindResult::Static(data) => Ok(VfsMetadata {
+                    file_type: VfsFileType::File,
+                    len: data.len() as u64,
+                    mtime: self.mtime,
+                }),
+                WadFindResult::Content(_, content) => Ok(VfsMetadata {
+                    file_type: VfsFileType::File,
+                    len: content.size.get(),
+                    mtime: self.mtime,
+                }),
+                WadFindResult::Window(_, len) => {
+                    Ok(VfsMetadata { file_type: VfsFileType::File, len, mtime: self.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) -> io::Result<()> {
+        match &mut self.inner {
+            WadContentInner::Stream(stream) => {
+                let pos = stream.stream_position()?;
+                stream.seek(SeekFrom::Start(0))?;
+                let mut data = <[u8]>::new_box_zeroed_with_elems(stream.len() as usize)
+                    .map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
+                stream.read_exact(&mut data)?;
+                let mut cursor = Cursor::new(Arc::from(data));
+                cursor.set_position(pos);
+                self.inner = WadContentInner::Mapped(cursor);
+            }
+            WadContentInner::Mapped(_) => {}
+        };
+        Ok(())
+    }
+}
+
+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)
+    }
+}