ncpn/src/app/api/files.rs

use std::path::PathBuf;

use axum::{
    extract::{BodyStream, State},
    routing::post,
    Json, Router,
};
use futures_util::TryStreamExt;
use serde::Serialize;
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use tokio::{
    fs::{self, File},
    io,
};
use tokio_util::io::StreamReader;
use tracing::{error, field, info, instrument};
use ulid::Ulid;

use crate::error::AppError;

pub fn router(db: PgPool) -> Router {
    Router::new().route("/", post(upload_file)).with_state(db)
}

#[derive(Debug, Serialize)]
struct UploadedFile {
    id: Ulid,
    hash: String,
}

#[instrument(skip(_db, body))]
async fn upload_file(
    State(_db): State<PgPool>,
    body: BodyStream,
) -> Result<Json<UploadedFile>, 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
            .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());
    }

    Ok(Json(UploadedFile {
        id: id_temp,
        hash: hash_hex,
    }))
}