From 6791dec74549ca3fe8732abbe4308210a56d6f94 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 13:51:03 +0100
Subject: [PATCH] [qmc2/rc4] fix: segment key calculation.

---
 um_crypto/qmc/src/lib.rs                |  2 +-
 um_crypto/qmc/src/v2_rc4/cipher.rs      |  4 ++--
 um_crypto/qmc/src/v2_rc4/segment_key.rs | 14 ++++++++++++--
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/um_crypto/qmc/src/lib.rs b/um_crypto/qmc/src/lib.rs
index e5e0883..98f1b42 100644
--- a/um_crypto/qmc/src/lib.rs
+++ b/um_crypto/qmc/src/lib.rs
@@ -29,7 +29,7 @@ impl QMCv2Cipher {
         let key = key.as_ref();
         let cipher = match key.len() {
             0 => Err(QmcCryptoError::QMCV2MapKeyEmpty)?,
-            ..=300 => QMCv2Cipher::MapL(QMC2Map::new(key)?),
+            1..=300 => QMCv2Cipher::MapL(QMC2Map::new(key)?),
             _ => QMCv2Cipher::RC4(QMC2RC4::new(key)),
         };
         Ok(cipher)
diff --git a/um_crypto/qmc/src/v2_rc4/cipher.rs b/um_crypto/qmc/src/v2_rc4/cipher.rs
index e23e02a..40e9d01 100644
--- a/um_crypto/qmc/src/v2_rc4/cipher.rs
+++ b/um_crypto/qmc/src/v2_rc4/cipher.rs
@@ -11,7 +11,7 @@ const RC4_STREAM_CACHE_SIZE: usize = OTHER_SEGMENT_SIZE + 512;
 pub struct QMC2RC4 {
     hash: f64,
     key: Box<[u8]>,
-    key_stream: [u8; RC4_STREAM_CACHE_SIZE],
+    key_stream: Box<[u8; RC4_STREAM_CACHE_SIZE]>,
 }
 
 impl QMC2RC4 {
@@ -23,7 +23,7 @@ impl QMC2RC4 {
         Self {
             hash: hash(key),
             key: key.into(),
-            key_stream,
+            key_stream: key_stream.into(),
         }
     }
 
diff --git a/um_crypto/qmc/src/v2_rc4/segment_key.rs b/um_crypto/qmc/src/v2_rc4/segment_key.rs
index 104af3a..f3010d4 100644
--- a/um_crypto/qmc/src/v2_rc4/segment_key.rs
+++ b/um_crypto/qmc/src/v2_rc4/segment_key.rs
@@ -1,9 +1,9 @@
-pub fn get_segment_key(id: u64, seed: u8, hash: f64) -> u32 {
+pub fn get_segment_key(id: u64, seed: u8, hash: f64) -> u64 {
     match seed {
         0 => 0,
         seed => {
             let result = hash / ((id + 1).wrapping_mul(seed.into()) as f64) * 100.0;
-            result as u32
+            result as u64
         }
     }
 }
@@ -21,4 +21,14 @@ mod tests {
     fn test_segment_key_123() {
         assert_eq!(get_segment_key(1, 123, 12345.0), 5018);
     }
+
+    #[test]
+    fn test_segment_key_large_1() {
+        assert_eq!(get_segment_key(51, 35, 516402887.0), 28373784);
+    }
+
+    #[test]
+    fn test_segment_key_large_2() {
+        assert_eq!(get_segment_key(0, 66, 3908240000.0), 5921575757);
+    }
 }