SharedWindowedReadStream -> FileStream & impl BufRead

This commit is contained in:
Luke Street 2024-09-10 23:19:57 -06:00
parent 6f3052e05d
commit 312dd6f080
11 changed files with 256 additions and 175 deletions

View File

@ -1,20 +1,17 @@
use std::{
cmp::min,
io,
io::{Read, Seek, SeekFrom},
io::{BufRead, Read, Seek, SeekFrom},
mem::size_of,
};
use zerocopy::{FromBytes, FromZeroes};
use super::{
ApploaderHeader, DiscHeader, DolHeader, FileStream, Node, PartitionBase, PartitionHeader,
PartitionMeta, BI2_SIZE, BOOT_SIZE, SECTOR_SIZE,
};
use crate::{
disc::{
ApploaderHeader, DiscHeader, DolHeader, PartitionBase, PartitionHeader, PartitionMeta,
BI2_SIZE, BOOT_SIZE, SECTOR_SIZE,
},
fst::{Node, NodeKind},
io::block::{Block, BlockIO},
streams::{ReadStream, SharedWindowedReadStream},
util::read::{read_box, read_box_slice, read_vec},
Result, ResultContext,
};
@ -63,8 +60,8 @@ impl PartitionGC {
pub fn into_inner(self) -> Box<dyn BlockIO> { self.io }
}
impl Read for PartitionGC {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
impl BufRead for PartitionGC {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
let sector = (self.pos / SECTOR_SIZE as u64) as u32;
let block_idx = (sector as u64 * SECTOR_SIZE as u64 / self.block_buf.len() as u64) as u32;
@ -86,9 +83,20 @@ impl Read for PartitionGC {
}
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(&self.sector_buf[offset..])
}
#[inline]
fn consume(&mut self, amt: usize) { self.pos += amt as u64; }
}
impl Read for PartitionGC {
#[inline]
fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
let buf = self.fill_buf()?;
let len = buf.len().min(out.len());
out[..len].copy_from_slice(&buf[..len]);
self.consume(len);
Ok(len)
}
}
@ -115,16 +123,19 @@ impl PartitionBase for PartitionGC {
read_part_meta(self, false)
}
fn open_file(&mut self, node: &Node) -> io::Result<SharedWindowedReadStream> {
assert_eq!(node.kind(), NodeKind::File);
self.new_window(node.offset(false), node.length())
fn open_file(&mut self, node: &Node) -> io::Result<FileStream> {
if !node.is_file() {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"Node is not a file".to_string(),
));
}
FileStream::new(self, node.offset(false), node.length())
}
fn ideal_buffer_size(&self) -> usize { SECTOR_SIZE }
}
pub(crate) fn read_part_meta(
reader: &mut dyn ReadStream,
reader: &mut dyn PartitionBase,
is_wii: bool,
) -> Result<Box<PartitionMeta>> {
// boot.bin

View File

@ -5,6 +5,7 @@ use std::{
ffi::CStr,
fmt::{Debug, Display, Formatter},
io,
io::{BufRead, Seek},
mem::size_of,
str::from_utf8,
};
@ -12,19 +13,19 @@ use std::{
use dyn_clone::DynClone;
use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes};
use crate::{
disc::wii::{Ticket, TmdHeader},
fst::Node,
static_assert,
streams::{ReadStream, SharedWindowedReadStream},
Fst, Result,
};
use crate::{static_assert, Result};
pub(crate) mod fst;
pub(crate) mod gcn;
pub(crate) mod hashes;
pub(crate) mod reader;
pub(crate) mod streams;
pub(crate) mod wii;
pub use fst::{Fst, Node, NodeKind};
pub use streams::FileStream;
pub use wii::{SignedHeader, Ticket, TicketLimit, TmdHeader};
/// Size in bytes of a disc sector.
pub const SECTOR_SIZE: usize = 0x8000;
@ -274,11 +275,11 @@ impl From<u32> for PartitionKind {
}
/// An open disc partition.
pub trait PartitionBase: DynClone + ReadStream + Send + Sync {
pub trait PartitionBase: DynClone + BufRead + Seek + Send + Sync {
/// Reads the partition header and file system table.
fn meta(&mut self) -> Result<Box<PartitionMeta>>;
/// Seeks the read stream to the specified file system node
/// Seeks the partition stream to the specified file system node
/// and returns a windowed stream.
///
/// # Examples
@ -306,12 +307,7 @@ pub trait PartitionBase: DynClone + ReadStream + Send + Sync {
/// Ok(())
/// }
/// ```
fn open_file(&mut self, node: &Node) -> io::Result<SharedWindowedReadStream>;
/// The ideal size for buffered reads from this partition.
/// GameCube discs have a data block size of 0x8000,
/// whereas Wii discs have a data block size of 0x7C00.
fn ideal_buffer_size(&self) -> usize;
fn open_file(&mut self, node: &Node) -> io::Result<FileStream>;
}
dyn_clone::clone_trait_object!(PartitionBase);

View File

@ -1,22 +1,21 @@
use std::{
cmp::min,
io,
io::{Read, Seek, SeekFrom},
io::{BufRead, Read, Seek, SeekFrom},
};
use zerocopy::FromZeroes;
use super::{
gcn::PartitionGC,
hashes::{rebuild_hashes, HashTable},
wii::{PartitionWii, WiiPartEntry, WiiPartGroup, WiiPartitionHeader, WII_PART_GROUP_OFF},
DiscHeader, PartitionBase, PartitionHeader, PartitionKind, DL_DVD_SIZE, MINI_DVD_SIZE,
SL_DVD_SIZE,
};
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,
DiscMeta, Error, OpenOptions, Result, ResultContext, SECTOR_SIZE,
};
#[derive(Debug, Eq, PartialEq, Copy, Clone)]
@ -150,8 +149,8 @@ impl DiscReader {
}
}
impl Read for DiscReader {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
impl BufRead for DiscReader {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
let block_idx = (self.pos / self.block_buf.len() as u64) as u32;
let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;
@ -199,9 +198,20 @@ impl Read for DiscReader {
// 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(&self.sector_buf[offset..])
}
#[inline]
fn consume(&mut self, amt: usize) { self.pos += amt as u64; }
}
impl Read for DiscReader {
#[inline]
fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
let buf = self.fill_buf()?;
let len = buf.len().min(out.len());
out[..len].copy_from_slice(&buf[..len]);
self.consume(len);
Ok(len)
}
}

83
nod/src/disc/streams.rs Normal file
View File

@ -0,0 +1,83 @@
//! Partition file read stream.
use std::{
io,
io::{BufRead, Read, Seek, SeekFrom},
};
use super::PartitionBase;
/// A file read stream for a [`PartitionBase`].
pub struct FileStream<'a> {
base: &'a mut dyn PartitionBase,
pos: u64,
begin: u64,
end: u64,
}
impl FileStream<'_> {
/// Creates a new file stream with offset and size.
///
/// Seeks underlying stream immediately.
#[inline]
pub(crate) fn new(
base: &mut dyn PartitionBase,
offset: u64,
size: u64,
) -> io::Result<FileStream> {
base.seek(SeekFrom::Start(offset))?;
Ok(FileStream { base, pos: offset, begin: offset, end: offset + size })
}
}
impl<'a> Read for FileStream<'a> {
#[inline]
fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
let buf = self.fill_buf()?;
let len = buf.len().min(out.len());
out[..len].copy_from_slice(&buf[..len]);
self.consume(len);
Ok(len)
}
}
impl<'a> BufRead for FileStream<'a> {
#[inline]
fn fill_buf(&mut self) -> io::Result<&[u8]> {
let limit = self.end.saturating_sub(self.pos);
if limit == 0 {
return Ok(&[]);
}
let buf = self.base.fill_buf()?;
let max = (buf.len() as u64).min(limit) as usize;
Ok(&buf[..max])
}
#[inline]
fn consume(&mut self, amt: usize) {
self.base.consume(amt);
self.pos += amt as u64;
}
}
impl<'a> Seek for FileStream<'a> {
#[inline]
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
let mut pos = match pos {
SeekFrom::Start(p) => self.begin + p,
SeekFrom::End(p) => self.end.saturating_add_signed(p),
SeekFrom::Current(p) => self.pos.saturating_add_signed(p),
};
if pos < self.begin {
pos = self.begin;
} else if pos > self.end {
pos = self.end;
}
let result = self.base.seek(SeekFrom::Start(pos))?;
self.pos = result;
Ok(result - self.begin)
}
#[inline]
fn stream_position(&mut self) -> io::Result<u64> { Ok(self.pos) }
}

