nod/io/
block.rs

1use std::{
2    fs, io,
3    io::{Read, Seek},
4    path::Path,
5};
6
7use dyn_clone::DynClone;
8use zerocopy::transmute_ref;
9
10use crate::{
11    array_ref,
12    disc::{
13        hashes::HashTable,
14        wii::{WiiPartitionHeader, HASHES_SIZE, SECTOR_DATA_SIZE},
15        DiscHeader, PartitionHeader, PartitionKind, GCN_MAGIC, SECTOR_SIZE, WII_MAGIC,
16    },
17    io::{
18        aes_decrypt, aes_encrypt, split::SplitFileReader, DiscMeta, Format, KeyBytes, MagicBytes,
19    },
20    util::{lfg::LaggedFibonacci, read::read_from},
21    Error, Result, ResultContext,
22};
23
24/// Required trait bounds for reading disc images.
25pub trait DiscStream: Read + Seek + DynClone + Send + Sync {}
26
27impl<T> DiscStream for T where T: Read + Seek + DynClone + Send + Sync + ?Sized {}
28
29dyn_clone::clone_trait_object!(DiscStream);
30
31/// Block I/O trait for reading disc images.
32pub trait BlockIO: DynClone + Send + Sync {
33    /// Reads a block from the disc image.
34    fn read_block_internal(
35        &mut self,
36        out: &mut [u8],
37        block: u32,
38        partition: Option<&PartitionInfo>,
39    ) -> io::Result<Block>;
40
41    /// Reads a full block from the disc image, combining smaller blocks if necessary.
42    fn read_block(
43        &mut self,
44        out: &mut [u8],
45        block: u32,
46        partition: Option<&PartitionInfo>,
47    ) -> io::Result<Block> {
48        let block_size_internal = self.block_size_internal();
49        let block_size = self.block_size();
50        if block_size_internal == block_size {
51            self.read_block_internal(out, block, partition)
52        } else {
53            let mut offset = 0usize;
54            let mut result = None;
55            let mut block_idx =
56                ((block as u64 * block_size as u64) / block_size_internal as u64) as u32;
57            while offset < block_size as usize {
58                let block = self.read_block_internal(
59                    &mut out[offset..offset + block_size_internal as usize],
60                    block_idx,
61                    partition,
62                )?;
63                if result.is_none() {
64                    result = Some(block);
65                } else if result != Some(block) {
66                    if block == Block::Zero {
67                        out[offset..offset + block_size_internal as usize].fill(0);
68                    } else {
69                        return Err(io::Error::new(
70                            io::ErrorKind::InvalidData,
71                            "Inconsistent block types in split block",
72                        ));
73                    }
74                }
75                offset += block_size_internal as usize;
76                block_idx += 1;
77            }
78            Ok(result.unwrap_or_default())
79        }
80    }
81
82    /// The format's block size in bytes. Can be smaller than the sector size (0x8000).
83    fn block_size_internal(&self) -> u32;
84
85    /// The block size used for processing. Must be a multiple of the sector size (0x8000).
86    fn block_size(&self) -> u32 { self.block_size_internal().max(SECTOR_SIZE as u32) }
87
88    /// Returns extra metadata included in the disc file format, if any.
89    fn meta(&self) -> DiscMeta;
90}
91
92dyn_clone::clone_trait_object!(BlockIO);
93
94/// Creates a new [`BlockIO`] instance from a stream.
95pub fn new(mut stream: Box<dyn DiscStream>) -> Result<Box<dyn BlockIO>> {
96    let io: Box<dyn BlockIO> = match detect(stream.as_mut()).context("Detecting file type")? {
97        Some(Format::Iso) => crate::io::iso::DiscIOISO::new(stream)?,
98        Some(Format::Ciso) => crate::io::ciso::DiscIOCISO::new(stream)?,
99        Some(Format::Gcz) => {
100            #[cfg(feature = "compress-zlib")]
101            {
102                crate::io::gcz::DiscIOGCZ::new(stream)?
103            }
104            #[cfg(not(feature = "compress-zlib"))]
105            return Err(Error::DiscFormat("GCZ support is disabled".to_string()));
106        }
107        Some(Format::Nfs) => {
108            return Err(Error::DiscFormat("NFS requires a filesystem path".to_string()))
109        }
110        Some(Format::Wbfs) => crate::io::wbfs::DiscIOWBFS::new(stream)?,
111        Some(Format::Wia | Format::Rvz) => crate::io::wia::DiscIOWIA::new(stream)?,
112        Some(Format::Tgc) => crate::io::tgc::DiscIOTGC::new(stream)?,
113        None => return Err(Error::DiscFormat("Unknown disc format".to_string())),
114    };
115    check_block_size(io.as_ref())?;
116    Ok(io)
117}
118
119/// Creates a new [`BlockIO`] instance from a filesystem path.
120pub fn open(filename: &Path) -> Result<Box<dyn BlockIO>> {
121    let path_result = fs::canonicalize(filename);
122    if let Err(err) = path_result {
123        return Err(Error::Io(format!("Failed to open {}", filename.display()), err));
124    }
125    let path = path_result.as_ref().unwrap();
126    let meta = fs::metadata(path);
127    if let Err(err) = meta {
128        return Err(Error::Io(format!("Failed to open {}", filename.display()), err));
129    }
130    if !meta.unwrap().is_file() {
131        return Err(Error::DiscFormat(format!("Input is not a file: {}", filename.display())));
132    }
133    let mut stream = Box::new(SplitFileReader::new(filename)?);
134    let io: Box<dyn BlockIO> = match detect(stream.as_mut()).context("Detecting file type")? {
135        Some(Format::Iso) => crate::io::iso::DiscIOISO::new(stream)?,
136        Some(Format::Ciso) => crate::io::ciso::DiscIOCISO::new(stream)?,
137        Some(Format::Gcz) => {
138            #[cfg(feature = "compress-zlib")]
139            {
140                crate::io::gcz::DiscIOGCZ::new(stream)?
141            }
142            #[cfg(not(feature = "compress-zlib"))]
143            return Err(Error::DiscFormat("GCZ support is disabled".to_string()));
144        }
145        Some(Format::Nfs) => match path.parent() {
146            Some(parent) if parent.is_dir() => {
147                crate::io::nfs::DiscIONFS::new(path.parent().unwrap())?
148            }
149            _ => {
150                return Err(Error::DiscFormat("Failed to locate NFS parent directory".to_string()));
151            }
152        },
153        Some(Format::Tgc) => crate::io::tgc::DiscIOTGC::new(stream)?,
154        Some(Format::Wbfs) => crate::io::wbfs::DiscIOWBFS::new(stream)?,
155        Some(Format::Wia | Format::Rvz) => crate::io::wia::DiscIOWIA::new(stream)?,
156        None => return Err(Error::DiscFormat("Unknown disc format".to_string())),
157    };
158    check_block_size(io.as_ref())?;
159    Ok(io)
160}
161
162pub const CISO_MAGIC: MagicBytes = *b"CISO";
163pub const GCZ_MAGIC: MagicBytes = [0x01, 0xC0, 0x0B, 0xB1];
164pub const NFS_MAGIC: MagicBytes = *b"EGGS";
165pub const TGC_MAGIC: MagicBytes = [0xae, 0x0f, 0x38, 0xa2];
166pub const WBFS_MAGIC: MagicBytes = *b"WBFS";
167pub const WIA_MAGIC: MagicBytes = *b"WIA\x01";
168pub const RVZ_MAGIC: MagicBytes = *b"RVZ\x01";
169
170pub fn detect<R: Read + ?Sized>(stream: &mut R) -> io::Result<Option<Format>> {
171    let data: [u8; 0x20] = match read_from(stream) {
172        Ok(magic) => magic,
173        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
174        Err(e) => return Err(e),
175    };
176    let out = match *array_ref!(data, 0, 4) {
177        CISO_MAGIC => Some(Format::Ciso),
178        GCZ_MAGIC => Some(Format::Gcz),
179        NFS_MAGIC => Some(Format::Nfs),
180        TGC_MAGIC => Some(Format::Tgc),
181        WBFS_MAGIC => Some(Format::Wbfs),
182        WIA_MAGIC => Some(Format::Wia),
183        RVZ_MAGIC => Some(Format::Rvz),
184        _ if *array_ref!(data, 0x18, 4) == WII_MAGIC || *array_ref!(data, 0x1C, 4) == GCN_MAGIC => {
185            Some(Format::Iso)
186        }
187        _ => None,
188    };
189    Ok(out)
190}
191
192fn check_block_size(io: &dyn BlockIO) -> Result<()> {
193    if io.block_size_internal() < SECTOR_SIZE as u32
194        && SECTOR_SIZE as u32 % io.block_size_internal() != 0
195    {
196        return Err(Error::DiscFormat(format!(
197            "Sector size {} is not divisible by block size {}",
198            SECTOR_SIZE,
199            io.block_size_internal(),
200        )));
201    }
202    if io.block_size() % SECTOR_SIZE as u32 != 0 {
203        return Err(Error::DiscFormat(format!(
204            "Block size {} is not a multiple of sector size {}",
205            io.block_size(),
206            SECTOR_SIZE
207        )));
208    }
209    Ok(())
210}
211
212/// Wii partition information.
213#[derive(Debug, Clone)]
214pub struct PartitionInfo {
215    /// The partition index.
216    pub index: usize,
217    /// The kind of disc partition.
218    pub kind: PartitionKind,
219    /// The start sector of the partition.
220    pub start_sector: u32,
221    /// The start sector of the partition's (encrypted) data.
222    pub data_start_sector: u32,
223    /// The end sector of the partition's (encrypted) data.
224    pub data_end_sector: u32,
225    /// The AES key for the partition, also known as the "title key".
226    pub key: KeyBytes,
227    /// The Wii partition header.
228    pub header: Box<WiiPartitionHeader>,
229    /// The disc header within the partition.
230    pub disc_header: Box<DiscHeader>,
231    /// The partition header within the partition.
232    pub partition_header: Box<PartitionHeader>,
233    /// The hash table for the partition, if rebuilt.
234    pub hash_table: Option<HashTable>,
235}
236
237/// The block kind returned by [`BlockIO::read_block`].
238#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
239pub enum Block {
240    /// Raw data or encrypted Wii partition data
241    Raw,
242    /// Decrypted Wii partition data
243    PartDecrypted {
244        /// Whether the sector has its hash block intact
245        has_hashes: bool,
246    },
247    /// Wii partition junk data
248    Junk,
249    /// All zeroes
250    #[default]
251    Zero,
252}
253
254impl Block {
255    /// Decrypts the block's data (if necessary) and writes it to the output buffer.
256    pub(crate) fn decrypt(
257        self,
258        out: &mut [u8; SECTOR_SIZE],
259        data: &[u8],
260        abs_sector: u32,
261        partition: &PartitionInfo,
262    ) -> io::Result<()> {
263        let part_sector = abs_sector - partition.data_start_sector;
264        match self {
265            Block::Raw => {
266                out.copy_from_slice(block_sector::<SECTOR_SIZE>(data, abs_sector)?);
267                decrypt_sector(out, partition);
268            }
269            Block::PartDecrypted { has_hashes } => {
270                out.copy_from_slice(block_sector::<SECTOR_SIZE>(data, abs_sector)?);
271                if !has_hashes {
272                    rebuild_hash_block(out, part_sector, partition);
273                }
274            }
275            Block::Junk => {
276                generate_junk(out, part_sector, Some(partition), &partition.disc_header);
277                rebuild_hash_block(out, part_sector, partition);
278            }
279            Block::Zero => {
280                out.fill(0);
281                rebuild_hash_block(out, part_sector, partition);
282            }
283        }
284        Ok(())
285    }
286
287    /// Encrypts the block's data (if necessary) and writes it to the output buffer.
288    pub(crate) fn encrypt(
289        self,
290        out: &mut [u8; SECTOR_SIZE],
291        data: &[u8],
292        abs_sector: u32,
293        partition: &PartitionInfo,
294    ) -> io::Result<()> {
295        let part_sector = abs_sector - partition.data_start_sector;
296        match self {
297            Block::Raw => {
298                out.copy_from_slice(block_sector::<SECTOR_SIZE>(data, abs_sector)?);
299            }
300            Block::PartDecrypted { has_hashes } => {
301                out.copy_from_slice(block_sector::<SECTOR_SIZE>(data, abs_sector)?);
302                if !has_hashes {
303                    rebuild_hash_block(out, part_sector, partition);
304                }
305                encrypt_sector(out, partition);
306            }
307            Block::Junk => {
308                generate_junk(out, part_sector, Some(partition), &partition.disc_header);
309                rebuild_hash_block(out, part_sector, partition);
310                encrypt_sector(out, partition);
311            }
312            Block::Zero => {
313                out.fill(0);
314                rebuild_hash_block(out, part_sector, partition);
315                encrypt_sector(out, partition);
316            }
317        }
318        Ok(())
319    }
320
321    /// Copies the block's raw data to the output buffer.
322    pub(crate) fn copy_raw(
323        self,
324        out: &mut [u8; SECTOR_SIZE],
325        data: &[u8],
326        abs_sector: u32,
327        disc_header: &DiscHeader,
328    ) -> io::Result<()> {
329        match self {
330            Block::Raw => {
331                out.copy_from_slice(block_sector::<SECTOR_SIZE>(data, abs_sector)?);
332            }
333            Block::PartDecrypted { .. } => {
334                return Err(io::Error::new(
335                    io::ErrorKind::InvalidData,
336                    "Cannot copy decrypted data as raw",
337                ));
338            }
339            Block::Junk => generate_junk(out, abs_sector, None, disc_header),
340            Block::Zero => out.fill(0),
341        }
342        Ok(())
343    }
344}
345
346#[inline(always)]
347fn block_sector<const N: usize>(data: &[u8], sector_idx: u32) -> io::Result<&[u8; N]> {
348    if data.len() % N != 0 {
349        return Err(io::Error::new(
350            io::ErrorKind::InvalidData,
351            format!("Expected block size {} to be a multiple of {}", data.len(), N),
352        ));
353    }
354    let rel_sector = sector_idx % (data.len() / N) as u32;
355    let offset = rel_sector as usize * N;
356    data.get(offset..offset + N)
357        .ok_or_else(|| {
358            io::Error::new(
359                io::ErrorKind::InvalidData,
360                format!(
361                    "Sector {} out of range (block size {}, sector size {})",
362                    rel_sector,
363                    data.len(),
364                    N
365                ),
366            )
367        })
368        .map(|v| unsafe { &*(v as *const [u8] as *const [u8; N]) })
369}
370
371fn generate_junk(
372    out: &mut [u8; SECTOR_SIZE],
373    sector: u32,
374    partition: Option<&PartitionInfo>,
375    disc_header: &DiscHeader,
376) {
377    let (pos, offset) = if partition.is_some() {
378        (sector as u64 * SECTOR_DATA_SIZE as u64, HASHES_SIZE)
379    } else {
380        (sector as u64 * SECTOR_SIZE as u64, 0)
381    };
382    out[..offset].fill(0);
383    let mut lfg = LaggedFibonacci::default();
384    lfg.fill_sector_chunked(
385        &mut out[offset..],
386        *array_ref![disc_header.game_id, 0, 4],
387        disc_header.disc_num,
388        pos,
389    );
390}
391
392fn rebuild_hash_block(out: &mut [u8; SECTOR_SIZE], part_sector: u32, partition: &PartitionInfo) {
393    let Some(hash_table) = partition.hash_table.as_ref() else {
394        return;
395    };
396    let sector_idx = part_sector as usize;
397    let h0_hashes: &[u8; 0x26C] =
398        transmute_ref!(array_ref![hash_table.h0_hashes, sector_idx * 31, 31]);
399    out[0..0x26C].copy_from_slice(h0_hashes);
400    let h1_hashes: &[u8; 0xA0] =
401        transmute_ref!(array_ref![hash_table.h1_hashes, sector_idx & !7, 8]);
402    out[0x280..0x320].copy_from_slice(h1_hashes);
403    let h2_hashes: &[u8; 0xA0] =
404        transmute_ref!(array_ref![hash_table.h2_hashes, (sector_idx / 8) & !7, 8]);
405    out[0x340..0x3E0].copy_from_slice(h2_hashes);
406}
407
408fn encrypt_sector(out: &mut [u8; SECTOR_SIZE], partition: &PartitionInfo) {
409    aes_encrypt(&partition.key, [0u8; 16], &mut out[..HASHES_SIZE]);
410    // Data IV from encrypted hash block
411    let iv = *array_ref![out, 0x3D0, 16];
412    aes_encrypt(&partition.key, iv, &mut out[HASHES_SIZE..]);
413}
414
415fn decrypt_sector(out: &mut [u8; SECTOR_SIZE], partition: &PartitionInfo) {
416    // Data IV from encrypted hash block
417    let iv = *array_ref![out, 0x3D0, 16];
418    aes_decrypt(&partition.key, [0u8; 16], &mut out[..HASHES_SIZE]);
419    aes_decrypt(&partition.key, iv, &mut out[HASHES_SIZE..]);
420}
nod/io/block.rs

nod/io/
block.rs
