docgen_server/

lib.rs

//! `docgen-server` is the **dev-only** server behind `docgen dev`: an axum app
//! (bound to `127.0.0.1` only) that serves the built site, watches `docs/` for
//! changes (debounced), rebuilds via [`docgen_build::build_site`], pushes a
//! live-reload signal over SSE, and exposes a path-guarded markdown write
//! endpoint for the in-browser editor.
//!
//! Nothing in this crate ships in a static `docgen build` dist: the editor UI,
//! the reload client, the write/SSE endpoints, and the vendored CodeMirror
//! assets exist ONLY while this server runs.

mod handlers;
mod watch;

use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use axum::Router;
use tokio::sync::broadcast;

/// Window after an editor-initiated write during which a watcher event for the
/// same on-disk change is treated as a duplicate and skipped. Must comfortably
/// cover the 200ms watcher debounce.
const SELF_WRITE_SUPPRESS: Duration = Duration::from_millis(750);

/// Dev-server configuration.
pub struct DevOptions {
    pub project_root: PathBuf,
    /// Loopback port. Default 4321.
    pub port: u16,
    /// Open a browser on start (off in tests/CI). Default false.
    pub open: bool,
}

/// One live-reload signal. Carried over the SSE channel.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReloadEvent {
    /// A rebuild finished; browsers should reload.
    Reload,
    /// The server is shutting down (Ctrl+C). Tells each live-reload SSE stream to
    /// END, so axum's graceful shutdown can drain the otherwise-eternal keep-alive
    /// connections instead of hanging on them.
    Shutdown,
}

/// Shared, cheaply-clonable state behind every handler. `Clone` bumps the
/// `Arc`/broadcast handle.
#[derive(Clone)]
pub struct AppState {
    pub project_root: PathBuf,
    pub out_dir: PathBuf,
    /// Canonicalized `docs/` dir — the write-guard root.
    pub docs_dir: PathBuf,
    /// Site `base` path (e.g. `/docs`), normalized to a leading-slash, no-trailing-slash
    /// form; empty when the site is served at the root. The built HTML prefixes every
    /// asset/nav/wikilink URL with this, so the dev server must strip it off incoming
    /// request paths before resolving them against `out_dir`.
    pub base: String,
    /// The loopback port the server is bound to. Used by the Host-header
    /// allowlist to defeat DNS-rebinding (see `handlers::loopback_guard`).
    pub port: u16,
    pub reload_tx: broadcast::Sender<ReloadEvent>,
    /// Shared "last editor write" clock (millis since `epoch`) used to suppress
    /// the duplicate watcher rebuild after an in-browser save. `0` = never.
    self_write_at_ms: Arc<AtomicU64>,
    /// Monotonic reference instant for `self_write_at_ms`.
    epoch: Instant,
    /// The incremental build engine, shared (and serialized) across the watcher
    /// thread and the editor-write handler so a single doc edit re-renders only
    /// the changed page instead of the whole O(n²) site. `None` in handler tests
    /// that construct an `AppState` without a live engine — those fall back to a
    /// full [`docgen_build::build_site`] in [`rebuild_and_reload`].
    engine: Option<Arc<std::sync::Mutex<docgen_build::DevState>>>,
}

impl AppState {
    /// Construct an `AppState`. Initializes the self-write suppression clock.
    pub fn new(
        project_root: PathBuf,
        out_dir: PathBuf,
        docs_dir: PathBuf,
        port: u16,
        reload_tx: broadcast::Sender<ReloadEvent>,
    ) -> Self {
        Self {
            project_root,
            out_dir,
            docs_dir,
            base: String::new(),
            port,
            reload_tx,
            self_write_at_ms: Arc::new(AtomicU64::new(0)),
            epoch: Instant::now(),
            engine: None,
        }
    }

    /// Attach the incremental build engine. Builder-style so the `new()` call
    /// sites (including tests) stay untouched. Only `serve_async` sets one.
    pub fn with_engine(mut self, engine: docgen_build::DevState) -> Self {
        self.engine = Some(Arc::new(std::sync::Mutex::new(engine)));
        self
    }

    /// Set the site `base` path, normalized to a leading-slash, no-trailing-slash
    /// form (`docs` / `/docs/` / `docs/` all become `/docs`; empty stays empty).
    /// Builder-style so the existing `new()` call sites stay untouched.
    pub fn with_base(mut self, base: &str) -> Self {
        self.base = normalize_base(base);
        self
    }

