From d75acc8d0540dc0e28d8b27c5c588677da49c7d9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tristan=20Dani=C3=ABl=20Maat?= <tm@tlater.net>
Date: Tue, 16 Aug 2022 23:16:39 +0100
Subject: [PATCH] Clean up error handling, logging and 404 correctly for static
 files

---
 Cargo.lock  |  21 +++++++++++
 Cargo.toml  |   2 ++
 src/main.rs | 102 +++++++++++++++++++++++++++++++++++++++++-----------
 3 files changed, 104 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 47f59d5..db16587 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -456,6 +456,19 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "env_logger"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b2cf0344971ee6c64c31be0d530793fba457d322dfec2810c453d0ef228f9c3"
+dependencies = [
+ "atty",
+ "humantime",
+ "log",
+ "regex",
+ "termcolor",
+]
+
 [[package]]
 name = "firestorm"
 version = "0.5.1"
@@ -623,6 +636,12 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c4a1e36c821dbe04574f602848a19f742f4fb3c98d40449f11bcad18d6b17421"
 
+[[package]]
+name = "humantime"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
+
 [[package]]
 name = "idna"
 version = "0.2.3"
@@ -1196,7 +1215,9 @@ dependencies = [
  "actix-files",
  "actix-web",
  "clap",
+ "env_logger",
  "handlebars",
+ "log",
 ]
 
 [[package]]
diff --git a/Cargo.toml b/Cargo.toml
index 58dda85..30665c6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,4 +7,6 @@ edition = "2021"
 actix-files = "0.6.2"
 actix-web = "4.1.0"
 clap = { version = "3.2.17", features = ["derive"] }
+env_logger = "0.9.0"
 handlebars = { version = "4.3.3", features = ["dir_source"] }
+log = "0.4.17"
diff --git a/src/main.rs b/src/main.rs
index 013bab2..1820455 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,46 +1,106 @@
-use std::io;
+use std::net::SocketAddr;
 
-use actix_files::Files;
-use actix_web::{web, App, HttpRequest, HttpResponse, HttpServer, Responder};
+use actix_files::NamedFile;
+use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder};
 use clap::Parser;
 use handlebars::Handlebars;
+use std::path::PathBuf;
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 struct Config {
-    template_directory: String,
-    address: String,
-    port: u16,
+    #[clap(long, value_parser)]
+    /// The directory from which to serve static content and
+    /// handlebars templates
+    template_directory: PathBuf,
+    #[clap(long, default_value = "127.0.0.1:8000", value_parser)]
+    /// The address on which to listen
+    address: SocketAddr,
 }
 
-async fn template(hb: web::Data<Handlebars<'_>>, req: HttpRequest) -> impl Responder {
-    let path: String = req
+struct SharedData<'a> {
+    handlebars: Handlebars<'a>,
+    config: Config,
+}
+
+#[get("/{filename:.*.html}")]
+async fn template(
+    shared: web::Data<SharedData<'_>>,
+    req: HttpRequest,
+) -> Result<impl Responder, Box<dyn std::error::Error>> {
+    let path = req
         .match_info()
         .query("filename")
-        .parse()
-        .expect("need a file name on the request");
+        .strip_suffix(".html")
+        .expect("only paths with this suffix should get here");
 
-    let body = hb.render(path.strip_suffix(".html").unwrap(), &()).unwrap();
+    if shared.handlebars.has_template(path) {
+        let body = shared.handlebars.render(path, &())?;
+        Ok(HttpResponse::Ok().body(body))
+    } else {
+        let body = shared
+            .handlebars
+            .render("404", &())
+            .expect("404 template not found");
+        Ok(HttpResponse::NotFound().body(body))
+    }
+}
 
-    HttpResponse::Ok().body(body)
+#[get("/{filename:.*}")]
+async fn static_file(
+    shared: web::Data<SharedData<'_>>,
+    req: HttpRequest,
+) -> Result<impl Responder, Box<dyn std::error::Error>> {
+    let requested = req.match_info().query("filename");
+
+    match shared
+        .config
+        .template_directory
+        .join(requested)
+        .canonicalize()
+    {
+        // We only want to serve paths that are both valid *and* in
+        // the template directory.
+        //
+        // i.e., don't serve up /etc/passwd
+        Ok(path) if path.starts_with(&shared.config.template_directory) => {
+            let file = NamedFile::open_async(path).await?;
+            Ok(file.use_last_modified(false).respond_to(&req))
+        }
+        // Any other cases should 404
+        _ => {
+            let body = shared
+                .handlebars
+                .render("404", &())
+                .expect("404 template not found");
+            Ok(HttpResponse::NotFound().body(body))
+        }
+    }
 }
 
 #[actix_web::main]
-async fn main() -> io::Result<()> {
-    let config = Config::parse();
+async fn main() -> Result<(), std::io::Error> {
+    let mut config = Config::parse();
+    config.template_directory = config.template_directory.canonicalize()?;
+
+    env_logger::init();
 
     let mut handlebars = Handlebars::new();
     handlebars
         .register_templates_directory(".html", config.template_directory.clone())
-        .expect("could not read template directory");
-    let handlebars_ref = web::Data::new(handlebars);
+        .expect("templates should compile correctly");
+
+    let shared_data = web::Data::new(SharedData {
+        handlebars,
+        config: config.clone(),
+    });
 
     HttpServer::new(move || {
         App::new()
-            .app_data(handlebars_ref.clone())
-            .route("/{filename:.*.html}", web::get().to(template))
-            .service(Files::new("/", config.template_directory.clone()))
+            .app_data(shared_data.clone())
+            .service(template)
+            .service(static_file)
     })
-    .bind((config.address, config.port))?
+    .bind(config.address)?
     .run()
     .await
 }