From 645c7f2d987f59a1c64bf8fe508f111f747fd738 Mon Sep 17 00:00:00 2001 From: Honbra Date: Sat, 3 Feb 2024 14:46:18 +0100 Subject: [PATCH] Add file uploads Good luck getting them, though --- .gitignore | 2 + Cargo.lock | 17 ++++++++ Cargo.toml | 6 ++- src/app/api/files.rs | 92 ++++++++++++++++++++++++++++++++++++++++++++ src/app/api/mod.rs | 5 ++- src/error.rs | 7 ++++ 6 files changed, 127 insertions(+), 2 deletions(-) create mode 100644 src/app/api/files.rs diff --git a/.gitignore b/.gitignore index e772320..70fbd0c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ /target /.direnv /config.toml +/temp +/files diff --git a/Cargo.lock b/Cargo.lock index 79e2f39..e09bbeb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -807,11 +807,15 @@ dependencies = [ "axum", "eyre", "figment", + "futures-util", + "hex", "http", "serde", + "sha2", "sqlx", "thiserror", "tokio", + "tokio-util", "tower-http", "tracing", "tracing-subscriber", @@ -1684,6 +1688,19 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-util" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5419f34732d9eb6ee4c3578b7989078579b7f039cbbb9ca2c4da015749371e15" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "toml" version = "0.8.8" diff --git a/Cargo.toml b/Cargo.toml index 087b06d..b0b1bf0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,11 +7,15 @@ edition = "2021" axum = { version = "0.6.20", default-features = false, features = ["http1", "json", "macros", "matched-path", "tokio", "tower-log", "tracing"] } eyre = "0.6.8" figment = { version = "0.10.11", features = ["env", "toml"] } +futures-util = { version = "0.3.30", default-features = false } +hex = "0.4.3" http = "0.2.9" serde = { version = "1.0.189", features = ["derive"] } +sha2 = "0.10.8" sqlx = { version = "0.7.3", features = ["runtime-tokio", "postgres", "uuid"] } thiserror = "1.0.51" -tokio = { version = "1.33.0", features = ["rt-multi-thread", "macros"] } +tokio = { version = "1.33.0", features = ["rt-multi-thread", "macros", "fs", "io-std"] } +tokio-util = { version = "0.7.10", features = ["io"] } tower-http = { version = "0.4.4", features = ["trace"] } tracing = "0.1.37" tracing-subscriber = "0.3.17" diff --git a/src/app/api/files.rs b/src/app/api/files.rs new file mode 100644 index 0000000..d720777 --- /dev/null +++ b/src/app/api/files.rs @@ -0,0 +1,92 @@ +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(body))] +async fn upload_file( + State(_db): State, + body: BodyStream, +) -> 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 + .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 rename", + ); + } + return Err(err.into()); + } + + Ok(Json(UploadedFile { + id: id_temp, + hash: hash_hex, + })) +} diff --git a/src/app/api/mod.rs b/src/app/api/mod.rs index 70f8a7f..7356308 100644 --- a/src/app/api/mod.rs +++ b/src/app/api/mod.rs @@ -1,8 +1,11 @@ +mod files; mod links; use axum::Router; use sqlx::PgPool; pub fn router(db: PgPool) -> Router { - Router::new().nest("/links", links::router(db)) + Router::new() + .nest("/files", files::router(db.clone())) + .nest("/links", links::router(db)) } diff --git a/src/error.rs b/src/error.rs index 0ec0d94..269ca3b 100644 --- a/src/error.rs +++ b/src/error.rs @@ -14,6 +14,8 @@ pub enum AppError { #[error("database error")] Database(#[from] sqlx::Error), #[error(transparent)] + Io(#[from] std::io::Error), + #[error(transparent)] Other(#[from] eyre::Report), } @@ -33,6 +35,11 @@ impl IntoResponse for AppError { "A database error has occured", ) .into_response(), + Self::Io(_) => ( + StatusCode::INTERNAL_SERVER_ERROR, + "An I/O error has occured", + ) + .into_response(), Self::Other(err) => ( StatusCode::INTERNAL_SERVER_ERROR, format!("An error has occured:\n{err:?}"),