use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use axum::{
extract::State,
http::{header, HeaderMap, StatusCode, Uri},
response::{IntoResponse, Response},
Router,
};
use include_dir::Dir;
use super::serving::{
contained_file, content_type, has_source_extension, has_traversal, is_source_file,
relative_under,
};
enum Source {
Embedded(&'static Dir<'static>),
Dir(PathBuf),
}
struct Root {
prefix: String,
source: Source,
#[cfg_attr(not(feature = "dev"), allow(dead_code))]
watch: bool,
}
impl Root {
fn new(prefix: &str, source: Source) -> Self {
Self {
prefix: prefix.trim_matches('/').to_string(),
source,
watch: true,
}
}
}
#[derive(Default)]
pub struct Frontend {
roots: Vec<Root>,
}
impl Frontend {
pub fn new() -> Self {
Self::default()
}
pub fn embedded(dir: &'static Dir<'static>) -> Self {
Self {
roots: vec![Root::new("", Source::Embedded(dir))],
}
}
pub fn dir(path: impl Into<PathBuf>) -> Self {
Self {
roots: vec![Root::new("", Source::Dir(path.into()))],
}
}
pub fn source(self, path: impl Into<PathBuf>) -> Self {
self.mount_dir("/", path)
}
pub fn mount_embedded(mut self, prefix: impl AsRef<str>, dir: &'static Dir<'static>) -> Self {
self.roots
.push(Root::new(prefix.as_ref(), Source::Embedded(dir)));
self
}
pub fn mount_dir(mut self, prefix: impl AsRef<str>, path: impl Into<PathBuf>) -> Self {
self.roots
.push(Root::new(prefix.as_ref(), Source::Dir(path.into())));
self
}
pub fn router(self) -> Router {
Router::new()
.fallback(serve_static)
.with_state(Arc::new(self.roots))
}
#[cfg(feature = "dev")]
pub fn dev(self) -> Router {
let mut mounts = Vec::new();
let mut embedded = None;
for root in self.roots {
match root.source {
Source::Dir(dir) => {
let mount = if root.prefix.is_empty() {
crate::mount::Mount::root(dir)
} else {
crate::mount::Mount::new(&root.prefix, dir)
};
mounts.push(mount.watched(root.watch));
}
Source::Embedded(dir) => embedded = embedded.or(Some(dir)),
}
}
match embedded {
Some(dir) => crate::dev::dev_router_mounted_with_embedded(mounts, dir),
None => crate::dev::dev_router_mounted(mounts),
}
}
pub fn auto(self) -> Router {
#[cfg(all(debug_assertions, feature = "dev"))]
return self.dev();
#[cfg(not(all(debug_assertions, feature = "dev")))]
return self.router();
}
}
async fn serve_static(
State(roots): State<Arc<Vec<Root>>>,
headers: HeaderMap,
uri: Uri,
) -> Response {
let raw = uri.path().trim_start_matches('/');
let requested = if raw.is_empty() || raw.ends_with('/') {
format!("{raw}index.html")
} else {
raw.to_string()
};
if has_traversal(&requested) {
return StatusCode::NOT_FOUND.into_response();
}
let gzip_ok = accepts_gzip(&headers);
let mut candidates: Vec<(&Root, String)> = roots
.iter()
.filter_map(|root| relative_under(&root.prefix, &requested).map(|rel| (root, rel)))
.collect();
candidates.sort_by_key(|(root, _)| std::cmp::Reverse(root.prefix.len()));
for (root, rel) in candidates {
if rel.is_empty() || is_source_file(&rel) {
continue;
}
let ct = content_type(&rel);
match &root.source {
Source::Embedded(dir) => {
if gzip_ok {
if let Some(gz) = dir.get_file(format!("{rel}.gz")) {
return gz_response(ct, gz.contents().to_vec());
}
}
if let Some(file) = dir.get_file(&rel) {
return ([(header::CONTENT_TYPE, ct)], file.contents()).into_response();
}
}
Source::Dir(path) => {
if gzip_ok {
if let Some(gz) = contained_file(path, &format!("{rel}.gz")) {
let degz_is_source = gz
.file_stem()
.is_some_and(|s| has_source_extension(Path::new(s)));
if !degz_is_source {
if let Ok(bytes) = std::fs::read(&gz) {
return gz_response(ct, bytes);
}
}
}
}
if let Some(file) = contained_file(path, &rel) {
if !has_source_extension(&file) {
return match std::fs::read(&file) {
Ok(bytes) => ([(header::CONTENT_TYPE, ct)], bytes).into_response(),
Err(e) => {
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
}
};
}
}
}
}
}
StatusCode::NOT_FOUND.into_response()
}
fn accepts_gzip(headers: &HeaderMap) -> bool {
headers
.get(header::ACCEPT_ENCODING)
.and_then(|v| v.to_str().ok())
.is_some_and(|v| {
v.split(',')
.any(|e| e.split(';').next().map(str::trim) == Some("gzip"))
})
}
fn gz_response(content_type: String, bytes: Vec<u8>) -> Response {
(
[
(header::CONTENT_TYPE, content_type),
(header::CONTENT_ENCODING, "gzip".to_string()),
],
bytes,
)
.into_response()
}
pub async fn serve(app: Router, addr: SocketAddr) -> std::io::Result<()> {
let listener = tokio::net::TcpListener::bind(addr).await?;
println!("web-modules: serving on http://{addr}/ (Ctrl-C to stop)");
axum::serve(listener, app).await
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn root_normalises_prefix() {
let r = Root::new("/ui/", Source::Dir(PathBuf::from("x")));
assert_eq!(r.prefix, "ui");
assert!(r.watch);
}
#[test]
fn default_constructors_place_one_root_at_slash() {
assert_eq!(Frontend::dir("web").roots.len(), 1);
assert_eq!(Frontend::dir("web").roots[0].prefix, "");
assert_eq!(Frontend::new().roots.len(), 0);
assert_eq!(
Frontend::new()
.source("web")
.mount_dir("/api", "api")
.roots
.len(),
2
);
}
}