ncpn/src/app/api/files.rs

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<Config>,
}

pub fn router(db: PgPool, config: Arc<Config>) -> 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<SharedState>,
    body: Body,
) -> 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
            .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)),
    }
}