use std::fs;
use std::path::{Path, PathBuf};
use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use tokio::sync::broadcast;
use tower::ServiceExt;
use docgen_build::{build_site, BuildMode, BuildOptions};
use docgen_server::{rebuild_and_reload, router, AppState, ReloadEvent};
fn setup_fixture(root: &Path) {
let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR")); let workspace = manifest.parent().unwrap().parent().unwrap();
let fixture = workspace.join("fixtures/site-basic");
fs::create_dir_all(root.join("docs/guide")).unwrap();
fs::copy(fixture.join("docs/index.md"), root.join("docs/index.md")).unwrap();
fs::copy(
fixture.join("docs/guide/intro.md"),
root.join("docs/guide/intro.md"),
)
.unwrap();
}
fn state_with_built_site() -> (tempfile::TempDir, tempfile::TempDir, AppState) {
let root = tempfile::tempdir().unwrap();
setup_fixture(root.path());
let out = tempfile::tempdir().unwrap();
build_site(&BuildOptions {
project_root: root.path(),
out_dir: out.path(),
mode: BuildMode::Dev,
})
.unwrap();
docgen_assets::emit(&docgen_assets::dev_assets(), out.path()).unwrap();
let docs_dir = root.path().join("docs").canonicalize().unwrap();
let (reload_tx, _rx) = broadcast::channel(16);
let state = AppState::new(
root.path().to_path_buf(),
out.path().to_path_buf(),
docs_dir,
4321,
reload_tx,
);
(root, out, state)
}
#[tokio::test]
async fn serves_built_index_with_injected_dev_html() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/")
.header("host", "127.0.0.1:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = resp.into_body().collect().await.unwrap().to_bytes();
let html = String::from_utf8(body.to_vec()).unwrap();
assert!(
html.contains("<title>Home</title>"),
"not the index: {html}"
);
assert!(html.contains("/__docgen/livereload.js"));
}
#[tokio::test]
async fn serves_static_asset_without_injection() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/docgen.css")
.header("host", "127.0.0.1:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let ct = resp
.headers()
.get("content-type")
.unwrap()
.to_str()
.unwrap()
.to_string();
assert!(ct.contains("css"), "unexpected content-type: {ct}");
let body = resp.into_body().collect().await.unwrap().to_bytes();
let css = String::from_utf8(body.to_vec()).unwrap();
assert!(!css.contains("/__docgen/livereload.js"));
}
#[tokio::test]
async fn unknown_path_404() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/does/not/exist")
.header("host", "127.0.0.1:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
let ct = resp
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
assert!(ct.contains("html"), "404 should be html, got {ct}");
let body = resp.into_body().collect().await.unwrap().to_bytes();
let html = String::from_utf8(body.to_vec()).unwrap();
assert!(html.contains("docgen-sidebar"), "404 lacks the nav shell");
assert!(html.contains("404"), "404 body missing the 404 marker");
assert!(
html.contains("/__docgen/livereload.js"),
"404 missing dev injection"
);
}
#[tokio::test]
async fn serves_dev_livereload_client_asset() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/__docgen/livereload.js")
.header("host", "127.0.0.1:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = resp.into_body().collect().await.unwrap().to_bytes();
let js = String::from_utf8(body.to_vec()).unwrap();
assert!(js.contains("EventSource('/__docgen/livereload')"));
}
#[tokio::test]
async fn rejects_foreign_host_header_dns_rebinding() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/")
.header("host", "attacker.com:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn rejects_missing_host_header() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn accepts_localhost_host_header() {
let (_root, _out, state) = state_with_built_site();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/")
.header("host", "localhost:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn rejects_cross_site_origin_on_mutating_put() {
let (root, _out, state) = state_with_built_site();
let secret_before = fs::read_to_string(root.path().join("docs/index.md")).unwrap();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.method("PUT")
.uri("/__docgen/source")
.header("host", "127.0.0.1:4321")
.header("origin", "http://attacker.com")
.header("content-type", "application/json")
.body(Body::from(
serde_json::json!({ "path": "index.md", "source": "HACKED" }).to_string(),
))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
assert_eq!(
fs::read_to_string(root.path().join("docs/index.md")).unwrap(),
secret_before
);
}
#[tokio::test]
async fn livereload_endpoint_streams_reload_event() {
use tokio::time::{timeout, Duration};
let (_root, _out, state) = state_with_built_site();
let tx = state.reload_tx.clone();
let app = router(state);
let resp = app
.oneshot(
Request::builder()
.uri("/__docgen/livereload")
.header("host", "127.0.0.1:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let ct = resp
.headers()
.get("content-type")
.unwrap()
.to_str()
.unwrap()
.to_string();
assert!(ct.contains("text/event-stream"), "not SSE: {ct}");
let mut body = resp.into_body().into_data_stream();
tx.send(ReloadEvent::Reload).unwrap();
use futures::StreamExt as _;
let chunk = timeout(Duration::from_secs(2), body.next())
.await
.expect("SSE frame timed out")
.expect("stream ended")
.expect("stream error");
let frame = String::from_utf8(chunk.to_vec()).unwrap();
assert!(frame.contains("event: reload"), "frame: {frame}");
assert!(frame.contains("data: now"), "frame: {frame}");
}
#[tokio::test]
async fn rebuild_broadcasts_reload() {
let (_root, _out, state) = state_with_built_site();
let mut rx = state.reload_tx.subscribe();
rebuild_and_reload(&state).unwrap();
assert_eq!(rx.recv().await.unwrap(), ReloadEvent::Reload);
}
#[tokio::test]
async fn rebuild_regenerates_changed_page() {
let (root, out, state) = state_with_built_site();
let before = fs::read_to_string(out.path().join("index/index.html")).unwrap();
assert!(before.contains("<title>Home</title>"));
fs::write(
root.path().join("docs/index.md"),
"---\ntitle: Renamed Home\n---\n\nfresh body\n",
)
.unwrap();
rebuild_and_reload(&state).unwrap();
let after = fs::read_to_string(out.path().join("index/index.html")).unwrap();
assert!(
after.contains("Renamed Home"),
"rebuild did not pick up the edit: {after}"
);
}
#[tokio::test]
async fn failed_build_does_not_broadcast() {
let (root, _out, state) = state_with_built_site();
let mut rx = state.reload_tx.subscribe();
fs::remove_dir_all(root.path().join("docs")).unwrap();
let res = rebuild_and_reload(&state);
assert!(res.is_err(), "expected a hard build failure");
assert!(rx.try_recv().is_err());
}
fn state_with_site(
files: &[(&str, &str)],
toml: &str,
) -> (tempfile::TempDir, tempfile::TempDir, AppState) {
let root = tempfile::tempdir().unwrap();
fs::write(root.path().join("docgen.toml"), toml).unwrap();
for (rel, body) in files {
let path = root.path().join("docs").join(rel);
fs::create_dir_all(path.parent().unwrap()).unwrap();
fs::write(path, body).unwrap();
}
let out = tempfile::tempdir().unwrap();
build_site(&BuildOptions {
project_root: root.path(),
out_dir: out.path(),
mode: BuildMode::Dev,
})
.unwrap();
docgen_assets::emit(&docgen_assets::dev_assets(), out.path()).unwrap();
let docs_dir = root.path().join("docs").canonicalize().unwrap();
let base = docgen_config::load(root.path()).unwrap().base;
let (reload_tx, _rx) = broadcast::channel(16);
let state = AppState::new(
root.path().to_path_buf(),
out.path().to_path_buf(),
docs_dir,
4321,
reload_tx,
)
.with_base(&base);
(root, out, state)
}
async fn get_status(app: axum::Router, uri: &str) -> StatusCode {
app.oneshot(
Request::builder()
.uri(uri)
.header("host", "127.0.0.1:4321")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap()
.status()
}
#[tokio::test]
async fn serves_non_ascii_slug_via_percent_encoded_request() {
let (_root, _out, state) = state_with_site(
&[("café/résumé.md", "---\ntitle: R\n---\n\nbody\n")],
"base = \"\"\n",
);
let app = router(state);
assert_eq!(
get_status(app, "/caf%C3%A9/r%C3%A9sum%C3%A9").await,
StatusCode::OK
);
}
#[tokio::test]
async fn strips_base_prefix_from_requests() {
let (_root, _out, state) = state_with_site(
&[
("index.md", "---\ntitle: H\n---\n\nhome\n"),
("guide.md", "---\ntitle: G\n---\n\nbody\n"),
],
"base = \"/docs\"\n",
);
let app = router(state.clone());
assert_eq!(
get_status(app.clone(), "/docs/docgen.css").await,
StatusCode::OK
);
assert_eq!(get_status(app.clone(), "/docs/guide").await, StatusCode::OK);
assert_eq!(get_status(app, "/docs/").await, StatusCode::OK);
}