From 50f54837cabbea5ca89267336da120a0d9d53828 Mon Sep 17 00:00:00 2001 From: spuds <71292624+bananaturtlesandwich@users.noreply.github.com> Date: Sun, 27 Aug 2023 19:46:36 +0100 Subject: [PATCH] use Key enum to reduce cfg attributes used --- repak/src/entry.rs | 11 ++++--- repak/src/lib.rs | 18 +++++++++++ repak/src/pak.rs | 75 +++++++++++++-------------------------------- repak/tests/test.rs | 8 ++--- 4 files changed, 50 insertions(+), 62 deletions(-) diff --git a/repak/src/entry.rs b/repak/src/entry.rs index 35d6224..a6476b0 100644 --- a/repak/src/entry.rs +++ b/repak/src/entry.rs @@ -3,13 +3,13 @@ use byteorder::{ReadBytesExt, WriteBytesExt, LE}; use std::io; #[derive(Debug, PartialEq, Clone, Copy)] -pub enum EntryLocation { +pub(crate) enum EntryLocation { Data, Index, } #[derive(Debug)] -pub struct Block { +pub(crate) struct Block { pub start: u64, pub end: u64, } @@ -35,7 +35,7 @@ fn align(offset: u64) -> u64 { } #[derive(Debug)] -pub struct Entry { +pub(crate) struct Entry { pub offset: u64, pub compressed: u64, pub uncompressed: u64, @@ -311,7 +311,7 @@ impl Entry { reader: &mut R, version: Version, compression: &[Compression], - #[cfg(feature = "encryption")] key: Option<&aes::Aes256>, + key: &super::Key, buf: &mut W, ) -> Result<(), super::Error> { reader.seek(io::SeekFrom::Start(self.offset))?; @@ -328,7 +328,7 @@ impl Entry { return Err(super::Error::Encryption); #[cfg(feature = "encryption")] { - let Some(key) = key else { + let super::Key::Some(key) = key else { return Err(super::Error::Encrypted); }; use aes::cipher::BlockDecrypt; @@ -423,6 +423,7 @@ impl Entry { */ #[allow(non_snake_case)] + #[allow(clippy::type_complexity)] let OodleLZ_Decompress: libloading::Symbol< extern "C" fn( compBuf: *mut u8, diff --git a/repak/src/lib.rs b/repak/src/lib.rs index e66c4e0..7c1f99b 100644 --- a/repak/src/lib.rs +++ b/repak/src/lib.rs @@ -122,3 +122,21 @@ pub enum Compression { Oodle, Zstd, } + +#[allow(clippy::large_enum_variant)] +#[derive(Debug)] +pub(crate) enum Key { + #[cfg(feature = "encryption")] + Some(aes::Aes256), + None, +} + +impl From> for Key { + fn from(value: Option) -> Self { + match value { + #[cfg(feature = "encryption")] + Some(key) => Self::Some(key), + _ => Self::None, + } + } +} diff --git a/repak/src/pak.rs b/repak/src/pak.rs index 8161b42..c5de066 100644 --- a/repak/src/pak.rs +++ b/repak/src/pak.rs @@ -1,7 +1,5 @@ use super::ext::{ReadExt, WriteExt}; use super::{Version, VersionMajor}; -#[cfg(feature = "encryption")] -use aes::Aes256; use byteorder::{ReadBytesExt, WriteBytesExt, LE}; use std::collections::BTreeMap; use std::io::{self, Read, Seek, Write}; @@ -9,20 +7,18 @@ use std::io::{self, Read, Seek, Write}; #[derive(Debug)] pub struct PakReader { pak: Pak, - #[cfg(feature = "encryption")] - key: Option, + key: super::Key, } #[derive(Debug)] pub struct PakWriter { pak: Pak, writer: W, - #[cfg(feature = "encryption")] - key: Option, + key: super::Key, } #[derive(Debug)] -pub struct Pak { +pub(crate) struct Pak { version: Version, mount_point: String, index_offset: Option, @@ -47,7 +43,7 @@ impl Pak { } #[derive(Debug, Default)] -pub struct Index { +pub(crate) struct Index { path_hash_seed: Option, entries: BTreeMap, } @@ -74,8 +70,8 @@ impl Index { } #[cfg(feature = "encryption")] -fn decrypt(key: Option<&aes::Aes256>, bytes: &mut [u8]) -> Result<(), super::Error> { - if let Some(key) = key { +fn decrypt(key: &super::Key, bytes: &mut [u8]) -> Result<(), super::Error> { + if let super::Key::Some(key) = key { use aes::cipher::BlockDecrypt; for chunk in bytes.chunks_mut(16) { key.decrypt_block(aes::Block::from_mut_slice(chunk)) @@ -89,7 +85,7 @@ fn decrypt(key: Option<&aes::Aes256>, bytes: &mut [u8]) -> Result<(), super::Err impl PakReader { #[cfg(not(feature = "encryption"))] pub fn new_any(reader: &mut R) -> Result { - Self::new_any_inner(reader) + Self::new_any_inner(reader, super::Key::None) } #[cfg(feature = "encryption")] @@ -97,30 +93,19 @@ impl PakReader { reader: &mut R, key: Option, ) -> Result { - Self::new_any_inner(reader, key) + Self::new_any_inner(reader, key.into()) } fn new_any_inner( reader: &mut R, - #[cfg(feature = "encryption")] key: Option, + key: super::Key, ) -> Result { use std::fmt::Write; let mut log = "\n".to_owned(); for ver in Version::iter() { - match Pak::read( - &mut *reader, - ver, - #[cfg(feature = "encryption")] - key.as_ref(), - ) { - Ok(pak) => { - return Ok(Self { - pak, - #[cfg(feature = "encryption")] - key, - }) - } + match Pak::read(&mut *reader, ver, &key) { + Ok(pak) => return Ok(Self { pak, key }), Err(err) => writeln!(log, "trying version {} failed: {}", ver, err)?, } } @@ -141,25 +126,15 @@ impl PakReader { version: super::Version, key: Option, ) -> Result { - Self::new_inner(reader, version, key) + Self::new_inner(reader, version, key.into()) } fn new_inner( reader: &mut R, version: super::Version, - #[cfg(feature = "encryption")] key: Option, + key: super::Key, ) -> Result { - Pak::read( - reader, - version, - #[cfg(feature = "encryption")] - key.as_ref(), - ) - .map(|pak| Self { - pak, - #[cfg(feature = "encryption")] - key, - }) + Pak::read(reader, version, &key).map(|pak| Self { pak, key }) } pub fn version(&self) -> super::Version { @@ -195,8 +170,7 @@ impl PakReader { reader, self.pak.version, &self.pak.compression, - #[cfg(feature = "encryption")] - self.key.as_ref(), + &self.key, writer, ), None => Err(super::Error::MissingEntry(path.to_owned())), @@ -214,7 +188,6 @@ impl PakReader { writer.seek(io::SeekFrom::Start(self.pak.index_offset.unwrap()))?; Ok(PakWriter { pak: self.pak, - #[cfg(feature = "encryption")] key: self.key, writer, }) @@ -232,6 +205,7 @@ impl PakWriter { PakWriter { pak: Pak::new(version, mount_point, path_hash_seed), writer, + key: super::Key::None, } } @@ -246,13 +220,13 @@ impl PakWriter { PakWriter { pak: Pak::new(version, mount_point, path_hash_seed), writer, - key, + key: key.into(), } } fn new_inner( writer: W, - #[cfg(feature = "encryption")] key: Option, + key: super::Key, version: Version, mount_point: String, path_hash_seed: Option, @@ -260,7 +234,6 @@ impl PakWriter { PakWriter { pak: Pak::new(version, mount_point, path_hash_seed), writer, - #[cfg(feature = "encryption")] key, } } @@ -305,11 +278,7 @@ impl PakWriter { } pub fn write_index(mut self) -> Result { - self.pak.write( - &mut self.writer, - #[cfg(feature = "encryption")] - self.key, - )?; + self.pak.write(&mut self.writer, &self.key)?; Ok(self.writer) } } @@ -318,7 +287,7 @@ impl Pak { fn read( reader: &mut R, version: super::Version, - #[cfg(feature = "encryption")] key: Option<&aes::Aes256>, + key: &super::Key, ) -> Result { // read footer to get index, encryption & compression info reader.seek(io::SeekFrom::End(-version.size()))?; @@ -468,7 +437,7 @@ impl Pak { fn write( &self, writer: &mut W, - #[cfg(feature = "encryption")] _key: Option, + _key: &super::Key, ) -> Result<(), super::Error> { let index_offset = writer.stream_position()?; @@ -696,7 +665,7 @@ fn pad_zeros_to_alignment(v: &mut Vec, alignment: usize) { } #[cfg(feature = "encryption")] -fn encrypt(key: Aes256, bytes: &mut [u8]) { +fn encrypt(key: aes::Aes256, bytes: &mut [u8]) { use aes::cipher::BlockEncrypt; for chunk in bytes.chunks_mut(16) { key.encrypt_block(aes::Block::from_mut_slice(chunk)) diff --git a/repak/tests/test.rs b/repak/tests/test.rs index 1a7da1a..1355ac5 100644 --- a/repak/tests/test.rs +++ b/repak/tests/test.rs @@ -96,7 +96,7 @@ fn test_read(version: repak::Version, _file_name: &str, bytes: &[u8]) { let len = inner_reader.seek(SeekFrom::End(0)).unwrap(); let mut reader = ReadCounter::new_size(inner_reader, len as usize); - let pak = repak::PakReader::new_any(&mut reader, Some(key)).unwrap(); + let pak = repak::PakReader::new_any_with_key(&mut reader, Some(key)).unwrap(); assert_eq!(pak.mount_point(), "../mount/point/root/"); assert_eq!(pak.version(), version); @@ -159,10 +159,10 @@ fn test_write(_version: repak::Version, _file_name: &str, bytes: &[u8]) { .unwrap(); let mut reader = std::io::Cursor::new(bytes); - let pak_reader = repak::PakReader::new_any(&mut reader, Some(key)).unwrap(); + let pak_reader = repak::PakReader::new_any_with_key(&mut reader, Some(key)).unwrap(); let writer = Cursor::new(vec![]); - let mut pak_writer = repak::PakWriter::new( + let mut pak_writer = repak::PakWriter::new_with_key( writer, None, pak_reader.version(), @@ -191,7 +191,7 @@ fn test_rewrite_index(_version: repak::Version, _file_name: &str, bytes: &[u8]) .unwrap(); let mut buf = std::io::Cursor::new(bytes.to_vec()); - let pak_reader = repak::PakReader::new_any(&mut buf, Some(key)).unwrap(); + let pak_reader = repak::PakReader::new_any_with_key(&mut buf, Some(key)).unwrap(); let rewrite = pak_reader .into_pakwriter(buf)