Added *very* basic paste support.
This commit is contained in:
48
src/main.rs
48
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()),
|
||||
);
|
||||
|
||||
14
src/store.rs
14
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())
|
||||
);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user