encounter
/
nod-rs
mirror of https://github.com/encounter/nod-rs.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Convert all formats to new BlockIO trait (WIP) 

Simplifies format handling by moving
common logic into a new `DiscReader`
type.
This commit is contained in:
Luke Street 2024-02-21 00:04:23 -07:00
parent ce9fbbf822
commit 7f97dac399
24 changed files with 1991 additions and 1575 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
Cargo.toml
View File

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

4
README.md
View File

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

222
src/bin.rs
View File

@ -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::<crc32fast::Hasher>(),
digest_thread::<md5::Md5>(),
digest_thread::<sha1::Sha1>(),
digest_thread::<xxhash_rust::xxh64::Xxh64>(),
];
let digest_threads = if md5 {
vec![
digest_thread::<crc32fast::Hasher>(),
digest_thread::<md5::Md5>(),
digest_thread::<sha1::Sha1>(),
digest_thread::<xxhash_rust::xxh64::Xxh64>(),
]
} else {
vec![
digest_thread::<crc32fast::Hasher>(),
digest_thread::<sha1::Sha1>(),
digest_thread::<xxhash_rust::xxh64::Xxh64>(),
]
};
let (w_tx, w_rx) = sync_channel::<Arc<[u8]>>(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 = <u8>::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<str>, 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<str>, 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<H>() -> (SyncSender<Arc<[u8]>>, JoinHandle<DigestResult>)
type DigestThread = (SyncSender<Arc<[u8]>>, JoinHandle<DigestResult>);
fn digest_thread<H>() -> DigestThread
where H: Hasher + Send + 'static {
let (tx, rx) = sync_channel::<Arc<[u8]>>(1);
let handle = thread::spawn(move || {

31
src/disc/gcn.rs
View File

@ -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<u64>,
) -> Result<DiscGCN> {
// stream.seek(SeekFrom::Start(size_of::<DiscHeader>() 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<u64> {
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<u64> { Ok(self.offset) }
}
impl<'a> ReadStream for PartitionGC<'a> {
fn stable_stream_len(&mut self) -> io::Result<u64> { 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<Box<PartitionMeta>> {
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<R>(reader: &mut R, is_wii: bool) -> Result<Box<PartitionMeta>>
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::<DiscHeader>()..]).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<u8> =
@ -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<u8> = 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,

203
src/disc/hashes.rs Normal file
View File

@ -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 = <u8>::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()
}

25
src/disc/mod.rs
View File

@ -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<u8>,
pub raw_apploader: Box<[u8]>,
/// File system table (fst.bin)
pub raw_fst: Vec<u8>,
pub raw_fst: Box<[u8]>,
/// Main binary (main.dol)
pub raw_dol: Vec<u8>,
pub raw_dol: Box<[u8]>,
/// Ticket (ticket.bin, Wii only)
pub raw_ticket: Option<Vec<u8>>,
pub raw_ticket: Option<Box<[u8]>>,
/// TMD (tmd.bin, Wii only)
pub raw_tmd: Option<Vec<u8>>,
pub raw_tmd: Option<Box<[u8]>>,
/// Certificate chain (cert.bin, Wii only)
pub raw_cert_chain: Option<Vec<u8>>,
pub raw_cert_chain: Option<Box<[u8]>>,
/// H3 hash table (h3.bin, Wii only)
pub raw_h3_table: Option<Vec<u8>>,
pub raw_h3_table: Option<Box<[u8]>>,
}
impl PartitionMeta {

177
src/disc/partition.rs Normal file
View File

@ -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<dyn BlockIO>,
partition: BPartitionInfo,
block: Option<Block>,
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: <u8>::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<dyn BlockIO>, partition: &BPartitionInfo) -> Result<Self> {
let block_size = inner.block_size();
Ok(Self {
io: inner,
partition: partition.clone(),
block: None,
block_buf: <u8>::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<usize> {
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<u64> {
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(())
}

273
src/disc/reader.rs Normal file
View File

@ -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<dyn BlockIO>,
block: Option<Block>,
block_buf: Box<[u8]>,
block_idx: u32,
sector_buf: Box<[u8; SECTOR_SIZE]>,
sector_idx: u32,
pos: u64,
mode: EncryptionMode,
disc_header: Box<DiscHeader>,
pub(crate) partitions: Vec<BPartitionInfo>,
hash_tables: Vec<HashTable>,
}
impl Clone for DiscReader {
fn clone(&self) -> Self {
Self {
io: self.io.clone(),
block: None,
block_buf: <u8>::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<dyn BlockIO>, mode: EncryptionMode) -> Result<Self> {
let block_size = inner.block_size();
let meta = inner.meta()?;
let mut reader = Self {
io: inner,
block: None,
block_buf: <u8>::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<DiscHeader> = 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<usize> {
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<u64> {
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<Vec<BPartitionInfo>> {
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<WiiPartEntry> = 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<WiiPartitionHeader> = 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
}
}

265
src/disc/wii.rs
View File

@ -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::<WiiPartEntry>() == 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::<WiiPartGroup>() == 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::<Ticket>() == 0x2A4);
impl Ticket {
pub fn decrypt_title_key(&self) -> Result<KeyBytes> {
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<u64>,
) -> Result<Self> {
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<Vec<WiiPartInfo>> {
pub(crate) fn read_partition_info(
stream: &mut dyn ReadStream,
disc_header: &DiscHeader,
) -> Result<Vec<WiiPartInfo>> {
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<Vec<Wii
stream
.seek(SeekFrom::Start(offset))
.with_context(|| format!("Seeking to partition data {group_idx}:{part_idx}"))?;
let mut header: WiiPartitionHeader = read_from(stream)
let header: WiiPartitionHeader = read_from(stream)
.with_context(|| format!("Reading partition header {group_idx}:{part_idx}"))?;
// Decrypt title key
let mut iv: KeyBytes = [0; 16];
iv[..8].copy_from_slice(&header.ticket.title_id);
let common_key =
COMMON_KEYS.get(header.ticket.common_key_idx as usize).ok_or(Error::DiscFormat(
format!("unknown common key index {}", header.ticket.common_key_idx),
))?;
aes_decrypt(common_key, iv, &mut header.ticket.title_key);
// Open partition stream and read junk data seed
let inner = stream
.new_window(offset + header.data_off(), header.data_size())
.context("Wrapping partition stream")?;
// let inner = stream
// .new_window(offset + header.data_off(), DL_DVD_SIZE) // header.data_size()
// .context("Wrapping partition stream")?;
let title_key = header.ticket.decrypt_title_key()?;
let part_offset = entry.offset() + header.data_off();
if part_offset % SECTOR_SIZE as u64 != 0 {
return Err(Error::DiscFormat(format!(
"Partition {group_idx}:{part_idx} offset is not sector aligned",
)));
}
let start_sector = (part_offset / SECTOR_SIZE as u64) as u32;
let mut stream = PartitionWii {
start_sector,
header: header.clone(),
tmd: vec![],
cert_chain: vec![],
h3_table: vec![],
stream: Box::new(inner),
key: Some(header.ticket.title_key),
stream: Box::new(stream.as_dyn()),
key: Some(title_key),
offset: 0,
cur_block: 0,
cur_block: u32::MAX,
buf: [0; SECTOR_SIZE],
has_hashes: disc_header.no_partition_hashes == 0,
validate_hashes: false,
};
let junk_id: [u8; 4] = read_from(&mut stream).context("Reading junk seed bytes")?;
@ -259,12 +307,13 @@ pub(crate) fn read_partition_info(stream: &mut dyn ReadStream) -> Result<Vec<Wii
read_from(&mut stream).context("Reading partition header")?;
let junk_start = part_header.fst_off(true) + part_header.fst_sz(true);
// log::debug!(
// "Partition: {:?} - {:?}: {:?}",
// offset + header.data_off(),
// header.data_size(),
// header.ticket.title_key
// );
log::debug!("Header: {:?}", header);
log::debug!(
"Partition: {:?} - {:?}: {:?}",
offset + header.data_off(),
header.data_size(),
header.ticket.title_key
);
part_info.push(WiiPartInfo {
group_idx: group_idx as u32,
@ -274,6 +323,7 @@ pub(crate) fn read_partition_info(stream: &mut dyn ReadStream) -> Result<Vec<Wii
header,
junk_id,
junk_start,
title_key,
});
}
}
@ -309,8 +359,6 @@ fn open_partition<'a>(
options: &OpenOptions,
header: &DiscHeader,
) -> Result<Box<dyn PartitionBase + 'a>> {
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<u8> = 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<u8>,
cert_chain: Vec<u8>,
@ -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<Box<PartitionMeta>> {
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<usize> {
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<u64> {
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<u64> { Ok(self.offset) }
}
impl<'a> ReadStream for PartitionWii<'a> {
fn stable_stream_len(&mut self) -> io::Result<u64> {
Ok(to_block_size(self.stream.stable_stream_len()?))
}
fn as_dyn(&mut self) -> &mut dyn ReadStream { self }
}

279
src/io/block.rs Normal file
View File

@ -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<Option<Block>>;
/// 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<DiscMeta>;
}
dyn_clone::clone_trait_object!(BlockIO);
/// Creates a new [`BlockIO`] instance.
pub fn open(filename: &Path, options: &OpenOptions) -> Result<Box<dyn BlockIO>> {
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<WiiPartitionHeader>,
pub disc_header: Box<DiscHeader>,
pub partition_header: Box<PartitionHeader>,
pub hash_table: Option<HashTable>,
}
#[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::<SECTOR_SIZE>(data, rel_sector)?);
decrypt_sector(out, partition);
}
Block::PartDecrypted { has_hashes } => {
out.copy_from_slice(block_sector::<SECTOR_SIZE>(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::<SECTOR_SIZE>(data, rel_sector)?);
}
Block::PartDecrypted { has_hashes } => {
out.copy_from_slice(block_sector::<SECTOR_SIZE>(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::<SECTOR_SIZE>(
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<const N: usize>(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..]);
}

251
src/io/ciso.rs
View File

@ -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<NKitHeader>,
junk_blocks: Option<Box<[u8]>>,
partitions: Vec<PartitionInfo>,
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<Self> {
let mut inner = BufReader::new(SplitFileReader::new(filename)?);
pub fn new(filename: &Path) -> Result<Box<Self>> {
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<dyn ReadStream> = 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<dyn DiscBase> = 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<Box<dyn ReadStream>> {
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<Option<Block>> {
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::<CISOHeader>() 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<DiscMeta> {
Ok(self.nkit_header.as_ref().map(DiscMeta::from).unwrap_or_default())
}
fn disc_size(&self) -> Option<u64> { self.nkit_header.as_ref().and_then(|h| h.size) }
}
struct CISOReadStream {
inner: BufReader<SplitFileReader>,
block_size: u32,
block_map: [u16; CISO_MAP_SIZE],
cur_block: u16,
pos: u64,
// Data for recreating junk data
junk_blocks: Option<Box<[u8]>>,
partitions: Vec<PartitionInfo>,
disc_num: u8,
}
impl CISOReadStream {
fn read_junk_data(&mut self, buf: &mut [u8]) -> io::Result<usize> {
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<usize> {
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::<CISOHeader>() 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<u64> {
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<u64> {
Ok(self.block_size as u64 * CISO_MAP_SIZE as u64)
}
fn as_dyn(&mut self) -> &mut dyn ReadStream { self }
}

57
src/io/iso.rs
View File

@ -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<Self> {
Ok(Self { inner: SplitFileReader::new(filename)? })
pub fn new(filename: &Path) -> Result<Box<Self>> {
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<Box<dyn ReadStream>> {
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<Option<Block>> {
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<u64> { Some(self.inner.len()) }
fn block_size(&self) -> u32 { SECTOR_SIZE as u32 }
fn meta(&self) -> Result<DiscMeta> {
Ok(DiscMeta { lossless: true, disc_size: Some(self.inner.len()), ..Default::default() })
}
}

55
src/io/mod.rs
View File

@ -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<u64>;
}
/// 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<u64>,
/// The original disc's CRC32 hash, if stored by the format.
pub crc32: Option<u32>,
/// 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<u64>,
}
/// Creates a new [`DiscIO`] instance.
pub fn open(filename: &Path, options: &OpenOptions) -> Result<Box<dyn DiscIO>> {
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};

236
src/io/nfs.rs
View File

@ -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::<NFSHeader>() == 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<DiscIONFS> {
let mut disc_io = DiscIONFS {
pub fn new(directory: &Path, options: &OpenOptions) -> Result<Box<Self>> {
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<DiscHeader>,
/// 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<PartitionInfo>,
}
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<Option<Block>> {
// 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::<NFSHeader>() 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<usize> {
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<u64> {
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<u64> { Ok(self.pos) }
}
fn block_size(&self) -> u32 { SECTOR_SIZE as u32 }
impl ReadStream for NFSReadStream {
fn stable_stream_len(&mut self) -> io::Result<u64> { Ok(self.disc_size) }
fn as_dyn(&mut self) -> &mut dyn ReadStream { self }
}
impl DiscIO for DiscIONFS {
fn open(&self) -> Result<Box<dyn ReadStream>> {
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<DiscMeta> {
Ok(DiscMeta { decrypted: true, lossless: true, ..Default::default() })
}
fn disc_size(&self) -> Option<u64> { None }
}
fn get_path<P>(directory: &Path, path: P) -> PathBuf

53
src/io/nkit.rs
View File

@ -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<u64>,
/// Bitstream of blocks that are junk data
pub junk_bits: Option<Vec<u8>>,
}
const VERSION_PREFIX: [u8; 7] = *b"NKIT v";
impl NKitHeader {
pub fn try_read_from<R>(reader: &mut R) -> Option<Self>
pub fn try_read_from<R>(reader: &mut R, block_size: u32, has_junk_bits: bool) -> Option<Self>
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<R>(reader: &mut R) -> io::Result<Self>
pub fn read_from<R>(reader: &mut R, block_size: u32, has_junk_bits: bool) -> io::Result<Self>
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<bool> {
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()
}
}
}

53
src/io/split.rs
View File

@ -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<Split<PathBuf>>,
open_file: Option<Split<File>>,
open_file: Option<Split<BufReader<File>>>,
pos: u64,
}
@ -108,30 +109,24 @@ impl SplitFileReader {
}
impl Read for SplitFileReader {
fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
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<usize> {
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<u64> { 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 } }
}

157
src/io/wbfs.rs
View File

@ -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<U16>,
block_table: Box<[U16]>,
/// Optional NKit header
nkit_header: Option<NKitHeader>,
}
impl DiscIOWBFS {
pub fn new(filename: &Path) -> Result<Self> {
let mut inner = BufReader::new(SplitFileReader::new(filename)?);
pub fn new(filename: &Path) -> Result<Box<Self>> {
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<u8> =
read_vec(&mut inner, header.sector_size() as usize - size_of::<WBFSHeader>())
let disc_table: Box<[u8]> =
read_box_slice(&mut inner, header.sector_size() as usize - size_of::<WBFSHeader>())
.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<U16> = 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<Box<dyn ReadStream>> {
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<Option<Block>> {
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<DiscMeta> {
Ok(self.nkit_header.as_ref().map(DiscMeta::from).unwrap_or_default())
}
fn disc_size(&self) -> Option<u64> { self.nkit_header.as_ref().and_then(|h| h.size) }
}
struct WBFSReadStream {
/// File reader
inner: BufReader<SplitFileReader>,
/// WBFS header
header: WBFSHeader,
/// Map of Wii LBAs to WBFS LBAs
wlba_table: Vec<U16>,
/// Current WBFS LBA
wlba: u32,
/// Current stream offset
pos: u64,
/// Optional known size
disc_size: Option<u64>,
}
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<usize> {
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<u64> {
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<u64> { Ok(self.disc_size()) }
fn as_dyn(&mut self) -> &mut dyn ReadStream { self }
}

852
src/io/wia.rs
View File

File diff suppressed because it is too large Load Diff

58
src/lib.rs
View File

@ -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<dyn DiscIO>,
base: Box<dyn DiscBase>,
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<P: AsRef<Path>>(path: P, options: &OpenOptions) -> Result<Disc> {
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<DiscMeta> { self.io.meta() }
pub fn meta(&self) -> Result<DiscMeta> { 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<PartitionInfo> { 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<Box<dyn ReadStream + '_>> { self.io.open() }
/// Opens a new, decrypted partition read stream for the specified partition index.
pub fn open_partition(&self, index: usize) -> Result<Box<dyn PartitionBase + '_>> {
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<Box<dyn PartitionBase + '_>> {
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<Box<dyn ReadStream + '_>> { self.io.open() }
//
// /// Opens a new, decrypted partition read stream for the specified partition index.
// pub fn open_partition(&self, index: usize) -> Result<Box<dyn PartitionBase + '_>> {
// 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<Box<dyn PartitionBase + '_>> {
// self.base.open_partition_kind(self.io.as_ref(), kind, &self.options)
// }
}

235
src/streams.rs
View File

@ -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<T>(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<T>(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.
///
/// <https://github.com/rust-lang/rust/issues/59359>
fn stable_stream_len(&mut self) -> io::Result<u64>;
/// 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<u64> {
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<T> ReadStream for BufReader<T>
where T: ReadStream
impl<T> ReadStream for T
where T: Read + Seek
{
fn stable_stream_len(&mut self) -> io::Result<u64> { 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<dyn ReadStream + 'a>,
/// 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<dyn ReadStream + 'a>,
offset: u64,
size: u64,
) -> io::Result<OwningWindowedReadStream<'a>> {
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<usize> {
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<u64> {
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<usize> { windowed_read(self, buf) }
}
impl<'a> Seek for OwningWindowedReadStream<'a> {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> { windowed_seek(self, pos) }
fn stream_position(&mut self) -> io::Result<u64> {
Ok(self.base.stream_position()? - self.begin)
}
}
impl<'a> ReadStream for OwningWindowedReadStream<'a> {
fn stable_stream_len(&mut self) -> io::Result<u64> { 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<usize> { windowed_read(self, buf) }
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
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<u64> { windowed_seek(self, pos) }
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
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<u64> {
Ok(self.base.stream_position()? - self.begin)
}
}
impl<'a> ReadStream for SharedWindowedReadStream<'a> {
fn stable_stream_len(&mut self) -> io::Result<u64> { 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<usize> {
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<u64> {
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<u64> { Ok(self.bytes.len() as u64) }
fn stream_position(&mut self) -> io::Result<u64> { Ok(self.position) }
}
impl ReadStream for ByteReadStream<'_> {
fn stable_stream_len(&mut self) -> io::Result<u64> { Ok(self.bytes.len() as u64) }
fn as_dyn(&mut self) -> &mut dyn ReadStream { self }
}
impl<'a> AsMut<dyn ReadStream + 'a> for ByteReadStream<'a> {
fn as_mut(&mut self) -> &mut (dyn ReadStream + 'a) { self as &mut (dyn ReadStream + 'a) }
}

28
src/util/compress.rs
View File

@ -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<liblzma::stream::LzmaOptions> {
pub fn lzma_props_decode(props: &[u8]) -> io::Result<liblzma::stream::LzmaOptions> {
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<liblzma::stream::LzmaOptions> {
/// 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<liblzma::stream::LzmaOptions> {
pub fn lzma2_props_decode(props: &[u8]) -> io::Result<liblzma::stream::LzmaOptions> {
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),

3
src/util/lfg.rs
View File

@ -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<R>(&mut self, reader: &mut R) -> io::Result<()>

34
src/util/mod.rs
View File

@ -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<Output = T> + Rem<Output = T> + 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<T>(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<T>(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);
};
}

11
src/util/reader.rs → src/util/read.rs
View File

@ -24,6 +24,17 @@ where
Ok(ret)
}
#[inline(always)]
pub fn read_box<T, R>(reader: &mut R) -> io::Result<Box<T>>
where
T: FromBytes + FromZeroes + AsBytes,
R: Read + ?Sized,
{
let mut ret = <T>::new_box_zeroed();
reader.read_exact(ret.as_mut().as_bytes_mut())?;
Ok(ret)
}
#[inline(always)]
pub fn read_box_slice<T, R>(reader: &mut R, count: usize) -> io::Result<Box<[T]>>
where