1 files changed, 103 insertions, 146 deletions

diff --git a/src/request.rs b/src/request.rs
index 953e7ec..1166ef2 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -1,51 +1,44 @@
-//! Provides the [`Parts`] and [`Body`] convenience wrappers.
+//! Provides the [`SputnikParts`] and [`SputnikBody`] traits.
 
 use cookie::Cookie;
-use header::CONTENT_TYPE;
-use mime::{APPLICATION_WWW_FORM_URLENCODED, Mime};
+use mime::Mime;
+use rand::{Rng, distributions::Alphanumeric};
+use security::CsrfToken;
 use serde::{Deserialize, de::DeserializeOwned};
 use hyper::{body::Bytes, header};
-use hyper::http::request::Parts as ReqParts;
-use std::collections::HashMap;
+use time::Duration;
+use std::{collections::HashMap, sync::Arc};
 
-use crate::security;
+use crate::{response::SputnikBuilder, security};
 
-use error::*;
-
-type HyperRequest = hyper::Request<hyper::Body>;
+pub trait SputnikParts {
+    /// Parses the query string of the request into a given struct.
+    fn query<X: DeserializeOwned>(&self) -> Result<X,QueryError>;
 
-/// Convenience wrapper around [`hyper::Body`].
-pub struct Body {
-    body: hyper::Body,
-    content_type: Option<header::HeaderValue>,
-}
+    /// Parses the cookies of the request.
+    fn cookies(&mut self) -> Arc<HashMap<String, Cookie<'static>>>;
 
-/// Convert [`hyper::Request`] to ([`Parts`], [`Body`])
-pub fn adapt<'a>(req: HyperRequest) -> (Parts, Body) {
-    let (parts, body) = req.into_parts();
-    let body = Body{body, content_type: parts.headers.get(CONTENT_TYPE).map(|x| x.to_owned())};
-    let parts = Parts{parts, cookies: None};
-    (parts, body)
-}
+    /// Retrieves the CSRF token from a `csrf` cookie or generates
+    /// a new token and stores it as a cookie if it doesn't exist.
+    fn csrf_token(&mut self, builder: &mut dyn SputnikBuilder) -> CsrfToken;
 
-/// Convenience wrapper around [`hyper::http::request::Parts`].
-pub struct Parts {
-    parts: ReqParts,
-    cookies: Option<HashMap<String,Cookie<'static>>>,
+    /// Enforces a specific Content-Type.
+    fn enforce_content_type(&self, mime: Mime) -> Result<(), WrongContentTypeError>;
 }
 
