The S3 upload needs to happen in a transaction!

src/paste.rs

@@ -1,6 +1,7 @@
 use crate::forms;
 use crate::runtime;
 use crate::s3;
+use crate::utils;
 use axum::http::StatusCode;
 use bigdecimal::BigDecimal;
 use chrono::Utc;
@@ -178,31 +179,32 @@ impl Paste {
             PasteFormat::Html(ref html) => html,
         };
 
-        // Determine file extension and content type for S3
+        // Determine file extension for S3
         let ext = match self.format {
             PasteFormat::Text(_) => "txt",
             PasteFormat::Html(_) => "html",
         };
 
+        // Determine content type for S3
         let content_type = match self.format {
             PasteFormat::Text(_) => "text/plain".to_string(),
             PasteFormat::Html(_) => "text/html".to_string(),
         };
 
-        let (s3_key, content_length) = match s3::upload(
-            s3_bucket,
-            &format!("{}{}.{}", s3_prefix, self.paste_id, ext),
-            content,
-            &content_type,
-            &self.title,
-            &format!("{}.{}", self.paste_id, ext),
-        )
-        .await
-        {
+        // Crunch crunch!
+        let (s3_content, content_encoding) = match utils::compress(content).await {
             Ok(response) => response,
-            Err(err) => return Err(err),
+            Err(err) => return Err(format!("Failed to compress content: {}", err)),
         };
 
+        // Let's append .br to the S3 key if we're using brotli compression
+        let mut s3_key = format!("{}{}.{}", s3_prefix, self.paste_id, ext);
+        if content_encoding == "br" {
+            s3_key.push_str(".br");
+        }
+
+        let content_length = s3_content.len() as i32;
+
         // Start a DB transaction
         let mut transaction = match db.begin().await {
             Ok(transaction) => transaction,
@@ -230,12 +232,31 @@ impl Paste {
         .await
         .map_err(|err| format!("Failed to insert paste: {}", err))?;
 
-        match transaction.commit().await {
-            Ok(_) => Ok(self.paste_id.clone()),
-            Err(err) => {
-                error!("Failed to commit transaction: {}", err);
-                Err(format!("Failed to commit transaction: {}", err))
-            }
+        match s3::upload(
+            s3_bucket,
+            &s3_key,
+            s3_content,
+            &content_type,
+            &content_encoding,
+            &self.title,
+            &format!("{}.{}", self.paste_id, ext),
+        )
+        .await
+        {
+            Ok(_) => match transaction.commit().await {
+                Ok(_) => Ok(self.paste_id.clone()),
+                Err(err) => {
+                    error!("Failed to commit transaction: {}", err);
+                    Err(format!("Failed to commit transaction: {}", err))
+                }
+            },
+            Err(err) => match transaction.rollback().await {
+                Ok(_) => Err(format!("Failed to upload to S3: {}", err)),
+                Err(err) => {
+                    error!("Failed to rollback transaction: {}", err);
+                    Err(format!("Failed to rollback transaction: {}", err))
+                }
+            },
         }
     }
 

src/s3.rs

@@ -1,16 +1,15 @@
 use aws_sdk_s3 as s3;
-use brotli::CompressorWriter;
-use std::io::Write;
 use tracing::error;
 
 pub async fn upload(
     bucket: &String,
     key: &str,
-    content: &String,
+    content: Vec<u8>,
     content_type: &String,
+    content_encoding: &str,
     title: &Option<String>,
     paste_id_w_ext: &String,
-) -> Result<(String, i32), String> {
+) -> Result<(), String> {
     let _config = aws_config::load_from_env().await;
 
     let config = s3::Config::from(&_config)
@@ -20,40 +19,25 @@ pub async fn upload(
 
     let client = s3::Client::from_conf(config);
 
-    let mut real_key = key.to_owned();
-    let body: Vec<u8>;
-    let content_encoding: String;
-    match compress(content) {
-        Ok(result) => {
-            body = result.0;
-            content_encoding = result.1;
-            if content_encoding == "br" {
-                real_key.push_str(".br");
-            }
-        }
-        Err(err) => {
-            return Err(err.to_string());
-        }
-    };
-    let content_length = body.len();
-
     let title = match title {
         Some(title) => title,
         None => "",
     };
 
+    let content_length = content.len() as i64;
+
     match client
         .put_object()
         .bucket(bucket)
-        .key(&real_key)
-        .body(body.into())
+        .key(key)
+        .body(content.into())
         .content_type(content_type)
         .content_encoding(content_encoding)
         .content_disposition(format!(
             "attachment; filename=\"{}\"; filename*=UTF-8''{}",
             paste_id_w_ext, paste_id_w_ext
         ))
-        .content_length(content_length as i64)
+        .content_length(content_length)
         .metadata("title", title)
         .send()
         .await
@@ -65,32 +49,7 @@ pub async fn upload(
         }
     };
 
-    Ok((real_key, content_length.try_into().unwrap()))
-}
-
-fn compress(content: &String) -> Result<(Vec<u8>, String), String> {
-    if content.len() < 1024 {
-        return Ok((content.as_bytes().to_vec(), "identity".to_string()));
-    }
-
-    let mut encoder = CompressorWriter::new(Vec::new(), 4096, 6, 22);
-    match encoder.write_all(content.as_bytes()) {
-        Ok(_) => {}
-        Err(err) => {
-            error!("Failed to compress content: {}", err);
-            return Err(format!("{}", err));
-        }
-    };
-
-    match encoder.flush() {
-        Ok(_) => {}
-        Err(err) => {
-            error!("Failed to compress content: {}", err);
-            return Err(format!("{}", err));
-        }
-    };
-
-    Ok((encoder.into_inner(), "br".to_string()))
+    Ok(())
 }
 
 pub async fn delete(bucket: &str, key: &str) -> Result<(), String> {

src/utils.rs

@@ -5,7 +5,10 @@ use axum::{
     middleware::Next,
     response::{IntoResponse, Response},
 };
+use brotli::CompressorWriter;
+use std::io::{Error, Write};
 use std::sync::Arc;
+use tracing::error;
 
 pub async fn extra_sugar(request: Request, next: Next) -> Result<impl IntoResponse, Response> {
     let headers = request.headers().clone();
@@ -91,3 +94,33 @@ pub async fn csp(
 
     Ok(response)
 }
+
+// Compress content using brotli, returning the compressed content and the content encoding
+pub async fn compress(content: &String) -> Result<(Vec<u8>, String), Error> {
+    if content.len() < 1024 {
+        return Ok((content.as_bytes().to_vec(), "identity".to_string()));
+    }
+
+    let mut encoder = CompressorWriter::new(Vec::new(), 4096, 6, 22);
+    match encoder.write_all(content.as_bytes()) {
+        Ok(_) => {}
+        Err(err) => {
+            return {
+                error!("Failed to write compressed content: {}", err);
+                Err(err)
+            };
+        }
+    };
+
+    match encoder.flush() {
+        Ok(_) => {}
+        Err(err) => {
+            return {
+                error!("Failed to flush compress content: {}", err);
+                Err(err)
+            };
+        }
+    };
+
+    Ok((encoder.into_inner(), "br".to_string()))
+}