From cda28bb618275cfc5e442f4c6b97afd7a1c00404 Mon Sep 17 00:00:00 2001
From: Manuel Woelker <github@manuel.woelker.org>
Date: Wed, 13 May 2020 14:45:35 +0200
Subject: [PATCH] Generate 404.html page (#539)

---
 book-example/src/404.md                      |  3 ++
 src/cmd/serve.rs                             |  8 +++-
 src/config.rs                                | 33 ++++++++++++++++
 src/renderer/html_handlebars/hbs_renderer.rs | 41 ++++++++++++++++++++
 4 files changed, 83 insertions(+), 2 deletions(-)
 create mode 100644 book-example/src/404.md

diff --git a/book-example/src/404.md b/book-example/src/404.md
new file mode 100644
index 00000000..a55db44e
--- /dev/null
+++ b/book-example/src/404.md
@@ -0,0 +1,3 @@
+# Document not found (404)
+
+This URL is invalid, sorry. Try the search instead!
\ No newline at end of file
diff --git a/src/cmd/serve.rs b/src/cmd/serve.rs
index 4e865f61..843b4e9f 100644
--- a/src/cmd/serve.rs
+++ b/src/cmd/serve.rs
@@ -142,7 +142,11 @@ async fn serve(build_dir: PathBuf, address: SocketAddr, reload_tx: broadcast::Se
             })
         });
     // A warp Filter that serves from the filesystem.
-    let book_route = warp::fs::dir(build_dir);
-    let routes = livereload.or(book_route);
+    let book_route = warp::fs::dir(build_dir.clone());
+    // The fallback route for 404 errors
+    let fallback_file = "404.html";
+    let fallback_route = warp::fs::file(build_dir.join(fallback_file))
+        .map(|reply| warp::reply::with_status(reply, warp::http::StatusCode::NOT_FOUND));
+    let routes = livereload.or(book_route).or(fallback_route);
     warp::serve(routes).run(address).await;
 }
diff --git a/src/config.rs b/src/config.rs
index a426199e..8a8f6ce5 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -507,6 +507,10 @@ pub struct HtmlConfig {
     /// FontAwesome icon class to use for the Git repository link.
     /// Defaults to `fa-github` if `None`.
     pub git_repository_icon: Option<String>,
+    /// Input path for the 404 file, defaults to 404.md
+    pub input_404: Option<String>,
+    /// Output path for 404.html file, defaults to 404.html, set to "" to disable 404 file output
+    pub output_404: Option<String>,
     /// This is used as a bit of a workaround for the `mdbook serve` command.
     /// Basically, because you set the websocket port from the command line, the
     /// `mdbook serve` command needs a way to let the HTML renderer know where
@@ -538,6 +542,8 @@ impl Default for HtmlConfig {
             search: None,
             git_repository_url: None,
             git_repository_icon: None,
+            input_404: None,
+            output_404: None,
             livereload_url: None,
             redirect: HashMap::new(),
         }
@@ -1004,4 +1010,31 @@ mod tests {
 
         assert_eq!(cfg.book.title, Some(should_be));
     }
+
+    #[test]
+    fn file_404_default() {
+        let src = r#"
+        [output.html]
+        destination = "my-book"
+        "#;
+
+        let got = Config::from_str(src).unwrap();
+        let html_config = got.html_config().unwrap();
+        assert_eq!(html_config.input_404, None);
+        assert_eq!(html_config.output_404, None);
+    }
+
+    #[test]
+    fn file_404_custom() {
+        let src = r#"
+        [output.html]
+        input-404= "missing.md"
+        output-404= "missing.html"
+        "#;
+
+        let got = Config::from_str(src).unwrap();
+        let html_config = got.html_config().unwrap();
+        assert_eq!(html_config.input_404, Some("missing.md".to_string()));
+        assert_eq!(html_config.output_404, Some("missing.html".to_string()));
+    }
 }
diff --git a/src/renderer/html_handlebars/hbs_renderer.rs b/src/renderer/html_handlebars/hbs_renderer.rs
index 4d3e95f2..b99b83a0 100644
--- a/src/renderer/html_handlebars/hbs_renderer.rs
+++ b/src/renderer/html_handlebars/hbs_renderer.rs
@@ -437,6 +437,47 @@ impl Renderer for HtmlHandlebars {
             is_index = false;
         }
 
+        // Render 404 page
+        if html_config.output_404 != Some("".to_string()) {
+            let default_404_location = src_dir.join("404.md");
+            let content_404 = if let Some(ref filename) = html_config.input_404 {
+                let path = src_dir.join(filename);
+                std::fs::read_to_string(&path).map_err(|failure| {
+                    std::io::Error::new(
+                        failure.kind(),
+                        format!("Unable to open 404 input file {:?}", &path),
+                    )
+                })?
+            } else if default_404_location.exists() {
+                std::fs::read_to_string(&default_404_location).map_err(|failure| {
+                    std::io::Error::new(
+                        failure.kind(),
+                        format!(
+                            "Unable to open default 404 input file {:?}",
+                            &default_404_location
+                        ),
+                    )
+                })?
+            } else {
+                "# 404 - Document not found\n\nUnfortunately, this URL is no longer valid, please use the navigation bar or search to continue.".to_string()
+            };
+            let html_content_404 = utils::render_markdown(&content_404, html_config.curly_quotes);
+
+            let mut data_404 = data.clone();
+            data_404.insert("path".to_owned(), json!("404.md"));
+            data_404.insert("content".to_owned(), json!(html_content_404));
+            let rendered = handlebars.render("index", &data_404)?;
+
+            let rendered =
+                self.post_process(rendered, &html_config.playpen, ctx.config.rust.edition);
+            let output_file = match &html_config.output_404 {
+                None => "404.html",
+                Some(file) => &file,
+            };
+            utils::fs::write_file(&destination, output_file, rendered.as_bytes())?;
+            debug!("Creating 404.html ✓");
+        }
+
         // Print version
         self.configure_print_version(&mut data, &print_content);
         if let Some(ref title) = ctx.config.book.title {