From 6df22507fd3233078e85f18964983cb3d298ad10 Mon Sep 17 00:00:00 2001
From: Micheal Smith <xulfer@cheapbsd.net>
Date: Fri, 7 Nov 2025 00:19:09 -0600
Subject: [PATCH] Added *very* basic paste support.

---
 src/main.rs  | 48 +++++++++++++++++++++++++++++++++++++-----------
 src/store.rs | 14 +++++++++-----
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 611c0dc..0ec5c28 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,7 @@
 use axum::{
     BoxError, Router,
     error_handling::HandleErrorLayer,
-    extract::{DefaultBodyLimit, Multipart},
+    extract::{DefaultBodyLimit, Multipart, Path, State, rejection::RawFormRejection},
     http::StatusCode,
     response::{Html, IntoResponse, Response},
     routing::{get, post},
@@ -13,6 +13,7 @@ use clap::Parser;
 use color_eyre::eyre::Result;
 use nanoid::nanoid;
 use std::{path::PathBuf, time::Duration};
+use store::Store;
 use tokio::fs;
 use tower::{ServiceBuilder, buffer::BufferLayer, limit::rate::RateLimitLayer};
 use tower_http::{compression::CompressionLayer, limit::RequestBodyLimitLayer, trace::TraceLayer};
@@ -43,11 +44,13 @@ struct Args {
     #[arg(short, long, default_value_t = 30)]
     requests_per_min: u64,
 
+    // TODO: Make this optional for temporary support.
     /// Path to paste store
     #[arg(short, long)]
     store_path: PathBuf,
 }
 
+// Should just display some default page.
 async fn default_handler() -> Response {
     match fs::read_to_string("index.html").await {
         Ok(content) => Html(content).into_response(),
@@ -59,6 +62,7 @@ async fn default_handler() -> Response {
     }
 }
 
+// NB: Anything that reaches this is pretty much unrecoverable.
 async fn handle_error(err: BoxError) -> (StatusCode, String) {
     if err.is::<tower::timeout::error::Elapsed>() {
         (StatusCode::REQUEST_TIMEOUT, "Request timed out".to_string())
@@ -75,18 +79,36 @@ async fn handle_error(err: BoxError) -> (StatusCode, String) {
     }
 }
 
-async fn upload_handler(mut multipart: Multipart) -> Response {
+// Handles multipart uploads where file=data.  Perhaps this should be
+// less rigid, but sticking with this for now.
+// To test:
+// curl -F'file=@some_file.txt' url
+//
+// Response is the id assigned to the paste.  Might do a redirect
+// to the url in the future.
+async fn upload_handler(State(mut state): State<Store>, mut multipart: Multipart) -> Response {
     while let Ok(Some(part)) = multipart.next_field().await {
         if let Some(name) = part.name()
             && name == "file"
         {
-            println!("Got: {}", name);
+            let id = nanoid!();
+            let data = part.bytes().await.unwrap();
+            let _ = state.put(id.clone(), data).await;
+            return (StatusCode::OK, id).into_response();
         }
     }
 
     (StatusCode::OK, nanoid!()).into_response()
 }
 
+// Fetch a paste based on id in the path.
+async fn paste_request(State(state): State<Store>, Path(paste_id): Path<String>) -> Response {
+    if let Ok(Some(data)) = state.get(&paste_id).await {
+        return Html(data).into_response();
+    }
+    (StatusCode::OK, format!("Couldn't fetch {paste_id}")).into_response()
+}
+
 // Defaults to info level tracing, but can be adjusted using the
 // RUST_LOG environment variable.
 fn install_tracing() {
@@ -116,19 +138,23 @@ async fn main() -> Result<()> {
     let request_limit = args.requests_per_min;
     let max_size = args.max_request_size;
     let listen_address = format!("{}:{}", args.listen_address, args.port);
+    let store = Store::new("/tmp").await?;
 
+    let layer = ServiceBuilder::new()
+        .layer(HandleErrorLayer::new(handle_error))
+        .layer(TraceLayer::new_for_http())
+        // Set to about the expected number of concurrent connections.
+        .layer(BufferLayer::new(64))
+        // Limit to 30 per minute
+        .layer(RateLimitLayer::new(request_limit, Duration::from_secs(60)))
+        .layer(RequestBodyLimitLayer::new(max_size * 1024));
     let app = Router::new()
         .route("/", get(default_handler))
         .route("/", post(upload_handler))
+        .route("/{paste_id}", get(paste_request))
+        .with_state(store)
         .layer(
-            ServiceBuilder::new()
-                .layer(HandleErrorLayer::new(handle_error))
-                .layer(TraceLayer::new_for_http())
-                // Set to about the expected number of concurrent connections.
-                .layer(BufferLayer::new(64))
-                // Limit to 30 per minute
-                .layer(RateLimitLayer::new(request_limit, Duration::from_secs(60)))
-                .layer(RequestBodyLimitLayer::new(max_size * 1024))
+            layer
                 .layer(DefaultBodyLimit::disable())
                 .layer(CompressionLayer::new()),
         );
diff --git a/src/store.rs b/src/store.rs
index b6d1b11..621b7d2 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -1,9 +1,11 @@
 use color_eyre::Result;
 use rocksdb::DB;
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 
+#[derive(Clone)]
 pub struct Store {
-    database: DB,
+    database: Arc<DB>,
     path: PathBuf,
 }
 
@@ -12,7 +14,7 @@ impl Store {
     pub async fn new(path: impl AsRef<Path>) -> Result<Store> {
         let db = DB::open_default(path.as_ref())?;
         Ok(Store {
-            database: db,
+            database: Arc::new(db),
             path: path.as_ref().to_path_buf(),
         })
     }
@@ -29,8 +31,8 @@ impl Store {
 
 #[cfg(test)]
 mod tests {
-    use rstest::*;
     use super::*;
+    use rstest::*;
     use tempfile::tempdir;
 
     #[fixture]
@@ -59,7 +61,9 @@ mod tests {
     async fn test_get(#[future] temp_store: Store) {
         let mut store = temp_store.await;
         store.put("one", "two").await.unwrap();
-        assert_eq!(store.get("one").await.unwrap(), Some("two".as_bytes().to_vec()));
+        assert_eq!(
+            store.get("one").await.unwrap(),
+            Some("two".as_bytes().to_vec())
+        );
     }
-
 }