use std::collections::HashMap;
use std::fmt::Write as _;
use std::net::SocketAddr;
use std::sync::Arc;
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::{Html, IntoResponse, Json, Response};
use axum::routing::get;
use axum::Router;
use tokio_util::sync::CancellationToken;
use crate::error::{DocsError, Result};
use crate::index::Index;
use crate::render;
use crate::search::SearchIndex;
pub struct DocsServer {
state: Arc<ServerState>,
}
struct ServerState {
index: Index,
pages: HashMap<String, String>, search: SearchIndex,
}
impl DocsServer {
pub fn new(index: Index, pages: HashMap<String, String>) -> Result<Self> {
let search_input: Vec<(String, String, String)> = index
.entries()
.map(|(_section, entry)| {
let body = pages.get(&entry.path).cloned().unwrap_or_default();
(entry.path.clone(), entry.title.clone(), body)
})
.collect();
let search = SearchIndex::build(&search_input)?;
Ok(Self { state: Arc::new(ServerState { index, pages, search }) })
}
pub async fn run(
self,
bind: SocketAddr,
bound: tokio::sync::oneshot::Sender<SocketAddr>,
cancel: CancellationToken,
) -> Result<()> {
let listener = tokio::net::TcpListener::bind(bind)
.await
.map_err(|e| DocsError::Server(format!("bind {bind}: {e}")))?;
let local_addr =
listener.local_addr().map_err(|e| DocsError::Server(format!("local_addr: {e}")))?;
let _ = bound.send(local_addr);
let app = self.router();
axum::serve(listener, app)
.with_graceful_shutdown(async move { cancel.cancelled().await })
.await
.map_err(|e| DocsError::Server(e.to_string()))
}
pub fn router(&self) -> Router {
let state = Arc::clone(&self.state);
Router::new()
.route("/", get(root_handler))
.route("/search", get(search_handler))
.route("/{*path}", get(page_handler))
.with_state(state)
}
}
async fn root_handler(State(state): State<Arc<ServerState>>) -> Html<String> {
let title = html_escape(&state.index.title);
let mut body = String::new();
let _ = writeln!(body, "<h1>{title}</h1>");
for section in &state.index.sections {
let _ = writeln!(body, "<h2>{}</h2>\n<ul>", html_escape(§ion.title));
for page in §ion.pages {
let href = page_href(&page.path);
let _ = writeln!(
body,
" <li><a href=\"{}\">{}</a></li>",
html_escape(&href),
html_escape(&page.title),
);
}
body.push_str("</ul>\n");
}
Html(render::to_html_document(&state.index.title, &markdownify(&body)))
}
fn markdownify(html: &str) -> String {
html.to_string()
}
async fn page_handler(
State(state): State<Arc<ServerState>>,
axum::extract::Path(path): axum::extract::Path<String>,
) -> Response {
let page_path = path.strip_suffix(".html").unwrap_or(&path).to_string();
if is_unsafe_path(&page_path) {
return (StatusCode::NOT_FOUND, Json(error_body("path not allowed"))).into_response();
}
let Some(body) =
state.pages.get(&format!("{page_path}.md")).or_else(|| state.pages.get(&page_path))
else {
return (StatusCode::NOT_FOUND, Json(error_body(&format!("page not found: {path}"))))
.into_response();
};
let title = first_heading(body).unwrap_or_else(|| page_path.clone());
let html = render::to_html_document(&title, body);
Html(html).into_response()
}
#[derive(Debug, serde::Deserialize)]
struct SearchQuery {
q: Option<String>,
#[serde(default = "default_limit")]
limit: usize,
}
const fn default_limit() -> usize {
10
}
#[derive(Debug, serde::Serialize)]
struct SearchResponse {
results: Vec<SearchHit>,
}
#[derive(Debug, serde::Serialize)]
struct SearchHit {
path: String,
title: String,
snippet: String,
score: f32,
}
async fn search_handler(
State(state): State<Arc<ServerState>>,
Query(params): Query<SearchQuery>,
) -> Response {
let q = params.q.unwrap_or_default();
match state.search.full_text_search(&q, params.limit) {
Ok(hits) => Json(SearchResponse {
results: hits
.into_iter()
.map(|h| SearchHit {
path: h.path,
title: h.title,
snippet: h.snippet,
score: h.score,
})
.collect(),
})
.into_response(),
Err(e) => {
(StatusCode::INTERNAL_SERVER_ERROR, Json(error_body(&format!("search failed: {e}"))))
.into_response()
}
}
}
fn error_body(msg: &str) -> serde_json::Value {
serde_json::json!({ "error": msg })
}
fn is_unsafe_path(path: &str) -> bool {
let p = std::path::Path::new(path);
p.is_absolute()
|| p.components()
.any(|c| matches!(c, std::path::Component::ParentDir | std::path::Component::Prefix(_)))
}
fn page_href(path: &str) -> String {
let stripped = path.strip_suffix(".md").unwrap_or(path);
format!("/{stripped}.html")
}
fn first_heading(body: &str) -> Option<String> {
body.lines().find_map(|l| l.strip_prefix("# ").map(|s| s.trim().to_string()))
}
fn html_escape(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for c in s.chars() {
match c {
'&' => out.push_str("&"),
'<' => out.push_str("<"),
'>' => out.push_str(">"),
'"' => out.push_str("""),
'\'' => out.push_str("'"),
other => out.push(other),
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::index::{IndexEntry, IndexSection};
fn sample_fixture() -> (Index, HashMap<String, String>) {
let index = Index {
title: "My Docs".into(),
sections: vec![IndexSection {
title: "Getting started".into(),
pages: vec![
IndexEntry { path: "intro.md".into(), title: "Introduction".into() },
IndexEntry { path: "install.md".into(), title: "Install".into() },
],
}],
};
let mut pages = HashMap::new();
pages.insert("intro.md".into(), "# Introduction\n\nWelcome to the tool.".into());
pages.insert("install.md".into(), "# Install\n\nRun cargo install widget.".into());
(index, pages)
}
#[tokio::test]
async fn root_lists_every_entry() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let router = server.router();
let request = axum::http::Request::builder()
.uri("/")
.body(axum::body::Body::empty())
.expect("build request");
let response = tower::ServiceExt::oneshot(router, request).await.expect("oneshot");
assert_eq!(response.status(), StatusCode::OK);
let body_bytes =
axum::body::to_bytes(response.into_body(), usize::MAX).await.expect("body");
let body = String::from_utf8_lossy(&body_bytes);
assert!(body.contains("My Docs"), "title: {body}");
assert!(body.contains("Introduction"), "intro link: {body}");
assert!(body.contains("Install"), "install link: {body}");
}
#[tokio::test]
async fn page_by_html_suffix_renders() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let router = server.router();
let request = axum::http::Request::builder()
.uri("/intro.html")
.body(axum::body::Body::empty())
.expect("build request");
let response = tower::ServiceExt::oneshot(router, request).await.expect("oneshot");
assert_eq!(response.status(), StatusCode::OK);
let body_bytes =
axum::body::to_bytes(response.into_body(), usize::MAX).await.expect("body");
let body = String::from_utf8_lossy(&body_bytes);
assert!(body.contains("<h1>Introduction</h1>"), "page body: {body}");
}
#[tokio::test]
async fn unknown_page_returns_404() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let router = server.router();
let request = axum::http::Request::builder()
.uri("/nope.html")
.body(axum::body::Body::empty())
.expect("build request");
let response = tower::ServiceExt::oneshot(router, request).await.expect("oneshot");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn traversal_attempt_is_404_not_filesystem_read() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let router = server.router();
let request = axum::http::Request::builder()
.uri("/intro.md")
.body(axum::body::Body::empty())
.expect("build request");
let response = tower::ServiceExt::oneshot(router, request).await.expect("oneshot");
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn search_endpoint_returns_json_hits() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let router = server.router();
let request = axum::http::Request::builder()
.uri("/search?q=welcome")
.body(axum::body::Body::empty())
.expect("build request");
let response = tower::ServiceExt::oneshot(router, request).await.expect("oneshot");
assert_eq!(response.status(), StatusCode::OK);
let body_bytes =
axum::body::to_bytes(response.into_body(), usize::MAX).await.expect("body");
let body: serde_json::Value = serde_json::from_slice(&body_bytes).expect("json");
let results = body["results"].as_array().expect("results array");
assert_eq!(results.len(), 1, "welcome should match intro.md only");
assert_eq!(results[0]["path"], "intro.md");
}
#[tokio::test]
async fn search_with_empty_query_returns_empty() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let router = server.router();
let request = axum::http::Request::builder()
.uri("/search?q=")
.body(axum::body::Body::empty())
.expect("build request");
let response = tower::ServiceExt::oneshot(router, request).await.expect("oneshot");
assert_eq!(response.status(), StatusCode::OK);
let body_bytes =
axum::body::to_bytes(response.into_body(), usize::MAX).await.expect("body");
let body: serde_json::Value = serde_json::from_slice(&body_bytes).expect("json");
assert_eq!(body["results"].as_array().expect("array").len(), 0);
}
#[tokio::test]
async fn server_binds_and_shuts_down_gracefully() {
let (index, pages) = sample_fixture();
let server = DocsServer::new(index, pages).expect("build");
let cancel = CancellationToken::new();
let (bound_tx, bound_rx) = tokio::sync::oneshot::channel();
let task = tokio::spawn({
let cancel = cancel.clone();
async move { server.run("127.0.0.1:0".parse().unwrap(), bound_tx, cancel).await }
});
let addr = bound_rx.await.expect("bound addr");
let resp = reqwest::get(format!("http://{addr}/")).await.expect("request");
assert_eq!(resp.status(), 200);
cancel.cancel();
task.await.expect("task join").expect("server exit");
}
}