1 files changed, 177 insertions, 0 deletions
diff --git a/vendor/mdbook/src/cmd/serve.rs b/vendor/mdbook/src/cmd/serve.rs
new file mode 100644
index 000000000..bafbfd52e
--- /dev/null
+++ b/vendor/mdbook/src/cmd/serve.rs
@@ -0,0 +1,177 @@
+#[cfg(feature = "watch")]
+use super::watch;
+use crate::{get_book_dir, open};
+use clap::{arg, App, Arg, ArgMatches};
+use futures_util::sink::SinkExt;
+use futures_util::StreamExt;
+use mdbook::errors::*;
+use mdbook::utils;
+use mdbook::utils::fs::get_404_output_file;
+use mdbook::MDBook;
+use std::net::{SocketAddr, ToSocketAddrs};
+use std::path::PathBuf;
+use tokio::sync::broadcast;
+use warp::ws::Message;
+use warp::Filter;
+
+/// The HTTP endpoint for the websocket used to trigger reloads when a file changes.
+const LIVE_RELOAD_ENDPOINT: &str = "__livereload";
+
+// Create clap subcommand arguments
+pub fn make_subcommand<'help>() -> App<'help> {
+    App::new("serve")
+        .about("Serves a book at http://localhost:3000, and rebuilds it on changes")
+        .arg(
+            Arg::new("dest-dir")
+                .short('d')
+                .long("dest-dir")
+                .value_name("dest-dir")
+                .help(
+                    "Output directory for the book{n}\
+                    Relative paths are interpreted relative to the book's root directory.{n}\
+                    If omitted, mdBook uses build.build-dir from book.toml or defaults to `./book`.",
+                ),
+        )
+        .arg(arg!([dir]
+            "Root directory for the book{n}\
+            (Defaults to the Current Directory when omitted)"
+        ))
+        .arg(
+            Arg::new("hostname")
+                .short('n')
+                .long("hostname")
+                .takes_value(true)
+                .default_value("localhost")
+                .forbid_empty_values(true)
+                .help("Hostname to listen on for HTTP connections"),
+        )
+        .arg(
+            Arg::new("port")
+                .short('p')
+                .long("port")
+                .takes_value(true)
+                .default_value("3000")
+                .forbid_empty_values(true)
+                .help("Port to use for HTTP connections"),
+        )
+        .arg(arg!(-o --open "Opens the compiled book in a web browser"))
+}
+
+// Serve command implementation
+pub fn execute(args: &ArgMatches) -> Result<()> {
+    let book_dir = get_book_dir(args);
+    let mut book = MDBook::load(&book_dir)?;
+
+    let port = args.value_of("port").unwrap();
+    let hostname = args.value_of("hostname").unwrap();
+    let open_browser = args.is_present("open");
+
+    let address = format!("{}:{}", hostname, port);
+
+    let update_config = |book: &mut MDBook| {
+        book.config
+            .set("output.html.live-reload-endpoint", &LIVE_RELOAD_ENDPOINT)
+            .expect("live-reload-endpoint update failed");
+        if let Some(dest_dir) = args.value_of("dest-dir") {
+            book.config.build.build_dir = dest_dir.into();
+        }
+        // Override site-url for local serving of the 404 file
+        book.config.set("output.html.site-url", "/").unwrap();
+    };
+    update_config(&mut book);
+    book.build()?;
+
+    let sockaddr: SocketAddr = address
+        .to_socket_addrs()?
+        .next()
+        .ok_or_else(|| anyhow::anyhow!("no address found for {}", address))?;
+    let build_dir = book.build_dir_for("html");
+    let input_404 = book
+        .config
+        .get("output.html.input-404")
+        .map(toml::Value::as_str)
+        .and_then(std::convert::identity) // flatten
+        .map(ToString::to_string);
+    let file_404 = get_404_output_file(&input_404);
+
+    // A channel used to broadcast to any websockets to reload when a file changes.
+    let (tx, _rx) = tokio::sync::broadcast::channel::<Message>(100);
+
+    let reload_tx = tx.clone();
+    let thread_handle = std::thread::spawn(move || {
+        serve(build_dir, sockaddr, reload_tx, &file_404);
+    });
+
+    let serving_url = format!("http://{}", address);
+    info!("Serving on: {}", serving_url);
+
+    if open_browser {
+        open(serving_url);
+    }
+
+    #[cfg(feature = "watch")]
+    watch::trigger_on_change(&book, move |paths, book_dir| {
+        info!("Files changed: {:?}", paths);
+        info!("Building book...");
+
+        // FIXME: This area is really ugly because we need to re-set livereload :(
+        let result = MDBook::load(&book_dir).and_then(|mut b| {
+            update_config(&mut b);
+            b.build()
+        });
+
+        if let Err(e) = result {
+            error!("Unable to load the book");
+            utils::log_backtrace(&e);
+        } else {
+            let _ = tx.send(Message::text("reload"));
+        }
+    });
+
+    let _ = thread_handle.join();
+
+    Ok(())
+}
+
+#[tokio::main]
+async fn serve(
+    build_dir: PathBuf,
+    address: SocketAddr,
+    reload_tx: broadcast::Sender<Message>,
+    file_404: &str,
+) {
+    // A warp Filter which captures `reload_tx` and provides an `rx` copy to
+    // receive reload messages.
+    let sender = warp::any().map(move || reload_tx.subscribe());
+
+    // A warp Filter to handle the livereload endpoint. This upgrades to a
+    // websocket, and then waits for any filesystem change notifications, and
+    // relays them over the websocket.
+    let livereload = warp::path(LIVE_RELOAD_ENDPOINT)
+        .and(warp::ws())
+        .and(sender)
+        .map(|ws: warp::ws::Ws, mut rx: broadcast::Receiver<Message>| {
+            ws.on_upgrade(move |ws| async move {
+                let (mut user_ws_tx, _user_ws_rx) = ws.split();
+                trace!("websocket got connection");
+                if let Ok(m) = rx.recv().await {
+                    trace!("notify of reload");
+                    let _ = user_ws_tx.send(m).await;
+                }
+            })
+        });
+    // A warp Filter that serves from the filesystem.
+    let book_route = warp::fs::dir(build_dir.clone());
+    // The fallback route for 404 errors
+    let fallback_route = warp::fs::file(build_dir.join(file_404))
+        .map(|reply| warp::reply::with_status(reply, warp::http::StatusCode::NOT_FOUND));
+    let routes = livereload.or(book_route).or(fallback_route);
+
+    std::panic::set_hook(Box::new(move |panic_info| {
+        // exit if serve panics
+        error!("Unable to serve: {}", panic_info);
+        std::process::exit(1);
+    }));
+
+    warp::serve(routes).run(address).await;
+}