use std::collections::HashMap;
use async_trait::async_trait;
use crate::{
security::kms::{
error::{KmsError, KmsResult},
models::{DataKeyPair, EncryptedData, KeyPurpose, KeyReference, RotationPolicy},
},
utils::clock::{Clock, SystemClock},
};
#[async_trait]
pub trait BaseKmsProvider: Send + Sync {
fn provider_name(&self) -> &str;
fn timestamp_secs(&self) -> i64 {
SystemClock.now_secs_i64()
}
async fn encrypt(
&self,
plaintext: &[u8],
key_id: &str,
context: Option<HashMap<String, String>>,
) -> KmsResult<EncryptedData> {
let ctx = context.unwrap_or_default();
let (ciphertext, algorithm) =
self.do_encrypt(plaintext, key_id, &ctx).await.map_err(|e| {
KmsError::EncryptionFailed {
message: format!("Provider encryption failed: {}", e),
}
})?;
Ok(EncryptedData::new(
ciphertext,
KeyReference::new(
self.provider_name().to_string(),
key_id.to_string(),
KeyPurpose::EncryptDecrypt,
self.timestamp_secs(),
),
algorithm,
self.timestamp_secs(),
ctx,
))
}
async fn decrypt(
&self,
encrypted: &EncryptedData,
context: Option<HashMap<String, String>>,
) -> KmsResult<Vec<u8>> {
let ctx = context.unwrap_or_else(|| encrypted.context.clone());
let key_id = &encrypted.key_reference.key_id;
self.do_decrypt(&encrypted.ciphertext, key_id, &ctx).await.map_err(|e| {
KmsError::DecryptionFailed {
message: format!("Provider decryption failed: {}", e),
}
})
}
async fn generate_data_key(
&self,
key_id: &str,
context: Option<HashMap<String, String>>,
) -> KmsResult<DataKeyPair> {
let ctx = context.unwrap_or_default();
let (plaintext_key, encrypted_key_bytes) = self
.do_generate_data_key(key_id, &ctx)
.await
.map_err(|e| KmsError::EncryptionFailed {
message: format!("Data key generation failed: {}", e),
})?;
let key_ref = KeyReference::new(
self.provider_name().to_string(),
key_id.to_string(),
KeyPurpose::EncryptDecrypt,
self.timestamp_secs(),
);
Ok(DataKeyPair::new(
plaintext_key,
EncryptedData::new(
encrypted_key_bytes,
key_ref.clone(),
"data-key".to_string(),
self.timestamp_secs(),
ctx,
),
key_ref,
))
}
async fn rotate_key(&self, key_id: &str) -> KmsResult<KeyReference> {
self.do_rotate_key(key_id).await.map_err(|e| KmsError::RotationFailed {
message: format!("Provider rotation failed: {}", e),
})?;
self.get_key_info(key_id).await
}
async fn get_key_info(&self, key_id: &str) -> KmsResult<KeyReference> {
let info = self.do_get_key_info(key_id).await.map_err(|_e| KmsError::KeyNotFound {
key_id: key_id.to_string(),
})?;
Ok(KeyReference::new(
self.provider_name().to_string(),
key_id.to_string(),
KeyPurpose::EncryptDecrypt,
info.created_at,
)
.with_alias(info.alias.unwrap_or_default()))
}
async fn get_rotation_policy(&self, key_id: &str) -> KmsResult<RotationPolicy> {
let policy =
self.do_get_rotation_policy(key_id).await.map_err(|_e| KmsError::KeyNotFound {
key_id: key_id.to_string(),
})?;
Ok(RotationPolicy {
enabled: policy.enabled,
rotation_period_days: policy.rotation_period_days,
last_rotation: policy.last_rotation,
next_rotation: policy.next_rotation,
})
}
async fn do_encrypt(
&self,
plaintext: &[u8],
key_id: &str,
context: &HashMap<String, String>,
) -> KmsResult<(String, String)>;
async fn do_decrypt(
&self,
ciphertext: &str,
key_id: &str,
context: &HashMap<String, String>,
) -> KmsResult<Vec<u8>>;
async fn do_generate_data_key(
&self,
key_id: &str,
context: &HashMap<String, String>,
) -> KmsResult<(Vec<u8>, String)>;
async fn do_rotate_key(&self, key_id: &str) -> KmsResult<()>;
async fn do_get_key_info(&self, key_id: &str) -> KmsResult<KeyInfo>;
async fn do_get_rotation_policy(&self, key_id: &str) -> KmsResult<RotationPolicyInfo>;
}
pub type ArcKmsProvider = std::sync::Arc<dyn BaseKmsProvider>;
#[derive(Debug, Clone)]
pub struct KeyInfo {
pub alias: Option<String>,
pub created_at: i64,
}
#[derive(Debug, Clone)]
pub struct RotationPolicyInfo {
pub enabled: bool,
pub rotation_period_days: u32,
pub last_rotation: Option<i64>,
pub next_rotation: Option<i64>,
}