From 7030bf5fe82dc3c23e61a6f0e7a6b89b53a03e6d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Augusto=20C=C3=A9sar=20Dias?= <augusto.c.dias@gmail.com>
Date: Mon, 26 Oct 2020 18:02:45 +0100
Subject: [PATCH] Adding app_data to ServiceConfig (#1758)

Co-authored-by: Rob Ede <robjtede@icloud.com>
Co-authored-by: Augusto <augusto@flowciety.de>
---
 CHANGES.md                   |  1 +
 actix-http/src/extensions.rs | 35 +++++++++++++++++++++++++++++++++++
 src/app.rs                   |  1 +
 src/config.rs                | 23 ++++++++++++++++++-----
 src/scope.rs                 |  3 +++
 5 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index baa462a5..af34c3b4 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -4,6 +4,7 @@
 ### Added
 * Implement `exclude_regex` for Logger middleware. [#1723]
 * Add request-local data extractor `web::ReqData`. [#1748]
+* Add `app_data` to `ServiceConfig`. [#1757]
 
 ### Changed
 * Print non-configured `Data<T>` type when attempting extraction. [#1743]
diff --git a/actix-http/src/extensions.rs b/actix-http/src/extensions.rs
index 96e01767..09f1b711 100644
--- a/actix-http/src/extensions.rs
+++ b/actix-http/src/extensions.rs
@@ -61,6 +61,11 @@ impl Extensions {
     pub fn clear(&mut self) {
         self.map.clear();
     }
+
+    /// Extends self with the items from another `Extensions`.
+    pub fn extend(&mut self, other: Extensions) {
+        self.map.extend(other.map);
+    }
 }
 
 impl fmt::Debug for Extensions {
@@ -178,4 +183,34 @@ mod tests {
         assert_eq!(extensions.get::<bool>(), None);
         assert_eq!(extensions.get(), Some(&MyType(10)));
     }
+
+    #[test]
+    fn test_extend() {
+        #[derive(Debug, PartialEq)]
+        struct MyType(i32);
+
+        let mut extensions = Extensions::new();
+
+        extensions.insert(5i32);
+        extensions.insert(MyType(10));
+
+        let mut other = Extensions::new();
+
+        other.insert(15i32);
+        other.insert(20u8);
+
+        extensions.extend(other);
+
+        assert_eq!(extensions.get(), Some(&15i32));
+        assert_eq!(extensions.get_mut(), Some(&mut 15i32));
+
+        assert_eq!(extensions.remove::<i32>(), Some(15i32));
+        assert!(extensions.get::<i32>().is_none());
+
+        assert_eq!(extensions.get::<bool>(), None);
+        assert_eq!(extensions.get(), Some(&MyType(10)));
+
+        assert_eq!(extensions.get(), Some(&20u8));
+        assert_eq!(extensions.get_mut(), Some(&mut 20u8));
+    }
 }
diff --git a/src/app.rs b/src/app.rs
index 6a4b97b6..8dd86f7e 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -183,6 +183,7 @@ where
         self.data.extend(cfg.data);
         self.services.extend(cfg.services);
         self.external.extend(cfg.external);
+        self.extensions.extend(cfg.extensions);
         self
     }
 
diff --git a/src/config.rs b/src/config.rs
index f7bebb4c..03ba8273 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -178,6 +178,7 @@ pub struct ServiceConfig {
     pub(crate) services: Vec<Box<dyn AppServiceFactory>>,
     pub(crate) data: Vec<Box<dyn DataFactory>>,
     pub(crate) external: Vec<ResourceDef>,
+    pub(crate) extensions: Extensions,
 }
 
 impl ServiceConfig {
@@ -186,6 +187,7 @@ impl ServiceConfig {
             services: Vec::new(),
             data: Vec::new(),
             external: Vec::new(),
+            extensions: Extensions::new(),
         }
     }
 
@@ -198,6 +200,14 @@ impl ServiceConfig {
         self
     }
 
+    /// Set arbitrary data item.
+    ///
+    /// This is same as `App::data()` method.
+    pub fn app_data<U: 'static>(&mut self, ext: U) -> &mut Self {
+        self.extensions.insert(ext);
+        self
+    }
+
     /// Configure route for a specific path.
     ///
     /// This is same as `App::route()` method.
@@ -254,13 +264,16 @@ mod tests {
     async fn test_data() {
         let cfg = |cfg: &mut ServiceConfig| {
             cfg.data(10usize);
+            cfg.app_data(15u8);
         };
 
-        let mut srv =
-            init_service(App::new().configure(cfg).service(
-                web::resource("/").to(|_: web::Data<usize>| HttpResponse::Ok()),
-            ))
-            .await;
+        let mut srv = init_service(App::new().configure(cfg).service(
+            web::resource("/").to(|_: web::Data<usize>, req: HttpRequest| {
+                assert_eq!(*req.app_data::<u8>().unwrap(), 15u8);
+                HttpResponse::Ok()
+            }),
+        ))
+        .await;
         let req = TestRequest::default().to_request();
         let resp = srv.call(req).await.unwrap();
         assert_eq!(resp.status(), StatusCode::OK);
diff --git a/src/scope.rs b/src/scope.rs
index 1c5d8700..681d142b 100644
--- a/src/scope.rs
+++ b/src/scope.rs
@@ -209,6 +209,9 @@ where
 
             self.data = Some(data);
         }
+        self.data
+            .get_or_insert_with(Extensions::new)
+            .extend(cfg.extensions);
         self
     }