quoter/src/main.rs

320 lines
9.8 KiB
Rust
Raw Normal View History

2024-08-30 18:05:51 +00:00
use actix_web::{
error::{ErrorInternalServerError, ErrorUnauthorized},
middleware::Logger,
web, App, HttpRequest, HttpResponse, HttpServer, Result,
};
2024-08-30 19:11:21 +00:00
use aws_config::BehaviorVersion;
2024-08-30 18:05:51 +00:00
use aws_sdk_s3::primitives::ByteStream;
2024-08-30 19:24:47 +00:00
use aws_sdk_s3::{config::Credentials, error::SdkError, Client as S3Client};
use image::{imageops::FilterType, DynamicImage};
2024-08-30 18:05:51 +00:00
use mime_guess::MimeGuess;
use redis::Client as RedisClient;
2024-08-30 19:24:47 +00:00
use redis::{aio::MultiplexedConnection, AsyncCommands};
2023-09-27 23:08:48 +00:00
use std::env;
2024-08-30 18:05:51 +00:00
use std::io::Cursor;
use std::path::Path;
2024-04-08 07:17:14 +00:00
2024-08-30 19:11:21 +00:00
const MAX_QUOTA_BYTES: u64 = 2 * 1024 * 1024 * 1024; // 2 GB per week
2023-10-02 13:47:22 +00:00
2023-10-06 14:57:54 +00:00
#[derive(Clone)]
struct AppState {
2024-08-30 19:11:21 +00:00
redis: MultiplexedConnection,
s3_client: S3Client,
s3_bucket: String,
2024-08-30 19:24:47 +00:00
aws_bucket: String,
2023-10-06 14:57:54 +00:00
}
2024-08-30 19:11:21 +00:00
impl AppState {
async fn new() -> Self {
let redis_url = env::var("REDIS_URL").expect("REDIS_URL must be set");
let redis_client = RedisClient::open(redis_url).expect("Invalid Redis URL");
2024-08-30 19:24:47 +00:00
let redis_connection = redis_client
.get_multiplexed_async_connection()
.await
.unwrap();
2024-08-30 19:11:21 +00:00
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").expect("STORJ_END_POINT must be set");
let s3_bucket = env::var("STORJ_BUCKET_NAME").expect("STORJ_BUCKET_NAME must be set");
2024-08-30 19:24:47 +00:00
let aws_bucket = env::var("AWS_BUCKET_NAME").expect("AWS_BUCKET_NAME must be set");
2024-08-30 19:11:21 +00:00
let 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-s3-client",
))
.load()
.await;
let s3_client = S3Client::new(&config);
AppState {
redis: redis_connection,
s3_client,
s3_bucket,
2024-08-30 19:24:47 +00:00
aws_bucket,
2024-08-30 19:11:21 +00:00
}
}
}
2024-08-30 18:05:51 +00:00
2024-08-30 19:11:21 +00:00
async fn generate_thumbnail(image: &DynamicImage, width: u32) -> Result<Vec<u8>, actix_web::Error> {
let k = image.width() / width;
let height = image.height() / k;
let thumbnail = image.resize(width, height, FilterType::Lanczos3);
2024-08-30 18:05:51 +00:00
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)
}
async fn upload_to_s3(
s3_client: &S3Client,
bucket: &str,
key: &str,
body: Vec<u8>,
content_type: &str,
) -> Result<String, actix_web::Error> {
let body_stream = ByteStream::from(body);
2024-08-30 19:24:47 +00:00
s3_client
.put_object()
2024-08-30 18:05:51 +00:00
.bucket(bucket)
.key(key)
.body(body_stream)
.content_type(content_type)
.send()
.await
.map_err(|_| ErrorInternalServerError("Failed to upload file to S3"))?;
2024-08-30 19:24:47 +00:00
Ok(key.to_string())
2023-10-19 14:43:00 +00:00
}
2024-08-30 19:24:47 +00:00
async fn check_file_exists(
s3_client: &S3Client,
bucket: &str,
key: &str,
) -> Result<bool, actix_web::Error> {
2024-08-30 18:21:30 +00:00
match s3_client.head_object().bucket(bucket).key(key).send().await {
Ok(_) => Ok(true),
2024-08-30 19:24:47 +00:00
Err(SdkError::ServiceError(service_error)) if service_error.err().is_not_found() => {
Ok(false)
}
2024-08-30 19:11:21 +00:00
Err(e) => Err(ErrorInternalServerError(e.to_string())),
2024-08-30 18:21:30 +00:00
}
}
2024-08-30 18:05:51 +00:00
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"));
}
2024-08-30 19:24:47 +00:00
redis
.incr(user_id, file_size)
.await
.map_err(|_| ErrorInternalServerError("Failed to update quota in Redis"))?;
2024-08-30 19:11:21 +00:00
Ok(())
}
2024-08-30 18:05:51 +00:00
2024-08-30 19:11:21 +00:00
async fn save_filename_in_redis(
redis: &mut MultiplexedConnection,
user_id: &str,
filename: &str,
) -> Result<(), actix_web::Error> {
2024-08-30 19:24:47 +00:00
redis
.sadd(user_id, filename)
.await
.map_err(|_| ErrorInternalServerError("Failed to save filename in Redis"))?;
Ok(())
}
async fn upload_files_from_aws(
aws_client: &S3Client,
aws_bucket: &str,
storj_client: &S3Client,
storj_bucket: &str,
) -> Result<(), actix_web::Error> {
let list_objects_v2 = aws_client.list_objects_v2();
let list_response = list_objects_v2
.bucket(aws_bucket)
.send()
.await
.map_err(|_| ErrorInternalServerError("Failed to list files from AWS S3"))?;
if let Some(objects) = list_response.contents {
for object in objects {
if let Some(key) = object.key {
// Get the object from AWS S3
let object_response = aws_client
.get_object()
.bucket(aws_bucket)
.key(&key)
.send()
.await
.map_err(|_| ErrorInternalServerError("Failed to get object from AWS S3"))?;
let body = object_response
.body
.collect()
.await
.map_err(|_| ErrorInternalServerError("Failed to read object body"))?;
let content_type = object_response
.content_type
.unwrap_or_else(|| "application/octet-stream".to_string());
// Upload the object to Storj S3
let storj_url = upload_to_s3(
storj_client,
storj_bucket,
&key,
body.into_bytes().to_vec(),
&content_type,
)
.await?;
println!("Uploaded {} to Storj at {}", key, storj_url);
}
}
}
2024-08-30 19:11:21 +00:00
Ok(())
2023-10-16 13:22:54 +00:00
}
2024-08-30 18:05:51 +00:00
async fn proxy_handler(
2023-10-11 20:52:33 +00:00
req: HttpRequest,
2024-08-30 18:05:51 +00:00
path: web::Path<String>,
2023-10-06 14:57:54 +00:00
state: web::Data<AppState>,
2023-10-03 13:29:31 +00:00
) -> Result<HttpResponse, actix_web::Error> {
2024-08-30 19:24:47 +00:00
let token = req
.headers()
.get("Authorization")
.and_then(|header_value| header_value.to_str().ok());
2024-08-30 18:05:51 +00:00
if token.is_none() {
return Err(ErrorUnauthorized("Unauthorized"));
}
2023-10-16 14:44:19 +00:00
2024-08-30 19:11:21 +00:00
let user_id = token.unwrap(); // Assuming the token is the user ID
2024-08-30 18:05:51 +00:00
let file_path = path.into_inner();
let mime_type = MimeGuess::from_path(&file_path).first_or_octet_stream();
2024-08-30 19:24:47 +00:00
let extension = Path::new(&file_path)
.extension()
.and_then(|ext| ext.to_str())
.unwrap_or("bin");
2024-08-30 18:05:51 +00:00
if mime_type.type_() == "image" {
2024-08-30 19:24:47 +00:00
let image = image::open(&file_path)
.map_err(|_| ErrorInternalServerError("Failed to open image"))?;
2024-08-30 18:05:51 +00:00
2024-08-30 19:11:21 +00:00
// Define thumbnail sizes
let thumbnail_sizes = vec![40, 110, 300, 600, 800];
for width in thumbnail_sizes {
let thumbnail_key = format!("{}_{}.jpg", file_path, width);
let thumbnail_data = generate_thumbnail(&image, width).await?;
// Check if thumbnail already exists
if !check_file_exists(&state.s3_client, &state.s3_bucket, &thumbnail_key).await? {
2024-08-30 19:24:47 +00:00
upload_to_s3(
&state.s3_client,
&state.s3_bucket,
&thumbnail_key,
thumbnail_data,
"image/jpeg",
)
.await?;
2024-08-30 18:21:30 +00:00
}
}
2024-08-30 18:05:51 +00:00
// Prepare original image data
let mut original_buffer = Vec::new();
2024-08-30 19:24:47 +00:00
image
.write_to(
&mut Cursor::new(&mut original_buffer),
image::ImageFormat::Jpeg,
)
2024-08-30 18:05:51 +00:00
.map_err(|_| ErrorInternalServerError("Failed to read image data"))?;
2024-08-30 19:11:21 +00:00
2024-08-30 18:05:51 +00:00
// Upload the original image
let image_key = format!("{}.{}", file_path, extension);
2024-08-30 19:24:47 +00:00
let image_url = upload_to_s3(
&state.s3_client,
&state.s3_bucket,
&image_key,
original_buffer.clone(),
mime_type.essence_str(),
)
.await?;
2024-08-30 18:05:51 +00:00
// Update quota and save filename
2024-08-30 19:24:47 +00:00
check_and_update_quota(
&mut state.redis.clone(),
user_id,
original_buffer.len() as u64,
)
.await?;
2024-08-30 18:05:51 +00:00
save_filename_in_redis(&mut state.redis.clone(), user_id, &image_key).await?;
2024-08-30 19:24:47 +00:00
return Ok(
HttpResponse::Ok().body(format!("Image and thumbnails uploaded to: {}", image_url))
);
2024-08-30 18:05:51 +00:00
}
// Handle non-image files
2024-08-30 19:24:47 +00:00
let file_data =
std::fs::read(&file_path).map_err(|_| ErrorInternalServerError("Failed to read file"))?;
2024-08-30 18:05:51 +00:00
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);
2024-08-30 19:24:47 +00:00
let file_url = upload_to_s3(
&state.s3_client,
&state.s3_bucket,
&file_key,
file_data,
mime_type.essence_str(),
)
.await?;
2024-08-30 18:05:51 +00:00
// 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)))
2023-09-27 23:08:48 +00:00
}
#[actix_web::main]
async fn main() -> std::io::Result<()> {
2024-08-30 19:11:21 +00:00
let app_state = AppState::new().await;
2024-08-30 19:24:47 +00:00
// Example of uploading files from AWS S3 to Storj
upload_files_from_aws(
&app_state.s3_client,
&app_state.aws_bucket,
&app_state.s3_client,
&app_state.s3_bucket,
)
.await
.expect("Failed to upload files from AWS to Storj");
2023-09-27 23:08:48 +00:00
HttpServer::new(move || {
App::new()
2024-08-30 19:11:21 +00:00
.app_data(web::Data::new(app_state.clone()))
2023-10-11 20:03:12 +00:00
.wrap(Logger::default())
2024-08-30 18:05:51 +00:00
.route("/{path:.*}", web::get().to(proxy_handler))
2023-09-27 23:08:48 +00:00
})
2024-08-30 18:05:51 +00:00
.bind("127.0.0.1:8080")?
2023-09-27 23:08:48 +00:00
.run()
.await
2024-08-30 19:24:47 +00:00
}