docs

src/app_state.rs

@@ -1,22 +1,22 @@
+use crate::s3_utils::get_s3_filelist;
 use actix_web::error::ErrorInternalServerError;
 use aws_config::BehaviorVersion;
 use aws_sdk_s3::{config::Credentials, Client as S3Client};
+use log::warn;
 use redis::{aio::MultiplexedConnection, AsyncCommands, Client as RedisClient};
 use std::env;
-use log::warn;
-use crate::s3_utils::get_s3_filelist;
-
 
 #[derive(Clone)]
 pub struct AppState {
     pub redis: MultiplexedConnection,
     pub storj_client: S3Client,
     pub aws_client: S3Client,
-    pub bucket: String
+    pub bucket: String,
 }
 
 const PATH_MAPPING_KEY: &str = "filepath_mapping"; // Ключ для хранения маппинга путей
-const WEEK_SECONDS: u64 = 604800;
+                                                   // Убираем TTL для квоты - она должна быть постоянной на пользователя
+const QUOTA_TTL: u64 = 0; // 0 означает отсутствие TTL
 
 impl AppState {
     /// Инициализация нового состояния приложения.
@@ -40,7 +40,7 @@ impl AppState {
         let aws_access_key = env::var("AWS_ACCESS_KEY").expect("AWS_ACCESS_KEY must be set");
         let aws_secret_key = env::var("AWS_SECRET_KEY").expect("AWS_SECRET_KEY must be set");
         let aws_endpoint =
-           env::var("AWS_END_POINT").unwrap_or_else(|_| "https://s3.amazonaws.com".to_string());
+            env::var("AWS_END_POINT").unwrap_or_else(|_| "https://s3.amazonaws.com".to_string());
 
         // Конфигурируем клиент S3 для Storj
         let storj_config = aws_config::defaults(BehaviorVersion::latest())
@@ -71,14 +71,14 @@ impl AppState {
             ))
             .load()
             .await;
-        
+
         let aws_client = S3Client::new(&aws_config);
 
         let app_state = AppState {
             redis: redis_connection,
             storj_client,
             aws_client,
-            bucket
+            bucket,
         };
 
         // Кэшируем список файлов из AWS при старте приложения
@@ -91,7 +91,7 @@ impl AppState {
     pub async fn cache_filelist(&self) {
         warn!("caching AWS filelist...");
         let mut redis = self.redis.clone();
-        
+
         // Запрашиваем список файлов из Storj S3
         let filelist = get_s3_filelist(&self.aws_client, &self.bucket).await;
 
@@ -105,7 +105,7 @@ impl AppState {
 
         warn!("cached {} files", filelist.len());
     }
-    
+
     /// Получает путь из ключа (имени файла) в Redis.
     pub async fn get_path(&self, filename: &str) -> Result<Option<String>, actix_web::Error> {
         let mut redis = self.redis.clone();
@@ -133,9 +133,9 @@ impl AppState {
         let quota: u64 = redis.get(&quota_key).await.unwrap_or(0);
 
         if quota == 0 {
-            // Если квота не найдена, устанавливаем её в 0 байт и задаем TTL на одну неделю
+            // Если квота не найдена, устанавливаем её в 0 байт без TTL (постоянная квота)
             redis
-                .set_ex::<&str, u64, ()>(&quota_key, 0, WEEK_SECONDS)
+                .set::<&str, u64, ()>(&quota_key, 0)
                 .await
                 .map_err(|_| {
                     ErrorInternalServerError("Failed to set initial user quota in Redis")
@@ -161,10 +161,10 @@ impl AppState {
             ErrorInternalServerError("Failed to check if user quota exists in Redis")
         })?;
 
-        // Если ключ не существует, создаем его с начальным значением и устанавливаем TTL
+        // Если ключ не существует, создаем его с начальным значением без TTL
         if !exists {
             redis
-                .set_ex::<_, u64, ()>(&quota_key, bytes, WEEK_SECONDS)
+                .set::<_, u64, ()>(&quota_key, bytes)
                 .await
                 .map_err(|_| {
                     ErrorInternalServerError("Failed to set initial user quota in Redis")
@@ -180,4 +180,42 @@ impl AppState {
 
         Ok(new_quota)
     }
+
+    /// Устанавливает квоту пользователя в байтах (позволяет увеличить или уменьшить)
+    pub async fn set_user_quota(&self, user_id: &str, bytes: u64) -> Result<u64, actix_web::Error> {
+        let mut redis = self.redis.clone();
+        let quota_key = format!("quota:{}", user_id);
+
+        // Устанавливаем новое значение квоты
+        redis
+            .set::<_, u64, ()>(&quota_key, bytes)
+            .await
+            .map_err(|_| ErrorInternalServerError("Failed to set user quota in Redis"))?;
+
+        Ok(bytes)
+    }
+
+    /// Увеличивает квоту пользователя на указанное количество байт
+    pub async fn increase_user_quota(
+        &self,
+        user_id: &str,
+        additional_bytes: u64,
+    ) -> Result<u64, actix_web::Error> {
+        let mut redis = self.redis.clone();
+        let quota_key = format!("quota:{}", user_id);
+
+        // Получаем текущую квоту
+        let current_quota: u64 = redis.get(&quota_key).await.unwrap_or(0);
+
+        // Вычисляем новую квоту
+        let new_quota = current_quota + additional_bytes;
+
+        // Устанавливаем новое значение
+        redis
+            .set::<_, u64, ()>(&quota_key, new_quota)
+            .await
+            .map_err(|_| ErrorInternalServerError("Failed to increase user quota in Redis"))?;
+
+        Ok(new_quota)
+    }
 }