use super::super::{
super::std::{collections::*, immutable::*},
headers::*,
};
use {
::axum::{extract::*, http::StatusCode, response::*, *},
tower::ServiceExt,
};
#[derive(Clone, Debug, Default)]
pub struct HostRouter {
pub routers: FastHashMap<ImmutableString, Router>,
pub fallback_host: Option<ImmutableString>,
}
impl HostRouter {
pub fn into_router(self) -> Option<Router> {
match self.routers.len() {
0 => None,
1 => self.routers.values().next().cloned(),
_ => Some(Router::default().fallback(host_router_handler).with_state(self)),
}
}
pub fn add(&mut self, host_and_optional_port: ImmutableString, router: Router) {
self.routers.insert(host_and_optional_port, router);
}
pub fn fallback_router(&mut self) -> Option<&mut Router> {
self.fallback_host.as_ref().and_then(|host_and_optional_port| self.routers.get_mut(host_and_optional_port))
}
pub async fn handle(&mut self, host_and_optional_port: ImmutableString, request: Request) -> Option<Response> {
let router = match self.routers.get_mut(&host_and_optional_port) {
Some(router) => {
tracing::debug!("router for host: {}", host_and_optional_port);
router
}
None => match self.fallback_router() {
Some(router) => {
tracing::debug!("using fallback, no router for host: {}", host_and_optional_port);
router
}
None => {
tracing::debug!("no fallback and no router for host: {}", host_and_optional_port);
return None;
}
},
};
Some(router.as_service().oneshot(request).await.expect("infallible"))
}
}
pub async fn host_router_handler(State(mut host_router): State<HostRouter>, request: Request) -> Response {
let host_and_optional_port = request
.headers()
.x_forwarded_host_or_host()
.or_else(|| request.uri().authority().map(|authority| authority.host().into()))
.unwrap_or_default();
host_router.handle(host_and_optional_port, request).await.unwrap_or_else(|| StatusCode::NOT_FOUND.into_response())
}