diff --git a/geckolib/src/iso/builder/mod.rs b/geckolib/src/iso/builder/mod.rs index e038087..e8b2ba2 100644 --- a/geckolib/src/iso/builder/mod.rs +++ b/geckolib/src/iso/builder/mod.rs @@ -27,7 +27,7 @@ use crate::vfs::{self, Directory, GeckoFS}; use crate::UPDATER; use crate::{framework_map, linker, warn}; -use super::disc::{DiscType, WiiDisc}; +use super::{disc::DiscType, read::DiscReader}; mod fs_source; @@ -42,7 +42,7 @@ pub struct IsoBuilder { config: Config, fs: FSSource, gfs: GeckoFS, - disc_info: Option, + reader: DiscReader, writer: W, } @@ -51,10 +51,10 @@ impl IsoBuilder { config: Config, zip: ZipArchive, gfs: GeckoFS, - disc_info: Option, + reader: DiscReader, writer: W, ) -> Self { - Self::internal_new(config, FSSource::Zip(Box::new(zip)), gfs, disc_info, writer) + Self::internal_new(config, FSSource::Zip(Box::new(zip)), gfs, reader, writer) } #[cfg(not(target_os = "unknown"))] @@ -62,14 +62,14 @@ impl IsoBuilder { config: Config, path: P, gfs: GeckoFS, - disc_info: Option, + reader: DiscReader, writer: W, ) -> Self { Self::internal_new( config, FSSource::FS(path.as_ref().to_path_buf()), gfs, - disc_info, + reader, writer, ) } @@ -78,14 +78,14 @@ impl IsoBuilder { config: Config, fs: FSSource, gfs: GeckoFS, - disc_info: Option, + reader: DiscReader, writer: W, ) -> Self { Self { config, fs, gfs, - disc_info, + reader, writer, } } @@ -185,7 +185,6 @@ where updater.set_message("Loading game...".into())?; } - let wii_disc_info = self.disc_info.clone(); let disc = &mut self.gfs; #[cfg(feature = "progress")] @@ -318,7 +317,7 @@ where )?; } - if wii_disc_info.is_none() { + if self.reader.get_type() == DiscType::Gamecube { #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { updater.set_message("".into())?; @@ -372,17 +371,16 @@ where // Finalize disc and write it back into a file - let out: DiscWriter = { DiscWriter::new(self.writer.clone(), wii_disc_info) }; - if let DiscWriter::Wii(wii_out) = out.clone() { - std::pin::pin!(wii_out).init().await?; - } + let out: DiscWriter = DiscWriter::from_reader(self.writer.clone(), &self.reader); + std::pin::pin!(out.clone()).init().await?; + // let out = DiscWriter::Gamecube(self.writer.clone()); let mut out = std::pin::pin!(out); - let is_wii = out.get_type() == DiscType::Wii; - disc.serialize(&mut out, is_wii).await?; + disc.serialize(&mut out).await?; #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { + updater.set_title("Finished".into())?; updater.finish()?; } @@ -482,8 +480,8 @@ impl Builder for PatchBuilder { #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { - updater.set_title("Storing replacement files...".into())?; updater.set_message("".into())?; + updater.set_title("Storing replacement files...".into())?; } let mut new_map = HashMap::new(); @@ -497,8 +495,8 @@ impl Builder for PatchBuilder { #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { - updater.set_title("Storing libraries...".into())?; updater.set_message("".into())?; + updater.set_title("Storing libraries...".into())?; } if let Some(link) = &mut config.link { @@ -543,8 +541,8 @@ impl Builder for PatchBuilder { #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { - updater.set_title("Storing banner...".into())?; updater.set_message("".into())?; + updater.set_title("Storing banner...".into())?; } write_file_to_zip(&mut zip, "banner.dat", &read(path).await?)?; @@ -554,8 +552,8 @@ impl Builder for PatchBuilder { #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { - updater.set_title("Storing patch index...".into())?; updater.set_message("".into())?; + updater.set_title("Storing patch index...".into())?; } config.src.iso = PathBuf::new(); @@ -568,6 +566,12 @@ impl Builder for PatchBuilder { zip.finish()?; + #[cfg(feature = "progress")] + if let Ok(mut updater) = UPDATER.lock() { + updater.set_title("Finished".into())?; + updater.finish()?; + } + Ok(()) } } diff --git a/geckolib/src/iso/disc.rs b/geckolib/src/iso/disc.rs index a34c372..fefd666 100644 --- a/geckolib/src/iso/disc.rs +++ b/geckolib/src/iso/disc.rs @@ -1,9 +1,10 @@ +use std::error::Error; use std::io::SeekFrom; +use std::mem::offset_of; use std::pin::Pin; use async_std::io::prelude::SeekExt; use async_std::io::ReadExt; -use eyre::Result; use num::Unsigned; use crate::crypto::aes_decrypt_inplace; @@ -17,6 +18,37 @@ use byteorder::{ByteOrder, BE}; use sha1_smol::Sha1; use std::convert::TryFrom; +#[derive(Debug)] +pub enum DiscError { + NoGamePartition, + EncryptionError(WiiCryptoError), + Io(std::io::Error), +} + +impl Error for DiscError {} + +impl std::fmt::Display for DiscError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DiscError::NoGamePartition => write!(f, "There is no game parition in this disc"), + DiscError::EncryptionError(e) => write!(f, "Encryption error: {}", e), + DiscError::Io(e) => write!(f, "IO error: {}", e), + } + } +} + +impl From for DiscError { + fn from(e: WiiCryptoError) -> Self { + DiscError::EncryptionError(e) + } +} + +impl From for DiscError { + fn from(e: std::io::Error) -> Self { + DiscError::Io(e) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum DiscType { Gamecube = 0, @@ -37,8 +69,8 @@ pub struct WiiDiscHeader { pub wii_magic: u32, pub gc_magic: u32, pub game_title: [u8; 64], - pub disable_hash_verif: u8, - pub disable_disc_encrypt: u8, + pub disable_hash_verif: bool, + pub disable_disc_encrypt: bool, pub padding: [u8; 0x39e], } @@ -88,8 +120,8 @@ pub fn disc_get_header(raw: &[u8]) -> WiiDiscHeader { wii_magic: BE::read_u32(&raw[0x18..0x1C]), gc_magic: BE::read_u32(&raw[0x1C..0x20]), game_title, - disable_hash_verif: raw[0x60], - disable_disc_encrypt: raw[0x61], + disable_hash_verif: raw[0x60] != 0, + disable_disc_encrypt: raw[0x61] != 0, padding, } } @@ -109,8 +141,8 @@ pub fn disc_set_header(buffer: &mut [u8], dh: &WiiDiscHeader) { BE::write_u32(&mut buffer[0x18..], dh.wii_magic); BE::write_u32(&mut buffer[0x1C..], dh.gc_magic); buffer[0x20..0x60].copy_from_slice(&dh.game_title[..]); - buffer[0x60] = dh.disable_hash_verif; - buffer[0x61] = dh.disable_disc_encrypt; + buffer[0x60] = dh.disable_hash_verif as u8; + buffer[0x61] = dh.disable_disc_encrypt as u8; buffer[0x62..0x400].copy_from_slice(&dh.padding); } @@ -128,9 +160,30 @@ pub struct WiiDiscRegionAgeRating { kr: u8, } +#[derive(Debug, Clone, Copy, Default)] +pub enum WiiDiscRegions { + #[default] + NTSCJ, + NTSCU, + PAL, + KOR, +} + +impl From for WiiDiscRegions { + fn from(r: u32) -> Self { + match r { + 0 => WiiDiscRegions::NTSCJ, + 1 => WiiDiscRegions::NTSCU, + 2 => WiiDiscRegions::PAL, + 3 => WiiDiscRegions::KOR, + _ => WiiDiscRegions::default(), + } + } +} + #[derive(Debug, Clone, Copy, Default)] pub struct WiiDiscRegion { - pub region: u32, + pub region: WiiDiscRegions, pub age_rating: WiiDiscRegionAgeRating, } @@ -141,7 +194,7 @@ impl Unpackable for WiiDiscRegion { impl WiiDiscRegion { pub fn parse(raw: &[u8]) -> Self { Self { - region: BE::read_u32(&raw[..4]), + region: BE::read_u32(&raw[..4]).into(), age_rating: WiiDiscRegionAgeRating { jp: raw[0x10], us: raw[0x11], @@ -158,7 +211,7 @@ impl WiiDiscRegion { } pub fn compose_into(&self, buf: &mut [u8]) { - BE::write_u32(&mut buf[..4], self.region); + buf[..4].copy_from_slice((self.region as u32).to_be_bytes().as_ref()); buf[0x10] = self.age_rating.jp; buf[0x11] = self.age_rating.us; buf[0x12] = self.age_rating.unknown1; @@ -186,7 +239,7 @@ pub struct PartInfo { pub async fn disc_get_part_info_async( reader: &mut Pin<&mut R>, -) -> Result { +) -> Result { crate::debug!("Parsing partition info (async)"); let mut entries: Vec = Vec::new(); let mut buf: [u8; 8] = [0u8; 8]; @@ -294,7 +347,7 @@ pub struct Ticket { pub time_limit: u32, pub fake_sign: [u8; 0x58], } -assert_eq_size!(Ticket, [u8; Ticket::BLOCK_SIZE]); +assert_eq_size!(Ticket, [u8; ::BLOCK_SIZE]); declare_tryfrom!(Ticket); impl Unpackable for Ticket { @@ -303,18 +356,18 @@ impl Unpackable for Ticket { impl Ticket { pub fn fake_sign(&mut self) -> Result<(), eyre::Report> { - let mut tik_buf = [0u8; Ticket::BLOCK_SIZE]; + let mut tik_buf = [0u8; ::BLOCK_SIZE]; self.sig.sig.fill(0); self.sig.sig_padding.fill(0); self.fake_sign.fill(0); - tik_buf[..Ticket::BLOCK_SIZE].copy_from_slice(&<[u8; Ticket::BLOCK_SIZE]>::from(&*self)); + tik_buf[..::BLOCK_SIZE].copy_from_slice(&<[u8; ::BLOCK_SIZE]>::from(&*self)); // start brute force crate::trace!("Ticket fake signing; starting brute force..."); let mut val = 0u32; let mut hash_0; let mut sha1 = Sha1::new(); loop { - BE::write_u32(&mut tik_buf[0x248..][..4], val); + BE::write_u32(&mut tik_buf[offset_of!(Self, time_limit)/* 0x248 */..][..size_of_val(&self.time_limit)], val); sha1.reset(); sha1.update(&tik_buf[0x140..]); hash_0 = sha1.digest().bytes()[0]; @@ -338,8 +391,8 @@ impl Ticket { } } -impl From<&[u8; Ticket::BLOCK_SIZE]> for Ticket { - fn from(buf: &[u8; Ticket::BLOCK_SIZE]) -> Self { +impl From<&[u8; ::BLOCK_SIZE]> for Ticket { + fn from(buf: &[u8; ::BLOCK_SIZE]) -> Self { let mut sig = [0_u8; 0x100]; let mut sig_padding = [0_u8; 0x3C]; let mut sig_issuer = [0_u8; 0x40]; @@ -389,9 +442,9 @@ impl From<&[u8; Ticket::BLOCK_SIZE]> for Ticket { } } -impl From<&Ticket> for [u8; Ticket::BLOCK_SIZE] { +impl From<&Ticket> for [u8; ::BLOCK_SIZE] { fn from(t: &Ticket) -> Self { - let mut buf = [0_u8; Ticket::BLOCK_SIZE]; + let mut buf = [0_u8; ::BLOCK_SIZE]; BE::write_u32(&mut buf[0x00..], t.sig.sig_type); buf[0x04..0x104].copy_from_slice(&t.sig.sig[..]); @@ -421,11 +474,12 @@ impl From<&Ticket> for [u8; Ticket::BLOCK_SIZE] { impl Default for Ticket { fn default() -> Self { - Ticket::from(&[0_u8; Ticket::BLOCK_SIZE]) + Ticket::from(&[0_u8; ::BLOCK_SIZE]) } } #[derive(Debug, Clone, Copy, Default)] +#[repr(C)] pub struct PartHeader { pub ticket: Ticket, pub tmd_size: usize, @@ -460,7 +514,7 @@ impl From<&[u8; PartHeader::BLOCK_SIZE]> for PartHeader { impl From<&PartHeader> for [u8; PartHeader::BLOCK_SIZE] { fn from(ph: &PartHeader) -> Self { let mut buf = [0_u8; PartHeader::BLOCK_SIZE]; - buf[..0x2A4].copy_from_slice(&<[u8; Ticket::BLOCK_SIZE]>::from(&ph.ticket)); + buf[..0x2A4].copy_from_slice(&<[u8; ::BLOCK_SIZE]>::from(&ph.ticket)); BE::write_u32(&mut buf[0x2A4..], ph.tmd_size as u32); BE::write_u32(&mut buf[0x2A8..], (ph.tmd_offset >> 2) as u32); BE::write_u32(&mut buf[0x2AC..], ph.cert_size as u32); diff --git a/geckolib/src/iso/read.rs b/geckolib/src/iso/read.rs index 33b0f1f..04080c8 100644 --- a/geckolib/src/iso/read.rs +++ b/geckolib/src/iso/read.rs @@ -6,14 +6,71 @@ use async_std::io::{Read as AsyncRead, ReadExt, Seek as AsyncSeek}; use async_std::sync::Mutex; use async_std::task::ready; use byteorder::{ByteOrder, BE}; -use eyre::Result; #[cfg(feature = "parallel")] use rayon::prelude::*; +use std::error::Error; +use std::fmt::Display; use std::io::SeekFrom; use std::pin::{pin, Pin}; use std::sync::Arc; use std::task::Poll; +// Error structures for the disc readers + +#[derive(Debug)] +pub enum WiiDiscReaderError { + InvalidWiiDisc { magic: u32 }, + NoGamePartition, + EncryptionError(WiiCryptoError), + ConvertError { name: String }, + Io(std::io::Error), +} + +impl Error for WiiDiscReaderError {} + +impl Display for WiiDiscReaderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + WiiDiscReaderError::InvalidWiiDisc { magic } => { + write!(f, "Invalid Wii disc: magic is {:#010X}", magic) + } + WiiDiscReaderError::NoGamePartition => write!(f, "There is no game parition in this disc"), + WiiDiscReaderError::EncryptionError(e) => write!(f, "Encryption error: {}", e), + WiiDiscReaderError::ConvertError { name } => { + write!(f, "The provided slice is too small to be converted into a {}.", name) + } + WiiDiscReaderError::Io(e) => write!(f, "I/O error: {}", e), + } + } +} + +impl From for WiiDiscReaderError { + fn from(e: std::io::Error) -> Self { + Self::Io(e) + } +} + +impl From for WiiDiscReaderError { + fn from(e: WiiCryptoError) -> Self { + match e { + WiiCryptoError::NotWiiDisc { magic } => Self::InvalidWiiDisc { magic }, + WiiCryptoError::NoGamePartition => Self::NoGamePartition, + WiiCryptoError::ConvertError { name } => Self::ConvertError { name }, + _ => Self::EncryptionError(e), + } + } +} + +impl From for WiiDiscReaderError { + fn from(e: DiscError) -> Self { + match e { + DiscError::NoGamePartition => Self::NoGamePartition, + DiscError::EncryptionError(e) => Self::EncryptionError(e), + DiscError::Io(e) => Self::Io(e), + } + } +} + #[derive(Debug, Clone)] enum WiiDiscReaderState { Seeking, @@ -37,7 +94,7 @@ pub struct WiiDiscReader { async fn get_partitions( reader: &mut Pin<&mut R>, part_info: &PartInfo, -) -> Result { +) -> Result { crate::debug!("Fetching partitions from reader"); let mut ret_vec: Vec = Vec::new(); let mut data_idx: Option = None; @@ -103,7 +160,7 @@ impl WiiDiscReader where R: AsyncRead + AsyncSeek + Unpin, { - pub async fn try_parse(reader: R) -> Result { + pub async fn try_parse(reader: R) -> Result { crate::debug!("Trying to parse a Wii Disc from the reader"); let mut this = Self { reader, @@ -273,6 +330,7 @@ where let mut data_pool: Vec<&mut [u8]> = buf2.chunks_exact_mut(consts::WII_SECTOR_SIZE).collect(); crate::trace!("data_pool size: {}", data_pool.len()); + let disable_disc_encrypt = this.disc.disc_header.disable_disc_encrypt; let decrypt_process = move |data: &mut &mut [u8]| { let mut iv = [0_u8; consts::WII_KEY_SIZE]; iv[..consts::WII_KEY_SIZE].copy_from_slice( @@ -280,17 +338,19 @@ where ); crate::trace!("iv: {:?}", iv); crate::trace!("before: {:?}", &data[consts::WII_SECTOR_HASH_SIZE..][..6]); - // Decrypt the hash to check if valid (not required here) - aes_decrypt_inplace( - &mut data[..consts::WII_SECTOR_HASH_SIZE], - &[0_u8; consts::WII_KEY_SIZE], - &part_key, - ); - aes_decrypt_inplace( - &mut data[consts::WII_SECTOR_HASH_SIZE..][..consts::WII_SECTOR_DATA_SIZE], - &iv, - &part_key, - ); + if !disable_disc_encrypt { + // Decrypt the hash to check if valid (not required here) + aes_decrypt_inplace( + &mut data[..consts::WII_SECTOR_HASH_SIZE], + &[0_u8; consts::WII_KEY_SIZE], + &part_key, + ); + aes_decrypt_inplace( + &mut data[consts::WII_SECTOR_HASH_SIZE..][..consts::WII_SECTOR_DATA_SIZE], + &iv, + &part_key, + ); + } crate::trace!("after: {:?}", &data[consts::WII_SECTOR_HASH_SIZE..][..6]); }; crate::trace!("Decrypting blocks"); @@ -327,6 +387,37 @@ where } } +#[derive(Debug)] +pub enum DiscReaderError { + NotDisc(u32, u32), + Io(std::io::Error), + Wii(WiiDiscReaderError), +} + +impl From for DiscReaderError { + fn from(e: std::io::Error) -> Self { + Self::Io(e) + } +} + +impl From for DiscReaderError { + fn from(e: WiiDiscReaderError) -> Self { + Self::Wii(e) + } +} + +impl Error for DiscReaderError {} + +impl Display for DiscReaderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DiscReaderError::NotDisc(gc_magic, wii_magic) => write!(f, "Not a disc (GC: {:#010X}; Wii: {:#010X})", gc_magic, wii_magic), + DiscReaderError::Io(e) => write!(f, "I/O error: {}", e), + DiscReaderError::Wii(e) => write!(f, "Wii error: {}", e), + } + } +} + #[derive(Debug)] pub enum DiscReader { Gamecube(R), @@ -381,7 +472,7 @@ impl DiscReader where R: AsyncRead + AsyncSeek + Unpin, { - pub async fn new(mut reader: R) -> Result { + pub async fn new(mut reader: R) -> Result { pin!(&mut reader).seek(SeekFrom::Start(0x18)).await?; let mut buf = [0u8; 8]; pin!(&mut reader).read(&mut buf).await?; @@ -393,10 +484,9 @@ where crate::debug!("Loading Wii disc"); Ok(Self::Wii(WiiDiscReader::try_parse(reader).await?)) } else { - Err(eyre::eyre!( - "Not a game disc (Wii: {:08X}; GC: {:08X})", + Err(DiscReaderError::NotDisc( + BE::read_u32(&buf[4..][..4]), BE::read_u32(&buf[..][..4]), - BE::read_u32(&buf[4..][..4]) )) } } diff --git a/geckolib/src/iso/write.rs b/geckolib/src/iso/write.rs index 76a0866..c8c6607 100644 --- a/geckolib/src/iso/write.rs +++ b/geckolib/src/iso/write.rs @@ -21,9 +21,9 @@ use crate::{ }, }; -use super::disc::{ +use super::{disc::{ decrypt_title_key, DiscType, WiiDisc, WiiGroup, WiiPartition, WiiSector, WiiSectorHash, -}; +}, read::DiscReader}; #[derive(Debug, Clone, Default)] enum WiiDiscWriterState { @@ -49,6 +49,7 @@ enum WiiDiscWriterState { #[derive(Debug)] struct WiiDiscWriterStatus { initialized: bool, + pub disc: WiiDisc, hashes: Vec<[u8; consts::WII_HASH_SIZE]>, // Virtual cursor which tracks where in the decrypted partition we are writing from. cursor: u64, @@ -60,7 +61,6 @@ struct WiiDiscWriterStatus { pub struct WiiDiscWriter { writer: W, status: Arc>, - pub disc: WiiDisc, } impl Clone for WiiDiscWriter @@ -71,7 +71,6 @@ where Self { writer: self.writer.clone(), status: self.status.clone(), - disc: self.disc.clone(), } } } @@ -176,8 +175,12 @@ fn fake_sign(part: &mut WiiPartition, hashes: &[[u8; consts::WII_HASH_SIZE]]) { .copy_from_slice(&Sha1::from(&hashes_).digest().bytes()); // Fake sign tmd - let _ = part.tmd.fake_sign(); - let _ = part.header.ticket.fake_sign(); + if let Err(err) = part.tmd.fake_sign() { + crate::warn!("Error while signing TMD: {}", err); + } + if let Err(err) = part.header.ticket.fake_sign() { + crate::warn!("Error while signing Ticket: {}", err); + } } fn encrypt_group(group: &mut WiiGroup, part_key: AesKey) { @@ -212,12 +215,12 @@ where writer, status: Arc::new(Mutex::new(WiiDiscWriterStatus { initialized: false, + disc, cursor: 0, state: WiiDiscWriterState::default(), group: Box::new(WiiGroup::default()), hashes: Vec::new(), })), - disc, } } @@ -232,7 +235,8 @@ where // Write ISO header let mut buf = vec![0u8; WiiDiscHeader::BLOCK_SIZE]; - disc_set_header(&mut buf, &this.disc.disc_header); + let disc = &mut state.disc; + disc_set_header(&mut buf, &disc.disc_header); this.writer.seek(SeekFrom::Start(0)).await?; this.writer.write_all(&buf).await?; @@ -242,15 +246,15 @@ where .await?; // Write Partition Info - let part_idx = this.disc.partitions.data_idx; + let part_idx = disc.partitions.data_idx; let mut buf = [0u8; 0x28]; BE::write_u32(&mut buf[..], 1); BE::write_u32(&mut buf[4..], 0x40020 >> 2); let offset: u64 = 0x50000; let i = 0; - let part_type: u32 = this.disc.partitions.partitions[part_idx].part_type.into(); + let part_type: u32 = disc.partitions.partitions[part_idx].part_type.into(); crate::debug!("part_type: {}", part_type); - this.disc.partitions.partitions[part_idx].part_offset = offset; + disc.partitions.partitions[part_idx].part_offset = offset; BE::write_u32(&mut buf[0x20 + (8 * i)..], (offset >> 2) as u32); BE::write_u32(&mut buf[0x20 + (8 * i) + 4..], part_type); this.writer.write_all(&buf).await?; @@ -262,7 +266,7 @@ where // Write Region area let mut buf = [0u8; 0x20]; - this.disc.disc_region.compose_into(&mut buf); + disc.disc_region.compose_into(&mut buf); this.writer.write_all(&buf).await?; // Get to Magic @@ -277,13 +281,13 @@ where this.writer.write_all(&buf).await?; // Make sure there is at least one content in the TitleMetaData - if this.disc.partitions.partitions[part_idx] + if disc.partitions.partitions[part_idx] .tmd .contents .is_empty() { crate::warn!("TMD has no content value. Generating new value"); - this.disc.partitions.partitions[part_idx] + disc.partitions.partitions[part_idx] .tmd .contents .push(TMDContent { @@ -297,7 +301,7 @@ where // Write (partial) partition header { - let part = &mut this.disc.partitions.partitions[part_idx]; + let part = &mut disc.partitions.partitions[part_idx]; part.header.data_size = 0; part.header.tmd_size = part.tmd.get_size(); part.header.tmd_offset = PartHeader::BLOCK_SIZE as u64; @@ -310,26 +314,26 @@ where align_addr(part.header.cert_offset + part.header.cert_size as u64, 17); } let buf = - <[u8; PartHeader::BLOCK_SIZE]>::from(&this.disc.partitions.partitions[part_idx].header); + <[u8; PartHeader::BLOCK_SIZE]>::from(&disc.partitions.partitions[part_idx].header); this.writer.write_all(&buf).await?; - let mut buf = vec![0u8; this.disc.partitions.partitions[part_idx].header.tmd_size]; - TitleMetaData::set_partition(&mut buf, 0, &this.disc.partitions.partitions[part_idx].tmd); + let mut buf = vec![0u8; disc.partitions.partitions[part_idx].header.tmd_size]; + TitleMetaData::set_partition(&mut buf, 0, &disc.partitions.partitions[part_idx].tmd); this.writer.write_all(&buf).await?; // Write certificate - let cert = this.disc.partitions.partitions[part_idx].cert.clone(); + let cert = disc.partitions.partitions[part_idx].cert.clone(); this.writer.write_all(&cert).await?; let padding_size = std::cmp::max( 0, - this.disc.partitions.partitions[part_idx].header.data_offset as i64 - - (this.disc.partitions.partitions[part_idx].header.h3_offset as i64), + disc.partitions.partitions[part_idx].header.data_offset as i64 + - (disc.partitions.partitions[part_idx].header.h3_offset as i64), ); if padding_size > 0 { let buf = vec![0u8; padding_size as usize]; this.writer.write_all(&buf).await?; } let pos = this.writer.seek(SeekFrom::Current(0)).await?; - let data_offset = this.disc.partitions.partitions[part_idx].part_offset - + this.disc.partitions.partitions[part_idx].header.data_offset; + let data_offset = disc.partitions.partitions[part_idx].part_offset + + disc.partitions.partitions[part_idx].header.data_offset; if data_offset > pos { this.writer .write_all(&vec![0u8; (data_offset - pos) as usize]) @@ -371,15 +375,15 @@ where buf: &[u8], ) -> Poll> { let this = self.get_mut(); - let mut state_guard = match this.status.try_lock_arc() { + let mut status = match this.status.try_lock_arc() { Some(state) => state, None => { cx.waker().wake_by_ref(); return Poll::Pending; } }; - let part_idx = this.disc.partitions.data_idx; - if let WiiDiscWriterState::Init = state_guard.state { + let part_idx = status.disc.partitions.data_idx; + if let WiiDiscWriterState::Init = status.state { crate::trace!("Pooling WiiDiscWriter for write ({} byte(s))", buf.len()); } // If the requested size is 0, or if we are done reading, return without changing buf. @@ -387,7 +391,7 @@ where return Poll::Ready(Ok(0)); } // The "virtual" start and end, in the sense that they are the positions within the decrypted partition. - let vstart = state_guard.cursor; + let vstart = status.cursor; let vend = vstart + buf.len() as u64; let start_blk_idx = (vstart / consts::WII_SECTOR_DATA_SIZE as u64) as usize; let end_blk_idx = ((vend - 1) / consts::WII_SECTOR_DATA_SIZE as u64) as usize; @@ -395,7 +399,7 @@ where let start_block_idx_in_group = start_blk_idx % 64; let end_group_idx = end_blk_idx / 64; let end_block_idx_in_group = end_blk_idx % 64; - if let WiiDiscWriterState::Init = state_guard.state { + if let WiiDiscWriterState::Init = status.state { crate::trace!( "Writing data from 0x{:08X} to 0x{:08X} (spanning {} block(s), from {} to {})", vstart, @@ -420,11 +424,11 @@ where //Data ---------- -------------- ----- // ^ start ^ end - let state = std::mem::take(&mut state_guard.state); + let state = std::mem::take(&mut status.state); match state { WiiDiscWriterState::Init => { - state_guard.state = - WiiDiscWriterState::Parse(state_guard.cursor, start_group_idx, buf.to_vec()); + status.state = + WiiDiscWriterState::Parse(status.cursor, start_group_idx, buf.to_vec()); cx.waker().wake_by_ref(); Poll::Pending } @@ -459,45 +463,47 @@ where ); let data; (data, curr_buf) = curr_buf.split_at(size); - state_guard.group.sub_groups[(i / 8) % 8].sectors[i % 8].data + status.group.sub_groups[(i / 8) % 8].sectors[i % 8].data [buffer_start as usize..buffer_end as usize] .copy_from_slice(data); if buffer_end == consts::WII_SECTOR_DATA_SIZE as u64 * 64 { crate::trace!("Reached end of group #{}", group_idx); } } - if (state_guard.cursor + (buf.len() - curr_buf.len()) as u64) + if (status.cursor + (buf.len() - curr_buf.len()) as u64) % (consts::WII_SECTOR_DATA_SIZE as u64 * 64) == 0 { // We are at the start of a group. We can hash and encrypt the group and write it. crate::trace!("Hashing and encrypting group #{}", group_idx); - if state_guard.hashes.len() <= group_idx { - state_guard + if status.hashes.len() <= group_idx { + status .hashes .resize(group_idx + 1, [0u8; consts::WII_HASH_SIZE]); } - let group_hash = hash_group(&mut state_guard.group); - state_guard.hashes[group_idx].copy_from_slice(&group_hash); + let group_hash = hash_group(&mut status.group); + status.hashes[group_idx].copy_from_slice(&group_hash); let part_key = - decrypt_title_key(&this.disc.partitions.partitions[part_idx].header.ticket); - encrypt_group(&mut state_guard.group, part_key); + decrypt_title_key(&status.disc.partitions.partitions[part_idx].header.ticket); + if !status.disc.disc_header.disable_disc_encrypt { + encrypt_group(&mut status.group, part_key); + } - state_guard.state = WiiDiscWriterState::Writing( + status.state = WiiDiscWriterState::Writing( cursor, group_idx, - state_guard.group.to_vec(), + status.group.to_vec(), curr_buf.to_vec(), ); - state_guard.group.reset(); + status.group.reset(); cx.waker().wake_by_ref(); Poll::Pending } else { // We are in the middle of a group. We need to read the rest of the group. cursor += in_buf.len() as u64; - state_guard.cursor = cursor; - state_guard.state = WiiDiscWriterState::Init; + status.cursor = cursor; + status.state = WiiDiscWriterState::Init; Poll::Ready(Ok(buf.len())) } } @@ -506,14 +512,14 @@ where let n_written = match pin!(&mut this.writer).poll_write(cx, &group_buf) { Poll::Ready(result) => result?, Poll::Pending => { - state_guard.state = + status.state = WiiDiscWriterState::Writing(cursor, group_idx, group_buf, curr_buf); return Poll::Pending; } }; crate::trace!("Writing succeeded"); if n_written < group_buf.len() { - state_guard.state = WiiDiscWriterState::Writing( + status.state = WiiDiscWriterState::Writing( cursor, group_idx, group_buf.split_at(n_written).1.to_vec(), @@ -532,12 +538,12 @@ where if curr_buf.is_empty() { // The write is done. crate::trace!("Write done at block #{}", group_idx); - state_guard.cursor = cursor; - state_guard.state = WiiDiscWriterState::Init; + status.cursor = cursor; + status.state = WiiDiscWriterState::Init; Poll::Ready(Ok(buf.len())) } else { // We need to write the rest of the buffer - state_guard.state = + status.state = WiiDiscWriterState::Parse(cursor, group_idx + 1, curr_buf); cx.waker().wake_by_ref(); Poll::Pending @@ -574,7 +580,7 @@ where return Poll::Pending; } }; - let part_idx = this.disc.partitions.data_idx; + let part_idx = status.disc.partitions.data_idx; let state = std::mem::take(&mut status.state); #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { @@ -584,14 +590,14 @@ where WiiDiscWriterState::Init => { crate::trace!("WiiDiscWriterFinalizeState::Init"); // Align the encrypted data size to 21 bits - this.disc.partitions.partitions[part_idx].header.data_size = + status.disc.partitions.partitions[part_idx].header.data_size = align_addr(to_raw_addr(status.cursor), 21); // Hash and encrypt the last group - let n_group = this.disc.partitions.partitions[part_idx].header.data_size + let n_group = status.disc.partitions.partitions[part_idx].header.data_size / consts::WII_SECTOR_SIZE as u64 / 64; - let group_idx = (this.disc.partitions.partitions[part_idx].header.data_size - 1) + let group_idx = (status.disc.partitions.partitions[part_idx].header.data_size - 1) / consts::WII_SECTOR_SIZE as u64 / 64; crate::trace!("Hashing and encrypting group #{}", group_idx); @@ -603,15 +609,17 @@ where let group_hash = hash_group(&mut status.group); status.hashes[group_idx as usize].copy_from_slice(&group_hash); let part_key = - decrypt_title_key(&this.disc.partitions.partitions[part_idx].header.ticket); - encrypt_group(&mut status.group, part_key); + decrypt_title_key(&status.disc.partitions.partitions[part_idx].header.ticket); + if !status.disc.disc_header.disable_disc_encrypt { + encrypt_group(&mut status.group, part_key); + } status.state = if status.cursor % (consts::WII_SECTOR_DATA_SIZE as u64 * 64) != 0 { WiiDiscWriterState::SeekToLastGroup(n_group - 1, status.group.to_vec()) } else { let hashes = status.hashes.clone(); WiiDiscWriterState::SeekToPartHeader(prepare_header( - &mut this.disc.partitions.partitions[part_idx], + &mut status.disc.partitions.partitions[part_idx], &hashes, )) }; @@ -623,8 +631,8 @@ where "WiiDiscWriterFinalizeState::SeekToLastGroup(group_idx=0x{:08X})", group_idx ); - let pos = this.disc.partitions.partitions[part_idx].part_offset - + this.disc.partitions.partitions[part_idx].header.data_offset + let pos = status.disc.partitions.partitions[part_idx].part_offset + + status.disc.partitions.partitions[part_idx].header.data_offset + group_idx * consts::WII_SECTOR_SIZE as u64 * 64; if pin!(&mut this.writer) .poll_seek(cx, SeekFrom::Start(pos)) @@ -659,7 +667,7 @@ where } else { let hashes = status.hashes.clone(); WiiDiscWriterState::SeekToPartHeader(prepare_header( - &mut this.disc.partitions.partitions[part_idx], + &mut status.disc.partitions.partitions[part_idx], &hashes, )) }; @@ -668,13 +676,17 @@ where } WiiDiscWriterState::SeekToPartHeader(buf) => { crate::trace!("WiiDiscWriterFinalizeState::SeekToPartHeader"); - if pin!(&mut this.writer) + if let Poll::Ready(result) = pin!(&mut this.writer) .poll_seek( cx, - SeekFrom::Start(this.disc.partitions.partitions[part_idx].part_offset), + SeekFrom::Start(status.disc.partitions.partitions[part_idx].part_offset), ) - .is_pending() { + match result { + Ok(new_pos) => {crate::trace!("Seeked to 0x{:08X}", new_pos);}, + Err(err) => return Poll::Ready(Err(err)), + } + } else { crate::trace!("Pending..."); status.state = WiiDiscWriterState::SeekToPartHeader(buf); return Poll::Pending; @@ -808,4 +820,33 @@ where Some(disc_info) => DiscWriter::new_wii(writer, disc_info), } } + + pub fn from_reader(writer: W, reader: &DiscReader) -> Self { + match reader { + DiscReader::Gamecube(_) => DiscWriter::new_gc(writer), + DiscReader::Wii(reader) => DiscWriter::new_wii(writer, reader.disc.to_owned()), + } + } +} + +impl DiscWriter +where + W: AsyncWrite + AsyncSeek + Clone + Unpin, +{ + pub async fn init(self: &mut Pin<&mut Self>) -> Result<()> { + match self.as_mut().as_wii_disc_mut() { + Some(writer) => std::pin::pin!(writer.clone()).init().await, + None => Ok(()), + } + } +} + +impl DiscWriter +{ + fn as_wii_disc_mut(&mut self) -> Option<&mut WiiDiscWriter> { + match self { + DiscWriter::Wii(writer) => Some(writer), + _ => None, + } + } } diff --git a/geckolib/src/lib.rs b/geckolib/src/lib.rs index 7ef5d55..fbb96fa 100644 --- a/geckolib/src/lib.rs +++ b/geckolib/src/lib.rs @@ -80,12 +80,11 @@ where }; let disc_reader = DiscReader::new(iso_reader).await?; - let disc_info = disc_reader.get_disc_info(); Ok(IsoBuilder::new_with_zip( config, zip, - GeckoFS::parse(disc_reader).await?, - disc_info, + GeckoFS::parse(disc_reader.clone()).await?, + disc_reader, writer, )) } @@ -109,9 +108,8 @@ pub async fn open_config_from_fs_iso( .open(&config.build.iso) .await?; let disc_reader = DiscReader::new(async_std::fs::File::open(&config.src.iso).await?).await?; - let disc_info = disc_reader.get_disc_info(); - let gfs = GeckoFS::parse(disc_reader).await?; - Ok(IsoBuilder::new_with_fs(config, PathBuf::new(), gfs, disc_info, writer)) + let gfs = GeckoFS::parse(disc_reader.clone()).await?; + Ok(IsoBuilder::new_with_fs(config, PathBuf::new(), gfs, disc_reader, writer)) } #[cfg(not(target_arch = "wasm32"))] diff --git a/geckolib/src/vfs/mod.rs b/geckolib/src/vfs/mod.rs index 8041fa7..d9a5483 100644 --- a/geckolib/src/vfs/mod.rs +++ b/geckolib/src/vfs/mod.rs @@ -2,6 +2,7 @@ use crate::crypto::Unpackable; use crate::iso::consts::OFFSET_DOL_OFFSET; use crate::iso::disc::{align_addr, DiscType}; use crate::iso::read::DiscReader; +use crate::iso::write::DiscWriter; use crate::iso::{consts, FstEntry, FstNode, FstNodeType}; #[cfg(feature = "progress")] use crate::UPDATER; @@ -356,11 +357,12 @@ where Ok(()) } - pub async fn serialize(&mut self, writer: &mut W, is_wii: bool) -> Result<()> + pub async fn serialize(&mut self, writer: &mut DiscWriter) -> Result<()> where - W: AsyncWrite + Unpin, + W: AsyncWrite + AsyncSeek + Unpin, { crate::debug!("Serializing the FileSystem"); + let is_wii = writer.get_type() == DiscType::Wii; let mut pos: u64 = 0; let header_size = self.sys().get_file("iso.hdr")?.len()?; let apploader_size = self.sys().get_file("AppLoader.ldr")?.len()?; diff --git a/gui/native/src/app.rs b/gui/native/src/app.rs index 1e6fe4a..a45e667 100644 --- a/gui/native/src/app.rs +++ b/gui/native/src/app.rs @@ -4,6 +4,8 @@ use async_std::fs; use egui::Vec2; use flume::{Receiver, Sender, TryRecvError, TrySendError}; use futures_lite::AsyncWriteExt; +use geckolib::iso::builder::Builder; +use geckolib::open_config_from_patch; use rfd::FileHandle; use std::path::PathBuf; @@ -87,7 +89,25 @@ pub struct PatcherApp { progress: Option, } -async fn reproc(file_path: PathBuf, save_path: PathBuf) -> Result<(), eyre::Error> { +async fn apply( + patch: PathBuf, + iso: PathBuf, + save: PathBuf, +) -> Result<(), eyre::Error> { + let patch = std::fs::OpenOptions::new().read(true).open(patch)?; + let iso = async_std::fs::OpenOptions::new().read(true).open(iso).await?; + let save = async_std::fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(save) + .await?; + let mut builder = open_config_from_patch(patch, iso, save).await?; + builder.build().await?; + Ok(()) +} + +async fn _reproc(file_path: PathBuf, save_path: PathBuf) -> Result<(), eyre::Error> { // let patch = fs::OpenOptions::new().read(true).open(patch_path).await; let file = fs::OpenOptions::new() .read(true) @@ -123,9 +143,8 @@ async fn reproc(file_path: PathBuf, save_path: PathBuf) -> Result<(), eyre::Erro let mut out = out; let mut fs = GeckoFS::parse(f).await?; { - let is_wii = out.get_type() == DiscType::Wii; - fs.serialize(&mut out, is_wii).await?; - if is_wii { + fs.serialize(&mut out).await?; + if out.get_type() == DiscType::Wii { log::info!("Encrypting the ISO"); } out.flush().await?; @@ -195,7 +214,7 @@ fn patcher_thread(snd: Sender, rcv: Receiver) { } }; } - ToAppMsg::PatchAndSave(_patch, iso) => { + ToAppMsg::PatchAndSave(patch, iso) => { match rfd::AsyncFileDialog::new() .add_filter("application/x-cd-image", &["iso"]) .add_filter("All Files", &["*"]) @@ -217,7 +236,7 @@ fn patcher_thread(snd: Sender, rcv: Receiver) { log::error!("Could not send Progress (update gui)"); return; } - if let Err(err) = reproc(iso, save).await { + if let Err(err) = apply(patch, iso, save).await { log::error!("{:?}", err); if sender .send_async(FromAppMsg::Progress( diff --git a/gui/web/src/bin/worker.rs b/gui/web/src/bin/worker.rs index 6ee283f..7b8bf9a 100644 --- a/gui/web/src/bin/worker.rs +++ b/gui/web/src/bin/worker.rs @@ -1,3 +1,4 @@ +use std::io::{Read, Seek}; use std::sync::Arc; use async_std::io::prelude::{ReadExt, SeekExt}; @@ -9,7 +10,8 @@ use geckolib::iso::read::DiscReader; use geckolib::iso::write::DiscWriter; use geckolib::update::UpdaterType; use geckolib::vfs::GeckoFS; -use geckolib::UPDATER; +use geckolib::{open_config_from_patch, UPDATER}; +use geckolib::iso::builder::Builder; use wasm_bindgen::prelude::*; use wasm_bindgen::JsValue; use web_gui_patcher::io::WebFile; @@ -19,7 +21,13 @@ use web_gui_patcher::io::WebFile; static ALLOC: wasm_tracing_allocator::WasmTracingAllocator = wasm_tracing_allocator::WasmTracingAllocator(std::alloc::System); -async fn reproc( +async fn apply(patch: R1, iso: R2, save: W) -> Result<(), eyre::Error> { + let mut builder = open_config_from_patch(patch, iso, save).await?; + builder.build().await?; + Ok(()) +} + +async fn _reproc( file: R, save: W, ) -> Result<(), eyre::Error> { @@ -56,7 +64,7 @@ async fn reproc std::io::Result { + let mut state = match self.state.try_lock_arc() { + Some(guard) => guard, + None => { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Failed to lock file state", + )) + } + }; + let mut options = web_sys::FileSystemReadWriteOptions::new(); + options.at(state.cursor as f64); + match state.handle.read_with_u8_array_and_options(buf, &options) { + Ok(n) => { + state.cursor += n as u64; + Ok(n as usize) + } + Err(err) => Err(std::io::Error::new( + std::io::ErrorKind::Other, + format!("{err:?}"), + )), + } + } +} + +impl Seek for WebFile { + fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result { + let mut state = match self.state.try_lock_arc() { + Some(guard) => guard, + None => { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Failed to lock file state", + )) + } + }; + let len = match state.handle.get_size() { + Ok(size) => size as u64, + Err(err) => { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + format!("{err:?}"), + )) + } + }; + match pos { + std::io::SeekFrom::Start(pos) => { + if pos > len { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Cursor past end of stream", + )); + } + state.cursor = pos; + } + std::io::SeekFrom::End(pos) => { + let new_pos = len as i64 + pos; + if !(0..=len as i64).contains(&new_pos) { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Cursor outside of stream range", + )); + } + state.cursor = new_pos as u64; + } + std::io::SeekFrom::Current(pos) => { + let new_pos = state.cursor as i64 + pos; + if !(0..=len as i64).contains(&new_pos) { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Cursor outside of stream range", + )); + } + state.cursor = new_pos as u64; + } + }; + Ok(state.cursor) + } +} + impl AsyncRead for WebFile { fn poll_read( self: std::pin::Pin<&mut Self>, diff --git a/romhack/examples/reproc.rs b/romhack/examples/reproc.rs index 232994a..ce8656f 100644 --- a/romhack/examples/reproc.rs +++ b/romhack/examples/reproc.rs @@ -5,7 +5,7 @@ use async_std::io::{prelude::SeekExt, ReadExt}; use clap::{arg, command, Parser, ValueHint}; use futures::AsyncWriteExt; use geckolib::{ - iso::{disc::DiscType, read::DiscReader, write::DiscWriter}, vfs::GeckoFS + iso::{read::DiscReader, write::DiscWriter}, vfs::GeckoFS }; #[cfg(feature = "progress")] use geckolib::UPDATER; @@ -74,14 +74,13 @@ fn main() -> color_eyre::eyre::Result<()> { } let mut fs = GeckoFS::parse(f).await?; - let is_wii = out.get_type() == DiscType::Wii; #[cfg(feature = "log")] log::info!("Encrypting the ISO"); #[cfg(feature = "progress")] if let Ok(mut updater) = UPDATER.lock() { updater.init(None)?; } - fs.serialize(&mut out, is_wii).await?; + fs.serialize(&mut out).await?; out.close().await?; >::Ok(()) })?;