use std::path::PathBuf; use axum::{ body::Body, extract::{Path, State}, Json, }; use axum_extra::{routing::Resource, TypedHeader}; use futures_util::TryStreamExt; use headers::ContentType; use mime::Mime; use serde::Serialize; use sha2::{Digest, Sha256}; use sqlx::query; use tokio::{fs, io}; use tokio_util::io::StreamReader; use tracing::{error, field, info, instrument}; use ulid::Ulid; use uuid::Uuid; use crate::{app::SharedState, error::AppError}; pub fn resource() -> Resource { Resource::named("files") .create(upload_file) .show(get_file_info) } #[derive(Serialize)] struct File { id: Ulid, hash: String, mime: String, keys: Vec, } #[instrument(skip(db, body))] async fn upload_file( State(SharedState { db, config }): State, TypedHeader(content_type): TypedHeader, body: Body, ) -> Result, AppError> { let id = Ulid::new(); let path_temp = config.file_temp_dir.join(id.to_string()); let mut hasher = Sha256::new(); { let mut file_temp = fs::File::create(&path_temp).await?; let better_body = body .into_data_stream() .inspect_ok(|b| hasher.update(b)) .map_err(|err| io::Error::new(io::ErrorKind::Other, err)); let mut reader = StreamReader::new(better_body); if let Err(err) = io::copy(&mut reader, &mut file_temp).await { error!( err = field::display(&err), file_path = field::debug(&path_temp), "failed to copy file, removing", ); drop(file_temp); if let Err(err) = fs::remove_file(path_temp).await { error!( err = field::display(err), "failed to remove failed upload file", ); } return Err(err.into()); } } let hash = hasher.finalize(); let hash_hex = hex::encode(hash); let path_hash = PathBuf::from("files").join(&hash_hex); if fs::try_exists(&path_hash).await? { info!(hash = hash_hex, "file already exists"); if let Err(err) = fs::remove_file(&path_temp).await { error!(err = field::display(&err), "failed to remove temp file"); } } else if let Err(err) = fs::rename(&path_temp, &path_hash).await { error!(err = field::display(&err), "failed to move finished file"); if let Err(err) = fs::remove_file(&path_temp).await { error!( err = field::display(&err), "failed to remove file after failed move", ); } return Err(err.into()); } let mime = Into::::into(content_type); let mime_str = mime.to_string(); match query!( "INSERT INTO file (id, hash, mime) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING", Uuid::from(id), &hash[..], mime_str, ) .execute(&db) .await? .rows_affected() { 0 | 1 => {} rows => return Err(AppError::ImpossibleAffectedRows(rows)), } let key = Ulid::new(); match query!( "INSERT INTO file_key (id, file_id) VALUES ($1, $2)", Uuid::from(key), Uuid::from(id), ) .execute(&db) .await? .rows_affected() { 1 => Ok(Json(File { id, hash: hash_hex, mime: mime_str, keys: vec![key], })), rows => Err(AppError::ImpossibleAffectedRows(rows)), } } async fn get_file_info( State(SharedState { db, .. }): State, Path(id): Path, ) -> Result, AppError> { let (file, keys) = tokio::try_join!( query!( "SELECT id, hash, mime FROM file WHERE id = $1", Uuid::from(id), ) .fetch_optional(&db), query!("SELECT id FROM file_key WHERE file_id = $1", Uuid::from(id)).fetch_all(&db), )?; match file { Some(r) => Ok(Json(File { id, hash: hex::encode(r.hash), mime: r.mime, keys: keys.into_iter().map(|r| r.id.into()).collect(), })), None => Err(AppError::FileNotFoundId(id)), } }