From 70081f8bc0ff74acdae961a16af064d0f52d61d6 Mon Sep 17 00:00:00 2001
From: Untone <anton.rewin@gmail.com>
Date: Fri, 30 Aug 2024 21:21:30 +0300
Subject: [PATCH] post-generation-thumbs

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

diff --git a/src/main.rs b/src/main.rs
index 2e97f1c..a2f7ffc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,7 @@ use actix_web::{
     web, App, HttpRequest, HttpResponse, HttpServer, Result,
 };
 use aws_config::{load_defaults, BehaviorVersion};
-use aws_sdk_s3::Client as S3Client;
+use aws_sdk_s3::{error::SdkError, operation::head_object::HeadObjectError, Client as S3Client};
 use aws_sdk_s3::primitives::ByteStream;
 use image::DynamicImage;
 use image::imageops::FilterType;
@@ -15,7 +15,7 @@ use std::env;
 use std::io::Cursor;
 use std::path::Path;
 
-const MAX_QUOTA_BYTES: u64 = 2 * 1024 * 1024 * 1024; // 2 GB per week
+const MAX_QUOTA_BYTES: u64 = 1024 * 1024 * 1024; // 1 GB per week
 
 #[derive(Clone)]
 struct AppState {
@@ -60,6 +60,15 @@ async fn upload_to_s3(
     Ok(format!("{}/{}", cdn_domain, key))
 }
 
+// Check if the original file exists in S3
+async fn check_file_exists(s3_client: &S3Client, bucket: &str, key: &str) -> Result<bool, SdkError<HeadObjectError>> {
+    match s3_client.head_object().bucket(bucket).key(key).send().await {
+        Ok(_) => Ok(true),
+        Err(SdkError::ServiceError(service_error)) if service_error.err().is_not_found() => Ok(false),
+        Err(e) => Err(e),
+    }
+}
+
 // Check and update the user's quota
 async fn check_and_update_quota(
     redis: &mut MultiplexedConnection,
@@ -107,15 +116,29 @@ async fn proxy_handler(
         let thumbnail_key = format!("thumbnail_{}.{}", file_path, "jpg");
 
         // Upload the thumbnail
-        upload_to_s3(
+        if let Err(_) = upload_to_s3(
             &state.s3_client,
             &state.s3_bucket,
             &thumbnail_key,
             thumbnail_data.clone(),
             "image/jpeg",
             &state.cdn_domain,
-        )
-        .await?;
+        ).await {
+            // If thumbnail upload fails, check if original exists
+            let original_key = format!("{}.{}", file_path, extension);
+            if check_file_exists(&state.s3_client, &state.s3_bucket, &original_key).await.unwrap_or_default() {
+                // Generate and upload the thumbnail again
+                let thumbnail_data = generate_thumbnail(&image)?;
+                upload_to_s3(
+                    &state.s3_client,
+                    &state.s3_bucket,
+                    &thumbnail_key,
+                    thumbnail_data,
+                    "image/jpeg",
+                    &state.cdn_domain,
+                ).await?;
+            }
+        }
 
         // Prepare original image data
         let mut original_buffer = Vec::new();