From 370d03fa9a95777b3b8ee704ce48437e47139d85 Mon Sep 17 00:00:00 2001 From: Luke Street Date: Wed, 2 Oct 2024 21:50:59 -0600 Subject: [PATCH] Add Disc::new_stream/new_stream_with_options Allows opening a disc image from a custom stream, rather than a filesystem path. --- nod/src/disc/fst.rs | 6 ++-- nod/src/disc/gcn.rs | 11 +++++++ nod/src/disc/mod.rs | 31 +++++++++++++++++- nod/src/disc/streams.rs | 41 +++++++++++++++--------- nod/src/disc/wii.rs | 11 +++++++ nod/src/io/block.rs | 69 +++++++++++++++++++++++++++++++---------- nod/src/io/ciso.rs | 22 ++++++------- nod/src/io/gcz.rs | 20 +++++------- nod/src/io/iso.rs | 25 +++++++-------- nod/src/io/tgc.rs | 27 ++++++++-------- nod/src/io/wbfs.rs | 35 +++++++-------------- nod/src/io/wia.rs | 23 ++++++-------- nod/src/lib.rs | 28 ++++++++++++++--- 13 files changed, 222 insertions(+), 127 deletions(-) diff --git a/nod/src/disc/fst.rs b/nod/src/disc/fst.rs index 243dc50..b08e69d 100644 --- a/nod/src/disc/fst.rs +++ b/nod/src/disc/fst.rs @@ -61,7 +61,7 @@ impl Node { /// For directories, this is the parent node index in the FST. #[inline] pub fn offset(&self, is_wii: bool) -> u64 { - if is_wii && self.kind == 0 { + if is_wii && self.is_file() { self.offset.get() as u64 * 4 } else { self.offset.get() as u64 @@ -108,7 +108,7 @@ impl<'a> Fst<'a> { /// Get the name of a node. #[allow(clippy::missing_inline_in_public_items)] - pub fn get_name(&self, node: &Node) -> Result, String> { + pub fn get_name(&self, node: &Node) -> Result, String> { let name_buf = self.string_table.get(node.name_offset() as usize..).ok_or_else(|| { format!( "FST: name offset {} out of bounds (string table size: {})", @@ -128,7 +128,7 @@ impl<'a> Fst<'a> { /// Finds a particular file or directory by path. #[allow(clippy::missing_inline_in_public_items)] - pub fn find(&self, path: &str) -> Option<(usize, &Node)> { + pub fn find(&self, path: &str) -> Option<(usize, &'a Node)> { let mut split = path.trim_matches('/').split('/'); let mut current = split.next()?; let mut idx = 1; diff --git a/nod/src/disc/gcn.rs b/nod/src/disc/gcn.rs index 0a495f7..a52b7b7 100644 --- a/nod/src/disc/gcn.rs +++ b/nod/src/disc/gcn.rs @@ -11,6 +11,7 @@ use super::{ PartitionMeta, BI2_SIZE, BOOT_SIZE, SECTOR_SIZE, }; use crate::{ + disc::streams::OwnedFileStream, io::block::{Block, BlockIO}, util::read::{read_box, read_box_slice, read_vec}, Result, ResultContext, @@ -132,6 +133,16 @@ impl PartitionBase for PartitionGC { } FileStream::new(self, node.offset(false), node.length()) } + + fn into_open_file(self: Box, node: &Node) -> io::Result { + if !node.is_file() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Node is not a file".to_string(), + )); + } + OwnedFileStream::new(self, node.offset(false), node.length()) + } } pub(crate) fn read_part_meta( diff --git a/nod/src/disc/mod.rs b/nod/src/disc/mod.rs index b800cea..10064c4 100644 --- a/nod/src/disc/mod.rs +++ b/nod/src/disc/mod.rs @@ -23,7 +23,7 @@ pub(crate) mod streams; pub(crate) mod wii; pub use fst::{Fst, Node, NodeKind}; -pub use streams::FileStream; +pub use streams::{FileStream, OwnedFileStream, WindowedStream}; pub use wii::{SignedHeader, Ticket, TicketLimit, TmdHeader}; /// Size in bytes of a disc sector. @@ -308,6 +308,35 @@ pub trait PartitionBase: DynClone + BufRead + Seek + Send + Sync { /// } /// ``` fn open_file(&mut self, node: &Node) -> io::Result; + + /// Consumes the partition instance and returns a windowed stream. + /// + /// # Examples + /// + /// ```no_run + /// use std::io::Read; + /// + /// use nod::{Disc, PartitionKind, OwnedFileStream}; + /// + /// fn main() -> nod::Result<()> { + /// let disc = Disc::new("path/to/file.iso")?; + /// let mut partition = disc.open_partition_kind(PartitionKind::Data)?; + /// let meta = partition.meta()?; + /// let fst = meta.fst()?; + /// if let Some((_, node)) = fst.find("/disc.tgc") { + /// let file: OwnedFileStream = partition + /// .clone() // Clone the Box + /// .into_open_file(node) // Get an OwnedFileStream + /// .expect("Failed to open file stream"); + /// // Open the inner disc image using the owned stream + /// let inner_disc = Disc::new_stream(Box::new(file)) + /// .expect("Failed to open inner disc"); + /// // ... + /// } + /// Ok(()) + /// } + /// ``` + fn into_open_file(self: Box, node: &Node) -> io::Result; } dyn_clone::clone_trait_object!(PartitionBase); diff --git a/nod/src/disc/streams.rs b/nod/src/disc/streams.rs index 0af2461..cbee2fb 100644 --- a/nod/src/disc/streams.rs +++ b/nod/src/disc/streams.rs @@ -7,30 +7,39 @@ use std::{ use super::PartitionBase; -/// A file read stream for a [`PartitionBase`]. -pub struct FileStream<'a> { - base: &'a mut dyn PartitionBase, +/// A file read stream borrowing a [`PartitionBase`]. +pub type FileStream<'a> = WindowedStream<&'a mut dyn PartitionBase>; + +/// A file read stream owning a [`PartitionBase`]. +pub type OwnedFileStream = WindowedStream>; + +/// A read stream with a fixed window. +#[derive(Clone)] +pub struct WindowedStream +where T: BufRead + Seek +{ + base: T, pos: u64, begin: u64, end: u64, } -impl FileStream<'_> { - /// Creates a new file stream with offset and size. +impl WindowedStream +where T: BufRead + Seek +{ + /// Creates a new windowed stream with offset and size. /// /// Seeks underlying stream immediately. #[inline] - pub(crate) fn new( - base: &mut dyn PartitionBase, - offset: u64, - size: u64, - ) -> io::Result { + pub fn new(mut base: T, offset: u64, size: u64) -> io::Result { base.seek(SeekFrom::Start(offset))?; - Ok(FileStream { base, pos: offset, begin: offset, end: offset + size }) + Ok(Self { base, pos: offset, begin: offset, end: offset + size }) } } -impl<'a> Read for FileStream<'a> { +impl Read for WindowedStream +where T: BufRead + Seek +{ #[inline] fn read(&mut self, out: &mut [u8]) -> io::Result { let buf = self.fill_buf()?; @@ -41,7 +50,9 @@ impl<'a> Read for FileStream<'a> { } } -impl<'a> BufRead for FileStream<'a> { +impl BufRead for WindowedStream +where T: BufRead + Seek +{ #[inline] fn fill_buf(&mut self) -> io::Result<&[u8]> { let limit = self.end.saturating_sub(self.pos); @@ -60,7 +71,9 @@ impl<'a> BufRead for FileStream<'a> { } } -impl<'a> Seek for FileStream<'a> { +impl Seek for WindowedStream +where T: BufRead + Seek +{ #[inline] fn seek(&mut self, pos: SeekFrom) -> io::Result { let mut pos = match pos { diff --git a/nod/src/disc/wii.rs b/nod/src/disc/wii.rs index 87b1794..586fd4e 100644 --- a/nod/src/disc/wii.rs +++ b/nod/src/disc/wii.rs @@ -14,6 +14,7 @@ use super::{ }; use crate::{ array_ref, + disc::streams::OwnedFileStream, io::{ aes_decrypt, block::{Block, BlockIO, PartitionInfo}, @@ -497,4 +498,14 @@ impl PartitionBase for PartitionWii { } FileStream::new(self, node.offset(true), node.length()) } + + fn into_open_file(self: Box, node: &Node) -> io::Result { + if !node.is_file() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Node is not a file".to_string(), + )); + } + OwnedFileStream::new(self, node.offset(true), node.length()) + } } diff --git a/nod/src/io/block.rs b/nod/src/io/block.rs index 98bdbba..794f353 100644 --- a/nod/src/io/block.rs +++ b/nod/src/io/block.rs @@ -1,4 +1,9 @@ -use std::{cmp::min, fs, fs::File, io, path::Path}; +use std::{ + cmp::min, + fs, io, + io::{Read, Seek}, + path::Path, +}; use dyn_clone::DynClone; use zerocopy::transmute_ref; @@ -10,11 +15,18 @@ use crate::{ wii::{WiiPartitionHeader, HASHES_SIZE, SECTOR_DATA_SIZE}, SECTOR_SIZE, }, - io::{aes_decrypt, aes_encrypt, KeyBytes, MagicBytes}, + io::{aes_decrypt, aes_encrypt, split::SplitFileReader, KeyBytes, MagicBytes}, util::{lfg::LaggedFibonacci, read::read_from}, DiscHeader, DiscMeta, Error, PartitionHeader, PartitionKind, Result, ResultContext, }; +/// Required trait bounds for reading disc images. +pub trait DiscStream: Read + Seek + DynClone + Send + Sync {} + +impl DiscStream for T where T: Read + Seek + DynClone + Send + Sync + ?Sized {} + +dyn_clone::clone_trait_object!(DiscStream); + /// Block I/O trait for reading disc images. pub trait BlockIO: DynClone + Send + Sync { /// Reads a block from the disc image. @@ -78,7 +90,27 @@ pub trait BlockIO: DynClone + Send + Sync { dyn_clone::clone_trait_object!(BlockIO); -/// Creates a new [`BlockIO`] instance. +/// Creates a new [`BlockIO`] instance from a stream. +pub fn new(mut stream: Box) -> Result> { + let magic: MagicBytes = read_from(stream.as_mut()).context("Reading magic bytes")?; + stream.seek(io::SeekFrom::Start(0)).context("Seeking to start")?; + let io: Box = match magic { + crate::io::ciso::CISO_MAGIC => crate::io::ciso::DiscIOCISO::new(stream)?, + #[cfg(feature = "compress-zlib")] + crate::io::gcz::GCZ_MAGIC => crate::io::gcz::DiscIOGCZ::new(stream)?, + crate::io::nfs::NFS_MAGIC => todo!(), + crate::io::wbfs::WBFS_MAGIC => crate::io::wbfs::DiscIOWBFS::new(stream)?, + crate::io::wia::WIA_MAGIC | crate::io::wia::RVZ_MAGIC => { + crate::io::wia::DiscIOWIA::new(stream)? + } + crate::io::tgc::TGC_MAGIC => crate::io::tgc::DiscIOTGC::new(stream)?, + _ => crate::io::iso::DiscIOISO::new(stream)?, + }; + check_block_size(io.as_ref())?; + Ok(io) +} + +/// Creates a new [`BlockIO`] instance from a filesystem path. pub fn open(filename: &Path) -> Result> { let path_result = fs::canonicalize(filename); if let Err(err) = path_result { @@ -92,16 +124,16 @@ pub fn open(filename: &Path) -> Result> { if !meta.unwrap().is_file() { return Err(Error::DiscFormat(format!("Input is not a file: {}", filename.display()))); } - let magic: MagicBytes = { - let mut file = - File::open(path).with_context(|| format!("Opening file {}", filename.display()))?; - read_from(&mut file) - .with_context(|| format!("Reading magic bytes from {}", filename.display()))? - }; + let mut stream = Box::new(SplitFileReader::new(filename)?); + let magic: MagicBytes = read_from(stream.as_mut()) + .with_context(|| format!("Reading magic bytes from {}", filename.display()))?; + stream + .seek(io::SeekFrom::Start(0)) + .with_context(|| format!("Seeking to start of {}", filename.display()))?; let io: Box = match magic { - crate::io::ciso::CISO_MAGIC => crate::io::ciso::DiscIOCISO::new(path)?, + crate::io::ciso::CISO_MAGIC => crate::io::ciso::DiscIOCISO::new(stream)?, #[cfg(feature = "compress-zlib")] - crate::io::gcz::GCZ_MAGIC => crate::io::gcz::DiscIOGCZ::new(path)?, + crate::io::gcz::GCZ_MAGIC => crate::io::gcz::DiscIOGCZ::new(stream)?, crate::io::nfs::NFS_MAGIC => match path.parent() { Some(parent) if parent.is_dir() => { crate::io::nfs::DiscIONFS::new(path.parent().unwrap())? @@ -110,13 +142,18 @@ pub fn open(filename: &Path) -> Result> { return Err(Error::DiscFormat("Failed to locate NFS parent directory".to_string())); } }, - crate::io::wbfs::WBFS_MAGIC => crate::io::wbfs::DiscIOWBFS::new(path)?, + crate::io::wbfs::WBFS_MAGIC => crate::io::wbfs::DiscIOWBFS::new(stream)?, crate::io::wia::WIA_MAGIC | crate::io::wia::RVZ_MAGIC => { - crate::io::wia::DiscIOWIA::new(path)? + crate::io::wia::DiscIOWIA::new(stream)? } - crate::io::tgc::TGC_MAGIC => crate::io::tgc::DiscIOTGC::new(path)?, - _ => crate::io::iso::DiscIOISO::new(path)?, + crate::io::tgc::TGC_MAGIC => crate::io::tgc::DiscIOTGC::new(stream)?, + _ => crate::io::iso::DiscIOISO::new(stream)?, }; + check_block_size(io.as_ref())?; + Ok(io) +} + +fn check_block_size(io: &dyn BlockIO) -> Result<()> { if io.block_size_internal() < SECTOR_SIZE as u32 && SECTOR_SIZE as u32 % io.block_size_internal() != 0 { @@ -133,7 +170,7 @@ pub fn open(filename: &Path) -> Result> { SECTOR_SIZE ))); } - Ok(io) + Ok(()) } /// Wii partition information. diff --git a/nod/src/io/ciso.rs b/nod/src/io/ciso.rs index 3df10af..3f61465 100644 --- a/nod/src/io/ciso.rs +++ b/nod/src/io/ciso.rs @@ -2,7 +2,6 @@ use std::{ io, io::{Read, Seek, SeekFrom}, mem::size_of, - path::Path, }; use zerocopy::{little_endian::*, AsBytes, FromBytes, FromZeroes}; @@ -10,9 +9,8 @@ use zerocopy::{little_endian::*, AsBytes, FromBytes, FromZeroes}; use crate::{ disc::SECTOR_SIZE, io::{ - block::{Block, BlockIO, PartitionInfo}, + block::{Block, BlockIO, DiscStream, PartitionInfo}, nkit::NKitHeader, - split::SplitFileReader, Format, MagicBytes, }, static_assert, @@ -36,18 +34,16 @@ static_assert!(size_of::() == SECTOR_SIZE); #[derive(Clone)] pub struct DiscIOCISO { - inner: SplitFileReader, + inner: Box, header: CISOHeader, block_map: [u16; CISO_MAP_SIZE], nkit_header: Option, } impl DiscIOCISO { - pub fn new(filename: &Path) -> Result> { - let mut inner = SplitFileReader::new(filename)?; - + pub fn new(mut inner: Box) -> Result> { // Read header - let header: CISOHeader = read_from(&mut inner).context("Reading CISO header")?; + let header: CISOHeader = read_from(inner.as_mut()).context("Reading CISO header")?; if header.magic != CISO_MAGIC { return Err(Error::DiscFormat("Invalid CISO magic".to_string())); } @@ -64,18 +60,18 @@ impl DiscIOCISO { } } let file_size = SECTOR_SIZE as u64 + block as u64 * header.block_size.get() as u64; - if file_size > inner.len() { + let len = inner.seek(SeekFrom::End(0)).context("Determining stream length")?; + if file_size > len { return Err(Error::DiscFormat(format!( "CISO file size mismatch: expected at least {} bytes, got {}", - file_size, - inner.len() + file_size, len ))); } // Read NKit header if present (after CISO data) - let nkit_header = if inner.len() > file_size + 4 { + let nkit_header = if len > file_size + 4 { inner.seek(SeekFrom::Start(file_size)).context("Seeking to NKit header")?; - NKitHeader::try_read_from(&mut inner, header.block_size.get(), true) + NKitHeader::try_read_from(inner.as_mut(), header.block_size.get(), true) } else { None }; diff --git a/nod/src/io/gcz.rs b/nod/src/io/gcz.rs index adf42ee..5230d0d 100644 --- a/nod/src/io/gcz.rs +++ b/nod/src/io/gcz.rs @@ -2,7 +2,6 @@ use std::{ io, io::{Read, Seek, SeekFrom}, mem::size_of, - path::Path, }; use adler::adler32_slice; @@ -12,8 +11,7 @@ use zstd::zstd_safe::WriteBuf; use crate::{ io::{ - block::{Block, BlockIO}, - split::SplitFileReader, + block::{Block, BlockIO, DiscStream}, MagicBytes, }, static_assert, @@ -38,7 +36,7 @@ struct GCZHeader { static_assert!(size_of::() == 32); pub struct DiscIOGCZ { - inner: SplitFileReader, + inner: Box, header: GCZHeader, block_map: Box<[U64]>, block_hashes: Box<[U32]>, @@ -60,21 +58,19 @@ impl Clone for DiscIOGCZ { } impl DiscIOGCZ { - pub fn new(filename: &Path) -> Result> { - let mut inner = SplitFileReader::new(filename)?; - + pub fn new(mut inner: Box) -> Result> { // Read header - let header: GCZHeader = read_from(&mut inner).context("Reading GCZ header")?; + let header: GCZHeader = read_from(inner.as_mut()).context("Reading GCZ header")?; if header.magic != GCZ_MAGIC { return Err(Error::DiscFormat("Invalid GCZ magic".to_string())); } // Read block map and hashes let block_count = header.block_count.get(); - let block_map = - read_box_slice(&mut inner, block_count as usize).context("Reading GCZ block map")?; - let block_hashes = - read_box_slice(&mut inner, block_count as usize).context("Reading GCZ block hashes")?; + let block_map = read_box_slice(inner.as_mut(), block_count as usize) + .context("Reading GCZ block map")?; + let block_hashes = read_box_slice(inner.as_mut(), block_count as usize) + .context("Reading GCZ block hashes")?; // header + block_count * (u64 + u32) let data_offset = size_of::() as u64 + block_count as u64 * 12; diff --git a/nod/src/io/iso.rs b/nod/src/io/iso.rs index 167a816..484b85b 100644 --- a/nod/src/io/iso.rs +++ b/nod/src/io/iso.rs @@ -1,28 +1,28 @@ use std::{ io, io::{Read, Seek, SeekFrom}, - path::Path, }; use crate::{ disc::SECTOR_SIZE, io::{ - block::{Block, BlockIO, PartitionInfo}, - split::SplitFileReader, + block::{Block, BlockIO, DiscStream, PartitionInfo}, Format, }, - DiscMeta, Result, + DiscMeta, Result, ResultContext, }; #[derive(Clone)] pub struct DiscIOISO { - inner: SplitFileReader, + inner: Box, + stream_len: u64, } impl DiscIOISO { - pub fn new(filename: &Path) -> Result> { - let inner = SplitFileReader::new(filename)?; - Ok(Box::new(Self { inner })) + pub fn new(mut inner: Box) -> Result> { + let stream_len = inner.seek(SeekFrom::End(0)).context("Determining stream length")?; + inner.seek(SeekFrom::Start(0)).context("Seeking to start")?; + Ok(Box::new(Self { inner, stream_len })) } } @@ -34,16 +34,15 @@ impl BlockIO for DiscIOISO { _partition: Option<&PartitionInfo>, ) -> io::Result { let offset = block as u64 * SECTOR_SIZE as u64; - let total_size = self.inner.len(); - if offset >= total_size { + if offset >= self.stream_len { // End of file return Ok(Block::Zero); } self.inner.seek(SeekFrom::Start(offset))?; - if offset + SECTOR_SIZE as u64 > total_size { + if offset + SECTOR_SIZE as u64 > self.stream_len { // If the last block is not a full sector, fill the rest with zeroes - let read = (total_size - offset) as usize; + let read = (self.stream_len - offset) as usize; self.inner.read_exact(&mut out[..read])?; out[read..].fill(0); } else { @@ -58,7 +57,7 @@ impl BlockIO for DiscIOISO { DiscMeta { format: Format::Iso, lossless: true, - disc_size: Some(self.inner.len()), + disc_size: Some(self.stream_len), ..Default::default() } } diff --git a/nod/src/io/tgc.rs b/nod/src/io/tgc.rs index 685b5f8..310fd52 100644 --- a/nod/src/io/tgc.rs +++ b/nod/src/io/tgc.rs @@ -2,7 +2,6 @@ use std::{ io, io::{Read, Seek, SeekFrom}, mem::size_of, - path::Path, }; use zerocopy::{big_endian::U32, AsBytes, FromBytes, FromZeroes}; @@ -10,8 +9,7 @@ use zerocopy::{big_endian::U32, AsBytes, FromBytes, FromZeroes}; use crate::{ disc::SECTOR_SIZE, io::{ - block::{Block, BlockIO, PartitionInfo}, - split::SplitFileReader, + block::{Block, BlockIO, DiscStream, PartitionInfo}, Format, MagicBytes, }, util::read::{read_box_slice, read_from}, @@ -56,17 +54,19 @@ struct TGCHeader { #[derive(Clone)] pub struct DiscIOTGC { - inner: SplitFileReader, + inner: Box, + stream_len: u64, header: TGCHeader, fst: Box<[u8]>, } impl DiscIOTGC { - pub fn new(filename: &Path) -> Result> { - let mut inner = SplitFileReader::new(filename)?; + pub fn new(mut inner: Box) -> Result> { + let stream_len = inner.seek(SeekFrom::End(0)).context("Determining stream length")?; + inner.seek(SeekFrom::Start(0)).context("Seeking to start")?; // Read header - let header: TGCHeader = read_from(&mut inner).context("Reading TGC header")?; + let header: TGCHeader = read_from(inner.as_mut()).context("Reading TGC header")?; if header.magic != TGC_MAGIC { return Err(Error::DiscFormat("Invalid TGC magic".to_string())); } @@ -75,7 +75,7 @@ impl DiscIOTGC { inner .seek(SeekFrom::Start(header.fst_offset.get() as u64)) .context("Seeking to TGC FST")?; - let mut fst = read_box_slice(&mut inner, header.fst_size.get() as usize) + let mut fst = read_box_slice(inner.as_mut(), header.fst_size.get() as usize) .context("Reading TGC FST")?; let root_node = Node::ref_from_prefix(&fst) .ok_or_else(|| Error::DiscFormat("Invalid TGC FST".to_string()))?; @@ -89,7 +89,7 @@ impl DiscIOTGC { } } - Ok(Box::new(Self { inner, header, fst })) + Ok(Box::new(Self { inner, stream_len, header, fst })) } } @@ -101,16 +101,15 @@ impl BlockIO for DiscIOTGC { _partition: Option<&PartitionInfo>, ) -> io::Result { let offset = self.header.header_offset.get() as u64 + block as u64 * SECTOR_SIZE as u64; - let total_size = self.inner.len(); - if offset >= total_size { + if offset >= self.stream_len { // End of file return Ok(Block::Zero); } self.inner.seek(SeekFrom::Start(offset))?; - if offset + SECTOR_SIZE as u64 > total_size { + if offset + SECTOR_SIZE as u64 > self.stream_len { // If the last block is not a full sector, fill the rest with zeroes - let read = (total_size - offset) as usize; + let read = (self.stream_len - offset) as usize; self.inner.read_exact(&mut out[..read])?; out[read..].fill(0); } else { @@ -149,7 +148,7 @@ impl BlockIO for DiscIOTGC { DiscMeta { format: Format::Tgc, lossless: true, - disc_size: Some(self.inner.len() - self.header.header_offset.get() as u64), + disc_size: Some(self.stream_len - self.header.header_offset.get() as u64), ..Default::default() } } diff --git a/nod/src/io/wbfs.rs b/nod/src/io/wbfs.rs index 11f584f..8749225 100644 --- a/nod/src/io/wbfs.rs +++ b/nod/src/io/wbfs.rs @@ -2,16 +2,14 @@ use std::{ io, io::{Read, Seek, SeekFrom}, mem::size_of, - path::Path, }; use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes}; use crate::{ io::{ - block::{Block, BlockIO, PartitionInfo}, + block::{Block, BlockIO, DiscStream, PartitionInfo}, nkit::NKitHeader, - split::SplitFileReader, DiscMeta, Format, MagicBytes, }, util::read::{read_box_slice, read_from}, @@ -35,18 +33,6 @@ impl WBFSHeader { fn block_size(&self) -> u32 { 1 << self.block_size_shift } - // fn align_lba(&self, x: u32) -> u32 { (x + self.sector_size() - 1) & !(self.sector_size() - 1) } - // - // fn num_wii_sectors(&self) -> u32 { - // (self.num_sectors.get() / SECTOR_SIZE as u32) * self.sector_size() - // } - // - // fn max_wii_sectors(&self) -> u32 { NUM_WII_SECTORS } - // - // fn num_wbfs_sectors(&self) -> u32 { - // self.num_wii_sectors() >> (self.wbfs_sector_size_shift - 15) - // } - fn max_blocks(&self) -> u32 { NUM_WII_SECTORS >> (self.block_size_shift - 15) } } @@ -55,7 +41,7 @@ const NUM_WII_SECTORS: u32 = 143432 * 2; // Double layer discs #[derive(Clone)] pub struct DiscIOWBFS { - inner: SplitFileReader, + inner: Box, /// WBFS header header: WBFSHeader, /// Map of Wii LBAs to WBFS LBAs @@ -65,14 +51,12 @@ pub struct DiscIOWBFS { } impl DiscIOWBFS { - pub fn new(filename: &Path) -> Result> { - let mut inner = SplitFileReader::new(filename)?; - - let header: WBFSHeader = read_from(&mut inner).context("Reading WBFS header")?; + pub fn new(mut inner: Box) -> Result> { + let header: WBFSHeader = read_from(inner.as_mut()).context("Reading WBFS header")?; if header.magic != WBFS_MAGIC { return Err(Error::DiscFormat("Invalid WBFS magic".to_string())); } - let file_len = inner.len(); + let file_len = inner.seek(SeekFrom::End(0)).context("Determining stream length")?; let expected_file_len = header.num_sectors.get() as u64 * header.sector_size() as u64; if file_len != expected_file_len { return Err(Error::DiscFormat(format!( @@ -81,8 +65,11 @@ impl DiscIOWBFS { ))); } + inner + .seek(SeekFrom::Start(size_of::() as u64)) + .context("Seeking to WBFS disc table")?; let disc_table: Box<[u8]> = - read_box_slice(&mut inner, header.sector_size() as usize - size_of::()) + read_box_slice(inner.as_mut(), header.sector_size() as usize - size_of::()) .context("Reading WBFS disc table")?; if disc_table[0] != 1 { return Err(Error::DiscFormat("WBFS doesn't contain a disc".to_string())); @@ -95,12 +82,12 @@ impl DiscIOWBFS { inner .seek(SeekFrom::Start(header.sector_size() as u64 + DISC_HEADER_SIZE as u64)) .context("Seeking to WBFS LBA table")?; // Skip header - let block_map: Box<[U16]> = read_box_slice(&mut inner, header.max_blocks() as usize) + let block_map: Box<[U16]> = read_box_slice(inner.as_mut(), header.max_blocks() as usize) .context("Reading WBFS LBA table")?; // Read NKit header if present (always at 0x10000) inner.seek(SeekFrom::Start(0x10000)).context("Seeking to NKit header")?; - let nkit_header = NKitHeader::try_read_from(&mut inner, header.block_size(), true); + let nkit_header = NKitHeader::try_read_from(inner.as_mut(), header.block_size(), true); Ok(Box::new(Self { inner, header, block_map, nkit_header })) } diff --git a/nod/src/io/wia.rs b/nod/src/io/wia.rs index 99fd144..1bfa8e0 100644 --- a/nod/src/io/wia.rs +++ b/nod/src/io/wia.rs @@ -2,7 +2,6 @@ use std::{ io, io::{Read, Seek, SeekFrom}, mem::size_of, - path::Path, }; use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes}; @@ -14,9 +13,8 @@ use crate::{ SECTOR_SIZE, }, io::{ - block::{Block, BlockIO, PartitionInfo}, + block::{Block, BlockIO, DiscStream, PartitionInfo}, nkit::NKitHeader, - split::SplitFileReader, Compression, Format, HashBytes, KeyBytes, MagicBytes, }, static_assert, @@ -502,7 +500,7 @@ impl Decompressor { } pub struct DiscIOWIA { - inner: SplitFileReader, + inner: Box, header: WIAFileHeader, disc: WIADisc, partitions: Box<[WIAPartition]>, @@ -549,17 +547,16 @@ fn verify_hash(buf: &[u8], expected: &HashBytes) -> Result<()> { } impl DiscIOWIA { - pub fn new(filename: &Path) -> Result> { - let mut inner = SplitFileReader::new(filename)?; - + pub fn new(mut inner: Box) -> Result> { // Load & verify file header - let header: WIAFileHeader = read_from(&mut inner).context("Reading WIA/RVZ file header")?; + let header: WIAFileHeader = + read_from(inner.as_mut()).context("Reading WIA/RVZ file header")?; header.validate()?; let is_rvz = header.is_rvz(); // log::debug!("Header: {:?}", header); // Load & verify disc header - let mut disc_buf: Vec = read_vec(&mut inner, header.disc_size.get() as usize) + let mut disc_buf: Vec = read_vec(inner.as_mut(), header.disc_size.get() as usize) .context("Reading WIA/RVZ disc header")?; verify_hash(&disc_buf, &header.disc_hash)?; disc_buf.resize(size_of::(), 0); @@ -576,14 +573,14 @@ impl DiscIOWIA { // log::debug!("Disc: {:?}", disc); // Read NKit header if present (after disc header) - let nkit_header = NKitHeader::try_read_from(&mut inner, disc.chunk_size.get(), false); + let nkit_header = NKitHeader::try_read_from(inner.as_mut(), disc.chunk_size.get(), false); // Load & verify partition headers inner .seek(SeekFrom::Start(disc.partition_offset.get())) .context("Seeking to WIA/RVZ partition headers")?; let partitions: Box<[WIAPartition]> = - read_box_slice(&mut inner, disc.num_partitions.get() as usize) + read_box_slice(inner.as_mut(), disc.num_partitions.get() as usize) .context("Reading WIA/RVZ partition headers")?; verify_hash(partitions.as_ref().as_bytes(), &disc.partition_hash)?; // log::debug!("Partitions: {:?}", partitions); @@ -597,7 +594,7 @@ impl DiscIOWIA { .seek(SeekFrom::Start(disc.raw_data_offset.get())) .context("Seeking to WIA/RVZ raw data headers")?; let mut reader = decompressor - .wrap((&mut inner).take(disc.raw_data_size.get() as u64)) + .wrap(inner.as_mut().take(disc.raw_data_size.get() as u64)) .context("Creating WIA/RVZ decompressor")?; read_box_slice(&mut reader, disc.num_raw_data.get() as usize) .context("Reading WIA/RVZ raw data headers")? @@ -621,7 +618,7 @@ impl DiscIOWIA { .seek(SeekFrom::Start(disc.group_offset.get())) .context("Seeking to WIA/RVZ group headers")?; let mut reader = decompressor - .wrap((&mut inner).take(disc.group_size.get() as u64)) + .wrap(inner.as_mut().take(disc.group_size.get() as u64)) .context("Creating WIA/RVZ decompressor")?; if is_rvz { read_box_slice(&mut reader, disc.num_groups.get() as usize) diff --git a/nod/src/lib.rs b/nod/src/lib.rs index 1a85d13..0371afa 100644 --- a/nod/src/lib.rs +++ b/nod/src/lib.rs @@ -64,11 +64,14 @@ use std::{ }; pub use disc::{ - ApploaderHeader, DiscHeader, DolHeader, FileStream, Fst, Node, NodeKind, PartitionBase, - PartitionHeader, PartitionKind, PartitionMeta, SignedHeader, Ticket, TicketLimit, TmdHeader, - BI2_SIZE, BOOT_SIZE, SECTOR_SIZE, + ApploaderHeader, DiscHeader, DolHeader, FileStream, Fst, Node, NodeKind, OwnedFileStream, + PartitionBase, PartitionHeader, PartitionKind, PartitionMeta, SignedHeader, Ticket, + TicketLimit, TmdHeader, WindowedStream, BI2_SIZE, BOOT_SIZE, SECTOR_SIZE, +}; +pub use io::{ + block::{DiscStream, PartitionInfo}, + Compression, DiscMeta, Format, KeyBytes, }; -pub use io::{block::PartitionInfo, Compression, DiscMeta, Format, KeyBytes}; mod disc; mod io; @@ -170,6 +173,23 @@ impl Disc { Ok(Disc { reader, options: options.clone() }) } + /// Opens a disc image from a read stream. + #[inline] + pub fn new_stream(stream: Box) -> Result { + Disc::new_stream_with_options(stream, &OpenOptions::default()) + } + + /// Opens a disc image from a read stream with custom options. + #[inline] + pub fn new_stream_with_options( + stream: Box, + options: &OpenOptions, + ) -> Result { + let io = io::block::new(stream)?; + let reader = disc::reader::DiscReader::new(io, options)?; + Ok(Disc { reader, options: options.clone() }) + } + /// The disc's primary header. #[inline] pub fn header(&self) -> &DiscHeader { self.reader.header() }