Implement Serve feature

Cargo.toml

@@ -25,6 +25,11 @@ notify = { version = "2.5.4", optional = true }
 time = { version = "0.1.34", optional = true }
 crossbeam = { version = "0.2.8", optional = true }
 
+# Serve feature
+iron = { version = "0.3", optional = true }
+staticfile = { version = "0.2", optional = true }
+websocket = { version = "0.16.1", optional = true}
+
 
 # Tests
 [dev-dependencies]
@@ -32,11 +37,12 @@ tempdir = "0.3.4"
 
 
 [features]
-default = ["output", "watch"]
+default = ["output", "watch", "serve"]
 debug = []
 output = []
 regenerate-css = []
 watch = ["notify", "time", "crossbeam"]
+serve = ["iron", "staticfile", "websocket"]
 
 [[bin]]
 doc = false

src/bin/livereload.rs

@@ -0,0 +1,210 @@
+extern crate websocket;
+extern crate crossbeam;
+
+use std::sync::mpsc::channel;
+use std::sync::mpsc;
+use std::io;
+use std::thread;
+use std::sync::{Arc, Mutex};
+use std::ops::Deref;
+use std::marker::PhantomData;
+
+use self::websocket::header::WebSocketProtocol;
+use self::websocket::ws::sender::Sender;
+use self::websocket::ws::receiver::Receiver;
+use self::websocket::message::Type;
+use self::websocket::{Server, Message};
+
+const WS_PROTOCOL: &'static str = "livereload";
+const RELOAD_COMMAND: &'static str = "reload";
+
+
+#[derive(Debug, Clone, PartialEq)]
+enum MessageType {
+    Reload,
+    Close,
+}
+
+
+#[derive(Clone)]
+struct ComparableSender<T> {
+    sender: mpsc::Sender<T>,
+    id: usize,
+}
+
+impl<T> PartialEq for ComparableSender<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id
+    }
+}
+
+impl<T> Deref for ComparableSender<T> {
+    type Target = mpsc::Sender<T>;
+
+    fn deref(&self) -> &mpsc::Sender<T> {
+        &self.sender
+    }
+}
+
+
+struct ComparableSenderFactory<T> {
+    next_id: usize,
+    sender_type: PhantomData<T>,
+}
+
+impl<T> ComparableSenderFactory<T> {
+    fn generate(&mut self, sender: mpsc::Sender<T>) -> ComparableSender<T> {
+        let tx = ComparableSender {
+            sender: sender,
+            id: self.next_id,
+        };
+        self.next_id += 1;
+        tx
+    }
+
+    fn new() -> ComparableSenderFactory<T> {
+        ComparableSenderFactory {
+            next_id: 0,
+            sender_type: PhantomData,
+        }
+    }
+}
+
+
+pub struct LiveReload {
+    senders: Arc<Mutex<Vec<ComparableSender<MessageType>>>>,
+}
+
+impl LiveReload {
+    pub fn new(address: &str) -> io::Result<LiveReload> {
+        let server = try!(Server::bind(address));
+
+        let senders: Arc<Mutex<Vec<ComparableSender<MessageType>>>> = Arc::new(Mutex::new(vec![]));
+        let senders_clone = senders.clone();
+
+        let mut factory = ComparableSenderFactory::new();
+
+        let lr = LiveReload { senders: senders_clone };
+
+        // handle connection attempts on a separate thread
+        thread::spawn(move || {
+            for connection in server {
+                let mut senders = senders.clone();
+                let (tx, rx) = channel();
+
+                let tx = factory.generate(tx);
+
+                senders.lock().unwrap().push(tx.clone());
+
+                // each connection gets a separate thread
+                thread::spawn(move || {
+                    let request = connection.unwrap().read_request().unwrap();
+                    let headers = request.headers.clone();
+
+                    let mut valid = false;
+                    if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
+                        if protocols.contains(&(WS_PROTOCOL.to_owned())) {
+                            valid = true;
+                        }
+                    }
+
+                    let client;
+                    if valid {
+                        let mut response = request.accept();
+                        response.headers.set(WebSocketProtocol(vec![WS_PROTOCOL.to_owned()]));
+                        client = response.send().unwrap();
+                    } else {
+                        request.fail().send().unwrap();
+                        println!("{:?}", "Rejecting invalid websocket request.");
+                        return;
+                    }
+
+                    let (mut ws_tx, mut ws_rx) = client.split();
+
+                    // handle receiving and sending (websocket) in two separate threads
+                    crossbeam::scope(|scope| {
+                        let tx_clone = tx.clone();
+                        scope.spawn(move || {
+                            let tx = tx_clone;
+                            loop {
+                                match rx.recv() {
+                                    Ok(msg) => {
+                                        match msg {
+                                            MessageType::Reload => {
+                                                let message: Message = Message::text(RELOAD_COMMAND.to_owned());
+                                                let mut senders = senders.clone();
+                                                if ws_tx.send_message(&message).is_err() {
+                                                    // the receiver isn't available anymore
+                                                    // remove the tx from senders and exit
+                                                    LiveReload::remove_sender(&mut senders, &tx);
+                                                    break;
+                                                }
+                                            },
+                                            MessageType::Close => {
+                                                LiveReload::remove_sender(&mut senders, &tx);
+                                                break;
+                                            },
+                                        }
+                                    },
+                                    Err(e) => {
+                                        println!("{:?}", e);
+                                        break;
+                                    },
+                                }
+                            }
+                        });
+
+                        for message in ws_rx.incoming_messages() {
+                            match message {
+                                Ok(message) => {
+                                    let message: Message = message;
+                                    match message.opcode {
+                                        Type::Close => {
+                                            tx.send(MessageType::Close).unwrap();
+                                            break;
+                                        },
+                                        // TODO ?
+                                        // Type::Ping => {
+                                        //     let message = websocket::Message::pong(message.payload);
+                                        //     ws_tx.send_message(&message).unwrap();
+                                        // },
+                                        _ => {
+                                            println!("{:?}", message.opcode);
+                                            unimplemented!()
+                                        },
+                                    }
+                                },
+                                Err(err) => {
+                                    println!("Error: {}", err);
+                                    break;
+                                },
+                            }
+                        }
+                    });
+                });
+            }
+        });
+
+        Ok(lr)
+    }
+
+    fn remove_sender(senders: &mut Arc<Mutex<Vec<ComparableSender<MessageType>>>>, el: &ComparableSender<MessageType>) {
+        let mut senders = senders.lock().unwrap();
+        let mut index = 0;
+        for i in 0..senders.len() {
+            if &senders[i] == el {
+                index = i;
+                break;
+            }
+        }
+        senders.remove(index);
+    }
+
+    pub fn trigger_reload(&self) {
+        let senders = self.senders.lock().unwrap();
+        println!("Reloading {} client(s).", senders.len());
+        for sender in senders.iter() {
+            sender.send(MessageType::Reload).unwrap();
+        }
+    }
+}

