use std::sync::Arc;
use base64::Engine as _;
use crate::error::{Error, Result};
use crate::manager::EncryptionManager;
const MASTER_KEY_ENV: &str = "ENCRYPTION_MASTER_KEY";
const ALGORITHM_ENV: &str = "ENCRYPTION_ALGORITHM";
const DEFAULT_ALGORITHM: &str = "aes-256-gcm";
#[derive(Debug, Clone)]
pub struct Guard {
manager: Arc<EncryptionManager>,
}
impl Guard {
pub fn new(manager: EncryptionManager) -> Self {
Self {
manager: Arc::new(manager),
}
}
pub fn from_master_key(master_key: &[u8], default_identifier: &str) -> Result<Self> {
Ok(Self::new(
crate::factory::EncryptionManagerFactory::from_master_key(
master_key,
default_identifier,
)?,
))
}
pub fn from_env() -> Result<Self> {
Self::from_env_with(|key| std::env::var(key).ok())
}
pub fn from_env_with<F: Fn(&str) -> Option<String>>(lookup: F) -> Result<Self> {
let raw = lookup(MASTER_KEY_ENV).ok_or(Error::MissingConfig {
key: MASTER_KEY_ENV,
})?;
let algorithm = lookup(ALGORITHM_ENV).unwrap_or_else(|| DEFAULT_ALGORITHM.to_owned());
Self::from_master_key(&decode_master_key(&raw)?, &algorithm)
}
pub fn manager(&self) -> &EncryptionManager {
&self.manager
}
pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>> {
self.manager.encrypt(plaintext)
}
pub fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>> {
self.manager.decrypt(ciphertext)
}
pub fn encrypt_with(&self, identifier: &str, plaintext: &[u8]) -> Result<Vec<u8>> {
self.manager.encrypt_with(identifier, plaintext)
}
pub fn decrypt_with(&self, identifier: &str, ciphertext: &[u8]) -> Result<Vec<u8>> {
self.manager.decrypt_with(identifier, ciphertext)
}
pub fn identifiers(&self) -> Vec<&str> {
self.manager.registry().identifiers()
}
}
fn decode_master_key(raw: &str) -> Result<[u8; 32]> {
let value = raw.trim();
let (is_hex, payload) = if let Some(hex) = value.strip_prefix("hex:") {
(true, hex)
} else if let Some(b64) = value.strip_prefix("base64:") {
(false, b64)
} else if value.len() == 64 && value.bytes().all(|b| b.is_ascii_hexdigit()) {
(true, value)
} else {
(false, value)
};
let decoded = if is_hex {
crate::internal::hex_decode(payload).ok()
} else {
base64::engine::general_purpose::STANDARD
.decode(payload)
.ok()
}
.ok_or(Error::InvalidMasterKey { got: 0 })?;
decoded
.try_into()
.map_err(|v: Vec<u8>| Error::InvalidMasterKey { got: v.len() })
}
#[cfg(test)]
mod tests {
use super::*;
fn guard() -> Guard {
Guard::from_master_key(&[7u8; 32], "aes-256-gcm").unwrap()
}
#[test]
fn wraps_the_factory_and_forwards() {
let guard = guard();
assert_eq!(guard.identifiers().len(), 5);
assert_eq!(guard.manager().default_identifier(), "aes-256-gcm");
let blob = guard.encrypt(b"phone").unwrap();
assert_eq!(guard.decrypt(&blob).unwrap(), b"phone");
let blob = guard.encrypt_with("sm4-cbc", b"phone").unwrap();
assert_eq!(guard.decrypt_with("sm4-cbc", &blob).unwrap(), b"phone");
}
#[test]
fn clones_share_one_manager() {
let guard = guard();
let clone = guard.clone();
let blob = guard.encrypt(b"shared").unwrap();
assert_eq!(clone.decrypt(&blob).unwrap(), b"shared");
assert_eq!(Arc::strong_count(&guard.manager), 2);
}
#[test]
fn invalid_master_key_fails_at_construction() {
assert!(Guard::from_master_key(&[0u8; 16], "aes-256-gcm").is_err());
assert!(Guard::from_master_key(&[0u8; 32], "no-such-algorithm").is_err());
}
fn fake_env(entries: Vec<(&str, String)>) -> impl Fn(&str) -> Option<String> {
move |key| {
entries
.iter()
.find(|(k, _)| *k == key)
.map(|(_, v)| v.clone())
}
}
fn roundtrip(guard: &Guard) {
let blob = guard.encrypt(b"phone").unwrap();
assert_eq!(guard.decrypt(&blob).unwrap(), b"phone");
}
#[test]
fn from_env_hex_prefix() {
let hex = "07".repeat(32);
let lookup = fake_env(vec![(MASTER_KEY_ENV, format!("hex:{hex}"))]);
let guard = Guard::from_env_with(lookup).unwrap();
assert_eq!(guard.manager().default_identifier(), "aes-256-gcm");
roundtrip(&guard);
}
#[test]
fn from_env_base64_prefix() {
let b64 = "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=";
let lookup = fake_env(vec![(MASTER_KEY_ENV, format!("base64:{b64}"))]);
let guard = Guard::from_env_with(lookup).unwrap();
roundtrip(&guard);
}
#[test]
fn from_env_bare_hex_auto_detected() {
let hex = "00".repeat(32);
let lookup = fake_env(vec![(MASTER_KEY_ENV, hex)]);
let guard = Guard::from_env_with(lookup).unwrap();
roundtrip(&guard);
}
#[test]
fn from_env_missing_master_key() {
let err = Guard::from_env_with(fake_env(vec![])).unwrap_err();
assert_eq!(
err,
Error::MissingConfig {
key: "ENCRYPTION_MASTER_KEY"
}
);
}
#[test]
fn from_env_wrong_length() {
let lookup = fake_env(vec![(MASTER_KEY_ENV, format!("hex:{}", "00".repeat(16)))]);
assert_eq!(
Guard::from_env_with(lookup).unwrap_err(),
Error::InvalidMasterKey { got: 16 }
);
let lookup = fake_env(vec![(MASTER_KEY_ENV, format!("base64:{}", "A".repeat(32)))]);
assert_eq!(
Guard::from_env_with(lookup).unwrap_err(),
Error::InvalidMasterKey { got: 24 }
);
let lookup = fake_env(vec![(MASTER_KEY_ENV, "hex:zz".to_owned())]);
assert_eq!(
Guard::from_env_with(lookup).unwrap_err(),
Error::InvalidMasterKey { got: 0 }
);
}
#[test]
fn from_env_respects_algorithm() {
let lookup = fake_env(vec![
(MASTER_KEY_ENV, "07".repeat(32)),
(ALGORITHM_ENV, "sm4-cbc".to_owned()),
]);
let guard = Guard::from_env_with(lookup).unwrap();
assert_eq!(guard.manager().default_identifier(), "sm4-cbc");
roundtrip(&guard);
}
#[test]
fn from_env_unknown_algorithm_bubbles() {
let lookup = fake_env(vec![
(MASTER_KEY_ENV, "07".repeat(32)),
(ALGORITHM_ENV, "no-such-algorithm".to_owned()),
]);
assert!(matches!(
Guard::from_env_with(lookup).unwrap_err(),
Error::UnknownIdentifier { .. }
));
}
}