1use std::{pin::Pin, sync::Arc};
2
3use axum::{
4 Router,
5 extract::{FromRequest, State},
6 http::HeaderMap,
7 routing::{MethodRouter, get},
8};
9use bytes::Bytes;
10use rand::distr::{Alphanumeric, SampleString};
11use wist::WistTunnelState;
12pub fn wist<F, S: Clone + 'static + Send + Sync>(
13 g: impl Fn(&S) -> &WistTunnelState<F> + Send + Sync + 'static,
14) -> MethodRouter<S> {
15 wist2::<F, S, Bytes>(g, |_, _| {
16 Box::pin(async move { Alphanumeric.sample_string(&mut rand::rng(), 24) })
17 })
18}
19pub fn wist2<F, S: Clone + 'static + Send + Sync, X: FromRequest<S> + Send + 'static>(
20 g: impl Fn(&S) -> &WistTunnelState<F> + Send + Sync + 'static,
21 s: impl Fn(X, S) -> Pin<Box<dyn Future<Output = String> + Send>> + Send + Sync + 'static,
22) -> MethodRouter<S> {
23 let g = Arc::new(g);
24 let s = Arc::new(s);
25 return get(|s2: State<S>, x: X| async move { s(x, s2.0).await }).post(
26 |s: State<S>, h: HeaderMap, body: Bytes| async move {
27 let s = g(&s.0);
28 let Some(h) = h.get("X-Instance-Id").and_then(|a| a.to_str().ok()) else {
29 return Bytes::default();
30 };
31 let h = s.handler(h.to_owned());
32 return h.process(&body, 1048576).await.into();
33 },
34 );
35}
36pub fn add_wist<F, S: Clone + 'static + Send + Sync>(
37 g: impl Fn(&S) -> &WistTunnelState<F> + Send + Sync + 'static,
38 path: &str,
39 r: Router<S>,
40) -> Router<S> {
41 r.route(&format!("{path}.wist"), wist(g))
42}
43pub fn add_wist2<F, S: Clone + 'static + Send + Sync, X: FromRequest<S> + Send + 'static>(
44 g: impl Fn(&S) -> &WistTunnelState<F> + Send + Sync + 'static,
45 s: impl Fn(X, S) -> Pin<Box<dyn Future<Output = String> + Send>> + Send + Sync + 'static,
46 path: &str,
47 r: Router<S>,
48) -> Router<S> {
49 r.route(&format!("{path}.wist"), wist2(g, s))
50}