diff --git a/src/disc/wii.rs b/src/disc/wii.rs index 49a1d88..0ab370f 100644 --- a/src/disc/wii.rs +++ b/src/disc/wii.rs @@ -230,13 +230,14 @@ impl DiscBase for DiscWii { .find(|v| v.part_type == WiiPartType::Data) .ok_or(Error::DiscFormat("Failed to locate data partition".to_string()))?; let data_off = part.part_header.data_off; + let has_crypto = disc_io.has_wii_crypto(); let result = Box::new(WiiPartReadStream { stream: wrap_windowed( disc_io.begin_read_stream(data_off)?, data_off, part.part_header.data_size, )?, - crypto: if disc_io.has_wii_crypto() { + crypto: if has_crypto { Aes128::new(&part.part_header.ticket.enc_key.into()).into() } else { Option::None diff --git a/src/io/iso.rs b/src/io/iso.rs index 62723cc..6a4a21a 100644 --- a/src/io/iso.rs +++ b/src/io/iso.rs @@ -5,7 +5,11 @@ use std::{ path::{Path, PathBuf}, }; -use crate::{io::DiscIO, streams::ReadStream, Result}; +use crate::{ + io::DiscIO, + streams::{ByteReadStream, ReadStream}, + Result, +}; pub(crate) struct DiscIOISO { pub(crate) filename: PathBuf, @@ -18,9 +22,28 @@ impl DiscIOISO { } impl DiscIO for DiscIOISO { - fn begin_read_stream(&self, offset: u64) -> io::Result> { + fn begin_read_stream(&mut self, offset: u64) -> io::Result> { let mut file = File::open(&*self.filename)?; file.seek(SeekFrom::Start(offset))?; io::Result::Ok(Box::from(file)) } } + +pub(crate) struct DiscIOISOStream { + pub(crate) stream: T, +} + +impl DiscIOISOStream { + pub(crate) fn new(stream: T) -> Result> { + Result::Ok(DiscIOISOStream { stream }) + } +} + +impl DiscIO for DiscIOISOStream { + fn begin_read_stream<'a>(&'a mut self, offset: u64) -> io::Result> { + let size = self.stream.stable_stream_len()?; + let mut stream = self.stream.new_window(0, size)?; + stream.seek(SeekFrom::Start(offset))?; + Ok(Box::from(stream)) + } +} diff --git a/src/io/mod.rs b/src/io/mod.rs index 83d8646..d534590 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -3,8 +3,11 @@ use std::{fs, io, path::Path}; use crate::{ - io::{iso::DiscIOISO, nfs::DiscIONFS}, - streams::ReadStream, + io::{ + iso::{DiscIOISO, DiscIOISOStream}, + nfs::DiscIONFS, + }, + streams::{ByteReadStream, ReadStream}, Error, Result, }; @@ -15,7 +18,7 @@ pub(crate) mod nfs; pub trait DiscIO { /// Opens a new read stream for the disc file(s). /// Generally does _not_ need to be used directly. - fn begin_read_stream(&self, offset: u64) -> io::Result>; + fn begin_read_stream(&mut self, offset: u64) -> io::Result>; /// If false, the file format does not use standard Wii partition encryption. (e.g. NFS) fn has_wii_crypto(&self) -> bool { true } @@ -66,6 +69,16 @@ pub fn new_disc_io(filename: &Path) -> Result> { } } +pub fn new_disc_io_from_buf(buf: &[u8]) -> Result> { + Ok(Box::from(DiscIOISOStream::new(ByteReadStream { bytes: buf, position: 0 })?)) +} + +pub fn new_disc_io_from_stream<'a, T: 'a + ReadStream + Sized>( + stream: T, +) -> Result> { + Ok(Box::from(DiscIOISOStream::new(stream)?)) +} + /// Helper function for checking a file extension. #[inline(always)] pub fn has_extension(filename: &Path, extension: &str) -> bool { diff --git a/src/io/nfs.rs b/src/io/nfs.rs index 5c1e705..22463c2 100644 --- a/src/io/nfs.rs +++ b/src/io/nfs.rs @@ -228,7 +228,7 @@ impl<'a> ReadStream for NFSReadStream<'a> { } impl DiscIO for DiscIONFS { - fn begin_read_stream(&self, offset: u64) -> io::Result> { + fn begin_read_stream(&mut self, offset: u64) -> io::Result> { io::Result::Ok(Box::from(NFSReadStream { disc_io: self, file: Option::None, diff --git a/src/streams.rs b/src/streams.rs index 3a666dd..1f6871e 100644 --- a/src/streams.rs +++ b/src/streams.rs @@ -173,3 +173,55 @@ impl<'a> WindowedReadStream for SharedWindowedReadStream<'a> { fn window(&self) -> (u64, u64) { (self.begin, self.end) } } + +pub struct ByteReadStream<'a> { + pub bytes: &'a [u8], + pub position: u64, +} + +impl Read for ByteReadStream<'_> { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let mut len = buf.len(); + let total = self.bytes.len(); + let pos = self.position as usize; + if len + pos > total { + if pos >= total { + return Err(io::Error::from(io::ErrorKind::UnexpectedEof)); + } + len = total - pos; + } + buf.copy_from_slice(&self.bytes[pos..pos + len]); + self.position += len as u64; + Ok(len) + } +} + +impl Seek for ByteReadStream<'_> { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + let new_pos = match pos { + SeekFrom::Start(v) => v, + SeekFrom::End(v) => (self.bytes.len() as i64 + v) as u64, + SeekFrom::Current(v) => (self.position as i64 + v) as u64, + }; + if new_pos > self.bytes.len() as u64 { + Err(io::Error::from(io::ErrorKind::UnexpectedEof)) + } else { + self.position = new_pos; + Ok(new_pos) + } + } + + // fn stream_len(&mut self) -> io::Result { Ok(self.bytes.len() as u64) } + + fn stream_position(&mut self) -> io::Result { Ok(self.position) } +} + +impl ReadStream for ByteReadStream<'_> { + fn stable_stream_len(&mut self) -> io::Result { Ok(self.bytes.len() as u64) } + + fn as_dyn(&mut self) -> &mut dyn ReadStream { self } +} + +impl<'a> AsMut for ByteReadStream<'a> { + fn as_mut(&mut self) -> &mut (dyn ReadStream + 'a) { self as &mut (dyn ReadStream + 'a) } +}