hyaenidae/accounts/src/extractors.rs

use crate::{store::User, State};
use actix_session::{Session, UserSession};
use actix_web::{
    dev::{Payload, Service, ServiceRequest, ServiceResponse, Transform},
    http::{header::LOCATION, StatusCode},
    web::Data,
    FromRequest, HttpMessage, HttpRequest, HttpResponse, ResponseError,
};
use event_listener::Event;
use futures::future::{ok, LocalBoxFuture, Ready};
use std::{
    cell::{Cell, RefCell},
    rc::Rc,
    task::{Context, Poll},
};
use uuid::Uuid;

#[derive(serde::Deserialize, serde::Serialize)]
pub(crate) struct CookieData {
    accepted: bool,
}

impl CookieData {
    pub(crate) fn set_accepted(session: &Session) -> Result<(), SessionError> {
        session
            .set("accepted", CookieData { accepted: true })
            .map_err(|_| SessionError::Set)
    }

    fn accepted(session: &Session) -> bool {
        session
            .get::<CookieData>("accepted")
            .map(|opt| opt.is_some())
            .unwrap_or(false)
    }
}

#[derive(serde::Deserialize, serde::Serialize)]
pub(crate) struct UserData {
    id: Uuid,
}

impl UserData {
    pub(crate) fn set_data(id: Uuid, session: &Session) -> Result<(), SessionError> {
        session
            .set("user-data", UserData { id })
            .map_err(|_| SessionError::Set)
    }

    pub(crate) fn remove(session: &Session) {
        session.remove("user-data")
    }

    fn data(session: &Session) -> Result<Option<UserData>, actix_web::Error> {
        Ok(session.get::<UserData>("user-data")?)
    }
}

pub struct AcceptedCookies(());

#[derive(Clone)]
pub struct Authenticated {
    user: User,
}

impl Authenticated {
    pub fn user(self) -> User {
        self.user
    }
}

#[derive(Debug, thiserror::Error)]
#[error("{0}")]
struct ServerError(String);

impl ResponseError for ServerError {
    fn status_code(&self) -> StatusCode {
        StatusCode::SEE_OTHER
    }

    fn error_response(&self) -> HttpResponse {
        HttpResponse::build(self.status_code())
            .header(LOCATION, self.0.clone())
            .finish()
    }
}

#[derive(Debug, thiserror::Error)]
pub enum SessionError {
    #[error("Failed to set sesion data")]
    Set,
}

impl FromRequest for AcceptedCookies {
    type Error = actix_web::Error;
    type Future = LocalBoxFuture<'static, Result<Self, Self::Error>>;
    type Config = ();

    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        let session = req.get_session();
        let state_fut = Data::<State>::extract(req);

        Box::pin(async move {
            let state = state_fut.await?;

            if CookieData::accepted(&session) {
                Ok(AcceptedCookies(()))
            } else {
                Err(ServerError(state.pages.not_found_path()).into())
            }
        })
    }
}

impl FromRequest for User {
    type Error = actix_web::Error;
    type Future = LocalBoxFuture<'static, Result<Self, Self::Error>>;
    type Config = ();

    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        let fut = Authenticated::extract(req);

        Box::pin(async move { fut.await.map(|auth| auth.user()) })
    }
}

impl FromRequest for Authenticated {
    type Error = actix_web::Error;
    type Future = LocalBoxFuture<'static, Result<Self, Self::Error>>;
    type Config = ();

    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        let opt = req
            .extensions()
            .get::<AuthState>()
            .map(|auth_state| (auth_state.clone(), auth_state.2.listen()));

        let state_fut = Data::<State>::extract(req);

        Box::pin(async move {
            let (auth_state, listen_fut) = match opt {
                Some(auth_state) => auth_state,
                None => {
                    let state = state_fut.await?;
                    return Err(ServerError(state.pages.not_found_path()).into());
                }
            };

            if let Some(auth) = auth_state.0.borrow().as_ref() {
                return Ok(auth.clone());
            }

            if !auth_state.1.get() {
                listen_fut.await;

                if let Some(auth) = auth_state.0.borrow().as_ref() {
                    return Ok(auth.clone());
                }
            }

            let state = state_fut.await?;
            Err(ServerError(state.pages.not_found_path()).into())
        })
    }
}

pub struct Auth(pub State);
pub struct AuthMiddleware<S>(S, State);
#[derive(Clone)]
struct AuthState(
    Rc<RefCell<Option<Authenticated>>>,
    Rc<Cell<bool>>,
    Rc<Event>,
);
struct DropGuard(Rc<Event>, Rc<Cell<bool>>);

fn auth_state() -> (AuthState, DropGuard) {
    let event = Rc::new(Event::new());
    let state = Rc::new(RefCell::new(None));
    let flag = Rc::new(Cell::new(false));

    (
        AuthState(state, Rc::clone(&flag), Rc::clone(&event)),
        DropGuard(event, flag),
    )
}

impl Drop for DropGuard {
    fn drop(&mut self) {
        self.1.set(true);
        self.0.notify(usize::MAX)
    }
}

impl<S, B> Transform<S> for Auth
where
    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
    S::Future: 'static,
{
    type Request = S::Request;
    type Response = S::Response;
    type Error = S::Error;
    type InitError = ();
    type Transform = AuthMiddleware<S>;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ok(AuthMiddleware(service, self.0.clone()))
    }
}

impl<S, B> Service for AuthMiddleware<S>
where
    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
    S::Future: 'static,
{
    type Request = S::Request;
    type Response = S::Response;
    type Error = S::Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

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

    fn call(&mut self, req: Self::Request) -> Self::Future {
        let session = req.get_session();
        let state = self.1.clone();

        let (auth_state, drop_guard) = auth_state();
        req.extensions_mut().insert(auth_state.clone());

        let fut = self.0.call(req);

        Box::pin(async move {
            if let Some(UserData { id }) = UserData::data(&session)? {
                let res = state.user_store.user_by_id(id).await;

                let user_opt = match res {
                    Ok(user_opt) => user_opt,
                    Err(_) => {
                        log::warn!("Error fetching user data from store");
                        return Err(ServerError(state.pages.internal_server_error_path()).into());
                    }
                };

                if let Some(user) = user_opt {
                    if !user.suspended() {
                        *auth_state.0.borrow_mut() = Some(Authenticated { user });
                    }
                }
            }

            drop(drop_guard);

            fut.await
        })
    }
}