forge/cmd/doc/
server.rs

1use axum::{routing::get_service, Router};
2use forge_doc::mdbook::{utils::fs::get_404_output_file, MDBook};
3use std::{
4    io,
5    net::{SocketAddr, ToSocketAddrs},
6    path::PathBuf,
7};
8use tower_http::services::{ServeDir, ServeFile};
9
10/// The HTTP endpoint for the websocket used to trigger reloads when a file changes.
11const LIVE_RELOAD_ENDPOINT: &str = "/__livereload";
12
13/// Basic mdbook server. Given a path, hostname and port, serves the mdbook.
14#[derive(Debug)]
15pub struct Server {
16    path: PathBuf,
17    hostname: String,
18    port: usize,
19    open: bool,
20}
21
22impl Default for Server {
23    fn default() -> Self {
24        Self { path: PathBuf::default(), hostname: "localhost".to_owned(), port: 3000, open: false }
25    }
26}
27
28impl Server {
29    /// Create a new instance.
30    pub fn new(path: PathBuf) -> Self {
31        Self { path, ..Default::default() }
32    }
33
34    /// Set the host to serve on.
35    pub fn with_hostname(mut self, hostname: String) -> Self {
36        self.hostname = hostname;
37        self
38    }
39
40    /// Set the port to serve on.
41    pub fn with_port(mut self, port: usize) -> Self {
42        self.port = port;
43        self
44    }
45
46    /// Set whether to open the browser after serving.
47    pub fn open(mut self, open: bool) -> Self {
48        self.open = open;
49        self
50    }
51
52    /// Serve the mdbook.
53    pub fn serve(self) -> eyre::Result<()> {
54        let mut book =
55            MDBook::load(&self.path).map_err(|err| eyre::eyre!("failed to load book: {err:?}"))?;
56
57        let reload = LIVE_RELOAD_ENDPOINT.strip_prefix('/').unwrap();
58        book.config.set("output.html.live-reload-endpoint", reload).unwrap();
59        // Override site-url for local serving of the 404 file
60        book.config.set("output.html.site-url", "/").unwrap();
61
62        book.build().map_err(|err| eyre::eyre!("failed to build book: {err:?}"))?;
63
64        let address = format!("{}:{}", self.hostname, self.port);
65        let sockaddr: SocketAddr = address
66            .to_socket_addrs()?
67            .next()
68            .ok_or_else(|| eyre::eyre!("no address found for {}", address))?;
69        let build_dir = book.build_dir_for("html");
70        let input_404 = book
71            .config
72            .get("output.html.input-404")
73            .and_then(|v| v.as_str())
74            .map(ToString::to_string);
75        let file_404 = get_404_output_file(&input_404);
76
77        let serving_url = format!("http://{address}");
78        sh_println!("Serving on: {serving_url}")?;
79
80        let thread_handle = std::thread::spawn(move || serve(build_dir, sockaddr, &file_404));
81
82        if self.open {
83            open(serving_url);
84        }
85
86        match thread_handle.join() {
87            Ok(r) => r.map_err(Into::into),
88            Err(e) => std::panic::resume_unwind(e),
89        }
90    }
91}
92
93#[tokio::main]
94async fn serve(build_dir: PathBuf, address: SocketAddr, file_404: &str) -> io::Result<()> {
95    let file_404 = build_dir.join(file_404);
96    let svc = ServeDir::new(build_dir).not_found_service(ServeFile::new(file_404));
97    let app = Router::new().nest_service("/", get_service(svc));
98    let tcp_listener = tokio::net::TcpListener::bind(address).await?;
99    axum::serve(tcp_listener, app.into_make_service()).await
100}
101
102fn open<P: AsRef<std::ffi::OsStr>>(path: P) {
103    info!("Opening web browser");
104    if let Err(e) = opener::open(path) {
105        error!("Error opening web browser: {}", e);
106    }
107}