184 lines
7.2 KiB
Rust
184 lines
7.2 KiB
Rust
use actix_web::error::ErrorInternalServerError;
|
||
use aws_config::BehaviorVersion;
|
||
use aws_sdk_s3::{config::Credentials, Client as S3Client};
|
||
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
|
||
}
|
||
|
||
const PATH_MAPPING_KEY: &str = "filepath_mapping"; // Ключ для хранения маппинга путей
|
||
const WEEK_SECONDS: u64 = 604800;
|
||
|
||
impl AppState {
|
||
/// Инициализация нового состояния приложения.
|
||
pub async fn new() -> Self {
|
||
// Получаем конфигурацию для Redis
|
||
let redis_url = env::var("REDIS_URL").expect("REDIS_URL must be set");
|
||
let redis_client = RedisClient::open(redis_url).expect("Invalid Redis URL");
|
||
let redis_connection = redis_client
|
||
.get_multiplexed_async_connection()
|
||
.await
|
||
.unwrap();
|
||
|
||
// Получаем конфигурацию для S3 (Storj)
|
||
let s3_access_key = env::var("STORJ_ACCESS_KEY").expect("STORJ_ACCESS_KEY must be set");
|
||
let s3_secret_key = env::var("STORJ_SECRET_KEY").expect("STORJ_SECRET_KEY must be set");
|
||
let s3_endpoint = env::var("STORJ_END_POINT")
|
||
.unwrap_or_else(|_| "https://gateway.storjshare.io".to_string());
|
||
let bucket = env::var("STORJ_BUCKET_NAME").unwrap_or_else(|_| "discours-io".to_string());
|
||
|
||
// Получаем конфигурацию для AWS S3
|
||
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());
|
||
|
||
// Конфигурируем клиент S3 для Storj
|
||
let storj_config = aws_config::defaults(BehaviorVersion::latest())
|
||
.region("eu-west-1")
|
||
.endpoint_url(s3_endpoint)
|
||
.credentials_provider(Credentials::new(
|
||
s3_access_key,
|
||
s3_secret_key,
|
||
None,
|
||
None,
|
||
"rust-storj-client",
|
||
))
|
||
.load()
|
||
.await;
|
||
|
||
let storj_client = S3Client::new(&storj_config);
|
||
|
||
// Конфигурируем клиент S3 для AWS
|
||
let aws_config = aws_config::defaults(BehaviorVersion::latest())
|
||
.region("eu-west-1")
|
||
.endpoint_url(aws_endpoint)
|
||
.credentials_provider(Credentials::new(
|
||
aws_access_key,
|
||
aws_secret_key,
|
||
None,
|
||
None,
|
||
"rust-aws-client",
|
||
))
|
||
.load()
|
||
.await;
|
||
|
||
let aws_client = S3Client::new(&aws_config);
|
||
|
||
let app_state = AppState {
|
||
redis: redis_connection,
|
||
storj_client,
|
||
aws_client,
|
||
bucket
|
||
};
|
||
|
||
// Кэшируем список файлов из AWS при старте приложения
|
||
app_state.cache_filelist().await;
|
||
|
||
app_state
|
||
}
|
||
|
||
/// Кэширует список файлов из Storj S3 в Redis.
|
||
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;
|
||
|
||
for [filename, filepath] in filelist.clone() {
|
||
// Сохраняем список файлов в Redis, используя HSET для каждого файла
|
||
let _: () = redis
|
||
.hset(PATH_MAPPING_KEY, filename.clone(), filepath)
|
||
.await
|
||
.expect(&format!("Failed to cache file {} in Redis", filename));
|
||
}
|
||
|
||
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();
|
||
let new_path: Option<String> = redis
|
||
.hget(PATH_MAPPING_KEY, filename)
|
||
.await
|
||
.map_err(|_| ErrorInternalServerError("Failed to get path mapping from Redis"))?;
|
||
Ok(new_path)
|
||
}
|
||
|
||
pub async fn set_path(&self, filename: &str, filepath: &str) {
|
||
let mut redis = self.redis.clone();
|
||
let _: () = redis
|
||
.hset(PATH_MAPPING_KEY, filename, filepath)
|
||
.await
|
||
.expect(&format!("Failed to cache file {} in Redis", filename));
|
||
}
|
||
|
||
/// создает или получает текущее значение квоты пользователя
|
||
pub async fn get_or_create_quota(&self, user_id: &str) -> Result<u64, actix_web::Error> {
|
||
let mut redis = self.redis.clone();
|
||
let quota_key = format!("quota:{}", user_id);
|
||
|
||
// Попытка получить квоту из Redis
|
||
let quota: u64 = redis.get("a_key).await.unwrap_or(0);
|
||
|
||
if quota == 0 {
|
||
// Если квота не найдена, устанавливаем её в 0 байт и задаем TTL на одну неделю
|
||
redis
|
||
.set_ex::<&str, u64, ()>("a_key, 0, WEEK_SECONDS)
|
||
.await
|
||
.map_err(|_| {
|
||
ErrorInternalServerError("Failed to set initial user quota in Redis")
|
||
})?;
|
||
|
||
Ok(0) // Возвращаем 0 как начальную квоту
|
||
} else {
|
||
Ok(quota)
|
||
}
|
||
}
|
||
|
||
/// инкрементирует значение квоты пользователя в байтах
|
||
pub async fn increment_uploaded_bytes(
|
||
&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
|
||
let exists: bool = redis.exists::<_, bool>("a_key).await.map_err(|_| {
|
||
ErrorInternalServerError("Failed to check if user quota exists in Redis")
|
||
})?;
|
||
|
||
// Если ключ не существует, создаем его с начальным значением и устанавливаем TTL
|
||
if !exists {
|
||
redis
|
||
.set_ex::<_, u64, ()>("a_key, bytes, WEEK_SECONDS)
|
||
.await
|
||
.map_err(|_| {
|
||
ErrorInternalServerError("Failed to set initial user quota in Redis")
|
||
})?;
|
||
return Ok(bytes);
|
||
}
|
||
|
||
// Если ключ существует, инкрементируем его значение на заданное количество байт
|
||
let new_quota: u64 = redis
|
||
.incr::<_, u64, u64>("a_key, bytes)
|
||
.await
|
||
.map_err(|_| ErrorInternalServerError("Failed to increment user quota in Redis"))?;
|
||
|
||
Ok(new_quota)
|
||
}
|
||
}
|