Clean up error handling, logging and 404 correctly for static files

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
 }