From 662dc668dbb40b465b4c2b72d2d31c7175db2d8e Mon Sep 17 00:00:00 2001
From: Sosthene <contact@sosthene.net>
Date: Mon, 18 Nov 2024 17:41:03 +0100
Subject: [PATCH] Add simple ttl

---
 src/main.rs | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index bfc50dd..84fba2e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,15 +1,19 @@
 use async_std::fs::{create_dir_all, File};
 use async_std::io::WriteExt;
 use serde::{Deserialize, Serialize};
+use serde_json::json;
 use tide::{Request, Response, StatusCode};
 use base64;
+use std::time::{Duration, SystemTime};
 
 const STORAGE_DIR: &str = "./storage";
+const DEFAULT_TTL: u64 = 86400;
 
 #[derive(Deserialize)]
 struct StoreRequest {
     key: String,
     value: String,
+    ttl: Option<u64>,
 }
 
 #[derive(Serialize)]
@@ -24,18 +28,21 @@ struct RetrieveResponse {
 }
 
 /// Store data on the filesystem
-async fn store_data(key: &str, value: &[u8]) -> Result<(), String> {
-    // Create directory based on the first 2 characters of the key
+async fn store_data(key: &str, value: &[u8], expires_at: SystemTime) -> Result<(), String> {
     let dir_name = format!("{}/{}", STORAGE_DIR, &key[..2]);
     let file_path = format!("{}/{}", dir_name, &key[2..]);
+    let metadata_path = format!("{}.meta", file_path);
 
-    // Create directory if it doesn't exist
     create_dir_all(&dir_name).await.map_err(|e| e.to_string())?;
 
-    // Write value to file
     let mut file = File::create(&file_path).await.map_err(|e| e.to_string())?;
     file.write_all(value).await.map_err(|e| e.to_string())?;
 
+    let metadata = serde_json::to_string(&json!({ "expires_at": expires_at }))
+        .map_err(|e| e.to_string())?;
+    let mut meta_file = File::create(&metadata_path).await.map_err(|e| e.to_string())?;
+    meta_file.write_all(metadata.as_bytes()).await.map_err(|e| e.to_string())?;
+
     Ok(())
 }
 
@@ -67,8 +74,12 @@ async fn handle_store(mut req: Request<()>) -> tide::Result<Response> {
         tide::Error::from_str(StatusCode::BadRequest, "Invalid Base64 value")
     })?;
 
+    let now = SystemTime::now();
+    let live_for = if let Some(ttl) = data.ttl { Duration::from_secs(ttl) } else { Duration::from_secs(DEFAULT_TTL) };
+    let expires_at = now.checked_add(live_for).ok_or(tide::Error::from_str(StatusCode::BadRequest, "Invalid ttl"))?;
+
     // Store the data
-    store_data(&data.key, &value_bytes).await.map_err(|e| {
+    store_data(&data.key, &value_bytes, expires_at).await.map_err(|e| {
         tide::Error::from_str(StatusCode::InternalServerError, e)
     })?;