From 58d71ad86b545291c80058bcb725d8d94d34cf41 Mon Sep 17 00:00:00 2001
From: ljoonal <ljoona@ljoonal.xyz>
Date: Mon, 17 Aug 2020 14:49:16 +0300
Subject: [PATCH] Configurable behaviour for NormalizePath

---
 src/middleware/normalize.rs | 70 ++++++++++++++++++++++++++++++++-----
 1 file changed, 61 insertions(+), 9 deletions(-)

diff --git a/src/middleware/normalize.rs b/src/middleware/normalize.rs
index a1021ed14..98a18ecfb 100644
--- a/src/middleware/normalize.rs
+++ b/src/middleware/normalize.rs
@@ -10,20 +10,29 @@ use regex::Regex;
 use crate::service::{ServiceRequest, ServiceResponse};
 use crate::Error;
 
-#[derive(Default, Clone, Copy)]
+/// To be used when constructing `NormalizePath` to define it's behavior.
+#[derive(Clone, Copy)]
+pub enum TrailingSlash {
+    /// Always add a trailing slash to the end of the path.
+    Always,
+    /// Trim trailing slashes from the end of the path.
+    Trim,
+}
+
+#[derive(Clone, Copy)]
 /// `Middleware` to normalize request's URI in place
 ///
 /// Performs following:
 ///
 /// - Merges multiple slashes into one.
-/// - Appends a trailing slash if one is not present.
+/// - Appends a trailing slash if one is not present, or removes one if present, depending on the supplied `TrailingSlash`.
 ///
 /// ```rust
 /// use actix_web::{web, http, middleware, App, HttpResponse};
 ///
 /// # fn main() {
 /// let app = App::new()
-///     .wrap(middleware::NormalizePath)
+///     .wrap(middleware::NormalizePath::default())
 ///     .service(
 ///         web::resource("/test")
 ///             .route(web::get().to(|| HttpResponse::Ok()))
@@ -32,7 +41,20 @@ use crate::Error;
 /// # }
 /// ```
 
-pub struct NormalizePath;
+pub struct NormalizePath(TrailingSlash);
+
+impl NormalizePath {
+    /// Create new `NormalizePath` middleware with the specified trailing slash style.
+    pub fn new(trailing_slash_style: TrailingSlash) -> Self {
+        NormalizePath(trailing_slash_style)
+    }
+}
+
+impl Default for NormalizePath {
+    fn default() -> Self {
+        NormalizePath::new(TrailingSlash::Always)
+    }
+}
 
 impl<S, B> Transform<S> for NormalizePath
 where
@@ -50,6 +72,10 @@ where
         ok(NormalizePathNormalization {
             service,
             merge_slash: Regex::new("//+").unwrap(),
+            trim_last_slash: match self.0 {
+                TrailingSlash::Trim => true,
+                TrailingSlash::Always => false,
+            },
         })
     }
 }
@@ -57,6 +83,7 @@ where
 pub struct NormalizePathNormalization<S> {
     service: S,
     merge_slash: Regex,
+    trim_last_slash: bool,
 }
 
 impl<S, B> Service for NormalizePathNormalization<S>
@@ -78,8 +105,11 @@ where
 
         let original_path = head.uri.path();
 
-        // always add trailing slash, might be an extra one
-        let path = original_path.to_string() + "/";
+        // Either adds a string to the end (duplicates will be removed anyways) or trims all slashes from the end
+        let path = match self.trim_last_slash {
+            false => original_path.to_string() + "/",
+            true => original_path.trim_end_matches(|c| c == '/').to_string(),
+        };
 
         // normalize multiple /'s to one /
         let path = self.merge_slash.replace_all(&path, "/");
@@ -150,6 +180,28 @@ mod tests {
         assert!(res4.status().is_success());
     }
 
+    #[actix_rt::test]
+    async fn trim_trailing_slashes() {
+        let mut app = init_service(
+            App::new()
+                .wrap(NormalizePath(TrailingSlash::Trim))
+                .service(web::resource("/v1/something").to(HttpResponse::Ok)),
+        )
+        .await;
+
+        let req4 = TestRequest::with_uri("/v1/something////").to_request();
+        let res4 = call_service(&mut app, req4).await;
+        assert!(res4.status().is_success());
+
+        let req4 = TestRequest::with_uri("/v1/something/").to_request();
+        let res4 = call_service(&mut app, req4).await;
+        assert!(res4.status().is_success());
+
+        let req4 = TestRequest::with_uri("//v1//something//").to_request();
+        let res4 = call_service(&mut app, req4).await;
+        assert!(res4.status().is_success());
+    }
+
     #[actix_rt::test]
     async fn test_in_place_normalization() {
         let srv = |req: ServiceRequest| {
@@ -157,7 +209,7 @@ mod tests {
             ok(req.into_response(HttpResponse::Ok().finish()))
         };
 
-        let mut normalize = NormalizePath
+        let mut normalize = NormalizePath::default()
             .new_transform(srv.into_service())
             .await
             .unwrap();
@@ -188,7 +240,7 @@ mod tests {
             ok(req.into_response(HttpResponse::Ok().finish()))
         };
 
-        let mut normalize = NormalizePath
+        let mut normalize = NormalizePath::default()
             .new_transform(srv.into_service())
             .await
             .unwrap();
@@ -207,7 +259,7 @@ mod tests {
             ok(req.into_response(HttpResponse::Ok().finish()))
         };
 
-        let mut normalize = NormalizePath
+        let mut normalize = NormalizePath::default()
             .new_transform(srv.into_service())
             .await
             .unwrap();