quoter/src/main.rs

use actix_web::{
    error::{ErrorInternalServerError, ErrorUnauthorized},
    middleware::Logger,
    web, App, HttpRequest, HttpResponse, HttpServer, Result,
};
use aws_config::{load_defaults, BehaviorVersion};
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;
use mime_guess::MimeGuess;
use redis::{aio::MultiplexedConnection, AsyncCommands};
use redis::Client as RedisClient;
use std::env;
use std::io::Cursor;
use std::path::Path;

const MAX_QUOTA_BYTES: u64 = 1024 * 1024 * 1024; // 1 GB per week

#[derive(Clone)]
struct AppState {
    redis: MultiplexedConnection, // Redis connection for managing quotas and file names
    s3_client: S3Client,          // S3 client for uploading files
    s3_bucket: String,            // S3 bucket name for storing files
    cdn_domain: String,           // CDN domain for generating URLs
}

// Generate a thumbnail for the image
fn generate_thumbnail(image: &DynamicImage) -> Result<Vec<u8>, actix_web::Error> {
    let thumbnail = image.resize(320, 320, FilterType::Lanczos3); // Размер миниатюры 320x320

    let mut buffer = Vec::new();
    thumbnail
        .write_to(&mut Cursor::new(&mut buffer), image::ImageFormat::Jpeg)
        .map_err(|_| ErrorInternalServerError("Failed to generate thumbnail"))?;

    Ok(buffer)
}

// Upload the file to S3 and return the URL
async fn upload_to_s3(
    s3_client: &S3Client,
    bucket: &str,
    key: &str,
    body: Vec<u8>,
    content_type: &str,
    cdn_domain: &str,
) -> Result<String, actix_web::Error> {
    let body_stream = ByteStream::from(body);

    s3_client.put_object()
        .bucket(bucket)
        .key(key)
        .body(body_stream)
        .content_type(content_type)
        .send()
        .await
        .map_err(|_| ErrorInternalServerError("Failed to upload file 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,
    user_id: &str,
    file_size: u64,
) -> Result<(), actix_web::Error> {
    let current_quota: u64 = redis.get(user_id).await.unwrap_or(0);

    if current_quota + file_size > MAX_QUOTA_BYTES {
        return Err(ErrorUnauthorized("Quota exceeded"));
    }

    redis.incr(user_id, file_size).await.map_err(|_| ErrorInternalServerError("Failed to update quota in Redis"))
}

// Proxy handler for serving static files and uploading them to S3
async fn proxy_handler(
    req: HttpRequest,
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> Result<HttpResponse, actix_web::Error> {
    let token = req.headers().get("Authorization").and_then(|header_value| header_value.to_str().ok());

    // Validate token (implementation needed)
    if token.is_none() {
        return Err(ErrorUnauthorized("Unauthorized"));
    }

    let user_id = token.unwrap(); // Assuming the token is the user ID, adjust as necessary

    // Load the file (implement your file loading logic)
    let file_path = path.into_inner();
    let mime_type = MimeGuess::from_path(&file_path).first_or_octet_stream();
    let extension = Path::new(&file_path)
        .extension()
        .and_then(|ext| ext.to_str())
        .unwrap_or("bin");

    // Handle image files: generate thumbnail and upload both
    if mime_type.type_() == "image" {
        let image = image::open(&file_path).map_err(|_| ErrorInternalServerError("Failed to open image"))?;

        // Generate thumbnail
        let thumbnail_data = generate_thumbnail(&image)?;
        let thumbnail_key = format!("thumbnail_{}.{}", file_path, "jpg");

        // Upload the thumbnail
        if let Err(_) = upload_to_s3(
            &state.s3_client,
            &state.s3_bucket,
            &thumbnail_key,
            thumbnail_data.clone(),
            "image/jpeg",
            &state.cdn_domain,
        ).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();
        image.write_to(&mut Cursor::new(&mut original_buffer), image::ImageFormat::Jpeg)
            .map_err(|_| ErrorInternalServerError("Failed to read image data"))?;
        
        // Upload the original image
        let image_key = format!("{}.{}", file_path, extension);
        let image_url = upload_to_s3(
            &state.s3_client,
            &state.s3_bucket,
            &image_key,
            original_buffer.clone(),
            mime_type.essence_str(),
            &state.cdn_domain,
        )
        .await?;

        // Update quota and save filename
        check_and_update_quota(&mut state.redis.clone(), user_id, original_buffer.len() as u64).await?;
        save_filename_in_redis(&mut state.redis.clone(), user_id, &image_key).await?;

        return Ok(HttpResponse::Ok().body(format!("Image and thumbnail uploaded to: {}", image_url)));
    }

    // Handle non-image files
    let file_data = std::fs::read(&file_path).map_err(|_| ErrorInternalServerError("Failed to read file"))?;
    let file_size = file_data.len() as u64;

    // Check and update the user's quota
    check_and_update_quota(&mut state.redis.clone(), user_id, file_size).await?;

    // Upload the file
    let file_key = format!("{}.{}", file_path, extension);
    let file_url = upload_to_s3(
        &state.s3_client,
        &state.s3_bucket,
        &file_key,
        file_data,
        mime_type.essence_str(),
        &state.cdn_domain,
    )
    .await?;

    // Save the filename in Redis for this user
    save_filename_in_redis(&mut state.redis.clone(), user_id, &file_key).await?;

    Ok(HttpResponse::Ok().body(format!("File uploaded to: {}", file_url)))
}

// Save filename in Redis for a specific user
async fn save_filename_in_redis(
    redis: &mut MultiplexedConnection,
    user_id: &str,
    filename: &str,
) -> Result<(), actix_web::Error> {
    redis.sadd(user_id, filename).await.map_err(|_| ErrorInternalServerError("Failed to save filename in Redis"))
}

// Main function to start the server
#[actix_web::main]
async fn main() -> std::io::Result<()> {
    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.ok().unwrap();

    // Initialize AWS S3 client
    let s3_bucket = env::var("S3_BUCKET").expect("S3_BUCKET must be set");
    let cdn_domain = env::var("CDN_DOMAIN").expect("CDN_DOMAIN must be set");
    let config = load_defaults(BehaviorVersion::latest()).await;
    let s3_client = S3Client::new(&config);

    // Create application state
    let app_state = web::Data::new(AppState {
        redis: redis_connection,
        s3_client,
        s3_bucket,
        cdn_domain,
    });

    // Start HTTP server
    HttpServer::new(move || {
        App::new()
            .app_data(app_state.clone())
            .wrap(Logger::default())
            .route("/{path:.*}", web::get().to(proxy_handler))
    })
    .bind("127.0.0.1:8080")?
    .run()
    .await
}