Support decrypted discs & decrypt/encrypt conversion

Cargo.toml

@@ -9,7 +9,7 @@ strip = "debuginfo"
 codegen-units = 1
 
 [workspace.package]
-version = "1.4.4"
+version = "2.0.0-alpha.1"
 edition = "2021"
 rust-version = "1.74"
 authors = ["Luke Street <[email protected]>"]

nod/src/disc/hashes.rs

@@ -16,7 +16,7 @@ use crate::{
     },
     io::HashBytes,
     util::read::read_box_slice,
-    OpenOptions, Result, ResultContext, SECTOR_SIZE,
+    PartitionOptions, Result, ResultContext, SECTOR_SIZE,
 };
 
 /// In a sector, following the 0x400 byte block of hashes, each 0x400 bytes of decrypted data is
@@ -81,7 +81,7 @@ pub fn rebuild_hashes(reader: &mut DiscReader) -> Result<()> {
 
     // Precompute hashes for zeroed sectors.
     const ZERO_H0_BYTES: &[u8] = &[0u8; HASHES_SIZE];
-    let zero_h0_hash = hash_bytes(ZERO_H0_BYTES);
+    let zero_h0_hash = sha1_hash(ZERO_H0_BYTES);
 
     let partitions = reader.partitions();
     let mut hash_tables = Vec::with_capacity(partitions.len());
@@ -97,8 +97,9 @@ pub fn rebuild_hashes(reader: &mut DiscReader) -> Result<()> {
 
         let group_count = hash_table.h3_hashes.len();
         let mutex = Arc::new(Mutex::new(hash_table));
+        let partition_options = PartitionOptions { validate_hashes: false };
         (0..group_count).into_par_iter().try_for_each_with(
-            (reader.open_partition(part.index, &OpenOptions::default())?, mutex.clone()),
+            (reader.open_partition(part.index, &partition_options)?, mutex.clone()),
             |(stream, mutex), h3_index| -> Result<()> {
                 let mut result = HashResult::new_box_zeroed()?;
                 let mut data_buf = <[u8]>::new_box_zeroed_with_elems(SECTOR_DATA_SIZE)?;
@@ -122,7 +123,7 @@ pub fn rebuild_hashes(reader: &mut DiscReader) -> Result<()> {
                                 .read_exact(&mut data_buf)
                                 .with_context(|| format!("Reading sector {}", part_sector))?;
                             for h0_index in 0..NUM_H0_HASHES {
-                                let h0_hash = hash_bytes(array_ref![
+                                let h0_hash = sha1_hash(array_ref![
                                     data_buf,
                                     h0_index * HASHES_SIZE,
                                     HASHES_SIZE
@@ -196,9 +197,6 @@ pub fn rebuild_hashes(reader: &mut DiscReader) -> Result<()> {
     Ok(())
 }
 
+/// Hashes a byte slice with SHA-1.
 #[inline]
-pub fn hash_bytes(buf: &[u8]) -> HashBytes {
-    let mut hasher = Sha1::new();
-    hasher.update(buf);
-    hasher.finalize().into()
-}
+pub fn sha1_hash(buf: &[u8]) -> HashBytes { HashBytes::from(Sha1::digest(buf)) }

nod/src/disc/mod.rs

@@ -24,7 +24,7 @@ pub(crate) mod wii;
 
 pub use fst::{Fst, Node, NodeKind};
 pub use streams::{FileStream, OwnedFileStream, WindowedStream};
-pub use wii::{SignedHeader, Ticket, TicketLimit, TmdHeader, REGION_SIZE};
+pub use wii::{ContentMetadata, SignedHeader, Ticket, TicketLimit, TmdHeader, REGION_SIZE};
 
 /// Size in bytes of a disc sector. (32 KiB)
 pub const SECTOR_SIZE: usize = 0x8000;
@@ -90,6 +90,14 @@ impl DiscHeader {
     /// Whether this is a Wii disc.
     #[inline]
     pub fn is_wii(&self) -> bool { self.wii_magic == WII_MAGIC }
+
+    /// Whether the disc has partition data hashes.
+    #[inline]
+    pub fn has_partition_hashes(&self) -> bool { self.no_partition_hashes == 0 }
+
+    /// Whether the disc has partition data encryption.
+    #[inline]
+    pub fn has_partition_encryption(&self) -> bool { self.no_partition_encryption == 0 }
 }
 
 /// A header describing the contents of a disc partition.
@@ -379,19 +387,23 @@ impl PartitionMeta {
     /// A view into the disc header.
     #[inline]
     pub fn header(&self) -> &DiscHeader {
-        DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()]).unwrap()
+        DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()])
+            .expect("Invalid header alignment")
     }
 
     /// A view into the partition header.
     #[inline]
     pub fn partition_header(&self) -> &PartitionHeader {
-        PartitionHeader::ref_from_bytes(&self.raw_boot[size_of::<DiscHeader>()..]).unwrap()
+        PartitionHeader::ref_from_bytes(&self.raw_boot[size_of::<DiscHeader>()..])
+            .expect("Invalid partition header alignment")
     }
 
     /// A view into the apploader header.
     #[inline]
     pub fn apploader_header(&self) -> &ApploaderHeader {
-        ApploaderHeader::ref_from_prefix(&self.raw_apploader).unwrap().0
+        ApploaderHeader::ref_from_prefix(&self.raw_apploader)
+            .expect("Invalid apploader alignment")
+            .0
     }
 
     /// A view into the file system table (FST).
@@ -400,18 +412,29 @@ impl PartitionMeta {
 
     /// A view into the DOL header.
     #[inline]
-    pub fn dol_header(&self) -> &DolHeader { DolHeader::ref_from_prefix(&self.raw_dol).unwrap().0 }
+    pub fn dol_header(&self) -> &DolHeader {
+        DolHeader::ref_from_prefix(&self.raw_dol).expect("Invalid DOL alignment").0
+    }
 
     /// A view into the ticket. (Wii only)
     #[inline]
     pub fn ticket(&self) -> Option<&Ticket> {
-        self.raw_ticket.as_ref().and_then(|v| Ticket::ref_from_bytes(v).ok())
+        let raw_ticket = self.raw_ticket.as_deref()?;
+        Some(Ticket::ref_from_bytes(raw_ticket).expect("Invalid ticket alignment"))
     }
 
     /// A view into the TMD. (Wii only)
     #[inline]
     pub fn tmd_header(&self) -> Option<&TmdHeader> {
-        self.raw_tmd.as_ref().and_then(|v| TmdHeader::ref_from_prefix(v).ok().map(|(v, _)| v))
+        let raw_tmd = self.raw_tmd.as_deref()?;
+        Some(TmdHeader::ref_from_prefix(raw_tmd).expect("Invalid TMD alignment").0)
+    }
+
+    /// A view into the TMD content metadata. (Wii only)
+    #[inline]
+    pub fn content_metadata(&self) -> Option<&[ContentMetadata]> {
+        let raw_cmd = &self.raw_tmd.as_deref()?[size_of::<TmdHeader>()..];
+        Some(<[ContentMetadata]>::ref_from_bytes(raw_cmd).expect("Invalid CMD alignment"))
     }
 }
 

nod/src/disc/reader.rs

@@ -3,7 +3,7 @@ use std::{
     io::{BufRead, Read, Seek, SeekFrom},
 };
 
-use zerocopy::FromZeros;
+use zerocopy::{FromBytes, FromZeros};
 
 use super::{
     gcn::PartitionGC,
@@ -16,15 +16,10 @@ use crate::{
     disc::wii::REGION_OFFSET,
     io::block::{Block, BlockIO, PartitionInfo},
     util::read::{read_box, read_from, read_vec},
-    DiscMeta, Error, OpenOptions, Result, ResultContext, SECTOR_SIZE,
+    DiscMeta, Error, OpenOptions, PartitionEncryptionMode, PartitionOptions, Result, ResultContext,
+    SECTOR_SIZE,
 };
 
-#[derive(Debug, Eq, PartialEq, Copy, Clone)]
-pub enum EncryptionMode {
-    Encrypted,
-    Decrypted,
-}
-
 pub struct DiscReader {
     io: Box<dyn BlockIO>,
     block: Block,
@@ -33,7 +28,7 @@ pub struct DiscReader {
     sector_buf: Box<[u8; SECTOR_SIZE]>,
     sector_idx: u32,
     pos: u64,
-    mode: EncryptionMode,
+    mode: PartitionEncryptionMode,
     disc_header: Box<DiscHeader>,
     pub(crate) partitions: Vec<PartitionInfo>,
     hash_tables: Vec<HashTable>,
@@ -71,11 +66,7 @@ impl DiscReader {
             sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed()?,
             sector_idx: u32::MAX,
             pos: 0,
-            mode: if options.rebuild_encryption {
-                EncryptionMode::Encrypted
-            } else {
-                EncryptionMode::Decrypted
-            },
+            mode: options.partition_encryption,
             disc_header: DiscHeader::new_box_zeroed()?,
             partitions: vec![],
             hash_tables: vec![],
@@ -84,11 +75,28 @@ impl DiscReader {
         let disc_header: Box<DiscHeader> = read_box(&mut reader).context("Reading disc header")?;
         reader.disc_header = disc_header;
         if reader.disc_header.is_wii() {
+            if reader.disc_header.has_partition_encryption()
+                && !reader.disc_header.has_partition_hashes()
+            {
+                return Err(Error::DiscFormat(
+                    "Wii disc is encrypted but has no partition hashes".to_string(),
+                ));
+            }
+            if !reader.disc_header.has_partition_hashes()
+                && options.partition_encryption == PartitionEncryptionMode::ForceEncrypted
+            {
+                return Err(Error::Other(
+                    "Unsupported: Rebuilding encryption for Wii disc without hashes".to_string(),
+                ));
+            }
             reader.seek(SeekFrom::Start(REGION_OFFSET)).context("Seeking to region info")?;
             reader.region = Some(read_from(&mut reader).context("Reading region info")?);
             reader.partitions = read_partition_info(&mut reader)?;
             // Rebuild hashes if the format requires it
-            if (options.rebuild_encryption || options.validate_hashes) && meta.needs_hash_recovery {
+            if options.partition_encryption != PartitionEncryptionMode::AsIs
+                && meta.needs_hash_recovery
+                && reader.disc_header.has_partition_hashes()
+            {
                 rebuild_hashes(&mut reader)?;
             }
         }
@@ -125,7 +133,7 @@ impl DiscReader {
     pub fn open_partition(
         &self,
         index: usize,
-        options: &OpenOptions,
+        options: &PartitionOptions,
     ) -> Result<Box<dyn PartitionBase>> {
         if self.disc_header.is_gamecube() {
             if index == 0 {
@@ -145,7 +153,7 @@ impl DiscReader {
     pub fn open_partition_kind(
         &self,
         kind: PartitionKind,
-        options: &OpenOptions,
+        options: &PartitionOptions,
     ) -> Result<Box<dyn PartitionBase>> {
         if self.disc_header.is_gamecube() {
             if kind == PartitionKind::Data {
@@ -182,30 +190,51 @@ impl BufRead for DiscReader {
 
         // Read new sector into buffer
         if abs_sector != self.sector_idx {
-            if let Some(partition) = partition {
-                match self.mode {
-                    EncryptionMode::Decrypted => self.block.decrypt(
+            match (self.mode, partition, self.disc_header.has_partition_encryption()) {
+                (PartitionEncryptionMode::Original, Some(partition), true)
+                | (PartitionEncryptionMode::ForceEncrypted, Some(partition), _) => {
+                    self.block.encrypt(
                         self.sector_buf.as_mut(),
                         self.block_buf.as_ref(),
                         abs_sector,
                         partition,
-                    )?,
-                    EncryptionMode::Encrypted => self.block.encrypt(
+                    )?;
+                }
+                (PartitionEncryptionMode::ForceDecrypted, Some(partition), _) => {
+                    self.block.decrypt(
                         self.sector_buf.as_mut(),
                         self.block_buf.as_ref(),
                         abs_sector,
                         partition,
-                    )?,
+                    )?;
+                }
+                (PartitionEncryptionMode::AsIs, _, _) | (_, None, _) | (_, _, false) => {
+                    self.block.copy_raw(
+                        self.sector_buf.as_mut(),
+                        self.block_buf.as_ref(),
+                        abs_sector,
+                        &self.disc_header,
+                    )?;
                 }
-            } else {
-                self.block.copy_raw(
-                    self.sector_buf.as_mut(),
-                    self.block_buf.as_ref(),
-                    abs_sector,
-                    &self.disc_header,
-                )?;
             }
             self.sector_idx = abs_sector;
+
+            if self.sector_idx == 0
+                && self.disc_header.is_wii()
+                && matches!(
+                    self.mode,
+                    PartitionEncryptionMode::ForceDecrypted
+                        | PartitionEncryptionMode::ForceEncrypted
+                )
+            {
+                let (disc_header, _) = DiscHeader::mut_from_prefix(self.sector_buf.as_mut())
+                    .expect("Invalid disc header alignment");
+                disc_header.no_partition_encryption = match self.mode {
+                    PartitionEncryptionMode::ForceDecrypted => 1,
+                    PartitionEncryptionMode::ForceEncrypted => 0,
+                    _ => unreachable!(),
+                };
+            }
         }
 
         // Read from sector buffer
@@ -273,8 +302,19 @@ fn read_partition_info(reader: &mut DiscReader) -> Result<Vec<PartitionInfo>> {
                     "Partition {group_idx}:{part_idx} offset is not sector aligned",
                 )));
             }
+
+            let disc_header = reader.header();
             let data_start_offset = entry.offset() + header.data_off();
-            let data_end_offset = data_start_offset + header.data_size();
+            let mut data_size = header.data_size();
+            if data_size == 0 {
+                // Read until next partition or end of disc
+                // TODO: handle multiple partition groups
+                data_size = entries
+                    .get(part_idx + 1)
+                    .map(|part| part.offset() - data_start_offset)
+                    .unwrap_or(reader.disc_size() - data_start_offset);
+            }
+            let data_end_offset = data_start_offset + data_size;
             if data_start_offset % SECTOR_SIZE as u64 != 0
                 || data_end_offset % SECTOR_SIZE as u64 != 0
             {
@@ -293,13 +333,15 @@ fn read_partition_info(reader: &mut DiscReader) -> Result<Vec<PartitionInfo>> {
                 disc_header: DiscHeader::new_box_zeroed()?,
                 partition_header: PartitionHeader::new_box_zeroed()?,
                 hash_table: None,
+                has_encryption: disc_header.has_partition_encryption(),
+                has_hashes: disc_header.has_partition_hashes(),
             };
 
             let mut partition_reader = PartitionWii::new(
                 reader.io.clone(),
                 reader.disc_header.clone(),
                 &info,
-                &OpenOptions::default(),
+                &PartitionOptions { validate_hashes: false },
             )?;
             info.disc_header = read_box(&mut partition_reader).context("Reading disc header")?;
             info.partition_header =