From b9d870645f1d25e1b981d0fa7c9afafbab601c74 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Tue, 15 May 2018 10:09:48 -0700
Subject: [PATCH] store cookies in extensions

---
 src/httprequest.rs | 17 ++++++++++-------
 src/test.rs        | 15 +++++----------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/src/httprequest.rs b/src/httprequest.rs
index fbca52caf..0c3ee31de 100644
--- a/src/httprequest.rs
+++ b/src/httprequest.rs
@@ -39,7 +39,6 @@ pub struct HttpInnerMessage {
     pub headers: HeaderMap,
     pub extensions: Extensions,
     pub params: Params<'static>,
-    pub cookies: Option<Vec<Cookie<'static>>>,
     pub addr: Option<SocketAddr>,
     pub payload: Option<Payload>,
     pub info: Option<ConnectionInfo<'static>>,
@@ -48,6 +47,7 @@ pub struct HttpInnerMessage {
 }
 
 struct Query(Params<'static>);
+struct Cookies(Vec<Cookie<'static>>);
 
 #[derive(Debug, Copy, Clone, PartialEq)]
 enum RouterResource {
@@ -65,7 +65,6 @@ impl Default for HttpInnerMessage {
             flags: MessageFlags::empty(),
             params: Params::new(),
             addr: None,
-            cookies: None,
             payload: None,
             extensions: Extensions::new(),
             info: None,
@@ -90,7 +89,6 @@ impl HttpInnerMessage {
         self.addr = None;
         self.info = None;
         self.flags = MessageFlags::empty();
-        self.cookies = None;
         self.payload = None;
         self.prefix = 0;
         self.resource = RouterResource::Notset;
@@ -124,7 +122,6 @@ impl HttpRequest<()> {
                 payload,
                 params: Params::new(),
                 extensions: Extensions::new(),
-                cookies: None,
                 addr: None,
                 info: None,
                 prefix: 0,
@@ -402,7 +399,7 @@ impl<S> HttpRequest<S> {
 
     /// Load request cookies.
     pub fn cookies(&self) -> Result<&Vec<Cookie<'static>>, CookieParseError> {
-        if self.as_ref().cookies.is_none() {
+        if let None = self.extensions().get::<Query>() {
             let msg = self.as_mut();
             let mut cookies = Vec::new();
             for hdr in msg.headers.get_all(header::COOKIE) {
@@ -413,9 +410,9 @@ impl<S> HttpRequest<S> {
                     }
                 }
             }
-            msg.cookies = Some(cookies);
+            msg.extensions.insert(Cookies(cookies));
         }
-        Ok(&self.as_ref().cookies.as_ref().unwrap())
+        Ok(&self.extensions().get::<Cookies>().unwrap().0)
     }
 
     /// Return request cookie.
@@ -430,6 +427,12 @@ impl<S> HttpRequest<S> {
         None
     }
 
+    pub(crate) fn set_cookies(&mut self, cookies: Option<Vec<Cookie<'static>>>) {
+        if let Some(cookies) = cookies {
+            self.extensions_mut().insert(Cookies(cookies));
+        }
+    }
+
     /// Get a reference to the Params object.
     ///
     /// Params is a container for url parameters.
diff --git a/src/test.rs b/src/test.rs
index 4e7398396..135135f7e 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -355,12 +355,7 @@ impl<S: 'static> TestApp<S> {
 
     /// Register handler for "/"
     pub fn handler<H: Handler<S>>(&mut self, handler: H) {
-        self.app = Some(
-            self.app
-                .take()
-                .unwrap()
-                .resource("/", |r| r.h(handler)),
-        );
+        self.app = Some(self.app.take().unwrap().resource("/", |r| r.h(handler)));
     }
 
     /// Register middleware
@@ -562,8 +557,8 @@ impl<S: 'static> TestRequest<S> {
             cookies,
             payload,
         } = self;
-        let req = HttpRequest::new(method, uri, version, headers, payload);
-        req.as_mut().cookies = cookies;
+        let mut req = HttpRequest::new(method, uri, version, headers, payload);
+        req.set_cookies(cookies);
         req.as_mut().params = params;
         let (router, _) = Router::new::<S>("/", ServerSettings::default(), Vec::new());
         req.with_state(Rc::new(state), router)
@@ -583,8 +578,8 @@ impl<S: 'static> TestRequest<S> {
             payload,
         } = self;
 
-        let req = HttpRequest::new(method, uri, version, headers, payload);
-        req.as_mut().cookies = cookies;
+        let mut req = HttpRequest::new(method, uri, version, headers, payload);
+        req.set_cookies(cookies);
         req.as_mut().params = params;
         req.with_state(Rc::new(state), router)
     }