use axum::{
Router,
body::Body,
extract::{Query, State},
response::IntoResponse,
routing::get,
};
use http_body_util::BodyExt;
use notify::{RecommendedWatcher, RecursiveMode, Watcher};
use std::{
collections::HashMap,
convert::Infallible,
path::{Path, PathBuf},
sync::Arc,
time::Duration,
};
use tokio::sync::broadcast;
use tower_http::services::ServeDir;
use tower_service::Service;
const SERVILE_JS: &str = r#"(function() {
function poll() {
var path = window.location.pathname;
fetch("/servile-reload?path=" + encodeURIComponent(path))
.then(function(r) {
if (r.ok) window.location.reload();
else setTimeout(poll, 1000);
})
.catch(function() { setTimeout(poll, 1000); });
}
poll();
})();"#;
const SCRIPT_TAG: &str = r#"<script src="/servile.js"></script>"#;
type LogFn = Arc<dyn Fn(&str) + Send + Sync>;
#[derive(Clone)]
struct AppState {
dir: PathBuf,
change_tx: broadcast::Sender<()>,
_watcher: Arc<RecommendedWatcher>,
logger: LogFn,
}
async fn servile_js() -> impl IntoResponse {
([("content-type", "application/javascript")], SERVILE_JS)
}
async fn servile_reload(
State(state): State<AppState>,
Query(params): Query<HashMap<String, String>>,
) -> impl IntoResponse {
let path = params
.get("path")
.map(|p| p.trim_start_matches('/').to_string())
.unwrap_or_else(|| "index.html".to_string());
let file_path = state.dir.join(&path);
let mut rx = state.change_tx.subscribe();
loop {
let _ = rx.recv().await;
if file_path.exists() {
return "ok";
}
}
}
fn is_html_content_type(headers: &axum::http::HeaderMap) -> bool {
headers
.get("content-type")
.and_then(|v| v.to_str().ok())
.is_some_and(|ct| ct.contains("text/html"))
}
fn start_watcher(dir: &Path, tx: broadcast::Sender<()>, debounce_ms: u64) -> RecommendedWatcher {
let (notify_tx, notify_rx) = std::sync::mpsc::channel();
let mut watcher =
RecommendedWatcher::new(notify_tx, notify::Config::default())
.expect("failed to create watcher");
watcher
.watch(dir, RecursiveMode::Recursive)
.expect("failed to watch directory");
let debounce = Duration::from_millis(debounce_ms);
std::thread::spawn(move || loop {
match notify_rx.recv() {
Ok(_) => {
while notify_rx.recv_timeout(debounce).is_ok() {}
let _ = tx.send(());
}
Err(_) => break,
}
});
watcher
}
pub fn build_router(dir: &Path, debounce_ms: u64) -> Router {
build_router_with_logger(dir, debounce_ms, |url| {
println!("{url}");
})
}
pub fn build_router_with_logger(
dir: &Path,
debounce_ms: u64,
logger: impl Fn(&str) + Send + Sync + 'static,
) -> Router {
let (change_tx, _) = broadcast::channel(16);
let watcher = start_watcher(dir, change_tx.clone(), debounce_ms);
let state = AppState {
dir: dir.to_path_buf(),
change_tx,
_watcher: Arc::new(watcher),
logger: Arc::new(logger),
};
let serve_dir = ServeDir::new(dir);
let log_fn = state.logger.clone();
let inject_layer = tower::service_fn(move |req: axum::http::Request<Body>| {
let mut inner = serve_dir.clone();
let logger = log_fn.clone();
async move {
let uri = req.uri().clone();
let host = req
.headers()
.get("host")
.and_then(|v| v.to_str().ok())
.unwrap_or("localhost");
let full_url = format!("http://{host}{uri}");
logger(&full_url);
let resp = match inner.call(req).await {
Ok(r) => r.map(Body::new),
Err(_) => axum::http::Response::builder()
.status(500)
.body(Body::empty())
.unwrap(),
};
let (mut parts, body) = resp.into_parts();
if !is_html_content_type(&parts.headers) {
return Ok::<_, Infallible>(axum::http::Response::from_parts(parts, body));
}
let bytes = body
.collect()
.await
.map(|b| b.to_bytes())
.unwrap_or_default();
let html = String::from_utf8_lossy(&bytes);
let injected = if html.contains("</body>") {
html.replace("</body>", &format!("{SCRIPT_TAG}</body>"))
} else {
format!("{html}{SCRIPT_TAG}")
};
parts.headers.insert(
"content-length",
axum::http::HeaderValue::from(injected.len()),
);
Ok(axum::http::Response::from_parts(parts, Body::from(injected)))
}
});
Router::new()
.route("/servile.js", get(servile_js))
.route("/servile-reload", get(servile_reload))
.fallback_service(inject_layer)
.with_state(state)
}