nod-rs/nod/src/disc/reader.rs

use std::{
    cmp::min,
    io,
    io::{Read, Seek, SeekFrom},
};

use zerocopy::FromZeroes;

use crate::{
    disc::{
        gcn::PartitionGC,
        hashes::{rebuild_hashes, HashTable},
        wii::{PartitionWii, WiiPartEntry, WiiPartGroup, WiiPartitionHeader, WII_PART_GROUP_OFF},
        DL_DVD_SIZE, MINI_DVD_SIZE, SL_DVD_SIZE,
    },
    io::block::{Block, BlockIO, PartitionInfo},
    util::read::{read_box, read_from, read_vec},
    DiscHeader, DiscMeta, Error, OpenOptions, PartitionBase, PartitionHeader, PartitionKind,
    Result, ResultContext, SECTOR_SIZE,
};

#[derive(Debug, Eq, PartialEq, Copy, Clone)]
pub enum EncryptionMode {
    Encrypted,
    Decrypted,
}

pub struct DiscReader {
    io: Box<dyn BlockIO>,
    block: Block,
    block_buf: Box<[u8]>,
    block_idx: u32,
    sector_buf: Box<[u8; SECTOR_SIZE]>,
    sector_idx: u32,
    pos: u64,
    mode: EncryptionMode,
    disc_header: Box<DiscHeader>,
    pub(crate) partitions: Vec<PartitionInfo>,
    hash_tables: Vec<HashTable>,
}

impl Clone for DiscReader {
    fn clone(&self) -> Self {
        Self {
            io: self.io.clone(),
            block: Block::default(),
            block_buf: <u8>::new_box_slice_zeroed(self.block_buf.len()),
            block_idx: u32::MAX,
            sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed(),
            sector_idx: u32::MAX,
            pos: 0,
            mode: self.mode,
            disc_header: self.disc_header.clone(),
            partitions: self.partitions.clone(),
            hash_tables: self.hash_tables.clone(),
        }
    }
}

impl DiscReader {
    pub fn new(inner: Box<dyn BlockIO>, options: &OpenOptions) -> Result<Self> {
        let block_size = inner.block_size();
        let meta = inner.meta();
        let mut reader = Self {
            io: inner,
            block: Block::default(),
            block_buf: <u8>::new_box_slice_zeroed(block_size as usize),
            block_idx: u32::MAX,
            sector_buf: <[u8; SECTOR_SIZE]>::new_box_zeroed(),
            sector_idx: u32::MAX,
            pos: 0,
            mode: if options.rebuild_encryption {
                EncryptionMode::Encrypted
            } else {
                EncryptionMode::Decrypted
            },
            disc_header: DiscHeader::new_box_zeroed(),
            partitions: vec![],
            hash_tables: vec![],
        };
        let disc_header: Box<DiscHeader> = read_box(&mut reader).context("Reading disc header")?;
        reader.disc_header = disc_header;
        if reader.disc_header.is_wii() {
            reader.partitions = read_partition_info(&mut reader)?;
            // Rebuild hashes if the format requires it
            if options.rebuild_encryption && meta.needs_hash_recovery {
                rebuild_hashes(&mut reader)?;
            }
        }
        reader.reset();
        Ok(reader)
    }

    pub fn reset(&mut self) {
        self.block = Block::default();
        self.block_buf.fill(0);
        self.block_idx = u32::MAX;
        self.sector_buf.fill(0);
        self.sector_idx = u32::MAX;
        self.pos = 0;
    }

    pub fn disc_size(&self) -> u64 {
        self.io.meta().disc_size.unwrap_or_else(|| guess_disc_size(&self.partitions))
    }

    pub fn header(&self) -> &DiscHeader { &self.disc_header }

    pub fn partitions(&self) -> &[PartitionInfo] { &self.partitions }

    pub fn meta(&self) -> DiscMeta { self.io.meta() }

    /// Opens a new, decrypted partition read stream for the specified partition index.
    pub fn open_partition(
        &self,
        index: usize,
        options: &OpenOptions,
    ) -> Result<Box<dyn PartitionBase>> {
        if self.disc_header.is_gamecube() {
            if index == 0 {
                Ok(PartitionGC::new(self.io.clone(), self.disc_header.clone())?)
            } else {
                Err(Error::DiscFormat("GameCube discs only have one partition".to_string()))
            }
        } else if let Some(part) = self.partitions.get(index) {
            Ok(PartitionWii::new(self.io.clone(), self.disc_header.clone(), part, options)?)
        } else {
            Err(Error::DiscFormat(format!("Partition {index} not found")))
        }
    }

    /// Opens a new, decrypted partition read stream for the first partition matching
    /// the specified type.
    pub fn open_partition_kind(
        &self,
        part_type: PartitionKind,
        options: &OpenOptions,
    ) -> Result<Box<dyn PartitionBase>> {
        if self.disc_header.is_gamecube() {
            if part_type == PartitionKind::Data {
                Ok(PartitionGC::new(self.io.clone(), self.disc_header.clone())?)
            } else {
                Err(Error::DiscFormat("GameCube discs only have a data partition".to_string()))
            }
        } else if let Some(part) = self.partitions.iter().find(|v| v.kind == part_type) {
            Ok(PartitionWii::new(self.io.clone(), self.disc_header.clone(), part, options)?)
        } else {
            Err(Error::DiscFormat(format!("Partition type {part_type} not found")))
        }
    }
}

impl Read for DiscReader {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let block_idx = (self.pos / self.block_buf.len() as u64) as u32;
        let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;

        let partition = if self.disc_header.is_wii() {
            self.partitions.iter().find(|part| {
                abs_sector >= part.data_start_sector && abs_sector < part.data_end_sector
            })
        } else {
            None
        };

