forge/cmd/doc/
server.rs

1use axum::{Router, routing::get_service};
2use forge_doc::mdbook_driver::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 file_404 = book
71            .config
72            .html_config()
73            .map(|c| c.get_404_output_file())
74            .unwrap_or_else(|| "404.html".to_string());
75
76        let serving_url = format!("http://{address}");
77        sh_println!("Serving on: {serving_url}")?;
78
79        let thread_handle = std::thread::spawn(move || serve(build_dir, sockaddr, &file_404));
80
81        if self.open {
82            open(serving_url);
83        }
84
85        match thread_handle.join() {
86            Ok(r) => r.map_err(Into::into),
87            Err(e) => std::panic::resume_unwind(e),
88        }
89    }
90}
91
92#[tokio::main]
93async fn serve(build_dir: PathBuf, address: SocketAddr, file_404: &str) -> io::Result<()> {
94    let file_404 = build_dir.join(file_404);
95    let svc = ServeDir::new(build_dir).not_found_service(ServeFile::new(file_404));
96    let app = Router::new().fallback_service(get_service(svc));
97    let tcp_listener = tokio::net::TcpListener::bind(address).await?;
98    axum::serve(tcp_listener, app.into_make_service()).await
99}
100
101fn open<P: AsRef<std::ffi::OsStr>>(path: P) {
102    info!("Opening web browser");
103    if let Err(e) = opener::open(path) {
104        error!("Error opening web browser: {}", e);
105    }
106}
forge/cmd/doc/server.rs

forge/cmd/doc/
server.rs
