From da806d7ad4d5130022c04a72e19650297a8304d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= Date: Sun, 15 Sep 2024 23:25:19 +0100 Subject: [PATCH] [kgm/ncm/kwm] refactor: improve performance for large file --- um_crypto/kgm/src/v2.rs | 5 ++--- um_crypto/kgm/src/v3.rs | 27 ++++++++++----------------- um_crypto/kuwo/src/kwm_v1.rs | 11 ++++------- um_crypto/ncm/src/header.rs | 5 ++--- 4 files changed, 18 insertions(+), 30 deletions(-) diff --git a/um_crypto/kgm/src/v2.rs b/um_crypto/kgm/src/v2.rs index b89c2a7..fff218c 100644 --- a/um_crypto/kgm/src/v2.rs +++ b/um_crypto/kgm/src/v2.rs @@ -15,9 +15,8 @@ impl DecipherV2 { impl Decipher for DecipherV2 { fn decrypt(&self, buffer: &mut [u8], offset: usize) { - let key_stream = self.key.iter().cycle().skip(offset % self.key.len()); - for (datum, &k) in buffer.iter_mut().zip(key_stream) { - *datum ^= k; + for (datum, offset) in buffer.iter_mut().zip(offset..) { + *datum ^= self.key[offset % self.key.len()]; } } } diff --git a/um_crypto/kgm/src/v3.rs b/um_crypto/kgm/src/v3.rs index 451ff47..0182950 100644 --- a/um_crypto/kgm/src/v3.rs +++ b/um_crypto/kgm/src/v3.rs @@ -25,34 +25,27 @@ impl DecipherV3 { Ok(Self { slot_key, file_key }) } + + fn offset_key(offset: usize) -> u8 { + let offset_key = (offset as u32).to_ne_bytes(); + offset_key[0] ^ offset_key[1] ^ offset_key[2] ^ offset_key[3] + } } impl Decipher for DecipherV3 { fn decrypt(&self, buffer: &mut [u8], offset: usize) { - let slot_key_stream = self - .slot_key - .iter() - .cycle() - .skip(offset % self.slot_key.len()); - let file_key_stream = self - .file_key - .iter() - .cycle() - .skip(offset % self.file_key.len()); + for (datum, offset) in buffer.iter_mut().zip(offset..) { + let offset_key = Self::offset_key(offset); - let mut offset = offset as u32; - let key_stream = slot_key_stream.zip(file_key_stream); - for (datum, (&slot_key, &file_key)) in buffer.iter_mut().zip(key_stream) { - let offset_key = offset.to_ne_bytes().iter().fold(0, |acc, &x| acc ^ x); + let file_key = self.file_key[offset % self.file_key.len()]; + let slot_key = self.slot_key[offset % self.slot_key.len()]; let mut temp = *datum; temp ^= file_key; - temp ^= temp.wrapping_shl(4); + temp ^= temp << 4; temp ^= slot_key; temp ^= offset_key; *datum = temp; - - offset = offset.wrapping_add(1); } } } diff --git a/um_crypto/kuwo/src/kwm_v1.rs b/um_crypto/kuwo/src/kwm_v1.rs index b530122..6b75196 100644 --- a/um_crypto/kuwo/src/kwm_v1.rs +++ b/um_crypto/kuwo/src/kwm_v1.rs @@ -11,10 +11,8 @@ const KEY: [u8; 0x20] = [ impl CipherV1 { pub fn new(resource_id: u32) -> Self { let mut key = KEY; - for (k, r) in key - .iter_mut() - .zip(resource_id.to_string().as_bytes().iter().cycle()) - { + let resource_id = resource_id.to_string(); + for (k, &r) in key.iter_mut().zip(resource_id.as_bytes().iter().cycle()) { *k ^= r; } @@ -26,10 +24,9 @@ impl CipherV1 { T: AsMut<[u8]> + ?Sized, { let data = data.as_mut(); - let key_stream = self.key.iter().cycle().skip(offset % self.key.len()); - for (datum, key) in data.iter_mut().zip(key_stream) { - *datum ^= *key; + for (datum, offset) in data.iter_mut().zip(offset..) { + *datum ^= self.key[offset % self.key.len()]; } } } diff --git a/um_crypto/ncm/src/header.rs b/um_crypto/ncm/src/header.rs index 2ab8bfd..a7df520 100644 --- a/um_crypto/ncm/src/header.rs +++ b/um_crypto/ncm/src/header.rs @@ -182,9 +182,8 @@ impl NCMFile { where T: AsMut<[u8]> + ?Sized, { - let key_stream = self.audio_rc4_key_stream.iter().cycle().skip(offset & 0xff); - for (datum, &key) in data.as_mut().iter_mut().zip(key_stream) { - *datum ^= key; + for (datum, offset) in data.as_mut().iter_mut().zip(offset..) { + *datum ^= self.audio_rc4_key_stream[offset % self.audio_rc4_key_stream.len()]; } } }