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::{
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 app = Router::new()
.route("/", get(default_handler))
.route("/", post(upload_handler))
.layer(
ServiceBuilder::new()
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))
.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(
layer
.layer(DefaultBodyLimit::disable())
.layer(CompressionLayer::new()),
);

View File

@@ -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())
);
}
}