Add Disc::detect for detecting disc image format

nod/src/disc/mod.rs

@@ -13,7 +13,7 @@ use std::{
 use dyn_clone::DynClone;
 use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes};
 
-use crate::{static_assert, Result};
+use crate::{io::MagicBytes, static_assert, Result};
 
 pub(crate) mod fst;
 pub(crate) mod gcn;
@@ -29,6 +29,12 @@ pub use wii::{SignedHeader, Ticket, TicketLimit, TmdHeader};
 /// Size in bytes of a disc sector.
 pub const SECTOR_SIZE: usize = 0x8000;
 
+/// Magic bytes for Wii discs. Located at offset 0x18.
+pub const WII_MAGIC: MagicBytes = [0x5D, 0x1C, 0x9E, 0xA3];
+
+/// Magic bytes for GameCube discs. Located at offset 0x1C.
+pub const GCN_MAGIC: MagicBytes = [0xC2, 0x33, 0x9F, 0x3D];
+
 /// Shared GameCube & Wii disc header.
 ///
 /// This header is always at the start of the disc image and within each Wii partition.
@@ -48,9 +54,9 @@ pub struct DiscHeader {
     /// Padding
     _pad1: [u8; 14],
     /// If this is a Wii disc, this will be 0x5D1C9EA3
-    pub wii_magic: U32,
+    pub wii_magic: MagicBytes,
     /// If this is a GameCube disc, this will be 0xC2339F3D
-    pub gcn_magic: U32,
+    pub gcn_magic: MagicBytes,
     /// Game title
     pub game_title: [u8; 64],
     /// If 1, disc omits partition hashes
@@ -79,11 +85,11 @@ impl DiscHeader {
 
     /// Whether this is a GameCube disc.
     #[inline]
-    pub fn is_gamecube(&self) -> bool { self.gcn_magic.get() == 0xC2339F3D }
+    pub fn is_gamecube(&self) -> bool { self.gcn_magic == GCN_MAGIC }
 
     /// Whether this is a Wii disc.
     #[inline]
-    pub fn is_wii(&self) -> bool { self.wii_magic.get() == 0x5D1C9EA3 }
+    pub fn is_wii(&self) -> bool { self.wii_magic == WII_MAGIC }
 }
 
 /// A header describing the contents of a disc partition.

nod/src/io/block.rs

@@ -13,11 +13,13 @@ use crate::{
     disc::{
         hashes::HashTable,
         wii::{WiiPartitionHeader, HASHES_SIZE, SECTOR_DATA_SIZE},
-        SECTOR_SIZE,
+        DiscHeader, PartitionHeader, PartitionKind, GCN_MAGIC, SECTOR_SIZE, WII_MAGIC,
+    },
+    io::{
+        aes_decrypt, aes_encrypt, split::SplitFileReader, DiscMeta, Format, KeyBytes, MagicBytes,
     },
-    io::{aes_decrypt, aes_encrypt, split::SplitFileReader, KeyBytes, MagicBytes},
     util::{lfg::LaggedFibonacci, read::read_from},
-    DiscHeader, DiscMeta, Error, PartitionHeader, PartitionKind, Result, ResultContext,
+    Error, Result, ResultContext,
 };
 
 /// Required trait bounds for reading disc images.
@@ -92,19 +94,26 @@ dyn_clone::clone_trait_object!(BlockIO);
 
 /// Creates a new [`BlockIO`] instance from a stream.
 pub fn new(mut stream: Box<dyn DiscStream>) -> Result<Box<dyn BlockIO>> {
-    let magic: MagicBytes = read_from(stream.as_mut()).context("Reading magic bytes")?;
+    let io: Box<dyn BlockIO> = match detect(stream.as_mut())? {
-    stream.seek(io::SeekFrom::Start(0)).context("Seeking to start")?;
+        Some(Format::Iso) => crate::io::iso::DiscIOISO::new(stream)?,
-    let io: Box<dyn BlockIO> = match magic {
+        Some(Format::Ciso) => crate::io::ciso::DiscIOCISO::new(stream)?,
-        crate::io::ciso::CISO_MAGIC => crate::io::ciso::DiscIOCISO::new(stream)?,
+        Some(Format::Gcz) => {
-        #[cfg(feature = "compress-zlib")]
+            #[cfg(feature = "compress-zlib")]
-        crate::io::gcz::GCZ_MAGIC => crate::io::gcz::DiscIOGCZ::new(stream)?,
+            {
-        crate::io::nfs::NFS_MAGIC => todo!(),
+                crate::io::gcz::DiscIOGCZ::new(stream)?
-        crate::io::wbfs::WBFS_MAGIC => crate::io::wbfs::DiscIOWBFS::new(stream)?,
+            }
-        crate::io::wia::WIA_MAGIC | crate::io::wia::RVZ_MAGIC => {
+            #[cfg(not(feature = "compress-zlib"))]
-            crate::io::wia::DiscIOWIA::new(stream)?
+            {
+                return Err(Error::DiscFormat("GCZ support is disabled".to_string()));
+            }
         }
-        crate::io::tgc::TGC_MAGIC => crate::io::tgc::DiscIOTGC::new(stream)?,
+        Some(Format::Nfs) => {
-        _ => crate::io::iso::DiscIOISO::new(stream)?,
+            return Err(Error::DiscFormat("NFS requires a filesystem path".to_string()))
+        }
+        Some(Format::Wbfs) => crate::io::wbfs::DiscIOWBFS::new(stream)?,
+        Some(Format::Wia | Format::Rvz) => crate::io::wia::DiscIOWIA::new(stream)?,
+        Some(Format::Tgc) => crate::io::tgc::DiscIOTGC::new(stream)?,
+        None => return Err(Error::DiscFormat("Unknown disc format".to_string())),
     };
     check_block_size(io.as_ref())?;
     Ok(io)
@@ -125,16 +134,20 @@ pub fn open(filename: &Path) -> Result<Box<dyn BlockIO>> {
         return Err(Error::DiscFormat(format!("Input is not a file: {}", filename.display())));
     }
     let mut stream = Box::new(SplitFileReader::new(filename)?);
-    let magic: MagicBytes = read_from(stream.as_mut())
+    let io: Box<dyn BlockIO> = match detect(stream.as_mut())? {
-        .with_context(|| format!("Reading magic bytes from {}", filename.display()))?;
+        Some(Format::Iso) => crate::io::iso::DiscIOISO::new(stream)?,
-    stream
+        Some(Format::Ciso) => crate::io::ciso::DiscIOCISO::new(stream)?,
-        .seek(io::SeekFrom::Start(0))
+        Some(Format::Gcz) => {
-        .with_context(|| format!("Seeking to start of {}", filename.display()))?;
+            #[cfg(feature = "compress-zlib")]
-    let io: Box<dyn BlockIO> = match magic {
+            {
-        crate::io::ciso::CISO_MAGIC => crate::io::ciso::DiscIOCISO::new(stream)?,
+                crate::io::gcz::DiscIOGCZ::new(stream)?
-        #[cfg(feature = "compress-zlib")]
+            }
-        crate::io::gcz::GCZ_MAGIC => crate::io::gcz::DiscIOGCZ::new(stream)?,
+            #[cfg(not(feature = "compress-zlib"))]
-        crate::io::nfs::NFS_MAGIC => match path.parent() {
+            {
+                return Err(Error::DiscFormat("GCZ support is disabled".to_string()));
+            }
+        }
+        Some(Format::Nfs) => match path.parent() {
             Some(parent) if parent.is_dir() => {
                 crate::io::nfs::DiscIONFS::new(path.parent().unwrap())?
             }
@@ -142,17 +155,45 @@ pub fn open(filename: &Path) -> Result<Box<dyn BlockIO>> {
                 return Err(Error::DiscFormat("Failed to locate NFS parent directory".to_string()));
             }
         },
-        crate::io::wbfs::WBFS_MAGIC => crate::io::wbfs::DiscIOWBFS::new(stream)?,
+        Some(Format::Tgc) => crate::io::tgc::DiscIOTGC::new(stream)?,
-        crate::io::wia::WIA_MAGIC | crate::io::wia::RVZ_MAGIC => {
+        Some(Format::Wbfs) => crate::io::wbfs::DiscIOWBFS::new(stream)?,
-            crate::io::wia::DiscIOWIA::new(stream)?
+        Some(Format::Wia | Format::Rvz) => crate::io::wia::DiscIOWIA::new(stream)?,
-        }
+        None => return Err(Error::DiscFormat("Unknown disc format".to_string())),
-        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)
 }
 
+pub const CISO_MAGIC: MagicBytes = *b"CISO";
+pub const GCZ_MAGIC: MagicBytes = [0x01, 0xC0, 0x0B, 0xB1];
+pub const NFS_MAGIC: MagicBytes = *b"EGGS";
+pub const TGC_MAGIC: MagicBytes = [0xae, 0x0f, 0x38, 0xa2];
+pub const WBFS_MAGIC: MagicBytes = *b"WBFS";
+pub const WIA_MAGIC: MagicBytes = *b"WIA\x01";
+pub const RVZ_MAGIC: MagicBytes = *b"RVZ\x01";
+
+pub fn detect<R: Read + ?Sized>(stream: &mut R) -> Result<Option<Format>> {
+    let data: [u8; 0x20] = match read_from(stream) {
+        Ok(magic) => magic,
+        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
+        Err(e) => return Err(e).context("Reading magic bytes"),
+    };
+    let out = match *array_ref!(data, 0, 4) {
+        CISO_MAGIC => Some(Format::Ciso),
+        GCZ_MAGIC => Some(Format::Gcz),
+        NFS_MAGIC => Some(Format::Nfs),
+        TGC_MAGIC => Some(Format::Tgc),
+        WBFS_MAGIC => Some(Format::Wbfs),
+        WIA_MAGIC => Some(Format::Wia),
+        RVZ_MAGIC => Some(Format::Rvz),
+        _ if *array_ref!(data, 0x18, 4) == WII_MAGIC || *array_ref!(data, 0x1C, 4) == GCN_MAGIC => {
+            Some(Format::Iso)
+        }
+        _ => None,
+    };
+    Ok(out)
+}
+
 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

nod/src/io/ciso.rs

@@ -9,7 +9,7 @@ use zerocopy::{little_endian::*, AsBytes, FromBytes, FromZeroes};
 use crate::{
     disc::SECTOR_SIZE,
     io::{
-        block::{Block, BlockIO, DiscStream, PartitionInfo},
+        block::{Block, BlockIO, DiscStream, PartitionInfo, CISO_MAGIC},
         nkit::NKitHeader,
         Format, MagicBytes,
     },
@@ -18,7 +18,6 @@ use crate::{
     DiscMeta, Error, Result, ResultContext,
 };
 
-pub const CISO_MAGIC: MagicBytes = *b"CISO";
 pub const CISO_MAP_SIZE: usize = SECTOR_SIZE - 8;
 
 /// CISO header (little endian)
@@ -43,6 +42,7 @@ pub struct DiscIOCISO {
 impl DiscIOCISO {
     pub fn new(mut inner: Box<dyn DiscStream>) -> Result<Box<Self>> {
         // Read header
+        inner.seek(SeekFrom::Start(0)).context("Seeking to start")?;
         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()));

nod/src/io/gcz.rs

@@ -11,7 +11,7 @@ use zstd::zstd_safe::WriteBuf;
 
 use crate::{
     io::{
-        block::{Block, BlockIO, DiscStream},
+        block::{Block, BlockIO, DiscStream, GCZ_MAGIC},
         MagicBytes,
     },
     static_assert,
@@ -19,8 +19,6 @@ use crate::{
     Compression, DiscMeta, Error, Format, PartitionInfo, Result, ResultContext,
 };
 
-pub const GCZ_MAGIC: MagicBytes = [0x01, 0xC0, 0x0B, 0xB1];
-
 /// GCZ header (little endian)
 #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)]
 #[repr(C, align(4))]
@@ -60,6 +58,7 @@ impl Clone for DiscIOGCZ {
 impl DiscIOGCZ {
     pub fn new(mut inner: Box<dyn DiscStream>) -> Result<Box<Self>> {
         // Read header
+        inner.seek(SeekFrom::Start(0)).context("Seeking to start")?;
         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()));

nod/src/io/nfs.rs

@@ -12,7 +12,7 @@ use crate::{
     disc::SECTOR_SIZE,
     io::{
         aes_decrypt,
-        block::{Block, BlockIO, PartitionInfo},
+        block::{Block, BlockIO, PartitionInfo, NFS_MAGIC},
         split::SplitFileReader,
         Format, KeyBytes, MagicBytes,
     },
@@ -21,7 +21,6 @@ use crate::{
     DiscMeta, Error, Result, ResultContext,
 };
 
-pub const NFS_MAGIC: MagicBytes = *b"EGGS";
 pub const NFS_END_MAGIC: MagicBytes = *b"SGGE";
 
 #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)]

nod/src/io/split.rs

@@ -136,8 +136,6 @@ impl Seek for SplitFileReader {
             if split.contains(self.pos) {
                 // Seek within the open file
                 split.inner.seek(SeekFrom::Start(self.pos - split.begin))?;
-            } else {
-                self.open_file = None;
             }
         }
         Ok(self.pos)

nod/src/io/tgc.rs

@@ -9,15 +9,13 @@ use zerocopy::{big_endian::U32, AsBytes, FromBytes, FromZeroes};
 use crate::{
     disc::SECTOR_SIZE,
     io::{
-        block::{Block, BlockIO, DiscStream, PartitionInfo},
+        block::{Block, BlockIO, DiscStream, PartitionInfo, TGC_MAGIC},
         Format, MagicBytes,
     },
     util::read::{read_box_slice, read_from},
     DiscHeader, DiscMeta, Error, Node, PartitionHeader, Result, ResultContext,
 };
 
-pub const TGC_MAGIC: MagicBytes = [0xae, 0x0f, 0x38, 0xa2];
-
 /// TGC header (big endian)
 #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)]
 #[repr(C, align(4))]

nod/src/io/wbfs.rs

@@ -8,7 +8,7 @@ use zerocopy::{big_endian::*, AsBytes, FromBytes, FromZeroes};
 
 use crate::{
     io::{
-        block::{Block, BlockIO, DiscStream, PartitionInfo},
+        block::{Block, BlockIO, DiscStream, PartitionInfo, WBFS_MAGIC},
         nkit::NKitHeader,
         DiscMeta, Format, MagicBytes,
     },
@@ -16,8 +16,6 @@ use crate::{
     Error, Result, ResultContext,
 };
 
-pub const WBFS_MAGIC: MagicBytes = *b"WBFS";
-
 #[derive(Debug, Clone, PartialEq, FromBytes, FromZeroes, AsBytes)]
 #[repr(C, align(4))]
 struct WBFSHeader {
@@ -52,6 +50,7 @@ pub struct DiscIOWBFS {
 
 impl DiscIOWBFS {
     pub fn new(mut inner: Box<dyn DiscStream>) -> Result<Box<Self>> {
+        inner.seek(SeekFrom::Start(0)).context("Seeking to start")?;
         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()));

nod/src/io/wia.rs

@@ -13,7 +13,7 @@ use crate::{
         SECTOR_SIZE,
     },
     io::{
-        block::{Block, BlockIO, DiscStream, PartitionInfo},
+        block::{Block, BlockIO, DiscStream, PartitionInfo, RVZ_MAGIC, WIA_MAGIC},
         nkit::NKitHeader,
         Compression, Format, HashBytes, KeyBytes, MagicBytes,
     },
@@ -27,9 +27,6 @@ use crate::{
     DiscMeta, Error, Result, ResultContext,
 };
 
-pub const WIA_MAGIC: MagicBytes = *b"WIA\x01";
-pub const RVZ_MAGIC: MagicBytes = *b"RVZ\x01";
-
 /// This struct is stored at offset 0x0 and is 0x48 bytes long. The wit source code says its format
 /// will never be changed.
 #[derive(Clone, Debug, PartialEq, FromBytes, FromZeroes, AsBytes)]
@@ -549,6 +546,7 @@ fn verify_hash(buf: &[u8], expected: &HashBytes) -> Result<()> {
 impl DiscIOWIA {
     pub fn new(mut inner: Box<dyn DiscStream>) -> Result<Box<Self>> {
         // Load & verify file header
+        inner.seek(SeekFrom::Start(0)).context("Seeking to start")?;
         let header: WIAFileHeader =
             read_from(inner.as_mut()).context("Reading WIA/RVZ file header")?;
         header.validate()?;

nod/src/lib.rs

@@ -66,11 +66,11 @@ use std::{
 pub use disc::{
     ApploaderHeader, DiscHeader, DolHeader, FileStream, Fst, Node, NodeKind, OwnedFileStream,
     PartitionBase, PartitionHeader, PartitionKind, PartitionMeta, SignedHeader, Ticket,
-    TicketLimit, TmdHeader, WindowedStream, BI2_SIZE, BOOT_SIZE, SECTOR_SIZE,
+    TicketLimit, TmdHeader, WindowedStream, BI2_SIZE, BOOT_SIZE, GCN_MAGIC, SECTOR_SIZE, WII_MAGIC,
 };
 pub use io::{
     block::{DiscStream, PartitionInfo},
-    Compression, DiscMeta, Format, KeyBytes,
+    Compression, DiscMeta, Format, KeyBytes, MagicBytes,
 };
 
 mod disc;
@@ -190,6 +190,13 @@ impl Disc {
         Ok(Disc { reader, options: options.clone() })
     }
 
+    /// Detects the format of a disc image from a read stream.
+    #[inline]
+    pub fn detect<R>(stream: &mut R) -> Result<Option<Format>>
+    where R: Read + ?Sized {
+        io::block::detect(stream)
+    }
+
     /// The disc's primary header.
     #[inline]
     pub fn header(&self) -> &DiscHeader { self.reader.header() }

nodtool/src/cmd/info.rs

@@ -92,11 +92,7 @@ fn info_file(path: &Path) -> nod::Result<()> {
     } else if header.is_gamecube() {
         // TODO
     } else {
-        println!(
+        println!("Invalid GC/Wii magic: {:#x?}/{:#x?}", header.gcn_magic, header.wii_magic);
-            "Invalid GC/Wii magic: {:#010X}/{:#010X}",
-            header.gcn_magic.get(),
-            header.wii_magic.get()
-        );
     }
     println!();
     Ok(())