use rama::{
Context, Layer,
http::{
Method, StatusCode,
layer::{
compression::CompressionLayer, trace::TraceLayer,
validate_request::ValidateRequestHeaderLayer,
},
matcher::HttpMatcher,
server::HttpServer,
service::web::{
IntoEndpointService, WebService,
extract::{Bytes, Path},
response::{IntoResponse, Json},
},
},
rt::Executor,
};
use serde::Deserialize;
use serde_json::json;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::RwLock;
use tracing::level_filters::LevelFilter;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{EnvFilter, fmt};
#[derive(Debug, Default)]
struct AppState {
db: RwLock<HashMap<String, bytes::Bytes>>,
}
#[derive(Debug, Deserialize)]
struct ItemParam {
key: String,
}
#[tokio::main]
async fn main() {
tracing_subscriber::registry()
.with(fmt::layer())
.with(
EnvFilter::builder()
.with_default_directive(LevelFilter::DEBUG.into())
.from_env_lossy(),
)
.init();
let addr = "127.0.0.1:62006";
tracing::info!("running service at: {addr}");
let exec = Executor::default();
HttpServer::auto(exec)
.listen_with_state(
Arc::new(AppState::default()),
addr,
TraceLayer::new_for_http()
.into_layer(
WebService::default()
.get("/", Json(json!({
"GET /": "show this API documentation in Json Format",
"GET /keys": "list all keys for which (bytes) data is stored",
"GET /item/:key": "return a 200 Ok containing the (bytes) data stored at <key>, and a 404 Not Found otherwise",
"HEAD /item/:key": "return a 200 Ok if found, and a 404 Not Found otherwise",
"POST /item/:key": "store the given request payload as the value referenced by <key>, returning a 400 Bad Request if no payload was defined",
"admin": {
"DELETE /keys": "clear all keys and their associated data",
"DELETE /item/:key": "remove the data stored at <key>, returning a 200 Ok if the key was found, and a 404 Not Found otherwise"
}
})))
.get("/keys", list_keys)
.nest("/admin", ValidateRequestHeaderLayer::bearer("secret-token")
.into_layer(WebService::default()
.delete("/keys", async |ctx: Context<Arc<AppState>>| {
ctx.state().db.write().await.clear();
})
.delete("/item/:key", async |Path(params): Path<ItemParam>, ctx: Context<Arc<AppState>>| {
match ctx.state().db.write().await.remove(¶ms.key) {
Some(_) => StatusCode::OK,
None => StatusCode::NOT_FOUND,
}
})))
.on(
HttpMatcher::method_get().or_method_head().and_path("/item/:key"),
CompressionLayer::new()
.into_layer((async |Path(params): Path<ItemParam>, method: Method, ctx: Context<Arc<AppState>>| {
match method {
Method::GET => {
match ctx.state().db.read().await.get(¶ms.key) {
Some(b) => b.clone().into_response(),
None => StatusCode::NOT_FOUND.into_response(),
}
}
Method::HEAD => {
if ctx.state().db.read().await.contains_key(¶ms.key) {
StatusCode::OK
} else {
StatusCode::NOT_FOUND
}.into_response()
}
_ => StatusCode::INTERNAL_SERVER_ERROR.into_response()
}
}).into_endpoint_service()),
)
.post("/items", async |ctx: Context<Arc<AppState>>, Json(dict): Json<HashMap<String, String>>| {
let mut db = ctx.state().db.write().await;
for (k, v) in dict {
db.insert(k, bytes::Bytes::from(v));
}
StatusCode::OK
})
.post("/item/:key", async |Path(params): Path<ItemParam>, ctx: Context<Arc<AppState>>, Bytes(value): Bytes| {
if value.is_empty() {
return StatusCode::BAD_REQUEST;
}
ctx.state().db.write().await.insert(params.key, value);
StatusCode::OK
}),
),
)
.await
.unwrap();
}
async fn list_keys(ctx: Context<Arc<AppState>>) -> impl IntoResponse {
ctx.state()
.db
.read()
.await
.keys()
.fold(String::new(), |a, b| {
if a.is_empty() {
b.clone()
} else {
format!("{a}, {b}")
}
})
}