refactor: replace EncryptionGuard with with_encryption for better async context handling

This commit is contained in:
Tunglies
2025-11-05 16:56:04 +08:00
parent 8c3e9c9ea9
commit ea41e71f72
2 changed files with 16 additions and 30 deletions

View File

@@ -6,9 +6,14 @@ use aes_gcm::{
use base64::{Engine, engine::general_purpose::STANDARD}; use base64::{Engine, engine::general_purpose::STANDARD};
use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::cell::Cell; use std::cell::Cell;
use std::future::Future;
const NONCE_LENGTH: usize = 12; 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 /// Encrypt data
#[allow(deprecated)] #[allow(deprecated)]
@@ -92,31 +97,14 @@ where
} }
} }
pub struct EncryptionGuard; pub async fn with_encryption<F, Fut, R>(f: F) -> R
where
impl EncryptionGuard { F: FnOnce() -> Fut,
pub fn new() -> Self { Fut: Future<Output = R>,
ENCRYPTION_CONTEXT.with(|ctx| { {
ctx.set(true); ENCRYPTION_ACTIVE.scope(Cell::new(true), f()).await
});
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);
});
}
} }
fn is_encryption_active() -> bool { fn is_encryption_active() -> bool {
ENCRYPTION_CONTEXT.with(|ctx| ctx.get()) ENCRYPTION_ACTIVE.try_with(|c| c.get()).unwrap_or(false)
} }

View File

@@ -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 anyhow::{Context, Result, anyhow, bail};
use nanoid::nanoid; use nanoid::nanoid;
use serde::{Serialize, de::DeserializeOwned}; 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()); bail!("file not found \"{}\"", path.display());
} }
let _guard = EncryptionGuard::new();
let yaml_str = tokio::fs::read_to_string(path).await?; 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 /// read mapping from yaml
@@ -66,8 +65,7 @@ pub async fn save_yaml<T: Serialize + Sync>(
data: &T, data: &T,
prefix: Option<&str>, prefix: Option<&str>,
) -> Result<()> { ) -> Result<()> {
let _guard = EncryptionGuard::new(); let data_str = with_encryption(|| async { serde_yaml_ng::to_string(data) }).await?;
let data_str = serde_yaml_ng::to_string(data)?;
let yaml_str = match prefix { let yaml_str = match prefix {
Some(prefix) => format!("{prefix}\n\n{data_str}"), Some(prefix) => format!("{prefix}\n\n{data_str}"),