appks/session/session.rs

use std::cell::{Ref, RefCell};
use std::future::{Ready, ready};
use std::mem;
use std::rc::Rc;

use actix_web::dev::{Extensions, Payload, ServiceRequest, ServiceResponse};
use actix_web::{FromRequest, HttpMessage, HttpRequest};
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::{Map, Value};
use uuid::Uuid;

use crate::error::AppError;

const SESSION_USER_KEY: &str = "session:user_uid";

/// Request-local session handle.
///
/// Clones share the same interior state through request extensions, so changes
/// made by extractors/handlers can be collected by `SessionMiddleware` at the
/// end of the request lifecycle.
#[derive(Clone)]
pub struct Session(Rc<RefCell<SessionInner>>);

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum SessionStatus {
    /// Session state has changed and must be persisted.
    Changed,
    /// Session and cookie must be deleted. Further mutations are ignored.
    Purged,
    /// Session key must be regenerated while preserving current state.
    Renewed,
    /// Session has not been modified during this request.
    #[default]
    Unchanged,
}

#[derive(Default)]
struct SessionInner {
    state: SessionState,
    status: SessionStatus,
}

impl Session {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn from_state(state: SessionState) -> Self {
        Self(Rc::new(RefCell::new(SessionInner {
            state,
            status: SessionStatus::Unchanged,
        })))
    }

    pub fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, AppError> {
        let value = self.0.borrow().state.get(key).cloned();
        value
            .map(serde_json::from_value)
            .transpose()
            .map_err(AppError::Json)
    }

    pub fn contains_key(&self, key: &str) -> bool {
        self.0.borrow().state.contains_key(key)
    }

    pub fn entries(&self) -> Ref<'_, SessionState> {
        Ref::map(self.0.borrow(), |inner| &inner.state)
    }

    pub fn status(&self) -> SessionStatus {
        self.0.borrow().status.clone()
    }

    pub fn is_empty(&self) -> bool {
        self.0.borrow().state.is_empty()
    }

    pub fn insert<T: Serialize>(&self, key: impl Into<String>, value: T) -> Result<(), AppError> {
        let value = serde_json::to_value(value)?;
        let mut inner = self.0.borrow_mut();

        if inner.status != SessionStatus::Purged {
            mark_changed(&mut inner.status);
            inner.state.insert(key.into(), value);
        }

        Ok(())
    }

    pub fn update<T: Serialize + DeserializeOwned, F>(
        &self,
        key: impl Into<String>,
        updater: F,
    ) -> Result<(), AppError>
    where
        F: FnOnce(T) -> T,
    {
        let mut inner = self.0.borrow_mut();
        let key = key.into();

        if inner.status == SessionStatus::Purged {
            return Ok(());
        }

        let Some(value) = inner.state.get(&key).cloned() else {
            return Ok(());
        };

        let value = serde_json::from_value(value)?;
        let updated = serde_json::to_value(updater(value))?;
        mark_changed(&mut inner.status);
        inner.state.insert(key, updated);
        Ok(())
    }

    pub fn update_or<T: Serialize + DeserializeOwned, F>(
        &self,
        key: &str,
        default: T,
        updater: F,
    ) -> Result<(), AppError>
    where
        F: FnOnce(T) -> T,
    {
        if self.contains_key(key) {
            self.update(key, updater)
        } else {
            self.insert(key, default)
        }
    }

    pub fn remove(&self, key: &str) -> Option<Value> {
        let mut inner = self.0.borrow_mut();

        if inner.status == SessionStatus::Purged {
            return None;
        }

        mark_changed(&mut inner.status);
        inner.state.remove(key)
    }

    pub fn remove_as<T: DeserializeOwned>(&self, key: &str) -> Option<Result<T, AppError>> {
        self.remove(key)
            .map(|value| serde_json::from_value(value).map_err(AppError::Json))
    }

    pub fn clear(&self) {
        let mut inner = self.0.borrow_mut();

        if inner.status != SessionStatus::Purged {
            mark_changed(&mut inner.status);
            inner.state.clear();
        }
    }

    pub fn purge(&self) {
        let mut inner = self.0.borrow_mut();
        inner.status = SessionStatus::Purged;
        inner.state.clear();
    }

    pub fn renew(&self) {
        let mut inner = self.0.borrow_mut();

        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Renewed;
        }
    }

    pub fn user(&self) -> Option<Uuid> {
        self.get::<Uuid>(SESSION_USER_KEY).ok().flatten()
    }

    pub fn set_user(&self, uid: Uuid) {
        let _ = self.insert(SESSION_USER_KEY, uid);
    }

    pub fn clear_user(&self) {
        let _ = self.remove(SESSION_USER_KEY);
    }

    pub fn take_state(&self) -> SessionState {
        mem::take(&mut self.0.borrow_mut().state)
    }

    pub fn mark_unchanged(&self) {
        self.0.borrow_mut().status = SessionStatus::Unchanged;
    }

    pub(crate) fn set_session(req: &mut ServiceRequest, state: SessionState) {
        let session = Self::get_session(&mut req.extensions_mut());
        let mut inner = session.0.borrow_mut();
        inner.state = state;
        inner.status = SessionStatus::Unchanged;
    }

    pub(crate) fn get_changes<B>(res: &mut ServiceResponse<B>) -> (SessionStatus, SessionState) {
        let Some(inner) = res
            .request()
            .extensions()
            .get::<Rc<RefCell<SessionInner>>>()
            .cloned()
        else {
            return (SessionStatus::Unchanged, SessionState::new());
        };

        let mut inner = inner.borrow_mut();
        let status = inner.status.clone();
        let state = mem::take(&mut inner.state);
        (status, state)
    }

    pub(crate) fn get_session(extensions: &mut Extensions) -> Session {
        if let Some(inner) = extensions.get::<Rc<RefCell<SessionInner>>>() {
            return Session(Rc::clone(inner));
        }

        let inner = Rc::new(RefCell::new(SessionInner::default()));
        extensions.insert(Rc::clone(&inner));
        Session(inner)
    }
}

impl Default for Session {
    fn default() -> Self {
        Self(Rc::new(RefCell::new(SessionInner::default())))
    }
}

pub type SessionState = Map<String, Value>;

#[derive(Debug, Clone, Copy)]
pub struct SessionUser(pub Uuid);

impl FromRequest for Session {
    type Error = AppError;
    type Future = Ready<Result<Self, Self::Error>>;

    fn from_request(req: &HttpRequest, _payload: &mut Payload) -> Self::Future {
        ready(Ok(Self::get_session(&mut req.extensions_mut())))
    }
}

fn mark_changed(status: &mut SessionStatus) {
    if *status != SessionStatus::Renewed {
        *status = SessionStatus::Changed;
    }
}