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;
const SELF_WRITE_SUPPRESS: Duration = Duration::from_millis(750);
pub struct DevOptions {
pub project_root: PathBuf,
pub port: u16,
pub open: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReloadEvent {
Reload,
Shutdown,
}
#[derive(Clone)]
pub struct AppState {
pub project_root: PathBuf,
pub out_dir: PathBuf,
pub docs_dir: PathBuf,
pub base: String,
pub port: u16,
pub reload_tx: broadcast::Sender<ReloadEvent>,
self_write_at_ms: Arc<AtomicU64>,
epoch: Instant,
engine: Option<Arc<std::sync::Mutex<docgen_build::DevState>>>,
}
impl AppState {
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,
}
}
pub fn with_engine(mut self, engine: docgen_build::DevState) -> Self {
self.engine = Some(Arc::new(std::sync::Mutex::new(engine)));
self
}
pub fn with_base(mut self, base: &str) -> Self {
self.base = normalize_base(base);
self
}
pub fn note_self_write(&self) {
let ms = self.epoch.elapsed().as_millis() as u64 + 1;
self.self_write_at_ms.store(ms, Ordering::SeqCst);
}
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
}
}
pub use docgen_config::normalize_base;
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,
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum PathGuardError {
NotMarkdown,
Absolute,
Traversal,
NotAFile,
NotFound,
}
pub fn resolve_doc_path(docs_dir: &Path, rel: &str) -> Result<PathBuf, PathGuardError> {
if rel.contains('\\') {
return Err(PathGuardError::Traversal);
}
if rel.starts_with('/') || Path::new(rel).is_absolute() {
return Err(PathGuardError::Absolute);
}
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, ".." => return Err(PathGuardError::Traversal),
other => kept.push(other),
}
}
if kept.is_empty() {
return Err(PathGuardError::Traversal);
}
let normalized = kept.join("/");
if !normalized.ends_with(".md") {
return Err(PathGuardError::NotMarkdown);
}
let candidate = docs_dir.join(&normalized);
if !candidate.starts_with(docs_dir) {
return Err(PathGuardError::Traversal);
}
let canonical = match candidate.canonicalize() {
Ok(p) => p,
Err(_) => return Err(PathGuardError::NotFound),
};
if !canonical.starts_with(docs_dir) {
return Err(PathGuardError::Traversal);
}
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)
}
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>
"#;
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
}
None => format!("{html}{DEV_HTML}"),
}
}
pub fn dev_bind_addr(port: u16) -> std::net::SocketAddr {
std::net::SocketAddr::from(([127, 0, 0, 1], port))
}
pub fn router(state: AppState) -> Router {
handlers::router(state)
}
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()?;
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)
};
let _ = state.reload_tx.send(ReloadEvent::Reload);
tracing::info!(
pages = page_count,
kind = ?kind,
elapsed_ms = start.elapsed().as_millis(),
"rebuilt + reloaded"
);
Ok(())
}
pub fn serve(opts: DevOptions) -> anyhow::Result<()> {
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());
let out_tmp = tempfile::tempdir()?;
let out_dir = out_tmp.path().to_path_buf();
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);
let start = std::time::Instant::now();
let (engine, first) = docgen_build::DevState::initial(&project_root, &out_dir)?;
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);
let _ = state.reload_tx.send(ReloadEvent::Reload);
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}"));
}
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(())
}
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>");
for marker in ["__docgen/livereload.js", "docgen-ctl--edit", "/edit/"] {
assert!(out.contains(marker), "missing injected marker {marker}");
}
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>"
);
}
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,
);
assert!(!state.take_self_write_suppression());
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() {
assert_eq!(strip_base("/docgen.css", ""), "/docgen.css");
assert_eq!(strip_base("/docs/docgen.css", "/docs"), "/docgen.css");
assert_eq!(strip_base("/docs/guide/intro", "/docs"), "/guide/intro");
assert_eq!(strip_base("/docs", "/docs"), "/");
assert_eq!(strip_base("/docs/", "/docs"), "/");
assert_eq!(strip_base("/docsxyz", "/docs"), "/docsxyz");
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);
}
}