View File

@ -1,30 +1,27 @@
use std::{
cmp::min,
ffi::CStr,
io,
io::{Read, Seek, SeekFrom},
io::{BufRead, Read, Seek, SeekFrom},
mem::size_of,
};
use sha1::{Digest, Sha1};
use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes};
use super::{
gcn::{read_part_meta, PartitionGC},
DiscHeader, FileStream, Node, PartitionBase, PartitionMeta, SECTOR_SIZE,
};
use crate::{
array_ref,
disc::{
gcn::{read_part_meta, PartitionGC},
PartitionBase, PartitionMeta, SECTOR_SIZE,
},
fst::{Node, NodeKind},
io::{
aes_decrypt,
block::{Block, BlockIO, PartitionInfo},
KeyBytes,
},
static_assert,
streams::{ReadStream, SharedWindowedReadStream},
util::{div_rem, read::read_box_slice},
DiscHeader, Error, OpenOptions, Result, ResultContext,
Error, OpenOptions, Result, ResultContext,
};
/// Size in bytes of the hashes block in a Wii disc sector
@ -85,6 +82,7 @@ impl WiiPartGroup {
pub(crate) fn part_entry_off(&self) -> u64 { (self.part_entry_off.get() as u64) << 2 }
}
/// Signed blob header
#[derive(Debug, Clone, PartialEq, FromBytes, FromZeroes, AsBytes)]
#[repr(C, align(4))]
pub struct SignedHeader {
@ -97,43 +95,64 @@ pub struct SignedHeader {
static_assert!(size_of::<SignedHeader>() == 0x140);
/// Ticket limit
#[derive(Debug, Clone, PartialEq, Default, FromBytes, FromZeroes, AsBytes)]
#[repr(C, align(4))]
pub struct TicketTimeLimit {
pub enable_time_limit: U32,
pub time_limit: U32,
pub struct TicketLimit {
/// Limit type
pub limit_type: U32,
/// Maximum value for the limit
pub max_value: U32,
}
static_assert!(size_of::<TicketTimeLimit>() == 8);
static_assert!(size_of::<TicketLimit>() == 8);
/// Wii ticket
#[derive(Debug, Clone, PartialEq, FromBytes, FromZeroes, AsBytes)]
#[repr(C, align(4))]
pub struct Ticket {
/// Signed blob header
pub header: SignedHeader,
/// Signature issuer
pub sig_issuer: [u8; 64],
/// ECDH data
pub ecdh: [u8; 60],
/// Ticket format version
pub version: u8,
_pad1: U16,
/// Title key (encrypted)
pub title_key: KeyBytes,
_pad2: u8,
/// Ticket ID
pub ticket_id: [u8; 8],
/// Console ID
pub console_id: [u8; 4],
/// Title ID
pub title_id: [u8; 8],
_pad3: U16,
/// Ticket title version
pub ticket_title_version: U16,
/// Permitted titles mask
pub permitted_titles_mask: U32,
/// Permit mask
pub permit_mask: U32,
/// Title export allowed
pub title_export_allowed: u8,
/// Common key index
pub common_key_idx: u8,
_pad4: [u8; 48],
/// Content access permissions
pub content_access_permissions: [u8; 64],
_pad5: [u8; 2],
pub time_limits: [TicketTimeLimit; 8],
/// Ticket limits
pub limits: [TicketLimit; 8],
}
static_assert!(size_of::<Ticket>() == 0x2A4);
impl Ticket {
/// Decrypts the ticket title key using the appropriate common key
#[allow(clippy::missing_inline_in_public_items)]
pub fn decrypt_title_key(&self) -> Result<KeyBytes> {
let mut iv: KeyBytes = [0; 16];
iv[..8].copy_from_slice(&self.title_id);
@ -158,29 +177,48 @@ impl Ticket {
}
}
/// Title metadata header
#[derive(Debug, Clone, PartialEq, FromBytes, FromZeroes, AsBytes)]
#[repr(C, align(4))]
pub struct TmdHeader {
/// Signed blob header
pub header: SignedHeader,
/// Signature issuer
pub sig_issuer: [u8; 64],
/// Version
pub version: u8,
/// CA CRL version
pub ca_crl_version: u8,
/// Signer CRL version
pub signer_crl_version: u8,
/// Is vWii title
pub is_vwii: u8,
/// IOS ID
pub ios_id: [u8; 8],
/// Title ID
pub title_id: [u8; 8],
/// Title type
pub title_type: u32,
/// Group ID
pub group_id: U16,
_pad1: [u8; 2],
/// Region
pub region: U16,
/// Ratings
pub ratings: KeyBytes,
_pad2: [u8; 12],
/// IPC mask
pub ipc_mask: [u8; 12],
_pad3: [u8; 18],
/// Access flags
pub access_flags: U32,
/// Title version
pub title_version: U16,
/// Number of contents
pub num_contents: U16,
/// Boot index
pub boot_idx: U16,
/// Minor version (unused)
pub minor_version: U16,
}
@ -301,12 +339,12 @@ impl PartitionWii {
}
}
impl Read for PartitionWii {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
impl BufRead for PartitionWii {
fn fill_buf(&mut self) -> io::Result<&[u8]> {
let part_sector = (self.pos / SECTOR_DATA_SIZE as u64) as u32;
let abs_sector = self.partition.data_start_sector + part_sector;
if abs_sector >= self.partition.data_end_sector {
return Ok(0);
return Ok(&[]);
}
let block_idx =
(abs_sector as u64 * SECTOR_SIZE as u64 / self.block_buf.len() as u64) as u32;
@ -333,10 +371,20 @@ impl Read for PartitionWii {
}
let offset = (self.pos % SECTOR_DATA_SIZE as u64) as usize;
let len = min(buf.len(), SECTOR_DATA_SIZE - offset);
buf[..len]
.copy_from_slice(&self.sector_buf[HASHES_SIZE + offset..HASHES_SIZE + offset + len]);
self.pos += len as u64;
Ok(&self.sector_buf[HASHES_SIZE + offset..])
}
#[inline]
fn consume(&mut self, amt: usize) { self.pos += amt as u64; }
}
impl Read for PartitionWii {
#[inline]
fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
let buf = self.fill_buf()?;
let len = buf.len().min(out.len());
out[..len].copy_from_slice(&buf[..len]);
self.consume(len);
Ok(len)
}
}
@ -440,10 +488,13 @@ impl PartitionBase for PartitionWii {
Ok(meta)
}
fn open_file(&mut self, node: &Node) -> io::Result<SharedWindowedReadStream> {
assert_eq!(node.kind(), NodeKind::File);
self.new_window(node.offset(true), node.length())
fn open_file(&mut self, node: &Node) -> io::Result<FileStream> {
if !node.is_file() {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"Node is not a file".to_string(),
));
}
FileStream::new(self, node.offset(true), node.length())
}
fn ideal_buffer_size(&self) -> usize { SECTOR_DATA_SIZE }
}

View File

@ -15,13 +15,13 @@ pub(crate) mod wbfs;
pub(crate) mod wia;
/// SHA-1 hash bytes
pub(crate) type HashBytes = [u8; 20];
pub type HashBytes = [u8; 20];
/// AES key bytes
pub(crate) type KeyBytes = [u8; 16];
pub type KeyBytes = [u8; 16];
/// Magic bytes
pub(crate) type MagicBytes = [u8; 4];
pub type MagicBytes = [u8; 4];
/// The disc file format.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]

View File

@ -192,7 +192,7 @@ impl DiscIONFS {
let resolved_path = key_path.unwrap();
File::open(resolved_path.as_path())
.map_err(|v| Error::Io(format!("Failed to open {}", resolved_path.display()), v))?
.read(&mut self.key)
.read_exact(&mut self.key)
.map_err(|v| Error::Io(format!("Failed to read {}", resolved_path.display()), v))?;
}

View File

@ -59,22 +59,19 @@
//! ```
use std::{
io::{Read, Seek},
io::{BufRead, Read, Seek},
path::Path,
};
pub use disc::{
ApploaderHeader, DiscHeader, DolHeader, PartitionBase, PartitionHeader, PartitionKind,
PartitionMeta, BI2_SIZE, BOOT_SIZE, SECTOR_SIZE,
ApploaderHeader, DiscHeader, DolHeader, FileStream, Fst, Node, NodeKind, PartitionBase,
PartitionHeader, PartitionKind, PartitionMeta, SignedHeader, Ticket, TicketLimit, TmdHeader,
BI2_SIZE, BOOT_SIZE, SECTOR_SIZE,
};
pub use fst::{Fst, Node, NodeKind};
pub use io::{block::PartitionInfo, Compression, DiscMeta, Format};
pub use streams::ReadStream;
pub use io::{block::PartitionInfo, Compression, DiscMeta, Format, KeyBytes};
mod disc;
mod fst;
mod io;
mod streams;
mod util;
/// Error types for nod.
@ -209,6 +206,14 @@ impl Disc {
}
}
impl BufRead for Disc {
#[inline]
fn fill_buf(&mut self) -> std::io::Result<&[u8]> { self.reader.fill_buf() }
#[inline]
fn consume(&mut self, amt: usize) { self.reader.consume(amt) }
}
impl Read for Disc {
#[inline]
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { self.reader.read(buf) }

View File

@ -1,82 +0,0 @@
//! Common stream types
use std::{
io,
io::{Read, Seek, SeekFrom},
};
/// A helper trait for seekable read streams.
pub trait ReadStream: Read + Seek {
/// Creates a windowed read sub-stream with offset and size.
///
/// Seeks underlying stream immediately.
#[inline]
fn new_window(&mut self, offset: u64, size: u64) -> io::Result<SharedWindowedReadStream> {
self.seek(SeekFrom::Start(offset))?;
Ok(SharedWindowedReadStream { base: self.as_dyn(), begin: offset, end: offset + size })
}
/// Retrieves a type-erased reference to the stream.
fn as_dyn(&mut self) -> &mut dyn ReadStream;
}
impl<T> ReadStream for T
where T: Read + Seek
{
#[inline]
fn as_dyn(&mut self) -> &mut dyn ReadStream { self }
}
/// A non-owning window into an existing [`ReadStream`].
pub struct SharedWindowedReadStream<'a> {
/// A reference to the base stream.
pub base: &'a mut dyn ReadStream,
/// The beginning of the window in bytes.
pub begin: u64,
/// The end of the window in bytes.
pub end: u64,
}
impl<'a> SharedWindowedReadStream<'a> {
/// Modifies the current window & seeks to the beginning of the window.
pub fn set_window(&mut self, begin: u64, end: u64) -> io::Result<()> {
self.base.seek(SeekFrom::Start(begin))?;
self.begin = begin;
self.end = end;
Ok(())
}
}
impl<'a> Read for SharedWindowedReadStream<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let pos = self.stream_position()?;
let size = self.end - self.begin;
if pos == size {
return Ok(0);
}
self.base.read(if pos + buf.len() as u64 > size {
&mut buf[..(size - pos) as usize]
} else {
buf
})
}
}
impl<'a> Seek for SharedWindowedReadStream<'a> {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
let result = self.base.seek(match pos {
SeekFrom::Start(p) => SeekFrom::Start(self.begin + p),
SeekFrom::End(p) => SeekFrom::End(self.end as i64 + p),
SeekFrom::Current(_) => pos,
})?;
if result < self.begin || result > self.end {
Err(io::Error::from(io::ErrorKind::UnexpectedEof))
} else {
Ok(result - self.begin)
}
}
fn stream_position(&mut self) -> io::Result<u64> {
Ok(self.base.stream_position()? - self.begin)
}
}