    /// Record that the editor just wrote a doc on disk. The watcher will skip
    /// the next change it sees within [`SELF_WRITE_SUPPRESS`] as a duplicate.
    pub fn note_self_write(&self) {
        // Store elapsed + 1 so the "never written" sentinel (0) is unambiguous
        // even when the write happens at elapsed == 0.
        let ms = self.epoch.elapsed().as_millis() as u64 + 1;
        self.self_write_at_ms.store(ms, Ordering::SeqCst);
    }

    /// Whether a watcher event right now should be suppressed as the echo of a
    /// recent editor write. Consumes the marker so only one rebuild is skipped.
    pub fn take_self_write_suppression(&self) -> bool {
        let marked = self.self_write_at_ms.swap(0, Ordering::SeqCst);
        if marked == 0 {
            return false;
        }
        let now = self.epoch.elapsed().as_millis() as u64 + 1;
        now.saturating_sub(marked) <= SELF_WRITE_SUPPRESS.as_millis() as u64
    }
}

/// Normalize a configured `base` into a leading-slash, no-trailing-slash form
/// for prefix matching: `""` -> `""`, `"docs"` / `"/docs/"` / `"docs/"` -> `"/docs"`.
/// Re-exported from `docgen-config` so the dev server's request-path stripping
/// and the build's URL prefixing share one canonicalization.
pub use docgen_config::normalize_base;

/// Strip the site `base` prefix from a (already percent-decoded) request path so
/// it can be resolved against `out_dir`. `base` must be normalized
/// (`normalize_base`). A request that does not fall under `base` is returned
/// unchanged — the caller then resolves it normally (and likely 404s), matching
/// production where the host only routes in-base requests to the site.
///
/// `/docs/x.css` -> `/x.css`; `/docs` and `/docs/` -> `/`; `/other` -> `/other`.
pub fn strip_base<'a>(path: &'a str, base: &str) -> &'a str {
    if base.is_empty() {
        return path;
    }
    match path.strip_prefix(base) {
        Some("") => "/",
        Some(rest) if rest.starts_with('/') => rest,
        _ => path,
    }
}

/// Errors from [`resolve_doc_path`]; each maps to an HTTP status (see `handlers`).
#[derive(Debug, PartialEq, Eq)]
pub enum PathGuardError {
    /// Not a `.md` path. (400)
    NotMarkdown,
    /// Absolute path or leading `/`. (400)
    Absolute,
    /// `..` component, backslash, or a realpath that escapes `docs/`. (403)
    Traversal,
    /// Resolves to something that is not a regular file. (400)
    NotAFile,
    /// In-bounds but the file does not exist. (404)
    NotFound,
}

/// Resolve a client-supplied doc-relative path (e.g. `"guide/intro.md"`) to a
/// canonical absolute path strictly inside `docs_dir`, or reject. `docs_dir`
/// MUST already be canonicalized by the caller. Layered checks mirror the
/// original `validateRepoDocPath`:
///
/// 1. backslash -> `Traversal`; absolute / leading `/` -> `Absolute`.
/// 2. strip leading `./`; any `..` component -> `Traversal`; empty -> `Traversal`.
/// 3. require a `.md` suffix -> else `NotMarkdown`.
/// 4. lexical: `docs_dir.join(rel)` must stay under `docs_dir`.
/// 5. `canonicalize()`: missing -> `NotFound`; realpath escaping `docs_dir`
///    (symlink escape) -> `Traversal`.
/// 6. the canonical target must be a regular file -> else `NotAFile`.
pub fn resolve_doc_path(docs_dir: &Path, rel: &str) -> Result<PathBuf, PathGuardError> {
    // (1) gross-shape rejections.
    if rel.contains('\\') {
        return Err(PathGuardError::Traversal);
    }
    if rel.starts_with('/') || Path::new(rel).is_absolute() {
        return Err(PathGuardError::Absolute);
    }

    // (2) normalize + component scan.
    let trimmed = rel.strip_prefix("./").unwrap_or(rel);
    if trimmed.is_empty() {
        return Err(PathGuardError::Traversal);
    }
    let mut kept: Vec<&str> = Vec::new();
    for comp in trimmed.split('/') {
        match comp {
            "" | "." => continue, // collapse `//` and `.` segments
            ".." => return Err(PathGuardError::Traversal),
            other => kept.push(other),
        }
    }
    if kept.is_empty() {
        return Err(PathGuardError::Traversal);
    }
    let normalized = kept.join("/");

    // (3) extension whitelist (markdown-only; the TS guard also allowed `.svx`,
    // which the Rust rewrite does not support).
    if !normalized.ends_with(".md") {
        return Err(PathGuardError::NotMarkdown);
    }

    // (4) lexical containment.
    let candidate = docs_dir.join(&normalized);
    if !candidate.starts_with(docs_dir) {
        return Err(PathGuardError::Traversal);
    }

    // (5) realpath check (catches symlink escapes).
    // Any canonicalize failure (missing path, permission, etc.) is reported as
    // NotFound for a dev write endpoint — the client cannot act on finer detail.
    let canonical = match candidate.canonicalize() {
        Ok(p) => p,
        Err(_) => return Err(PathGuardError::NotFound),
    };
    if !canonical.starts_with(docs_dir) {
        return Err(PathGuardError::Traversal);
    }

    // (6) must be a regular file (not a dir, not a symlink-to-dir).
    let meta = match std::fs::symlink_metadata(&canonical) {
        Ok(m) => m,
        Err(_) => return Err(PathGuardError::NotFound),
    };
    if !meta.is_file() {
        return Err(PathGuardError::NotAFile);
    }

    Ok(canonical)
}

