Added *very* basic paste support.

This commit is contained in:
2025-11-07 00:19:09 -06:00
parent 05b46e7b1d
commit 6df22507fd
2 changed files with 46 additions and 16 deletions

View File

@@ -4,7 +4,7 @@
use axum::{ use axum::{
BoxError, Router, BoxError, Router,
error_handling::HandleErrorLayer, error_handling::HandleErrorLayer,
extract::{DefaultBodyLimit, Multipart}, extract::{DefaultBodyLimit, Multipart, Path, State, rejection::RawFormRejection},
http::StatusCode, http::StatusCode,
response::{Html, IntoResponse, Response}, response::{Html, IntoResponse, Response},
routing::{get, post}, routing::{get, post},
@@ -13,6 +13,7 @@ use clap::Parser;
use color_eyre::eyre::Result; use color_eyre::eyre::Result;
use nanoid::nanoid; use nanoid::nanoid;
use std::{path::PathBuf, time::Duration}; use std::{path::PathBuf, time::Duration};
use store::Store;
use tokio::fs; use tokio::fs;
use tower::{ServiceBuilder, buffer::BufferLayer, limit::rate::RateLimitLayer}; use tower::{ServiceBuilder, buffer::BufferLayer, limit::rate::RateLimitLayer};
use tower_http::{compression::CompressionLayer, limit::RequestBodyLimitLayer, trace::TraceLayer}; use tower_http::{compression::CompressionLayer, limit::RequestBodyLimitLayer, trace::TraceLayer};
@@ -43,11 +44,13 @@ struct Args {
#[arg(short, long, default_value_t = 30)] #[arg(short, long, default_value_t = 30)]
requests_per_min: u64, requests_per_min: u64,
// TODO: Make this optional for temporary support.
/// Path to paste store /// Path to paste store
#[arg(short, long)] #[arg(short, long)]
store_path: PathBuf, store_path: PathBuf,
} }
// Should just display some default page.
async fn default_handler() -> Response { async fn default_handler() -> Response {
match fs::read_to_string("index.html").await { match fs::read_to_string("index.html").await {
Ok(content) => Html(content).into_response(), 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) { async fn handle_error(err: BoxError) -> (StatusCode, String) {
if err.is::<tower::timeout::error::Elapsed>() { if err.is::<tower::timeout::error::Elapsed>() {
(StatusCode::REQUEST_TIMEOUT, "Request timed out".to_string()) (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 { while let Ok(Some(part)) = multipart.next_field().await {
if let Some(name) = part.name() if let Some(name) = part.name()
&& name == "file" && 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() (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 // Defaults to info level tracing, but can be adjusted using the
// RUST_LOG environment variable. // RUST_LOG environment variable.
fn install_tracing() { fn install_tracing() {
@@ -116,19 +138,23 @@ async fn main() -> Result<()> {
let request_limit = args.requests_per_min; let request_limit = args.requests_per_min;
let max_size = args.max_request_size; let max_size = args.max_request_size;
let listen_address = format!("{}:{}", args.listen_address, args.port); 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() let app = Router::new()
.route("/", get(default_handler)) .route("/", get(default_handler))
.route("/", post(upload_handler)) .route("/", post(upload_handler))
.route("/{paste_id}", get(paste_request))
.with_state(store)
.layer( .layer(
ServiceBuilder::new() layer
.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(DefaultBodyLimit::disable()) .layer(DefaultBodyLimit::disable())
.layer(CompressionLayer::new()), .layer(CompressionLayer::new()),
); );

View File

@@ -1,9 +1,11 @@
use color_eyre::Result; use color_eyre::Result;
use rocksdb::DB; use rocksdb::DB;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::Arc;
#[derive(Clone)]
pub struct Store { pub struct Store {
database: DB, database: Arc<DB>,
path: PathBuf, path: PathBuf,
} }
@@ -12,7 +14,7 @@ impl Store {
pub async fn new(path: impl AsRef<Path>) -> Result<Store> { pub async fn new(path: impl AsRef<Path>) -> Result<Store> {
let db = DB::open_default(path.as_ref())?; let db = DB::open_default(path.as_ref())?;
Ok(Store { Ok(Store {
database: db, database: Arc::new(db),
path: path.as_ref().to_path_buf(), path: path.as_ref().to_path_buf(),
}) })
} }
@@ -29,8 +31,8 @@ impl Store {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use rstest::*;
use super::*; use super::*;
use rstest::*;
use tempfile::tempdir; use tempfile::tempdir;
#[fixture] #[fixture]
@@ -59,7 +61,9 @@ mod tests {
async fn test_get(#[future] temp_store: Store) { async fn test_get(#[future] temp_store: Store) {
let mut store = temp_store.await; let mut store = temp_store.await;
store.put("one", "two").await.unwrap(); 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())
);
} }
} }