use ais_keystore_lib::{grpc, key_store, shared};
use dusa_collection_utils::core::functions::current_timestamp;
use dusa_collection_utils::core::logger::LogLevel;
use dusa_collection_utils::core::types::rwarc::LockWithTimeout;
use dusa_collection_utils::log;
use key_store::store::{generate_new_key, KeyStore};
use shared::consts::{ADDR, PORT};
use shared::messages::{
AppMessage, CommandResponse as StoreResponse,
};
use std::time::Duration;
use tokio::time::sleep;
use tonic::{transport::Server, Request, Response, Status};
use grpc::keystore::key_store_server::{KeyStore as GrpcKeyStoreTrait, KeyStoreServer};
use grpc::keystore::{AddKeyRequest, CommandResponse, GenerateKeyRequest, GenerateKeyResponse, GetKeyRequest, Key, KeyResponse, KeyStatus as GrpcKeyStatus, MigrateKeyRequest};
#[derive(Clone)]
pub struct SharedKeyStore {
inner: LockWithTimeout<KeyStore>,
}
impl SharedKeyStore {
pub async fn new() -> Self {
Self {
inner: LockWithTimeout::new(KeyStore::new().await),
}
}
pub fn spawn_purge_thread(&self, retention_days: i64) {
let store = self.inner.clone();
tokio::spawn(async move {
loop {
{
let mut store = match store.try_write().await {
Ok(data) => data,
Err(err) => {
log!(LogLevel::Error, "{}", err);
continue;
}
};
store.purge_old_archives(retention_days);
}
sleep(std::time::Duration::from_secs(24 * 60 * 60)).await; }
});
}
pub fn spawn_status_update_thread(&self) {
let store = self.inner.clone();
tokio::spawn(async move {
loop {
{
let mut store = match store.try_write().await {
Ok(data) => data,
Err(err) => {
log!(LogLevel::Error, "{}", err);
continue;
}
};
let rotated = store.rotate_expiring_keys();
if !rotated.is_empty() {
log!(
LogLevel::Info,
"Rotated {} key(s) nearing expiration: {}",
rotated.len(),
rotated.join(", ")
);
if let Err(err) = store.save_keys().await {
log!(LogLevel::Error, "Failed to persist rotated keys: {}", err);
}
}
}
sleep(std::time::Duration::from_secs(60 * 60)).await; }
});
}
pub fn spawn_save_keys_thread(&self) {
let key_store = self.inner.clone();
tokio::spawn(async move {
loop {
sleep(Duration::from_secs(300)).await;
let store = match key_store.try_write().await {
Ok(data) => data,
Err(err) => {
log!(LogLevel::Error, "{}", err);
continue;
}
};
match store.save_keys().await {
Ok(_) => {
log!(LogLevel::Info, "Keys saved @ {}", current_timestamp());
continue;
}
Err(err) => {
log!(LogLevel::Error, "Err saving key: {}", err);
continue;
}
}
}
});
}
pub async fn handle_command(&self, command: AppMessage) -> StoreResponse {
let mut store = match self.inner.try_write().await {
Ok(data) => data,
Err(err) => {
log!(LogLevel::Error, "{}", err);
return StoreResponse {
success: false,
message: format!("{}", err),
};
}
};
match command {
AppMessage::AddKey(mut cmd) => {
let mut updated_ttl: bool = false;
if let Some(ttl) = cmd.ttl {
if ttl < 60 {
cmd.ttl = Some(60);
updated_ttl = true;
}
}
store.add_key(cmd.id, cmd.value, cmd.ttl);
return match store.save_keys().await {
Ok(_) => match updated_ttl {
true => StoreResponse {
success: true,
message: "WARN: key ttl increased to 1 min, minimun time".into(),
},
false => StoreResponse {
success: true,
message: "Key added successfully".into(),
},
},
Err(err) => StoreResponse {
success: false,
message: format!("Key added, store saving failed: {}", err),
},
};
}
AppMessage::GetKey(cmd) => {
if let Some(key) = store.get_key(&cmd.id, cmd.version) {
StoreResponse {
success: true,
message: serde_json::to_string(&key).unwrap(),
}
} else {
StoreResponse {
success: false,
message: "Key not found".into(),
}
}
}
AppMessage::MigrateKey(cmd) => {
store.migrate_key(&cmd.id);
StoreResponse {
success: true,
message: "Key migrated successfully".into(),
}
}
AppMessage::GenerateKey(cmd) => {
let key_length = cmd.length.unwrap_or(32);
let new_key = generate_new_key(Some(key_length));
StoreResponse {
success: true,
message: hex::encode(new_key),
}
}
AppMessage::Response(_) => StoreResponse {
success: false,
message: "Invalid command received as a response".into(),
},
}
}
}
pub struct GrpcService {
store: SharedKeyStore,
}
impl GrpcService {
pub fn new(store: SharedKeyStore) -> Self {
Self { store }
}
}
#[tonic::async_trait]
impl GrpcKeyStoreTrait for GrpcService {
async fn add_key(
&self,
request: Request<AddKeyRequest>,
) -> Result<Response<CommandResponse>, Status> {
let req = request.into_inner();
let ttl = if req.ttl == 0 { None } else { Some(req.ttl) };
let msg = AppMessage::AddKey(shared::messages::AddKeyCommand {
id: req.id,
value: req.value,
ttl,
});
let resp = self.store.handle_command(msg).await;
Ok(Response::new(CommandResponse { success: resp.success, message: resp.message }))
}
async fn get_key(
&self,
request: Request<GetKeyRequest>,
) -> Result<Response<KeyResponse>, Status> {
let req = request.into_inner();
let msg = AppMessage::GetKey(shared::messages::GetKeyCommand {
id: req.id,
version: if req.version == 0 { None } else { Some(req.version) },
});
let resp = self.store.handle_command(msg).await;
if !resp.success {
return Ok(Response::new(KeyResponse { success: false, message: resp.message, key: None }));
}
let key: key_store::key::Key = serde_json::from_str(&resp.message)
.map_err(|e| Status::internal(format!("{}", e)))?;
let grpc_key = Key {
id: key.id,
value: key.value,
ttl: key.ttl,
created_at: key.created_at,
expires_at: key.expires_at,
version: key.version,
status: match key.status {
key_store::key::KeyStatus::Active => GrpcKeyStatus::Active as i32,
key_store::key::KeyStatus::NearExpiration => GrpcKeyStatus::NearExpiration as i32,
key_store::key::KeyStatus::Expired => GrpcKeyStatus::Expired as i32,
key_store::key::KeyStatus::Invalid => GrpcKeyStatus::Invalid as i32,
},
};
Ok(Response::new(KeyResponse { success: true, message: String::new(), key: Some(grpc_key) }))
}
async fn migrate_key(
&self,
request: Request<MigrateKeyRequest>,
) -> Result<Response<CommandResponse>, Status> {
let req = request.into_inner();
let msg = AppMessage::MigrateKey(shared::messages::MigrateKeyCommand { id: req.id });
let resp = self.store.handle_command(msg).await;
Ok(Response::new(CommandResponse { success: resp.success, message: resp.message }))
}
async fn generate_key(
&self,
request: Request<GenerateKeyRequest>,
) -> Result<Response<GenerateKeyResponse>, Status> {
let req = request.into_inner();
let msg = AppMessage::GenerateKey(shared::messages::GenerateKeyCommand { length: Some(req.length as usize) });
let resp = self.store.handle_command(msg).await;
if !resp.success {
return Ok(Response::new(GenerateKeyResponse { success: false, key: String::new() }));
}
Ok(Response::new(GenerateKeyResponse { success: true, key: resp.message }))
}
}
#[tokio::main]
async fn main() {
let retention_days = 30;
let store: SharedKeyStore = SharedKeyStore::new().await;
store.spawn_purge_thread(retention_days);
store.spawn_status_update_thread();
store.spawn_save_keys_thread();
let addr = format!("{}:{}", ADDR, PORT).parse().expect("Invalid address");
let service = GrpcService::new(store.clone());
log!(LogLevel::Info, "gRPC KeyStore server listening on {}", addr);
Server::builder()
.add_service(KeyStoreServer::new(service))
.serve(addr)
.await
.unwrap();
}