use std::{collections::HashMap, str::Split};
use jing_core::{Adhoc, AdhocRequest, AdhocResponse, Responder};
use worker::kv::KvStore;
use worker::wasm_bindgen::JsValue;
use worker::{event, Context, Env};
use worker::{Fetch, Fetcher, Method, Request, RequestInit, Response};
use http::StatusCode;
use nanorand::{Rng, WyRand};
use tera::Tera;
use url::Url;
trait Respond {
async fn respond<I>(&self, req: &Request, rewrite: I) -> worker::Result<Response>
where
I: IntoIterator,
I::Item: AsRef<str>;
}
impl Respond for Responder {
async fn respond<I>(&self, req: &Request, rewrite: I) -> worker::Result<Response>
where
I: IntoIterator,
I::Item: AsRef<str>,
{
match self {
Self::Redirect { base } => {
let mut location = base.clone();
if let Ok(mut base_segs) = location.path_segments_mut() {
base_segs.extend(rewrite);
}
let body = HashMap::from([("location", location.as_str())]);
Response::builder()
.with_status(http::StatusCode::TEMPORARY_REDIRECT.into())
.with_header("location", location.as_str())?
.from_json(&body)
}
Self::Proxy { base } => {
let mut proxied_url = base.clone();
if let Ok(mut base_segs) = proxied_url.path_segments_mut() {
base_segs.extend(rewrite);
}
let url = req.url()?;
let query = url.query();
proxied_url.set_query(query);
let mut proxied_headers = req.headers().clone();
for key in req.headers().keys() {
let key = key.to_lowercase();
if key.starts_with("cf-") {
proxied_headers.delete(&key)?;
}
}
if let Some(host) = base.host_str() {
proxied_headers.set("host", host)?;
}
let body: Option<JsValue> = req.inner().body().map(|body| body.into());
let proxied = Request::new_with_init(
proxied_url.as_str(),
RequestInit::new()
.with_method(req.method())
.with_headers(proxied_headers)
.with_body(body),
)?;
let res = Fetch::Request(proxied).send().await?;
Ok(res)
}
}
}
}
pub struct App {
assets: Fetcher,
kv: KvStore,
rng: WyRand,
}
impl App {
fn new(env: Env, _ctx: Context) -> worker::Result<Self> {
Ok(Self {
assets: env.assets("ASSETS")?,
kv: env.kv("KV")?,
rng: WyRand::new(),
})
}
fn generate_tag<const N: usize>(&mut self) -> String {
static CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789";
(0..N)
.map(|_| {
let idx = self.rng.generate_range(0..CHARSET.len());
CHARSET[idx] as char
})
.collect::<String>()
}
fn head_segments<'a>(url: &'a Url) -> Option<(&'a str, Split<'a, char>)> {
url.path_segments().and_then(|mut segs| {
let first = segs.next();
first.filter(|s| !s.is_empty()).map(|f| (f, segs))
})
}
async fn respond(&mut self, mut req: Request) -> worker::Result<Response> {
let url = req.url()?;
match url.path() {
"/" => {
let mut tera = Tera::default();
let mut res: Response = self
.assets
.fetch(url.join("templates/index.html")?, None)
.await?
.try_into()?;
let template = res.text().await?;
if let Err(e) = tera.add_raw_template("index", &template) {
Response::error(e.to_string(), StatusCode::INTERNAL_SERVER_ERROR.into())
} else {
use tera::Context;
let mut context = Context::new();
let list = self.kv.list().prefix("responder:".into()).execute().await?;
let entries = list
.keys
.iter()
.filter_map(|k| k.name.strip_prefix("responder:"))
.collect::<Vec<_>>();
context.insert("entries", &entries);
context.insert("authority", url.authority());
context.insert("path", url.path());
Response::from_html(tera.render("index", &context).unwrap())
}
}
route if route.starts_with("/s/") => {
let tag = route.strip_prefix("/s/").unwrap();
let rest = url.path_segments().unwrap().skip(2);
if let Some(adhoc) = self
.kv
.get(&format!("adhoc:{}", tag))
.cache_ttl(60)
.json::<Adhoc>()
.await?
{
let responder = Responder::Proxy {
base: adhoc.url().clone(),
};
if adhoc.rewrite() {
responder.respond(&req, rest).await
} else {
responder.respond(&req, None::<&str>).await
}
} else {
Response::error("adhoc not found or expired", StatusCode::NOT_FOUND.into())
}
}
route if route.starts_with("/api/v1/") => {
let route = route.strip_prefix("/api/v1").unwrap();
match route {
"/ping" => Response::from_json(&HashMap::from([("pong", true)])),
"/adhoc" => {
let body: AdhocRequest = req.json().await?;
let ttl = body.ttl().unwrap_or(14400);
let tag = self.generate_tag::<8>();
self.kv
.put(&format!("adhoc:{}", tag), body.adhoc())?
.expiration_ttl(ttl)
.execute()
.await?;
Response::from_json(&AdhocResponse {
at: url.join(&format!("/s/{}", tag))?,
tag,
ttl,
})
}
route if route.starts_with("/responders/") => {
let name = route.strip_prefix("/responders/").unwrap();
let key = format!("responder:{}", name);
if req.method() == Method::Delete {
self.kv.delete(&key).await?;
Response::ok("ok")
} else if req.method() == Method::Post {
let body: Responder = req.json().await?;
self.kv.put(&key, body)?.execute().await?;
Response::ok("ok")
} else {
Response::error(
"method not allowed",
StatusCode::METHOD_NOT_ALLOWED.into(),
)
}
}
_ => Response::empty(),
}
}
route if Self::head_segments(&url).is_some() => {
let (head, segments) = Self::head_segments(&url).unwrap();
let key = format!("responder:{}", head);
if let Some(responder) = self.kv.get(&key).json::<Responder>().await? {
responder.respond(&req, segments).await
} else {
Response::error(format!("not found: {}", head), StatusCode::NOT_FOUND.into())
}
}
_ => Response::error("not found", StatusCode::NOT_FOUND.into()),
}
}
}
#[event(fetch)]
async fn fetch(req: Request, env: Env, ctx: Context) -> worker::Result<Response> {
console_error_panic_hook::set_once();
App::new(env, ctx)?.respond(req).await
}