From d595dd850e73f4ce075390c82f40745982be1180 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Thu, 7 Dec 2017 18:00:20 -0800
Subject: [PATCH] load cookies automatically

---
 src/httprequest.rs         | 70 ++++++++++++++++----------------------
 src/middlewares/session.rs |  2 +-
 src/router.rs              |  3 +-
 tests/test_httprequest.rs  | 11 +++---
 tests/test_httpresponse.rs |  4 +--
 5 files changed, 38 insertions(+), 52 deletions(-)

diff --git a/src/httprequest.rs b/src/httprequest.rs
index 36dcf22c4..44e13f875 100644
--- a/src/httprequest.rs
+++ b/src/httprequest.rs
@@ -23,12 +23,10 @@ struct HttpMessage {
     version: Version,
     method: Method,
     uri: Uri,
-    prefix: usize,
     headers: HeaderMap,
     extensions: Extensions,
     params: Params<'static>,
-    cookies: Vec<Cookie<'static>>,
-    cookies_loaded: bool,
+    cookies: Option<Vec<Cookie<'static>>>,
     addr: Option<SocketAddr>,
     payload: Payload,
     info: Option<ConnectionInfo<'static>>,
@@ -41,12 +39,10 @@ impl Default for HttpMessage {
         HttpMessage {
             method: Method::GET,
             uri: Uri::default(),
-            prefix: 0,
             version: Version::HTTP_11,
             headers: HeaderMap::new(),
             params: Params::default(),
-            cookies: Vec::new(),
-            cookies_loaded: false,
+            cookies: None,
             addr: None,
             payload: Payload::empty(),
             extensions: Extensions::new(),
@@ -68,12 +64,10 @@ impl HttpRequest<()> {
             Rc::new(HttpMessage {
                 method: method,
                 uri: uri,
-                prefix: 0,
                 version: version,
                 headers: headers,
                 params: Params::default(),
-                cookies: Vec::new(),
-                cookies_loaded: false,
+                cookies: None,
                 addr: None,
                 payload: payload,
                 extensions: Extensions::new(),
@@ -119,14 +113,13 @@ impl<S> HttpRequest<S> {
         &mut self.as_mut().extensions
     }
 
-    #[inline]
-    pub(crate) fn set_prefix(&mut self, idx: usize) {
-        self.as_mut().prefix = idx;
-    }
-
     #[doc(hidden)]
     pub fn prefix_len(&self) -> usize {
-        self.0.prefix
+        if let Some(router) = self.router() {
+            router.prefix().len()
+        } else {
+            0
+        }
     }
 
     /// Read the Request Uri.
@@ -225,37 +218,34 @@ impl<S> HttpRequest<S> {
         }
     }
 
-    /// Return request cookies.
+    /// Load request cookies.
     #[inline]
-    pub fn cookies(&self) -> &Vec<Cookie<'static>> {
-        &self.0.cookies
-    }
-
-    /// Return request cookie.
-    pub fn cookie(&self, name: &str) -> Option<&Cookie> {
-        for cookie in &self.0.cookies {
-            if cookie.name() == name {
-                return Some(cookie)
-            }
-        }
-        None
-    }
-
-    /// Load cookies
-    pub fn load_cookies(&mut self) -> Result<&Vec<Cookie<'static>>, CookieParseError>
-    {
-        if !self.0.cookies_loaded {
+    pub fn cookies(&self) -> Result<&Vec<Cookie<'static>>, CookieParseError> {
+        if self.0.cookies.is_none() {
             let msg = self.as_mut();
-            msg.cookies_loaded = true;
+            let mut cookies = Vec::new();
             if let Some(val) = msg.headers.get(header::COOKIE) {
                 let s = str::from_utf8(val.as_bytes())
                     .map_err(CookieParseError::from)?;
                 for cookie in s.split("; ") {
-                    msg.cookies.push(Cookie::parse_encoded(cookie)?.into_owned());
+                    cookies.push(Cookie::parse_encoded(cookie)?.into_owned());
+                }
+            }
+            msg.cookies = Some(cookies)
+        }
+        Ok(self.0.cookies.as_ref().unwrap())
+    }
+
+    /// Return request cookie.
+    pub fn cookie(&self, name: &str) -> Option<&Cookie> {
+        if let Ok(cookies) = self.cookies() {
+            for cookie in cookies {
+                if cookie.name() == name {
+                    return Some(cookie)
                 }
             }
         }
-        Ok(&self.0.cookies)
+        None
     }
 
     /// Get a reference to the Params object.
@@ -535,7 +525,7 @@ mod tests {
         let mut headers = HeaderMap::new();
         headers.insert(header::HOST,
                        header::HeaderValue::from_static("www.rust-lang.org"));
-        let mut req = HttpRequest::new(
+        let req = HttpRequest::new(
             Method::GET, Uri::from_str("/").unwrap(),
             Version::HTTP_11, headers, Payload::empty());
 
@@ -550,7 +540,7 @@ mod tests {
         assert_eq!(req.url_for("unknown", &["test"]),
                    Err(UrlGenerationError::RouterNotAvailable));
 
-        let mut req = req.with_state(Rc::new(()), router);
+        let req = req.with_state(Rc::new(()), router);
 
         assert_eq!(req.url_for("unknown", &["test"]),
                    Err(UrlGenerationError::ResourceNotFound));
@@ -573,7 +563,7 @@ mod tests {
         let router = Router::new("", map);
         assert!(!router.has_route("https://youtube.com/watch/unknown"));
 
-        let mut req = req.with_state(Rc::new(()), router);
+        let req = req.with_state(Rc::new(()), router);
         let url = req.url_for("youtube", &["oHg5SJYRHA0"]);
         assert_eq!(url.ok().unwrap().as_str(), "https://youtube.com/watch/oHg5SJYRHA0");
     }
diff --git a/src/middlewares/session.rs b/src/middlewares/session.rs
index 1ed02ee0c..f962171a7 100644
--- a/src/middlewares/session.rs
+++ b/src/middlewares/session.rs
@@ -259,7 +259,7 @@ impl CookieSessionInner {
     }
 
     fn load(&self, req: &mut HttpRequest) -> HashMap<String, String> {
-        if let Ok(cookies) = req.load_cookies() {
+        if let Ok(cookies) = req.cookies() {
             for cookie in cookies {
                 if cookie.name() == self.name {
                     let mut jar = CookieJar::new();
diff --git a/src/router.rs b/src/router.rs
index b90f79c56..3d03e11a0 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -54,7 +54,7 @@ impl<S> Router<S> {
 
     /// Router prefix
     #[inline]
-    pub(crate) fn prefix(&self) -> &str {
+    pub fn prefix(&self) -> &str {
         &self.0.prefix
     }
 
@@ -74,7 +74,6 @@ impl<S> Router<S> {
 
         if let Some(idx) = idx {
             let path: &str = unsafe{ mem::transmute(&req.path()[self.0.prefix.len()..]) };
-            req.set_prefix(self.prefix().len());
             self.0.patterns[idx].update_match_info(path, req);
             return Some(&self.0.resources[idx])
         } else {
diff --git a/tests/test_httprequest.rs b/tests/test_httprequest.rs
index f3519dbe6..b6fecce57 100644
--- a/tests/test_httprequest.rs
+++ b/tests/test_httprequest.rs
@@ -20,12 +20,10 @@ fn test_debug() {
 
 #[test]
 fn test_no_request_cookies() {
-    let mut req = HttpRequest::new(
+    let req = HttpRequest::new(
         Method::GET, Uri::from_str("/").unwrap(),
         Version::HTTP_11, HeaderMap::new(), Payload::empty());
-    assert!(req.cookies().is_empty());
-    let _ = req.load_cookies();
-    assert!(req.cookies().is_empty());
+    assert!(req.cookies().unwrap().is_empty());
 }
 
 #[test]
@@ -34,12 +32,11 @@ fn test_request_cookies() {
     headers.insert(header::COOKIE,
                    header::HeaderValue::from_static("cookie1=value1; cookie2=value2"));
 
-    let mut req = HttpRequest::new(
+    let req = HttpRequest::new(
         Method::GET, Uri::from_str("/").unwrap(),
         Version::HTTP_11, headers, Payload::empty());
-    assert!(req.cookies().is_empty());
     {
-        let cookies = req.load_cookies().unwrap();
+        let cookies = req.cookies().unwrap();
         assert_eq!(cookies.len(), 2);
         assert_eq!(cookies[0].name(), "cookie1");
         assert_eq!(cookies[0].value(), "value1");
diff --git a/tests/test_httpresponse.rs b/tests/test_httpresponse.rs
index 79e629ed9..53b1149b9 100644
--- a/tests/test_httpresponse.rs
+++ b/tests/test_httpresponse.rs
@@ -14,9 +14,9 @@ fn test_response_cookies() {
     headers.insert(header::COOKIE,
                    header::HeaderValue::from_static("cookie1=value1; cookie2=value2"));
 
-    let mut req = HttpRequest::new(
+    let req = HttpRequest::new(
         Method::GET, Uri::from_str("/").unwrap(), Version::HTTP_11, headers, Payload::empty());
-    let cookies = req.load_cookies().unwrap();
+    let cookies = req.cookies().unwrap();
 
     let resp = httpcodes::HTTPOk
         .build()