From ba422e3ccef9fe07de08b36987307f4487fef479 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: Tue, 25 Feb 2025 04:51:44 +0900
Subject: [PATCH] feat(kgm): move slot_key logic out

---
 um_crypto/kgm/src/lib.rs       | 11 +++--------
 um_crypto/kgm/src/slot_keys.rs |  9 +++++++++
 um_crypto/kgm/src/v2.rs        |  7 ++++---
 um_crypto/kgm/src/v3.rs        |  5 +++--
 4 files changed, 19 insertions(+), 13 deletions(-)
 create mode 100644 um_crypto/kgm/src/slot_keys.rs

diff --git a/um_crypto/kgm/src/lib.rs b/um_crypto/kgm/src/lib.rs
index a8f88cc..6f65ed8 100644
--- a/um_crypto/kgm/src/lib.rs
+++ b/um_crypto/kgm/src/lib.rs
@@ -1,5 +1,6 @@
 pub mod header;
 mod pc_db_decrypt;
+mod slot_keys;
 pub mod v2;
 pub mod v3;
 mod v5;
@@ -68,15 +69,9 @@ impl Decipher {
     }
 
     pub fn new_v5(header: &Header, ekey: Option<String>) -> Result<Self, KugouError> {
-        let slot_key: &[u8] = match header.key_slot {
-            1 => b"l,/'",
-            -1 => b"", // unused, kgm v5 (kgg)
-            slot => Err(KugouError::UnsupportedKeySlot(slot))?,
-        };
-
         let decipher = match header.crypto_version {
-            2 => Decipher::V2(DecipherV2::new(header, slot_key)?),
-            3 => Decipher::V3(DecipherV3::new(header, slot_key)?),
+            2 => Decipher::V2(DecipherV2::new(header)?),
+            3 => Decipher::V3(DecipherV3::new(header)?),
             5 => match ekey {
                 Some(ekey) => Decipher::V5(DecipherV5::new(&ekey)?),
                 _ => Err(KugouError::V5EKeyRequired)?,
diff --git a/um_crypto/kgm/src/slot_keys.rs b/um_crypto/kgm/src/slot_keys.rs
new file mode 100644
index 0000000..26d8375
--- /dev/null
+++ b/um_crypto/kgm/src/slot_keys.rs
@@ -0,0 +1,9 @@
+use crate::header::Header;
+use crate::KugouError;
+
+pub fn get_slot_key(header: &Header) -> Result<&'static [u8], KugouError> {
+    match header.key_slot {
+        1 => Ok(b"l,/'"),
+        slot => Err(KugouError::UnsupportedKeySlot(slot)),
+    }
+}
diff --git a/um_crypto/kgm/src/v2.rs b/um_crypto/kgm/src/v2.rs
index 84fc587..9cbc684 100644
--- a/um_crypto/kgm/src/v2.rs
+++ b/um_crypto/kgm/src/v2.rs
@@ -1,4 +1,5 @@
 use crate::header::Header;
+use crate::slot_keys::get_slot_key;
 use crate::KugouError;
 
 pub struct DecipherV2 {
@@ -6,9 +7,9 @@ pub struct DecipherV2 {
 }
 
 impl DecipherV2 {
-    pub fn new(_header: &Header, slot_key: &[u8]) -> Result<Self, KugouError> {
+    pub fn new(header: &Header) -> Result<Self, KugouError> {
         let mut key = [0u8; 4];
-        key.copy_from_slice(slot_key);
+        key.copy_from_slice(get_slot_key(&header)?);
         Ok(Self { key })
     }
 
@@ -22,7 +23,7 @@ impl DecipherV2 {
 #[test]
 fn test_v2_init() -> Result<(), KugouError> {
     let hdr = Header::from_buffer(include_bytes!("__fixtures__/kgm_v2_hdr.bin"))?;
-    DecipherV2::new(&hdr, b"1234")?;
+    DecipherV2::new(&hdr)?;
 
     Ok(())
 }
diff --git a/um_crypto/kgm/src/v3.rs b/um_crypto/kgm/src/v3.rs
index 687a495..53c2301 100644
--- a/um_crypto/kgm/src/v3.rs
+++ b/um_crypto/kgm/src/v3.rs
@@ -1,4 +1,5 @@
 use crate::header::Header;
+use crate::slot_keys::get_slot_key;
 use crate::KugouError;
 
 pub struct DecipherV3 {
@@ -17,8 +18,8 @@ impl DecipherV3 {
         result
     }
 
-    pub fn new(header: &Header, slot_key: &[u8]) -> Result<Self, KugouError> {
-        let slot_key = Self::hash_key(slot_key);
+    pub fn new(header: &Header) -> Result<Self, KugouError> {
+        let slot_key = Self::hash_key(get_slot_key(header)?);
 
         let mut file_key = [0x6b; 17];
         file_key[..16].copy_from_slice(&Self::hash_key(header.file_key));