proxy-reworked

2024-10-22 19:34:08 +03:00 · 2024-10-22 19:34:08 +03:00 · 1c27714455
commit 1c27714455
parent 30a2672ca8
7 changed files with 125 additions and 116 deletions
--- a/src/app_state.rs
+++ b/src/app_state.rs
@ -3,7 +3,6 @@ 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 std::collections::HashMap;
 use log::info;
 use crate::s3_utils::get_s3_filelist;

@ -12,7 +11,8 @@ use crate::s3_utils::get_s3_filelist;
 pub struct AppState {
    pub redis: MultiplexedConnection,
    pub storj_client: S3Client,
-    pub storj_bucket: String
+    pub aws_client: S3Client,
+    pub bucket: String
 }

 const PATH_MAPPING_KEY: &str = "filepath_mapping"; // Ключ для хранения маппинга путей
@ -34,14 +34,13 @@ impl AppState {
        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 storj_bucket = env::var("STORJ_BUCKET_NAME").unwrap_or_else(|_| "discours-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());
-        // let aws_bucket = env::var("AWS_BUCKET_NAME").unwrap_or_else(|_| "discours-io".to_string());
+        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())
@ -60,7 +59,7 @@ impl AppState {
        let storj_client = S3Client::new(&storj_config);

        // Конфигурируем клиент S3 для AWS
-        /* let aws_config = aws_config::defaults(BehaviorVersion::latest())
+        let aws_config = aws_config::defaults(BehaviorVersion::latest())
            .region("eu-west-1")
            .endpoint_url(aws_endpoint)
            .credentials_provider(Credentials::new(
@ -72,54 +71,42 @@ impl AppState {
            ))
            .load()
            .await;
-        */
-        // let aws_client = S3Client::new(&aws_config);
+        
+        let aws_client = S3Client::new(&aws_config);

        let app_state = AppState {
            redis: redis_connection,
            storj_client,
-            storj_bucket
+            aws_client,
+            bucket
        };

-        // Кэшируем список файлов из Storj S3 при старте приложения
-        app_state.cache_storj_filelist().await;
+        // Кэшируем список файлов из AWS при старте приложения
+        app_state.cache_filelist().await;

        app_state
    }

    /// Кэширует список файлов из Storj S3 в Redis.
-    pub async fn cache_storj_filelist(&self) {
+    pub async fn cache_filelist(&self) {
        info!("caching storj filelist...");
        let mut redis = self.redis.clone();
        
        // Запрашиваем список файлов из Storj S3
-        let filelist = get_s3_filelist(&self.storj_client, &self.storj_bucket).await;
+        let filelist = get_s3_filelist(&self.aws_client, &self.bucket).await;

        for [filename, filepath] in filelist.clone() {
-            if !filepath.starts_with("development") {
-                // Сохраняем список файлов в Redis, используя HSET для каждого файла
-                let _: () = redis
-                    .hset(PATH_MAPPING_KEY, filename.clone(), filepath)
-                    .await
-                    .expect(&format!("Failed to cache file {} in Redis", filename));
-            }
+            // Сохраняем список файлов в Redis, используя HSET для каждого файла
+            let _: () = redis
+                .hset(PATH_MAPPING_KEY, filename.clone(), filepath)
+                .await
+                .expect(&format!("Failed to cache file {} in Redis", filename));
        }

        info!("cached {} files", filelist.len());
    }
    
-    /// Получает кэшированный список файлов из Redis.
-    pub async fn get_cached_file_list(&self) -> Vec<String> {
-        let mut redis = self.redis.clone();
-
-        // Пытаемся получить кэшированный список из Redis
-        let cached_list: HashMap<String, String> = redis.hgetall(PATH_MAPPING_KEY).await.unwrap_or_default();
-
-        // Преобразуем HashMap в Vec<String>, используя значения (пути файлов)
-        cached_list.into_values().collect()
-    }
-    
-    /// Получает путь в Storj из ключа (имени файла) в Redis.
+    /// Получает путь из ключа (имени файла) в 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
--- a/src/handlers/proxy.rs
+++ b/src/handlers/proxy.rs
@ -1,9 +1,8 @@
 use actix_web::{error::ErrorInternalServerError, web, HttpRequest, HttpResponse, Result};
-use log::{info, warn, error};

 use crate::app_state::AppState;
-use crate::thumbnail::{parse_thumbnail_request, find_closest_width, generate_thumbnails};
-use crate::s3_utils::{load_file_from_s3, upload_to_s3};
+use crate::thumbnail::{find_closest_width, generate_thumbnails, parse_image_request};
+use crate::s3_utils::{check_file_exists, load_file_from_s3, upload_to_s3};
 use crate::handlers::serve_file::serve_file;

 /// Обработчик для скачивания файла и генерации миниатюры, если она недоступна.
@ -12,82 +11,105 @@ pub async fn proxy_handler(
    requested_res: web::Path<String>,
    state: web::Data<AppState>,
 ) -> Result<HttpResponse, actix_web::Error> {
-    info!("requested_path: {}", requested_res);
-    let requested_path = requested_res.replace("/webp", "");
-    let parts = requested_path.split('/').collect::<Vec<&str>>(); // Explicit type annotation
-    let filename = parts[parts.len()-1];
-
-    let stored_path = match state.get_path(&filename).await {
-        Ok(Some(path)) => path,
-        Ok(None) => {
-            warn!("wrong filename: {}", filename);
-            return Ok(HttpResponse::NotFound().finish());
-        }
-        Err(e) => {
-            warn!("error: {}", e);
-            return Ok(HttpResponse::InternalServerError().finish());
-        }
+    let normalized_path = match requested_res.ends_with("/webp") {
+        true => requested_res.replace("/webp", ""),
+        false => requested_res.to_string(),
    };
-    info!("stored path: {}", stored_path);

-    // Проверяем, запрошена ли миниатюра
-    if let Some((base_filename, requested_width, extension)) =
-        parse_thumbnail_request(&requested_res)
-    {
-        info!("thumbnail requested: {}  width: {} ext: {}", base_filename, requested_width, extension);
+    // парсим GET запрос
+    if let Some((base_filename, requested_width, extension)) = parse_image_request(&normalized_path) {
+        let filekey = format!("{}.{}", base_filename, extension);
+        let content_type = match extension.as_str() {
+            "jpg" | "jpeg" => "image/jpeg",
+            "png" => "image/png",
+            "webp" => "image/webp",
+            "gif" => "image/gif",
+            "mp3" => "audio/mpeg",
+            "wav" => "audio/x-wav",
+            "ogg" => "audio/ogg",
+            "aac" => "audio/aac",
+            "m4a" => "audio/m4a",
+            "flac" => "audio/flac",
+            _ => return Err(ErrorInternalServerError("unsupported file format"))
+        };

-        // Находим ближайший подходящий размер
-        let closest_width = find_closest_width(requested_width);
-        let thumb_filename = format!("{}_{}.jpg", base_filename, closest_width);
-        info!("closest width: {}, thumb_filename: {}", closest_width, thumb_filename);
+        return match state.get_path(&filekey).await {
+            Ok(Some(stored_path)) => {

-        // Проверяем наличие миниатюры в кэше
-        let cached_files = state.get_cached_file_list().await;
-        if !cached_files.contains(&thumb_filename) {
-            info!("no thumb found");
-            if cached_files.contains(&base_filename) {
-                info!("no original file found");
-                // Загружаем оригинальный файл из S3
-                let original_data: Vec<u8> =
-                    load_file_from_s3(&state.storj_client, &state.storj_bucket, &base_filename).await?;
+                // we have stored file path in storj
+                if check_file_exists(&state.storj_client, &state.bucket, &stored_path).await? {
+                    if content_type.starts_with("image") {
+                        return match requested_width == 0 {
+                            true => serve_file(&stored_path, &state).await,
+                            false => {
+                                // find closest thumb width
+                                let closest: u32 = find_closest_width(requested_width as u32);
+                                let thumb_filename = &format!("{}_{}.{}", base_filename, closest, extension);
    
-                // Генерируем миниатюру для ближайшего подходящего размера
-                let image = image::load_from_memory(&original_data).map_err(|_| {
-                    ErrorInternalServerError("Failed to load image for thumbnail generation")
-                })?;
-                let thumbnails_bytes =
-                    generate_thumbnails(&image).await?;
-                let thumbnail_bytes = thumbnails_bytes[&closest_width].clone();
+                                return match check_file_exists(&state.storj_client, &state.bucket, thumb_filename).await {
+                                    Ok(true) => serve_file(thumb_filename, &state).await,
+                                    Ok(false) => {
+                                        if let Ok(filedata) = load_file_from_s3(
+                                            &state.storj_client, &state.bucket, &stored_path).await {
+                                                thumbdata_save(&filedata, &state, &base_filename, content_type).await;
+                                                serve_file(thumb_filename, &state).await
+                                            } else {
+                                                Err(ErrorInternalServerError("cannot generate thumbnail"))
+                                            }
+                                    }
+                                    Err(_) => Err(ErrorInternalServerError("failed to load thumbnail"))
+                                }
+                            }
+                        }
+                    }
+                    // not image passing thumb generation
+                }

-                // Загружаем миниатюру в S3
-                upload_to_s3(
-                    &state.storj_client,
-                    &state.storj_bucket,
-                    &thumb_filename,
-                    thumbnail_bytes.clone(),
-                    "image/jpeg",
-                )
-                .await?;
-                info!("thumb was saved in storj");
-                return Ok(HttpResponse::Ok()
-                    .content_type("image/jpeg")
-                    .body(thumbnail_bytes));
-            } else {
-                warn!("original was not found");
-            }
-        } else {
-            info!("thumb was found");
-            return serve_file(&thumb_filename, &state).await;
+                // we need to download what stored_path keeping in aws
+                return match load_file_from_s3(
+                    &state.aws_client, 
+                    &state.bucket, 
+                    &stored_path).await {
+                        Ok(filedata) => {
+                            let _ = upload_to_s3(&state.storj_client, &state.bucket, &filekey, filedata.clone(), content_type).await;
+                            thumbdata_save(&filedata, &state, &base_filename, content_type).await;
+                            
+                            Ok(HttpResponse::Ok()
+                                .content_type(content_type)
+                                .body(filedata)
+                            )
+                        },
+                        Err(err) => Err(ErrorInternalServerError(err)),
+                    }
+            },
+            Ok(None) => Err(ErrorInternalServerError("requested file path was not found")),
+            Err(e) => Err(ErrorInternalServerError(e))
        }
    }

-    // Если запрошен целый файл
-    info!("serving full file: {}", requested_path);
-    match serve_file(&requested_path, &state).await {
-        Ok(response) => Ok(response),
-        Err(e) => {
-            error!("error: {}", e);
-            Err(e)
+    Err(ErrorInternalServerError("invalid file key"))
+    
+}
+
+async fn thumbdata_save(original_data: &[u8], state: &AppState, original_filename: &str, content_type: &str) {
+    if content_type.starts_with("image") {
+        let ext = original_filename.split('.').last().unwrap();
+        let img = image::load_from_memory(&original_data).unwrap();
+        
+        if let Ok(thumbnails_bytes) = generate_thumbnails(&img).await {
+            for (thumb_width, thumbnail) in thumbnails_bytes {
+                let thumb_filename = &format!("{}_{}.{}", original_filename, thumb_width, ext);
+                let thumbnail_bytes = thumbnail.clone();
+                // Загружаем миниатюру в S3
+                let _ = upload_to_s3(
+                    &state.storj_client,
+                    &state.bucket,
+                    thumb_filename,
+                    thumbnail_bytes,
+                    content_type,
+                )
+                .await;
+            }
        }
    }
 }
--- a/src/handlers/serve_file.rs
+++ b/src/handlers/serve_file.rs
@ -11,7 +11,7 @@ pub async fn serve_file(filepath: &str, state: &AppState) -> Result<HttpResponse
    }

    // Проверяем наличие файла в Storj S3
-    let exists = check_file_exists(&state.storj_client, &state.storj_bucket, &filepath).await?;
+    let exists = check_file_exists(&state.storj_client, &state.bucket, &filepath).await?;
    if !exists {
        return Err(ErrorInternalServerError(format!("File {} not found in Storj", filepath)));
    }
@ -20,7 +20,7 @@ pub async fn serve_file(filepath: &str, state: &AppState) -> Result<HttpResponse
    let get_object_output = state
        .storj_client
        .get_object()
-        .bucket(&state.storj_bucket)
+        .bucket(&state.bucket)
        .key(filepath)
        .send()
        .await
--- a/src/handlers/upload.rs
+++ b/src/handlers/upload.rs
@ -62,7 +62,7 @@ pub async fn upload_handler(
            // Загружаем файл в S3
            upload_to_s3(
                &state.storj_client,
-                &state.storj_bucket,
+                &state.bucket,
                &body,
                file_bytes,
                &content_type,
--- a/src/main.rs
+++ b/src/main.rs
@ -26,7 +26,7 @@ async fn main() -> std::io::Result<()> {
    spawn_blocking(move || {
        let rt = tokio::runtime::Handle::current();
        rt.block_on(async move {
-            app_state_clone.cache_storj_filelist().await;
+            app_state_clone.cache_filelist().await;
        });
    });

--- a/src/s3_utils.rs
+++ b/src/s3_utils.rs
@ -42,11 +42,11 @@ pub async fn check_file_exists(

 /// Загружает файл из S3.
 pub async fn load_file_from_s3(
-    storj_client: &S3Client,
+    s3_client: &S3Client,
    bucket: &str,
    key: &str,
 ) -> Result<Vec<u8>, actix_web::Error> {
-    let get_object_output = storj_client
+    let get_object_output = s3_client
        .get_object()
        .bucket(bucket)
        .key(key)
--- a/src/thumbnail.rs
+++ b/src/thumbnail.rs
@ -6,14 +6,14 @@ pub const THUMB_WIDTHS: [u32; 6] = [10, 40, 110, 300, 600, 800];

 /// Парсит запрос на миниатюру, извлекая оригинальное имя файла и требуемую ширину.
 /// Пример: "filename_150.ext" -> ("filename.ext", 150)
-pub fn parse_thumbnail_request(path: &str) -> Option<(String, u32, String)> {
+pub fn parse_image_request(path: &str) -> Option<(String, u32, String)> {
    if let Some((name_part, ext_part)) = path.rsplit_once('.') {
        if let Some((base_name, width_str)) = name_part.rsplit_once('_') {
            if let Ok(width) = width_str.parse::<u32>() {
                return Some((base_name.to_string(), width, ext_part.to_string()));
            }
        }
-        return Some((name_part.to_string(), 0, ext_part.to_string()))
+        return Some((name_part.to_string(), 0, ext_part.to_string().to_lowercase()));
    }
    None
 }