use axum::{ body::StreamBody, extract::{Extension, Multipart, Path}, http::{HeaderMap, StatusCode}, response::IntoResponse, routing::{delete, get, post}, Json, Router, }; use sqlx::PgPool; use std::{path::PathBuf, sync::Arc}; use tokio::fs::File; use tokio_util::io::ReaderStream; use uuid::Uuid; use crate::config::Config; use crate::error::AppError; use crate::models::file::{CreateDirectoryDto, CreateFileDto, File as FileModel, FileType}; use crate::services::auth::AuthUser; use crate::services::encryption::EncryptionService; use crate::services::storage::StorageService; pub fn routes() -> Router { Router::new() .route("/files", get(list_files)) .route("/files/:id", get(get_file)) .route("/files/:id/download", get(download_file)) .route("/files/upload", post(upload_file)) .route("/files/directory", post(create_directory)) .route("/files/:id", delete(delete_file)) } async fn list_files( auth_user: AuthUser, Extension(pool): Extension, Path(parent_id): Option>, ) -> Result>, AppError> { let files = match parent_id { Some(parent_id) => FileModel::list_by_parent(&pool, parent_id, auth_user.id).await?, None => FileModel::list_root_directory(&pool, auth_user.id).await?, }; Ok(Json(files)) } async fn get_file( auth_user: AuthUser, Extension(pool): Extension, Path(id): Path, ) -> Result, AppError> { let file = FileModel::find_by_id(&pool, id).await?; // Check if user has access to this file if file.owner_id != auth_user.id { return Err(AppError::AccessDenied("You don't have access to this file".to_string())); } Ok(Json(file)) } async fn upload_file( auth_user: AuthUser, Extension(pool): Extension, Extension(config): Extension>, Extension(encryption_service): Extension, mut multipart: Multipart, ) -> Result, AppError> { // Extract file data from multipart form let mut file_name = None; let mut file_data = None; let mut parent_id = None; while let Some(field) = multipart.next_field().await.map_err(|e| AppError::InvalidInput(e.to_string()))? { let name = field.name().unwrap_or("").to_string(); if name == "file" { file_name = field.file_name().map(|s| s.to_string()); file_data = Some(field.bytes().await.map_err(|e| AppError::InvalidInput(e.to_string()))?); } else if name == "parent_id" { let parent_id_str = field.text().await.map_err(|e| AppError::InvalidInput(e.to_string()))?; if !parent_id_str.is_empty() { parent_id = Some(Uuid::parse_str(&parent_id_str).map_err(|e| AppError::InvalidInput(e.to_string()))?); } } } let file_name = file_name.ok_or_else(|| AppError::InvalidInput("File name is required".to_string()))?; let file_data = file_data.ok_or_else(|| AppError::InvalidInput("File data is required".to_string()))?; // Check user storage quota let user = crate::models::user::User::find_by_id(&pool, auth_user.id).await?; if user.storage_used + file_data.len() as i64 > user.storage_quota { return Err(AppError::StorageQuotaExceeded("Storage quota exceeded".to_string())); } let storage_service = StorageService::new(&config.storage_path); // Create file record in database let create_file_dto = CreateFileDto { name: file_name, parent_id, size: file_data.len() as i64, mime_type: mime_guess::from_path(&file_name).first_or_octet_stream().to_string(), }; let file = FileModel::create(&pool, auth_user.id, create_file_dto).await?; // Encrypt and save file to disk let encrypted_data = encryption_service.encrypt(&file_data)?; storage_service.save_file(auth_user.id, file.id, &encrypted_data).await?; // Update user storage used crate::models::user::User::update_storage_used(&pool, auth_user.id, file_data.len() as i64).await?; Ok(Json(file)) } async fn download_file( auth_user: AuthUser, Extension(pool): Extension, Extension(config): Extension>, Extension(encryption_service): Extension, Path(id): Path, ) -> Result { let file = FileModel::find_by_id(&pool, id).await?; if file.owner_id != auth_user.id { return Err(AppError::AccessDenied("You don't have access to this file".to_string())); } if file.file_type != FileType::File { return Err(AppError::InvalidInput("Cannot download a directory".to_string())); } let storage_service = StorageService::new(&config.storage_path); let encrypted_data = storage_service.read_file(auth_user.id, file.id).await?; let decrypted_data = encryption_service.decrypt(&encrypted_data)?; let mut headers = HeaderMap::new(); headers.insert( axum::http::header::CONTENT_TYPE, file.mime_type.parse().unwrap_or_else(|_| "application/octet-stream".parse().unwrap()), ); headers.insert( axum::http::header::CONTENT_DISPOSITION, format!("attachment; filename=\"{}\"", file.name).parse().unwrap(), ); let stream = tokio_util::io::ReaderStream::new(std::io::Cursor::new(decrypted_data)); let body = StreamBody::new(stream); Ok((StatusCode::OK, headers, body)) } async fn create_directory( auth_user: AuthUser, Extension(pool): Extension, Json(create_dir_dto): Json, ) -> Result, AppError> { let directory = FileModel::create_directory(&pool, auth_user.id, create_dir_dto).await?; Ok(Json(directory)) } async fn delete_file( auth_user: AuthUser, Extension(pool): Extension, Extension(config): Extension>, Path(id): Path, ) -> Result { let file = FileModel::find_by_id(&pool, id).await?; if file.owner_id != auth_user.id { return Err(AppError::AccessDenied("You don't have access to this file".to_string())); } let storage_service = StorageService::new(&config.storage_path); if file.file_type == FileType::Directory { FileModel::delete_directory_recursive(&pool, id, auth_user.id, &storage_service).await?; } else { storage_service.delete_file(auth_user.id, file.id).await?; FileModel::delete(&pool, id).await?; crate::models::user::User::update_storage_used(&pool, auth_user.id, -file.size).await?; } Ok(StatusCode::NO_CONTENT) }