From 7f97dac3994323886274c50de94090b1bfdc7617 Mon Sep 17 00:00:00 2001 From: Luke Street Date: Wed, 21 Feb 2024 00:04:23 -0700 Subject: [PATCH] Convert all formats to new BlockIO trait (WIP) Simplifies format handling by moving common logic into a new `DiscReader` type. --- Cargo.toml | 4 + README.md | 4 +- src/bin.rs | 222 +++++---- src/disc/gcn.rs | 31 +- src/disc/hashes.rs | 203 ++++++++ src/disc/mod.rs | 25 +- src/disc/partition.rs | 177 +++++++ src/disc/reader.rs | 273 ++++++++++ src/disc/wii.rs | 265 ++++++---- src/io/block.rs | 279 +++++++++++ src/io/ciso.rs | 251 +++------- src/io/iso.rs | 57 ++- src/io/mod.rs | 55 +-- src/io/nfs.rs | 236 +++------ src/io/nkit.rs | 53 +- src/io/split.rs | 53 +- src/io/wbfs.rs | 157 ++---- src/io/wia.rs | 852 ++++++++++++-------------------- src/lib.rs | 58 +-- src/streams.rs | 235 +-------- src/util/compress.rs | 28 +- src/util/lfg.rs | 3 +- src/util/mod.rs | 34 +- src/util/{reader.rs => read.rs} | 11 + 24 files changed, 1991 insertions(+), 1575 deletions(-) create mode 100644 src/disc/hashes.rs create mode 100644 src/disc/partition.rs create mode 100644 src/disc/reader.rs create mode 100644 src/io/block.rs rename src/util/{reader.rs => read.rs} (85%) diff --git a/Cargo.toml b/Cargo.toml index c356dfd..9be9b00 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,9 @@ build = "build.rs" name = "nodtool" path = "src/bin.rs" +[profile.release] +debug = true + [profile.release-lto] inherits = "release" lto = "thin" @@ -41,6 +44,7 @@ bzip2 = { version = "0.4.4", features = ["static"], optional = true } cbc = "0.1.2" crc32fast = "1.4.0" digest = "0.10.7" +dyn-clone = "1.0.16" enable-ansi-support = "0.2.1" encoding_rs = "0.8.33" file-size = "1.0.3" diff --git a/README.md b/README.md index 230badf..e8463bd 100644 --- a/README.md +++ b/README.md @@ -16,8 +16,8 @@ but does not currently support authoring. Currently supported file formats: - ISO (GCM) - WIA / RVZ -- WBFS -- CISO +- WBFS (+ NKit 2 lossless) +- CISO (+ NKit 2 lossless) - NFS (Wii U VC) ## CLI tool diff --git a/src/bin.rs b/src/bin.rs index d710108..516fd40 100644 --- a/src/bin.rs +++ b/src/bin.rs @@ -2,6 +2,7 @@ mod argp_version; use std::{ borrow::Cow, + cmp::min, env, error::Error, ffi::OsStr, @@ -26,7 +27,7 @@ use indicatif::{ProgressBar, ProgressState, ProgressStyle}; use itertools::Itertools; use nod::{ Disc, DiscHeader, Fst, Node, OpenOptions, PartitionBase, PartitionKind, PartitionMeta, Result, - ResultContext, + ResultContext, SECTOR_SIZE, }; use supports_color::Stream; use tracing::level_filters::LevelFilter; @@ -100,6 +101,9 @@ struct ConvertArgs { #[argp(positional)] /// output ISO file out: PathBuf, + #[argp(switch)] + /// enable MD5 hashing (slower) + md5: bool, } #[derive(FromArgs, Debug)] @@ -109,6 +113,9 @@ struct VerifyArgs { #[argp(positional)] /// path to disc image file: PathBuf, + #[argp(switch)] + /// enable MD5 hashing (slower) + md5: bool, } #[derive(Debug, Eq, PartialEq, Copy, Clone)] @@ -243,55 +250,60 @@ fn info(args: InfoArgs) -> Result<()> { if header.is_wii() { for (idx, info) in disc.partitions().iter().enumerate() { println!(); - println!("Partition {}:{}", info.group_index, info.part_index); + println!("Partition {}", idx); println!("\tType: {}", info.kind); - println!("\tPartition offset: {:#X}", info.part_offset); + let offset = info.start_sector as u64 * SECTOR_SIZE as u64; + println!("\tStart sector: {} (offset {:#X})", info.start_sector, offset); + let data_size = + (info.data_end_sector - info.data_start_sector) as u64 * SECTOR_SIZE as u64; println!( "\tData offset / size: {:#X} / {:#X} ({})", - info.part_offset + info.data_offset, - info.data_size, - file_size::fit_4(info.data_size) + info.data_start_sector as u64 * SECTOR_SIZE as u64, + data_size, + file_size::fit_4(data_size) + ); + println!( + "\tTMD offset / size: {:#X} / {:#X}", + offset + info.header.tmd_off(), + info.header.tmd_size() + ); + println!( + "\tCert offset / size: {:#X} / {:#X}", + offset + info.header.cert_chain_off(), + info.header.cert_chain_size() + ); + println!( + "\tH3 offset / size: {:#X} / {:#X}", + offset + info.header.h3_table_off(), + info.header.h3_table_size() ); - if let Some(header) = &info.header { - println!( - "\tTMD offset / size: {:#X} / {:#X}", - info.part_offset + header.tmd_off(), - header.tmd_size() - ); - println!( - "\tCert offset / size: {:#X} / {:#X}", - info.part_offset + header.cert_chain_off(), - header.cert_chain_size() - ); - println!( - "\tH3 offset / size: {:#X} / {:#X}", - info.part_offset + header.h3_table_off(), - header.h3_table_size() - ); - } - let mut partition = disc.open_partition(idx)?; - let meta = partition.meta()?; - let header = meta.header(); - let tmd = meta.tmd_header(); - let title_id_str = if let Some(tmd) = tmd { - format!( - "{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}", - tmd.title_id[0], - tmd.title_id[1], - tmd.title_id[2], - tmd.title_id[3], - tmd.title_id[4], - tmd.title_id[5], - tmd.title_id[6], - tmd.title_id[7] - ) - } else { - "N/A".to_string() - }; - println!("\tName: {}", header.game_title_str()); - println!("\tGame ID: {} ({})", header.game_id_str(), title_id_str); - println!("\tDisc {}, Revision {}", header.disc_num + 1, header.disc_version); + // let mut partition = disc.open_partition(idx)?; + // let meta = partition.meta()?; + // let header = meta.header(); + // let tmd = meta.tmd_header(); + // let title_id_str = if let Some(tmd) = tmd { + // format!( + // "{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}", + // tmd.title_id[0], + // tmd.title_id[1], + // tmd.title_id[2], + // tmd.title_id[3], + // tmd.title_id[4], + // tmd.title_id[5], + // tmd.title_id[6], + // tmd.title_id[7] + // ) + // } else { + let title_id_str = "N/A".to_string(); + // }; + println!("\tName: {}", info.disc_header.game_title_str()); + println!("\tGame ID: {} ({})", info.disc_header.game_id_str(), title_id_str); + println!( + "\tDisc {}, Revision {}", + info.disc_header.disc_num + 1, + info.disc_header.disc_version + ); } } else if header.is_gamecube() { // TODO @@ -305,13 +317,15 @@ fn info(args: InfoArgs) -> Result<()> { Ok(()) } -fn convert(args: ConvertArgs) -> Result<()> { convert_and_verify(&args.file, Some(&args.out)) } +fn convert(args: ConvertArgs) -> Result<()> { + convert_and_verify(&args.file, Some(&args.out), args.md5) +} -fn verify(args: VerifyArgs) -> Result<()> { convert_and_verify(&args.file, None) } +fn verify(args: VerifyArgs) -> Result<()> { convert_and_verify(&args.file, None, args.md5) } -fn convert_and_verify(in_file: &Path, out_file: Option<&Path>) -> Result<()> { +fn convert_and_verify(in_file: &Path, out_file: Option<&Path>, md5: bool) -> Result<()> { println!("Loading {}", in_file.display()); - let disc = Disc::new_with_options(in_file, &OpenOptions { + let mut disc = Disc::new_with_options(in_file, &OpenOptions { rebuild_hashes: true, validate_hashes: false, rebuild_encryption: true, @@ -320,7 +334,7 @@ fn convert_and_verify(in_file: &Path, out_file: Option<&Path>) -> Result<()> { print_header(header); let meta = disc.meta()?; - let mut stream = disc.open()?.take(disc.disc_size()); + let disc_size = disc.disc_size(); let mut file = if let Some(out_file) = out_file { Some( @@ -332,7 +346,7 @@ fn convert_and_verify(in_file: &Path, out_file: Option<&Path>) -> Result<()> { }; println!("\nHashing..."); - let pb = ProgressBar::new(stream.limit()); + let pb = ProgressBar::new(disc_size); pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})") .unwrap() .with_key("eta", |state: &ProgressState, w: &mut dyn std::fmt::Write| { @@ -341,12 +355,20 @@ fn convert_and_verify(in_file: &Path, out_file: Option<&Path>) -> Result<()> { .progress_chars("#>-")); const BUFFER_SIZE: usize = 1015808; // LCM(0x8000, 0x7C00) - let digest_threads = [ - digest_thread::(), - digest_thread::(), - digest_thread::(), - digest_thread::(), - ]; + let digest_threads = if md5 { + vec![ + digest_thread::(), + digest_thread::(), + digest_thread::(), + digest_thread::(), + ] + } else { + vec![ + digest_thread::(), + digest_thread::(), + digest_thread::(), + ] + }; let (w_tx, w_rx) = sync_channel::>(1); let w_thread = thread::spawn(move || { @@ -370,13 +392,11 @@ fn convert_and_verify(in_file: &Path, out_file: Option<&Path>) -> Result<()> { let mut total_read = 0u64; let mut buf = ::new_box_slice_zeroed(BUFFER_SIZE); - loop { - let read = stream.read(buf.as_mut()).with_context(|| { + while total_read < disc_size { + let read = min(BUFFER_SIZE as u64, disc_size - total_read) as usize; + disc.reader.read_exact(&mut buf[..read]).with_context(|| { format!("Reading {} bytes at disc offset {}", BUFFER_SIZE, total_read) })?; - if read == 0 { - break; - } let arc = Arc::<[u8]>::from(&buf[..read]); for (tx, _) in &digest_threads { @@ -394,7 +414,7 @@ fn convert_and_verify(in_file: &Path, out_file: Option<&Path>) -> Result<()> { } println!(); - for (tx, handle) in digest_threads.into_iter() { + for (tx, handle) in digest_threads { drop(tx); // Close channel match handle.join().unwrap() { DigestResult::Crc32(crc) => { @@ -475,48 +495,48 @@ fn extract(args: ExtractArgs) -> Result<()> { rebuild_encryption: false, })?; let is_wii = disc.header().is_wii(); - let mut partition = disc.open_partition_kind(PartitionKind::Data)?; - let meta = partition.meta()?; - extract_sys_files(meta.as_ref(), &output_dir.join("sys"), args.quiet)?; - - // Extract FST - let files_dir = output_dir.join("files"); - let fst = Fst::new(&meta.raw_fst)?; - let mut path_segments = Vec::<(Cow, usize)>::new(); - for (idx, node, name) in fst.iter() { - // Remove ended path segments - let mut new_size = 0; - for (_, end) in path_segments.iter() { - if *end == idx { - break; - } - new_size += 1; - } - path_segments.truncate(new_size); - - // Add the new path segment - let end = if node.is_dir() { node.length(false) as usize } else { idx + 1 }; - path_segments.push((name?, end)); - - let path = path_segments.iter().map(|(name, _)| name.as_ref()).join("/"); - if node.is_dir() { - fs::create_dir_all(files_dir.join(&path)) - .with_context(|| format!("Creating directory {}", path))?; - } else { - extract_node(node, partition.as_mut(), &files_dir, &path, is_wii, args.quiet)?; - } - } + // let mut partition = disc.open_partition_kind(PartitionKind::Data)?; + // let meta = partition.meta()?; + // extract_sys_files(meta.as_ref(), &output_dir.join("sys"), args.quiet)?; + // + // // Extract FST + // let files_dir = output_dir.join("files"); + // let fst = Fst::new(&meta.raw_fst)?; + // let mut path_segments = Vec::<(Cow, usize)>::new(); + // for (idx, node, name) in fst.iter() { + // // Remove ended path segments + // let mut new_size = 0; + // for (_, end) in path_segments.iter() { + // if *end == idx { + // break; + // } + // new_size += 1; + // } + // path_segments.truncate(new_size); + // + // // Add the new path segment + // let end = if node.is_dir() { node.length(false) as usize } else { idx + 1 }; + // path_segments.push((name?, end)); + // + // let path = path_segments.iter().map(|(name, _)| name.as_ref()).join("/"); + // if node.is_dir() { + // fs::create_dir_all(files_dir.join(&path)) + // .with_context(|| format!("Creating directory {}", path))?; + // } else { + // extract_node(node, partition.as_mut(), &files_dir, &path, is_wii, args.quiet)?; + // } + // } Ok(()) } fn extract_sys_files(data: &PartitionMeta, out_dir: &Path, quiet: bool) -> Result<()> { fs::create_dir_all(out_dir) .with_context(|| format!("Creating output directory {}", out_dir.display()))?; - extract_file(&data.raw_boot, &out_dir.join("boot.bin"), quiet)?; - extract_file(&data.raw_bi2, &out_dir.join("bi2.bin"), quiet)?; - extract_file(&data.raw_apploader, &out_dir.join("apploader.img"), quiet)?; - extract_file(&data.raw_fst, &out_dir.join("fst.bin"), quiet)?; - extract_file(&data.raw_dol, &out_dir.join("main.dol"), quiet)?; + extract_file(data.raw_boot.as_ref(), &out_dir.join("boot.bin"), quiet)?; + extract_file(data.raw_bi2.as_ref(), &out_dir.join("bi2.bin"), quiet)?; + extract_file(data.raw_apploader.as_ref(), &out_dir.join("apploader.img"), quiet)?; + extract_file(data.raw_fst.as_ref(), &out_dir.join("fst.bin"), quiet)?; + extract_file(data.raw_dol.as_ref(), &out_dir.join("main.dol"), quiet)?; Ok(()) } @@ -564,7 +584,9 @@ fn extract_node( Ok(()) } -fn digest_thread() -> (SyncSender>, JoinHandle) +type DigestThread = (SyncSender>, JoinHandle); + +fn digest_thread() -> DigestThread where H: Hasher + Send + 'static { let (tx, rx) = sync_channel::>(1); let handle = thread::spawn(move || { diff --git a/src/disc/gcn.rs b/src/disc/gcn.rs index 819d46e..e5401d7 100644 --- a/src/disc/gcn.rs +++ b/src/disc/gcn.rs @@ -17,7 +17,7 @@ use crate::{ streams::{ReadStream, SharedWindowedReadStream}, util::{ div_rem, - reader::{read_from, read_vec}, + read::{read_box, read_box_slice, read_vec}, }, Error, OpenOptions, Result, ResultContext, }; @@ -25,7 +25,6 @@ use crate::{ pub(crate) struct DiscGCN { pub(crate) header: DiscHeader, pub(crate) disc_size: u64, - // pub(crate) junk_start: u64, } impl DiscGCN { @@ -34,10 +33,7 @@ impl DiscGCN { header: DiscHeader, disc_size: Option, ) -> Result { - // stream.seek(SeekFrom::Start(size_of::() as u64)).context("Seeking to partition header")?; - // let partition_header: PartitionHeader = read_from(stream).context("Reading partition header")?; - // let junk_start = partition_header.fst_off(false) + partition_header.fst_sz(false); - Ok(DiscGCN { header, disc_size: disc_size.unwrap_or(MINI_DVD_SIZE) /*, junk_start*/ }) + Ok(DiscGCN { header, disc_size: disc_size.unwrap_or(MINI_DVD_SIZE) }) } } @@ -140,7 +136,12 @@ impl<'a> Seek for PartitionGC<'a> { fn seek(&mut self, pos: SeekFrom) -> io::Result { self.offset = match pos { SeekFrom::Start(v) => v, - SeekFrom::End(v) => self.stable_stream_len()?.saturating_add_signed(v), + SeekFrom::End(_) => { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "PartitionGC: SeekFrom::End is not supported", + )); + } SeekFrom::Current(v) => self.offset.saturating_add_signed(v), }; let block = self.offset / SECTOR_SIZE as u64; @@ -154,12 +155,6 @@ impl<'a> Seek for PartitionGC<'a> { fn stream_position(&mut self) -> io::Result { Ok(self.offset) } } -impl<'a> ReadStream for PartitionGC<'a> { - fn stable_stream_len(&mut self) -> io::Result { self.stream.stable_stream_len() } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - impl<'a> PartitionBase for PartitionGC<'a> { fn meta(&mut self) -> Result> { self.seek(SeekFrom::Start(0)).context("Seeking to partition header")?; @@ -177,11 +172,11 @@ impl<'a> PartitionBase for PartitionGC<'a> { pub(crate) fn read_part_header(reader: &mut R, is_wii: bool) -> Result> where R: Read + Seek + ?Sized { // boot.bin - let raw_boot: [u8; BOOT_SIZE] = read_from(reader).context("Reading boot.bin")?; + let raw_boot: Box<[u8; BOOT_SIZE]> = read_box(reader).context("Reading boot.bin")?; let partition_header = PartitionHeader::ref_from(&raw_boot[size_of::()..]).unwrap(); // bi2.bin - let raw_bi2: [u8; BI2_SIZE] = read_from(reader).context("Reading bi2.bin")?; + let raw_bi2: Box<[u8; BI2_SIZE]> = read_box(reader).context("Reading bi2.bin")?; // apploader.bin let mut raw_apploader: Vec = @@ -201,7 +196,7 @@ where R: Read + Seek + ?Sized { reader .seek(SeekFrom::Start(partition_header.fst_off(is_wii))) .context("Seeking to FST offset")?; - let raw_fst: Vec = read_vec(reader, partition_header.fst_sz(is_wii) as usize) + let raw_fst: Box<[u8]> = read_box_slice(reader, partition_header.fst_sz(is_wii) as usize) .with_context(|| { format!( "Reading partition FST (offset {}, size {})", @@ -236,9 +231,9 @@ where R: Read + Seek + ?Sized { Ok(Box::new(PartitionMeta { raw_boot, raw_bi2, - raw_apploader, + raw_apploader: raw_apploader.into_boxed_slice(), raw_fst, - raw_dol, + raw_dol: raw_dol.into_boxed_slice(), raw_ticket: None, raw_tmd: None, raw_cert_chain: None, diff --git a/src/disc/hashes.rs b/src/disc/hashes.rs new file mode 100644 index 0000000..2f44d67 --- /dev/null +++ b/src/disc/hashes.rs @@ -0,0 +1,203 @@ +use std::{ + io::{Read, Seek, SeekFrom}, + sync::{Arc, Mutex}, + time::Instant, +}; + +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use sha1::{Digest, Sha1}; +use zerocopy::FromZeroes; + +use crate::{ + array_ref, array_ref_mut, + disc::{ + partition::PartitionReader, + reader::DiscReader, + wii::{HASHES_SIZE, SECTOR_DATA_SIZE}, + }, + io::HashBytes, + util::read::read_box_slice, + Result, ResultContext, SECTOR_SIZE, +}; + +/// In a sector, following the 0x400 byte block of hashes, each 0x400 bytes of decrypted data is +/// hashed, yielding 31 H0 hashes. +/// Then, 8 sectors are aggregated into a subgroup, and the 31 H0 hashes for each sector are hashed, +/// yielding 8 H1 hashes. +/// Then, 8 subgroups are aggregated into a group, and the 8 H1 hashes for each subgroup are hashed, +/// yielding 8 H2 hashes. +/// Finally, the 8 H2 hashes for each group are hashed, yielding 1 H3 hash. +/// The H3 hashes for each group are stored in the partition's H3 table. +#[derive(Clone, Debug)] +pub struct HashTable { + /// SHA-1 hash of each 0x400 byte block of decrypted data. + pub h0_hashes: Box<[HashBytes]>, + /// SHA-1 hash of the 31 H0 hashes for each sector. + pub h1_hashes: Box<[HashBytes]>, + /// SHA-1 hash of the 8 H1 hashes for each subgroup. + pub h2_hashes: Box<[HashBytes]>, + /// SHA-1 hash of the 8 H2 hashes for each group. + pub h3_hashes: Box<[HashBytes]>, +} + +#[derive(Clone, FromZeroes)] +struct HashResult { + h0_hashes: [HashBytes; 1984], + h1_hashes: [HashBytes; 64], + h2_hashes: [HashBytes; 8], + h3_hash: HashBytes, +} + +impl HashTable { + fn new(num_sectors: u32) -> Self { + let num_sectors = num_sectors.next_multiple_of(64) as usize; + let num_data_hashes = num_sectors * 31; + let num_subgroups = num_sectors / 8; + let num_groups = num_subgroups / 8; + Self { + h0_hashes: HashBytes::new_box_slice_zeroed(num_data_hashes), + h1_hashes: HashBytes::new_box_slice_zeroed(num_sectors), + h2_hashes: HashBytes::new_box_slice_zeroed(num_subgroups), + h3_hashes: HashBytes::new_box_slice_zeroed(num_groups), + } + } + + fn extend(&mut self, group_index: usize, result: &HashResult) { + *array_ref_mut![self.h0_hashes, group_index * 1984, 1984] = result.h0_hashes; + *array_ref_mut![self.h1_hashes, group_index * 64, 64] = result.h1_hashes; + *array_ref_mut![self.h2_hashes, group_index * 8, 8] = result.h2_hashes; + self.h3_hashes[group_index] = result.h3_hash; + } +} + +pub fn rebuild_hashes(reader: &mut DiscReader) -> Result<()> { + const NUM_H0_HASHES: usize = SECTOR_DATA_SIZE / HASHES_SIZE; + + log::info!( + "Rebuilding hashes for Wii partition data (using {} threads)", + rayon::current_num_threads() + ); + + let start = Instant::now(); + + // Precompute hashes for zeroed sectors. + const ZERO_H0_BYTES: &[u8] = &[0u8; HASHES_SIZE]; + let zero_h0_hash = hash_bytes(ZERO_H0_BYTES); + let mut zero_h1_hash = Sha1::new(); + for _ in 0..NUM_H0_HASHES { + zero_h1_hash.update(zero_h0_hash); + } + + let mut hash_tables = Vec::with_capacity(reader.partitions.len()); + for part in &reader.partitions { + let part_sectors = part.data_end_sector - part.data_start_sector; + let hash_table = HashTable::new(part_sectors); + log::debug!( + "Rebuilding hashes: {} sectors, {} subgroups, {} groups", + hash_table.h1_hashes.len(), + hash_table.h2_hashes.len(), + hash_table.h3_hashes.len() + ); + + let group_count = hash_table.h3_hashes.len(); + let mutex = Arc::new(Mutex::new(hash_table)); + (0..group_count).into_par_iter().try_for_each_with( + (PartitionReader::new(reader.io.clone(), part)?, mutex.clone()), + |(stream, mutex), h3_index| -> Result<()> { + let mut result = HashResult::new_box_zeroed(); + let mut data_buf = ::new_box_slice_zeroed(SECTOR_DATA_SIZE); + let mut h3_hasher = Sha1::new(); + for h2_index in 0..8 { + let mut h2_hasher = Sha1::new(); + for h1_index in 0..8 { + let sector = h1_index + h2_index * 8; + let part_sector = sector as u32 + h3_index as u32 * 64; + let mut h1_hasher = Sha1::new(); + if part_sector >= part_sectors { + for h0_index in 0..NUM_H0_HASHES { + result.h0_hashes[h0_index + sector * 31] = zero_h0_hash; + h1_hasher.update(zero_h0_hash); + } + } else { + stream + .seek(SeekFrom::Start(part_sector as u64 * SECTOR_DATA_SIZE as u64)) + .with_context(|| format!("Seeking to sector {}", part_sector))?; + stream + .read_exact(&mut data_buf) + .with_context(|| format!("Reading sector {}", part_sector))?; + for h0_index in 0..NUM_H0_HASHES { + let h0_hash = hash_bytes(array_ref![ + data_buf, + h0_index * HASHES_SIZE, + HASHES_SIZE + ]); + result.h0_hashes[h0_index + sector * 31] = h0_hash; + h1_hasher.update(h0_hash); + } + }; + let h1_hash = h1_hasher.finalize().into(); + result.h1_hashes[sector] = h1_hash; + h2_hasher.update(h1_hash); + } + let h2_hash = h2_hasher.finalize().into(); + result.h2_hashes[h2_index] = h2_hash; + h3_hasher.update(h2_hash); + } + result.h3_hash = h3_hasher.finalize().into(); + let mut hash_table = mutex.lock().map_err(|_| "Failed to lock mutex")?; + hash_table.extend(h3_index, &result); + Ok(()) + }, + )?; + + let hash_table = Arc::try_unwrap(mutex) + .map_err(|_| "Failed to unwrap Arc")? + .into_inner() + .map_err(|_| "Failed to lock mutex")?; + hash_tables.push(hash_table); + } + + // Verify against H3 table + for (part, hash_table) in reader.partitions.clone().iter().zip(hash_tables.iter()) { + log::debug!( + "Verifying H3 table for partition {} (count {})", + part.index, + hash_table.h3_hashes.len() + ); + reader + .seek(SeekFrom::Start( + part.start_sector as u64 * SECTOR_SIZE as u64 + part.header.h3_table_off(), + )) + .context("Seeking to H3 table")?; + let h3_table: Box<[HashBytes]> = + read_box_slice(reader, hash_table.h3_hashes.len()).context("Reading H3 table")?; + for (idx, (expected_hash, h3_hash)) in + h3_table.iter().zip(hash_table.h3_hashes.iter()).enumerate() + { + if expected_hash != h3_hash { + let mut got_bytes = [0u8; 40]; + let got = base16ct::lower::encode_str(h3_hash, &mut got_bytes).unwrap(); + let mut expected_bytes = [0u8; 40]; + let expected = + base16ct::lower::encode_str(expected_hash, &mut expected_bytes).unwrap(); + log::warn!( + "Partition {} H3 table does not match:\n\tindex {}\n\texpected: {}\n\tgot: {}", + part.index, idx, expected, got + ); + } + } + } + + for (part, hash_table) in reader.partitions.iter_mut().zip(hash_tables) { + part.hash_table = Some(hash_table); + } + log::info!("Rebuilt hashes in {:?}", start.elapsed()); + Ok(()) +} + +#[inline] +fn hash_bytes(buf: &[u8]) -> HashBytes { + let mut hasher = Sha1::new(); + hasher.update(buf); + hasher.finalize().into() +} diff --git a/src/disc/mod.rs b/src/disc/mod.rs index 0cd3182..de234a2 100644 --- a/src/disc/mod.rs +++ b/src/disc/mod.rs @@ -20,14 +20,17 @@ use crate::{ io::DiscIO, static_assert, streams::{ReadStream, SharedWindowedReadStream}, - util::reader::read_from, + util::read::read_from, Error, Fst, OpenOptions, Result, ResultContext, }; pub(crate) mod gcn; +pub(crate) mod hashes; +pub mod partition; +pub mod reader; pub(crate) mod wii; -pub(crate) const SECTOR_SIZE: usize = 0x8000; +pub const SECTOR_SIZE: usize = 0x8000; /// Shared GameCube & Wii disc header #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)] @@ -372,23 +375,23 @@ pub const BI2_SIZE: usize = 0x2000; #[derive(Clone, Debug)] pub struct PartitionMeta { /// Disc and partition header (boot.bin) - pub raw_boot: [u8; BOOT_SIZE], + pub raw_boot: Box<[u8; BOOT_SIZE]>, /// Debug and region information (bi2.bin) - pub raw_bi2: [u8; BI2_SIZE], + pub raw_bi2: Box<[u8; BI2_SIZE]>, /// Apploader (apploader.bin) - pub raw_apploader: Vec, + pub raw_apploader: Box<[u8]>, /// File system table (fst.bin) - pub raw_fst: Vec, + pub raw_fst: Box<[u8]>, /// Main binary (main.dol) - pub raw_dol: Vec, + pub raw_dol: Box<[u8]>, /// Ticket (ticket.bin, Wii only) - pub raw_ticket: Option>, + pub raw_ticket: Option>, /// TMD (tmd.bin, Wii only) - pub raw_tmd: Option>, + pub raw_tmd: Option>, /// Certificate chain (cert.bin, Wii only) - pub raw_cert_chain: Option>, + pub raw_cert_chain: Option>, /// H3 hash table (h3.bin, Wii only) - pub raw_h3_table: Option>, + pub raw_h3_table: Option>, } impl PartitionMeta { diff --git a/src/disc/partition.rs b/src/disc/partition.rs new file mode 100644 index 0000000..3ef43ef --- /dev/null +++ b/src/disc/partition.rs @@ -0,0 +1,177 @@ +use std::{ + cmp::min, + io, + io::{Read, Seek, SeekFrom}, +}; + +use sha1::{Digest, Sha1}; +use zerocopy::FromZeroes; + +use crate::{ + array_ref, + disc::wii::{as_digest, HASHES_SIZE, SECTOR_DATA_SIZE}, + io::block::{BPartitionInfo, Block, BlockIO}, + util::div_rem, + Result, SECTOR_SIZE, +}; + +pub struct PartitionReader { + io: Box, + partition: BPartitionInfo, + block: Option, + block_buf: Box<[u8]>, + block_idx: u32, + sector_buf: Box<[u8; SECTOR_SIZE]>, + sector: u32, + pos: u64, + verify: bool, +} + +impl Clone for PartitionReader { + fn clone(&self) -> Self { + Self { + io: self.io.clone(), + partition: self.partition.clone(), + block: None, + block_buf: ::new_box_slice_zeroed(self.block_buf.len()), + block_idx: u32::MAX, + sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed(), + sector: u32::MAX, + pos: 0, + verify: self.verify, + } + } +} + +impl PartitionReader { + pub fn new(inner: Box, partition: &BPartitionInfo) -> Result { + let block_size = inner.block_size(); + Ok(Self { + io: inner, + partition: partition.clone(), + block: None, + block_buf: ::new_box_slice_zeroed(block_size as usize), + block_idx: u32::MAX, + sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed(), + sector: u32::MAX, + pos: 0, + verify: false, + }) + } +} + +impl Read for PartitionReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let partition_sector = (self.pos / SECTOR_DATA_SIZE as u64) as u32; + let sector = self.partition.data_start_sector + partition_sector; + if sector >= self.partition.data_end_sector { + return Ok(0); + } + let block_idx = (sector as u64 * SECTOR_SIZE as u64 / self.block_buf.len() as u64) as u32; + + // Read new block if necessary + if block_idx != self.block_idx { + self.block = + self.io.read_block(self.block_buf.as_mut(), block_idx, Some(&self.partition))?; + self.block_idx = block_idx; + } + + // Decrypt sector if necessary + if sector != self.sector { + let Some(block) = &self.block else { + return Ok(0); + }; + block.decrypt( + &mut self.sector_buf, + self.block_buf.as_ref(), + block_idx, + sector, + &self.partition, + )?; + + if self.verify { + verify_hashes(&self.sector_buf, sector)?; + } + + self.sector = sector; + } + + let offset = (self.pos % SECTOR_DATA_SIZE as u64) as usize; + let len = min(buf.len(), SECTOR_DATA_SIZE - offset); + buf[..len] + .copy_from_slice(&self.sector_buf[HASHES_SIZE + offset..HASHES_SIZE + offset + len]); + self.pos += len as u64; + Ok(len) + } +} + +impl Seek for PartitionReader { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + self.pos = match pos { + SeekFrom::Start(v) => v, + SeekFrom::End(_) => { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "PartitionReader: SeekFrom::End is not supported".to_string(), + )); + } + SeekFrom::Current(v) => self.pos.saturating_add_signed(v), + }; + Ok(self.pos) + } +} + +fn verify_hashes(buf: &[u8; SECTOR_SIZE], sector: u32) -> io::Result<()> { + let (mut group, sub_group) = div_rem(sector as usize, 8); + group %= 8; + + // H0 hashes + for i in 0..31 { + let mut hash = Sha1::new(); + hash.update(array_ref![buf, (i + 1) * 0x400, 0x400]); + let expected = as_digest(array_ref![buf, i * 20, 20]); + let output = hash.finalize(); + if output != expected { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid H0 hash! (block {:?}) {:x}\n\texpected {:x}", i, output, expected), + )); + } + } + + // H1 hash + { + let mut hash = Sha1::new(); + hash.update(array_ref![buf, 0, 0x26C]); + let expected = as_digest(array_ref![buf, 0x280 + sub_group * 20, 20]); + let output = hash.finalize(); + if output != expected { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Invalid H1 hash! (subgroup {:?}) {:x}\n\texpected {:x}", + sub_group, output, expected + ), + )); + } + } + + // H2 hash + { + let mut hash = Sha1::new(); + hash.update(array_ref![buf, 0x280, 0xA0]); + let expected = as_digest(array_ref![buf, 0x340 + group * 20, 20]); + let output = hash.finalize(); + if output != expected { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Invalid H2 hash! (group {:?}) {:x}\n\texpected {:x}", + group, output, expected + ), + )); + } + } + // TODO H3 hash + Ok(()) +} diff --git a/src/disc/reader.rs b/src/disc/reader.rs new file mode 100644 index 0000000..50c8747 --- /dev/null +++ b/src/disc/reader.rs @@ -0,0 +1,273 @@ +use std::{ + cmp::min, + io, + io::{Read, Seek, SeekFrom}, +}; + +use zerocopy::FromZeroes; + +use crate::{ + disc::{ + hashes::{rebuild_hashes, HashTable}, + partition::PartitionReader, + wii::{WiiPartEntry, WiiPartGroup, WiiPartitionHeader, WII_PART_GROUP_OFF}, + DL_DVD_SIZE, MINI_DVD_SIZE, SL_DVD_SIZE, + }, + io::block::{BPartitionInfo, Block, BlockIO}, + util::read::{read_box, read_from, read_vec}, + DiscHeader, Error, PartitionHeader, PartitionKind, Result, ResultContext, SECTOR_SIZE, +}; + +#[derive(Debug, Eq, PartialEq, Copy, Clone)] +pub enum EncryptionMode { + Encrypted, + Decrypted, +} + +pub struct DiscReader { + pub(crate) io: Box, + block: Option, + block_buf: Box<[u8]>, + block_idx: u32, + sector_buf: Box<[u8; SECTOR_SIZE]>, + sector_idx: u32, + pos: u64, + mode: EncryptionMode, + disc_header: Box, + pub(crate) partitions: Vec, + hash_tables: Vec, +} + +impl Clone for DiscReader { + fn clone(&self) -> Self { + Self { + io: self.io.clone(), + block: None, + block_buf: ::new_box_slice_zeroed(self.block_buf.len()), + block_idx: u32::MAX, + sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed(), + sector_idx: u32::MAX, + pos: 0, + mode: self.mode, + disc_header: self.disc_header.clone(), + partitions: self.partitions.clone(), + hash_tables: self.hash_tables.clone(), + } + } +} + +impl DiscReader { + pub fn new(inner: Box, mode: EncryptionMode) -> Result { + let block_size = inner.block_size(); + let meta = inner.meta()?; + let mut reader = Self { + io: inner, + block: None, + block_buf: ::new_box_slice_zeroed(block_size as usize), + block_idx: u32::MAX, + sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed(), + sector_idx: u32::MAX, + pos: 0, + mode, + disc_header: DiscHeader::new_box_zeroed(), + partitions: vec![], + hash_tables: vec![], + }; + let disc_header: Box = read_box(&mut reader).context("Reading disc header")?; + reader.disc_header = disc_header; + if reader.disc_header.is_wii() { + reader.partitions = read_partition_info(&mut reader)?; + // Rebuild hashes if the format requires it + if mode == EncryptionMode::Encrypted && meta.needs_hash_recovery { + rebuild_hashes(&mut reader)?; + } + } + reader.reset(); + Ok(reader) + } + + pub fn reset(&mut self) { + self.block = None; + self.block_buf.fill(0); + self.block_idx = u32::MAX; + self.sector_buf.fill(0); + self.sector_idx = u32::MAX; + self.pos = 0; + } + + pub fn disc_size(&self) -> u64 { + self.io + .meta() + .ok() + .and_then(|m| m.disc_size) + .unwrap_or_else(|| guess_disc_size(&self.partitions)) + } + + pub fn header(&self) -> &DiscHeader { &self.disc_header } + + pub fn partitions(&self) -> &[BPartitionInfo] { &self.partitions } +} + +impl Read for DiscReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let block_idx = (self.pos / self.block_buf.len() as u64) as u32; + let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32; + + let partition = if self.disc_header.is_wii() { + self.partitions.iter().find(|part| { + abs_sector >= part.data_start_sector && abs_sector < part.data_end_sector + }) + } else { + None + }; + + // Read new block + if block_idx != self.block_idx { + self.block = self.io.read_block(self.block_buf.as_mut(), block_idx, partition)?; + self.block_idx = block_idx; + } + + // Read new sector into buffer + if abs_sector != self.sector_idx { + let Some(block) = &self.block else { + return Ok(0); + }; + if let Some(partition) = partition { + match self.mode { + EncryptionMode::Decrypted => block.decrypt( + &mut self.sector_buf, + self.block_buf.as_ref(), + block_idx, + abs_sector, + partition, + )?, + EncryptionMode::Encrypted => block.encrypt( + &mut self.sector_buf, + self.block_buf.as_ref(), + block_idx, + abs_sector, + partition, + )?, + } + } else { + block.copy_raw( + &mut self.sector_buf, + self.block_buf.as_ref(), + block_idx, + abs_sector, + &self.disc_header, + )?; + } + self.sector_idx = abs_sector; + } + + // Read from sector buffer + let offset = (self.pos % SECTOR_SIZE as u64) as usize; + let len = min(buf.len(), SECTOR_SIZE - offset); + buf[..len].copy_from_slice(&self.sector_buf[offset..offset + len]); + self.pos += len as u64; + Ok(len) + } +} + +impl Seek for DiscReader { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + self.pos = match pos { + SeekFrom::Start(v) => v, + SeekFrom::End(_) => { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "BlockIOReader: SeekFrom::End is not supported".to_string(), + )); + } + SeekFrom::Current(v) => self.pos.saturating_add_signed(v), + }; + Ok(self.pos) + } +} + +fn read_partition_info(stream: &mut DiscReader) -> crate::Result> { + stream.seek(SeekFrom::Start(WII_PART_GROUP_OFF)).context("Seeking to partition groups")?; + let part_groups: [WiiPartGroup; 4] = read_from(stream).context("Reading partition groups")?; + let mut part_info = Vec::new(); + for (group_idx, group) in part_groups.iter().enumerate() { + let part_count = group.part_count.get(); + if part_count == 0 { + continue; + } + stream + .seek(SeekFrom::Start(group.part_entry_off())) + .with_context(|| format!("Seeking to partition group {group_idx}"))?; + let entries: Vec = read_vec(stream, part_count as usize) + .with_context(|| format!("Reading partition group {group_idx}"))?; + for (part_idx, entry) in entries.iter().enumerate() { + let offset = entry.offset(); + stream + .seek(SeekFrom::Start(offset)) + .with_context(|| format!("Seeking to partition data {group_idx}:{part_idx}"))?; + let header: Box = read_box(stream) + .with_context(|| format!("Reading partition header {group_idx}:{part_idx}"))?; + + let key = header.ticket.decrypt_title_key()?; + let start_offset = entry.offset(); + if start_offset % SECTOR_SIZE as u64 != 0 { + return Err(Error::DiscFormat(format!( + "Partition {group_idx}:{part_idx} offset is not sector aligned", + ))); + } + let data_start_offset = entry.offset() + header.data_off(); + let data_end_offset = data_start_offset + header.data_size(); + if data_start_offset % SECTOR_SIZE as u64 != 0 + || data_end_offset % SECTOR_SIZE as u64 != 0 + { + return Err(Error::DiscFormat(format!( + "Partition {group_idx}:{part_idx} data is not sector aligned", + ))); + } + let mut info = BPartitionInfo { + index: part_info.len() as u32, + kind: entry.kind.get().into(), + start_sector: (start_offset / SECTOR_SIZE as u64) as u32, + data_start_sector: (data_start_offset / SECTOR_SIZE as u64) as u32, + data_end_sector: (data_end_offset / SECTOR_SIZE as u64) as u32, + key, + header, + disc_header: DiscHeader::new_box_zeroed(), + partition_header: PartitionHeader::new_box_zeroed(), + hash_table: None, + }; + + let mut partition_reader = PartitionReader::new(stream.io.clone(), &info)?; + info.disc_header = read_box(&mut partition_reader).context("Reading disc header")?; + info.partition_header = + read_box(&mut partition_reader).context("Reading partition header")?; + + part_info.push(info); + } + } + Ok(part_info) +} + +fn guess_disc_size(part_info: &[BPartitionInfo]) -> u64 { + let max_offset = part_info + .iter() + .flat_map(|v| { + let offset = v.start_sector as u64 * SECTOR_SIZE as u64; + [ + offset + v.header.tmd_off() + v.header.tmd_size(), + offset + v.header.cert_chain_off() + v.header.cert_chain_size(), + offset + v.header.h3_table_off() + v.header.h3_table_size(), + offset + v.header.data_off() + v.header.data_size(), + ] + }) + .max() + .unwrap_or(0x50000); + if max_offset <= MINI_DVD_SIZE && !part_info.iter().any(|v| v.kind == PartitionKind::Data) { + // Datel disc + MINI_DVD_SIZE + } else if max_offset < SL_DVD_SIZE { + SL_DVD_SIZE + } else { + DL_DVD_SIZE + } +} diff --git a/src/disc/wii.rs b/src/disc/wii.rs index 2e490df..20bfcdb 100644 --- a/src/disc/wii.rs +++ b/src/disc/wii.rs @@ -1,4 +1,6 @@ use std::{ + cmp::min, + ffi::CStr, io, io::{Read, Seek, SeekFrom}, mem::size_of, @@ -16,36 +18,52 @@ use crate::{ fst::{Node, NodeKind}, io::{aes_decrypt, KeyBytes}, static_assert, - streams::{wrap_windowed, ReadStream, SharedWindowedReadStream}, + streams::{ReadStream, SharedWindowedReadStream}, util::{ div_rem, - reader::{read_from, read_vec}, + read::{read_from, read_vec}, }, Error, OpenOptions, PartitionHeader, Result, ResultContext, }; pub(crate) const HASHES_SIZE: usize = 0x400; -pub(crate) const BLOCK_SIZE: usize = SECTOR_SIZE - HASHES_SIZE; // 0x7C00 +pub(crate) const SECTOR_DATA_SIZE: usize = SECTOR_SIZE - HASHES_SIZE; // 0x7C00 +// ppki (Retail) +const RVL_CERT_ISSUER_PPKI_TICKET: &str = "Root-CA00000001-XS00000003"; #[rustfmt::skip] -const COMMON_KEYS: [KeyBytes; 2] = [ - /* Normal */ +const RETAIL_COMMON_KEYS: [KeyBytes; 3] = [ + /* RVL_KEY_RETAIL */ [0xeb, 0xe4, 0x2a, 0x22, 0x5e, 0x85, 0x93, 0xe4, 0x48, 0xd9, 0xc5, 0x45, 0x73, 0x81, 0xaa, 0xf7], - /* Korean */ + /* RVL_KEY_KOREAN */ [0x63, 0xb8, 0x2b, 0xb4, 0xf4, 0x61, 0x4e, 0x2e, 0x13, 0xf2, 0xfe, 0xfb, 0xba, 0x4c, 0x9b, 0x7e], + /* vWii_KEY_RETAIL */ + [0x30, 0xbf, 0xc7, 0x6e, 0x7c, 0x19, 0xaf, 0xbb, 0x23, 0x16, 0x33, 0x30, 0xce, 0xd7, 0xc2, 0x8d], +]; + +// dpki (Debug) +const RVL_CERT_ISSUER_DPKI_TICKET: &str = "Root-CA00000002-XS00000006"; +#[rustfmt::skip] +const DEBUG_COMMON_KEYS: [KeyBytes; 3] = [ + /* RVL_KEY_DEBUG */ + [0xa1, 0x60, 0x4a, 0x6a, 0x71, 0x23, 0xb5, 0x29, 0xae, 0x8b, 0xec, 0x32, 0xc8, 0x16, 0xfc, 0xaa], + /* RVL_KEY_KOREAN_DEBUG */ + [0x67, 0x45, 0x8b, 0x6b, 0xc6, 0x23, 0x7b, 0x32, 0x69, 0x98, 0x3c, 0x64, 0x73, 0x48, 0x33, 0x66], + /* vWii_KEY_DEBUG */ + [0x2f, 0x5c, 0x1b, 0x29, 0x44, 0xe7, 0xfd, 0x6f, 0xc3, 0x97, 0x96, 0x4b, 0x05, 0x76, 0x91, 0xfa], ]; #[derive(Debug, PartialEq, FromBytes, FromZeroes, AsBytes)] #[repr(C, align(4))] -struct WiiPartEntry { - offset: U32, - kind: U32, +pub(crate) struct WiiPartEntry { + pub(crate) offset: U32, + pub(crate) kind: U32, } static_assert!(size_of::() == 8); impl WiiPartEntry { - fn offset(&self) -> u64 { (self.offset.get() as u64) << 2 } + pub(crate) fn offset(&self) -> u64 { (self.offset.get() as u64) << 2 } } #[derive(Debug, PartialEq)] @@ -57,21 +75,22 @@ pub(crate) struct WiiPartInfo { pub(crate) header: WiiPartitionHeader, pub(crate) junk_id: [u8; 4], pub(crate) junk_start: u64, + pub(crate) title_key: KeyBytes, } -const WII_PART_GROUP_OFF: u64 = 0x40000; +pub(crate) const WII_PART_GROUP_OFF: u64 = 0x40000; #[derive(Debug, PartialEq, FromBytes, FromZeroes, AsBytes)] #[repr(C, align(4))] -struct WiiPartGroup { - part_count: U32, - part_entry_off: U32, +pub(crate) struct WiiPartGroup { + pub(crate) part_count: U32, + pub(crate) part_entry_off: U32, } static_assert!(size_of::() == 8); impl WiiPartGroup { - fn part_entry_off(&self) -> u64 { (self.part_entry_off.get() as u64) << 2 } + pub(crate) fn part_entry_off(&self) -> u64 { (self.part_entry_off.get() as u64) << 2 } } #[derive(Debug, Clone, PartialEq, FromBytes, FromZeroes, AsBytes)] @@ -122,6 +141,31 @@ pub struct Ticket { static_assert!(size_of::() == 0x2A4); +impl Ticket { + pub fn decrypt_title_key(&self) -> Result { + let mut iv: KeyBytes = [0; 16]; + iv[..8].copy_from_slice(&self.title_id); + let cert_issuer_ticket = + CStr::from_bytes_until_nul(&self.sig_issuer).ok().and_then(|c| c.to_str().ok()); + let common_keys = match cert_issuer_ticket { + Some(RVL_CERT_ISSUER_PPKI_TICKET) => &RETAIL_COMMON_KEYS, + Some(RVL_CERT_ISSUER_DPKI_TICKET) => &DEBUG_COMMON_KEYS, + Some(v) => { + return Err(Error::DiscFormat(format!("unknown certificate issuer {:?}", v))); + } + None => { + return Err(Error::DiscFormat("failed to parse certificate issuer".to_string())); + } + }; + let common_key = common_keys.get(self.common_key_idx as usize).ok_or(Error::DiscFormat( + format!("unknown common key index {}", self.common_key_idx), + ))?; + let mut title_key = self.title_key; + aes_decrypt(common_key, iv, &mut title_key); + Ok(title_key) + } +} + #[derive(Debug, Clone, PartialEq, FromBytes, FromZeroes, AsBytes)] #[repr(C, align(4))] pub struct TmdHeader { @@ -197,14 +241,17 @@ impl DiscWii { header: DiscHeader, disc_size: Option, ) -> Result { - let part_info = read_partition_info(stream)?; + let part_info = read_partition_info(stream, &header)?; // Guess disc size if not provided let disc_size = disc_size.unwrap_or_else(|| guess_disc_size(&part_info)); Ok(Self { header, part_info, disc_size }) } } -pub(crate) fn read_partition_info(stream: &mut dyn ReadStream) -> Result> { +pub(crate) fn read_partition_info( + stream: &mut dyn ReadStream, + disc_header: &DiscHeader, +) -> Result> { stream.seek(SeekFrom::Start(WII_PART_GROUP_OFF)).context("Seeking to partition groups")?; let part_groups: [WiiPartGroup; 4] = read_from(stream).context("Reading partition groups")?; let mut part_info = Vec::new(); @@ -223,32 +270,33 @@ pub(crate) fn read_partition_info(stream: &mut dyn ReadStream) -> Result Result Result( options: &OpenOptions, header: &DiscHeader, ) -> Result> { - let data_off = part.offset + part.header.data_off(); - let has_crypto = header.no_partition_encryption == 0; let mut base = disc_io.open()?; base.seek(SeekFrom::Start(part.offset + part.header.tmd_off())) @@ -327,19 +375,31 @@ fn open_partition<'a>( .context("Seeking to H3 table offset")?; let h3_table: Vec = read_vec(&mut base, H3_TABLE_SIZE).context("Reading H3 table")?; - let stream = wrap_windowed(base, data_off, part.header.data_size()).with_context(|| { - format!("Wrapping {}:{} partition stream", part.group_idx, part.part_idx) - })?; + let key = if header.no_partition_encryption == 0 { + Some(part.header.ticket.decrypt_title_key()?) + } else { + None + }; + let data_off = part.offset + part.header.data_off(); + if data_off % SECTOR_SIZE as u64 != 0 { + return Err(Error::DiscFormat(format!( + "Partition {}:{} offset is not sector aligned", + part.group_idx, part.part_idx + ))); + } + let start_sector = (data_off / SECTOR_SIZE as u64) as u32; Ok(Box::new(PartitionWii { + start_sector, header: part.header.clone(), tmd, cert_chain, h3_table, - stream: Box::new(stream), - key: has_crypto.then_some(part.header.ticket.title_key), + stream: base, + key, offset: 0, cur_block: u32::MAX, buf: [0; SECTOR_SIZE], + has_hashes: header.no_partition_hashes == 0, validate_hashes: options.validate_hashes && header.no_partition_hashes == 0, })) } @@ -392,6 +452,7 @@ impl DiscBase for DiscWii { } struct PartitionWii<'a> { + start_sector: u32, header: WiiPartitionHeader, tmd: Vec, cert_chain: Vec, @@ -402,6 +463,7 @@ struct PartitionWii<'a> { offset: u64, cur_block: u32, buf: [u8; SECTOR_SIZE], + has_hashes: bool, validate_hashes: bool, } @@ -409,10 +471,10 @@ impl<'a> PartitionBase for PartitionWii<'a> { fn meta(&mut self) -> Result> { self.seek(SeekFrom::Start(0)).context("Seeking to partition header")?; let mut meta = read_part_header(self, true)?; - meta.raw_ticket = Some(self.header.ticket.as_bytes().to_vec()); - meta.raw_tmd = Some(self.tmd.clone()); - meta.raw_cert_chain = Some(self.cert_chain.clone()); - meta.raw_h3_table = Some(self.h3_table.clone()); + meta.raw_ticket = Some(Box::from(self.header.ticket.as_bytes())); + meta.raw_tmd = Some(Box::from(self.tmd.as_slice())); + meta.raw_cert_chain = Some(Box::from(self.cert_chain.as_slice())); + meta.raw_h3_table = Some(Box::from(self.h3_table.as_slice())); Ok(meta) } @@ -421,7 +483,13 @@ impl<'a> PartitionBase for PartitionWii<'a> { self.new_window(node.offset(true), node.length(true)) } - fn ideal_buffer_size(&self) -> usize { BLOCK_SIZE } + fn ideal_buffer_size(&self) -> usize { + if self.has_hashes { + SECTOR_DATA_SIZE + } else { + SECTOR_SIZE + } + } } #[inline(always)] @@ -495,65 +563,76 @@ fn decrypt_block(part: &mut PartitionWii, cluster: u32) -> io::Result<()> { impl<'a> Read for PartitionWii<'a> { fn read(&mut self, buf: &mut [u8]) -> io::Result { - let (block, block_offset) = div_rem(self.offset, BLOCK_SIZE as u64); - let mut block = block as u32; - let mut block_offset = block_offset as usize; - - let mut rem = buf.len(); - let mut read: usize = 0; - - while rem > 0 { - if block != self.cur_block { - decrypt_block(self, block)?; - self.cur_block = block; - } - - let mut cache_size = rem; - if cache_size + block_offset > BLOCK_SIZE { - cache_size = BLOCK_SIZE - block_offset; - } - - buf[read..read + cache_size].copy_from_slice( - &self.buf[HASHES_SIZE + block_offset..HASHES_SIZE + block_offset + cache_size], - ); - read += cache_size; - rem -= cache_size; - block_offset = 0; - block += 1; + let block_size = self.ideal_buffer_size() as u64; + let (block, block_offset) = div_rem(self.offset, block_size); + let block = block as u32; + if block != self.cur_block { + self.stream + .seek(SeekFrom::Start((self.start_sector + block) as u64 * SECTOR_SIZE as u64))?; + decrypt_block(self, block)?; + self.cur_block = block; } - self.offset += buf.len() as u64; - Ok(buf.len()) + let offset = (SECTOR_SIZE - block_size as usize) + block_offset as usize; + let read = min(buf.len(), block_size as usize - block_offset as usize); + buf[..read].copy_from_slice(&self.buf[offset..offset + read]); + self.offset += read as u64; + Ok(read) + + // let mut block = block as u32; + // + // let mut rem = buf.len(); + // let mut read: usize = 0; + // + // while rem > 0 { + // if block != self.cur_block { + // decrypt_block(self, block)?; + // self.cur_block = block; + // } + // + // let mut cache_size = rem; + // if cache_size as u64 + block_offset > block_size { + // cache_size = (block_size - block_offset) as usize; + // } + // + // let hashes_size = SECTOR_SIZE - block_size as usize; + // let start = hashes_size + block_offset as usize; + // buf[read..read + cache_size].copy_from_slice(&self.buf[start..start + cache_size]); + // read += cache_size; + // rem -= cache_size; + // block_offset = 0; + // block += 1; + // } + // + // self.offset += buf.len() as u64; + // Ok(buf.len()) } } #[inline(always)] fn to_block_size(v: u64) -> u64 { - (v / SECTOR_SIZE as u64) * BLOCK_SIZE as u64 + (v % SECTOR_SIZE as u64) + (v / SECTOR_SIZE as u64) * SECTOR_DATA_SIZE as u64 + (v % SECTOR_SIZE as u64) } impl<'a> Seek for PartitionWii<'a> { fn seek(&mut self, pos: SeekFrom) -> io::Result { self.offset = match pos { SeekFrom::Start(v) => v, - SeekFrom::End(v) => self.stable_stream_len()?.saturating_add_signed(v), + SeekFrom::End(_) => { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "PartitionWii: SeekFrom::End is not supported", + )); + } SeekFrom::Current(v) => self.offset.saturating_add_signed(v), }; - let block = self.offset / BLOCK_SIZE as u64; - if block as u32 != self.cur_block { - self.stream.seek(SeekFrom::Start(block * SECTOR_SIZE as u64))?; - self.cur_block = u32::MAX; - } + // let block = self.offset / self.ideal_buffer_size() as u64; + // if block as u32 != self.cur_block { + // self.stream.seek(SeekFrom::Start((self.start_sector + block) * SECTOR_SIZE as u64))?; + // self.cur_block = u32::MAX; + // } Ok(self.offset) } fn stream_position(&mut self) -> io::Result { Ok(self.offset) } } - -impl<'a> ReadStream for PartitionWii<'a> { - fn stable_stream_len(&mut self) -> io::Result { - Ok(to_block_size(self.stream.stable_stream_len()?)) - } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} diff --git a/src/io/block.rs b/src/io/block.rs new file mode 100644 index 0000000..c46a92c --- /dev/null +++ b/src/io/block.rs @@ -0,0 +1,279 @@ +use std::{cmp::min, fs, fs::File, io, path::Path}; + +use dyn_clone::DynClone; +use zerocopy::transmute_ref; + +use crate::{ + array_ref, + disc::{ + hashes::HashTable, + wii::{WiiPartitionHeader, HASHES_SIZE, SECTOR_DATA_SIZE}, + SECTOR_SIZE, + }, + io::{aes_decrypt, aes_encrypt, ciso, iso, nfs, wbfs, wia, KeyBytes, MagicBytes}, + util::{lfg::LaggedFibonacci, read::read_from}, + DiscHeader, DiscMeta, Error, OpenOptions, PartitionHeader, PartitionKind, Result, + ResultContext, +}; + +/// Block I/O trait for reading disc images. +pub trait BlockIO: DynClone + Send + Sync { + /// Reads a block from the disc image. + fn read_block( + &mut self, + out: &mut [u8], + block: u32, + partition: Option<&BPartitionInfo>, + ) -> io::Result>; + + /// The format's block size in bytes. Must be a multiple of the sector size (0x8000). + fn block_size(&self) -> u32; + + /// Returns extra metadata included in the disc file format, if any. + fn meta(&self) -> Result; +} + +dyn_clone::clone_trait_object!(BlockIO); + +/// Creates a new [`BlockIO`] instance. +pub fn open(filename: &Path, options: &OpenOptions) -> Result> { + let path_result = fs::canonicalize(filename); + if let Err(err) = path_result { + return Err(Error::Io(format!("Failed to open {}", filename.display()), err)); + } + let path = path_result.as_ref().unwrap(); + let meta = fs::metadata(path); + if let Err(err) = meta { + return Err(Error::Io(format!("Failed to open {}", filename.display()), err)); + } + if !meta.unwrap().is_file() { + return Err(Error::DiscFormat(format!("Input is not a file: {}", filename.display()))); + } + let magic: MagicBytes = { + let mut file = + File::open(path).with_context(|| format!("Opening file {}", filename.display()))?; + read_from(&mut file) + .with_context(|| format!("Reading magic bytes from {}", filename.display()))? + }; + match magic { + ciso::CISO_MAGIC => Ok(ciso::DiscIOCISO::new(path)?), + nfs::NFS_MAGIC => match path.parent() { + Some(parent) if parent.is_dir() => { + Ok(nfs::DiscIONFS::new(path.parent().unwrap(), options)?) + } + _ => Err(Error::DiscFormat("Failed to locate NFS parent directory".to_string())), + }, + wbfs::WBFS_MAGIC => Ok(wbfs::DiscIOWBFS::new(path)?), + wia::WIA_MAGIC | wia::RVZ_MAGIC => Ok(wia::DiscIOWIA::new(path, options)?), + _ => Ok(iso::DiscIOISO::new(path)?), + } +} + +#[derive(Debug, Clone)] +pub struct BPartitionInfo { + pub index: u32, + pub kind: PartitionKind, + pub start_sector: u32, + pub data_start_sector: u32, + pub data_end_sector: u32, + pub key: KeyBytes, + pub header: Box, + pub disc_header: Box, + pub partition_header: Box, + pub hash_table: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Block { + /// Raw data or encrypted Wii partition data + Raw, + /// Decrypted Wii partition data + PartDecrypted { + /// Whether the sector has its hash block intact + has_hashes: bool, + }, + /// Wii partition junk data + Junk, + /// All zeroes + Zero, +} + +impl Block { + pub(crate) fn decrypt( + self, + out: &mut [u8; SECTOR_SIZE], + data: &[u8], + block_idx: u32, + abs_sector: u32, + partition: &BPartitionInfo, + ) -> io::Result<()> { + let rel_sector = abs_sector - self.start_sector(block_idx, data.len()); + match self { + Block::Raw => { + out.copy_from_slice(block_sector::(data, rel_sector)?); + decrypt_sector(out, partition); + } + Block::PartDecrypted { has_hashes } => { + out.copy_from_slice(block_sector::(data, rel_sector)?); + if !has_hashes { + rebuild_hash_block(out, abs_sector, partition); + } + } + Block::Junk => { + generate_junk(out, abs_sector, Some(partition), &partition.disc_header); + rebuild_hash_block(out, abs_sector, partition); + } + Block::Zero => { + out.fill(0); + rebuild_hash_block(out, abs_sector, partition); + } + } + Ok(()) + } + + pub(crate) fn encrypt( + self, + out: &mut [u8; SECTOR_SIZE], + data: &[u8], + block_idx: u32, + abs_sector: u32, + partition: &BPartitionInfo, + ) -> io::Result<()> { + let rel_sector = abs_sector - self.start_sector(block_idx, data.len()); + match self { + Block::Raw => { + out.copy_from_slice(block_sector::(data, rel_sector)?); + } + Block::PartDecrypted { has_hashes } => { + out.copy_from_slice(block_sector::(data, rel_sector)?); + if !has_hashes { + rebuild_hash_block(out, abs_sector, partition); + } + encrypt_sector(out, partition); + } + Block::Junk => { + generate_junk(out, abs_sector, Some(partition), &partition.disc_header); + rebuild_hash_block(out, abs_sector, partition); + encrypt_sector(out, partition); + } + Block::Zero => { + out.fill(0); + rebuild_hash_block(out, abs_sector, partition); + encrypt_sector(out, partition); + } + } + Ok(()) + } + + pub(crate) fn copy_raw( + self, + out: &mut [u8; SECTOR_SIZE], + data: &[u8], + block_idx: u32, + abs_sector: u32, + disc_header: &DiscHeader, + ) -> io::Result<()> { + match self { + Block::Raw => { + out.copy_from_slice(block_sector::( + data, + abs_sector - self.start_sector(block_idx, data.len()), + )?); + } + Block::PartDecrypted { .. } => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Cannot copy decrypted data as raw", + )); + } + Block::Junk => generate_junk(out, abs_sector, None, disc_header), + Block::Zero => out.fill(0), + } + Ok(()) + } + + /// Returns the start sector of the block. + fn start_sector(&self, index: u32, block_size: usize) -> u32 { + (index as u64 * block_size as u64 / SECTOR_SIZE as u64) as u32 + } +} + +#[inline(always)] +fn block_sector(data: &[u8], sector_idx: u32) -> io::Result<&[u8; N]> { + if data.len() % N != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Expected block size {} to be a multiple of {}", data.len(), N), + )); + } + let offset = sector_idx as usize * N; + data.get(offset..offset + N) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Sector {} out of range (block size {}, sector size {})", + sector_idx, + data.len(), + N + ), + ) + }) + .map(|v| unsafe { &*(v as *const [u8] as *const [u8; N]) }) +} + +fn generate_junk( + out: &mut [u8; SECTOR_SIZE], + sector: u32, + partition: Option<&BPartitionInfo>, + disc_header: &DiscHeader, +) { + let mut pos = if let Some(partition) = partition { + (sector - partition.data_start_sector) as u64 * SECTOR_DATA_SIZE as u64 + } else { + sector as u64 * SECTOR_SIZE as u64 + }; + let mut offset = if partition.is_some() { HASHES_SIZE } else { 0 }; + out[..offset].fill(0); + while offset < SECTOR_SIZE { + // The LFG spans a single sector of the decrypted data, + // so we may need to initialize it multiple times + let mut lfg = LaggedFibonacci::default(); + lfg.init_with_seed(*array_ref![disc_header.game_id, 0, 4], disc_header.disc_num, pos); + let sector_end = (pos + SECTOR_SIZE as u64) & !(SECTOR_SIZE as u64 - 1); + let len = min(SECTOR_SIZE - offset, (sector_end - pos) as usize); + lfg.fill(&mut out[offset..offset + len]); + pos += len as u64; + offset += len; + } +} + +fn rebuild_hash_block(out: &mut [u8; SECTOR_SIZE], sector: u32, partition: &BPartitionInfo) { + let Some(hash_table) = partition.hash_table.as_ref() else { + return; + }; + let sector_idx = (sector - partition.data_start_sector) as usize; + let h0_hashes: &[u8; 0x26C] = + transmute_ref!(array_ref![hash_table.h0_hashes, sector_idx * 31, 31]); + out[0..0x26C].copy_from_slice(h0_hashes); + let h1_hashes: &[u8; 0xA0] = + transmute_ref!(array_ref![hash_table.h1_hashes, sector_idx & !7, 8]); + out[0x280..0x320].copy_from_slice(h1_hashes); + let h2_hashes: &[u8; 0xA0] = + transmute_ref!(array_ref![hash_table.h2_hashes, (sector_idx / 8) & !7, 8]); + out[0x340..0x3E0].copy_from_slice(h2_hashes); +} + +fn encrypt_sector(out: &mut [u8; SECTOR_SIZE], partition: &BPartitionInfo) { + aes_encrypt(&partition.key, [0u8; 16], &mut out[..HASHES_SIZE]); + // Data IV from encrypted hash block + let iv = *array_ref![out, 0x3D0, 16]; + aes_encrypt(&partition.key, iv, &mut out[HASHES_SIZE..]); +} + +fn decrypt_sector(out: &mut [u8; SECTOR_SIZE], partition: &BPartitionInfo) { + // Data IV from encrypted hash block + let iv = *array_ref![out, 0x3D0, 16]; + aes_decrypt(&partition.key, [0u8; 16], &mut out[..HASHES_SIZE]); + aes_decrypt(&partition.key, iv, &mut out[HASHES_SIZE..]); +} diff --git a/src/io/ciso.rs b/src/io/ciso.rs index 4eeae39..07b632c 100644 --- a/src/io/ciso.rs +++ b/src/io/ciso.rs @@ -1,7 +1,6 @@ use std::{ - cmp::min, io, - io::{BufReader, Read, Seek, SeekFrom}, + io::{Read, Seek, SeekFrom}, mem::size_of, path::Path, }; @@ -9,14 +8,16 @@ use std::{ use zerocopy::{little_endian::*, AsBytes, FromBytes, FromZeroes}; use crate::{ - disc::{gcn::DiscGCN, wii::DiscWii, DiscBase, DL_DVD_SIZE, SECTOR_SIZE}, - io::{nkit::NKitHeader, split::SplitFileReader, DiscIO, MagicBytes}, - static_assert, - util::{ - lfg::LaggedFibonacci, - reader::{read_box_slice, read_from}, + disc::SECTOR_SIZE, + io::{ + block::{BPartitionInfo, Block, BlockIO}, + nkit::NKitHeader, + split::SplitFileReader, + MagicBytes, }, - DiscHeader, DiscMeta, Error, PartitionInfo, ReadStream, Result, ResultContext, + static_assert, + util::read::read_from, + DiscMeta, Error, Result, ResultContext, }; pub const CISO_MAGIC: MagicBytes = *b"CISO"; @@ -38,14 +39,22 @@ pub struct DiscIOCISO { header: CISOHeader, block_map: [u16; CISO_MAP_SIZE], nkit_header: Option, - junk_blocks: Option>, - partitions: Vec, - disc_num: u8, +} + +impl Clone for DiscIOCISO { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + header: self.header.clone(), + block_map: self.block_map, + nkit_header: self.nkit_header.clone(), + } + } } impl DiscIOCISO { - pub fn new(filename: &Path) -> Result { - let mut inner = BufReader::new(SplitFileReader::new(filename)?); + pub fn new(filename: &Path) -> Result> { + let mut inner = SplitFileReader::new(filename)?; // Read header let header: CISOHeader = read_from(&mut inner).context("Reading CISO header")?; @@ -65,203 +74,63 @@ impl DiscIOCISO { } } let file_size = SECTOR_SIZE as u64 + block as u64 * header.block_size.get() as u64; - if file_size > inner.get_ref().len() { + if file_size > inner.len() { return Err(Error::DiscFormat(format!( "CISO file size mismatch: expected at least {} bytes, got {}", file_size, - inner.get_ref().len() + inner.len() ))); } // Read NKit header if present (after CISO data) - let nkit_header = if inner.get_ref().len() > file_size + 4 { + let nkit_header = if inner.len() > file_size + 4 { inner.seek(SeekFrom::Start(file_size)).context("Seeking to NKit header")?; - NKitHeader::try_read_from(&mut inner) + NKitHeader::try_read_from(&mut inner, header.block_size.get(), true) } else { None }; - // Read junk data bitstream if present (after NKit header) - let junk_blocks = if nkit_header.is_some() { - let n = 1 + DL_DVD_SIZE / header.block_size.get() as u64 / 8; - Some(read_box_slice(&mut inner, n as usize).context("Reading NKit bitstream")?) - } else { - None - }; - - let (partitions, disc_num) = if junk_blocks.is_some() { - let mut stream: Box = Box::new(CISOReadStream { - inner: BufReader::new(inner.get_ref().clone()), - block_size: header.block_size.get(), - block_map, - cur_block: u16::MAX, - pos: 0, - junk_blocks: None, - partitions: vec![], - disc_num: 0, - }); - let header: DiscHeader = read_from(stream.as_mut()).context("Reading disc header")?; - let disc_num = header.disc_num; - let disc_base: Box = if header.is_wii() { - Box::new(DiscWii::new(stream.as_mut(), header, None)?) - } else if header.is_gamecube() { - Box::new(DiscGCN::new(stream.as_mut(), header, None)?) - } else { - return Err(Error::DiscFormat(format!( - "Invalid GC/Wii magic: {:#010X}/{:#010X}", - header.gcn_magic.get(), - header.wii_magic.get() - ))); - }; - (disc_base.partitions(), disc_num) - } else { - (vec![], 0) - }; - // Reset reader - let mut inner = inner.into_inner(); inner.reset(); - Ok(Self { inner, header, block_map, nkit_header, junk_blocks, partitions, disc_num }) + Ok(Box::new(Self { inner, header, block_map, nkit_header })) } } -impl DiscIO for DiscIOCISO { - fn open(&self) -> Result> { - Ok(Box::new(CISOReadStream { - inner: BufReader::new(self.inner.clone()), - block_size: self.header.block_size.get(), - block_map: self.block_map, - cur_block: u16::MAX, - pos: 0, - junk_blocks: self.junk_blocks.clone(), - partitions: self.partitions.clone(), - disc_num: self.disc_num, - })) +impl BlockIO for DiscIOCISO { + fn read_block( + &mut self, + out: &mut [u8], + block: u32, + _partition: Option<&BPartitionInfo>, + ) -> io::Result> { + if block >= CISO_MAP_SIZE as u32 { + // Out of bounds + return Ok(None); + } + + // Find the block in the map + let phys_block = self.block_map[block as usize]; + if phys_block == u16::MAX { + // Check if block is junk data + if self.nkit_header.as_ref().is_some_and(|h| h.is_junk_block(block).unwrap_or(false)) { + return Ok(Some(Block::Junk)); + }; + + // Otherwise, read zeroes + return Ok(Some(Block::Zero)); + } + + // Read block + let file_offset = size_of::() as u64 + + phys_block as u64 * self.header.block_size.get() as u64; + self.inner.seek(SeekFrom::Start(file_offset))?; + self.inner.read_exact(out)?; + Ok(Some(Block::Raw)) } + fn block_size(&self) -> u32 { self.header.block_size.get() } + fn meta(&self) -> Result { Ok(self.nkit_header.as_ref().map(DiscMeta::from).unwrap_or_default()) } - - fn disc_size(&self) -> Option { self.nkit_header.as_ref().and_then(|h| h.size) } -} - -struct CISOReadStream { - inner: BufReader, - block_size: u32, - block_map: [u16; CISO_MAP_SIZE], - cur_block: u16, - pos: u64, - - // Data for recreating junk data - junk_blocks: Option>, - partitions: Vec, - disc_num: u8, -} - -impl CISOReadStream { - fn read_junk_data(&mut self, buf: &mut [u8]) -> io::Result { - let Some(junk_blocks) = self.junk_blocks.as_deref() else { - return Ok(0); - }; - let block_size = self.block_size as u64; - let block = (self.pos / block_size) as u16; - if junk_blocks[(block / 8) as usize] & (1 << (7 - (block & 7))) == 0 { - return Ok(0); - } - let Some(partition) = self.partitions.iter().find(|p| { - let start = p.part_offset + p.data_offset; - start <= self.pos && self.pos < start + p.data_size - }) else { - log::warn!("No partition found for junk data at offset {:#x}", self.pos); - return Ok(0); - }; - let offset = self.pos - (partition.part_offset + partition.data_offset); - let to_read = min( - buf.len(), - // The LFG is only valid for a single sector - SECTOR_SIZE - (offset % SECTOR_SIZE as u64) as usize, - ); - let mut lfg = LaggedFibonacci::default(); - lfg.init_with_seed(partition.lfg_seed, self.disc_num, offset); - lfg.fill(&mut buf[..to_read]); - self.pos += to_read as u64; - Ok(to_read) - } -} - -impl Read for CISOReadStream { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - let block_size = self.block_size as u64; - let block = (self.pos / block_size) as u16; - let block_offset = self.pos & (block_size - 1); - if block != self.cur_block { - if block >= CISO_MAP_SIZE as u16 { - return Ok(0); - } - - // Find the block in the map - let phys_block = self.block_map[block as usize]; - if phys_block == u16::MAX { - // Try to recreate junk data - let read = self.read_junk_data(buf)?; - if read > 0 { - return Ok(read); - } - - // Otherwise, read zeroes - let to_read = min(buf.len(), (block_size - block_offset) as usize); - buf[..to_read].fill(0); - self.pos += to_read as u64; - return Ok(to_read); - } - - // Seek to the new block - let file_offset = - size_of::() as u64 + phys_block as u64 * block_size + block_offset; - self.inner.seek(SeekFrom::Start(file_offset))?; - self.cur_block = block; - } - - let to_read = min(buf.len(), (block_size - block_offset) as usize); - let read = self.inner.read(&mut buf[..to_read])?; - self.pos += read as u64; - Ok(read) - } -} - -impl Seek for CISOReadStream { - fn seek(&mut self, pos: SeekFrom) -> io::Result { - let new_pos = match pos { - SeekFrom::Start(v) => v, - SeekFrom::End(_) => { - return Err(io::Error::new( - io::ErrorKind::Unsupported, - "CISOReadStream: SeekFrom::End is not supported", - )); - } - SeekFrom::Current(v) => self.pos.saturating_add_signed(v), - }; - - let block_size = self.block_size as u64; - let new_block = (self.pos / block_size) as u16; - if new_block == self.cur_block { - // Seek within the same block - self.inner.seek(SeekFrom::Current(new_pos as i64 - self.pos as i64))?; - } else { - // Seek to a different block, handled by next read - self.cur_block = u16::MAX; - } - - self.pos = new_pos; - Ok(new_pos) - } -} - -impl ReadStream for CISOReadStream { - fn stable_stream_len(&mut self) -> io::Result { - Ok(self.block_size as u64 * CISO_MAP_SIZE as u64) - } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } diff --git a/src/io/iso.rs b/src/io/iso.rs index 6b971fb..1bc6362 100644 --- a/src/io/iso.rs +++ b/src/io/iso.rs @@ -1,25 +1,56 @@ -use std::{io::BufReader, path::Path}; - -use crate::{ - io::{split::SplitFileReader, DiscIO}, - streams::ReadStream, - Result, +use std::{ + io, + io::{Read, Seek}, + path::Path, }; +use crate::{ + disc::SECTOR_SIZE, + io::{ + block::{BPartitionInfo, Block, BlockIO}, + split::SplitFileReader, + }, + DiscMeta, Error, Result, +}; + +#[derive(Clone)] pub struct DiscIOISO { - pub inner: SplitFileReader, + inner: SplitFileReader, } impl DiscIOISO { - pub fn new(filename: &Path) -> Result { - Ok(Self { inner: SplitFileReader::new(filename)? }) + pub fn new(filename: &Path) -> Result> { + let inner = SplitFileReader::new(filename)?; + if inner.len() % SECTOR_SIZE as u64 != 0 { + return Err(Error::DiscFormat( + "ISO size is not a multiple of sector size (0x8000 bytes)".to_string(), + )); + } + Ok(Box::new(Self { inner })) } } -impl DiscIO for DiscIOISO { - fn open(&self) -> Result> { - Ok(Box::new(BufReader::new(self.inner.clone()))) +impl BlockIO for DiscIOISO { + fn read_block( + &mut self, + out: &mut [u8], + block: u32, + _partition: Option<&BPartitionInfo>, + ) -> io::Result> { + let offset = block as u64 * SECTOR_SIZE as u64; + if offset >= self.inner.len() { + // End of file + return Ok(None); + } + + self.inner.seek(io::SeekFrom::Start(offset))?; + self.inner.read_exact(out)?; + Ok(Some(Block::Raw)) } - fn disc_size(&self) -> Option { Some(self.inner.len()) } + fn block_size(&self) -> u32 { SECTOR_SIZE as u32 } + + fn meta(&self) -> Result { + Ok(DiscMeta { lossless: true, disc_size: Some(self.inner.len()), ..Default::default() }) + } } diff --git a/src/io/mod.rs b/src/io/mod.rs index 5b80e64..efe50e1 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -1,11 +1,8 @@ //! Disc file format related logic (CISO, NFS, WBFS, WIA, etc.) -use std::{fs, fs::File, path::Path}; - -use crate::{ - streams::ReadStream, util::reader::read_from, Error, OpenOptions, Result, ResultContext, -}; +use crate::{streams::ReadStream, Result}; +pub(crate) mod block; pub(crate) mod ciso; pub(crate) mod iso; pub(crate) mod nfs; @@ -36,49 +33,27 @@ pub trait DiscIO: Send + Sync { fn disc_size(&self) -> Option; } -/// Extra metadata included in some disc file formats. +/// Extra metadata about the underlying disc file format. #[derive(Debug, Clone, Default)] pub struct DiscMeta { + /// Whether Wii partitions are stored decrypted in the format. + pub decrypted: bool, + /// Whether the format omits Wii partition data hashes. + pub needs_hash_recovery: bool, + /// Whether the format supports recovering the original disc data losslessly. + pub lossless: bool, + /// The original disc's size in bytes, if stored by the format. + pub disc_size: Option, + /// The original disc's CRC32 hash, if stored by the format. pub crc32: Option, + /// The original disc's MD5 hash, if stored by the format. pub md5: Option<[u8; 16]>, + /// The original disc's SHA-1 hash, if stored by the format. pub sha1: Option<[u8; 20]>, + /// The original disc's XXH64 hash, if stored by the format. pub xxhash64: Option, } -/// Creates a new [`DiscIO`] instance. -pub fn open(filename: &Path, options: &OpenOptions) -> Result> { - let path_result = fs::canonicalize(filename); - if let Err(err) = path_result { - return Err(Error::Io(format!("Failed to open {}", filename.display()), err)); - } - let path = path_result.as_ref().unwrap(); - let meta = fs::metadata(path); - if let Err(err) = meta { - return Err(Error::Io(format!("Failed to open {}", filename.display()), err)); - } - if !meta.unwrap().is_file() { - return Err(Error::DiscFormat(format!("Input is not a file: {}", filename.display()))); - } - let magic: MagicBytes = { - let mut file = - File::open(path).with_context(|| format!("Opening file {}", filename.display()))?; - read_from(&mut file) - .with_context(|| format!("Reading magic bytes from {}", filename.display()))? - }; - match magic { - ciso::CISO_MAGIC => Ok(Box::new(ciso::DiscIOCISO::new(path)?)), - nfs::NFS_MAGIC => match path.parent() { - Some(parent) if parent.is_dir() => { - Ok(Box::new(nfs::DiscIONFS::new(path.parent().unwrap(), options)?)) - } - _ => Err(Error::DiscFormat("Failed to locate NFS parent directory".to_string())), - }, - wbfs::WBFS_MAGIC => Ok(Box::new(wbfs::DiscIOWBFS::new(path)?)), - wia::WIA_MAGIC | wia::RVZ_MAGIC => Ok(Box::new(wia::DiscIOWIA::new(path, options)?)), - _ => Ok(Box::new(iso::DiscIOISO::new(path)?)), - } -} - /// Encrypts data in-place using AES-128-CBC with the given key and IV. pub(crate) fn aes_encrypt(key: &KeyBytes, iv: KeyBytes, data: &mut [u8]) { use aes::cipher::{block_padding::NoPadding, BlockEncryptMut, KeyIvInit}; diff --git a/src/io/nfs.rs b/src/io/nfs.rs index e3edea6..94930b5 100644 --- a/src/io/nfs.rs +++ b/src/io/nfs.rs @@ -9,16 +9,16 @@ use std::{ use zerocopy::{big_endian::U32, AsBytes, FromBytes, FromZeroes}; use crate::{ - array_ref, - disc::{ - wii::{read_partition_info, HASHES_SIZE}, - SECTOR_SIZE, + disc::SECTOR_SIZE, + io::{ + aes_decrypt, + block::{BPartitionInfo, Block, BlockIO}, + split::SplitFileReader, + KeyBytes, MagicBytes, }, - io::{aes_decrypt, aes_encrypt, split::SplitFileReader, DiscIO, KeyBytes, MagicBytes}, static_assert, - streams::ReadStream, - util::reader::read_from, - DiscHeader, Error, OpenOptions, Result, ResultContext, + util::read::read_from, + DiscMeta, Error, OpenOptions, Result, ResultContext, }; pub const NFS_MAGIC: MagicBytes = *b"EGGS"; @@ -26,27 +26,27 @@ pub const NFS_END_MAGIC: MagicBytes = *b"SGGE"; #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)] #[repr(C, align(4))] -pub struct LBARange { - pub start_sector: U32, - pub num_sectors: U32, +struct LBARange { + start_sector: U32, + num_sectors: U32, } #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)] #[repr(C, align(4))] -pub struct NFSHeader { - pub magic: MagicBytes, - pub version: U32, - pub unk1: U32, - pub unk2: U32, - pub num_lba_ranges: U32, - pub lba_ranges: [LBARange; 61], - pub end_magic: MagicBytes, +struct NFSHeader { + magic: MagicBytes, + version: U32, + unk1: U32, + unk2: U32, + num_lba_ranges: U32, + lba_ranges: [LBARange; 61], + end_magic: MagicBytes, } static_assert!(size_of::() == 0x200); impl NFSHeader { - pub fn validate(&self) -> Result<()> { + fn validate(&self) -> Result<()> { if self.magic != NFS_MAGIC { return Err(Error::DiscFormat("Invalid NFS magic".to_string())); } @@ -59,11 +59,9 @@ impl NFSHeader { Ok(()) } - pub fn lba_ranges(&self) -> &[LBARange] { - &self.lba_ranges[..self.num_lba_ranges.get() as usize] - } + fn lba_ranges(&self) -> &[LBARange] { &self.lba_ranges[..self.num_lba_ranges.get() as usize] } - pub fn calculate_num_files(&self) -> u32 { + fn calculate_num_files(&self) -> u32 { let sector_count = self.lba_ranges().iter().fold(0u32, |acc, range| acc + range.num_sectors.get()); (((sector_count as u64) * (SECTOR_SIZE as u64) @@ -71,7 +69,7 @@ impl NFSHeader { / 0xFA00000u64) as u32 } - pub fn phys_sector(&self, sector: u32) -> u32 { + fn phys_sector(&self, sector: u32) -> u32 { let mut cur_sector = 0u32; for range in self.lba_ranges().iter() { if sector >= range.start_sector.get() @@ -86,190 +84,82 @@ impl NFSHeader { } pub struct DiscIONFS { - pub inner: SplitFileReader, - pub header: NFSHeader, - pub raw_size: u64, - pub disc_size: u64, - pub key: KeyBytes, - pub encrypt: bool, + inner: SplitFileReader, + header: NFSHeader, + raw_size: u64, + disc_size: u64, + key: KeyBytes, + encrypt: bool, +} + +impl Clone for DiscIONFS { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + header: self.header.clone(), + raw_size: self.raw_size, + disc_size: self.disc_size, + key: self.key, + encrypt: self.encrypt, + } + } } impl DiscIONFS { - pub fn new(directory: &Path, options: &OpenOptions) -> Result { - let mut disc_io = DiscIONFS { + pub fn new(directory: &Path, options: &OpenOptions) -> Result> { + let mut disc_io = Box::new(Self { inner: SplitFileReader::empty(), header: NFSHeader::new_zeroed(), raw_size: 0, disc_size: 0, key: [0; 16], encrypt: options.rebuild_encryption, - }; + }); disc_io.load_files(directory)?; Ok(disc_io) } } -pub struct NFSReadStream { - /// Underlying file reader - inner: SplitFileReader, - /// NFS file header - header: NFSHeader, - /// Inner disc header - disc_header: Option, - /// Estimated disc size - disc_size: u64, - /// Current offset - pos: u64, - /// Current sector - sector: u32, - /// Current decrypted sector - buf: [u8; SECTOR_SIZE], - /// AES key - key: KeyBytes, - /// Wii partition info - part_info: Vec, -} - -struct PartitionInfo { - start_sector: u32, - end_sector: u32, - key: KeyBytes, -} - -impl NFSReadStream { - fn read_sector(&mut self, sector: u32) -> io::Result<()> { +impl BlockIO for DiscIONFS { + fn read_block( + &mut self, + out: &mut [u8], + sector: u32, + partition: Option<&BPartitionInfo>, + ) -> io::Result> { // Calculate physical sector let phys_sector = self.header.phys_sector(sector); if phys_sector == u32::MAX { // Logical zero sector - self.buf.fill(0u8); - return Ok(()); + return Ok(Some(Block::Zero)); } // Read sector let offset = size_of::() as u64 + phys_sector as u64 * SECTOR_SIZE as u64; self.inner.seek(SeekFrom::Start(offset))?; - self.inner.read_exact(&mut self.buf)?; + self.inner.read_exact(out)?; // Decrypt let iv_bytes = sector.to_be_bytes(); #[rustfmt::skip] - let iv: KeyBytes = [ + let iv: KeyBytes = [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, iv_bytes[0], iv_bytes[1], iv_bytes[2], iv_bytes[3], ]; - aes_decrypt(&self.key, iv, &mut self.buf); + aes_decrypt(&self.key, iv, out); - if sector == 0 { - if let Some(header) = &self.disc_header { - // Replace disc header in buffer - let header_bytes = header.as_bytes(); - self.buf[..header_bytes.len()].copy_from_slice(header_bytes); - } + if partition.is_some() { + Ok(Some(Block::PartDecrypted { has_hashes: true })) + } else { + Ok(Some(Block::Raw)) } - - // Re-encrypt if needed - if let Some(part) = self - .part_info - .iter() - .find(|part| sector >= part.start_sector && sector < part.end_sector) - { - // Encrypt hashes - aes_encrypt(&part.key, [0u8; 16], &mut self.buf[..HASHES_SIZE]); - // Encrypt data using IV from H2 - aes_encrypt(&part.key, *array_ref![self.buf, 0x3d0, 16], &mut self.buf[HASHES_SIZE..]); - } - - Ok(()) - } -} - -impl Read for NFSReadStream { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - let sector = (self.pos / SECTOR_SIZE as u64) as u32; - let sector_off = (self.pos % SECTOR_SIZE as u64) as usize; - if sector != self.sector { - self.read_sector(sector)?; - self.sector = sector; - } - - let read = buf.len().min(SECTOR_SIZE - sector_off); - buf[..read].copy_from_slice(&self.buf[sector_off..sector_off + read]); - self.pos += read as u64; - Ok(read) - } -} - -impl Seek for NFSReadStream { - fn seek(&mut self, pos: SeekFrom) -> io::Result { - self.pos = match pos { - SeekFrom::Start(v) => v, - SeekFrom::End(_) => { - return Err(io::Error::new( - io::ErrorKind::Unsupported, - "NFSReadStream: SeekFrom::End is not supported", - )); - } - SeekFrom::Current(v) => self.pos.saturating_add_signed(v), - }; - Ok(self.pos) } - fn stream_position(&mut self) -> io::Result { Ok(self.pos) } -} + fn block_size(&self) -> u32 { SECTOR_SIZE as u32 } -impl ReadStream for NFSReadStream { - fn stable_stream_len(&mut self) -> io::Result { Ok(self.disc_size) } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - -impl DiscIO for DiscIONFS { - fn open(&self) -> Result> { - let mut stream = NFSReadStream { - inner: self.inner.clone(), - header: self.header.clone(), - disc_header: None, - disc_size: self.disc_size, - pos: 0, - sector: u32::MAX, - buf: [0; SECTOR_SIZE], - key: self.key, - part_info: vec![], - }; - let mut disc_header: DiscHeader = read_from(&mut stream).context("Reading disc header")?; - if !self.encrypt { - // If we're not re-encrypting, disable partition encryption in disc header - disc_header.no_partition_encryption = 1; - } - - // Read partition info so we can re-encrypt - if self.encrypt && disc_header.is_wii() { - for part in read_partition_info(&mut stream)? { - let start = part.offset + part.header.data_off(); - let end = start + part.header.data_size(); - if start % SECTOR_SIZE as u64 != 0 || end % SECTOR_SIZE as u64 != 0 { - return Err(Error::DiscFormat(format!( - "Partition start / end not aligned to sector size: {} / {}", - start, end - ))); - } - stream.part_info.push(PartitionInfo { - start_sector: (start / SECTOR_SIZE as u64) as u32, - end_sector: (end / SECTOR_SIZE as u64) as u32, - key: part.header.ticket.title_key, - }); - } - } - - stream.disc_header = Some(disc_header); - // Reset stream position - stream.pos = 0; - stream.sector = u32::MAX; - Ok(Box::new(stream)) + fn meta(&self) -> Result { + Ok(DiscMeta { decrypted: true, lossless: true, ..Default::default() }) } - - fn disc_size(&self) -> Option { None } } fn get_path