-#[derive(Deserialize)]
-struct CsrfData {
-    csrf: String,
-}
+impl SputnikParts for hyper::http::request::Parts {
+    fn query<T: DeserializeOwned>(&self) -> Result<T,QueryError> {
+        serde_urlencoded::from_str::<T>(self.uri.query().unwrap_or("")).map_err(QueryError)
+    }
 
-impl Parts {
-    pub fn cookies(&mut self) -> &HashMap<String,Cookie> {
-        if let Some(ref cookies) = self.cookies {
-            return cookies
+    fn cookies(&mut self) -> Arc<HashMap<String, Cookie<'static>>> {
+        let cookies: Option<&Arc<HashMap<String, Cookie>>> = self.extensions.get();
+        if let Some(cookies) = cookies {
+            return cookies.clone();
         }
+
         let mut cookies = HashMap::new();
-        for header in self.parts.headers.get_all(header::COOKIE) {
+        for header in self.headers.get_all(header::COOKIE) {
             let raw_str = match std::str::from_utf8(header.as_bytes()) {
                 Ok(string) => string,
                 Err(_) => continue
@@ -57,149 +50,113 @@ impl Parts {
                 }
             }
         }
-        self.cookies = Some(cookies);
-        &self.cookies.as_ref().unwrap()
-    }
-
-    pub fn method(&self) -> &hyper::Method {
-        &self.parts.method
+        let cookies = Arc::new(cookies);
+        self.extensions.insert(cookies.clone());
+        cookies
     }
 
-    pub fn headers(&self) -> &hyper::HeaderMap<header::HeaderValue> {
-        &self.parts.headers
-    }
-
-    pub fn uri(&self) -> &hyper::Uri {
-        &self.parts.uri
+    fn csrf_token(&mut self, builder: &mut dyn SputnikBuilder) -> CsrfToken {
+        if let Some(cookie) = self.cookies().get("csrf") {
+            return CsrfToken{token: cookie.value().to_string(), from_client: true}
+        }
+        let val: String = rand::thread_rng().sample_iter(Alphanumeric).take(16).collect();
+        let mut c = Cookie::new("csrf", val.clone());
+        c.set_secure(Some(true));
+        c.set_max_age(Some(Duration::hours(1)));
+        builder.set_cookie(c);
+        CsrfToken{token: val, from_client: false}
     }
 
-    /// Parses the query string of the request into a given struct.
-    pub fn query<T: DeserializeOwned>(&self) -> Result<T,QueryError> {
-        serde_urlencoded::from_str::<T>(self.parts.uri.query().unwrap_or("")).map_err(QueryError)
+    fn enforce_content_type(&self, mime: Mime) -> Result<(), WrongContentTypeError> {
+        if let Some(content_type) = self.headers.get(header::CONTENT_TYPE) {
+            if *content_type == mime.to_string() {
+                return Ok(())
+            }
+        }
+        Err(WrongContentTypeError{expected: mime, received: self.headers.get(header::CONTENT_TYPE).as_ref().and_then(|h| h.to_str().ok().map(|s| s.to_owned()))})
     }
 }
 
-impl Body {
-    pub async fn into_bytes(self) -> Result<Bytes, BodyError> {
-        hyper::body::to_bytes(self.body).await.map_err(BodyError)
-    }
+use async_trait::async_trait;
+
+#[async_trait]
+pub trait SputnikBody {
+    async fn into_bytes(self) -> Result<Bytes, BodyError>;
 
     /// Parses a `application/x-www-form-urlencoded` request body into a given struct.
     ///
     /// This does make you vulnerable to CSRF, so you normally want to use
-    /// [`Body::into_form_csrf()`] instead.
-    ///
-    /// # Example
-    ///
-    /// ```
-    /// use hyper::{Response};
-    /// use sputnik::request::{Body, error::FormError};
-    /// use serde::Deserialize;
-    ///
-    /// #[derive(Deserialize)]
-    /// struct Message {text: String, year: i64}
-    ///
-    /// async fn greet(body: Body) -> Result<Response<hyper::Body>, FormError> {
-    ///     let msg: Message = body.into_form().await?;
-    ///     Ok(Response::new(format!("hello {}", msg.text).into()))
-    /// }
-    /// ```
-    pub async fn into_form<T: DeserializeOwned>(self) -> Result<T, FormError> {
-        self.enforce_content_type(APPLICATION_WWW_FORM_URLENCODED)?;
-        let full_body = self.into_bytes().await?;
-        serde_urlencoded::from_bytes::<T>(&full_body).map_err(FormError::Deserialize)
-    }
+    /// [`SputnikBody::into_form_csrf()`] instead.
+    async fn into_form<T: DeserializeOwned>(self) -> Result<T, FormError>;
 
     /// Parses a `application/x-www-form-urlencoded` request body into a given struct.
     /// Protects from CSRF by checking that the request body contains the same token retrieved from the cookies.
     ///
     /// The CSRF parameter is expected as the `csrf` parameter in the request body.
     /// This means for HTML forms you need to embed the token as a hidden input.
-    ///
-    /// # Example
-    ///
-    /// ```
-    /// use hyper::{Method};
-    /// use sputnik::{request::{Parts, Body, error::CsrfProtectedFormError}, response::Response};
-    /// use sputnik::security::CsrfToken;
-    /// use serde::Deserialize;
-    ///
-    /// #[derive(Deserialize)]
-    /// struct Message {text: String}
-    ///
-    /// async fn greet(req: &mut Parts, body: Body) -> Result<Response, CsrfProtectedFormError> {
-    ///     let mut response = Response::new();
-    ///     let csrf_token = CsrfToken::from_parts(req, &mut response);
-    ///     let msg: Message = body.into_form_csrf(&csrf_token).await?;
-    ///     *response.body() = format!("hello {}", msg.text).into();
-    ///     Ok(response)
-    /// }
-    /// ```
-    pub async fn into_form_csrf<T: DeserializeOwned>(self, csrf_token: &security::CsrfToken) -> Result<T, CsrfProtectedFormError> {
-        self.enforce_content_type(APPLICATION_WWW_FORM_URLENCODED)?;
+    async fn into_form_csrf<T: DeserializeOwned>(self, csrf_token: &security::CsrfToken) -> Result<T, CsrfProtectedFormError>;
+}
+
+#[async_trait]
+impl SputnikBody for hyper::Body {
+    async fn into_bytes(self) -> Result<Bytes, BodyError> {
+        hyper::body::to_bytes(self).await.map_err(BodyError)
+    }
+
+    async fn into_form<T: DeserializeOwned>(self) -> Result<T, FormError> {
+        let full_body = self.into_bytes().await?;
+        Ok(serde_urlencoded::from_bytes::<T>(&full_body)?)
+    }
+
+    async fn into_form_csrf<T: DeserializeOwned>(self, csrf_token: &CsrfToken) -> Result<T, CsrfProtectedFormError> {
         let full_body = self.into_bytes().await?;
         let csrf_data = serde_urlencoded::from_bytes::<CsrfData>(&full_body).map_err(|_| CsrfProtectedFormError::NoCsrf)?;
         csrf_token.matches(csrf_data.csrf)?;
         serde_urlencoded::from_bytes::<T>(&full_body).map_err(CsrfProtectedFormError::Deserialize)
     }
-
-    fn enforce_content_type(&self, mime: Mime) -> Result<(), WrongContentTypeError> {
-        if let Some(content_type) = &self.content_type {
-            if *content_type == mime.to_string() {
-                return Ok(())
-            }
-        }
-        Err(WrongContentTypeError{expected: mime, received: self.content_type.as_ref().and_then(|h| h.to_str().ok().map(|s| s.to_owned()))})
-    }
 }
 
-pub mod error {
-    use mime::Mime;
-    use thiserror::Error;
-
-    use crate::security::CsrfError;
-    #[derive(Error, Debug)]
-    #[error("query deserialize error: {0}")]
-    pub struct QueryError(pub serde_urlencoded::de::Error);
-
-    #[derive(Error, Debug)]
-    #[error("failed to read body")]
-    pub struct BodyError(pub hyper::Error);
+#[derive(Deserialize)]
+struct CsrfData {
+    csrf: String,
+}
 
-    #[derive(Error, Debug)]
-    #[error("expected Content-Type {expected} but received {}", received.as_ref().unwrap_or(&"nothing".to_owned()))]
-    pub struct WrongContentTypeError {
-        pub expected: Mime,
-        pub received: Option<String>,
-    }
+use crate::security::CsrfError;
+#[derive(thiserror::Error, Debug)]
+#[error("query deserialize error: {0}")]
+pub struct QueryError(pub serde_urlencoded::de::Error);
 
-    #[derive(Error, Debug)]
-    pub enum FormError {
-        #[error("{0}")]
-        ContentType(#[from] WrongContentTypeError),
+#[derive(thiserror::Error, Debug)]
+#[error("failed to read body")]
+pub struct BodyError(pub hyper::Error);
 
-        #[error("{0}")]
-        Body(#[from] BodyError),
+#[derive(thiserror::Error, Debug)]
+#[error("expected Content-Type {expected} but received {}", received.as_ref().unwrap_or(&"nothing".to_owned()))]
+pub struct WrongContentTypeError {
+    pub expected: Mime,
+    pub received: Option<String>,
+}
 
-        #[error("form deserialize error: {0}")]
-        Deserialize(#[from] serde_urlencoded::de::Error),
-    }
+#[derive(thiserror::Error, Debug)]
+pub enum FormError {
+    #[error("{0}")]
+    Body(#[from] BodyError),
 
-    #[derive(Error, Debug)]
-    pub enum CsrfProtectedFormError {
-        #[error("{0}")]
-        ContentType(#[from] WrongContentTypeError),
+    #[error("form deserialize error: {0}")]
+    Deserialize(#[from] serde_urlencoded::de::Error),
+}
 
-        #[error("{0}")]
-        Body(#[from] BodyError),
+#[derive(thiserror::Error, Debug)]
+pub enum CsrfProtectedFormError {
+    #[error("{0}")]
+    Body(#[from] BodyError),
 
-        #[error("form deserialize error: {0}")]
-        Deserialize(#[from] serde_urlencoded::de::Error),
+    #[error("form deserialize error: {0}")]
+    Deserialize(#[from] serde_urlencoded::de::Error),
 
-        #[error("no csrf token in form data")]
-        NoCsrf,
+    #[error("no csrf token in form data")]
+    NoCsrf,
 
-        #[error("{0}")]
-        Csrf(#[from] CsrfError),
-    }
+    #[error("{0}")]
+    Csrf(#[from] CsrfError),
 }
\ No newline at end of file