vendor/mdbook/src/cmd/serve.rs

   1 use super::command_prelude::*;
   2 #[cfg(feature = "watch")]
   3 use super::watch;
   4 use crate::{get_book_dir, open};
   5 use clap::builder::NonEmptyStringValueParser;
   6 use futures_util::sink::SinkExt;
   7 use futures_util::StreamExt;
   8 use mdbook::errors::*;
   9 use mdbook::utils;
  10 use mdbook::utils::fs::get_404_output_file;
  11 use mdbook::MDBook;
  12 use std::net::{SocketAddr, ToSocketAddrs};
  13 use std::path::PathBuf;
  14 use tokio::sync::broadcast;
  15 use warp::ws::Message;
  16 use warp::Filter;
  17
  18 /// The HTTP endpoint for the websocket used to trigger reloads when a file changes.
  19 const LIVE_RELOAD_ENDPOINT: &str = "__livereload";
  20
  21 // Create clap subcommand arguments
  22 pub fn make_subcommand() -> Command {
  23     Command::new("serve")
  24         .about("Serves a book at http://localhost:3000, and rebuilds it on changes")
  25         .arg_dest_dir()
  26         .arg_root_dir()
  27         .arg(
  28             Arg::new("hostname")
  29                 .short('n')
  30                 .long("hostname")
  31                 .num_args(1)
  32                 .default_value("localhost")
  33                 .value_parser(NonEmptyStringValueParser::new())
  34                 .help("Hostname to listen on for HTTP connections"),
  35         )
  36         .arg(
  37             Arg::new("port")
  38                 .short('p')
  39                 .long("port")
  40                 .num_args(1)
  41                 .default_value("3000")
  42                 .value_parser(NonEmptyStringValueParser::new())
  43                 .help("Port to use for HTTP connections"),
  44         )
  45         .arg_open()
  46 }
  47
  48 // Serve command implementation
  49 pub fn execute(args: &ArgMatches) -> Result<()> {
  50     let book_dir = get_book_dir(args);
  51     let mut book = MDBook::load(book_dir)?;
  52
  53     let port = args.get_one::<String>("port").unwrap();
  54     let hostname = args.get_one::<String>("hostname").unwrap();
  55     let open_browser = args.get_flag("open");
  56
  57     let address = format!("{}:{}", hostname, port);
  58
  59     let update_config = |book: &mut MDBook| {
  60         book.config
  61             .set("output.html.live-reload-endpoint", LIVE_RELOAD_ENDPOINT)
  62             .expect("live-reload-endpoint update failed");
  63         if let Some(dest_dir) = args.get_one::<PathBuf>("dest-dir") {
  64             book.config.build.build_dir = dest_dir.into();
  65         }
  66         // Override site-url for local serving of the 404 file
  67         book.config.set("output.html.site-url", "/").unwrap();
  68     };
  69     update_config(&mut book);
  70     book.build()?;
  71
  72     let sockaddr: SocketAddr = address
  73         .to_socket_addrs()?
  74         .next()
  75         .ok_or_else(|| anyhow::anyhow!("no address found for {}", address))?;
  76     let build_dir = book.build_dir_for("html");
  77     let input_404 = book
  78         .config
  79         .get("output.html.input-404")
  80         .and_then(toml::Value::as_str)
  81         .map(ToString::to_string);
  82     let file_404 = get_404_output_file(&input_404);
  83
  84     // A channel used to broadcast to any websockets to reload when a file changes.
  85     let (tx, _rx) = tokio::sync::broadcast::channel::<Message>(100);
  86
  87     let reload_tx = tx.clone();
  88     let thread_handle = std::thread::spawn(move || {
  89         serve(build_dir, sockaddr, reload_tx, &file_404);
  90     });
  91
  92     let serving_url = format!("http://{}", address);
  93     info!("Serving on: {}", serving_url);
  94
  95     if open_browser {
  96         open(serving_url);
  97     }
  98
  99     #[cfg(feature = "watch")]
 100     watch::trigger_on_change(&book, move |paths, book_dir| {
 101         info!("Files changed: {:?}", paths);
 102         info!("Building book...");
 103
 104         // FIXME: This area is really ugly because we need to re-set livereload :(
 105         let result = MDBook::load(book_dir).and_then(|mut b| {
 106             update_config(&mut b);
 107             b.build()
 108         });
 109
 110         if let Err(e) = result {
 111             error!("Unable to load the book");
 112             utils::log_backtrace(&e);
 113         } else {
 114             let _ = tx.send(Message::text("reload"));
 115         }
 116     });
 117
 118     let _ = thread_handle.join();
 119
 120     Ok(())
 121 }
 122
 123 #[tokio::main]
 124 async fn serve(
 125     build_dir: PathBuf,
 126     address: SocketAddr,
 127     reload_tx: broadcast::Sender<Message>,
 128     file_404: &str,
 129 ) {
 130     // A warp Filter which captures `reload_tx` and provides an `rx` copy to
 131     // receive reload messages.
 132     let sender = warp::any().map(move || reload_tx.subscribe());
 133
 134     // A warp Filter to handle the livereload endpoint. This upgrades to a
 135     // websocket, and then waits for any filesystem change notifications, and
 136     // relays them over the websocket.
 137     let livereload = warp::path(LIVE_RELOAD_ENDPOINT)
 138         .and(warp::ws())
 139         .and(sender)
 140         .map(|ws: warp::ws::Ws, mut rx: broadcast::Receiver<Message>| {
 141             ws.on_upgrade(move |ws| async move {
 142                 let (mut user_ws_tx, _user_ws_rx) = ws.split();
 143                 trace!("websocket got connection");
 144                 if let Ok(m) = rx.recv().await {
 145                     trace!("notify of reload");
 146                     let _ = user_ws_tx.send(m).await;
 147                 }
 148             })
 149         });
 150     // A warp Filter that serves from the filesystem.
 151     let book_route = warp::fs::dir(build_dir.clone());
 152     // The fallback route for 404 errors
 153     let fallback_route = warp::fs::file(build_dir.join(file_404))
 154         .map(|reply| warp::reply::with_status(reply, warp::http::StatusCode::NOT_FOUND));
 155     let routes = livereload.or(book_route).or(fallback_route);
 156
 157     std::panic::set_hook(Box::new(move |panic_info| {
 158         // exit if serve panics
 159         error!("Unable to serve: {}", panic_info);
 160         std::process::exit(1);
 161     }));
 162
 163     warp::serve(routes).run(address).await;
 164 }