use axum::{ body::Bytes, extract::{ConnectInfo, Path}, http::{HeaderMap, StatusCode}, response::{IntoResponse, Response}, Extension, Json, }; use chrono::Utc; use log::{error, info, warn}; use rand::Rng; use sqlx::AnyPool; use std::path::Path as PathBuf; use tokio::fs; use uuid::Uuid; use crate::data; use std::net::SocketAddr; use serde_json::json; /// Handler to return all files as JSON /// This function retrieves all files from the database /// and returns them as a JSON response. /// It also logs the IP address of the client making the request. /// example request: curl -X GET http://localhost:3000/all_files /// requires no parameters /// returns a JSON array of files /// TODO: add user authentication pub async fn all_files( Extension(pool): Extension, ConnectInfo(addr): ConnectInfo, ) -> impl IntoResponse { //log the IP address of the client and the call let ip = addr.ip().to_string(); info!("Received an all_files request from IP: {}", ip); // build the query and map the result to the File struct // and return the result as JSON if successful // or return an error message if not match sqlx::query_as::<_, data::File>( r#" SELECT * FROM files "#, ) .fetch_all(&pool) .await { Ok(files) => { info!("DB select all success"); (StatusCode::OK, Json(files)).into_response() } Err(e) => { warn!("DB select all error: {}", e); ( StatusCode::INTERNAL_SERVER_ERROR, "Database select all error", ) .into_response() } } } /// Handler to upload a file /// This function handles the file upload process. /// It receives the file data in the request body, /// saves it to the server's file system, /// and stores the file metadata in the database. /// It also logs the IP address of the client making the request. /// example request: curl -X POST -H "key: " -H "file_name: " -H "content-type: " -H "download_limit: " --data-binary @ http://localhost:3000/upload /// requires the following headers: /// - key: the key of the user (not optional) /// - file_name: the name of the file (optional) /// - content-type: the content type of the file (optional) /// - download_limit: the download limit of the file (optional) pub async fn upload( Extension(pool): Extension, ConnectInfo(addr): ConnectInfo, Extension(config): Extension, headers: HeaderMap, body: Bytes, ) -> Response { //log the IP address of the client and the call let ip = addr.ip().to_string(); info!("Received update from IP: {}", ip); //get the key from the headers let key = match headers.get("key") { Some(hv) => hv.to_str().unwrap_or("unknown").to_string(), None => { return ( axum::http::StatusCode::BAD_REQUEST, "Key header not supplied", ) .into_response(); } }; //check if the user exists let owner = sqlx::query_as::<_, data::user>( r#" SELECT * FROM users WHERE key = ? "#, ) .bind(&key) .fetch_one(&pool) .await; let owner = match owner { Ok(user) => { info!("User found in DB: {}", key); user.username } Err(e) => { error!("DB select error {}: {} Most likely because the Key is not valid", key, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Your key is not valid", ) .into_response(); } }; // gets the content type from the headers let content_type = headers .get("content-type") .and_then(|hv| hv.to_str().ok()) .unwrap_or("unknown") .to_string(); // gets the download limit from the headers let download_limit = headers .get("download_limit") // Option<&HeaderValue> .and_then(|hv| hv.to_str().ok()) // Option<&str> .and_then(|s| s.parse::().ok()) // Option .unwrap_or(1); // u32 //get filename from the headers let file_name = headers .get("file_name") .and_then(|hv| hv.to_str().ok()) .unwrap_or("unknown") .to_string(); //generate a random UUID for the file ID let id = { // Fallback to random UUID if body is too small let mut rng = rand::rng(); Uuid::from_u128(rng.random::()).to_string() }; //create the directory if it doesn't exist let dir = PathBuf::new(&config.data_path); if let Err(e) = fs::create_dir_all(dir).await { warn!("could not make dir at {} error: {}", &config.data_path, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Directory creation error", ) .into_response(); } //create the file path // the file path is the directory + the file ID + file type if file type is not application/x-executable info!("File type is {}", content_type); let file_path = dir.join(&id); if let Err(e) = fs::write(&file_path, &body).await { warn!("write error {}: {}", id, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "File write error", ) .into_response(); } let file_size = body.len() as i64; let upload_time = Utc::now().timestamp(); // i64 let download_count = 0; let download_url = match config.use_tls { true => format!("https://{}/download/{}", config.base_url, id), false => format!("http://{}/download/{}", config.base_url, id), }; if let Err(e) = sqlx::query( r#" INSERT INTO files (id, content_type, upload_time, download_limit, download_count, file_size, download_url, file_name, owner) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&id) .bind(&content_type) .bind(&upload_time) .bind(download_limit) .bind(download_count) .bind(file_size as i64) .bind(&download_url) .bind(&file_name) .bind(&owner) .execute(&pool) .await { error!("DB insert error {}: {}", id, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Database insert error", ) .into_response(); } let uploaded_file = data::File { id, file_name, content_type, upload_time, download_limit, download_count, file_size, download_url, owner, }; Json(uploaded_file).into_response() } /// This is The file Download handler /// This function handles the file download process. /// It retrieves the file metadata from the database /// and returns the file as a response. /// It also logs the IP address of the client making the request. /// example request: curl -X GET http://localhost:3000/download/ /// requires the following path parameter: /// - uuid: the UUID of the file (not optional) pub async fn download_file( Path(uuid): Path, // Add this extractor Extension(pool): Extension, ConnectInfo(addr): ConnectInfo, Extension(config): Extension, // Remove body: Bytes, // <-- GET handler shouldn't have a body ) -> Response { // Get UUID directly from path info!("Download request for UUID: {}", uuid); // Log the IP address of the client and the call let ip = addr.ip().to_string(); info!("Received download request for {} from IP: {}", uuid, ip); // find file by uuid in the config.data_path let file_path = PathBuf::new(&config.data_path).join(&uuid); if !file_path.exists() { error!("File not found: {}", file_path.display()); return ( axum::http::StatusCode::NOT_FOUND, "File not found", ) .into_response(); } // Check if the file exists in the database let file = sqlx::query_as::<_, data::File>( r#" SELECT * FROM files WHERE id = ? "#, ) .bind(&uuid) .fetch_one(&pool) .await; let file = match file { Ok(file) => { info!("File found in DB: {}", uuid); file } Err(e) => { error!("DB select error {}: {}", uuid, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Database select error", ) .into_response(); } }; //update download count if let Err(e) = sqlx::query( r#" UPDATE files SET download_count = download_count + 1 WHERE id = ? "#, ) .bind(&uuid) .execute(&pool) .await { error!("DB update error {}: {}", uuid, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Database update error", ) .into_response(); } info!("Update Download Count Sucess for UUID: {}", uuid); //rutn file to axum::body::Bytes let file_bytes = match fs::read(&file_path).await { Ok(file) => file, Err(e) => { error!("File read error {}: {}", uuid, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "File read error", ) .into_response(); } }; //if download count is greater or equal to download limit delete the file and remove it from the database if (file.download_count) >= file.download_limit { if let Err(e) = fs::remove_file(&file_path).await { error!("File delete error {}: {}", uuid, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "File delete error", ) .into_response(); } if let Err(e) = sqlx::query( r#" DELETE FROM files WHERE id = ? "#, ) .bind(&uuid) .execute(&pool) .await { error!("DB delete error {}: {}", uuid, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Database delete error", ) .into_response(); } info!("File deleted from DB because max download limit was reached: {}", uuid); } // return the file as a response return ( axum::http::StatusCode::OK, axum::response::IntoResponse::into_response( axum::response::Response::builder() .header("Content-Disposition", format!("attachment; filename=\"{}\"", uuid)) .header("Content-Type", format!("{}", &file.content_type )) .header("Content-Length", file.file_size) .header("filename", file.file_name) .body(axum::body::Body::from(file_bytes)) .unwrap(), ), ) .into_response() } /// Handler to upload a file /// This function registers a new user. /// It receives the user data in the request headers, /// saves it to the database, /// and returns the user data as a JSON response. /// It also logs the IP address of the client making the request. /// example request: curl -X POST -H "username: " -H "password: " http://localhost:3000/register /// requires the following headers: /// - username: the username of the user (not optional) /// - password: the password of the user (not optional) pub async fn register_user ( Extension(pool): Extension, ConnectInfo(addr): ConnectInfo, Extension(config): Extension, headers: HeaderMap, body: Bytes, ) -> Response { //log the IP address of the client and the call let ip = addr.ip().to_string(); info!("Received update from IP: {}", ip); //check if registration is allowed if !config.allow_register { return ( axum::http::StatusCode::FORBIDDEN, "Registration is not allowed", ) .into_response(); } // gets the content type from the headers return error if header is not suplyde let username = match headers .get("username") { Some(hv) => hv.to_str().unwrap_or("unknown").to_string(), None => { return ( axum::http::StatusCode::BAD_REQUEST, "Username header not supplied", ) .into_response(); } }; let password = match headers .get("password") { Some(hv) => hv.to_str().unwrap_or("unknown").to_string(), None => { return ( axum::http::StatusCode::BAD_REQUEST, "Password header not supplied", ) .into_response(); } }; //generate a random UUID for the user key let key = { // Fallback to random UUID if body is too small let mut rng = rand::rng(); Uuid::from_u128(rng.random::()).to_string() }; // check if the user already exists let user = sqlx::query_as::<_, data::user>( r#" SELECT * FROM users WHERE username = ? "#, ) .bind(&username) .fetch_one(&pool) .await; match user { Ok(_) => { info!("User already exists: {}", username); return ( axum::http::StatusCode::BAD_REQUEST, "User already exists", ) .into_response(); } Err(e) => { warn!("DB select error {}: {}", username, e); } } //add the user to the database if let Err(e) = sqlx::query( r#" INSERT INTO users (key, username, password) VALUES (?, ?, ?) "#, ) .bind(&key) .bind(&username) .bind(&password) .execute(&pool) .await { error!("DB insert error {}: {}", key, e); return ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, "Database insert error", ) .into_response(); } info!("User registered: {}", username); //return the user as a response let registered_user = json!({ "key": key, "username": username, }); Json(registered_user) .into_response() }