ncpn/src/app/api/files.rs

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<SharedState> {
    Resource::named("files")
        .create(upload_file)
        .show(get_file_info)
}

#[derive(Serialize)]
struct File {
    id: Ulid,
    hash: String,
    mime: String,
    keys: Vec<Ulid>,
}

#[instrument(skip(db, body))]
async fn upload_file(
    State(SharedState { db, config }): State<SharedState>,
    TypedHeader(content_type): TypedHeader<ContentType>,
    body: Body,
) -> Result<Json<File>, 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::<Mime>::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<SharedState>,
    Path(id): Path<Ulid>,
) -> Result<Json<File>, 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)),
    }
}