refactor: replace EncryptionGuard with with_encryption for better async context handling
This commit is contained in:
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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}"),
|
||||
|
||||
Reference in New Issue
Block a user