From ea41e71f726d3cbf4353961bd772978f9187c03d Mon Sep 17 00:00:00 2001
From: Tunglies <77394545+Tunglies@users.noreply.github.com>
Date: Wed, 5 Nov 2025 16:56:04 +0800
Subject: [PATCH] refactor: replace EncryptionGuard with with_encryption for
 better async context handling

---
 src-tauri/src/config/encrypt.rs | 38 +++++++++++----------------------
 src-tauri/src/utils/help.rs     |  8 +++----
 2 files changed, 16 insertions(+), 30 deletions(-)

diff --git a/src-tauri/src/config/encrypt.rs b/src-tauri/src/config/encrypt.rs
index 4e6b9dec..44ac85b1 100644
--- a/src-tauri/src/config/encrypt.rs
+++ b/src-tauri/src/config/encrypt.rs
@@ -6,9 +6,14 @@ use aes_gcm::{
 use base64::{Engine, engine::general_purpose::STANDARD};
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::cell::Cell;
+use std::future::Future;
 
 const NONCE_LENGTH: usize = 12;
-thread_local!(static ENCRYPTION_CONTEXT: Cell<bool> = const { Cell::new(false) });
+
+// Use task-local context so the flag follows the async task across threads
+tokio::task_local! {
+    static ENCRYPTION_ACTIVE: Cell<bool>;
+}
 
 /// Encrypt data
 #[allow(deprecated)]
@@ -92,31 +97,14 @@ where
     }
 }
 
-pub struct EncryptionGuard;
-
-impl EncryptionGuard {
-    pub fn new() -> Self {
-        ENCRYPTION_CONTEXT.with(|ctx| {
-            ctx.set(true);
-        });
-        EncryptionGuard
-    }
-}
-
-impl Default for EncryptionGuard {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl Drop for EncryptionGuard {
-    fn drop(&mut self) {
-        ENCRYPTION_CONTEXT.with(|ctx| {
-            ctx.set(false);
-        });
-    }
+pub async fn with_encryption<F, Fut, R>(f: F) -> R
+where
+    F: FnOnce() -> Fut,
+    Fut: Future<Output = R>,
+{
+    ENCRYPTION_ACTIVE.scope(Cell::new(true), f()).await
 }
 
 fn is_encryption_active() -> bool {
-    ENCRYPTION_CONTEXT.with(|ctx| ctx.get())
+    ENCRYPTION_ACTIVE.try_with(|c| c.get()).unwrap_or(false)
 }
diff --git a/src-tauri/src/utils/help.rs b/src-tauri/src/utils/help.rs
index 258fea2c..1fc58d03 100644
--- a/src-tauri/src/utils/help.rs
+++ b/src-tauri/src/utils/help.rs
@@ -1,4 +1,4 @@
-use crate::{config::EncryptionGuard, enhance::seq::SeqMap, logging, utils::logging::Type};
+use crate::{config::with_encryption, enhance::seq::SeqMap, logging, utils::logging::Type};
 use anyhow::{Context, Result, anyhow, bail};
 use nanoid::nanoid;
 use serde::{Serialize, de::DeserializeOwned};
@@ -11,10 +11,9 @@ pub async fn read_yaml<T: DeserializeOwned>(path: &PathBuf) -> Result<T> {
         bail!("file not found \"{}\"", path.display());
     }
 
-    let _guard = EncryptionGuard::new();
     let yaml_str = tokio::fs::read_to_string(path).await?;
 
-    Ok(serde_yaml_ng::from_str::<T>(&yaml_str)?)
+    Ok(with_encryption(|| async { serde_yaml_ng::from_str::<T>(&yaml_str) }).await?)
 }
 
 /// read mapping from yaml
@@ -66,8 +65,7 @@ pub async fn save_yaml<T: Serialize + Sync>(
     data: &T,
     prefix: Option<&str>,
 ) -> Result<()> {
-    let _guard = EncryptionGuard::new();
-    let data_str = serde_yaml_ng::to_string(data)?;
+    let data_str = with_encryption(|| async { serde_yaml_ng::to_string(data) }).await?;
 
     let yaml_str = match prefix {
         Some(prefix) => format!("{prefix}\n\n{data_str}"),