use async_trait::async_trait;
use zeroize::Zeroizing;
use crate::crypto::encryption::Encryptor;
use crate::error::{AppError, AppResult};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WrappedDek {
pub provider: String,
pub key_id: String,
pub key_version: String,
pub ciphertext: Vec<u8>,
}
#[async_trait]
pub trait KeyManager: Send + Sync {
async fn wrap_dek(&self, dek: &[u8]) -> AppResult<WrappedDek>;
async fn unwrap_dek(&self, wrapped: &WrappedDek) -> AppResult<Zeroizing<Vec<u8>>>;
fn provider(&self) -> &str;
}
pub struct LocalDevKms {
kek: Encryptor,
key_id: String,
key_version: String,
}
impl LocalDevKms {
pub fn from_master_key_base64(master_key_base64: &str) -> AppResult<Self> {
let kek = Encryptor::from_base64(master_key_base64)?;
Ok(Self {
kek,
key_id: "env:NOETL_ENCRYPTION_KEY".to_string(),
key_version: "v1".to_string(),
})
}
}
#[async_trait]
impl KeyManager for LocalDevKms {
async fn wrap_dek(&self, dek: &[u8]) -> AppResult<WrappedDek> {
let ciphertext = self.kek.encrypt(dek)?;
Ok(WrappedDek {
provider: "local".to_string(),
key_id: self.key_id.clone(),
key_version: self.key_version.clone(),
ciphertext,
})
}
async fn unwrap_dek(&self, wrapped: &WrappedDek) -> AppResult<Zeroizing<Vec<u8>>> {
if wrapped.provider != "local" {
return Err(AppError::Encryption(format!(
"LocalDevKms cannot unwrap a DEK from provider '{}'",
wrapped.provider
)));
}
let dek = self.kek.decrypt(&wrapped.ciphertext)?;
Ok(Zeroizing::new(dek))
}
fn provider(&self) -> &str {
"local"
}
}
#[cfg(test)]
mod tests {
use super::*;
fn kms() -> LocalDevKms {
let key = Encryptor::generate_key_base64();
LocalDevKms::from_master_key_base64(&key).unwrap()
}
#[tokio::test]
async fn wrap_unwrap_round_trips() {
let km = kms();
let dek = Encryptor::generate_key(); let wrapped = km.wrap_dek(&dek).await.unwrap();
assert_eq!(wrapped.provider, "local");
assert_ne!(wrapped.ciphertext, dek);
let unwrapped = km.unwrap_dek(&wrapped).await.unwrap();
assert_eq!(unwrapped.as_slice(), dek.as_slice());
}
#[tokio::test]
async fn unwrap_with_a_different_kek_fails() {
let km1 = kms();
let km2 = kms(); let dek = Encryptor::generate_key();
let wrapped = km1.wrap_dek(&dek).await.unwrap();
assert!(km2.unwrap_dek(&wrapped).await.is_err());
}
#[tokio::test]
async fn rejects_foreign_provider() {
let km = kms();
let dek = Encryptor::generate_key();
let mut wrapped = km.wrap_dek(&dek).await.unwrap();
wrapped.provider = "gcp-kms".to_string();
assert!(km.unwrap_dek(&wrapped).await.is_err());
}
}