/// The dev-only HTML injected before `</body>` of every served page: the editor
/// css, a tiny inline script that inserts the dev-only edit icon into the
/// topbar's `.docgen-btn-strip` (next to the diff/full-width controls — so the
/// static `docgen build` output never contains it), the editor island panel
/// element, the vendored CodeMirror UMD scripts (loaded in dependency order:
/// core -> xml mode -> overlay addon -> markdown mode), the editor island JS,
/// and the live-reload client.
///
/// These non-`defer` scripts execute at parse time — before the page's deferred
/// Alpine script fires `alpine:init` — so `editor.js`'s `docgen.island(...)`
/// registration lands before Alpine runs the registry. Injected ONLY by the dev
/// server (`inject_dev_html`); never written to disk by `docgen build`.
const DEV_HTML: &str = r#"
<script>(function(){
  var strip=document.querySelector('.docgen-btn-strip');
  if(!strip)return;
  // The full-page CM6 editor lives at /edit/<slug>; the pencil links to it.
  var slug=location.pathname.replace(/^\/+|\/+$/g,'');
  if(slug==='')slug='index';
  var a=document.createElement('a');
  a.className='icon-only docgen-ctl--edit';
  a.setAttribute('href','/edit/'+slug);
  a.setAttribute('aria-label','Edit this page');
  a.setAttribute('title','Edit this page (dev)');
  a.innerHTML='<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" aria-hidden="true"><path d="M12 20h9"/><path d="M16.5 3.5a2.1 2.1 0 0 1 3 3L7 19l-4 1 1-4Z"/></svg>';
  var fw=strip.querySelector('.docgen-ctl--fullwidth');
  if(fw)strip.insertBefore(a,fw);else strip.appendChild(a);
})();</script>
<script src="/__docgen/livereload.js"></script>
"#;

/// Post-process a served HTML body: inject the reload-client script + editor
/// toggle + editor island scripts/styles immediately before `</body>`. Dev-only;
/// never run by `docgen build`. Pure string fn so it is unit-testable.
pub fn inject_dev_html(html: &str) -> String {
    match html.rfind("</body>") {
        Some(i) => {
            let mut s = String::with_capacity(html.len() + DEV_HTML.len());
            s.push_str(&html[..i]);
            s.push_str(DEV_HTML);
            s.push_str(&html[i..]);
            s
        }
        // Graceful: append if there is no closing body tag.
        None => format!("{html}{DEV_HTML}"),
    }
}

/// The loopback bind address for the dev server. NEVER `0.0.0.0` — the dev
/// server (editor + write endpoint) must not be reachable off-host.
pub fn dev_bind_addr(port: u16) -> std::net::SocketAddr {
    std::net::SocketAddr::from(([127, 0, 0, 1], port))
}

/// Build the axum router (NO listener) for the given state. Split out so handler
/// tests can `oneshot` requests without binding a port.
pub fn router(state: AppState) -> Router {
    handlers::router(state)
}

