vendor/mdbook/src/cmd/serve.rs

   1 #[cfg(feature = "watch")]
   2 use super::watch;
   3 use crate::{get_book_dir, open};
   4 use clap::{App, Arg, ArgMatches, SubCommand};
   5 use iron::headers;
   6 use iron::{status, AfterMiddleware, Chain, Iron, IronError, IronResult, Request, Response, Set};
   7 use mdbook::errors::*;
   8 use mdbook::utils;
   9 use mdbook::MDBook;
  10
  11 struct ErrorRecover;
  12
  13 struct NoCache;
  14
  15 // Create clap subcommand arguments
  16 pub fn make_subcommand<'a, 'b>() -> App<'a, 'b> {
  17     SubCommand::with_name("serve")
  18         .about("Serves a book at http://localhost:3000, and rebuilds it on changes")
  19         .arg_from_usage(
  20             "-d, --dest-dir=[dest-dir] 'Output directory for the book{n}\
  21              Relative paths are interpreted relative to the book's root directory.{n}\
  22              If omitted, mdBook uses build.build-dir from book.toml or defaults to `./book`.'",
  23         )
  24         .arg_from_usage(
  25             "[dir] 'Root directory for the book{n}\
  26              (Defaults to the Current Directory when omitted)'",
  27         )
  28         .arg(
  29             Arg::with_name("hostname")
  30                 .short("n")
  31                 .long("hostname")
  32                 .takes_value(true)
  33                 .default_value("localhost")
  34                 .empty_values(false)
  35                 .help("Hostname to listen on for HTTP connections"),
  36         )
  37         .arg(
  38             Arg::with_name("port")
  39                 .short("p")
  40                 .long("port")
  41                 .takes_value(true)
  42                 .default_value("3000")
  43                 .empty_values(false)
  44                 .help("Port to use for HTTP connections"),
  45         )
  46         .arg(
  47             Arg::with_name("websocket-hostname")
  48                 .long("websocket-hostname")
  49                 .takes_value(true)
  50                 .empty_values(false)
  51                 .help(
  52                     "Hostname to connect to for WebSockets connections (Defaults to the HTTP hostname)",
  53                 ),
  54         )
  55         .arg(
  56             Arg::with_name("websocket-port")
  57                 .short("w")
  58                 .long("websocket-port")
  59                 .takes_value(true)
  60                 .default_value("3001")
  61                 .empty_values(false)
  62                 .help("Port to use for WebSockets livereload connections"),
  63         )
  64         .arg_from_usage("-o, --open 'Opens the book server in a web browser'")
  65 }
  66
  67 // Watch command implementation
  68 pub fn execute(args: &ArgMatches) -> Result<()> {
  69     let book_dir = get_book_dir(args);
  70     let mut book = MDBook::load(&book_dir)?;
  71
  72     let port = args.value_of("port").unwrap();
  73     let ws_port = args.value_of("websocket-port").unwrap();
  74     let hostname = args.value_of("hostname").unwrap();
  75     let public_address = args.value_of("websocket-hostname").unwrap_or(hostname);
  76     let open_browser = args.is_present("open");
  77
  78     let address = format!("{}:{}", hostname, port);
  79     let ws_address = format!("{}:{}", hostname, ws_port);
  80
  81     let livereload_url = format!("ws://{}:{}", public_address, ws_port);
  82     book.config
  83         .set("output.html.livereload-url", &livereload_url)?;
  84
  85     if let Some(dest_dir) = args.value_of("dest-dir") {
  86         book.config.build.build_dir = dest_dir.into();
  87     }
  88
  89     book.build()?;
  90
  91     let mut chain = Chain::new(staticfile::Static::new(book.build_dir_for("html")));
  92     chain.link_after(NoCache);
  93     chain.link_after(ErrorRecover);
  94     let _iron = Iron::new(chain)
  95         .http(&*address)
  96         .chain_err(|| "Unable to launch the server")?;
  97
  98     let ws_server =
  99         ws::WebSocket::new(|_| |_| Ok(())).chain_err(|| "Unable to start the websocket")?;
 100
 101     let broadcaster = ws_server.broadcaster();
 102
 103     std::thread::spawn(move || {
 104         ws_server.listen(&*ws_address).unwrap();
 105     });
 106
 107     let serving_url = format!("http://{}", address);
 108     info!("Serving on: {}", serving_url);
 109
 110     if open_browser {
 111         open(serving_url);
 112     }
 113
 114     #[cfg(feature = "watch")]
 115     watch::trigger_on_change(&book, move |paths, book_dir| {
 116         info!("Files changed: {:?}", paths);
 117         info!("Building book...");
 118
 119         // FIXME: This area is really ugly because we need to re-set livereload :(
 120
 121         let result = MDBook::load(&book_dir)
 122             .and_then(|mut b| {
 123                 b.config
 124                     .set("output.html.livereload-url", &livereload_url)?;
 125                 Ok(b)
 126             })
 127             .and_then(|b| b.build());
 128
 129         if let Err(e) = result {
 130             error!("Unable to load the book");
 131             utils::log_backtrace(&e);
 132         } else {
 133             let _ = broadcaster.send("reload");
 134         }
 135     });
 136
 137     Ok(())
 138 }
 139
 140 impl AfterMiddleware for NoCache {
 141     fn after(&self, _: &mut Request, mut res: Response) -> IronResult<Response> {
 142         res.headers.set(headers::CacheControl(vec![
 143             headers::CacheDirective::NoStore,
 144             headers::CacheDirective::MaxAge(0u32),
 145         ]));
 146
 147         Ok(res)
 148     }
 149 }
 150
 151 impl AfterMiddleware for ErrorRecover {
 152     fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
 153         match err.response.status {
 154             // each error will result in 404 response
 155             Some(_) => Ok(err.response.set(status::NotFound)),
 156             _ => Err(err),
 157         }
 158     }
 159 }