From da806d7ad4d5130022c04a72e19650297a8304d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= <lu.shuren@um-react.app>
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()];
         }
     }
 }