/// Rebuild the site into `state.out_dir` (Dev mode + dev-asset emission), then
/// broadcast a reload. Called on every debounced fs change AND after a successful
/// editor write. Returns `Err` only on a hard build failure; the caller logs and
/// keeps serving the last good build.
///
/// When the incremental [`engine`](AppState::engine) is present (the live dev
/// server), each rebuild re-renders only the doc(s) that actually changed and
/// leaves the rest of the site untouched — turning a ~10s full rebuild of a large
/// corpus into milliseconds. The dev-only assets (CodeMirror, editor island,
/// reload client) are re-emitted only after a *full* rebuild, since the atomic
/// staging swap that backs a full build wipes `out_dir`; an incremental rebuild
/// writes pages in place and leaves the dev assets intact.
///
/// Without an engine (handler tests constructing a bare `AppState`), this falls
/// back to a full atomic [`docgen_build::build_site`] + dev-asset emit — a failed
/// rebuild leaves the served dir (the previous good build) untouched.
pub fn rebuild_and_reload(state: &AppState) -> anyhow::Result<()> {
    let start = std::time::Instant::now();

    let (page_count, kind) = if let Some(engine) = &state.engine {
        let mut engine = engine
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let rebuilt = engine.rebuild()?;
        // A full fallback re-creates `out_dir` from scratch, wiping the dev-only
        // assets; re-emit them. The incremental path leaves them in place.
        if rebuilt.kind == docgen_build::RebuildKind::Full {
            docgen_assets::emit(&docgen_assets::dev_assets(), &state.out_dir)?;
        }
        (rebuilt.page_count, Some(rebuilt.kind))
    } else {
        let outcome = docgen_build::build_site(&docgen_build::BuildOptions {
            project_root: &state.project_root,
            out_dir: &state.out_dir,
            mode: docgen_build::BuildMode::Dev,
        })?;
        docgen_assets::emit(&docgen_assets::dev_assets(), &state.out_dir)?;
        (outcome.page_count, None)
    };

    // Ignore "no subscribers" — a reload with nobody listening is fine.
    let _ = state.reload_tx.send(ReloadEvent::Reload);
    tracing::info!(
        pages = page_count,
        kind = ?kind,
        elapsed_ms = start.elapsed().as_millis(),
        "rebuilt + reloaded"
    );
    Ok(())
}

/// Run the dev server: initial build, spawn the debounced watcher, bind
/// `127.0.0.1`, serve until Ctrl-C. Blocking entry point the `docgen dev` CLI
/// arm calls. Owns its own tokio runtime so the `docgen` bin's `main` stays a
/// plain `fn main() -> Result<()>`.
pub fn serve(opts: DevOptions) -> anyhow::Result<()> {
    // Idempotent: a second `serve` in-process (tests) won't panic.
    let _ = tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init();

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()?;
    runtime.block_on(serve_async(opts))
}

async fn serve_async(opts: DevOptions) -> anyhow::Result<()> {
    let project_root = opts.project_root.clone();
    let docs_dir = project_root.join("docs");
    let docs_canon = docs_dir.canonicalize().unwrap_or_else(|_| docs_dir.clone());

    // A process-owned output dir; kept alive for the whole run, auto-cleaned.
    let out_tmp = tempfile::tempdir()?;
    let out_dir = out_tmp.path().to_path_buf();

    // The built HTML prefixes every URL with the *resolved* `base`; the server
    // must strip that same prefix off incoming requests. Resolve it exactly as
    // `build_site` does (DOCGEN_BASE / CI Pages env override docgen.toml) so dev
    // and build never disagree — otherwise `DOCGEN_BASE=/x docgen dev`, or a dev
    // run under CI, would serve links the request router can't strip and 404.
    // A malformed config here is non-fatal for serving (the build step reports
    // it) — fall back to "served at root".
    let base = docgen_config::load(&project_root)
        .map(|c| docgen_config::resolve_base(&c.base))
        .unwrap_or_default();

    let (reload_tx, _rx) = broadcast::channel(16);

    // Initial full build (Dev mode), which also seeds the incremental engine.
    // Subsequent rebuilds (watcher / editor write) go through the engine and
    // re-render only what changed. The user accepts a slower first build for
    // blazingly fast incremental updates.
    let start = std::time::Instant::now();
    let (engine, first) = docgen_build::DevState::initial(&project_root, &out_dir)?;
    // Dev-only assets (editor island, CodeMirror, reload client) — emitted once
    // after the initial build; the incremental path leaves them in place.
    docgen_assets::emit(&docgen_assets::dev_assets(), &out_dir)?;
    tracing::info!(
        pages = first.page_count,
        elapsed_ms = start.elapsed().as_millis(),
        "initial build"
    );

    let state = AppState::new(
        project_root,
        out_dir,
        docs_canon.clone(),
        opts.port,
        reload_tx,
    )
    .with_base(&base)
    .with_engine(engine);

    // Tell any (re)connecting browser the first build is ready.
    let _ = state.reload_tx.send(ReloadEvent::Reload);

    // Spawn the debounced fs watcher; it rebuilds + reloads on every change.
    let _watcher = watch::spawn_watcher(state.clone(), &docs_canon)?;

    let addr = dev_bind_addr(opts.port);
    let listener = tokio::net::TcpListener::bind(addr).await?;
    tracing::info!("docgen dev server: http://{addr}");
    if opts.open {
        let _ = open_browser(&format!("http://{addr}"));
    }

    // Clone the reload sender before `state` moves into the router; the shutdown
    // hook uses it to terminate the open live-reload SSE streams so the graceful
    // drain can complete (otherwise the 15s keep-alive holds the process open).
    let shutdown_tx = state.reload_tx.clone();
    axum::serve(listener, router(state))
        .with_graceful_shutdown(async move {
            let _ = tokio::signal::ctrl_c().await;
            let _ = shutdown_tx.send(ReloadEvent::Shutdown);
        })
        .await?;
    Ok(())
}

