From da3d33f81d396beffc8df72a8186bf6e1ad8654f Mon Sep 17 00:00:00 2001
From: Naim A <227396+naim94a@users.noreply.github.com>
Date: Tue, 8 Oct 2019 00:47:10 +0300
Subject: [PATCH] [actix-files] added custom guard checks to FilesService

---
 actix-files/src/lib.rs | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/actix-files/src/lib.rs b/actix-files/src/lib.rs
index c99d3265f..ffb773ae0 100644
--- a/actix-files/src/lib.rs
+++ b/actix-files/src/lib.rs
@@ -15,8 +15,10 @@ use actix_web::dev::{
     AppService, HttpServiceFactory, Payload, ResourceDef, ServiceRequest,
     ServiceResponse,
 };
+use actix_web::guard::Guard;
 use actix_web::error::{BlockingError, Error, ErrorInternalServerError};
-use actix_web::http::header::DispositionType;
+use actix_web::http::header::{self, DispositionType};
+use actix_web::http::Method;
 use actix_web::{web, FromRequest, HttpRequest, HttpResponse, Responder};
 use bytes::Bytes;
 use futures::future::{ok, Either, FutureResult};
@@ -235,6 +237,7 @@ pub struct Files {
     renderer: Rc<DirectoryRenderer>,
     mime_override: Option<Rc<MimeOverride>>,
     file_flags: named::Flags,
+    guards: Option<Rc<Box<dyn Guard>>>,
 }
 
 impl Clone for Files {
@@ -248,6 +251,7 @@ impl Clone for Files {
             file_flags: self.file_flags,
             path: self.path.clone(),
             mime_override: self.mime_override.clone(),
+            guards: self.guards.clone(),
         }
     }
 }
@@ -273,6 +277,7 @@ impl Files {
             renderer: Rc::new(directory_listing),
             mime_override: None,
             file_flags: named::Flags::default(),
+            guards: None,
         }
     }
 
@@ -331,6 +336,15 @@ impl Files {
         self
     }
 
+    /// Specifies custom guards to use for directory listings and files.
+    ///
+    /// Default behaviour allows GET and HEAD.
+    #[inline]
+    pub fn use_guards<G: Guard + 'static>(mut self, guards: G) -> Self {
+        self.guards = Some(Rc::new(Box::new(guards)));
+        self
+    }
+
     /// Disable `Content-Disposition` header.
     ///
     /// By default Content-Disposition` header is enabled.
@@ -392,6 +406,7 @@ impl NewService for Files {
             renderer: self.renderer.clone(),
             mime_override: self.mime_override.clone(),
             file_flags: self.file_flags,
+            guards: self.guards.clone(),
         };
 
         if let Some(ref default) = *self.default.borrow() {
@@ -418,6 +433,7 @@ pub struct FilesService {
     renderer: Rc<DirectoryRenderer>,
     mime_override: Option<Rc<MimeOverride>>,
     file_flags: named::Flags,
+    guards: Option<Rc<Box<dyn Guard>>>,
 }
 
 impl FilesService {
@@ -454,6 +470,25 @@ impl Service for FilesService {
     fn call(&mut self, req: ServiceRequest) -> Self::Future {
         // let (req, pl) = req.into_parts();
 
+        let is_method_valid = if let Some(guard) = &self.guards {
+            // execute user defined guards
+            (**guard).check(req.head())
+        } else {
+            // default behaviour
+            match *req.method() {
+                Method::HEAD | Method::GET => true,
+                _ => false,
+            }
+        };
+
+        if !is_method_valid {
+            return Either::A(ok(req.into_response(
+                actix_web::HttpResponse::MethodNotAllowed()
+                    .header(header::CONTENT_TYPE, "text/plain")
+                    .body("Request did not meet this resource's requirements.")
+            )));
+        }
+
         let real_path = match PathBufWrp::get_pathbuf(req.match_info().path()) {
             Ok(item) => item,
             Err(e) => return Either::A(ok(req.error_response(e))),