It's safe guys

safety is our #1 priority

src/app/api/files.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 
 use axum::{
     body::Body,
-    extract::{Path, State},
+    extract::{Path, Query, State},
     routing::{delete, get, post},
     Json, Router,
 };
@@ -11,7 +11,7 @@ use futures_util::TryStreamExt;
 use headers::ContentType;
 use http::StatusCode;
 use mime::Mime;
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use sqlx::query;
 use tokio::{fs, io};
@@ -39,12 +39,27 @@ struct File {
     keys: Vec<Ulid>,
 }
 
+#[derive(Serialize)]
+struct NewFile {
+    id: Ulid,
+    hash: String,
+    mime: String,
+    key: Option<Ulid>,
+}
+
+#[derive(Deserialize)]
+struct UploadFileOptions {
+    #[serde(default)]
+    create_key: bool,
+}
+
 #[instrument(skip_all)]
 async fn upload_file(
     State(SharedState { db, config }): State<SharedState>,
+    Query(UploadFileOptions { create_key }): Query<UploadFileOptions>,
     TypedHeader(content_type): TypedHeader<ContentType>,
     body: Body,
-) -> Result<Json<File>, AppError> {
+) -> Result<Json<NewFile>, AppError> {
     let id = Ulid::new();
     let path_temp = config.file_temp_dir.join(id.to_string());
     let mut hasher = Sha256::new();
@@ -100,13 +115,15 @@ async fn upload_file(
     let mime = Into::<Mime>::into(content_type);
     let mime_str = mime.to_string();
 
+    let mut tx = db.begin().await?;
+
     match query!(
         "INSERT INTO file (id, hash, mime) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING",
         Uuid::from(id),
         &hash[..],
         mime_str,
     )
-    .execute(&db)
+    .execute(&mut *tx)
     .await?
     .rows_affected()
     {
@@ -114,25 +131,42 @@ async fn upload_file(
         rows => return Err(AppError::ImpossibleAffectedRows(rows)),
     }
 
-    let key = Ulid::new();
+    // `ON CONFLICT DO NOTHING RETURNING id` only works when there *isn't* a
+    // conflict
+    let id = query!("SELECT id FROM file WHERE hash = $1", &hash[..])
+        .fetch_one(&mut *tx)
+        .await?
+        .id
+        .into();
 
-    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)),
+    let mut key_opt = None;
+
+    if create_key {
+        let key = Ulid::new();
+        key_opt = Some(key);
+        match query!(
+            "INSERT INTO file_key (id, file_id) VALUES ($1, $2)",
+            Uuid::from(key),
+            Uuid::from(id),
+        )
+        .execute(&mut *tx)
+        .await?
+        .rows_affected()
+        {
+            1 => {}
+            0 => return Err(AppError::UlidConflict(key)),
+            rows => return Err(AppError::ImpossibleAffectedRows(rows)),
+        }
     }
+
+    tx.commit().await?;
+
+    Ok(Json(NewFile {
+        id,
+        hash: hash_hex,
+        mime: mime_str,
+        key: key_opt,
+    }))
 }
 
 async fn get_file_info(

src/error.rs

@@ -21,6 +21,8 @@ pub enum AppError {
     FileMissing(PathBuf),
     #[error("database returned an impossible number of affected rows ({0})")]
     ImpossibleAffectedRows(u64),
+    #[error("ulid conflict")]
+    UlidConflict(Ulid),
     #[error("database error: {0}")]
     Database(#[from] sqlx::Error),
     #[error(transparent)]
@@ -44,6 +46,7 @@ impl IntoResponse for AppError {
                 StatusCode::INTERNAL_SERVER_ERROR,
                 "Database returned an impossible number of affected rows",
             ),
+            Self::UlidConflict(_) => (StatusCode::INTERNAL_SERVER_ERROR, "ULID conflict real???"),
             Self::Database(_) => (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 "A database error has occured",