From 35388db4589e046e2365da8f849e866abb78001d Mon Sep 17 00:00:00 2001 From: Luke Street Date: Wed, 1 Sep 2021 16:10:22 -0400 Subject: [PATCH] Cleanup & add as_dyn to ReadStream --- src/bin.rs | 4 ++-- src/disc/gcn.rs | 6 +++++- src/disc/mod.rs | 6 +++--- src/disc/wii.rs | 12 ++++++++---- src/io/iso.rs | 22 +++++++++++----------- src/io/mod.rs | 30 +++++++++++++++--------------- src/io/nfs.rs | 14 +++++++++----- src/streams.rs | 26 +++++++++++++++++++++++--- 8 files changed, 76 insertions(+), 44 deletions(-) diff --git a/src/bin.rs b/src/bin.rs index f547a47..c4aee1f 100644 --- a/src/bin.rs +++ b/src/bin.rs @@ -7,9 +7,9 @@ use std::{ use clap::{clap_app, AppSettings}; use file_size; use nod::{ - disc::{new_disc_base, PartReadStream}, + disc::{new_disc_base, DiscBase, PartReadStream}, fst::NodeType, - io::{has_extension, new_disc_io}, + io::{has_extension, new_disc_io, DiscIO}, Result, }; diff --git a/src/disc/gcn.rs b/src/disc/gcn.rs index 15bfd0e..a9c7f3f 100644 --- a/src/disc/gcn.rs +++ b/src/disc/gcn.rs @@ -17,7 +17,9 @@ pub(crate) struct DiscGCN { pub(crate) header: Header, } -pub(crate) fn new_disc_gcn(header: Header) -> Result { Result::Ok(DiscGCN { header }) } +impl DiscGCN { + pub(crate) fn new(header: Header) -> Result { Result::Ok(DiscGCN { header }) } +} impl DiscBase for DiscGCN { fn get_header(&self) -> &Header { &self.header } @@ -92,6 +94,8 @@ impl<'a> Seek for GCPartReadStream<'a> { impl<'a> ReadStream for GCPartReadStream<'a> { fn stable_stream_len(&mut self) -> io::Result { self.stream.stable_stream_len() } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } impl<'a> PartReadStream for GCPartReadStream<'a> { diff --git a/src/disc/mod.rs b/src/disc/mod.rs index 0412293..287d768 100644 --- a/src/disc/mod.rs +++ b/src/disc/mod.rs @@ -5,7 +5,7 @@ use std::{fmt::Debug, io}; use binread::{prelude::*, BinReaderExt, NullString}; use crate::{ - disc::{gcn::new_disc_gcn, wii::new_disc_wii}, + disc::{gcn::DiscGCN, wii::DiscWii}, fst::{Node, NodeType}, io::DiscIO, streams::{ReadStream, SharedWindowedReadStream}, @@ -113,9 +113,9 @@ pub fn new_disc_base(disc_io: &mut dyn DiscIO) -> Result> { let mut stream = disc_io.begin_read_stream(0)?; let header: Header = stream.read_be()?; if header.wii_magic == 0x5D1C9EA3 { - Result::Ok(Box::from(new_disc_wii(stream.as_mut(), header)?)) + Result::Ok(Box::from(DiscWii::new(stream.as_mut(), header)?)) } else if header.gcn_magic == 0xC2339F3D { - Result::Ok(Box::from(new_disc_gcn(header)?)) + Result::Ok(Box::from(DiscGCN::new(header)?)) } else { Result::Err(Error::DiscFormat("Invalid GC/Wii magic".to_string())) } diff --git a/src/disc/wii.rs b/src/disc/wii.rs index 63e2770..49a1d88 100644 --- a/src/disc/wii.rs +++ b/src/disc/wii.rs @@ -192,10 +192,12 @@ pub(crate) struct DiscWii { part_info: WiiPartInfo, } -pub(crate) fn new_disc_wii(mut stream: &mut dyn ReadStream, header: Header) -> Result { - let mut disc = DiscWii { header, part_info: stream.read_be()? }; - disc.decrypt_partition_keys()?; - Result::Ok(disc) +impl DiscWii { + pub(crate) fn new(mut stream: &mut dyn ReadStream, header: Header) -> Result { + let mut disc = DiscWii { header, part_info: stream.read_be()? }; + disc.decrypt_partition_keys()?; + Result::Ok(disc) + } } impl DiscWii { @@ -401,6 +403,8 @@ impl<'a> ReadStream for WiiPartReadStream<'a> { fn stable_stream_len(&mut self) -> io::Result { io::Result::Ok(to_block_size(self.stream.stable_stream_len()?)) } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } #[derive(Clone, Debug, PartialEq, BinRead)] diff --git a/src/io/iso.rs b/src/io/iso.rs index 43417d7..62723cc 100644 --- a/src/io/iso.rs +++ b/src/io/iso.rs @@ -1,20 +1,20 @@ -use std::fs::File; -use std::io::{Seek, SeekFrom}; -use std::io; -use std::path::{Path, PathBuf}; +use std::{ + fs::File, + io, + io::{Seek, SeekFrom}, + path::{Path, PathBuf}, +}; -use crate::io::DiscIO; -use crate::Result; -use crate::streams::ReadStream; +use crate::{io::DiscIO, streams::ReadStream, Result}; pub(crate) struct DiscIOISO { pub(crate) filename: PathBuf, } -pub(crate) fn new_disc_io_iso(filename: &Path) -> Result { - Result::Ok(DiscIOISO { - filename: filename.to_owned(), - }) +impl DiscIOISO { + pub(crate) fn new(filename: &Path) -> Result { + Result::Ok(DiscIOISO { filename: filename.to_owned() }) + } } impl DiscIO for DiscIOISO { diff --git a/src/io/mod.rs b/src/io/mod.rs index ce479f6..83d8646 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -3,7 +3,7 @@ use std::{fs, io, path::Path}; use crate::{ - io::{iso::new_disc_io_iso, nfs::new_disc_io_nfs}, + io::{iso::DiscIOISO, nfs::DiscIONFS}, streams::ReadStream, Error, Result, }; @@ -21,18 +21,6 @@ pub trait DiscIO { fn has_wii_crypto(&self) -> bool { true } } -/// Helper function for checking a file extension. -#[inline(always)] -pub fn has_extension(filename: &Path, extension: &str) -> bool { - if let Some(ext) = filename.extension() { - // TODO use with Rust 1.53+ - // ext.eq_ignore_ascii_case(extension) - ext.to_str().unwrap_or("").eq_ignore_ascii_case(extension) - } else { - false - } -} - /// Creates a new [`DiscIO`] instance. /// /// # Examples @@ -66,10 +54,10 @@ pub fn new_disc_io(filename: &Path) -> Result> { ))); } if has_extension(path, "iso") { - Result::Ok(Box::from(new_disc_io_iso(path)?)) + Result::Ok(Box::from(DiscIOISO::new(path)?)) } else if has_extension(path, "nfs") { if matches!(path.parent(), Some(parent) if parent.is_dir()) { - Result::Ok(Box::from(new_disc_io_nfs(path.parent().unwrap())?)) + Result::Ok(Box::from(DiscIONFS::new(path.parent().unwrap())?)) } else { Result::Err(Error::DiscFormat("Failed to locate NFS parent directory".to_string())) } @@ -77,3 +65,15 @@ pub fn new_disc_io(filename: &Path) -> Result> { Result::Err(Error::DiscFormat("Unknown file type".to_string())) } } + +/// Helper function for checking a file extension. +#[inline(always)] +pub fn has_extension(filename: &Path, extension: &str) -> bool { + if let Some(ext) = filename.extension() { + // TODO use with Rust 1.53+ + // ext.eq_ignore_ascii_case(extension) + ext.to_str().unwrap_or("").eq_ignore_ascii_case(extension) + } else { + false + } +} diff --git a/src/io/nfs.rs b/src/io/nfs.rs index e1419da..17f962a 100644 --- a/src/io/nfs.rs +++ b/src/io/nfs.rs @@ -83,11 +83,13 @@ pub(crate) struct DiscIONFS { pub(crate) header: Option, } -pub(crate) fn new_disc_io_nfs(directory: &Path) -> Result { - let mut disc_io = - DiscIONFS { directory: directory.to_owned(), key: [0; 16], header: Option::None }; - disc_io.validate_files()?; - Result::Ok(disc_io) +impl DiscIONFS { + pub(crate) fn new(directory: &Path) -> Result { + let mut disc_io = + DiscIONFS { directory: directory.to_owned(), key: [0; 16], header: Option::None }; + disc_io.validate_files()?; + Result::Ok(disc_io) + } } pub(crate) struct NFSReadStream<'a> { @@ -221,6 +223,8 @@ impl<'a> Seek for NFSReadStream<'a> { impl<'a> ReadStream for NFSReadStream<'a> { fn stable_stream_len(&mut self) -> io::Result { todo!() } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } impl DiscIO for DiscIONFS { diff --git a/src/streams.rs b/src/streams.rs index a3f4345..3a666dd 100644 --- a/src/streams.rs +++ b/src/streams.rs @@ -40,11 +40,16 @@ pub trait ReadStream: Read + Seek { /// Creates a windowed read sub-stream with offset and size. /// /// Seeks underlying stream immediately. - fn new_window(&mut self, offset: u64, size: u64) -> io::Result - where Self: Sized { + fn new_window(&mut self, offset: u64, size: u64) -> io::Result { self.seek(SeekFrom::Start(offset))?; - io::Result::Ok(SharedWindowedReadStream { base: self, begin: offset, end: offset + size }) + io::Result::Ok(SharedWindowedReadStream { + base: self.as_dyn(), + begin: offset, + end: offset + size, + }) } + + fn as_dyn(&mut self) -> &mut dyn ReadStream; } impl ReadStream for File { @@ -55,6 +60,8 @@ impl ReadStream for File { self.seek(SeekFrom::Start(before))?; result } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } trait WindowedReadStream: ReadStream { @@ -84,6 +91,15 @@ pub struct SharedWindowedReadStream<'a> { pub end: u64, } +impl<'a> SharedWindowedReadStream<'a> { + pub fn set_window(&mut self, begin: u64, end: u64) -> io::Result<()> { + self.base.seek(SeekFrom::Start(begin))?; + self.begin = begin; + self.end = end; + io::Result::Ok(()) + } +} + #[inline(always)] fn windowed_read(stream: &mut dyn WindowedReadStream, buf: &mut [u8]) -> io::Result { let pos = stream.stream_position()?; @@ -124,6 +140,8 @@ impl<'a> Seek for OwningWindowedReadStream<'a> { impl<'a> ReadStream for OwningWindowedReadStream<'a> { fn stable_stream_len(&mut self) -> io::Result { Result::Ok(self.end - self.begin) } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } impl<'a> WindowedReadStream for OwningWindowedReadStream<'a> { @@ -146,6 +164,8 @@ impl<'a> Seek for SharedWindowedReadStream<'a> { impl<'a> ReadStream for SharedWindowedReadStream<'a> { fn stable_stream_len(&mut self) -> io::Result { Result::Ok(self.end - self.begin) } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } } impl<'a> WindowedReadStream for SharedWindowedReadStream<'a> {