4 files changed, 142 insertions, 213 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
diff --git a/src/response.rs b/src/response.rs
index 39e6c98..9c92e4c 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -1,60 +1,50 @@
-//! Provides the [`Response`] convenience wrapper.
+//! Provides the [`SputnikBuilder`] trait.
 
 use cookie::Cookie;
-use hyper::{StatusCode, header::{self, HeaderName, HeaderValue}};
+use hyper::{StatusCode, header::{self, HeaderValue}, http};
 use time::{Duration, OffsetDateTime};
+use hyper::http::response::Builder;
 
-type HyperResponse = hyper::Response<hyper::Body>;
+pub trait SputnikBuilder {
+    /// Adds a Set-Cookie header.
+    fn set_cookie(&mut self, cookie: Cookie);
 
-/// Convenience wrapper around [`hyper::Response`].
-pub struct Response {
-    res: HyperResponse
-}
+    /// Adds a Set-Cookie header to delete a cookie.
+    fn delete_cookie(&mut self, name: &str);
 
-impl Into<HyperResponse> for Response {
-    fn into(self) -> HyperResponse {
-        self.res
-    }
+    /// Sets the Content-Type.
+    fn content_type(self, mime: mime::Mime) -> Builder;
 }
 
-impl Response {
-    pub fn new() -> Self {
-        Response{res: HyperResponse::new(hyper::Body::empty())}
-    }
-
-    pub fn status(&mut self) -> &mut StatusCode {
-        self.res.status_mut()
-    }
-
-    pub fn body(&mut self) -> &mut hyper::Body {
-        self.res.body_mut()
-    }
 
-    pub fn headers(&mut self) -> &mut hyper::HeaderMap<header::HeaderValue> {
-        self.res.headers_mut()
-    }
+pub fn redirect(location: &str, code: StatusCode) -> Builder {
+    Builder::new().status(code).header(header::LOCATION, location)
+}
 
-    pub fn set_header<S: AsRef<str>>(&mut self, header: HeaderName, value: S) {
-        self.res.headers_mut().insert(header, HeaderValue::from_str(value.as_ref()).unwrap());
+impl SputnikBuilder for Builder {
+    fn set_cookie(&mut self, cookie: Cookie) {
+        self.headers_mut().unwrap().insert(header::SET_COOKIE, HeaderValue::from_str(&cookie.encoded().to_string()).unwrap());
     }
 
-    pub fn set_content_type(&mut self, mime: mime::Mime) {
-        self.res.headers_mut().insert(header::CONTENT_TYPE, mime.to_string().parse().unwrap());
+    fn delete_cookie(&mut self, name: &str) {
+        let mut cookie = Cookie::new(name, "");
+        cookie.set_max_age(Duration::seconds(0));
+        cookie.set_expires(OffsetDateTime::now_utc() - Duration::days(365));
+        self.set_cookie(cookie);
     }
 
-    pub fn redirect<S: AsRef<str>>(&mut self, location: S, code: StatusCode) {
-        *self.res.status_mut() = code;
-        self.set_header(header::LOCATION, location);
+    fn content_type(self, mime: mime::Mime) -> Self {
+        self.header(header::CONTENT_TYPE, mime.to_string())
     }
+}
 
-    pub fn set_cookie(&mut self, cookie: Cookie) {
-        self.res.headers_mut().append(header::SET_COOKIE, cookie.encoded().to_string().parse().unwrap());
-    }
+pub trait EmptyBuilder<B> {
+    /// Consume the builder with an empty body.
+    fn empty(self) -> http::Result<http::response::Response<B>>;
+}
 
-    pub fn delete_cookie(&mut self, name: &str) {
-        let mut cookie = Cookie::new(name, "");
-        cookie.set_max_age(Duration::seconds(0));
-        cookie.set_expires(OffsetDateTime::now_utc() - Duration::days(365));
-        self.set_cookie(cookie);
+impl EmptyBuilder<hyper::Body> for Builder {
+    fn empty(self) -> http::Result<http::response::Response<hyper::Body>> {
+        self.body(hyper::Body::empty())
     }
 }
\ No newline at end of file
diff --git a/src/security.rs b/src/security.rs
index c12a97f..5247d9e 100644
--- a/src/security.rs
+++ b/src/security.rs
@@ -1,18 +1,14 @@
 //! [`CsrfToken`], [`Key`] and functions to encode & decode expiring claims.
 
-use rand::{Rng, distributions::Alphanumeric};
-use ::cookie::Cookie;
-use time::{Duration, OffsetDateTime};
+use time::OffsetDateTime;
 use thiserror::Error;
 
 pub use crate::signed::Key;
 
-use crate::{request::Parts, response::Response};
-
-/// A cookie-based CSRF token to be used with [`crate::request::Body::into_form_csrf`].
+/// A cookie-based CSRF token to be used with [`crate::request::SputnikBody::into_form_csrf`].
 pub struct CsrfToken {
-    token: String,
-    from_client: bool,
+    pub(crate) token: String,
+    pub(crate) from_client: bool,
 }
 
 #[derive(Error, Debug)]
@@ -25,20 +21,6 @@ pub enum CsrfError {
 }
 
 impl CsrfToken {
-    /// Retrieves the CSRF token from a `csrf` cookie or generates
-    /// a new token and stores it as a cookie if it doesn't exist.
-    pub fn from_parts(request: &mut Parts, response: &mut Response) -> Self {
-        if let Some(cookie) = request.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)));
-        response.set_cookie(c);
-        CsrfToken{token: val, from_client: false}
-    }
-
     /// Wraps the token in a hidden HTML input.
     pub fn html_input(&self) -> String {
         format!("<input name=csrf type=hidden value=\"{}\">", self.token)
diff --git a/src/signed.rs b/src/signed.rs
index 9f0d184..4da6760 100644
--- a/src/signed.rs
+++ b/src/signed.rs
@@ -1,7 +1,3 @@
-// This code was adapted from the cookie crate which does not make the sign and verify functions public
-// forcing the use of CookieJars, which are akward to work with without a high-level framework.
-// Thanks to Sergio Benitez for writing the original code and releasing it under MIT!
-
 use hmac::{Hmac,NewMac,Mac};
 use sha2::Sha256;
 
@@ -9,6 +5,10 @@ const SIGNED_KEY_LEN: usize = 32;
 const BASE64_DIGEST_LEN: usize = 44;
 
 /// A convenience wrapper around HMAC.
+///
+/// This code was adapted from the [`cookie`] crate which does not make the sign and verify functions public
+/// forcing the use of [`CookieJar`](cookie::CookieJar)s, which are akward to work with without a high-level framework.
+// Thanks to Sergio Benitez for writing the original code and releasing it under MIT!
 pub struct Key (pub [u8; SIGNED_KEY_LEN]);
 
 impl Key {