krz/orgo

  1//! `serve`: a development server over the built site, with browser live reload.
  2//!
  3//! `watch` rebuilds but leaves you to serve the output and press reload yourself. This
  4//! closes that loop: build, watch, serve, and push a reload to the browser when a
  5//! rebuild lands.
  6//!
  7//! Three decisions shape it:
  8//!
  9//! 1. **Loopback by default.** A development server binds `127.0.0.1`, not `0.0.0.0`.
 10//!    It serves unreviewed drafts off someone's laptop, and exposing that to the local
 11//!    network should be a thing you ask for (`--host`), never a thing you get.
 12//! 2. **The reload script is injected at serve time**, never written to disk. The built
 13//!    site is what you deploy, and it must not carry a dev server's JavaScript.
 14//! 3. **Long-polling, not WebSockets or SSE.** The browser asks "has anything changed
 15//!    since generation N?" and the server holds the request open until something has.
 16//!    That is instant like a push, needs no protocol beyond ordinary HTTP, and — unlike
 17//!    a streamed response — completes, which is what makes it work at all: tiny_http
 18//!    buffers a response until its body ends, so a body that never ends never reaches
 19//!    the client. Long-polling was the version of this that worked.
 20
 21use std::io;
 22use std::sync::{Arc, Condvar, Mutex};
 23use std::time::Duration;
 24
 25use anyhow::Result;
 26use camino::{Utf8Path, Utf8PathBuf};
 27use tiny_http::{Header, Request, Response, Server, StatusCode};
 28
 29use crate::site::{build_site, BuildOptions};
 30
 31/// Where the browser subscribes for reload events. Namespaced so it cannot collide with
 32/// a real page.
 33pub const RELOAD_PATH: &str = "/__orgo/reload";
 34
 35/// How long a poll waits before answering "nothing yet". Long enough that an idle tab is
 36/// nearly silent, short enough to stay under any proxy or browser idle timeout.
 37const POLL_TIMEOUT: Duration = Duration::from_secs(25);
 38
 39/// The script injected into served HTML, carrying the generation the page was built
 40/// from.
 41///
 42/// Baking the generation in is what makes this race-free: if a rebuild lands between the
 43/// page being served and the first poll going out, the server answers immediately rather
 44/// than the tab sitting on stale content until the *next* edit.
 45fn reload_script(generation: u64) -> String {
 46    format!(
 47        "\n<script>(function p(n){{fetch(\"{RELOAD_PATH}?since=\"+n)\
 48         .then(function(r){{return r.json()}})\
 49         .then(function(g){{g>n?location.reload():p(g)}})\
 50         .catch(function(){{setTimeout(function(){{p(n)}},1000)}})}})({generation})</script>\n"
 51    )
 52}
 53
 54/// A build counter that event streams wait on.
 55#[derive(Default)]
 56struct BuildSignal {
 57    generation: Mutex<u64>,
 58    changed: Condvar,
 59}
 60
 61impl BuildSignal {
 62    fn bump(&self) {
 63        *self.generation.lock().expect("build signal") += 1;
 64        self.changed.notify_all();
 65    }
 66}
 67
 68/// Run the development server until interrupted.
 69pub fn run(
 70    src: &Utf8Path,
 71    out: &Utf8Path,
 72    opts: &BuildOptions,
 73    host: &str,
 74    port: u16,
 75) -> Result<()> {
 76    if !src.is_dir() {
 77        anyhow::bail!("serve requires a source directory: serve <src-dir> -o <out-dir>");
 78    }
 79    let report = build_site(src, out, opts)?;
 80
 81    let address = format!("{host}:{port}");
 82    let server = Server::http(&address).map_err(|e| {
 83        anyhow::anyhow!("cannot listen on {address}: {e}. Is something already using port {port}?")
 84    })?;
 85    let server = Arc::new(server);
 86    let signal = Arc::new(BuildSignal::default());
 87    let root: Utf8PathBuf = out.to_owned();
 88
 89    println!(
 90        "serving {} page(s) from {out} at http://{address}/ — Ctrl-C to stop.",
 91        report.pages.len()
 92    );
 93
 94    // Rebuild in the background; the main thread serves.
 95    {
 96        let (src, out, opts, signal) = (
 97            src.to_owned(),
 98            out.to_owned(),
 99            opts.clone(),
100            Arc::clone(&signal),
101        );
102        std::thread::spawn(move || {
103            let result = crate::watch::run_with(&src, &out, &opts, |built| {
104                // Reload on a *successful* rebuild only. Reloading onto a stale page
105                // because the build just failed tells the author nothing; the error is
106                // already on their terminal.
107                if built.is_ok() {
108                    signal.bump();
109                }
110            });
111            if let Err(e) = result {
112                eprintln!("watch stopped: {e:#}");
113            }
114        });
115    }
116
117    // A thread per request. The volume is one developer's browser, and an event stream
118    // occupies its thread for as long as the tab is open — which a fixed pool would let
119    // starve everything else.
120    for request in server.incoming_requests() {
121        let root = root.clone();
122        let signal = Arc::clone(&signal);
123        std::thread::spawn(move || {
124            if let Err(e) = handle(request, &root, &signal) {
125                // A browser closing a tab mid-response is routine, not a problem.
126                if e.kind() != io::ErrorKind::BrokenPipe {
127                    eprintln!("serve: {e}");
128                }
129            }
130        });
131    }
132    Ok(())
133}
134
135fn handle(request: Request, root: &Utf8Path, signal: &Arc<BuildSignal>) -> io::Result<()> {
136    let url = request.url().to_string();
137    if url.split(['?', '#']).next() == Some(RELOAD_PATH) {
138        let since = since_parameter(&url);
139        return serve_poll(request, signal, since);
140    }
141    match resolve(root, &url) {
142        Some(path) => serve_file(request, &path, signal),
143        None => request.respond(
144            Response::from_string("404 not found")
145                .with_status_code(StatusCode(404))
146                .with_header(header("Content-Type", "text/plain; charset=utf-8")),
147        ),
148    }
149}
150
151fn serve_file(request: Request, path: &Utf8Path, signal: &Arc<BuildSignal>) -> io::Result<()> {
152    let Ok(bytes) = std::fs::read(path) else {
153        return request.respond(
154            Response::from_string("404 not found").with_status_code(StatusCode(404)),
155        );
156    };
157    let mime = mime_type(path);
158    let bytes = if mime.starts_with("text/html") {
159        let generation = *signal.generation.lock().expect("build signal");
160        inject_reload_script(&bytes, generation)
161    } else {
162        bytes
163    };
164    request.respond(
165        Response::from_data(bytes)
166            .with_header(header("Content-Type", mime))
167            // A dev server must never be cached, or an edit appears not to have landed.
168            .with_header(header("Cache-Control", "no-store")),
169    )
170}
171
172/// Put the reload script just before `</body>`, or at the end if there is none.
173pub fn inject_reload_script(bytes: &[u8], generation: u64) -> Vec<u8> {
174    let Ok(text) = std::str::from_utf8(bytes) else {
175        return bytes.to_vec();
176    };
177    let script = reload_script(generation);
178    match text.rfind("</body>") {
179        Some(at) => format!("{}{script}{}", &text[..at], &text[at..]).into_bytes(),
180        None => format!("{text}{script}").into_bytes(),
181    }
182}
183
184/// Answer a poll: block until the build generation passes `since`, then report it.
185///
186/// A timeout answers with the *current* generation, which the client compares itself —
187/// so a slow answer is indistinguishable from a fast one and no event can be missed.
188fn serve_poll(request: Request, signal: &Arc<BuildSignal>, since: u64) -> io::Result<()> {
189    let guard = signal.generation.lock().expect("build signal");
190    let (guard, _) = signal
191        .changed
192        .wait_timeout_while(guard, POLL_TIMEOUT, |generation| *generation <= since)
193        .expect("build signal");
194    let generation = *guard;
195    drop(guard);
196
197    request.respond(
198        Response::from_string(generation.to_string())
199            .with_header(header("Content-Type", "application/json"))
200            .with_header(header("Cache-Control", "no-store")),
201    )
202}
203
204/// The `since=N` parameter of a poll request.
205pub fn since_parameter(url: &str) -> u64 {
206    url.split_once('?')
207        .map(|(_, query)| query)
208        .into_iter()
209        .flat_map(|query| query.split('&'))
210        .find_map(|pair| pair.strip_prefix("since="))
211        .and_then(|value| value.parse().ok())
212        .unwrap_or(0)
213}
214
215/// Map a request URL onto a file inside `root`, or `None` if it does not name one.
216///
217/// This is the server's security boundary, so it is a pure function with its own tests.
218/// A URL is attacker-controlled input even on a development server: `..` segments,
219/// percent-encoded `..`, absolute paths and backslashes all have to resolve to nothing
220/// rather than to somewhere outside the output directory.
221pub fn resolve(root: &Utf8Path, url: &str) -> Option<Utf8PathBuf> {
222    let path = url.split(['?', '#']).next().unwrap_or("");
223    let decoded = percent_decode(path);
224
225    // Build the path from scratch out of accepted segments. Normalizing a joined path
226    // afterwards is the version of this that has bugs: it is far easier to reason about
227    // a list that never contained a `..` than about removing one correctly.
228    let mut segments: Vec<&str> = Vec::new();
229    for segment in decoded.split(['/', '\\']) {
230        match segment {
231            "" | "." => {}
232            ".." => return None,
233            // A NUL or a path separator that survived decoding is not a filename.
234            s if s.contains('\0') => return None,
235            s => segments.push(s),
236        }
237    }
238
239    let mut candidate = root.to_owned();
240    for segment in &segments {
241        candidate.push(segment);
242    }
243    // Directories, and the bare root, serve their index.
244    if decoded.ends_with('/') || segments.is_empty() || candidate.is_dir() {
245        candidate.push("index.html");
246    }
247    // Defence in depth: whatever the segment logic did, the result must be inside root.
248    if !candidate.starts_with(root) {
249        return None;
250    }
251    candidate.is_file().then_some(candidate)
252}
253
254/// Decode `%XX` escapes. `+` is left alone: it means a space in a query string, not in a
255/// path, and turning `a+b.html` into `a b.html` would break a real filename.
256fn percent_decode(input: &str) -> String {
257    let bytes = input.as_bytes();
258    let mut out: Vec<u8> = Vec::with_capacity(bytes.len());
259    let mut i = 0;
260    while i < bytes.len() {
261        if bytes[i] == b'%' && i + 2 < bytes.len() {
262            let hex = std::str::from_utf8(&bytes[i + 1..i + 3]).ok();
263            if let Some(byte) = hex.and_then(|h| u8::from_str_radix(h, 16).ok()) {
264                out.push(byte);
265                i += 3;
266                continue;
267            }
268        }
269        out.push(bytes[i]);
270        i += 1;
271    }
272    String::from_utf8_lossy(&out).into_owned()
273}
274
275fn mime_type(path: &Utf8Path) -> &'static str {
276    match path.extension().unwrap_or("").to_ascii_lowercase().as_str() {
277        "html" | "htm" => "text/html; charset=utf-8",
278        "css" => "text/css; charset=utf-8",
279        "js" => "text/javascript; charset=utf-8",
280        "json" => "application/json",
281        "xml" | "rss" | "atom" => "application/xml; charset=utf-8",
282        "txt" => "text/plain; charset=utf-8",
283        "svg" => "image/svg+xml",
284        "png" => "image/png",
285        "jpg" | "jpeg" => "image/jpeg",
286        "gif" => "image/gif",
287        "webp" => "image/webp",
288        "avif" => "image/avif",
289        "ico" => "image/x-icon",
290        "woff2" => "font/woff2",
291        "woff" => "font/woff",
292        "ttf" => "font/ttf",
293        "pdf" => "application/pdf",
294        _ => "application/octet-stream",
295    }
296}
297
298fn header(name: &str, value: &str) -> Header {
299    Header::from_bytes(name.as_bytes(), value.as_bytes()).expect("static header is valid")
300}