tower-surf 0.3.0

use futures_util::future::BoxFuture;
use http::{HeaderValue, Request, Response};
use secstr::SecStr;
use std::{
    sync::Arc,
    task::{Context, Poll},
};
use tower_cookies::{
    cookie::{Expiration, SameSite},
    CookieManager, Cookies,
};
use tower_layer::Layer;
use tower_service::Service;

use crate::{guard::GuardService, Error, Token};

#[derive(Clone)]
pub(crate) struct Config {
    pub(crate) secret: SecStr,
    pub(crate) cookie_name: String,
    pub(crate) expires: Expiration,
    pub(crate) header_name: String,
    pub(crate) hsts: bool,
    pub(crate) http_only: bool,
    pub(crate) prefix: bool,
    pub(crate) preload: bool,
    pub(crate) same_site: SameSite,
    pub(crate) secure: bool,
}

impl Config {
    pub(crate) fn cookie_name(&self) -> String {
        if self.prefix {
            format!("__HOST-{}", self.cookie_name)
        } else {
            self.cookie_name.clone()
        }
    }
}

/// A layer providing the [`Token`] extension.
///
/// On every request, it will create a cookie for the current visitor if
/// one has not already been created. The session ID used is an `i128`
/// generated by the [`rand`] crate.
#[derive(Clone)]
pub struct Surf {
    pub(crate) config: Config,
}

impl Surf {
    /// Creates a new [`Surf`] layer with the provided secret and default token configuration.
    pub fn new(secret: impl Into<String>) -> Self {
        Self {
            config: Config {
                secret: SecStr::from(secret.into()),
                cookie_name: "csrf_token".into(),
                expires: Expiration::Session,
                header_name: "X-CSRF-Token".into(),
                hsts: true,
                http_only: true,
                prefix: true,
                preload: false,
                same_site: SameSite::Strict,
                secure: true,
            },
        }
    }

    /// Sets the cookie name. Note that this will be previed with `__HOST-` unless
    /// you have disabled it with [prefix](`Surf::prefix`). The default value is `csrf_token`.
    pub fn cookie_name(mut self, cookie_name: impl Into<String>) -> Self {
        self.config.cookie_name = cookie_name.into();

        self
    }

    /// Sets the cookie's expiration. The default value is `Expiration::Session`.
    pub fn expires(mut self, expires: Expiration) -> Self {
        self.config.expires = expires;

        self
    }

    /// Sets the header name used when validating the request. The default
    /// value is `X-CSRF-Token`.
    pub fn header_name(mut self, header_name: impl Into<String>) -> Self {
        self.config.header_name = header_name.into();

        self
    }

    /// Sets whether to send the `Strict-Transport-Security` header.
    ///
    /// See: [HTTP Strict Transport Security Cheat Sheet](https://cheatsheetseries.owasp.org/cheatsheets/HTTP_Strict_Transport_Security_Cheat_Sheet.html)
    pub fn hsts(mut self, hsts: bool) -> Self {
        self.config.hsts = hsts;

        self
    }

    /// Sets the `HTTPOnly` attribute of the cookie. The default value is `true`.
    ///
    /// ⚠️ **Warning**: This should generally _not_ be set to false.
    /// See: [HttpOnly Cookie Attribute](https://owasp.org/www-community/HttpOnly).
    pub fn http_only(mut self, http_only: bool) -> Self {
        self.config.http_only = http_only;

        self
    }

    /// Sets whether to prefix the cookie name with `__HOST-`. The default
    /// value is `true`.
    ///
    /// See: [Cookie Name](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#cookie-namecookie-value).
    pub fn prefix(mut self, prefix: bool) -> Self {
        self.config.prefix = prefix;

        self
    }

    /// Sets whether to append the [hsts](`Surf::hsts`) header with `preload`.
    ///
    /// See: [HTTP Strict Transport Security Cheat Sheet](https://cheatsheetseries.owasp.org/cheatsheets/HTTP_Strict_Transport_Security_Cheat_Sheet.html)
    pub fn preload(mut self, preload: bool) -> Self {
        self.config.preload = preload;

        self
    }

    /// Sets the `SameSite` attribute of the cookie. The default value is [`SameSite::Strict`].
    ///
    /// See: [SameSite Cookie Attribute](https://owasp.org/www-community/SameSite).
    pub fn same_site(mut self, same_site: SameSite) -> Self {
        self.config.same_site = same_site;

        self
    }

    /// Sets the `secure` attribute of the cookie. Note that this is required to
    /// be `false` for cookies to work on `localhost`. The default value is `true`.
    ///
    /// See: [Secure Cookie Attribute](https://owasp.org/www-community/controls/SecureCookieAttribute).
    pub fn secure(mut self, secure: bool) -> Self {
        self.config.secure = secure;

        self
    }
}

impl<S> Layer<S> for Surf {
    type Service = CookieManager<SurfService<GuardService<S>>>;

    fn layer(&self, inner: S) -> Self::Service {
        CookieManager::new(SurfService {
            config: Arc::new(self.config.clone()),
            inner: GuardService::new(inner),
        })
    }
}

#[derive(Clone)]
pub struct SurfService<S> {
    config: Arc<Config>,
    inner: S,
}

impl<S, Q, R> Service<Request<Q>> for SurfService<S>
where
    S: Service<Request<Q>, Response = Response<R>> + Send + 'static,
    S::Future: Send + 'static,
    Q: Send + 'static,
    R: Default + Send,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut request: Request<Q>) -> Self::Future {
        let cookies = match request
            .extensions()
            .get::<Cookies>()
            .ok_or(Error::ExtensionNotFound("Cookies".into()))
        {
            Ok(cookies) => cookies,
            Err(err) => return Box::pin(async move { Error::make_layer_error(err) }),
        };

        let token = Token {
            config: self.config.clone(),
            cookies: cookies.clone(),
        };

        if cookies.get(&self.config.cookie_name()).is_none() {
            if let Err(err) = token.create() {
                return Box::pin(async move { Error::make_layer_error(err) });
            };
        }

        request.extensions_mut().insert(self.config.clone());
        request.extensions_mut().insert(token);

        let config = self.config.clone();

        if config.hsts {
            let future = self.inner.call(request);

            Box::pin(async move {
                let mut response = future.await?;

                let mut value = "max-age=31536000; includeSubDomains".to_owned();

                if config.preload {
                    value.push_str("; preload");
                }

                let value = match HeaderValue::from_str(&value) {
                    Ok(value) => value,
                    Err(err) => return Error::make_layer_error(err),
                };

                response
                    .headers_mut()
                    .insert("Strict-Transport-Security", value);

                Ok(response)
            })
        } else {
            Box::pin(self.inner.call(request))
        }
    }
}