use std::{path::PathBuf, sync::Arc}; use axum::{body::Body, extract::State, routing::post, Json, Router}; use futures_util::TryStreamExt; use serde::Serialize; use sha2::{Digest, Sha256}; use sqlx::{query, PgPool}; use tokio::{ fs::{self, File}, io, }; use tokio_util::io::StreamReader; use tracing::{error, field, info, instrument}; use ulid::Ulid; use uuid::Uuid; use crate::{config::Config, error::AppError}; #[derive(Clone)] struct SharedState { db: PgPool, config: Arc, } pub fn router(db: PgPool, config: Arc) -> Router { Router::new() .route("/", post(upload_file)) .with_state(SharedState { db, config }) } #[derive(Debug, Serialize)] struct UploadedFile { key: Ulid, hash: String, } #[instrument(skip(db, body))] async fn upload_file( State(SharedState { db, config }): State, body: Body, ) -> Result, AppError> { let id_temp = Ulid::new(); let file_path_temp = PathBuf::from("temp").join(id_temp.to_string()); let mut hasher = Sha256::new(); { let mut file_temp = File::create(&file_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(&file_path_temp), "failed to copy file, removing", ); drop(file_temp); if let Err(err) = fs::remove_file(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 file_path_hash = PathBuf::from("files").join(&hash_hex); if fs::try_exists(&file_path_hash).await? { info!(hash = hash_hex, "file already exists"); if let Err(err) = fs::remove_file(&file_path_temp).await { error!(err = field::display(&err), "failed to remove temp file"); } } else if let Err(err) = fs::rename(&file_path_temp, &file_path_hash).await { error!(err = field::display(&err), "failed to move finished file"); if let Err(err) = fs::remove_file(&file_path_temp).await { error!( err = field::display(&err), "failed to remove file after failed move", ); } return Err(err.into()); } let key = Ulid::new(); query!( "INSERT INTO file (hash, mime) VALUES ($1, $2) ON CONFLICT DO NOTHING", &hash[..], "video/mp4", // I was testing with a video lol ) .execute(&db) .await?; let result = query!( "INSERT INTO file_key (id, file_hash) VALUES ($1, $2)", Uuid::from(key), &hash[..], ) .execute(&db) .await?; match result.rows_affected() { 1 => Ok(Json(UploadedFile { key, hash: hash_hex, })), rows => Err(AppError::ImpossibleAffectedRows(rows)), } }