kutil_http/axum/
host_router.rs

1use {
2    ::axum::{extract::*, http::StatusCode, response::*, *},
3    axum_extra::extract::*,
4    bytestring::*,
5    kutil_std::collections::*,
6    tower::ServiceExt,
7};
8
9//
10// HostRouter
11//
12
13/// Host router.
14#[derive(Clone, Debug, Default)]
15pub struct HostRouter {
16    /// Routers by host (with optional port).
17    pub routers: FastHashMap<ByteString, Router>,
18
19    /// Fallback host (with optional port).
20    ///
21    /// If a router is not found for a host will fallback to the router for this host instead.
22    ///
23    /// There must be a router mapped to this host.
24    pub fallback_host: Option<ByteString>,
25}
26
27impl HostRouter {
28    /// Into [Router].
29    pub fn into_router(self) -> Option<Router> {
30        match self.routers.len() {
31            0 => None,
32            1 => self.routers.values().next().cloned(),
33            _ => Some(Router::new().fallback(host_routers_handler).with_state(self)),
34        }
35    }
36
37    /// Add.
38    pub fn add(&mut self, host_and_optional_port: ByteString, router: Router) {
39        self.routers.insert(host_and_optional_port, router);
40    }
41
42    /// Fallback [Router].
43    pub fn fallback_router(&mut self) -> Option<&mut Router> {
44        match &self.fallback_host {
45            Some(host_and_optional_port) => self.routers.get_mut(host_and_optional_port),
46            None => None,
47        }
48    }
49
50    /// Handle a [Request] by forwarding it to a router if possible.
51    pub async fn handle(&mut self, host_and_optional_port: ByteString, request: Request) -> Option<Response> {
52        let router = match self.routers.get_mut(&host_and_optional_port) {
53            Some(router) => {
54                tracing::debug!("router for host: {}", host_and_optional_port);
55                router
56            }
57
58            None => match self.fallback_router() {
59                Some(router) => {
60                    tracing::debug!("using fallback, no router for host: {}", host_and_optional_port);
61                    router
62                }
63
64                None => {
65                    tracing::debug!("no fallback and no router for host: {}", host_and_optional_port);
66                    return None;
67                }
68            },
69        };
70
71        Some(router.as_service().oneshot(request).await.expect("infallible"))
72    }
73}
74
75/// Axum request handler that calls [HostRouters::handle].
76///
77/// Expects the [HostRouters] to be available as state. See
78/// [Router::with_state](::axum::Router::with_state).
79pub async fn host_routers_handler(
80    State(mut host_routers): State<HostRouter>,
81    Host(host_and_optional_port): Host,
82    request: Request,
83) -> Response {
84    host_routers
85        .handle(host_and_optional_port.into(), request)
86        .await
87        .unwrap_or_else(|| StatusCode::NOT_FOUND.into_response())
88}
kutil_http/axum/host_router.rs

kutil_http/axum/
host_router.rs