View File

@ -2,8 +2,7 @@ use std::{
borrow::Cow,
fs,
fs::File,
io,
io::{BufWriter, Write},
io::{BufRead, Write},
path::{Path, PathBuf},
};
@ -197,9 +196,8 @@ fn extract_node(
Size::from_bytes(node.length()).format().with_base(Base::Base10)
);
}
let file = File::create(&file_path)
let mut file = File::create(&file_path)
.with_context(|| format!("Creating file {}", display(&file_path)))?;
let mut w = BufWriter::with_capacity(partition.ideal_buffer_size(), file);
let mut r = partition.open_file(node).with_context(|| {
format!(
"Opening file {} on disc for reading (offset {}, size {})",
@ -208,7 +206,16 @@ fn extract_node(
node.length()
)
})?;
io::copy(&mut r, &mut w).with_context(|| format!("Extracting file {}", display(&file_path)))?;
w.flush().with_context(|| format!("Flushing file {}", display(&file_path)))?;
loop {
let buf =
r.fill_buf().with_context(|| format!("Extracting file {}", display(&file_path)))?;
let len = buf.len();
if len == 0 {
break;
}
file.write_all(buf).with_context(|| format!("Writing file {}", display(&file_path)))?;
r.consume(len);
}
file.flush().with_context(|| format!("Flushing file {}", display(&file_path)))?;
Ok(())
}