(directory: &Path, path: P) -> PathBuf diff --git a/src/io/nkit.rs b/src/io/nkit.rs index 12e997d..acebb84 100644 --- a/src/io/nkit.rs +++ b/src/io/nkit.rs @@ -4,8 +4,9 @@ use std::{ }; use crate::{ + disc::DL_DVD_SIZE, io::MagicBytes, - util::reader::{read_from, read_u16_be, read_u32_be, read_u64_be, read_vec}, + util::read::{read_from, read_u16_be, read_u32_be, read_u64_be, read_vec}, DiscMeta, }; @@ -65,17 +66,19 @@ pub struct NKitHeader { pub md5: Option<[u8; 16]>, pub sha1: Option<[u8; 20]>, pub xxhash64: Option, + /// Bitstream of blocks that are junk data + pub junk_bits: Option>, } const VERSION_PREFIX: [u8; 7] = *b"NKIT v"; impl NKitHeader { - pub fn try_read_from(reader: &mut R) -> Option + pub fn try_read_from(reader: &mut R, block_size: u32, has_junk_bits: bool) -> Option where R: Read + Seek + ?Sized { let magic: MagicBytes = read_from(reader).ok()?; if magic == *b"NKIT" { reader.seek(SeekFrom::Current(-4)).ok()?; - match NKitHeader::read_from(reader) { + match NKitHeader::read_from(reader, block_size, has_junk_bits) { Ok(header) => Some(header), Err(e) => { log::warn!("Failed to read NKit header: {}", e); @@ -87,7 +90,7 @@ impl NKitHeader { } } - pub fn read_from(reader: &mut R) -> io::Result + pub fn read_from(reader: &mut R, block_size: u32, has_junk_bits: bool) -> io::Result where R: Read + ?Sized { let version_string: [u8; 8] = read_from(reader)?; if version_string[0..7] != VERSION_PREFIX @@ -117,30 +120,54 @@ impl NKitHeader { remaining_header_size -= 2; } let header_bytes = read_vec(reader, remaining_header_size)?; - let mut reader = &header_bytes[..]; + let mut inner = &header_bytes[..]; - let flags = if version == 1 { NKIT_HEADER_V1_FLAGS } else { read_u16_be(&mut reader)? }; + let flags = if version == 1 { NKIT_HEADER_V1_FLAGS } else { read_u16_be(&mut inner)? }; let size = (flags & NKitHeaderFlags::Size as u16 != 0) - .then(|| read_u64_be(&mut reader)) + .then(|| read_u64_be(&mut inner)) .transpose()?; let crc32 = (flags & NKitHeaderFlags::Crc32 as u16 != 0) - .then(|| read_u32_be(&mut reader)) + .then(|| read_u32_be(&mut inner)) .transpose()?; let md5 = (flags & NKitHeaderFlags::Md5 as u16 != 0) - .then(|| read_from::<[u8; 16], _>(&mut reader)) + .then(|| read_from::<[u8; 16], _>(&mut inner)) .transpose()?; let sha1 = (flags & NKitHeaderFlags::Sha1 as u16 != 0) - .then(|| read_from::<[u8; 20], _>(&mut reader)) + .then(|| read_from::<[u8; 20], _>(&mut inner)) .transpose()?; let xxhash64 = (flags & NKitHeaderFlags::Xxhash64 as u16 != 0) - .then(|| read_u64_be(&mut reader)) + .then(|| read_u64_be(&mut inner)) .transpose()?; - Ok(Self { version, flags, size, crc32, md5, sha1, xxhash64 }) + + let junk_bits = if has_junk_bits { + let n = DL_DVD_SIZE.div_ceil(block_size as u64).div_ceil(8); + Some(read_vec(reader, n as usize)?) + } else { + None + }; + + Ok(Self { version, flags, size, crc32, md5, sha1, xxhash64, junk_bits }) + } + + pub fn is_junk_block(&self, block: u32) -> Option { + self.junk_bits + .as_ref() + .and_then(|v| v.get((block / 8) as usize)) + .map(|&b| b & (1 << (7 - (block & 7))) != 0) } } impl From<&NKitHeader> for DiscMeta { fn from(value: &NKitHeader) -> Self { - Self { crc32: value.crc32, md5: value.md5, sha1: value.sha1, xxhash64: value.xxhash64 } + Self { + needs_hash_recovery: value.junk_bits.is_some(), + lossless: value.size.is_some() && value.junk_bits.is_some(), + disc_size: value.size, + crc32: value.crc32, + md5: value.md5, + sha1: value.sha1, + xxhash64: value.xxhash64, + ..Default::default() + } } } diff --git a/src/io/split.rs b/src/io/split.rs index 8b8fb2b..b25444b 100644 --- a/src/io/split.rs +++ b/src/io/split.rs @@ -1,16 +1,17 @@ use std::{ + cmp::min, fs::File, io, - io::{Read, Seek, SeekFrom}, + io::{BufReader, Read, Seek, SeekFrom}, path::{Path, PathBuf}, }; -use crate::{ErrorContext, ReadStream, Result, ResultContext}; +use crate::{ErrorContext, Result, ResultContext}; #[derive(Debug)] pub struct SplitFileReader { files: Vec>, - open_file: Option>, + open_file: Option>>, pos: u64, } @@ -108,30 +109,24 @@ impl SplitFileReader { } impl Read for SplitFileReader { - fn read(&mut self, mut buf: &mut [u8]) -> io::Result { - let mut total = 0; - while !buf.is_empty() { - if let Some(split) = &mut self.open_file { - let n = buf.len().min((split.begin + split.size - self.pos) as usize); - if n == 0 { - self.open_file = None; - continue; - } - split.inner.read_exact(&mut buf[..n])?; - total += n; - self.pos += n as u64; - buf = &mut buf[n..]; - } else if let Some(split) = self.files.iter().find(|f| f.contains(self.pos)) { - let mut file = File::open(&split.inner)?; - if self.pos > split.begin { - file.seek(SeekFrom::Start(self.pos - split.begin))?; - } - self.open_file = Some(Split { inner: file, begin: split.begin, size: split.size }); + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if self.open_file.is_none() || !self.open_file.as_ref().unwrap().contains(self.pos) { + self.open_file = if let Some(split) = self.files.iter().find(|f| f.contains(self.pos)) { + let mut file = BufReader::new(File::open(&split.inner)?); + // log::info!("Opened file {} at pos {}", split.inner.display(), self.pos); + file.seek(SeekFrom::Start(self.pos - split.begin))?; + Some(Split { inner: file, begin: split.begin, size: split.size }) } else { - break; - } + None + }; } - Ok(total) + let Some(split) = self.open_file.as_mut() else { + return Ok(0); + }; + let to_read = min(buf.len(), (split.begin + split.size - self.pos) as usize); + let read = split.inner.read(&mut buf[..to_read])?; + self.pos += read as u64; + Ok(read) } } @@ -154,12 +149,6 @@ impl Seek for SplitFileReader { } } -impl ReadStream for SplitFileReader { - fn stable_stream_len(&mut self) -> io::Result { Ok(self.len()) } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - impl Clone for SplitFileReader { - fn clone(&self) -> Self { Self { files: self.files.clone(), open_file: None, pos: self.pos } } + fn clone(&self) -> Self { Self { files: self.files.clone(), open_file: None, pos: 0 } } } diff --git a/src/io/wbfs.rs b/src/io/wbfs.rs index f935585..c7d14cb 100644 --- a/src/io/wbfs.rs +++ b/src/io/wbfs.rs @@ -1,7 +1,6 @@ use std::{ - cmp::min, io, - io::{BufReader, Read, Seek, SeekFrom}, + io::{Read, Seek, SeekFrom}, mem::size_of, path::Path, }; @@ -9,10 +8,14 @@ use std::{ use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes}; use crate::{ - disc::SECTOR_SIZE, - io::{nkit::NKitHeader, split::SplitFileReader, DiscIO, DiscMeta, MagicBytes}, - util::reader::{read_from, read_vec}, - Error, ReadStream, Result, ResultContext, + io::{ + block::{BPartitionInfo, Block, BlockIO}, + nkit::NKitHeader, + split::SplitFileReader, + DiscMeta, MagicBytes, + }, + util::read::{read_box_slice, read_from}, + Error, Result, ResultContext, }; pub const WBFS_MAGIC: MagicBytes = *b"WBFS"; @@ -23,14 +26,14 @@ struct WBFSHeader { magic: MagicBytes, num_sectors: U32, sector_size_shift: u8, - wbfs_sector_size_shift: u8, + block_size_shift: u8, _pad: [u8; 2], } impl WBFSHeader { fn sector_size(&self) -> u32 { 1 << self.sector_size_shift } - fn wbfs_sector_size(&self) -> u32 { 1 << self.wbfs_sector_size_shift } + fn block_size(&self) -> u32 { 1 << self.block_size_shift } // fn align_lba(&self, x: u32) -> u32 { (x + self.sector_size() - 1) & !(self.sector_size() - 1) } // @@ -44,34 +47,32 @@ impl WBFSHeader { // self.num_wii_sectors() >> (self.wbfs_sector_size_shift - 15) // } - fn max_wbfs_sectors(&self) -> u32 { NUM_WII_SECTORS >> (self.wbfs_sector_size_shift - 15) } + fn max_blocks(&self) -> u32 { NUM_WII_SECTORS >> (self.block_size_shift - 15) } } const DISC_HEADER_SIZE: usize = 0x100; const NUM_WII_SECTORS: u32 = 143432 * 2; // Double layer discs +#[derive(Clone)] pub struct DiscIOWBFS { - pub inner: SplitFileReader, + inner: SplitFileReader, /// WBFS header header: WBFSHeader, /// Map of Wii LBAs to WBFS LBAs - wlba_table: Vec, + block_table: Box<[U16]>, /// Optional NKit header nkit_header: Option, } impl DiscIOWBFS { - pub fn new(filename: &Path) -> Result { - let mut inner = BufReader::new(SplitFileReader::new(filename)?); + pub fn new(filename: &Path) -> Result> { + let mut inner = SplitFileReader::new(filename)?; let header: WBFSHeader = read_from(&mut inner).context("Reading WBFS header")?; if header.magic != WBFS_MAGIC { return Err(Error::DiscFormat("Invalid WBFS magic".to_string())); } - // log::debug!("{:?}", header); - // log::debug!("sector_size: {}", header.sector_size()); - // log::debug!("wbfs_sector_size: {}", header.wbfs_sector_size()); - let file_len = inner.stable_stream_len().context("Getting WBFS file size")?; + let file_len = inner.len(); let expected_file_len = header.num_sectors.get() as u64 * header.sector_size() as u64; if file_len != expected_file_len { return Err(Error::DiscFormat(format!( @@ -80,8 +81,8 @@ impl DiscIOWBFS { ))); } - let disc_table: Vec = - read_vec(&mut inner, header.sector_size() as usize - size_of::()) + let disc_table: Box<[u8]> = + read_box_slice(&mut inner, header.sector_size() as usize - size_of::()) .context("Reading WBFS disc table")?; if disc_table[0] != 1 { return Err(Error::DiscFormat("WBFS doesn't contain a disc".to_string())); @@ -94,110 +95,46 @@ impl DiscIOWBFS { inner .seek(SeekFrom::Start(header.sector_size() as u64 + DISC_HEADER_SIZE as u64)) .context("Seeking to WBFS LBA table")?; // Skip header - let wlba_table: Vec = read_vec(&mut inner, header.max_wbfs_sectors() as usize) + let block_table: Box<[U16]> = read_box_slice(&mut inner, header.max_blocks() as usize) .context("Reading WBFS LBA table")?; // Read NKit header if present (always at 0x10000) inner.seek(SeekFrom::Start(0x10000)).context("Seeking to NKit header")?; - let nkit_header = NKitHeader::try_read_from(&mut inner); + let nkit_header = NKitHeader::try_read_from(&mut inner, header.block_size(), true); // Reset reader - let mut inner = inner.into_inner(); inner.reset(); - Ok(Self { inner, header, wlba_table, nkit_header }) + Ok(Box::new(Self { inner, header, block_table, nkit_header })) } } -impl DiscIO for DiscIOWBFS { - fn open(&self) -> Result> { - Ok(Box::new(WBFSReadStream { - inner: BufReader::new(self.inner.clone()), - header: self.header.clone(), - wlba_table: self.wlba_table.clone(), - wlba: u32::MAX, - pos: 0, - disc_size: self.nkit_header.as_ref().and_then(|h| h.size), - })) +impl BlockIO for DiscIOWBFS { + fn read_block( + &mut self, + out: &mut [u8], + block: u32, + _partition: Option<&BPartitionInfo>, + ) -> io::Result> { + let block_size = self.header.block_size(); + if block >= self.header.max_blocks() { + return Ok(None); + } + + // Check if block is junk data + if self.nkit_header.as_ref().is_some_and(|h| h.is_junk_block(block).unwrap_or(false)) { + return Ok(Some(Block::Junk)); + } + + // Read block + let block_start = block_size as u64 * self.block_table[block as usize].get() as u64; + self.inner.seek(SeekFrom::Start(block_start))?; + self.inner.read_exact(out)?; + Ok(Some(Block::Raw)) } + fn block_size(&self) -> u32 { self.header.block_size() } + fn meta(&self) -> Result { Ok(self.nkit_header.as_ref().map(DiscMeta::from).unwrap_or_default()) } - - fn disc_size(&self) -> Option { self.nkit_header.as_ref().and_then(|h| h.size) } -} - -struct WBFSReadStream { - /// File reader - inner: BufReader, - /// WBFS header - header: WBFSHeader, - /// Map of Wii LBAs to WBFS LBAs - wlba_table: Vec, - /// Current WBFS LBA - wlba: u32, - /// Current stream offset - pos: u64, - /// Optional known size - disc_size: Option, -} - -impl WBFSReadStream { - fn disc_size(&self) -> u64 { - self.disc_size.unwrap_or(NUM_WII_SECTORS as u64 * SECTOR_SIZE as u64) - } -} - -impl Read for WBFSReadStream { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - let wlba = (self.pos >> self.header.wbfs_sector_size_shift) as u32; - let wlba_size = self.header.wbfs_sector_size() as u64; - let wlba_offset = self.pos & (wlba_size - 1); - if wlba != self.wlba { - if self.pos >= self.disc_size() || wlba >= self.header.max_wbfs_sectors() { - return Ok(0); - } - let wlba_start = wlba_size * self.wlba_table[wlba as usize].get() as u64; - self.inner.seek(SeekFrom::Start(wlba_start + wlba_offset))?; - self.wlba = wlba; - } - - let to_read = min(buf.len(), (wlba_size - wlba_offset) as usize); - let read = self.inner.read(&mut buf[..to_read])?; - self.pos += read as u64; - Ok(read) - } -} - -impl Seek for WBFSReadStream { - fn seek(&mut self, pos: SeekFrom) -> io::Result { - let new_pos = match pos { - SeekFrom::Start(v) => v, - SeekFrom::End(_) => { - return Err(io::Error::new( - io::ErrorKind::Unsupported, - "WBFSReadStream: SeekFrom::End is not supported", - )); - } - SeekFrom::Current(v) => self.pos.saturating_add_signed(v), - }; - - let new_wlba = (self.pos >> self.header.wbfs_sector_size_shift) as u32; - if new_wlba == self.wlba { - // Seek within the same WBFS LBA - self.inner.seek(SeekFrom::Current(new_pos as i64 - self.pos as i64))?; - } else { - // Seek to a different WBFS LBA, handled by next read - self.wlba = u32::MAX; - } - - self.pos = new_pos; - Ok(new_pos) - } -} - -impl ReadStream for WBFSReadStream { - fn stable_stream_len(&mut self) -> io::Result { Ok(self.disc_size()) } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } diff --git a/src/io/wia.rs b/src/io/wia.rs index 88134cf..53ad81a 100644 --- a/src/io/wia.rs +++ b/src/io/wia.rs @@ -1,34 +1,32 @@ use std::{ - cmp::min, - fs::File, io, - io::{BufReader, Read, Seek, SeekFrom}, + io::{Read, Seek, SeekFrom}, mem::size_of, - path::{Path, PathBuf}, - sync::{Arc, Mutex}, - time::Instant, + path::Path, }; -use rayon::iter::{IntoParallelIterator, ParallelIterator}; use sha1::{Digest, Sha1}; use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes}; use crate::{ - array_ref, array_ref_mut, disc::{ - wii::{BLOCK_SIZE, HASHES_SIZE}, + wii::{HASHES_SIZE, SECTOR_DATA_SIZE}, SECTOR_SIZE, }, - io::{aes_encrypt, nkit::NKitHeader, DiscIO, DiscMeta, HashBytes, KeyBytes, MagicBytes}, + io::{ + block::{BPartitionInfo, Block, BlockIO}, + nkit::NKitHeader, + split::SplitFileReader, + HashBytes, KeyBytes, MagicBytes, + }, static_assert, - streams::ReadStream, util::{ compress::{lzma2_props_decode, lzma_props_decode, new_lzma2_decoder, new_lzma_decoder}, lfg::LaggedFibonacci, - reader::{read_from, read_u16_be, read_vec}, + read::{read_box_slice, read_from, read_u16_be, read_vec}, take_seek::TakeSeekExt, }, - Error, OpenOptions, Result, ResultContext, + DiscMeta, Error, OpenOptions, Result, ResultContext, }; pub const WIA_MAGIC: MagicBytes = *b"WIA\x01"; @@ -262,6 +260,13 @@ pub struct WIAPartitionData { static_assert!(size_of::() == 0x10); +impl WIAPartitionData { + pub fn contains(&self, sector: u32) -> bool { + let start = self.first_sector.get(); + sector >= start && sector < start + self.num_sectors.get() + } +} + /// This struct is used for keeping track of Wii partition data that on the actual disc is encrypted /// and hashed. This does not include the unencrypted area at the beginning of partitions that /// contains the ticket, TMD, certificate chain, and H3 table. So for a typical game partition, @@ -313,6 +318,20 @@ pub struct WIARawData { pub num_groups: U32, } +impl WIARawData { + pub fn start_offset(&self) -> u64 { self.raw_data_offset.get() & !(SECTOR_SIZE as u64 - 1) } + + pub fn start_sector(&self) -> u32 { (self.start_offset() / SECTOR_SIZE as u64) as u32 } + + pub fn end_offset(&self) -> u64 { self.raw_data_offset.get() + self.raw_data_size.get() } + + pub fn end_sector(&self) -> u32 { (self.end_offset() / SECTOR_SIZE as u64) as u32 } + + pub fn contains(&self, sector: u32) -> bool { + sector >= self.start_sector() && sector < self.end_sector() + } +} + /// This struct points directly to the actual disc data, stored compressed. /// /// The data is interpreted differently depending on whether the [WIAGroup] is referenced by a @@ -360,8 +379,8 @@ impl RVZGroup { pub fn is_compressed(&self) -> bool { self.data_size_and_flag.get() & 0x80000000 != 0 } } -impl From for RVZGroup { - fn from(value: WIAGroup) -> Self { +impl From<&WIAGroup> for RVZGroup { + fn from(value: &WIAGroup) -> Self { Self { data_offset: value.data_offset, data_size_and_flag: U32::new(value.data_size.get() | 0x80000000), @@ -428,31 +447,32 @@ pub struct WIAException { /// end offset of the last [WIAExceptionList] is not evenly divisible by 4, padding is inserted /// after it so that the data afterwards will start at a 4 byte boundary. This padding is not /// inserted for the other compression methods. -type WIAExceptionList = Vec; +type WIAExceptionList = Box<[WIAException]>; +#[derive(Clone)] pub enum Decompressor { None, #[cfg(feature = "compress-bzip2")] Bzip2, #[cfg(feature = "compress-lzma")] - Lzma(liblzma::stream::LzmaOptions), + Lzma(Box<[u8]>), #[cfg(feature = "compress-lzma")] - Lzma2(liblzma::stream::LzmaOptions), + Lzma2(Box<[u8]>), #[cfg(feature = "compress-zstd")] Zstandard, } impl Decompressor { pub fn new(disc: &WIADisc) -> Result { - let compr_data = &disc.compr_data[..disc.compr_data_len as usize]; + let data = &disc.compr_data[..disc.compr_data_len as usize]; match disc.compression() { Compression::None => Ok(Self::None), #[cfg(feature = "compress-bzip2")] Compression::Bzip2 => Ok(Self::Bzip2), #[cfg(feature = "compress-lzma")] - Compression::Lzma => Ok(Self::Lzma(lzma_props_decode(compr_data)?)), + Compression::Lzma => Ok(Self::Lzma(Box::from(data))), #[cfg(feature = "compress-lzma")] - Compression::Lzma2 => Ok(Self::Lzma2(lzma2_props_decode(compr_data)?)), + Compression::Lzma2 => Ok(Self::Lzma2(Box::from(data))), #[cfg(feature = "compress-zstd")] Compression::Zstandard => Ok(Self::Zstandard), comp => Err(Error::DiscFormat(format!("Unsupported WIA/RVZ compression: {:?}", comp))), @@ -466,86 +486,51 @@ impl Decompressor { #[cfg(feature = "compress-bzip2")] Decompressor::Bzip2 => Box::new(bzip2::read::BzDecoder::new(reader)), #[cfg(feature = "compress-lzma")] - Decompressor::Lzma(options) => Box::new(new_lzma_decoder(reader, options)?), + Decompressor::Lzma(data) => { + let options = lzma_props_decode(data)?; + Box::new(new_lzma_decoder(reader, &options)?) + } #[cfg(feature = "compress-lzma")] - Decompressor::Lzma2(options) => Box::new(new_lzma2_decoder(reader, options)?), + Decompressor::Lzma2(data) => { + let options = lzma2_props_decode(data)?; + Box::new(new_lzma2_decoder(reader, &options)?) + } #[cfg(feature = "compress-zstd")] Decompressor::Zstandard => Box::new(zstd::stream::Decoder::new(reader)?), }) } } -/// In a sector, following the 0x400 byte block of hashes, each 0x400 bytes of decrypted data is -/// hashed, yielding 31 H0 hashes. -/// Then, 8 sectors are aggregated into a subgroup, and the 31 H0 hashes for each sector are hashed, -/// yielding 8 H1 hashes. -/// Then, 8 subgroups are aggregated into a group, and the 8 H1 hashes for each subgroup are hashed, -/// yielding 8 H2 hashes. -/// Finally, the 8 H2 hashes for each group are hashed, yielding 1 H3 hash. -/// The H3 hashes for each group are stored in the partition's H3 table. -pub struct HashTable { - /// SHA-1 hash of the 31 H0 hashes for each sector. - pub h1_hashes: Vec, - /// SHA-1 hash of the 8 H1 hashes for each subgroup. - pub h2_hashes: Vec, - /// SHA-1 hash of the 8 H2 hashes for each group. - pub h3_hashes: Vec, +pub struct DiscIOWIA { + inner: SplitFileReader, + header: WIAFileHeader, + disc: WIADisc, + partitions: Box<[WIAPartition]>, + raw_data: Box<[WIARawData]>, + groups: Box<[RVZGroup]>, + nkit_header: Option, + decompressor: Decompressor, + group: u32, + group_data: Vec, + exception_lists: Vec, } -struct HashResult { - h1_hashes: [HashBytes; 64], - h2_hashes: [HashBytes; 8], - h3_hash: HashBytes, -} - -impl HashTable { - fn new(num_sectors: u32) -> Self { - let num_sectors = num_sectors.next_multiple_of(64) as usize; - let num_subgroups = num_sectors / 8; - let num_groups = num_subgroups / 8; +impl Clone for DiscIOWIA { + fn clone(&self) -> Self { Self { - h1_hashes: HashBytes::new_vec_zeroed(num_sectors), - h2_hashes: HashBytes::new_vec_zeroed(num_subgroups), - h3_hashes: HashBytes::new_vec_zeroed(num_groups), + header: self.header.clone(), + disc: self.disc.clone(), + partitions: self.partitions.clone(), + raw_data: self.raw_data.clone(), + groups: self.groups.clone(), + inner: self.inner.clone(), + nkit_header: self.nkit_header.clone(), + decompressor: self.decompressor.clone(), + group: u32::MAX, + group_data: Vec::new(), + exception_lists: Vec::new(), } } - - fn extend(&mut self, group_index: usize, result: &HashResult) { - let h1_start = group_index * 64; - self.h1_hashes[h1_start..h1_start + 64].copy_from_slice(&result.h1_hashes); - let h2_start = group_index * 8; - self.h2_hashes[h2_start..h2_start + 8].copy_from_slice(&result.h2_hashes); - self.h3_hashes[group_index] = result.h3_hash; - } -} - -pub struct DiscIOWIA { - pub header: WIAFileHeader, - pub disc: WIADisc, - pub partitions: Vec, - pub raw_data: Vec, - pub groups: Vec, - pub filename: PathBuf, - pub encrypt: bool, - pub hash_tables: Vec, - pub nkit_header: Option, -} - -#[derive(Debug)] -struct GroupResult { - /// Offset of the group in the raw disc image. - disc_offset: u64, - /// Data offset of the group within a partition, excluding hashes. - /// Same as `disc_offset` for raw data or GameCube discs. - partition_offset: u64, - /// The group. - group: RVZGroup, - /// The index of the Wii partition that this group belongs to. - partition_index: Option, - /// Chunk size, differs between Wii and raw data. - chunk_size: u32, - /// End offset for the partition or raw data. - partition_end: u64, } #[inline] @@ -571,296 +556,127 @@ fn verify_hash(buf: &[u8], expected: &HashBytes) -> Result<()> { } impl DiscIOWIA { - pub fn new(filename: &Path, options: &OpenOptions) -> Result { - let mut file = BufReader::new( - File::open(filename).with_context(|| format!("Opening file {}", filename.display()))?, - ); + pub fn new(filename: &Path, _options: &OpenOptions) -> Result> { + let mut inner = SplitFileReader::new(filename)?; // Load & verify file header - let header: WIAFileHeader = read_from(&mut file).context("Reading WIA/RVZ file header")?; + let header: WIAFileHeader = read_from(&mut inner).context("Reading WIA/RVZ file header")?; header.validate()?; let is_rvz = header.is_rvz(); // log::debug!("Header: {:?}", header); // Load & verify disc header - let mut disc_buf: Vec = read_vec(&mut file, header.disc_size.get() as usize) + let mut disc_buf: Vec = read_vec(&mut inner, header.disc_size.get() as usize) .context("Reading WIA/RVZ disc header")?; verify_hash(&disc_buf, &header.disc_hash)?; disc_buf.resize(size_of::(), 0); - let mut disc = WIADisc::read_from(disc_buf.as_slice()).unwrap(); + let disc = WIADisc::read_from(disc_buf.as_slice()).unwrap(); disc.validate()?; - if !options.rebuild_encryption { - // If we're not re-encrypting, disable partition encryption in disc header - disc.disc_head[0x61] = 1; - } + // if !options.rebuild_hashes { + // // If we're not rebuilding hashes, disable partition hashes in disc header + // disc.disc_head[0x60] = 1; + // } + // if !options.rebuild_encryption { + // // If we're not re-encrypting, disable partition encryption in disc header + // disc.disc_head[0x61] = 1; + // } // log::debug!("Disc: {:?}", disc); // Read NKit header if present (after disc header) - let nkit_header = NKitHeader::try_read_from(&mut file); + let nkit_header = NKitHeader::try_read_from(&mut inner, disc.chunk_size.get(), false); // Load & verify partition headers - file.seek(SeekFrom::Start(disc.partition_offset.get())) + inner + .seek(SeekFrom::Start(disc.partition_offset.get())) .context("Seeking to WIA/RVZ partition headers")?; - let partitions: Vec = read_vec(&mut file, disc.num_partitions.get() as usize) - .context("Reading WIA/RVZ partition headers")?; - verify_hash(partitions.as_slice().as_bytes(), &disc.partition_hash)?; + let partitions: Box<[WIAPartition]> = + read_box_slice(&mut inner, disc.num_partitions.get() as usize) + .context("Reading WIA/RVZ partition headers")?; + verify_hash(partitions.as_ref().as_bytes(), &disc.partition_hash)?; // log::debug!("Partitions: {:?}", partitions); // Create decompressor let mut decompressor = Decompressor::new(&disc)?; // Load raw data headers - let raw_data: Vec = { - file.seek(SeekFrom::Start(disc.raw_data_offset.get())) + let raw_data: Box<[WIARawData]> = { + inner + .seek(SeekFrom::Start(disc.raw_data_offset.get())) .context("Seeking to WIA/RVZ raw data headers")?; let mut reader = decompressor - .wrap((&mut file).take(disc.raw_data_size.get() as u64)) + .wrap((&mut inner).take(disc.raw_data_size.get() as u64)) .context("Creating WIA/RVZ decompressor")?; - read_vec(&mut reader, disc.num_raw_data.get() as usize) + read_box_slice(&mut reader, disc.num_raw_data.get() as usize) .context("Reading WIA/RVZ raw data headers")? }; + // Validate raw data alignment + for (idx, rd) in raw_data.iter().enumerate() { + let start_offset = rd.start_offset(); + let end_offset = rd.end_offset(); + if (start_offset % SECTOR_SIZE as u64) != 0 || (end_offset % SECTOR_SIZE as u64) != 0 { + return Err(Error::DiscFormat(format!( + "WIA/RVZ raw data {} not aligned to sector: {:#X}..{:#X}", + idx, start_offset, end_offset + ))); + } + } // log::debug!("Raw data: {:?}", raw_data); // Load group headers let groups = { - file.seek(SeekFrom::Start(disc.group_offset.get())) + inner + .seek(SeekFrom::Start(disc.group_offset.get())) .context("Seeking to WIA/RVZ group headers")?; let mut reader = decompressor - .wrap((&mut file).take(disc.group_size.get() as u64)) + .wrap((&mut inner).take(disc.group_size.get() as u64)) .context("Creating WIA/RVZ decompressor")?; if is_rvz { - read_vec(&mut reader, disc.num_groups.get() as usize) + read_box_slice(&mut reader, disc.num_groups.get() as usize) .context("Reading WIA/RVZ group headers")? } else { - let wia_groups: Vec = - read_vec(&mut reader, disc.num_groups.get() as usize) + let wia_groups: Box<[WIAGroup]> = + read_box_slice(&mut reader, disc.num_groups.get() as usize) .context("Reading WIA/RVZ group headers")?; - wia_groups.into_iter().map(RVZGroup::from).collect() + wia_groups.iter().map(RVZGroup::from).collect() } // log::debug!("Groups: {:?}", groups); }; - let mut disc_io = Self { + Ok(Box::new(Self { header, disc, partitions, raw_data, groups, - filename: filename.to_owned(), - encrypt: options.rebuild_encryption, - hash_tables: vec![], + inner, nkit_header, - }; - if options.rebuild_hashes { - disc_io.rebuild_hashes()?; - } - Ok(disc_io) + decompressor, + group: u32::MAX, + group_data: vec![], + exception_lists: vec![], + })) } - - fn group_for_offset(&self, offset: u64) -> Option { - if let Some((p_idx, pd)) = self.partitions.iter().enumerate().find_map(|(p_idx, p)| { - p.partition_data - .iter() - .find(|pd| { - let start = pd.first_sector.get() as u64 * SECTOR_SIZE as u64; - let end = start + pd.num_sectors.get() as u64 * SECTOR_SIZE as u64; - offset >= start && offset < end - }) - .map(|pd| (p_idx, pd)) - }) { - let start = pd.first_sector.get() as u64 * SECTOR_SIZE as u64; - let group_index = (offset - start) / self.disc.chunk_size.get() as u64; - if group_index >= pd.num_groups.get() as u64 { - return None; - } - let disc_offset = start + group_index * self.disc.chunk_size.get() as u64; - let chunk_size = - (self.disc.chunk_size.get() as u64 * BLOCK_SIZE as u64) / SECTOR_SIZE as u64; - let partition_offset = group_index * chunk_size; - let partition_end = pd.num_sectors.get() as u64 * BLOCK_SIZE as u64; - self.groups.get(pd.group_index.get() as usize + group_index as usize).map(|g| { - GroupResult { - disc_offset, - partition_offset, - group: g.clone(), - partition_index: Some(p_idx), - chunk_size: chunk_size as u32, - partition_end, - } - }) - } else if let Some(d) = self.raw_data.iter().find(|d| { - let start = d.raw_data_offset.get() & !0x7FFF; - let end = d.raw_data_offset.get() + d.raw_data_size.get(); - offset >= start && offset < end - }) { - let start = d.raw_data_offset.get() & !0x7FFF; - let end = d.raw_data_offset.get() + d.raw_data_size.get(); - let group_index = (offset - start) / self.disc.chunk_size.get() as u64; - if group_index >= d.num_groups.get() as u64 { - return None; - } - let disc_offset = start + group_index * self.disc.chunk_size.get() as u64; - self.groups.get(d.group_index.get() as usize + group_index as usize).map(|g| { - GroupResult { - disc_offset, - partition_offset: disc_offset, - group: g.clone(), - partition_index: None, - chunk_size: self.disc.chunk_size.get(), - partition_end: end, - } - }) - } else { - None - } - } - - pub fn rebuild_hashes(&mut self) -> Result<()> { - const NUM_H0_HASHES: usize = BLOCK_SIZE / HASHES_SIZE; - const H0_HASHES_SIZE: usize = size_of::() * NUM_H0_HASHES; - - let start = Instant::now(); - - // Precompute hashes for zeroed sectors. - const ZERO_H0_BYTES: &[u8] = &[0u8; HASHES_SIZE]; - let zero_h0_hash = hash_bytes(ZERO_H0_BYTES); - let mut zero_h1_hash = Sha1::new(); - for _ in 0..NUM_H0_HASHES { - zero_h1_hash.update(zero_h0_hash); - } - let zero_h1_hash: HashBytes = zero_h1_hash.finalize().into(); - - let mut hash_tables = Vec::with_capacity(self.partitions.len()); - for part in &self.partitions { - let first_sector = part.partition_data[0].first_sector.get(); - if first_sector + part.partition_data[0].num_sectors.get() - != part.partition_data[1].first_sector.get() - { - return Err(Error::DiscFormat(format!( - "Partition data is not contiguous: {}..{} != {}", - first_sector, - first_sector + part.partition_data[0].num_sectors.get(), - part.partition_data[1].first_sector.get() - ))); - } - - let part_sectors = - part.partition_data[0].num_sectors.get() + part.partition_data[1].num_sectors.get(); - let hash_table = HashTable::new(part_sectors); - log::debug!( - "Rebuilding hashes: {} sectors, {} subgroups, {} groups", - hash_table.h1_hashes.len(), - hash_table.h2_hashes.len(), - hash_table.h3_hashes.len() - ); - - let group_count = hash_table.h3_hashes.len(); - let mutex = Arc::new(Mutex::new(hash_table)); - (0..group_count).into_par_iter().try_for_each_init( - || (WIAReadStream::new(self, false), mutex.clone()), - |(stream, mutex), h3_index| -> Result<()> { - let stream = stream.as_mut().map_err(|_| { - Error::DiscFormat("Failed to create read stream".to_string()) - })?; - let mut result = HashResult { - h1_hashes: [HashBytes::default(); 64], - h2_hashes: [HashBytes::default(); 8], - h3_hash: HashBytes::default(), - }; - let mut h0_buf = [0u8; H0_HASHES_SIZE]; - let mut h3_hasher = Sha1::new(); - for h2_index in 0..8 { - let mut h2_hasher = Sha1::new(); - for h1_index in 0..8 { - let part_sector = - h1_index as u32 + h2_index as u32 * 8 + h3_index as u32 * 64; - let h1_hash = if part_sector >= part_sectors { - zero_h1_hash - } else { - let sector = first_sector + part_sector; - stream - .seek(SeekFrom::Start(sector as u64 * SECTOR_SIZE as u64)) - .with_context(|| format!("Seeking to sector {}", sector))?; - stream - .read_exact(&mut h0_buf) - .with_context(|| format!("Reading sector {}", sector))?; - hash_bytes(&h0_buf) - }; - result.h1_hashes[h1_index + h2_index * 8] = h1_hash; - h2_hasher.update(h1_hash); - } - let h2_hash = h2_hasher.finalize().into(); - result.h2_hashes[h2_index] = h2_hash; - h3_hasher.update(h2_hash); - } - result.h3_hash = h3_hasher.finalize().into(); - let mut hash_table = mutex.lock().map_err(|_| "Failed to lock mutex")?; - hash_table.extend(h3_index, &result); - Ok(()) - }, - )?; - - let hash_table = Arc::try_unwrap(mutex) - .map_err(|_| "Failed to unwrap Arc")? - .into_inner() - .map_err(|_| "Failed to lock mutex")?; - hash_tables.push(hash_table); - } - self.hash_tables = hash_tables; - log::info!("Rebuilt hashes in {:?}", start.elapsed()); - Ok(()) - } -} - -impl DiscIO for DiscIOWIA { - fn open(&self) -> Result> { - Ok(Box::new(WIAReadStream::new(self, self.encrypt)?)) - } - - fn meta(&self) -> Result { - Ok(self.nkit_header.as_ref().map(DiscMeta::from).unwrap_or_default()) - } - - fn disc_size(&self) -> Option { Some(self.header.iso_file_size.get()) } -} - -pub struct WIAReadStream<'a> { - /// The disc IO. - disc_io: &'a DiscIOWIA, - /// The currently open file handle. - file: BufReader, - /// The data read offset. - offset: u64, - /// The data offset of the current group. - group_offset: u64, - /// The current group data. - group_data: Vec, - /// Exception lists for the current group. - exception_lists: Vec, - /// The decompressor data. - decompressor: Decompressor, - /// Whether to re-encrypt Wii partition data. - encrypt: bool, } fn read_exception_lists( reader: &mut R, - partition_index: Option, + in_partition: bool, chunk_size: u32, ) -> io::Result> where R: Read + ?Sized, { - if partition_index.is_none() { + if !in_partition { return Ok(vec![]); } + // One exception list for each 2 MiB of data let num_exception_list = (chunk_size as usize).div_ceil(0x200000); // log::debug!("Num exception list: {:?}", num_exception_list); let mut exception_lists = Vec::with_capacity(num_exception_list); for i in 0..num_exception_list { let num_exceptions = read_u16_be(reader)?; - let exceptions: Vec = read_vec(reader, num_exceptions as usize)?; + let exceptions: Box<[WIAException]> = read_box_slice(reader, num_exceptions as usize)?; if !exceptions.is_empty() { log::debug!("Exception list {}: {:?}", i, exceptions); } @@ -869,244 +685,218 @@ where Ok(exception_lists) } -impl<'a> WIAReadStream<'a> { - pub fn new(disc_io: &'a DiscIOWIA, encrypt: bool) -> Result { - let file = BufReader::new( - File::open(&disc_io.filename) - .with_context(|| format!("Opening file {}", disc_io.filename.display()))?, - ); - let decompressor = Decompressor::new(&disc_io.disc)?; - let stream = Self { - disc_io, - file, - offset: 0, - group_offset: u64::MAX, - group_data: Vec::new(), - exception_lists: vec![], - decompressor, - encrypt, - }; - Ok(stream) - } +impl BlockIO for DiscIOWIA { + fn read_block( + &mut self, + out: &mut [u8], + sector: u32, + partition: Option<&BPartitionInfo>, + ) -> io::Result> { + let mut chunk_size = self.disc.chunk_size.get(); + let sectors_per_chunk = chunk_size / SECTOR_SIZE as u32; + let disc_offset = sector as u64 * SECTOR_SIZE as u64; + let mut partition_offset = disc_offset; + if let Some(partition) = partition { + // Within a partition, hashes are excluded from the data size + chunk_size = (chunk_size * SECTOR_DATA_SIZE as u32) / SECTOR_SIZE as u32; + partition_offset = + (sector - partition.data_start_sector) as u64 * SECTOR_DATA_SIZE as u64; + } - /// If the current group does not contain the current offset, load the new group. - /// Returns false if the offset is not in the disc. - fn check_group(&mut self) -> io::Result { - if self.offset < self.group_offset - || self.offset >= self.group_offset + self.group_data.len() as u64 - { - let Some(result) = self.disc_io.group_for_offset(self.offset) else { - return Ok(false); - }; - if result.disc_offset == self.group_offset { + let (group_index, group_sector) = if let Some(partition) = partition { + // Find the partition + let Some(wia_part) = self.partitions.get(partition.index as usize) else { return Err(io::Error::new( - io::ErrorKind::InvalidData, - "Group offset did not change", + io::ErrorKind::InvalidInput, + format!("Couldn't find WIA/RVZ partition index {}", partition.index), + )); + }; + + // Sanity check partition sector ranges + let wia_part_start = wia_part.partition_data[0].first_sector.get(); + let wia_part_end = wia_part.partition_data[1].first_sector.get() + + wia_part.partition_data[1].num_sectors.get(); + if partition.data_start_sector != wia_part_start + || partition.data_end_sector != wia_part_end + { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "WIA/RVZ partition sector mismatch: {}..{} != {}..{}", + wia_part_start, + wia_part_end, + partition.data_start_sector, + partition.data_end_sector + ), )); } - self.group_offset = result.disc_offset; - self.read_group(result)?; - } - Ok(true) - } - /// Reads new group data into the buffer, handling decompression and RVZ packing. - fn read_group(&mut self, result: GroupResult) -> io::Result<()> { - // Special case for all-zero data - if result.group.data_size() == 0 { - self.exception_lists.clear(); - let size = min(result.chunk_size as u64, result.partition_end - result.partition_offset) - as usize; - self.group_data = vec![0u8; size]; - self.recalculate_hashes(result)?; - return Ok(()); - } + // Find the partition data for the sector + let Some(pd) = wia_part.partition_data.iter().find(|pd| pd.contains(sector)) else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("Couldn't find WIA/RVZ partition data for sector {}", sector), + )); + }; - self.group_data = Vec::with_capacity(result.chunk_size as usize); - let group_data_start = result.group.data_offset.get() as u64 * 4; - self.file.seek(SeekFrom::Start(group_data_start))?; - - let mut reader = (&mut self.file).take_seek(result.group.data_size() as u64); - let uncompressed_exception_lists = - matches!(self.disc_io.disc.compression(), Compression::None | Compression::Purge) - || !result.group.is_compressed(); - if uncompressed_exception_lists { - self.exception_lists = read_exception_lists( - &mut reader, - result.partition_index, - self.disc_io.disc.chunk_size.get(), // result.chunk_size? - )?; - // Align to 4 - let rem = reader.stream_position()? % 4; - if rem != 0 { - reader.seek(SeekFrom::Current((4 - rem) as i64))?; + // Find the group index for the sector + let part_data_sector = sector - pd.first_sector.get(); + let part_group_index = part_data_sector / sectors_per_chunk; + let part_group_sector = part_data_sector % sectors_per_chunk; + if part_group_index >= pd.num_groups.get() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "WIA/RVZ partition group index out of range: {} >= {}", + part_group_index, + pd.num_groups.get() + ), + )); } - } - let mut reader: Box = if result.group.is_compressed() { - self.decompressor.wrap(reader)? + + (pd.group_index.get() + part_group_index, part_group_sector) } else { - Box::new(reader) + let Some(rd) = self.raw_data.iter().find(|d| d.contains(sector)) else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("Couldn't find WIA/RVZ raw data for sector {}", sector), + )); + }; + + // Find the group index for the sector + let data_sector = sector - (rd.raw_data_offset.get() / SECTOR_SIZE as u64) as u32; + let group_index = data_sector / sectors_per_chunk; + let group_sector = data_sector % sectors_per_chunk; + if group_index >= rd.num_groups.get() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "WIA/RVZ raw data group index out of range: {} >= {}", + group_index, + rd.num_groups.get() + ), + )); + } + + (rd.group_index.get() + group_index, group_sector) }; - if !uncompressed_exception_lists { - self.exception_lists = read_exception_lists( - reader.as_mut(), - result.partition_index, - self.disc_io.disc.chunk_size.get(), // result.chunk_size? - )?; + + // Fetch the group + let Some(group) = self.groups.get(group_index as usize) else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("Couldn't find WIA/RVZ group index {}", group_index), + )); + }; + + // Special case for all-zero data + if group.data_size() == 0 { + self.exception_lists.clear(); + return Ok(Some(Block::Zero)); } - if result.group.rvz_packed_size.get() > 0 { - // Decode RVZ packed data - let mut lfg = LaggedFibonacci::default(); - loop { - let mut size_bytes = [0u8; 4]; - match reader.read_exact(&mut size_bytes) { - Ok(_) => {} - Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break, - Err(e) => { - return Err(io::Error::new(e.kind(), "Failed to read RVZ packed size")); + // Read group data if necessary + if group_index != self.group { + self.group_data = Vec::with_capacity(chunk_size as usize); + let group_data_start = group.data_offset.get() as u64 * 4; + self.inner.seek(SeekFrom::Start(group_data_start))?; + + let mut reader = (&mut self.inner).take_seek(group.data_size() as u64); + let uncompressed_exception_lists = + matches!(self.disc.compression(), Compression::None | Compression::Purge) + || !group.is_compressed(); + if uncompressed_exception_lists { + self.exception_lists = read_exception_lists( + &mut reader, + partition.is_some(), + self.disc.chunk_size.get(), + )?; + // Align to 4 + let rem = reader.stream_position()? % 4; + if rem != 0 { + reader.seek(SeekFrom::Current((4 - rem) as i64))?; + } + } + let mut reader: Box = if group.is_compressed() { + self.decompressor.wrap(reader)? + } else { + Box::new(reader) + }; + if !uncompressed_exception_lists { + self.exception_lists = read_exception_lists( + reader.as_mut(), + partition.is_some(), + self.disc.chunk_size.get(), + )?; + } + + if group.rvz_packed_size.get() > 0 { + // Decode RVZ packed data + let mut lfg = LaggedFibonacci::default(); + loop { + let mut size_bytes = [0u8; 4]; + match reader.read_exact(&mut size_bytes) { + Ok(_) => {} + Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break, + Err(e) => { + return Err(io::Error::new(e.kind(), "Failed to read RVZ packed size")); + } + } + let size = u32::from_be_bytes(size_bytes); + let cur_data_len = self.group_data.len(); + if size & 0x80000000 != 0 { + // Junk data + let size = size & 0x7FFFFFFF; + lfg.init_with_reader(reader.as_mut())?; + lfg.skip( + ((partition_offset + cur_data_len as u64) % SECTOR_SIZE as u64) + as usize, + ); + self.group_data.resize(cur_data_len + size as usize, 0); + lfg.fill(&mut self.group_data[cur_data_len..]); + } else { + // Real data + self.group_data.resize(cur_data_len + size as usize, 0); + reader.read_exact(&mut self.group_data[cur_data_len..])?; } } - let size = u32::from_be_bytes(size_bytes); - let cur_data_len = self.group_data.len(); - if size & 0x80000000 != 0 { - // Junk data - let size = size & 0x7FFFFFFF; - lfg.init_with_reader(reader.as_mut())?; - lfg.skip( - ((result.partition_offset + cur_data_len as u64) % SECTOR_SIZE as u64) - as usize, - ); - self.group_data.resize(cur_data_len + size as usize, 0); - lfg.fill(&mut self.group_data[cur_data_len..]); - } else { - // Real data - self.group_data.resize(cur_data_len + size as usize, 0); - reader.read_exact(&mut self.group_data[cur_data_len..])?; - } + } else { + // Read and decompress data + reader.read_to_end(&mut self.group_data)?; } + + self.group = group_index; + } + + // Read sector from cached group data + if partition.is_some() { + let sector_data_start = group_sector as usize * SECTOR_DATA_SIZE; + let sector_data = + &self.group_data[sector_data_start..sector_data_start + SECTOR_DATA_SIZE]; + out[..HASHES_SIZE].fill(0); + out[HASHES_SIZE..SECTOR_SIZE].copy_from_slice(sector_data); + Ok(Some(Block::PartDecrypted { has_hashes: false })) } else { - // Read and decompress data - reader.read_to_end(&mut self.group_data)?; - } - - drop(reader); - self.recalculate_hashes(result)?; - Ok(()) - } - - fn recalculate_hashes(&mut self, result: GroupResult) -> io::Result<()> { - let Some(partition_index) = result.partition_index else { - // Data not inside a Wii partition - return Ok(()); - }; - let hash_table = self.disc_io.hash_tables.get(partition_index); - - if self.group_data.len() % BLOCK_SIZE != 0 { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("Invalid group data size: {:#X}", self.group_data.len()), - )); - } - - // WIA/RVZ excludes the hash data for each sector, instead storing all data contiguously. - // We need to add space for the hash data, and then recalculate the hashes. - let num_sectors = self.group_data.len() / BLOCK_SIZE; - let mut out = vec![0u8; num_sectors * SECTOR_SIZE]; - for i in 0..num_sectors { - let data = array_ref![self.group_data, i * BLOCK_SIZE, BLOCK_SIZE]; - let out = array_ref_mut![out, i * SECTOR_SIZE, SECTOR_SIZE]; - - // Rebuild H0 hashes - for n in 0..31 { - let hash = hash_bytes(array_ref![data, n * 0x400, 0x400]); - array_ref_mut![out, n * 20, 20].copy_from_slice(&hash); - } - - // Copy data - array_ref_mut![out, 0x400, BLOCK_SIZE].copy_from_slice(data); - - // Rebuild H1 and H2 hashes if available - if let Some(hash_table) = hash_table { - let partition = &self.disc_io.partitions[partition_index]; - let part_sector = (result.disc_offset / SECTOR_SIZE as u64) as usize + i - - partition.partition_data[0].first_sector.get() as usize; - let h1_start = part_sector & !7; - for i in 0..8 { - array_ref_mut![out, 0x280 + i * 20, 20] - .copy_from_slice(&hash_table.h1_hashes[h1_start + i]); - } - let h2_start = (h1_start / 8) & !7; - for i in 0..8 { - array_ref_mut![out, 0x340 + i * 20, 20] - .copy_from_slice(&hash_table.h2_hashes[h2_start + i]); - } - - if self.encrypt { - // Re-encrypt hashes and data - aes_encrypt(&partition.partition_key, [0u8; 16], &mut out[..HASHES_SIZE]); - let iv = *array_ref![out, 0x3d0, 16]; - aes_encrypt(&partition.partition_key, iv, &mut out[HASHES_SIZE..]); - } - } - } - - self.group_data = out; - Ok(()) - } -} - -impl<'a> Read for WIAReadStream<'a> { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - let mut rem = buf.len(); - let mut read: usize = 0; - - // Special case: First 0x80 bytes are stored in the disc header - if self.offset < DISC_HEAD_SIZE as u64 { - let to_read = min(rem, DISC_HEAD_SIZE); - buf[read..read + to_read].copy_from_slice( - &self.disc_io.disc.disc_head[self.offset as usize..self.offset as usize + to_read], + let sector_data_start = group_sector as usize * SECTOR_SIZE; + out.copy_from_slice( + &self.group_data[sector_data_start..sector_data_start + SECTOR_SIZE], ); - rem -= to_read; - read += to_read; - self.offset += to_read as u64; + Ok(Some(Block::Raw)) } - - // Decompress groups and read data - while rem > 0 { - if !self.check_group()? { - break; - } - let group_offset = (self.offset - self.group_offset) as usize; - let to_read = min(rem, self.group_data.len() - group_offset); - buf[read..read + to_read] - .copy_from_slice(&self.group_data[group_offset..group_offset + to_read]); - rem -= to_read; - read += to_read; - self.offset += to_read as u64; - } - Ok(read) - } -} - -impl<'a> Seek for WIAReadStream<'a> { - fn seek(&mut self, pos: SeekFrom) -> io::Result { - self.offset = match pos { - SeekFrom::Start(v) => v, - SeekFrom::End(v) => self.disc_io.header.iso_file_size.get().saturating_add_signed(v), - SeekFrom::Current(v) => self.offset.saturating_add_signed(v), - }; - self.check_group()?; - Ok(self.offset) } - fn stream_position(&mut self) -> io::Result { Ok(self.offset) } -} - -impl<'a> ReadStream for WIAReadStream<'a> { - fn stable_stream_len(&mut self) -> io::Result { - Ok(self.disc_io.header.iso_file_size.get()) + fn block_size(&self) -> u32 { + // WIA/RVZ chunks aren't always the full size, so we'll consider the + // block size to be one sector, and handle the complexity ourselves. + SECTOR_SIZE as u32 } - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } + fn meta(&self) -> Result { + let mut meta = self.nkit_header.as_ref().map(DiscMeta::from).unwrap_or_default(); + meta.decrypted = true; + meta.needs_hash_recovery = true; + meta.lossless = true; + meta.disc_size = Some(self.header.iso_file_size.get()); + Ok(meta) + } } diff --git a/src/lib.rs b/src/lib.rs index 4f619a6..51b7bdd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,19 +38,20 @@ use std::path::Path; -use disc::DiscBase; pub use disc::{ AppLoaderHeader, DiscHeader, DolHeader, PartitionBase, PartitionHeader, PartitionInfo, - PartitionKind, PartitionMeta, BI2_SIZE, BOOT_SIZE, + PartitionKind, PartitionMeta, BI2_SIZE, BOOT_SIZE, SECTOR_SIZE, }; pub use fst::{Fst, Node, NodeKind}; -use io::DiscIO; pub use io::DiscMeta; +use io::{block, block::BPartitionInfo}; pub use streams::ReadStream; +use crate::disc::reader::{DiscReader, EncryptionMode}; + mod disc; mod fst; -mod io; +pub mod io; mod streams; mod util; @@ -126,8 +127,7 @@ pub struct OpenOptions { } pub struct Disc { - io: Box, - base: Box, + pub reader: DiscReader, options: OpenOptions, } @@ -139,39 +139,39 @@ impl Disc { /// Opens a disc image from a file path with custom options. pub fn new_with_options>(path: P, options: &OpenOptions) -> Result { - let mut io = io::open(path.as_ref(), options)?; - let base = disc::new(io.as_mut())?; - Ok(Disc { io, base, options: options.clone() }) + let io = block::open(path.as_ref(), options)?; + let reader = DiscReader::new(io, EncryptionMode::Encrypted)?; + Ok(Disc { reader, options: options.clone() }) } /// The disc's header. - pub fn header(&self) -> &DiscHeader { self.base.header() } + pub fn header(&self) -> &DiscHeader { self.reader.header() } /// Returns extra metadata included in the disc file format, if any. - pub fn meta(&self) -> Result { self.io.meta() } + pub fn meta(&self) -> Result { self.reader.io.meta() } /// The disc's size in bytes or an estimate if not stored by the format. - pub fn disc_size(&self) -> u64 { self.base.disc_size() } + pub fn disc_size(&self) -> u64 { self.reader.disc_size() } /// A list of partitions on the disc. /// /// For GameCube discs, this will return a single data partition spanning the entire disc. - pub fn partitions(&self) -> Vec { self.base.partitions() } + pub fn partitions(&self) -> &[BPartitionInfo] { self.reader.partitions() } - /// Opens a new read stream for the base disc image. - /// - /// Generally does _not_ need to be used directly. Opening a partition will provide a - /// decrypted stream instead. - pub fn open(&self) -> Result> { self.io.open() } - - /// Opens a new, decrypted partition read stream for the specified partition index. - pub fn open_partition(&self, index: usize) -> Result> { - self.base.open_partition(self.io.as_ref(), index, &self.options) - } - - /// Opens a new partition read stream for the first partition matching - /// the specified type. - pub fn open_partition_kind(&self, kind: PartitionKind) -> Result> { - self.base.open_partition_kind(self.io.as_ref(), kind, &self.options) - } + // /// Opens a new read stream for the base disc image. + // /// + // /// Generally does _not_ need to be used directly. Opening a partition will provide a + // /// decrypted stream instead. + // pub fn open(&self) -> Result> { self.io.open() } + // + // /// Opens a new, decrypted partition read stream for the specified partition index. + // pub fn open_partition(&self, index: usize) -> Result> { + // self.base.open_partition(self.io.as_ref(), index, &self.options) + // } + // + // /// Opens a new partition read stream for the first partition matching + // /// the specified type. + // pub fn open_partition_kind(&self, kind: PartitionKind) -> Result> { + // self.base.open_partition_kind(self.io.as_ref(), kind, &self.options) + // } } diff --git a/src/streams.rs b/src/streams.rs index cd08974..9b6b9a1 100644 --- a/src/streams.rs +++ b/src/streams.rs @@ -1,50 +1,12 @@ //! Common stream types use std::{ - fs::File, io, - io::{BufReader, Read, Seek, SeekFrom}, + io::{Read, Seek, SeekFrom}, }; -/// Creates a fixed-size array reference from a slice. -#[macro_export] -macro_rules! array_ref { - ($slice:expr, $offset:expr, $size:expr) => {{ - #[inline] - fn to_array(slice: &[T]) -> &[T; $size] { - unsafe { &*(slice.as_ptr() as *const [_; $size]) } - } - to_array(&$slice[$offset..$offset + $size]) - }}; -} - -/// Creates a mutable fixed-size array reference from a slice. -#[macro_export] -macro_rules! array_ref_mut { - ($slice:expr, $offset:expr, $size:expr) => {{ - #[inline] - fn to_array(slice: &mut [T]) -> &mut [T; $size] { - unsafe { &mut *(slice.as_ptr() as *mut [_; $size]) } - } - to_array(&mut $slice[$offset..$offset + $size]) - }}; -} - -/// Compile-time assertion. -#[macro_export] -macro_rules! static_assert { - ($condition:expr) => { - const _: () = core::assert!($condition); - }; -} - /// A helper trait for seekable read streams. pub trait ReadStream: Read + Seek { - /// Replace with [`Read.stream_len`] when stabilized. - /// - /// - fn stable_stream_len(&mut self) -> io::Result; - /// Creates a windowed read sub-stream with offset and size. /// /// Seeks underlying stream immediately. @@ -57,54 +19,12 @@ pub trait ReadStream: Read + Seek { fn as_dyn(&mut self) -> &mut dyn ReadStream; } -impl ReadStream for File { - fn stable_stream_len(&mut self) -> io::Result { - let before = self.stream_position()?; - let result = self.seek(SeekFrom::End(0)); - // Try to restore position even if the above failed - let seek_result = self.seek(SeekFrom::Start(before)); - if seek_result.is_err() { - return if result.is_err() { result } else { seek_result }; - } - result - } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - -impl ReadStream for BufReader -where T: ReadStream +impl ReadStream for T +where T: Read + Seek { - fn stable_stream_len(&mut self) -> io::Result { self.get_mut().stable_stream_len() } - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } -trait WindowedReadStream: ReadStream { - fn base_stream(&mut self) -> &mut dyn ReadStream; - fn window(&self) -> (u64, u64); -} - -/// A window into an existing [`ReadStream`], with ownership of the underlying stream. -pub struct OwningWindowedReadStream<'a> { - /// The base stream. - pub base: Box, - /// The beginning of the window in bytes. - pub begin: u64, - /// The end of the window in bytes. - pub end: u64, -} - -/// Takes ownership of & wraps a read stream into a windowed read stream. -pub fn wrap_windowed<'a>( - mut base: Box, - offset: u64, - size: u64, -) -> io::Result> { - base.seek(SeekFrom::Start(offset))?; - Ok(OwningWindowedReadStream { base, begin: offset, end: offset + size }) -} - /// A non-owning window into an existing [`ReadStream`]. pub struct SharedWindowedReadStream<'a> { /// A reference to the base stream. @@ -125,137 +45,36 @@ impl<'a> SharedWindowedReadStream<'a> { } } -#[inline(always)] -fn windowed_read(stream: &mut impl WindowedReadStream, buf: &mut [u8]) -> io::Result { - let pos = stream.stream_position()?; - let size = stream.stable_stream_len()?; - if pos == size { - return Ok(0); - } - stream.base_stream().read(if pos + buf.len() as u64 > size { - &mut buf[..(size - pos) as usize] - } else { - buf - }) -} - -#[inline(always)] -fn windowed_seek(stream: &mut impl WindowedReadStream, pos: SeekFrom) -> io::Result { - let (begin, end) = stream.window(); - let result = stream.base_stream().seek(match pos { - SeekFrom::Start(p) => SeekFrom::Start(begin + p), - SeekFrom::End(p) => SeekFrom::End(end as i64 + p), - SeekFrom::Current(_) => pos, - })?; - if result < begin || result > end { - Err(io::Error::from(io::ErrorKind::UnexpectedEof)) - } else { - Ok(result - begin) - } -} - -impl<'a> Read for OwningWindowedReadStream<'a> { - fn read(&mut self, buf: &mut [u8]) -> io::Result { windowed_read(self, buf) } -} - -impl<'a> Seek for OwningWindowedReadStream<'a> { - fn seek(&mut self, pos: SeekFrom) -> io::Result { windowed_seek(self, pos) } - - fn stream_position(&mut self) -> io::Result { - Ok(self.base.stream_position()? - self.begin) - } -} - -impl<'a> ReadStream for OwningWindowedReadStream<'a> { - fn stable_stream_len(&mut self) -> io::Result { Ok(self.end - self.begin) } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - -impl<'a> WindowedReadStream for OwningWindowedReadStream<'a> { - fn base_stream(&mut self) -> &mut dyn ReadStream { self.base.as_dyn() } - - fn window(&self) -> (u64, u64) { (self.begin, self.end) } -} - impl<'a> Read for SharedWindowedReadStream<'a> { - fn read(&mut self, buf: &mut [u8]) -> io::Result { windowed_read(self, buf) } + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let pos = self.stream_position()?; + let size = self.end - self.begin; + if pos == size { + return Ok(0); + } + self.base.read(if pos + buf.len() as u64 > size { + &mut buf[..(size - pos) as usize] + } else { + buf + }) + } } impl<'a> Seek for SharedWindowedReadStream<'a> { - fn seek(&mut self, pos: SeekFrom) -> io::Result { windowed_seek(self, pos) } + fn seek(&mut self, pos: SeekFrom) -> io::Result { + let result = self.base.seek(match pos { + SeekFrom::Start(p) => SeekFrom::Start(self.begin + p), + SeekFrom::End(p) => SeekFrom::End(self.end as i64 + p), + SeekFrom::Current(_) => pos, + })?; + if result < self.begin || result > self.end { + Err(io::Error::from(io::ErrorKind::UnexpectedEof)) + } else { + Ok(result - self.begin) + } + } fn stream_position(&mut self) -> io::Result { Ok(self.base.stream_position()? - self.begin) } } - -impl<'a> ReadStream for SharedWindowedReadStream<'a> { - fn stable_stream_len(&mut self) -> io::Result { Ok(self.end - self.begin) } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - -impl<'a> WindowedReadStream for SharedWindowedReadStream<'a> { - fn base_stream(&mut self) -> &mut dyn ReadStream { self.base } - - fn window(&self) -> (u64, u64) { (self.begin, self.end) } -} - -/// A non-owning [`ReadStream`] wrapping a byte slice reference. -pub struct ByteReadStream<'a> { - /// A reference to the underlying byte slice. - pub bytes: &'a [u8], - /// The current position in the stream. - pub position: u64, -} - -impl Read for ByteReadStream<'_> { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - let mut len = buf.len(); - let total = self.bytes.len(); - let pos = self.position as usize; - if len + pos > total { - #[allow(clippy::comparison_chain)] - if pos > total { - return Err(io::Error::from(io::ErrorKind::UnexpectedEof)); - } else if pos == total { - return Ok(0); - } - len = total - pos; - } - buf.copy_from_slice(&self.bytes[pos..pos + len]); - self.position += len as u64; - Ok(len) - } -} - -impl Seek for ByteReadStream<'_> { - fn seek(&mut self, pos: SeekFrom) -> io::Result { - let new_pos = match pos { - SeekFrom::Start(v) => v, - SeekFrom::End(v) => (self.bytes.len() as u64).saturating_add_signed(v), - SeekFrom::Current(v) => self.position.saturating_add_signed(v), - }; - if new_pos > self.bytes.len() as u64 { - Err(io::Error::from(io::ErrorKind::UnexpectedEof)) - } else { - self.position = new_pos; - Ok(new_pos) - } - } - - // fn stream_len(&mut self) -> io::Result { Ok(self.bytes.len() as u64) } - - fn stream_position(&mut self) -> io::Result { Ok(self.position) } -} - -impl ReadStream for ByteReadStream<'_> { - fn stable_stream_len(&mut self) -> io::Result { Ok(self.bytes.len() as u64) } - - fn as_dyn(&mut self) -> &mut dyn ReadStream { self } -} - -impl<'a> AsMut for ByteReadStream<'a> { - fn as_mut(&mut self) -> &mut (dyn ReadStream + 'a) { self as &mut (dyn ReadStream + 'a) } -} diff --git a/src/util/compress.rs b/src/util/compress.rs index 9f1b172..e0d990c 100644 --- a/src/util/compress.rs +++ b/src/util/compress.rs @@ -1,14 +1,15 @@ use std::{io, io::Read}; -use crate::{Error, Result}; - /// Decodes the LZMA Properties byte (lc/lp/pb). /// See `lzma_lzma_lclppb_decode` in `liblzma/lzma/lzma_decoder.c`. #[cfg(feature = "compress-lzma")] -pub fn lzma_lclppb_decode(options: &mut liblzma::stream::LzmaOptions, byte: u8) -> Result<()> { +pub fn lzma_lclppb_decode(options: &mut liblzma::stream::LzmaOptions, byte: u8) -> io::Result<()> { let mut d = byte as u32; if d >= (9 * 5 * 5) { - return Err(Error::DiscFormat(format!("Invalid LZMA props byte: {}", d))); + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid LZMA props byte: {}", d), + )); } options.literal_context_bits(d % 9); d /= 9; @@ -20,10 +21,13 @@ pub fn lzma_lclppb_decode(options: &mut liblzma::stream::LzmaOptions, byte: u8) /// Decodes LZMA properties. /// See `lzma_lzma_props_decode` in `liblzma/lzma/lzma_decoder.c`. #[cfg(feature = "compress-lzma")] -pub fn lzma_props_decode(props: &[u8]) -> Result { +pub fn lzma_props_decode(props: &[u8]) -> io::Result { use crate::array_ref; if props.len() != 5 { - return Err(Error::DiscFormat(format!("Invalid LZMA props length: {}", props.len()))); + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid LZMA props length: {}", props.len()), + )); } let mut options = liblzma::stream::LzmaOptions::new(); lzma_lclppb_decode(&mut options, props[0])?; @@ -34,16 +38,22 @@ pub fn lzma_props_decode(props: &[u8]) -> Result { /// Decodes LZMA2 properties. /// See `lzma_lzma2_props_decode` in `liblzma/lzma/lzma2_decoder.c`. #[cfg(feature = "compress-lzma")] -pub fn lzma2_props_decode(props: &[u8]) -> Result { +pub fn lzma2_props_decode(props: &[u8]) -> io::Result { use std::cmp::Ordering; if props.len() != 1 { - return Err(Error::DiscFormat(format!("Invalid LZMA2 props length: {}", props.len()))); + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid LZMA2 props length: {}", props.len()), + )); } let d = props[0] as u32; let mut options = liblzma::stream::LzmaOptions::new(); options.dict_size(match d.cmp(&40) { Ordering::Greater => { - return Err(Error::DiscFormat(format!("Invalid LZMA2 props byte: {}", d))); + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid LZMA2 props byte: {}", d), + )); } Ordering::Equal => u32::MAX, Ordering::Less => (2 | (d & 1)) << (d / 2 + 11), diff --git a/src/util/lfg.rs b/src/util/lfg.rs index 46119df..15c5039 100644 --- a/src/util/lfg.rs +++ b/src/util/lfg.rs @@ -46,6 +46,7 @@ impl LaggedFibonacci { init[0].wrapping_add(init[1]), ]) ^ disc_num as u32; let sector = (partition_offset / SECTOR_SIZE as u64) as u32; + let sector_offset = partition_offset % SECTOR_SIZE as u64; let mut n = seed.wrapping_mul(0x260BCD5) ^ sector.wrapping_mul(0x1EF29123); for i in 0..SEED_SIZE { let mut v = 0u32; @@ -58,7 +59,7 @@ impl LaggedFibonacci { self.buffer[16] ^= self.buffer[0] >> 9 ^ self.buffer[16] << 23; self.position = 0; self.init(); - self.skip((partition_offset % SECTOR_SIZE as u64) as usize); + self.skip(sector_offset as usize); } pub fn init_with_reader(&mut self, reader: &mut R) -> io::Result<()> diff --git a/src/util/mod.rs b/src/util/mod.rs index 681ca76..69f60d9 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -2,7 +2,7 @@ use std::ops::{Div, Rem}; pub(crate) mod compress; pub(crate) mod lfg; -pub(crate) mod reader; +pub(crate) mod read; pub(crate) mod take_seek; #[inline(always)] @@ -12,3 +12,35 @@ where T: Div + Rem + Copy { let rem = x % y; (quot, rem) } + +/// Creates a fixed-size array reference from a slice. +#[macro_export] +macro_rules! array_ref { + ($slice:expr, $offset:expr, $size:expr) => {{ + #[inline] + fn to_array(slice: &[T]) -> &[T; $size] { + unsafe { &*(slice.as_ptr() as *const [_; $size]) } + } + to_array(&$slice[$offset..$offset + $size]) + }}; +} + +/// Creates a mutable fixed-size array reference from a slice. +#[macro_export] +macro_rules! array_ref_mut { + ($slice:expr, $offset:expr, $size:expr) => {{ + #[inline] + fn to_array(slice: &mut [T]) -> &mut [T; $size] { + unsafe { &mut *(slice.as_ptr() as *mut [_; $size]) } + } + to_array(&mut $slice[$offset..$offset + $size]) + }}; +} + +/// Compile-time assertion. +#[macro_export] +macro_rules! static_assert { + ($condition:expr) => { + const _: () = core::assert!($condition); + }; +} diff --git a/src/util/reader.rs b/src/util/read.rs similarity index 85% rename from src/util/reader.rs rename to src/util/read.rs index d8e063f..df1d4dc 100644 --- a/src/util/reader.rs +++ b/src/util/read.rs @@ -24,6 +24,17 @@ where Ok(ret) } +#[inline(always)] +pub fn read_box(reader: &mut R) -> io::Result> +where + T: FromBytes + FromZeroes + AsBytes, + R: Read + ?Sized, +{ + let mut ret = ::new_box_zeroed(); + reader.read_exact(ret.as_mut().as_bytes_mut())?; + Ok(ret) +} + #[inline(always)] pub fn read_box_slice(reader: &mut R, count: usize) -> io::Result> where