servile_cli/
lib.rs

1use axum::{
2    Router,
3    body::Body,
4    extract::{Query, State},
5    response::IntoResponse,
6    routing::get,
7};
8use http_body_util::BodyExt;
9use notify::{RecommendedWatcher, RecursiveMode, Watcher};
10use std::{
11    collections::HashMap,
12    convert::Infallible,
13    path::{Path, PathBuf},
14    sync::Arc,
15    time::Duration,
16};
17use tokio::sync::broadcast;
18use tower_http::services::ServeDir;
19use tower_service::Service;
20
21const SERVILE_JS: &str = r#"(function() {
22    function poll() {
23        var path = window.location.pathname;
24        fetch("/servile-reload?path=" + encodeURIComponent(path))
25            .then(function(r) {
26                if (r.ok) window.location.reload();
27                else setTimeout(poll, 1000);
28            })
29            .catch(function() { setTimeout(poll, 1000); });
30    }
31    poll();
32})();"#;
33
34const SCRIPT_TAG: &str = r#"<script src="/servile.js"></script>"#;
35
36type LogFn = Arc<dyn Fn(&str) + Send + Sync>;
37
38#[derive(Clone)]
39struct AppState {
40    dir: PathBuf,
41    change_tx: broadcast::Sender<()>,
42    _watcher: Arc<RecommendedWatcher>,
43    logger: LogFn,
44}
45
46async fn servile_js() -> impl IntoResponse {
47    ([("content-type", "application/javascript")], SERVILE_JS)
48}
49
50async fn servile_reload(
51    State(state): State<AppState>,
52    Query(params): Query<HashMap<String, String>>,
53) -> impl IntoResponse {
54    let path = params
55        .get("path")
56        .map(|p| p.trim_start_matches('/').to_string())
57        .unwrap_or_else(|| "index.html".to_string());
58
59    let file_path = state.dir.join(&path);
60    let mut rx = state.change_tx.subscribe();
61
62    loop {
63        let _ = rx.recv().await;
64        if file_path.exists() {
65            return "ok";
66        }
67    }
68}
69
70fn is_html_content_type(headers: &axum::http::HeaderMap) -> bool {
71    headers
72        .get("content-type")
73        .and_then(|v| v.to_str().ok())
74        .is_some_and(|ct| ct.contains("text/html"))
75}
76
77fn start_watcher(dir: &Path, tx: broadcast::Sender<()>, debounce_ms: u64) -> RecommendedWatcher {
78    let (notify_tx, notify_rx) = std::sync::mpsc::channel();
79
80    let mut watcher =
81        RecommendedWatcher::new(notify_tx, notify::Config::default())
82            .expect("failed to create watcher");
83
84    watcher
85        .watch(dir, RecursiveMode::Recursive)
86        .expect("failed to watch directory");
87
88    let debounce = Duration::from_millis(debounce_ms);
89    std::thread::spawn(move || loop {
90        match notify_rx.recv() {
91            Ok(_) => {
92                while notify_rx.recv_timeout(debounce).is_ok() {}
93                let _ = tx.send(());
94            }
95            Err(_) => break,
96        }
97    });
98
99    watcher
100}
101
102pub fn build_router(dir: &Path, debounce_ms: u64) -> Router {
103    build_router_with_logger(dir, debounce_ms, |url| {
104        println!("{url}");
105    })
106}
107
108pub fn build_router_with_logger(
109    dir: &Path,
110    debounce_ms: u64,
111    logger: impl Fn(&str) + Send + Sync + 'static,
112) -> Router {
113    let (change_tx, _) = broadcast::channel(16);
114    let watcher = start_watcher(dir, change_tx.clone(), debounce_ms);
115
116    let state = AppState {
117        dir: dir.to_path_buf(),
118        change_tx,
119        _watcher: Arc::new(watcher),
120        logger: Arc::new(logger),
121    };
122
123    let serve_dir = ServeDir::new(dir);
124    let log_fn = state.logger.clone();
125
126    let inject_layer = tower::service_fn(move |req: axum::http::Request<Body>| {
127        let mut inner = serve_dir.clone();
128        let logger = log_fn.clone();
129        async move {
130            let uri = req.uri().clone();
131            let host = req
132                .headers()
133                .get("host")
134                .and_then(|v| v.to_str().ok())
135                .unwrap_or("localhost");
136            let full_url = format!("http://{host}{uri}");
137            logger(&full_url);
138
139            let resp = match inner.call(req).await {
140                Ok(r) => r.map(Body::new),
141                Err(_) => axum::http::Response::builder()
142                    .status(500)
143                    .body(Body::empty())
144                    .unwrap(),
145            };
146            let (mut parts, body) = resp.into_parts();
147
148            if !is_html_content_type(&parts.headers) {
149                return Ok::<_, Infallible>(axum::http::Response::from_parts(parts, body));
150            }
151
152            let bytes = body
153                .collect()
154                .await
155                .map(|b| b.to_bytes())
156                .unwrap_or_default();
157            let html = String::from_utf8_lossy(&bytes);
158
159            let injected = if html.contains("</body>") {
160                html.replace("</body>", &format!("{SCRIPT_TAG}</body>"))
161            } else {
162                format!("{html}{SCRIPT_TAG}")
163            };
164
165            parts.headers.insert(
166                "content-length",
167                axum::http::HeaderValue::from(injected.len()),
168            );
169
170            Ok(axum::http::Response::from_parts(parts, Body::from(injected)))
171        }
172    });
173
174    Router::new()
175        .route("/servile.js", get(servile_js))
176        .route("/servile-reload", get(servile_reload))
177        .fallback_service(inject_layer)
178        .with_state(state)
179}
servile_cli/lib.rs

servile_cli/
lib.rs