src/bin/mdbook.rs

@@ -10,6 +10,14 @@ extern crate notify;
 #[cfg(feature = "watch")]
 extern crate time;
 
+// Dependencies for the Serve feature
+#[cfg(feature = "serve")]
+extern crate iron;
+#[cfg(feature = "serve")]
+extern crate staticfile;
+
+#[cfg(feature = "serve")]
+mod livereload;
 
 use std::env;
 use std::error::Error;
@@ -24,6 +32,10 @@ use notify::Watcher;
 #[cfg(feature = "watch")]
 use std::sync::mpsc::channel;
 
+// Uses for the Serve feature
+#[cfg(feature = "serve")]
+use livereload::LiveReload;
+
 
 use mdbook::MDBook;
 
@@ -50,6 +62,11 @@ fn main() {
                     .subcommand(SubCommand::with_name("watch")
                         .about("Watch the files for changes")
                         .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
+                    .subcommand(SubCommand::with_name("serve")
+                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
+                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
+                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
+                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'"))
                     .subcommand(SubCommand::with_name("test")
                         .about("Test that code samples compile"))
                     .get_matches();
@@ -60,6 +77,8 @@ fn main() {
         ("build", Some(sub_matches)) => build(sub_matches),
         #[cfg(feature = "watch")]
         ("watch", Some(sub_matches)) => watch(sub_matches),
+        #[cfg(feature = "serve")]
+        ("serve", Some(sub_matches)) => serve(sub_matches),
         ("test", Some(sub_matches)) => test(sub_matches),
         (_, _) => unreachable!(),
     };
@@ -148,76 +167,70 @@ fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
 #[cfg(feature = "watch")]
 fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
     let book_dir = get_book_dir(args);
-    let book = MDBook::new(&book_dir).read_config();
+    let mut book = MDBook::new(&book_dir).read_config();
 
-    // Create a channel to receive the events.
+    trigger_on_change(&mut book, |event, book| {
-    let (tx, rx) = channel();
+        if let Some(path) = event.path {
-
+            println!("File changed: {:?}\nBuilding book...\n", path);
-    let w: Result<notify::RecommendedWatcher, notify::Error> = notify::Watcher::new(tx);
+            match book.build() {
-
+                Err(e) => println!("Error while building: {:?}", e),
-    match w {
+                _ => {},
-        Ok(mut watcher) => {
-
-            // Add the source directory to the watcher
-            if let Err(e) = watcher.watch(book.get_src()) {
-                println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
-                ::std::process::exit(0);
-            };
-
-            // Add the book.json file to the watcher if it exists, because it's not
-            // located in the source directory
-            if let Err(_) = watcher.watch(book_dir.join("book.json")) {
-                // do nothing if book.json is not found
             }
-
+            println!("");
-            let mut previous_time = time::get_time();
+        }
-
+    });
-            crossbeam::scope(|scope| {
-                loop {
-                    match rx.recv() {
-                        Ok(event) => {
-
-                            // Skip the event if an event has already been issued in the last second
-                            let time = time::get_time();
-                            if time - previous_time < time::Duration::seconds(1) {
-                                continue;
-                            } else {
-                                previous_time = time;
-                            }
-
-                            if let Some(path) = event.path {
-                                // Trigger the build process in a new thread (to keep receiving events)
-                                scope.spawn(move || {
-                                    println!("File changed: {:?}\nBuilding book...\n", path);
-                                    match build(args) {
-                                        Err(e) => println!("Error while building: {:?}", e),
-                                        _ => {},
-                                    }
-                                    println!("");
-                                });
-
-                            } else {
-                                continue;
-                            }
-                        },
-                        Err(e) => {
-                            println!("An error occured: {:?}", e);
-                        },
-                    }
-                }
-            });
-
-        },
-        Err(e) => {
-            println!("Error while trying to watch the files:\n\n\t{:?}", e);
-            ::std::process::exit(0);
-        },
-    }
 
     Ok(())
 }
 
 
+// Watch command implementation
+#[cfg(feature = "serve")]
+fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
+    let book_dir = get_book_dir(args);
+    let mut book = MDBook::new(&book_dir).read_config();
+    let port = args.value_of("port").unwrap_or("3000");
+    let ws_port = args.value_of("ws-port").unwrap_or("3001");
+
+    book.set_livereload(format!(r#"
+        <script type="text/javascript">
+            var socket = new WebSocket("ws://localhost:{}", "livereload");
+            socket.onmessage = function (event) {{
+                if (event.data === "reload") {{
+                    socket.close();
+                    location.reload(true); // force reload from server (not from cache)
+                }}
+            }};
+
+            window.onbeforeunload = function() {{
+                socket.close();
+            }}
+        </script>
+    "#, ws_port).to_owned());
+
+    try!(book.build());
+
+    let staticfile = staticfile::Static::new(book.get_dest());
+    let iron = iron::Iron::new(staticfile);
+    let _iron = iron.http(&*format!("localhost:{}", port)).unwrap();
+
+    let lr = LiveReload::new(&format!("localhost:{}", ws_port)).unwrap();
+
+    println!("{:?}", "Registering change trigger");
+    trigger_on_change(&mut book, move |event, book| {
+        if let Some(path) = event.path {
+            println!("File changed: {:?}\nBuilding book...\n", path);
+            match book.build() {
+                Err(e) => println!("Error while building: {:?}", e),
+                _ => lr.trigger_reload(),
+            }
+            println!("");
+        }
+    });
+
+    Ok(())
+}
+
 
 fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
     let book_dir = get_book_dir(args);
@@ -229,7 +242,6 @@ fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
 }
 
 
-
 fn get_book_dir(args: &ArgMatches) -> PathBuf {
     if let Some(dir) = args.value_of("dir") {
         // Check if path is relative from current dir, or absolute...
@@ -243,3 +255,55 @@ fn get_book_dir(args: &ArgMatches) -> PathBuf {
         env::current_dir().unwrap()
     }
 }
+
+
+// Calls the closure when a book source file is changed. This is blocking!
+fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
+    where F: Fn(notify::Event, &mut MDBook) -> ()
+{
+    // Create a channel to receive the events.
+    let (tx, rx) = channel();
+
+    let w: Result<notify::RecommendedWatcher, notify::Error> = notify::Watcher::new(tx);
+
+    match w {
+        Ok(mut watcher) => {
+            // Add the source directory to the watcher
+            if let Err(e) = watcher.watch(book.get_src()) {
+                println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
+                ::std::process::exit(0);
+            };
+
+            // Add the book.json file to the watcher if it exists, because it's not
+            // located in the source directory
+            if let Err(_) = watcher.watch(book.get_root().join("book.json")) {
+                // do nothing if book.json is not found
+            }
+
+            let mut previous_time = time::get_time();
+
+            loop {
+                match rx.recv() {
+                    Ok(event) => {
+                        // Skip the event if an event has already been issued in the last second
+                        let time = time::get_time();
+                        if time - previous_time < time::Duration::seconds(1) {
+                            continue;
+                        } else {
+                            previous_time = time;
+                        }
+
+                        closure(event, book);
+                    },
+                    Err(e) => {
+                        println!("An error occured: {:?}", e);
+                    },
+                }
+            }
+        },
+        Err(e) => {
+            println!("Error while trying to watch the files:\n\n\t{:?}", e);
+            ::std::process::exit(0);
+        },
+    }
+}

src/book/mdbook.rs

@@ -15,6 +15,8 @@ pub struct MDBook {
     config: BookConfig,
     pub content: Vec<BookItem>,
     renderer: Box<Renderer>,
+    #[cfg(feature = "serve")]
+    livereload: Option<String>,
 }
 
 impl MDBook {
@@ -38,6 +40,7 @@ impl MDBook {
                         .set_dest(&root.join("book"))
                         .to_owned(),
             renderer: Box::new(HtmlHandlebars::new()),
+            livereload: None,
         }
     }
 
@@ -398,6 +401,23 @@ impl MDBook {
         &self.config.description
     }
 
+    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
+        self.livereload = Some(livereload);
+        self
+    }
+
+    pub fn unset_livereload(&mut self) -> &Self {
+        self.livereload = None;
+        self
+    }
+
+    pub fn get_livereload(&self) -> Option<&String> {
+        match self.livereload {
+            Some(ref livereload) => Some(&livereload),
+            None => None,
+        }
+    }
+
     // Construct book
     fn parse_summary(&mut self) -> Result<(), Box<Error>> {
         // When append becomes stable, use self.content.append() ...

src/renderer/html_handlebars/hbs_renderer.rs

@@ -287,6 +287,9 @@ fn make_data(book: &MDBook) -> Result<BTreeMap<String, Json>, Box<Error>> {
     data.insert("title".to_owned(), book.get_title().to_json());
     data.insert("description".to_owned(), book.get_description().to_json());
     data.insert("favicon".to_owned(), "favicon.png".to_json());
+    if let Some(livereload) = book.get_livereload() {
+        data.insert("livereload".to_owned(), livereload.to_json());
+    }
 
     let mut chapters = vec![];
 

src/theme/index.hbs

@@ -107,6 +107,9 @@
             }
         </script>
 
+        <!-- Livereload script (if served using the cli tool) -->
+        {{{livereload}}}
+
         <script src="highlight.js"></script>
         <script src="book.js"></script>
     </body>

src/utils/fs.rs

@@ -150,7 +150,7 @@ pub fn copy_files_except_ext(from: &Path, to: &Path, recursive: bool, ext_blackl
                 debug!("[*] creating path for file: {:?}",
                        &to.join(entry.path().file_name().expect("a file should have a file name...")));
 
-                output!("[*] copying file: {:?}\n    to {:?}",
+                output!("[*] Copying file: {:?}\n    to {:?}",
                         entry.path(),
                         &to.join(entry.path().file_name().expect("a file should have a file name...")));
                 try!(fs::copy(entry.path(),