bug fix + format

This commit is contained in:
Sosthene 2024-12-09 16:05:29 +01:00
parent ccb7c1353e
commit f4a3f38096

View File

@ -1,15 +1,15 @@
use async_std::fs::{create_dir_all, read_dir, read_to_string, remove_file, File};
use std::time::{SystemTime, UNIX_EPOCH, Duration};
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use async_std::io::WriteExt;
use async_std::path::Path;
use async_std::stream::{IntoStream, StreamExt};
use async_std::task;
use base64;
use sdk_common::log;
use serde_json::json;
use tide::{Request, Response, StatusCode};
use base64;
const STORAGE_DIR: &str = "./storage";
const MIN_TTL: u64 = 60; // 1 minute
@ -68,7 +68,8 @@ struct Metadata {
/// Converts a `SystemTime` to a UNIX timestamp (seconds since UNIX epoch).
fn system_time_to_unix(system_time: SystemTime) -> u64 {
system_time.duration_since(UNIX_EPOCH)
system_time
.duration_since(UNIX_EPOCH)
.expect("SystemTime before UNIX_EPOCH!")
.as_secs()
}
@ -110,23 +111,41 @@ async fn store_data(key: &str, value: &[u8], expires_at: SystemTime) -> Result<(
// Check if key exists
if file_path.exists().await {
return Err(tide::Error::from_str(StatusCode::Conflict, "Key already exists"));
return Err(tide::Error::from_str(
StatusCode::Conflict,
"Key already exists",
));
}
create_dir_all(file_path).await.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
create_dir_all(file_path.parent().ok_or(tide::Error::from_str(
StatusCode::InternalServerError,
"File path doesn't have parent",
))?)
.await
.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
let metadata_path = format!("{}.meta", file_name);
let mut file = File::create(&file_path).await.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
file.write_all(value).await.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
let mut file = File::create(&file_path)
.await
.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
file.write_all(value)
.await
.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
let metadata = Metadata {
expires_at: system_time_to_unix(expires_at)
expires_at: system_time_to_unix(expires_at),
};
let metadata_json = serde_json::to_string(&metadata).map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
let mut meta_file = File::create(&metadata_path).await.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
meta_file.write_all(metadata_json.as_bytes()).await.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
let metadata_json = serde_json::to_string(&metadata)
.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
let mut meta_file = File::create(&metadata_path)
.await
.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
meta_file
.write_all(metadata_json.as_bytes())
.await
.map_err(|e| tide::Error::new(StatusCode::InternalServerError, e))?;
Ok(())
}
@ -134,9 +153,13 @@ async fn store_data(key: &str, value: &[u8], expires_at: SystemTime) -> Result<(
async fn retrieve_data(key: &str) -> Result<Vec<u8>, String> {
let file_path = format!("{}/{}/{}", STORAGE_DIR, &key[..2], &key[2..]);
let mut file = File::open(&file_path).await.map_err(|_| "Key not found.".to_string())?;
let mut file = File::open(&file_path)
.await
.map_err(|_| "Key not found.".to_string())?;
let mut buffer = Vec::new();
async_std::io::ReadExt::read_to_end(&mut file, &mut buffer).await.map_err(|e| e.to_string())?;
async_std::io::ReadExt::read_to_end(&mut file, &mut buffer)
.await
.map_err(|e| e.to_string())?;
Ok(buffer)
}
@ -163,7 +186,10 @@ async fn handle_store(mut req: Request<()>) -> tide::Result<Response> {
let live_for = if let Some(ttl) = data.ttl {
if ttl < MIN_TTL {
return Ok(Response::builder(StatusCode::BadRequest)
.body(format!("Invalid ttl: must be at least {} seconds.", MIN_TTL))
.body(format!(
"Invalid ttl: must be at least {} seconds.",
MIN_TTL
))
.build());
} else if ttl > MAX_TTL {
return Ok(Response::builder(StatusCode::BadRequest)
@ -176,7 +202,9 @@ async fn handle_store(mut req: Request<()>) -> tide::Result<Response> {
};
let now = SystemTime::now();
let expires_at = now.checked_add(live_for).ok_or(tide::Error::from_str(StatusCode::BadRequest, "Invalid ttl"))?;
let expires_at = now
.checked_add(live_for)
.ok_or(tide::Error::from_str(StatusCode::BadRequest, "Invalid ttl"))?;
// Decode the value from Base64
let value_bytes = match base64::decode(&data.value) {
@ -190,20 +218,16 @@ async fn handle_store(mut req: Request<()>) -> tide::Result<Response> {
// Store the data
match store_data(&data.key, &value_bytes, expires_at).await {
Ok(()) => {
Ok(Response::builder(StatusCode::Ok)
Ok(()) => Ok(Response::builder(StatusCode::Ok)
.body(serde_json::to_value(&ApiResponse {
message: "Data stored successfully.".to_string(),
})?)
.build())
}
Err(e) => {
Ok(Response::builder(e.status())
.build()),
Err(e) => Ok(Response::builder(e.status())
.body(serde_json::to_value(&ApiResponse {
message: e.to_string(),
})?)
.build())
}
.build()),
}
}
@ -220,24 +244,32 @@ async fn handle_retrieve(req: Request<()>) -> tide::Result<Response> {
Ok(value) => {
let encoded_value = base64::encode(&value);
Ok(Response::builder(StatusCode::Ok)
.body(serde_json::to_value(&RetrieveResponse { key, value: encoded_value })?)
.body(serde_json::to_value(&RetrieveResponse {
key,
value: encoded_value,
})?)
.build())
}
Err(e) => Ok(Response::builder(StatusCode::NotFound)
.body(e)
.build()),
Err(e) => Ok(Response::builder(StatusCode::NotFound).body(e).build()),
}
}
/// Checks a metadata file and deletes the associated data file if expired
async fn handle_file_cleanup(now: u64, meta_path: &Path) -> Result<(), String> {
let meta_content = read_to_string(meta_path).await.map_err(|e| format!("Failed to read metadata: {}", e.to_string()))?;
let metadata: Metadata = serde_json::from_str(&meta_content).map_err(|e| format!("Failed to parse metadata: {}", e.to_string()))?;
let meta_content = read_to_string(meta_path)
.await
.map_err(|e| format!("Failed to read metadata: {}", e.to_string()))?;
let metadata: Metadata = serde_json::from_str(&meta_content)
.map_err(|e| format!("Failed to parse metadata: {}", e.to_string()))?;
if metadata.expires_at < now {
let data_file_path = meta_path.with_extension("");
remove_file(&data_file_path).await.map_err(|e| format!("Failed to remove data file: {}", e.to_string()))?;
remove_file(meta_path).await.map_err(|e| format!("Failed to remove metadata file: {}", e.to_string()))?;
remove_file(&data_file_path)
.await
.map_err(|e| format!("Failed to remove data file: {}", e.to_string()))?;
remove_file(meta_path)
.await
.map_err(|e| format!("Failed to remove metadata file: {}", e.to_string()))?;
log::debug!("Removed expired file: {:?}", data_file_path);
}
Ok(())
@ -246,7 +278,9 @@ async fn handle_file_cleanup(now: u64, meta_path: &Path) -> Result<(), String> {
#[async_std::main]
async fn main() -> tide::Result<()> {
sdk_common::env_logger::init();
create_dir_all(STORAGE_DIR).await.expect("Failed to create storage directory.");
create_dir_all(STORAGE_DIR)
.await
.expect("Failed to create storage directory.");
task::spawn(cleanup_expired_files());