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