From 4431bff9e61a1ec2e94a2af5444801f719207b6a Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 10 May 2025 00:17:32 +0800
Subject: [PATCH] feat: update zeroize

---
 src/pbeutil.rs | 30 ++++++++++++------------------
 1 file changed, 12 insertions(+), 18 deletions(-)

diff --git a/src/pbeutil.rs b/src/pbeutil.rs
index 6f37a45..799a658 100644
--- a/src/pbeutil.rs
+++ b/src/pbeutil.rs
@@ -4,6 +4,7 @@ use crate::util::{base64_decode, base64_encode, base64_encode_url_safe_no_pad};
 use aes_gcm_stream::{Aes256GcmStreamDecryptor, Aes256GcmStreamEncryptor};
 use rand::random;
 use rust_util::XResult;
+use secrecy::Zeroize;
 
 const PBE_ENC_PREFIX: &str = "pbe_enc:";
 
@@ -17,42 +18,35 @@ pub fn simple_pbe_decrypt_with_prompt_to_string(pin_opt: &mut Option<String>, ci
 }
 
 pub fn simple_pbe_encrypt_with_prompt(iteration: u32, plaintext: &[u8], password_opt: &mut Option<String>, password_double_check: bool) -> XResult<String> {
-    let pin = match password_opt {
+    let mut pin = match password_opt {
         None => {
             let pin1 = opt_value_result!(pinutil::get_pin(None), "Simple PBE password required");
             if password_double_check {
-                let pin2 = opt_value_result!(pinutil::get_pin(None), "Simple PBE password required");
+                let mut pin2 = opt_value_result!(pinutil::get_pin(None), "Simple PBE password required");
                 if pin1 != pin2 {
                     return simple_error!("Two PINs mismatch");
                 }
+                pin2.zeroize();
             }
             *password_opt = Some(pin1.clone());
             pin1
         }
         Some(pin) => pin.clone(),
     };
-    simple_pbe_encrypt(&pin, iteration, plaintext)
+    let encrypt_result = simple_pbe_encrypt(&pin, iteration, plaintext);
+    pin.zeroize();
+    encrypt_result
 }
 
 pub fn simple_pbe_decrypt_with_prompt(pin_opt: &mut Option<String>, ciphertext: &str) -> XResult<Vec<u8>> {
-    let pin = opt_value_result!(pinutil::get_pin(pin_opt.clone().as_deref()), "Simple PBE password required");
+    let mut pin = opt_value_result!(pinutil::get_pin(pin_opt.clone().as_deref()), "Simple PBE password required");
+    pin_opt.zeroize();
     *pin_opt = Some(pin.clone());
-    simple_pbe_decrypt(&pin, ciphertext)
+    let decrypt_result = simple_pbe_decrypt(&pin, ciphertext);
+    pin.zeroize();
+    decrypt_result
 }
 
-// pub fn simple_pbe_encrypt_from_string(
-//     password: &str,
-//     iteration: u32,
-//     plaintext: &str,
-// ) -> XResult<String> {
-//     simple_pbe_encrypt(password, iteration, plaintext.as_bytes())
-// }
-
-// pub fn simple_pbe_decrypt_to_string(password: &str, ciphertext: &str) -> XResult<String> {
-//     let plaintext = simple_pbe_decrypt(password, ciphertext)?;
-//     Ok(String::from_utf8(plaintext)?)
-// }
-
 pub fn simple_pbe_encrypt(password: &str, iteration: u32, plaintext: &[u8]) -> XResult<String> {
     let pbe_salt: [u8; 16] = random();
     let key = simple_pbe_kdf(password, &pbe_salt, iteration)?;