        // Read new block
        if block_idx != self.block_idx {
            self.block = self.io.read_block(self.block_buf.as_mut(), block_idx, partition)?;
            self.block_idx = block_idx;
        }

        // Read new sector into buffer
        if abs_sector != self.sector_idx {
            if let Some(partition) = partition {
                match self.mode {
                    EncryptionMode::Decrypted => self.block.decrypt(
                        self.sector_buf.as_mut(),
                        self.block_buf.as_ref(),
                        block_idx,
                        abs_sector,
                        partition,
                    )?,
                    EncryptionMode::Encrypted => self.block.encrypt(
                        self.sector_buf.as_mut(),
                        self.block_buf.as_ref(),
                        block_idx,
                        abs_sector,
                        partition,
                    )?,
                }
            } else {
                self.block.copy_raw(
                    self.sector_buf.as_mut(),
                    self.block_buf.as_ref(),
                    block_idx,
                    abs_sector,
                    &self.disc_header,
                )?;
            }
            self.sector_idx = abs_sector;
        }

        // Read from sector buffer
        let offset = (self.pos % SECTOR_SIZE as u64) as usize;
        let len = min(buf.len(), SECTOR_SIZE - offset);
        buf[..len].copy_from_slice(&self.sector_buf[offset..offset + len]);
        self.pos += len as u64;
        Ok(len)
    }
}

impl Seek for DiscReader {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        self.pos = match pos {
            SeekFrom::Start(v) => v,
            SeekFrom::End(_) => {
                return Err(io::Error::new(
                    io::ErrorKind::Unsupported,
                    "BlockIOReader: SeekFrom::End is not supported".to_string(),
                ));
            }
            SeekFrom::Current(v) => self.pos.saturating_add_signed(v),
        };
        Ok(self.pos)
    }
}

fn read_partition_info(reader: &mut DiscReader) -> crate::Result<Vec<PartitionInfo>> {
    reader.seek(SeekFrom::Start(WII_PART_GROUP_OFF)).context("Seeking to partition groups")?;
    let part_groups: [WiiPartGroup; 4] = read_from(reader).context("Reading partition groups")?;
    let mut part_info = Vec::new();
    for (group_idx, group) in part_groups.iter().enumerate() {
        let part_count = group.part_count.get();
        if part_count == 0 {
            continue;
        }
        reader
            .seek(SeekFrom::Start(group.part_entry_off()))
            .with_context(|| format!("Seeking to partition group {group_idx}"))?;
        let entries: Vec<WiiPartEntry> = read_vec(reader, part_count as usize)
            .with_context(|| format!("Reading partition group {group_idx}"))?;
        for (part_idx, entry) in entries.iter().enumerate() {
            let offset = entry.offset();
            reader
                .seek(SeekFrom::Start(offset))
                .with_context(|| format!("Seeking to partition data {group_idx}:{part_idx}"))?;
            let header: Box<WiiPartitionHeader> = read_box(reader)
                .with_context(|| format!("Reading partition header {group_idx}:{part_idx}"))?;

            let key = header.ticket.decrypt_title_key()?;
            let start_offset = entry.offset();
            if start_offset % SECTOR_SIZE as u64 != 0 {
                return Err(Error::DiscFormat(format!(
                    "Partition {group_idx}:{part_idx} offset is not sector aligned",
                )));
            }
            let data_start_offset = entry.offset() + header.data_off();
            let data_end_offset = data_start_offset + header.data_size();
            if data_start_offset % SECTOR_SIZE as u64 != 0
                || data_end_offset % SECTOR_SIZE as u64 != 0
            {
                return Err(Error::DiscFormat(format!(
                    "Partition {group_idx}:{part_idx} data is not sector aligned",
                )));
            }
            let mut info = PartitionInfo {
                index: part_info.len(),
                kind: entry.kind.get().into(),
                start_sector: (start_offset / SECTOR_SIZE as u64) as u32,
                data_start_sector: (data_start_offset / SECTOR_SIZE as u64) as u32,
                data_end_sector: (data_end_offset / SECTOR_SIZE as u64) as u32,
                key,
                header,
                disc_header: DiscHeader::new_box_zeroed(),
                partition_header: PartitionHeader::new_box_zeroed(),
                hash_table: None,
            };

            let mut partition_reader = PartitionWii::new(
                reader.io.clone(),
                reader.disc_header.clone(),
                &info,
                &OpenOptions::default(),
            )?;
            info.disc_header = read_box(&mut partition_reader).context("Reading disc header")?;
            info.partition_header =
                read_box(&mut partition_reader).context("Reading partition header")?;

            part_info.push(info);
        }
    }
    Ok(part_info)
}

fn guess_disc_size(part_info: &[PartitionInfo]) -> u64 {
    let max_offset = part_info
        .iter()
        .flat_map(|v| {
            let offset = v.start_sector as u64 * SECTOR_SIZE as u64;
            [
                offset + v.header.tmd_off() + v.header.tmd_size(),
                offset + v.header.cert_chain_off() + v.header.cert_chain_size(),
                offset + v.header.h3_table_off() + v.header.h3_table_size(),
                offset + v.header.data_off() + v.header.data_size(),
            ]
        })
        .max()
        .unwrap_or(0x50000);
    if max_offset <= MINI_DVD_SIZE && !part_info.iter().any(|v| v.kind == PartitionKind::Data) {
        // Datel disc
        MINI_DVD_SIZE
    } else if max_offset < SL_DVD_SIZE {
        SL_DVD_SIZE
    } else {
        DL_DVD_SIZE
    }
}