From cbcd319aaf7ea749f6997357e827ff085c8fc2c1 Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Tue, 20 Sep 2022 13:41:09 +0000
Subject: [PATCH] Update

---
 src/server/auth/pubkey.rs | 36 ++++++++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/src/server/auth/pubkey.rs b/src/server/auth/pubkey.rs
index a134f76..f966bbf 100644
--- a/src/server/auth/pubkey.rs
+++ b/src/server/auth/pubkey.rs
@@ -22,22 +22,23 @@ impl RSAKey {
     }
 
     pub fn decrypt(&self, data: &[u8]) -> Result<BytesMut, openssl::error::ErrorStack> {
-        let tosize = data.len();
+        let tosize = self.key.size() as usize;
         let mut buf = BytesMut::with_capacity(tosize);
         buf.resize(tosize, 0);
-        let mut i = 0;
-        let mut real = 0;
-        while i < tosize {
-            let i2 = i + self.key.size() as usize;
-            real += self
-                .key
-                .private_decrypt(&data[i..i2], &mut buf[i..i2], Padding::PKCS1)?;
-            i = i2;
-        }
+        let real = self.key.private_decrypt(&data, &mut buf, Padding::PKCS1)?;
         buf.truncate(real);
         Ok(buf)
     }
 
+    #[cfg(test)]
+    pub fn encrypt(&self, data: &[u8]) -> Result<BytesMut, openssl::error::ErrorStack> {
+        let tosize = self.key.size() as usize;
+        let mut buf = BytesMut::with_capacity(tosize);
+        buf.resize(tosize, 0);
+        self.key.public_encrypt(&data, &mut buf, Padding::PKCS1)?;
+        Ok(buf)
+    }
+
     pub fn is_too_old(&self) -> bool {
         self.generated_time < Utc::now() - chrono::Duration::hours(1)
     }
@@ -120,3 +121,18 @@ impl MatchRoute<Body, Body> for AuthPubkeyRoute {
         }
     }
 }
+
+#[test]
+fn test_rsa_decrypt() {
+    let key = RSAKey::new().unwrap();
+    let data = b"Hello, world!";
+    let enc = key.encrypt(data).unwrap();
+    let dec = key.decrypt(&enc).unwrap();
+    assert_eq!(data, &dec[..]);
+    let mut data = BytesMut::with_capacity(256);
+    data.resize(256, 0);
+    openssl::rand::rand_bytes(&mut data).unwrap();
+    let enc = key.encrypt(&data).unwrap();
+    let dec = key.decrypt(&enc).unwrap();
+    assert_eq!(data, dec);
+}