Rename PartitionHeader -> BootHeader & various fixes

This commit is contained in:
2025-03-04 22:54:09 -07:00
parent 73eebfe90b
commit fb3542f445
22 changed files with 694 additions and 402 deletions

View File

@@ -11,11 +11,11 @@ use zerocopy::{FromZeros, IntoBytes};
use crate::{
disc::{
fst::{Fst, FstBuilder},
DiscHeader, PartitionHeader, BI2_SIZE, BOOT_SIZE, GCN_MAGIC, MINI_DVD_SIZE, SECTOR_SIZE,
BootHeader, DiscHeader, BI2_SIZE, BOOT_SIZE, GCN_MAGIC, MINI_DVD_SIZE, SECTOR_SIZE,
WII_MAGIC,
},
read::DiscStream,
util::{align_up_64, array_ref, array_ref_mut, lfg::LaggedFibonacci},
util::{array_ref, array_ref_mut, lfg::LaggedFibonacci, Align},
Error, Result, ResultContext,
};
@@ -33,7 +33,7 @@ pub struct FileInfo {
pub struct GCPartitionBuilder {
disc_header: Box<DiscHeader>,
partition_header: Box<PartitionHeader>,
boot_header: Box<BootHeader>,
user_files: Vec<FileInfo>,
overrides: PartitionOverrides,
junk_files: Vec<String>,
@@ -97,7 +97,7 @@ impl GCPartitionBuilder {
}
Self {
disc_header,
partition_header: PartitionHeader::new_box_zeroed().unwrap(),
boot_header: BootHeader::new_box_zeroed().unwrap(),
user_files: Vec::new(),
overrides,
junk_files: Vec::new(),
@@ -110,8 +110,8 @@ impl GCPartitionBuilder {
}
#[inline]
pub fn set_partition_header(&mut self, partition_header: Box<PartitionHeader>) {
self.partition_header = partition_header;
pub fn set_boot_header(&mut self, boot_header: Box<BootHeader>) {
self.boot_header = boot_header;
}
pub fn add_file(&mut self, info: FileInfo) -> Result<()> {
@@ -139,8 +139,8 @@ impl GCPartitionBuilder {
layout.locate_sys_files(sys_file_callback)?;
layout.apply_overrides(&self.overrides)?;
let write_info = layout.layout_files()?;
let disc_size = layout.partition_header.user_offset.get() as u64
+ layout.partition_header.user_size.get() as u64;
let disc_size =
layout.boot_header.user_offset.get() as u64 + layout.boot_header.user_size.get() as u64;
let junk_id = layout.junk_id();
Ok(GCPartitionWriter::new(write_info, disc_size, junk_id, self.disc_header.disc_num))
}
@@ -148,7 +148,7 @@ impl GCPartitionBuilder {
struct GCPartitionLayout {
disc_header: Box<DiscHeader>,
partition_header: Box<PartitionHeader>,
boot_header: Box<BootHeader>,
user_files: Vec<FileInfo>,
apploader_file: Option<FileInfo>,
dol_file: Option<FileInfo>,
@@ -162,7 +162,7 @@ impl GCPartitionLayout {
fn new(builder: &GCPartitionBuilder) -> Self {
GCPartitionLayout {
disc_header: builder.disc_header.clone(),
partition_header: builder.partition_header.clone(),
boot_header: builder.boot_header.clone(),
user_files: builder.user_files.clone(),
apploader_file: None,
dol_file: None,
@@ -194,9 +194,7 @@ impl GCPartitionLayout {
)));
}
self.disc_header.as_mut_bytes().copy_from_slice(&data[..size_of::<DiscHeader>()]);
self.partition_header
.as_mut_bytes()
.copy_from_slice(&data[size_of::<DiscHeader>()..]);
self.boot_header.as_mut_bytes().copy_from_slice(&data[size_of::<DiscHeader>()..]);
*handled = true;
continue;
}
@@ -228,7 +226,7 @@ impl GCPartitionLayout {
// Locate other system files
let is_wii = self.disc_header.is_wii();
for (info, handled) in self.user_files.iter().zip(handled.iter_mut()) {
let dol_offset = self.partition_header.dol_offset(is_wii);
let dol_offset = self.boot_header.dol_offset(is_wii);
if (dol_offset != 0 && info.offset == Some(dol_offset)) || info.name == "sys/main.dol" {
let mut info = info.clone();
if info.alignment.is_none() {
@@ -239,7 +237,7 @@ impl GCPartitionLayout {
continue;
}
let fst_offset = self.partition_header.fst_offset(is_wii);
let fst_offset = self.boot_header.fst_offset(is_wii);
if (fst_offset != 0 && info.offset == Some(fst_offset)) || info.name == "sys/fst.bin" {
let mut data = Vec::with_capacity(info.size as usize);
file_callback(&mut data, &info.name)
@@ -293,7 +291,7 @@ impl GCPartitionLayout {
if let Some(audio_stream_buf_size) = overrides.audio_stream_buf_size {
self.disc_header.audio_stream_buf_size = audio_stream_buf_size;
}
let set_bi2 = self.raw_bi2.is_none() && overrides.region.is_some();
let set_bi2 = self.raw_bi2.is_none() || overrides.region.is_some();
let raw_bi2 = self.raw_bi2.get_or_insert_with(|| {
<[u8]>::new_box_zeroed_with_elems(BI2_SIZE).expect("Failed to allocate BI2")
});
@@ -383,11 +381,11 @@ impl GCPartitionLayout {
}
}
let raw_fst = builder.finalize();
if raw_fst.len() != self.partition_header.fst_size(is_wii) as usize {
if raw_fst.len() != self.boot_header.fst_size(is_wii) as usize {
return Err(Error::Other(format!(
"FST size mismatch: {} != {}",
raw_fst.len(),
self.partition_header.fst_size(is_wii)
self.boot_header.fst_size(is_wii)
)));
}
Ok(Arc::from(raw_fst))
@@ -411,7 +409,7 @@ impl GCPartitionLayout {
let mut boot = <[u8]>::new_box_zeroed_with_elems(BOOT_SIZE)?;
boot[..size_of::<DiscHeader>()].copy_from_slice(self.disc_header.as_bytes());
boot[size_of::<DiscHeader>()..].copy_from_slice(self.partition_header.as_bytes());
boot[size_of::<DiscHeader>()..].copy_from_slice(self.boot_header.as_bytes());
write_info.push(WriteInfo {
kind: WriteKind::Static(Arc::from(boot), "[BOOT]"),
size: BOOT_SIZE as u64,
@@ -433,21 +431,21 @@ impl GCPartitionLayout {
// Update DOL and FST offsets if not set
let is_wii = self.disc_header.is_wii();
let mut dol_offset = self.partition_header.dol_offset(is_wii);
let mut dol_offset = self.boot_header.dol_offset(is_wii);
if dol_offset == 0 {
dol_offset = align_up_64(last_offset, dol_file.alignment.unwrap() as u64);
self.partition_header.set_dol_offset(dol_offset, is_wii);
dol_offset = last_offset.align_up(dol_file.alignment.unwrap() as u64);
self.boot_header.set_dol_offset(dol_offset, is_wii);
}
let mut fst_offset = self.partition_header.fst_offset(is_wii);
let mut fst_offset = self.boot_header.fst_offset(is_wii);
if fst_offset == 0 {
// TODO handle DOL in user data
fst_offset = align_up_64(dol_offset + dol_file.size, 128);
self.partition_header.set_fst_offset(fst_offset, is_wii);
fst_offset = (dol_offset + dol_file.size).align_up(128);
self.boot_header.set_fst_offset(fst_offset, is_wii);
}
let fst_size = self.calculate_fst_size()?;
self.partition_header.set_fst_size(fst_size, is_wii);
if self.partition_header.fst_max_size(is_wii) < fst_size {
self.partition_header.set_fst_max_size(fst_size, is_wii);
self.boot_header.set_fst_size(fst_size, is_wii);
if self.boot_header.fst_max_size(is_wii) < fst_size {
self.boot_header.set_fst_max_size(fst_size, is_wii);
}
if dol_offset < fst_offset {
@@ -474,10 +472,10 @@ impl GCPartitionLayout {
let mut last_offset = self.layout_system_data(&mut system_write_info)?;
// Layout user data
let mut user_offset = self.partition_header.user_offset.get() as u64;
let mut user_offset = self.boot_header.user_offset.get() as u64;
if user_offset == 0 {
user_offset = align_up_64(last_offset, SECTOR_SIZE as u64);
self.partition_header.user_offset.set(user_offset as u32);
user_offset = last_offset.align_up(SECTOR_SIZE as u64);
self.boot_header.user_offset.set(user_offset as u32);
} else if user_offset < last_offset {
return Err(Error::Other(format!(
"User offset {:#X} is before FST {:#X}",
@@ -488,7 +486,7 @@ impl GCPartitionLayout {
for info in &self.user_files {
let offset = info
.offset
.unwrap_or_else(|| align_up_64(last_offset, info.alignment.unwrap_or(32) as u64));
.unwrap_or_else(|| last_offset.align_up(info.alignment.unwrap_or(32) as u64));
write_info.push(WriteInfo {
kind: WriteKind::File(info.name.clone()),
offset,
@@ -504,7 +502,7 @@ impl GCPartitionLayout {
write_info.push(WriteInfo {
kind: WriteKind::Static(fst_data, "[FST]"),
size: fst_size,
offset: self.partition_header.fst_offset(is_wii),
offset: self.boot_header.fst_offset(is_wii),
});
// Add system files to write info
write_info.extend(system_write_info);
@@ -512,17 +510,17 @@ impl GCPartitionLayout {
sort_files(&mut write_info)?;
// Update user size if not set
if self.partition_header.user_size.get() == 0 {
if self.boot_header.user_size.get() == 0 {
let user_end = if self.disc_header.is_wii() {
align_up_64(last_offset, SECTOR_SIZE as u64)
last_offset.align_up(SECTOR_SIZE as u64)
} else {
MINI_DVD_SIZE
};
self.partition_header.user_size.set((user_end - user_offset) as u32);
self.boot_header.user_size.set((user_end - user_offset) as u32);
}
// Insert junk data
let write_info = insert_junk_data(write_info, &self.partition_header);
let write_info = insert_junk_data(write_info, &self.boot_header);
Ok(write_info)
}
@@ -534,11 +532,11 @@ impl GCPartitionLayout {
pub(crate) fn insert_junk_data(
write_info: Vec<WriteInfo>,
partition_header: &PartitionHeader,
boot_header: &BootHeader,
) -> Vec<WriteInfo> {
let mut new_write_info = Vec::with_capacity(write_info.len());
let fst_end = partition_header.fst_offset(false) + partition_header.fst_size(false);
let fst_end = boot_header.fst_offset(false) + boot_header.fst_size(false);
let file_gap = find_file_gap(&write_info, fst_end);
let mut last_file_end = 0;
for info in write_info {
@@ -550,7 +548,7 @@ pub(crate) fn insert_junk_data(
// FST, and the junk data in between the inner and outer rim files. This attempts to
// determine the correct alignment, but is not 100% accurate.
let junk_start = if file_gap == Some(last_file_end) {
align_up_64(last_file_end, 4)
last_file_end.align_up(4)
} else {
aligned_end
};
@@ -565,8 +563,7 @@ pub(crate) fn insert_junk_data(
new_write_info.push(info);
}
let aligned_end = gcm_align(last_file_end);
let user_end =
partition_header.user_offset.get() as u64 + partition_header.user_size.get() as u64;
let user_end = boot_header.user_offset.get() as u64 + boot_header.user_size.get() as u64;
if aligned_end < user_end && aligned_end >= fst_end {
new_write_info.push(WriteInfo {
kind: WriteKind::Junk,

View File

@@ -6,8 +6,10 @@ use zerocopy::FromBytes;
use crate::{
disc::{
fst::Fst, wii::WiiPartitionHeader, DiscHeader, PartitionHeader, BOOT_SIZE, SECTOR_SIZE,
fst::Fst, wii::WiiPartitionHeader, BootHeader, DebugHeader, DiscHeader, BB2_OFFSET,
BOOT_SIZE, SECTOR_SIZE,
},
util::array_ref,
Error, Result,
};
@@ -308,7 +310,7 @@ pub struct PartitionInfo {
pub has_encryption: bool,
/// Whether the partition data hashes are present
pub has_hashes: bool,
/// Disc and partition header (boot.bin)
/// Disc and boot header (boot.bin)
pub raw_boot: Arc<[u8; BOOT_SIZE]>,
/// File system table (fst.bin), or `None` if partition is invalid
pub raw_fst: Option<Arc<[u8]>>,
@@ -330,15 +332,26 @@ impl PartitionInfo {
/// A view into the disc header.
#[inline]
pub fn disc_header(&self) -> &DiscHeader {
DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()])
DiscHeader::ref_from_bytes(array_ref![self.raw_boot, 0, size_of::<DiscHeader>()])
.expect("Invalid disc header alignment")
}
/// A view into the partition header.
/// A view into the debug header.
#[inline]
pub fn partition_header(&self) -> &PartitionHeader {
PartitionHeader::ref_from_bytes(&self.raw_boot[size_of::<DiscHeader>()..])
.expect("Invalid partition header alignment")
pub fn debug_header(&self) -> &DebugHeader {
DebugHeader::ref_from_bytes(array_ref![
self.raw_boot,
size_of::<DiscHeader>(),
size_of::<DebugHeader>()
])
.expect("Invalid debug header alignment")
}
/// A view into the boot header.
#[inline]
pub fn boot_header(&self) -> &BootHeader {
BootHeader::ref_from_bytes(array_ref![self.raw_boot, BB2_OFFSET, size_of::<BootHeader>()])
.expect("Invalid boot header alignment")
}
/// A view into the file system table (FST).

View File

@@ -14,6 +14,7 @@ use crate::{
Result,
};
#[derive(Clone)]
pub enum DirectDiscReaderMode {
Raw,
Partition { disc_header: Arc<DiscHeader>, data_start_sector: u32, key: KeyBytes },
@@ -31,6 +32,19 @@ pub struct DirectDiscReader {
mode: DirectDiscReaderMode,
}
impl Clone for DirectDiscReader {
fn clone(&self) -> Self {
Self {
io: self.io.clone(),
block: Block::default(),
block_buf: <[u8]>::new_box_zeroed_with_elems(self.block_buf.len()).unwrap(),
block_decrypted: false,
pos: 0,
mode: self.mode.clone(),
}
}
}
impl DirectDiscReader {
pub fn new(inner: Box<dyn BlockReader>) -> Result<Box<Self>> {
let block_size = inner.block_size() as usize;

View File

@@ -1,6 +1,6 @@
use std::{
io,
io::{BufRead, Read, Seek, SeekFrom},
io::{BufRead, Seek, SeekFrom},
mem::size_of,
sync::Arc,
};
@@ -10,11 +10,11 @@ use zerocopy::{FromBytes, FromZeros, IntoBytes};
use crate::{
disc::{
preloader::{fetch_sector_group, Preloader, SectorGroup, SectorGroupRequest},
ApploaderHeader, DiscHeader, DolHeader, PartitionHeader, BI2_SIZE, BOOT_SIZE,
SECTOR_GROUP_SIZE, SECTOR_SIZE,
ApploaderHeader, BootHeader, DolHeader, BB2_OFFSET, BI2_SIZE, BOOT_SIZE, SECTOR_GROUP_SIZE,
SECTOR_SIZE,
},
io::block::BlockReader,
read::{PartitionEncryption, PartitionMeta, PartitionReader},
read::{DiscStream, PartitionEncryption, PartitionMeta, PartitionReader},
util::{
impl_read_for_bufread,
read::{read_arc, read_arc_slice, read_from},
@@ -69,12 +69,12 @@ impl BufRead for PartitionReaderGC {
let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;
let group_idx = abs_sector / 64;
let abs_group_sector = group_idx * 64;
let max_groups = self.disc_size.div_ceil(SECTOR_GROUP_SIZE as u64) as u32;
let request = SectorGroupRequest {
group_idx,
partition_idx: None,
mode: PartitionEncryption::Original,
force_rehash: false,
};
// Load sector group
@@ -82,7 +82,7 @@ impl BufRead for PartitionReaderGC {
fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;
// Calculate the number of consecutive sectors in the group
let group_sector = abs_sector - abs_group_sector;
let group_sector = abs_sector - sector_group.start_sector;
let consecutive_sectors = sector_group.consecutive_sectors(group_sector);
if consecutive_sectors == 0 {
return Ok(&[]);
@@ -90,7 +90,7 @@ impl BufRead for PartitionReaderGC {
let num_sectors = group_sector + consecutive_sectors;
// Read from sector group buffer
let group_start = abs_group_sector as u64 * SECTOR_SIZE as u64;
let group_start = sector_group.start_sector as u64 * SECTOR_SIZE as u64;
let offset = (self.pos - group_start) as usize;
let end =
(num_sectors as u64 * SECTOR_SIZE as u64).min(self.disc_size - group_start) as usize;
@@ -131,12 +131,14 @@ impl PartitionReader for PartitionReaderGC {
}
pub(crate) fn read_dol(
reader: &mut dyn PartitionReader,
partition_header: &PartitionHeader,
// TODO: replace with &dyn mut DiscStream when trait_upcasting is stabilized
// https://github.com/rust-lang/rust/issues/65991
reader: &mut (impl DiscStream + ?Sized),
boot_header: &BootHeader,
is_wii: bool,
) -> Result<Arc<[u8]>> {
reader
.seek(SeekFrom::Start(partition_header.dol_offset(is_wii)))
.seek(SeekFrom::Start(boot_header.dol_offset(is_wii)))
.context("Seeking to DOL offset")?;
let dol_header: DolHeader = read_from(reader).context("Reading DOL header")?;
let dol_size = (dol_header.text_offs.iter().zip(&dol_header.text_sizes))
@@ -150,30 +152,28 @@ pub(crate) fn read_dol(
Ok(Arc::from(raw_dol))
}
pub(crate) fn read_fst<R>(
reader: &mut R,
partition_header: &PartitionHeader,
pub(crate) fn read_fst(
// TODO: replace with &dyn mut DiscStream when trait_upcasting is stabilized
// https://github.com/rust-lang/rust/issues/65991
reader: &mut (impl DiscStream + ?Sized),
boot_header: &BootHeader,
is_wii: bool,
) -> Result<Arc<[u8]>>
where
R: Read + Seek + ?Sized,
{
) -> Result<Arc<[u8]>> {
reader
.seek(SeekFrom::Start(partition_header.fst_offset(is_wii)))
.seek(SeekFrom::Start(boot_header.fst_offset(is_wii)))
.context("Seeking to FST offset")?;
let raw_fst: Arc<[u8]> = read_arc_slice(reader, partition_header.fst_size(is_wii) as usize)
let raw_fst: Arc<[u8]> = read_arc_slice(reader, boot_header.fst_size(is_wii) as usize)
.with_context(|| {
format!(
"Reading partition FST (offset {}, size {})",
partition_header.fst_offset(is_wii),
partition_header.fst_size(is_wii)
boot_header.fst_offset(is_wii),
boot_header.fst_size(is_wii)
)
})?;
Ok(raw_fst)
}
pub(crate) fn read_apploader<R>(reader: &mut R) -> Result<Arc<[u8]>>
where R: Read + Seek + ?Sized {
pub(crate) fn read_apploader(reader: &mut dyn DiscStream) -> Result<Arc<[u8]>> {
reader
.seek(SeekFrom::Start(BOOT_SIZE as u64 + BI2_SIZE as u64))
.context("Seeking to apploader offset")?;
@@ -190,14 +190,10 @@ where R: Read + Seek + ?Sized {
Ok(Arc::from(raw_apploader))
}
pub(crate) fn read_part_meta(
reader: &mut dyn PartitionReader,
is_wii: bool,
) -> Result<PartitionMeta> {
pub(crate) fn read_part_meta(reader: &mut dyn DiscStream, is_wii: bool) -> Result<PartitionMeta> {
// boot.bin
let raw_boot: Arc<[u8; BOOT_SIZE]> = read_arc(reader).context("Reading boot.bin")?;
let partition_header =
PartitionHeader::ref_from_bytes(&raw_boot[size_of::<DiscHeader>()..]).unwrap();
let boot_header = BootHeader::ref_from_bytes(&raw_boot[BB2_OFFSET..]).unwrap();
// bi2.bin
let raw_bi2: Arc<[u8; BI2_SIZE]> = read_arc(reader).context("Reading bi2.bin")?;
@@ -206,10 +202,10 @@ pub(crate) fn read_part_meta(
let raw_apploader = read_apploader(reader)?;
// fst.bin
let raw_fst = read_fst(reader, partition_header, is_wii)?;
let raw_fst = read_fst(reader, boot_header, is_wii)?;
// main.dol
let raw_dol = read_dol(reader, partition_header, is_wii)?;
let raw_dol = read_dol(reader, boot_header, is_wii)?;
Ok(PartitionMeta {
raw_boot,

View File

@@ -42,7 +42,10 @@ impl GroupHashes {
pub const NUM_H0_HASHES: usize = SECTOR_DATA_SIZE / HASHES_SIZE;
#[instrument(skip_all)]
pub fn hash_sector_group(sector_group: &[u8; SECTOR_GROUP_SIZE]) -> Box<GroupHashes> {
pub fn hash_sector_group(
sector_group: &[u8; SECTOR_GROUP_SIZE],
ignore_existing: bool,
) -> Box<GroupHashes> {
let mut result = GroupHashes::new_box_zeroed().unwrap();
for (h2_index, h2_hash) in result.h2_hashes.iter_mut().enumerate() {
let out_h1_hashes = array_ref_mut![result.h1_hashes, h2_index * 8, 8];
@@ -50,7 +53,9 @@ pub fn hash_sector_group(sector_group: &[u8; SECTOR_GROUP_SIZE]) -> Box<GroupHas
let sector = h1_index + h2_index * 8;
let out_h0_hashes =
array_ref_mut![result.h0_hashes, sector * NUM_H0_HASHES, NUM_H0_HASHES];
if array_ref![sector_group, sector * SECTOR_SIZE, 20].iter().any(|&v| v != 0) {
if !ignore_existing
&& array_ref![sector_group, sector * SECTOR_SIZE, 20].iter().any(|&v| v != 0)
{
// Hash block already present, use it
out_h0_hashes.as_mut_bytes().copy_from_slice(array_ref![
sector_group,

View File

@@ -27,10 +27,13 @@ pub const WII_MAGIC: MagicBytes = [0x5D, 0x1C, 0x9E, 0xA3];
/// Magic bytes for GameCube discs. Located at offset 0x1C.
pub const GCN_MAGIC: MagicBytes = [0xC2, 0x33, 0x9F, 0x3D];
/// Size in bytes of the disc header and partition header. (boot.bin)
pub const BOOT_SIZE: usize = size_of::<DiscHeader>() + size_of::<PartitionHeader>();
/// Offset in bytes of the boot block within a disc partition.
pub const BB2_OFFSET: usize = 0x420;
/// Size in bytes of the debug and region information. (bi2.bin)
/// Size in bytes of the disc header, debug block and boot block. (boot.bin)
pub const BOOT_SIZE: usize = 0x440;
/// Size in bytes of the DVD Boot Info (debug and region information, bi2.bin)
pub const BI2_SIZE: usize = 0x2000;
/// The size of a single-layer MiniDVD. (1.4 GB)
@@ -114,20 +117,28 @@ impl DiscHeader {
pub fn has_partition_encryption(&self) -> bool { self.no_partition_encryption == 0 }
}
/// A header describing the contents of a disc partition.
/// The debug block of a disc partition.
///
/// **GameCube**: Always follows the disc header.
///
/// **Wii**: Follows the disc header within each partition.
/// Located at offset 0x400 (following the disc header) within each partition.
#[derive(Clone, Debug, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)]
#[repr(C, align(4))]
pub struct PartitionHeader {
pub struct DebugHeader {
/// Debug monitor offset
pub debug_mon_offset: U32,
/// Debug monitor load address
pub debug_load_address: U32,
/// Padding
_pad1: [u8; 0x18],
}
static_assert!(size_of::<DebugHeader>() == 0x20);
/// The boot block (BB2) of a disc partition.
///
/// Located at offset 0x420 (following the debug block) within each partition.
#[derive(Clone, Debug, PartialEq, FromBytes, IntoBytes, Immutable, KnownLayout)]
#[repr(C, align(4))]
pub struct BootHeader {
/// Offset to main DOL (Wii: >> 2)
pub dol_offset: U32,
/// Offset to file system table (Wii: >> 2)
@@ -146,9 +157,12 @@ pub struct PartitionHeader {
_pad2: [u8; 4],
}
static_assert!(size_of::<PartitionHeader>() == 0x40);
static_assert!(size_of::<BootHeader>() == 0x20);
static_assert!(
size_of::<DiscHeader>() + size_of::<DebugHeader>() + size_of::<BootHeader>() == BOOT_SIZE
);
impl PartitionHeader {
impl BootHeader {
/// Offset within the partition to the main DOL.
#[inline]
pub fn dol_offset(&self, is_wii: bool) -> u64 {

View File

@@ -1,5 +1,6 @@
use std::{
collections::HashMap,
fmt::{Display, Formatter},
io,
num::NonZeroUsize,
sync::{Arc, Mutex},
@@ -19,7 +20,9 @@ use zerocopy::FromZeros;
use crate::{
common::PartitionInfo,
disc::{
hashes::hash_sector_group, wii::HASHES_SIZE, DiscHeader, SECTOR_GROUP_SIZE, SECTOR_SIZE,
hashes::{hash_sector_group, GroupHashes},
wii::HASHES_SIZE,
DiscHeader, SECTOR_GROUP_SIZE, SECTOR_SIZE,
},
io::{
block::{Block, BlockKind, BlockReader},
@@ -30,6 +33,7 @@ use crate::{
aes::{decrypt_sector, encrypt_sector},
array_ref_mut,
},
IoResultContext,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -37,14 +41,27 @@ pub struct SectorGroupRequest {
pub group_idx: u32,
pub partition_idx: Option<u8>,
pub mode: PartitionEncryption,
pub force_rehash: bool,
}
impl Display for SectorGroupRequest {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self.partition_idx {
Some(idx) => write!(f, "Partition {} group {}", idx, self.group_idx),
None => write!(f, "Group {}", self.group_idx),
}
}
}
#[derive(Clone)]
pub struct SectorGroup {
pub request: SectorGroupRequest,
pub start_sector: u32,
pub data: Bytes,
pub sector_bitmap: u64,
pub io_duration: Option<Duration>,
#[allow(unused)] // TODO WIA hash exceptions
pub group_hashes: Option<Arc<GroupHashes>>,
}
impl SectorGroup {
@@ -208,9 +225,12 @@ fn preloader_thread(
let end = Instant::now();
last_request_end = Some(end);
let req_time = end - start;
stat_tx
if stat_tx
.send(PreloaderThreadStats { thread_id, wait_time, req_time, io_time })
.expect("Failed to send preloader stats");
.is_err()
{
break;
}
}
})
.expect("Failed to spawn preloader thread")
@@ -329,6 +349,18 @@ impl Clone for SectorGroupLoader {
}
}
#[derive(Default)]
struct LoadedSectorGroup {
/// Start sector of the group
start_sector: u32,
/// Bitmap of sectors that were read
sector_bitmap: u64,
/// Total duration of I/O operations
io_duration: Option<Duration>,
/// Calculated sector group hashes
group_hashes: Option<Arc<GroupHashes>>,
}
impl SectorGroupLoader {
pub fn new(
io: Box<dyn BlockReader>,
@@ -344,13 +376,21 @@ impl SectorGroupLoader {
let mut sector_group_buf = BytesMut::zeroed(SECTOR_GROUP_SIZE);
let out = array_ref_mut![sector_group_buf, 0, SECTOR_GROUP_SIZE];
let (sector_bitmap, io_duration) = if request.partition_idx.is_some() {
self.load_partition_group(request, out)?
} else {
self.load_raw_group(request, out)?
};
let LoadedSectorGroup { start_sector, sector_bitmap, io_duration, group_hashes } =
if request.partition_idx.is_some() {
self.load_partition_group(request, out)?
} else {
self.load_raw_group(request, out)?
};
Ok(SectorGroup { request, data: sector_group_buf.freeze(), sector_bitmap, io_duration })
Ok(SectorGroup {
request,
start_sector,
data: sector_group_buf.freeze(),
sector_bitmap,
io_duration,
group_hashes,
})
}
/// Load a sector group from a partition.
@@ -360,16 +400,16 @@ impl SectorGroupLoader {
&mut self,
request: SectorGroupRequest,
sector_group_buf: &mut [u8; SECTOR_GROUP_SIZE],
) -> io::Result<(u64, Option<Duration>)> {
) -> io::Result<LoadedSectorGroup> {
let Some(partition) =
request.partition_idx.and_then(|idx| self.partitions.get(idx as usize))
else {
return Ok((0, None));
return Ok(LoadedSectorGroup::default());
};
let abs_group_sector = partition.data_start_sector + request.group_idx * 64;
if abs_group_sector >= partition.data_end_sector {
return Ok((0, None));
return Ok(LoadedSectorGroup::default());
}
// Bitmap of sectors that were read
@@ -382,6 +422,8 @@ impl SectorGroupLoader {
let mut hash_exceptions = Vec::<WIAException>::new();
// Total duration of I/O operations
let mut io_duration = None;
// Calculated sector group hashes
let mut group_hashes = None;
// Read sector group
for sector in 0..64 {
@@ -397,7 +439,10 @@ impl SectorGroupLoader {
// Read new block
if !self.block.contains(abs_sector) {
self.block = self.io.read_block(self.block_buf.as_mut(), abs_sector)?;
self.block = self
.io
.read_block(self.block_buf.as_mut(), abs_sector)
.io_with_context(|| format!("Reading block for sector {abs_sector}"))?;
if let Some(duration) = self.block.io_duration {
*io_duration.get_or_insert_with(Duration::default) += duration;
}
@@ -408,16 +453,21 @@ impl SectorGroupLoader {
}
// Add hash exceptions
self.block.append_hash_exceptions(abs_sector, sector, &mut hash_exceptions)?;
self.block
.append_hash_exceptions(abs_sector, sector, &mut hash_exceptions)
.io_with_context(|| format!("Appending hash exceptions for sector {abs_sector}"))?;
// Read new sector into buffer
let (encrypted, has_hashes) = self.block.copy_sector(
sector_data,
self.block_buf.as_mut(),
abs_sector,
partition.disc_header(),
Some(partition),
)?;
let (encrypted, has_hashes) = self
.block
.copy_sector(
sector_data,
self.block_buf.as_mut(),
abs_sector,
partition.disc_header(),
Some(partition),
)
.io_with_context(|| format!("Copying sector {abs_sector} from block"))?;
if !encrypted {
decrypted_sectors |= 1 << sector;
}
@@ -428,7 +478,9 @@ impl SectorGroupLoader {
}
// Recover hashes
if request.mode != PartitionEncryption::ForceDecryptedNoHashes && hash_recovery_sectors != 0
if request.force_rehash
|| (request.mode != PartitionEncryption::ForceDecryptedNoHashes
&& hash_recovery_sectors != 0)
{
// Decrypt any encrypted sectors
if decrypted_sectors != u64::MAX {
@@ -443,16 +495,19 @@ impl SectorGroupLoader {
}
// Recover hashes
let group_hashes = hash_sector_group(sector_group_buf);
let hashes = hash_sector_group(sector_group_buf, request.force_rehash);
// Apply hashes
for sector in 0..64 {
let sector_data =
array_ref_mut![sector_group_buf, sector * SECTOR_SIZE, SECTOR_SIZE];
if (hash_recovery_sectors >> sector) & 1 == 1 {
group_hashes.apply(sector_data, sector);
hashes.apply(sector_data, sector);
}
}
// Persist hashes
group_hashes = Some(Arc::from(hashes));
}
// Apply hash exceptions
@@ -508,7 +563,12 @@ impl SectorGroupLoader {
}
}
Ok((sector_bitmap, io_duration))
Ok(LoadedSectorGroup {
start_sector: abs_group_sector,
sector_bitmap,
io_duration,
group_hashes,
})
}
/// Loads a non-partition sector group.
@@ -516,7 +576,7 @@ impl SectorGroupLoader {
&mut self,
request: SectorGroupRequest,
sector_group_buf: &mut [u8; SECTOR_GROUP_SIZE],
) -> io::Result<(u64, Option<Duration>)> {
) -> io::Result<LoadedSectorGroup> {
let abs_group_sector = request.group_idx * 64;
// Bitmap of sectors that were read
@@ -534,7 +594,10 @@ impl SectorGroupLoader {
// Read new block
if !self.block.contains(abs_sector) {
self.block = self.io.read_block(self.block_buf.as_mut(), abs_sector)?;
self.block = self
.io
.read_block(self.block_buf.as_mut(), abs_sector)
.io_with_context(|| format!("Reading block for sector {abs_sector}"))?;
if let Some(duration) = self.block.io_duration {
*io_duration.get_or_insert_with(Duration::default) += duration;
}
@@ -544,17 +607,24 @@ impl SectorGroupLoader {
}
// Read new sector into buffer
self.block.copy_sector(
sector_data,
self.block_buf.as_mut(),
abs_sector,
self.disc_header.as_ref(),
None,
)?;
self.block
.copy_sector(
sector_data,
self.block_buf.as_mut(),
abs_sector,
self.disc_header.as_ref(),
None,
)
.io_with_context(|| format!("Copying sector {abs_sector} from block"))?;
sector_bitmap |= 1 << sector;
}
Ok((sector_bitmap, io_duration))
Ok(LoadedSectorGroup {
start_sector: abs_group_sector,
sector_bitmap,
io_duration,
group_hashes: None,
})
}
}

View File

@@ -5,6 +5,7 @@ use std::{
};
use bytes::Bytes;
use polonius_the_crab::{polonius, polonius_return};
use tracing::warn;
use zerocopy::{FromBytes, IntoBytes};
@@ -21,13 +22,13 @@ use crate::{
PartitionReaderWii, WiiPartEntry, WiiPartGroup, WiiPartitionHeader, REGION_OFFSET,
REGION_SIZE, WII_PART_GROUP_OFF,
},
DiscHeader, PartitionHeader, BOOT_SIZE, DL_DVD_SIZE, MINI_DVD_SIZE, SECTOR_GROUP_SIZE,
SECTOR_SIZE, SL_DVD_SIZE,
BootHeader, DiscHeader, BB2_OFFSET, BOOT_SIZE, DL_DVD_SIZE, MINI_DVD_SIZE,
SECTOR_GROUP_SIZE, SECTOR_SIZE, SL_DVD_SIZE,
},
io::block::BlockReader,
read::{DiscMeta, DiscOptions, PartitionEncryption, PartitionOptions, PartitionReader},
util::{
impl_read_for_bufread,
array_ref, impl_read_for_bufread,
read::{read_arc, read_from, read_vec},
},
Error, Result, ResultContext,
@@ -210,15 +211,18 @@ impl DiscReader {
}
}
/// A reference to the disc's boot header (BB2) for GameCube discs.
/// For Wii discs, use the boot header from the appropriate [PartitionInfo].
#[inline]
pub fn partition_header(&self) -> Option<&PartitionHeader> {
pub fn boot_header(&self) -> Option<&BootHeader> {
match &self.disc_data {
DiscReaderData::GameCube { .. } => Some(
PartitionHeader::ref_from_bytes(
&self.raw_boot[size_of::<DiscHeader>()
..size_of::<DiscHeader>() + size_of::<PartitionHeader>()],
)
.expect("Invalid partition header alignment"),
BootHeader::ref_from_bytes(array_ref![
self.raw_boot,
BB2_OFFSET,
size_of::<BootHeader>()
])
.expect("Invalid boot header alignment"),
),
_ => None,
}
@@ -306,48 +310,62 @@ impl DiscReader {
}
}
pub fn fill_buf_internal(&mut self) -> io::Result<Bytes> {
if self.pos >= self.size {
return Ok(Bytes::new());
}
// Read from modified disc header
if self.pos < size_of::<DiscHeader>() as u64 {
if let Some(alt_disc_header) = &self.alt_disc_header {
return Ok(Bytes::copy_from_slice(
&alt_disc_header.as_bytes()[self.pos as usize..],
));
}
}
// Build sector group request
let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;
let (request, abs_group_sector, max_groups) = if let Some(partition) =
pub fn load_sector_group(
&mut self,
abs_sector: u32,
force_rehash: bool,
) -> io::Result<(&SectorGroup, bool)> {
let (request, max_groups) = if let Some(partition) =
self.orig_partitions().iter().find(|part| part.data_contains_sector(abs_sector))
{
let group_idx = (abs_sector - partition.data_start_sector) / 64;
let abs_group_sector = partition.data_start_sector + group_idx * 64;
let max_groups = (partition.data_end_sector - partition.data_start_sector).div_ceil(64);
let request = SectorGroupRequest {
group_idx,
partition_idx: Some(partition.index as u8),
mode: self.mode,
force_rehash,
};
(request, abs_group_sector, max_groups)
(request, max_groups)
} else {
let group_idx = abs_sector / 64;
let abs_group_sector = group_idx * 64;
let max_groups = self.size.div_ceil(SECTOR_GROUP_SIZE as u64) as u32;
let request = SectorGroupRequest { group_idx, partition_idx: None, mode: self.mode };
(request, abs_group_sector, max_groups)
let request = SectorGroupRequest {
group_idx,
partition_idx: None,
mode: self.mode,
force_rehash,
};
(request, max_groups)
};
// Load sector group
let (sector_group, _updated) =
let (sector_group, updated) =
fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;
Ok((sector_group, updated))
}
pub fn fill_buf_internal(&mut self) -> io::Result<Bytes> {
let pos = self.pos;
let size = self.size;
if pos >= size {
return Ok(Bytes::new());
}
// Read from modified disc header
if pos < size_of::<DiscHeader>() as u64 {
if let Some(alt_disc_header) = &self.alt_disc_header {
return Ok(Bytes::copy_from_slice(&alt_disc_header.as_bytes()[pos as usize..]));
}
}
// Load sector group
let abs_sector = (pos / SECTOR_SIZE as u64) as u32;
let (sector_group, _updated) = self.load_sector_group(abs_sector, false)?;
// Calculate the number of consecutive sectors in the group
let group_sector = abs_sector - abs_group_sector;
let group_sector = abs_sector - sector_group.start_sector;
let consecutive_sectors = sector_group.consecutive_sectors(group_sector);
if consecutive_sectors == 0 {
return Ok(Bytes::new());
@@ -355,54 +373,37 @@ impl DiscReader {
let num_sectors = group_sector + consecutive_sectors;
// Read from sector group buffer
let group_start = abs_group_sector as u64 * SECTOR_SIZE as u64;
let offset = (self.pos - group_start) as usize;
let end = (num_sectors as u64 * SECTOR_SIZE as u64).min(self.size - group_start) as usize;
let group_start = sector_group.start_sector as u64 * SECTOR_SIZE as u64;
let offset = (pos - group_start) as usize;
let end = (num_sectors as u64 * SECTOR_SIZE as u64).min(size - group_start) as usize;
Ok(sector_group.data.slice(offset..end))
}
}
impl BufRead for DiscReader {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
if self.pos >= self.size {
let pos = self.pos;
let size = self.size;
if pos >= size {
return Ok(&[]);
}
// Read from modified disc header
if self.pos < size_of::<DiscHeader>() as u64 {
if let Some(alt_disc_header) = &self.alt_disc_header {
return Ok(&alt_disc_header.as_bytes()[self.pos as usize..]);
let mut this = self;
polonius!(|this| -> io::Result<&'polonius [u8]> {
// Read from modified disc header
if pos < size_of::<DiscHeader>() as u64 {
if let Some(alt_disc_header) = &this.alt_disc_header {
polonius_return!(Ok(&alt_disc_header.as_bytes()[pos as usize..]));
}
}
}
// Build sector group request
let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;
let (request, abs_group_sector, max_groups) = if let Some(partition) =
self.orig_partitions().iter().find(|part| part.data_contains_sector(abs_sector))
{
let group_idx = (abs_sector - partition.data_start_sector) / 64;
let abs_group_sector = partition.data_start_sector + group_idx * 64;
let max_groups = (partition.data_end_sector - partition.data_start_sector).div_ceil(64);
let request = SectorGroupRequest {
group_idx,
partition_idx: Some(partition.index as u8),
mode: self.mode,
};
(request, abs_group_sector, max_groups)
} else {
let group_idx = abs_sector / 64;
let abs_group_sector = group_idx * 64;
let max_groups = self.size.div_ceil(SECTOR_GROUP_SIZE as u64) as u32;
let request = SectorGroupRequest { group_idx, partition_idx: None, mode: self.mode };
(request, abs_group_sector, max_groups)
};
});
// Load sector group
let (sector_group, _updated) =
fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;
let abs_sector = (pos / SECTOR_SIZE as u64) as u32;
let (sector_group, _updated) = this.load_sector_group(abs_sector, false)?;
// Calculate the number of consecutive sectors in the group
let group_sector = abs_sector - abs_group_sector;
let group_sector = abs_sector - sector_group.start_sector;
let consecutive_sectors = sector_group.consecutive_sectors(group_sector);
if consecutive_sectors == 0 {
return Ok(&[]);
@@ -410,9 +411,9 @@ impl BufRead for DiscReader {
let num_sectors = group_sector + consecutive_sectors;
// Read from sector group buffer
let group_start = abs_group_sector as u64 * SECTOR_SIZE as u64;
let offset = (self.pos - group_start) as usize;
let end = (num_sectors as u64 * SECTOR_SIZE as u64).min(self.size - group_start) as usize;
let group_start = sector_group.start_sector as u64 * SECTOR_SIZE as u64;
let offset = (pos - group_start) as usize;
let end = (num_sectors as u64 * SECTOR_SIZE as u64).min(size - group_start) as usize;
Ok(&sector_group.data[offset..end])
}
@@ -490,37 +491,43 @@ fn read_partition_info(
data_start_sector,
key,
});
let raw_boot: Arc<[u8; BOOT_SIZE]> =
read_arc(reader).context("Reading partition headers")?;
let raw_boot: Arc<[u8; BOOT_SIZE]> = read_arc(reader).context("Reading boot data")?;
let partition_disc_header =
DiscHeader::ref_from_bytes(&raw_boot[..size_of::<DiscHeader>()])
DiscHeader::ref_from_bytes(array_ref![raw_boot, 0, size_of::<DiscHeader>()])
.expect("Invalid disc header alignment");
let partition_header =
PartitionHeader::ref_from_bytes(&raw_boot[size_of::<DiscHeader>()..])
.expect("Invalid partition header alignment");
let boot_header = BootHeader::ref_from_bytes(&raw_boot[BB2_OFFSET..])
.expect("Invalid boot header alignment");
let raw_fst = if partition_disc_header.is_wii() {
let raw_fst = read_fst(reader, partition_header, true)?;
let fst = Fst::new(&raw_fst)?;
let max_fst_offset = fst
.nodes
.iter()
.filter_map(|n| match n.kind() {
NodeKind::File => Some(n.offset(true) + n.length() as u64),
_ => None,
})
.max()
.unwrap_or(0);
if max_fst_offset > data_size {
if data_size == 0 {
// Guess data size for decrypted partitions
data_end_sector = max_fst_offset.div_ceil(SECTOR_SIZE as u64) as u32;
} else {
return Err(Error::DiscFormat(format!(
"Partition {group_idx}:{part_idx} FST exceeds data size",
)));
let raw_fst = read_fst(reader, boot_header, true)?;
match Fst::new(&raw_fst) {
Ok(fst) => {
let max_fst_offset = fst
.nodes
.iter()
.filter_map(|n| match n.kind() {
NodeKind::File => Some(n.offset(true) + n.length() as u64),
_ => None,
})
.max()
.unwrap_or(0);
if max_fst_offset > data_size {
if data_size == 0 {
// Guess data size for decrypted partitions
data_end_sector =
max_fst_offset.div_ceil(SECTOR_SIZE as u64) as u32;
} else {
return Err(Error::DiscFormat(format!(
"Partition {group_idx}:{part_idx} FST exceeds data size",
)));
}
}
Some(raw_fst)
}
Err(e) => {
warn!("Partition {group_idx}:{part_idx} FST is not valid: {e}");
None
}
}
Some(raw_fst)
} else {
warn!("Partition {group_idx}:{part_idx} is not valid");
None

View File

@@ -378,6 +378,7 @@ impl BufRead for PartitionReaderWii {
} else {
PartitionEncryption::ForceDecryptedNoHashes
},
force_rehash: false,
};
// Load sector group

View File

@@ -107,7 +107,8 @@ pub const WBFS_MAGIC: MagicBytes = *b"WBFS";
pub const WIA_MAGIC: MagicBytes = *b"WIA\x01";
pub const RVZ_MAGIC: MagicBytes = *b"RVZ\x01";
pub fn detect<R: Read + ?Sized>(stream: &mut R) -> io::Result<Option<Format>> {
pub fn detect<R>(stream: &mut R) -> io::Result<Option<Format>>
where R: Read + ?Sized {
let data: [u8; 0x20] = match read_from(stream) {
Ok(magic) => magic,
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),

View File

@@ -15,14 +15,14 @@ use crate::{
gcn::{read_dol, read_fst},
reader::DiscReader,
writer::{DataCallback, DiscWriter},
DiscHeader, PartitionHeader, SECTOR_SIZE,
BootHeader, DiscHeader, BB2_OFFSET, SECTOR_SIZE,
},
io::block::{Block, BlockKind, BlockReader, TGC_MAGIC},
read::{DiscMeta, DiscStream, PartitionOptions, PartitionReader},
util::{
align_up_32, array_ref,
array_ref,
read::{read_arc, read_arc_slice, read_from, read_with_zero_fill},
static_assert,
static_assert, Align,
},
write::{DiscFinalization, DiscWriterWeight, FormatOptions, ProcessOptions},
Error, Result, ResultContext,
@@ -82,19 +82,21 @@ impl BlockReaderTGC {
}
let disc_size = (header.gcm_files_start.get() + header.user_size.get()) as u64;
// Read disc header and partition header
// Read GCM header
inner
.seek(SeekFrom::Start(header.header_offset.get() as u64))
.context("Seeking to GCM header")?;
let raw_header =
read_arc::<[u8; GCM_HEADER_SIZE], _>(inner.as_mut()).context("Reading GCM header")?;
let (disc_header, remain) = DiscHeader::ref_from_prefix(raw_header.as_ref())
.expect("Invalid disc header alignment");
let disc_header =
DiscHeader::ref_from_bytes(array_ref![raw_header, 0, size_of::<DiscHeader>()])
.expect("Invalid disc header alignment");
let disc_header = disc_header.clone();
let (partition_header, _) =
PartitionHeader::ref_from_prefix(remain).expect("Invalid partition header alignment");
let partition_header = partition_header.clone();
let boot_header =
BootHeader::ref_from_bytes(array_ref![raw_header, BB2_OFFSET, size_of::<BootHeader>()])
.expect("Invalid boot header alignment");
let boot_header = boot_header.clone();
// Read DOL
inner.seek(SeekFrom::Start(header.dol_offset.get() as u64)).context("Seeking to DOL")?;
@@ -116,12 +118,12 @@ impl BlockReaderTGC {
write_info.push(WriteInfo {
kind: WriteKind::Static(raw_dol, "sys/main.dol"),
size: header.dol_size.get() as u64,
offset: partition_header.dol_offset(false),
offset: boot_header.dol_offset(false),
});
write_info.push(WriteInfo {
kind: WriteKind::Static(raw_fst.clone(), "sys/fst.bin"),
size: header.fst_size.get() as u64,
offset: partition_header.fst_offset(false),
offset: boot_header.fst_offset(false),
});
// Collect files
@@ -136,7 +138,7 @@ impl BlockReaderTGC {
});
}
write_info.sort_unstable_by(|a, b| a.offset.cmp(&b.offset).then(a.size.cmp(&b.size)));
let write_info = insert_junk_data(write_info, &partition_header);
let write_info = insert_junk_data(write_info, &boot_header);
let file_callback = FileCallbackTGC::new(inner, raw_fst, header);
let disc_id = *array_ref![disc_header.game_id, 0, 4];
@@ -225,14 +227,13 @@ impl DiscWriterTGC {
// Read GCM header
let mut raw_header = <[u8; GCM_HEADER_SIZE]>::new_box_zeroed()?;
inner.read_exact(raw_header.as_mut()).context("Reading GCM header")?;
let (_, remain) = DiscHeader::ref_from_prefix(raw_header.as_ref())
.expect("Invalid disc header alignment");
let (partition_header, _) =
PartitionHeader::ref_from_prefix(remain).expect("Invalid partition header alignment");
let boot_header =
BootHeader::ref_from_bytes(array_ref![raw_header, BB2_OFFSET, size_of::<BootHeader>()])
.expect("Invalid boot header alignment");
// Read DOL
let raw_dol = read_dol(inner.as_mut(), partition_header, false)?;
let raw_fst = read_fst(inner.as_mut(), partition_header, false)?;
let raw_dol = read_dol(inner.as_mut(), boot_header, false)?;
let raw_fst = read_fst(inner.as_mut(), boot_header, false)?;
// Parse FST
let fst = Fst::new(&raw_fst)?;
@@ -249,11 +250,10 @@ impl DiscWriterTGC {
// Layout system files
let gcm_header_offset = SECTOR_SIZE as u32;
let fst_offset = gcm_header_offset + GCM_HEADER_SIZE as u32;
let dol_offset = align_up_32(fst_offset + partition_header.fst_size.get(), 32);
let dol_offset = (fst_offset + boot_header.fst_size.get()).align_up(32);
let user_size =
partition_header.user_offset.get() + partition_header.user_size.get() - gcm_files_start;
let user_end =
align_up_32(dol_offset + raw_dol.len() as u32 + user_size, SECTOR_SIZE as u32);
boot_header.user_offset.get() + boot_header.user_size.get() - gcm_files_start;
let user_end = (dol_offset + raw_dol.len() as u32 + user_size).align_up(SECTOR_SIZE as u32);
let user_offset = user_end - user_size;
let header = TGCHeader {
@@ -262,8 +262,8 @@ impl DiscWriterTGC {
header_offset: gcm_header_offset.into(),
header_size: (GCM_HEADER_SIZE as u32).into(),
fst_offset: fst_offset.into(),
fst_size: partition_header.fst_size,
fst_max_size: partition_header.fst_max_size,
fst_size: boot_header.fst_size,
fst_max_size: boot_header.fst_max_size,
dol_offset: dol_offset.into(),
dol_size: (raw_dol.len() as u32).into(),
user_offset: user_offset.into(),

View File

@@ -19,7 +19,7 @@ use crate::{
reader::DiscReader,
wii::{HASHES_SIZE, SECTOR_DATA_SIZE},
writer::{par_process, read_block, BlockProcessor, BlockResult, DataCallback, DiscWriter},
DiscHeader, PartitionHeader, SECTOR_SIZE,
BootHeader, DiscHeader, SECTOR_SIZE,
},
io::{
block::{Block, BlockKind, BlockReader, RVZ_MAGIC, WIA_MAGIC},
@@ -28,15 +28,15 @@ use crate::{
read::{DiscMeta, DiscStream},
util::{
aes::decrypt_sector_data_b2b,
align_up_32, align_up_64, array_ref, array_ref_mut,
array_ref, array_ref_mut,
compress::{Compressor, DecompressionKind, Decompressor},
digest::{sha1_hash, xxh64_hash, DigestManager},
lfg::{LaggedFibonacci, SEED_SIZE, SEED_SIZE_BYTES},
read::{read_arc_slice, read_from, read_vec},
static_assert,
static_assert, Align,
},
write::{DiscFinalization, DiscWriterWeight, FormatOptions, ProcessOptions},
Error, Result, ResultContext,
Error, IoResultContext, Result, ResultContext,
};
const WIA_VERSION: u32 = 0x01000000;
@@ -395,7 +395,7 @@ pub struct WIARawData {
}
impl WIARawData {
pub fn start_offset(&self) -> u64 { self.raw_data_offset.get() & !(SECTOR_SIZE as u64 - 1) }
pub fn start_offset(&self) -> u64 { self.raw_data_offset.get().align_down(SECTOR_SIZE as u64) }
pub fn start_sector(&self) -> u32 { (self.start_offset() / SECTOR_SIZE as u64) as u32 }
@@ -457,19 +457,21 @@ pub struct RVZGroup {
pub rvz_packed_size: U32,
}
const COMPRESSED_BIT: u32 = 1 << 31;
impl RVZGroup {
#[inline]
pub fn data_size(&self) -> u32 { self.data_size_and_flag.get() & 0x7FFFFFFF }
pub fn data_size(&self) -> u32 { self.data_size_and_flag.get() & !COMPRESSED_BIT }
#[inline]
pub fn is_compressed(&self) -> bool { self.data_size_and_flag.get() & 0x80000000 != 0 }
pub fn is_compressed(&self) -> bool { self.data_size_and_flag.get() & COMPRESSED_BIT != 0 }
}
impl From<&WIAGroup> for RVZGroup {
fn from(value: &WIAGroup) -> Self {
Self {
data_offset: value.data_offset,
data_size_and_flag: U32::new(value.data_size.get() | 0x80000000),
data_size_and_flag: U32::new(value.data_size.get() | COMPRESSED_BIT),
rvz_packed_size: U32::new(0),
}
}
@@ -886,23 +888,40 @@ impl BlockReader for BlockReaderWIA {
|| !group.is_compressed();
let mut exception_lists = vec![];
if info.in_partition && uncompressed_exception_lists {
exception_lists = read_exception_lists(&mut group_data, chunk_size, true)?;
exception_lists = read_exception_lists(&mut group_data, chunk_size, true)
.io_context("Reading uncompressed exception lists")?;
}
let mut decompressed = if group.is_compressed() {
let mut decompressed = BytesMut::zeroed(chunk_size as usize);
let len = self.decompressor.decompress(group_data.as_ref(), decompressed.as_mut())?;
let max_decompressed_size =
self.decompressor.get_content_size(group_data.as_ref())?.unwrap_or_else(|| {
if info.in_partition && !uncompressed_exception_lists {
// Add room for potential hash exceptions. See [WIAExceptionList] for details on the
// maximum number of exceptions.
chunk_size as usize
+ (2 + 3328 * size_of::<WIAException>())
* (chunk_size as usize).div_ceil(0x200000)
} else {
chunk_size as usize
}
});
let mut decompressed = BytesMut::zeroed(max_decompressed_size);
let len = self
.decompressor
.decompress(group_data.as_ref(), decompressed.as_mut())
.io_context("Decompressing group data")?;
decompressed.truncate(len);
decompressed.freeze()
} else {
group_data
};
if info.in_partition && !uncompressed_exception_lists {
exception_lists = read_exception_lists(&mut decompressed, chunk_size, false)?;
exception_lists = read_exception_lists(&mut decompressed, chunk_size, false)
.io_context("Reading compressed exception lists")?;
}
if group.rvz_packed_size.get() > 0 {
// Decode RVZ packed data
rvz_unpack(&mut decompressed, out, &info)?;
rvz_unpack(&mut decompressed, out, &info).io_context("Unpacking RVZ group data")?;
} else {
// Read and decompress data
if decompressed.remaining() != info.size as usize {
@@ -975,9 +994,9 @@ fn rvz_unpack(data: &mut impl Buf, out: &mut [u8], info: &GroupInfo) -> io::Resu
while data.remaining() >= 4 {
let size = data.get_u32();
let remain = out.len() - read;
if size & 0x80000000 != 0 {
if size & COMPRESSED_BIT != 0 {
// Junk data
let size = size & 0x7FFFFFFF;
let size = size & !COMPRESSED_BIT;
if size as usize > remain {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
@@ -1057,14 +1076,13 @@ impl JunkInfo {
start_sector: u32,
end_sector: u32,
disc_header: &DiscHeader,
partition_header: Option<&PartitionHeader>,
boot_header: Option<&BootHeader>,
fst: Option<Fst>,
) -> Self {
let is_wii = disc_header.is_wii();
let mut file_ends = BTreeSet::new();
if let Some(partition_header) = partition_header {
file_ends
.insert(partition_header.fst_offset(is_wii) + partition_header.fst_size(is_wii));
if let Some(boot_header) = boot_header {
file_ends.insert(boot_header.fst_offset(is_wii) + boot_header.fst_size(is_wii));
}
if let Some(fst) = fst {
for entry in fst.nodes.iter().filter(|n| n.is_file()) {
@@ -1145,7 +1163,7 @@ impl BlockProcessor for BlockProcessorWIA {
};
let uncompressed_size =
align_up_32(hash_exception_data.len() as u32, 4) + group_data.len() as u32;
(hash_exception_data.len() as u32).align_up(4) + group_data.len() as u32;
if hash_exception_data.as_ref().iter().all(|&b| b == 0)
&& group_data.as_ref().iter().all(|&b| b == 0)
{
@@ -1167,7 +1185,7 @@ impl BlockProcessor for BlockProcessorWIA {
if is_rvz {
if let Some(packed_data) = self.try_rvz_pack(group_data.as_ref(), &info) {
meta.data_size =
align_up_32(hash_exception_data.len() as u32, 4) + packed_data.len() as u32;
(hash_exception_data.len() as u32).align_up(4) + packed_data.len() as u32;
meta.rvz_packed_size = packed_data.len() as u32;
group_data = packed_data;
}
@@ -1185,9 +1203,9 @@ impl BlockProcessor for BlockProcessorWIA {
let compressed_size = self.compressor.buffer.len() as u32;
// For WIA, we must always store compressed data.
// For RVZ, only store compressed data if it's smaller than uncompressed.
if !is_rvz || align_up_32(compressed_size, 4) < meta.data_size {
if !is_rvz || compressed_size.align_up(4) < meta.data_size {
// Align resulting block data to 4 bytes
let mut buf = BytesMut::zeroed(align_up_32(compressed_size, 4) as usize);
let mut buf = BytesMut::zeroed(compressed_size.align_up(4) as usize);
buf[..compressed_size as usize].copy_from_slice(&self.compressor.buffer);
meta.is_compressed = true;
// Data size does not include end alignment
@@ -1214,9 +1232,9 @@ impl BlockProcessor for BlockProcessorWIA {
}
// Store uncompressed group, aligned to 4 bytes after hash exceptions and at the end
let mut buf = BytesMut::zeroed(align_up_32(meta.data_size, 4) as usize);
let mut buf = BytesMut::zeroed(meta.data_size.align_up(4) as usize);
buf[..hash_exception_data.len()].copy_from_slice(hash_exception_data.as_ref());
let offset = align_up_32(hash_exception_data.len() as u32, 4) as usize;
let offset = (hash_exception_data.len() as u32).align_up(4) as usize;
buf[offset..offset + group_data.len()].copy_from_slice(group_data.as_ref());
meta.data_hash = xxh64_hash(buf.as_ref());
Ok(BlockResult { block_idx: group_idx, disc_data, block_data: buf.freeze(), meta })
@@ -1336,7 +1354,7 @@ impl BlockProcessorWIA {
sector,
);
}
*array_ref_mut![out, offset, 4] = (len as u32 | 0x80000000).to_be_bytes();
*array_ref_mut![out, offset, 4] = (len as u32 | COMPRESSED_BIT).to_be_bytes();
array_ref_mut![out, offset + 4, SEED_SIZE_BYTES].copy_from_slice(seed_out.as_bytes());
}
@@ -1528,13 +1546,11 @@ impl DiscWriterWIA {
let partition_data_start = partition.data_start_sector as u64 * SECTOR_SIZE as u64;
let partition_end = partition.data_end_sector as u64 * SECTOR_SIZE as u64;
let partition_header = partition.partition_header();
let management_data_end = align_up_64(
partition_offset_to_raw(
partition_header.fst_offset(true) + partition_header.fst_size(true),
),
0x200000, // Align to 2 MiB
);
let boot_header = partition.boot_header();
let management_data_end =
partition_offset_to_raw(boot_header.fst_offset(true) + boot_header.fst_size(true))
// Align to 2 MiB
.align_up(0x200000);
let management_end_sector = ((partition_data_start + management_data_end)
.min(partition_end)
/ SECTOR_SIZE as u64) as u32;
@@ -1629,7 +1645,7 @@ impl DiscWriterWIA {
partition.data_start_sector,
partition.data_end_sector,
partition.disc_header(),
Some(partition.partition_header()),
Some(partition.boot_header()),
partition.fst(),
));
}
@@ -1637,7 +1653,7 @@ impl DiscWriterWIA {
0,
disc_size.div_ceil(SECTOR_SIZE as u64) as u32,
disc_header,
inner.partition_header(),
inner.boot_header(),
inner.fst(),
));
@@ -1713,7 +1729,7 @@ impl DiscWriter for DiscWriterWIA {
groups[group_idx as usize] = RVZGroup {
data_offset: data_offset.into(),
data_size_and_flag: (group.meta.data_size
| if group.meta.is_compressed { 0x80000000 } else { 0 })
| if group.meta.is_compressed { COMPRESSED_BIT } else { 0 })
.into(),
rvz_packed_size: group.meta.rvz_packed_size.into(),
};

View File

@@ -160,7 +160,7 @@ pub enum Error {
#[error("disc format error: {0}")]
DiscFormat(String),
/// A general I/O error.
#[error("I/O error: {0}")]
#[error("{0}")]
Io(String, #[source] std::io::Error),
/// An unknown error.
#[error("error: {0}")]
@@ -225,3 +225,34 @@ where E: ErrorContext
self.map_err(|e| e.context(f()))
}
}
pub(crate) trait IoErrorContext {
fn io_context(self, context: impl Into<String>) -> std::io::Error;
}
impl IoErrorContext for std::io::Error {
#[inline]
fn io_context(self, context: impl Into<String>) -> std::io::Error {
std::io::Error::new(self.kind(), self.context(context))
}
}
pub(crate) trait IoResultContext<T> {
fn io_context(self, context: impl Into<String>) -> std::io::Result<T>;
fn io_with_context<F>(self, f: F) -> std::io::Result<T>
where F: FnOnce() -> String;
}
impl<T> IoResultContext<T> for std::io::Result<T> {
#[inline]
fn io_context(self, context: impl Into<String>) -> std::io::Result<T> {
self.map_err(|e| e.io_context(context))
}
#[inline]
fn io_with_context<F>(self, f: F) -> std::io::Result<T>
where F: FnOnce() -> String {
self.map_err(|e| e.io_context(f()))
}
}

View File

@@ -14,10 +14,11 @@ use crate::{
disc::{
fst::{Fst, Node},
wii::{ContentMetadata, Ticket, TmdHeader, H3_TABLE_SIZE, REGION_SIZE},
ApploaderHeader, DiscHeader, DolHeader, PartitionHeader, BI2_SIZE, BOOT_SIZE,
ApploaderHeader, BootHeader, DebugHeader, DiscHeader, DolHeader, BB2_OFFSET, BI2_SIZE,
BOOT_SIZE,
},
io::block,
util::WindowedReader,
util::{array_ref, WindowedReader},
Result,
};
@@ -198,7 +199,7 @@ pub struct DiscMeta {
}
/// An open disc partition.
pub trait PartitionReader: DynClone + BufRead + Seek + Send + Sync {
pub trait PartitionReader: BufRead + DiscStream {
/// Whether this is a Wii partition. (GameCube otherwise)
fn is_wii(&self) -> bool;
@@ -306,7 +307,7 @@ dyn_clone::clone_trait_object!(PartitionReader);
/// Extra disc partition data. (DOL, FST, etc.)
#[derive(Clone, Debug)]
pub struct PartitionMeta {
/// Disc and partition header (boot.bin)
/// Disc and boot header (boot.bin)
pub raw_boot: Arc<[u8; BOOT_SIZE]>,
/// Debug and region information (bi2.bin)
pub raw_bi2: Arc<[u8; BI2_SIZE]>,
@@ -329,16 +330,27 @@ pub struct PartitionMeta {
impl PartitionMeta {
/// A view into the disc header.
#[inline]
pub fn header(&self) -> &DiscHeader {
DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()])
.expect("Invalid header alignment")
pub fn disc_header(&self) -> &DiscHeader {
DiscHeader::ref_from_bytes(array_ref![self.raw_boot, 0, size_of::<DiscHeader>()])
.expect("Invalid disc header alignment")
}
/// A view into the partition header.
/// A view into the debug header.
#[inline]
pub fn partition_header(&self) -> &PartitionHeader {
PartitionHeader::ref_from_bytes(&self.raw_boot[size_of::<DiscHeader>()..])
.expect("Invalid partition header alignment")
pub fn debug_header(&self) -> &DebugHeader {
DebugHeader::ref_from_bytes(array_ref![
self.raw_boot,
size_of::<DiscHeader>(),
size_of::<DebugHeader>()
])
.expect("Invalid debug header alignment")
}
/// A view into the boot header.
#[inline]
pub fn boot_header(&self) -> &BootHeader {
BootHeader::ref_from_bytes(array_ref![self.raw_boot, BB2_OFFSET, size_of::<BootHeader>()])
.expect("Invalid boot header alignment")
}
/// A view into the apploader header.

View File

@@ -10,6 +10,7 @@ use crate::{
pub struct Decompressor {
pub kind: DecompressionKind,
#[allow(unused)] // if compression features are disabled
pub cache: DecompressorCache,
}
@@ -38,7 +39,13 @@ impl Decompressor {
pub fn decompress(&mut self, buf: &[u8], out: &mut [u8]) -> io::Result<usize> {
match &self.kind {
DecompressionKind::None => {
out.copy_from_slice(buf);
if buf.len() > out.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("Decompressed data too large: {} > {}", buf.len(), out.len()),
));
}
out[..buf.len()].copy_from_slice(buf);
Ok(buf.len())
}
#[cfg(feature = "compress-zlib")]
@@ -130,6 +137,18 @@ impl Decompressor {
}
}
}
pub fn get_content_size(&self, buf: &[u8]) -> io::Result<Option<usize>> {
match &self.kind {
DecompressionKind::None => Ok(Some(buf.len())),
#[cfg(feature = "compress-zstd")]
DecompressionKind::Zstandard => zstd_safe::get_frame_content_size(buf)
.map(|n| n.map(|n| n as usize))
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string())),
#[allow(unreachable_patterns)] // if compression features are disabled
_ => Ok(None),
}
}
}
#[derive(Debug, Clone)]
@@ -227,6 +246,14 @@ impl Compressor {
pub fn compress(&mut self, buf: &[u8]) -> io::Result<bool> {
self.buffer.clear();
match self.kind {
Compression::None => {
if self.buffer.capacity() >= buf.len() {
self.buffer.extend_from_slice(buf);
Ok(true)
} else {
Ok(false)
}
}
#[cfg(feature = "compress-zlib")]
Compression::Deflate(level) => {
let compressor = match &mut self.cache {
@@ -288,6 +315,8 @@ impl Compressor {
_ => {
let mut ctx = zstd_safe::CCtx::create();
ctx.init(level as i32).map_err(zstd_util::map_error_code)?;
ctx.set_parameter(zstd_safe::CParameter::ContentSizeFlag(true))
.map_err(zstd_util::map_error_code)?;
self.cache = CompressorCache::Zstandard(ctx);
match &mut self.cache {
CompressorCache::Zstandard(compressor) => compressor,
@@ -302,6 +331,7 @@ impl Compressor {
Err(e) => Err(zstd_util::map_error_code(e)),
}
}
#[allow(unreachable_patterns)] // if compression is disabled
_ => Err(io::Error::new(
io::ErrorKind::Other,
format!("Unsupported compression: {:?}", self.kind),

View File

@@ -127,18 +127,36 @@ where T: Div<Output = T> + Rem<Output = T> + Copy {
(x / y, x % y)
}
#[inline]
pub(crate) fn align_up_32(n: u32, align: u32) -> u32 { (n + align - 1) & !(align - 1) }
pub(crate) trait Align {
fn align_up(self, align: Self) -> Self;
#[inline]
pub(crate) fn align_up_64(n: u64, align: u64) -> u64 { (n + align - 1) & !(align - 1) }
fn align_down(self, align: Self) -> Self;
}
macro_rules! impl_align {
($ty:ident) => {
impl Align for $ty {
#[inline(always)]
fn align_up(self, align: Self) -> Self { (self + (align - 1)) & !(align - 1) }
#[inline(always)]
fn align_down(self, align: Self) -> Self { self & !(align - 1) }
}
};
}
impl_align!(u8);
impl_align!(u16);
impl_align!(u32);
impl_align!(u64);
impl_align!(usize);
/// Creates a fixed-size array reference from a slice.
macro_rules! array_ref {
($slice:expr, $offset:expr, $size:expr) => {{
#[inline(always)]
fn to_array<T>(slice: &[T]) -> &[T; $size] {
unsafe { &*(slice.as_ptr() as *const [_; $size]) }
unsafe { &*(slice as *const [T] as *const [T; $size]) }
}
to_array(&$slice[$offset..$offset + $size])
}};
@@ -150,7 +168,7 @@ macro_rules! array_ref_mut {
($slice:expr, $offset:expr, $size:expr) => {{
#[inline(always)]
fn to_array<T>(slice: &mut [T]) -> &mut [T; $size] {
unsafe { &mut *(slice.as_ptr() as *mut [_; $size]) }
unsafe { &mut *(slice as *mut [T] as *mut [T; $size]) }
}
to_array(&mut $slice[$offset..$offset + $size])
}};