From 32e08f95430dbcfcf14362804e7223e5212afe0b Mon Sep 17 00:00:00 2001 From: Luke Street Date: Fri, 18 Oct 2024 00:04:11 -0600 Subject: [PATCH] Export LaggedFibonacci & add more helper methods --- nod/src/io/block.rs | 21 +++---- nod/src/lib.rs | 1 + nod/src/util/lfg.rs | 144 +++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 145 insertions(+), 21 deletions(-) diff --git a/nod/src/io/block.rs b/nod/src/io/block.rs index b98d115..a322fb2 100644 --- a/nod/src/io/block.rs +++ b/nod/src/io/block.rs @@ -1,5 +1,4 @@ use std::{ - cmp::min, fs, io, io::{Read, Seek}, path::Path, @@ -375,23 +374,19 @@ fn generate_junk( partition: Option<&PartitionInfo>, disc_header: &DiscHeader, ) { - let (mut pos, mut offset) = if partition.is_some() { + let (pos, offset) = if partition.is_some() { (sector as u64 * SECTOR_DATA_SIZE as u64, HASHES_SIZE) } else { (sector as u64 * SECTOR_SIZE as u64, 0) }; out[..offset].fill(0); - while offset < SECTOR_SIZE { - // The LFG spans a single sector of the decrypted data, - // so we may need to initialize it multiple times - let mut lfg = LaggedFibonacci::default(); - lfg.init_with_seed(*array_ref![disc_header.game_id, 0, 4], disc_header.disc_num, pos); - let sector_end = (pos + SECTOR_SIZE as u64) & !(SECTOR_SIZE as u64 - 1); - let len = min(SECTOR_SIZE - offset, (sector_end - pos) as usize); - lfg.fill(&mut out[offset..offset + len]); - pos += len as u64; - offset += len; - } + let mut lfg = LaggedFibonacci::default(); + lfg.fill_sector_chunked( + &mut out[offset..], + *array_ref![disc_header.game_id, 0, 4], + disc_header.disc_num, + pos, + ); } fn rebuild_hash_block(out: &mut [u8; SECTOR_SIZE], part_sector: u32, partition: &PartitionInfo) { diff --git a/nod/src/lib.rs b/nod/src/lib.rs index 4ae0f66..f895071 100644 --- a/nod/src/lib.rs +++ b/nod/src/lib.rs @@ -73,6 +73,7 @@ pub use io::{ block::{DiscStream, PartitionInfo}, Compression, DiscMeta, Format, KeyBytes, MagicBytes, }; +pub use util::lfg::LaggedFibonacci; mod disc; mod io; diff --git a/nod/src/util/lfg.rs b/nod/src/util/lfg.rs index 211c84b..1e43abd 100644 --- a/nod/src/util/lfg.rs +++ b/nod/src/util/lfg.rs @@ -1,4 +1,8 @@ -use std::{cmp::min, io, io::Read}; +use std::{ + cmp::min, + io, + io::{Read, Write}, +}; use zerocopy::{transmute_ref, IntoBytes}; @@ -8,7 +12,7 @@ pub const LFG_K: usize = 521; pub const LFG_J: usize = 32; pub const SEED_SIZE: usize = 17; -/// Lagged Fibonacci generator for Wii partition junk data. +/// Lagged Fibonacci generator for GC / Wii partition junk data. /// /// References (license CC0-1.0): /// https://github.com/dolphin-emu/dolphin/blob/a0f555648c27ec0c928f6b1e1fcad5e2d7c4d0c4/docs/WiaAndRvz.md @@ -19,6 +23,7 @@ pub struct LaggedFibonacci { } impl Default for LaggedFibonacci { + #[inline] fn default() -> Self { Self { buffer: [0u32; LFG_K], position: 0 } } } @@ -38,12 +43,16 @@ impl LaggedFibonacci { } } - pub fn init_with_seed(&mut self, init: [u8; 4], disc_num: u8, partition_offset: u64) { + /// Initializes the LFG with the standard seed for a given disc ID, disc number, and sector. + /// The partition offset is used to determine the sector and how many bytes to skip within the + /// sector. + #[allow(clippy::missing_inline_in_public_items)] + pub fn init_with_seed(&mut self, disc_id: [u8; 4], disc_num: u8, partition_offset: u64) { let seed = u32::from_be_bytes([ - init[2], - init[1], - init[3].wrapping_add(init[2]), - init[0].wrapping_add(init[1]), + disc_id[2], + disc_id[1], + disc_id[3].wrapping_add(disc_id[2]), + disc_id[0].wrapping_add(disc_id[1]), ]) ^ disc_num as u32; let sector = (partition_offset / SECTOR_SIZE as u64) as u32; let sector_offset = partition_offset % SECTOR_SIZE as u64; @@ -62,6 +71,9 @@ impl LaggedFibonacci { self.skip(sector_offset as usize); } + /// Initializes the LFG with the seed read from a reader. The seed is assumed to be big-endian. + /// This is used for rebuilding junk data in WIA/RVZ files. + #[allow(clippy::missing_inline_in_public_items)] pub fn init_with_reader(&mut self, reader: &mut R) -> io::Result<()> where R: Read + ?Sized { reader.read_exact(self.buffer[..SEED_SIZE].as_mut_bytes())?; @@ -73,7 +85,8 @@ impl LaggedFibonacci { Ok(()) } - pub fn forward(&mut self) { + /// Advances the LFG by one step. + fn forward(&mut self) { for i in 0..LFG_J { self.buffer[i] ^= self.buffer[i + LFG_K - LFG_J]; } @@ -82,6 +95,8 @@ impl LaggedFibonacci { } } + /// Skips `n` bytes of junk data. + #[allow(clippy::missing_inline_in_public_items)] pub fn skip(&mut self, n: usize) { self.position += n; while self.position >= LFG_K * 4 { @@ -90,6 +105,8 @@ impl LaggedFibonacci { } } + /// Fills the buffer with junk data. + #[allow(clippy::missing_inline_in_public_items)] pub fn fill(&mut self, mut buf: &mut [u8]) { while !buf.is_empty() { let len = min(buf.len(), LFG_K * 4 - self.position); @@ -103,6 +120,68 @@ impl LaggedFibonacci { } } } + + /// Writes junk data to the output stream. + #[allow(clippy::missing_inline_in_public_items)] + pub fn write(&mut self, w: &mut W, mut len: u64) -> io::Result<()> + where W: Write + ?Sized { + while len > 0 { + let write_len = min(len, LFG_K as u64 * 4 - self.position as u64) as usize; + let bytes: &[u8; LFG_K * 4] = transmute_ref!(&self.buffer); + w.write_all(&bytes[self.position..self.position + write_len])?; + self.position += write_len; + len -= write_len as u64; + if self.position == LFG_K * 4 { + self.forward(); + self.position = 0; + } + } + Ok(()) + } + + /// The junk data on GC / Wii discs is reinitialized every 32KB. This functions handles the + /// wrapping logic and reinitializes the LFG at sector boundaries. + #[allow(clippy::missing_inline_in_public_items)] + pub fn fill_sector_chunked( + &mut self, + mut buf: &mut [u8], + disc_id: [u8; 4], + disc_num: u8, + mut partition_offset: u64, + ) { + while !buf.is_empty() { + self.init_with_seed(disc_id, disc_num, partition_offset); + let len = + (SECTOR_SIZE - (partition_offset % SECTOR_SIZE as u64) as usize).min(buf.len()); + self.fill(&mut buf[..len]); + buf = &mut buf[len..]; + partition_offset += len as u64; + } + } + + /// The junk data on GC / Wii discs is reinitialized every 32KB. This functions handles the + /// wrapping logic and reinitializes the LFG at sector boundaries. + #[allow(clippy::missing_inline_in_public_items)] + pub fn write_sector_chunked( + &mut self, + w: &mut W, + mut len: u64, + disc_id: [u8; 4], + disc_num: u8, + mut partition_offset: u64, + ) -> io::Result<()> + where + W: Write + ?Sized, + { + while len > 0 { + self.init_with_seed(disc_id, disc_num, partition_offset); + let write_len = (SECTOR_SIZE as u64 - (partition_offset % SECTOR_SIZE as u64)).min(len); + self.write(w, write_len)?; + len -= write_len; + partition_offset += write_len; + } + Ok(()) + } } #[cfg(test)] @@ -132,4 +211,53 @@ mod tests { 0xEA, 0xD0 ]); } + + #[test] + fn test_init_with_seed_3() { + let mut lfg = LaggedFibonacci::default(); + lfg.init_with_seed([0x47, 0x50, 0x49, 0x45], 0, 0x322904); + let mut buf = [0u8; 16]; + lfg.fill(&mut buf); + assert_eq!(buf, [ + 0x97, 0xD8, 0x23, 0x0B, 0x12, 0xAA, 0x20, 0x45, 0xC2, 0xBD, 0x71, 0x8C, 0x30, 0x32, + 0xC5, 0x2F + ]); + } + + #[test] + fn test_write() { + let mut lfg = LaggedFibonacci::default(); + lfg.init_with_seed([0x47, 0x50, 0x49, 0x45], 0, 0x322904); + let mut buf = [0u8; 16]; + lfg.write(&mut buf.as_mut_slice(), 16).unwrap(); + assert_eq!(buf, [ + 0x97, 0xD8, 0x23, 0x0B, 0x12, 0xAA, 0x20, 0x45, 0xC2, 0xBD, 0x71, 0x8C, 0x30, 0x32, + 0xC5, 0x2F + ]); + } + + #[test] + fn test_fill_sector_chunked() { + let mut lfg = LaggedFibonacci::default(); + let mut buf = [0u8; 32]; + lfg.fill_sector_chunked(&mut buf, [0x47, 0x4D, 0x38, 0x45], 0, 0x27FF0); + assert_eq!(buf, [ + 0xAD, 0x6F, 0x21, 0xBE, 0x05, 0x57, 0x10, 0xED, 0xEA, 0xB0, 0x8E, 0xFD, 0x91, 0x58, + 0xA2, 0x0E, 0xDC, 0x0D, 0x59, 0xC0, 0x02, 0x98, 0xA5, 0x00, 0x39, 0x5B, 0x68, 0xA6, + 0x5D, 0x53, 0x2D, 0xB6 + ]); + } + + #[test] + fn test_write_sector_chunked() { + let mut lfg = LaggedFibonacci::default(); + let mut buf = [0u8; 32]; + lfg.write_sector_chunked(&mut buf.as_mut_slice(), 32, [0x47, 0x4D, 0x38, 0x45], 0, 0x27FF0) + .unwrap(); + assert_eq!(buf, [ + 0xAD, 0x6F, 0x21, 0xBE, 0x05, 0x57, 0x10, 0xED, 0xEA, 0xB0, 0x8E, 0xFD, 0x91, 0x58, + 0xA2, 0x0E, 0xDC, 0x0D, 0x59, 0xC0, 0x02, 0x98, 0xA5, 0x00, 0x39, 0x5B, 0x68, 0xA6, + 0x5D, 0x53, 0x2D, 0xB6 + ]); + } }