/// Best-effort browser open (dev convenience; failures are non-fatal).
fn open_browser(url: &str) -> std::io::Result<()> {
    #[cfg(target_os = "macos")]
    let cmd = "open";
    #[cfg(all(unix, not(target_os = "macos")))]
    let cmd = "xdg-open";
    #[cfg(windows)]
    let cmd = "explorer";
    std::process::Command::new(cmd).arg(url).spawn().map(|_| ())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn inject_dev_html_inserts_before_body() {
        let out = inject_dev_html("<html><body><p>hi</p></body></html>");
        // Reload client + the edit pencil (a link to the /edit/<slug> CM6 editor).
        for marker in ["__docgen/livereload.js", "docgen-ctl--edit", "/edit/"] {
            assert!(out.contains(marker), "missing injected marker {marker}");
        }
        // Every injected marker precedes the closing body tag.
        let body = out.rfind("</body>").unwrap();
        for marker in ["__docgen/livereload.js", "docgen-ctl--edit"] {
            assert!(
                out.find(marker).unwrap() < body,
                "{marker} not before </body>"
            );
        }
        // The static build never contains the dev edit affordance.
        assert!(!"<html><body></body></html>".contains("docgen-ctl--edit"));
    }

    #[test]
    fn inject_dev_html_no_body_appends() {
        let out = inject_dev_html("<p>no body tag here</p>");
        assert!(out.contains("__docgen/livereload.js"));
        assert!(out.contains("docgen-ctl--edit"));
    }

    #[test]
    fn self_write_suppression_is_one_shot() {
        let (tx, _rx) = broadcast::channel(4);
        let state = AppState::new(
            PathBuf::from("/x"),
            PathBuf::from("/x/out"),
            PathBuf::from("/x/docs"),
            4321,
            tx,
        );
        // No write yet -> nothing to suppress.
        assert!(!state.take_self_write_suppression());
        // After a self-write, exactly one watcher event is suppressed.
        state.note_self_write();
        assert!(state.take_self_write_suppression());
        assert!(!state.take_self_write_suppression());
    }

    #[test]
    fn normalize_base_canonicalizes() {
        assert_eq!(normalize_base(""), "");
        assert_eq!(normalize_base("/"), "");
        assert_eq!(normalize_base("docs"), "/docs");
        assert_eq!(normalize_base("/docs"), "/docs");
        assert_eq!(normalize_base("/docs/"), "/docs");
        assert_eq!(normalize_base("docs/"), "/docs");
        assert_eq!(normalize_base("/a/b"), "/a/b");
    }

    #[test]
    fn strip_base_handles_prefix_and_misses() {
        // Empty base: everything passes through unchanged.
        assert_eq!(strip_base("/docgen.css", ""), "/docgen.css");
        // In-base requests get the prefix removed.
        assert_eq!(strip_base("/docs/docgen.css", "/docs"), "/docgen.css");
        assert_eq!(strip_base("/docs/guide/intro", "/docs"), "/guide/intro");
        // The base root itself maps to "/".
        assert_eq!(strip_base("/docs", "/docs"), "/");
        assert_eq!(strip_base("/docs/", "/docs"), "/");
        // A path that merely shares a textual prefix is NOT in-base.
        assert_eq!(strip_base("/docsxyz", "/docs"), "/docsxyz");
        // Out-of-base requests are returned unchanged.
        assert_eq!(strip_base("/other", "/docs"), "/other");
    }

    #[test]
    fn bind_addr_is_loopback() {
        assert!(dev_bind_addr(4321).ip().is_loopback());
        assert_eq!(dev_bind_addr(4321).port(), 4321);
    }
}