appks/session/middleware.rs

use std::fmt;
use std::future::{Future, Ready, ready};
use std::pin::Pin;
use std::rc::Rc;

use actix_web::HttpResponse;
use actix_web::body::MessageBody;
use actix_web::cookie::time::Duration;
use actix_web::cookie::{Cookie, CookieJar, Key, SameSite};
use actix_web::dev::{
    ResponseHead, Service, ServiceRequest, ServiceResponse, Transform, forward_ready,
};
use actix_web::http::header::{HeaderValue, SET_COOKIE};

use crate::config::AppConfig;
use crate::error::{AppError, AppResult};
use crate::session::storage::{SessionKey, SessionStore};
use crate::session::{Session, SessionState, SessionStatus};

/// Controls when the TTL of server-side session state is refreshed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TtlExtensionPolicy {
    /// Refresh TTL on every request associated with a valid session key.
    OnEveryRequest,
    /// Refresh TTL only when state changes or the session key is renewed.
    OnStateChanges,
}

/// Controls how the session cookie value (the server-side session key) is protected.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CookieContentSecurity {
    /// Encrypt and authenticate the cookie value.
    Private,
    /// Sign the cookie value without encrypting it.
    Signed,
}

/// Actix Web middleware that loads, persists and expires `Session` state.
///
/// The cookie only stores a signed/encrypted session key; the full state is kept
/// in the configured `SessionStore` implementation (Redis in production).
#[derive(Clone)]
pub struct SessionMiddleware<Store: SessionStore> {
    storage_backend: Rc<Store>,
    configuration: Rc<Configuration>,
}

impl<Store: SessionStore> SessionMiddleware<Store> {
    pub fn new(store: Store, key: Key) -> Self {
        Self::builder(store, key).build()
    }

    pub fn builder(store: Store, key: Key) -> SessionMiddlewareBuilder<Store> {
        SessionMiddlewareBuilder::new(store, default_configuration(key))
    }

    pub fn from_app_config(store: Store, key: Key, app_config: &AppConfig) -> AppResult<Self> {
        let ttl_secs = app_config.session_ttl_secs()?;
        if ttl_secs == 0 {
            return Err(AppError::Config(
                "APP_SESSION_TTL_SECS must be greater than 0".into(),
            ));
        }

        Ok(Self::builder(store, key)
            .cookie_name(app_config.session_cookie_name()?)
            .cookie_secure(app_config.session_cookie_secure()?)
            .cookie_http_only(app_config.session_cookie_http_only()?)
            .cookie_same_site(parse_same_site(&app_config.session_cookie_same_site()?)?)
            .cookie_path(app_config.session_cookie_path()?)
            .cookie_domain(app_config.session_cookie_domain()?)
            .cookie_max_age_secs(app_config.session_max_age_secs()?)
            .state_ttl_secs(ttl_secs)
            .build())
    }

    pub(crate) fn from_parts(store: Store, configuration: Configuration) -> Self {
        Self {
            storage_backend: Rc::new(store),
            configuration: Rc::new(configuration),
        }
    }
}

impl<S, B, Store> Transform<S, ServiceRequest> for SessionMiddleware<Store>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
    S::Future: 'static,
    B: MessageBody + 'static,
    Store: SessionStore + 'static,
{
    type Response = ServiceResponse<B>;
    type Error = actix_web::Error;
    type Transform = InnerSessionMiddleware<S, Store>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(InnerSessionMiddleware {
            service: Rc::new(service),
            configuration: Rc::clone(&self.configuration),
            storage_backend: Rc::clone(&self.storage_backend),
        }))
    }
}

#[doc(hidden)]
pub struct InnerSessionMiddleware<S, Store: SessionStore + 'static> {
    service: Rc<S>,
    configuration: Rc<Configuration>,
    storage_backend: Rc<Store>,
}

impl<S, B, Store> Service<ServiceRequest> for InnerSessionMiddleware<S, Store>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
    S::Future: 'static,
    B: MessageBody + 'static,
    Store: SessionStore + 'static,
{
    type Response = ServiceResponse<B>;
    type Error = actix_web::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;

    forward_ready!(service);

    fn call(&self, mut req: ServiceRequest) -> Self::Future {
        let service = Rc::clone(&self.service);
        let storage_backend = Rc::clone(&self.storage_backend);
        let configuration = Rc::clone(&self.configuration);

        Box::pin(async move {
            let session_key = extract_session_key(&req, &configuration.cookie);
            let (session_key, session_state) =
                load_session_state(session_key, storage_backend.as_ref()).await?;

            Session::set_session(&mut req, session_state);
            let mut res = service.call(req).await?;
            let (status, session_state) = Session::get_changes(&mut res);

            persist_session_changes(
                res.response_mut().head_mut(),
                session_key,
                status,
                session_state,
                storage_backend.as_ref(),
                &configuration,
            )
            .await?;

            Ok(res)
        })
    }
}

/// Fluent builder for `SessionMiddleware`.
pub struct SessionMiddlewareBuilder<Store: SessionStore> {
    storage_backend: Store,
    configuration: Configuration,
}

impl<Store: SessionStore> SessionMiddlewareBuilder<Store> {
    fn new(store: Store, configuration: Configuration) -> Self {
        Self {
            storage_backend: store,
            configuration,
        }
    }

    pub fn cookie_name(mut self, name: String) -> Self {
        self.configuration.cookie.name = name;
        self
    }

    pub fn cookie_secure(mut self, secure: bool) -> Self {
        self.configuration.cookie.secure = secure;
        self
    }

    pub fn cookie_http_only(mut self, http_only: bool) -> Self {
        self.configuration.cookie.http_only = http_only;
        self
    }

    pub fn cookie_same_site(mut self, same_site: SameSite) -> Self {
        self.configuration.cookie.same_site = same_site;
        self
    }

    pub fn cookie_path(mut self, path: String) -> Self {
        self.configuration.cookie.path = path;
        self
    }

    pub fn cookie_domain(mut self, domain: Option<String>) -> Self {
        self.configuration.cookie.domain = domain;
        self
    }

    pub fn cookie_max_age_secs(mut self, max_age_secs: Option<u64>) -> Self {
        self.configuration.cookie.max_age = max_age_secs.map(duration_from_secs);
        self
    }

    pub fn state_ttl_secs(mut self, ttl_secs: u64) -> Self {
        self.configuration.session.state_ttl_secs = ttl_secs;
        self
    }

    pub fn ttl_extension_policy(mut self, policy: TtlExtensionPolicy) -> Self {
        self.configuration.ttl_extension_policy = policy;
        self
    }

    pub fn cookie_content_security(mut self, security: CookieContentSecurity) -> Self {
        self.configuration.cookie.content_security = security;
        self
    }

    pub fn build(self) -> SessionMiddleware<Store> {
        SessionMiddleware::from_parts(self.storage_backend, self.configuration)
    }
}

#[derive(Clone)]
pub(crate) struct Configuration {
    cookie: CookieConfiguration,
    session: SessionConfiguration,
    ttl_extension_policy: TtlExtensionPolicy,
}

#[derive(Clone)]
pub(crate) struct SessionConfiguration {
    state_ttl_secs: u64,
}

#[derive(Clone)]
pub(crate) struct CookieConfiguration {
    secure: bool,
    http_only: bool,
    name: String,
    same_site: SameSite,
    path: String,
    domain: Option<String>,
    max_age: Option<Duration>,
    content_security: CookieContentSecurity,
    key: Key,
}

fn default_configuration(key: Key) -> Configuration {
    Configuration {
        cookie: CookieConfiguration {
            secure: true,
            http_only: true,
            name: "sid".into(),
            same_site: SameSite::Lax,
            path: "/".into(),
            domain: None,
            max_age: None,
            content_security: CookieContentSecurity::Private,
            key,
        },
        session: SessionConfiguration {
            state_ttl_secs: 86_400,
        },
        ttl_extension_policy: TtlExtensionPolicy::OnStateChanges,
    }
}

async fn persist_session_changes<Store: SessionStore>(
    response: &mut ResponseHead,
    session_key: Option<SessionKey>,
    status: SessionStatus,
    session_state: SessionState,
    storage_backend: &Store,
    configuration: &Configuration,
) -> Result<(), actix_web::Error> {
    match session_key {
        None => persist_new_session(response, session_state, storage_backend, configuration).await,
        Some(session_key) => {
            persist_existing_session(
                response,
                session_key,
                status,
                session_state,
                storage_backend,
                configuration,
            )
            .await
        }
    }
}

async fn persist_new_session<Store: SessionStore>(
    response: &mut ResponseHead,
    session_state: SessionState,
    storage_backend: &Store,
    configuration: &Configuration,
) -> Result<(), actix_web::Error> {
    if session_state.is_empty() {
        return Ok(());
    }

    let session_key = storage_backend
        .save(session_state, configuration.session.state_ttl_secs)
        .await
        .map_err(e500)?;
    set_session_cookie(response, session_key, &configuration.cookie).map_err(e500)
}

async fn persist_existing_session<Store: SessionStore>(
    response: &mut ResponseHead,
    session_key: SessionKey,
    status: SessionStatus,
    session_state: SessionState,
    storage_backend: &Store,
    configuration: &Configuration,
) -> Result<(), actix_web::Error> {
    match status {
        SessionStatus::Changed => {
            let session_key = storage_backend
                .update(
                    session_key,
                    session_state,
                    configuration.session.state_ttl_secs,
                )
                .await
                .map_err(e500)?;
            set_session_cookie(response, session_key, &configuration.cookie).map_err(e500)
        }
        SessionStatus::Purged => {
            storage_backend.delete(&session_key).await.map_err(e500)?;
            delete_session_cookie(response, &configuration.cookie).map_err(e500)
        }
        SessionStatus::Renewed => {
            storage_backend.delete(&session_key).await.map_err(e500)?;
            let session_key = storage_backend
                .save(session_state, configuration.session.state_ttl_secs)
                .await
                .map_err(e500)?;
            set_session_cookie(response, session_key, &configuration.cookie).map_err(e500)
        }
        SessionStatus::Unchanged => {
            refresh_ttl_if_needed(response, session_key, storage_backend, configuration).await
        }
    }
}

async fn refresh_ttl_if_needed<Store: SessionStore>(
    response: &mut ResponseHead,
    session_key: SessionKey,
    storage_backend: &Store,
    configuration: &Configuration,
) -> Result<(), actix_web::Error> {
    if configuration.ttl_extension_policy != TtlExtensionPolicy::OnEveryRequest {
        return Ok(());
    }

    storage_backend
        .update_ttl(&session_key, configuration.session.state_ttl_secs)
        .await
        .map_err(e500)?;

    if configuration.cookie.max_age.is_some() {
        set_session_cookie(response, session_key, &configuration.cookie).map_err(e500)?;
    }

    Ok(())
}

fn extract_session_key(req: &ServiceRequest, config: &CookieConfiguration) -> Option<SessionKey> {
    let cookies = req.cookies().ok()?;
    let session_cookie = cookies.iter().find(|cookie| cookie.name() == config.name)?;

    let mut jar = CookieJar::new();
    jar.add_original(session_cookie.clone());

    let verified = match config.content_security {
        CookieContentSecurity::Signed => jar.signed(&config.key).get(&config.name),
        CookieContentSecurity::Private => jar.private(&config.key).get(&config.name),
    };

    if verified.is_none() {
        tracing::warn!("session cookie failed signature verification/decryption");
    }

    match verified?.value().to_owned().try_into() {
        Ok(session_key) => Some(session_key),
        Err(err) => {
            tracing::warn!(error = %err, "invalid session key, ignoring cookie");
            None
        }
    }
}

async fn load_session_state<Store: SessionStore>(
    session_key: Option<SessionKey>,
    storage_backend: &Store,
) -> Result<(Option<SessionKey>, SessionState), actix_web::Error> {
    let Some(session_key) = session_key else {
        return Ok((None, SessionState::new()));
    };

    match storage_backend.load(&session_key).await {
        Ok(Some(state)) => Ok((Some(session_key), state)),
        Ok(None) => {
            tracing::info!("session key is valid but state is missing; starting a fresh session");
            Ok((None, SessionState::new()))
        }
        Err(err) if is_invalid_state_error(&err) => {
            tracing::warn!(error = %err, "invalid session state; using an empty state");
            Ok((Some(session_key), SessionState::new()))
        }
        Err(err) => Err(e500(err)),
    }
}

fn set_session_cookie(
    response: &mut ResponseHead,
    session_key: SessionKey,
    config: &CookieConfiguration,
) -> AppResult<()> {
    let mut cookie = Cookie::new(config.name.clone(), String::from(session_key));
    cookie.set_secure(config.secure);
    cookie.set_http_only(config.http_only);
    cookie.set_same_site(config.same_site);
    cookie.set_path(config.path.clone());

    if let Some(max_age) = config.max_age {
        cookie.set_max_age(max_age);
    }
    if let Some(domain) = &config.domain {
        cookie.set_domain(domain.clone());
    }

    let mut jar = CookieJar::new();
    match config.content_security {
        CookieContentSecurity::Signed => jar.signed_mut(&config.key).add(cookie),
        CookieContentSecurity::Private => jar.private_mut(&config.key).add(cookie),
    }

    let cookie = jar
        .delta()
        .next()
        .ok_or_else(|| AppError::InternalServerError("failed to build session cookie".into()))?;
    append_set_cookie(response, cookie.encoded().to_string())
}

fn delete_session_cookie(
    response: &mut ResponseHead,
    config: &CookieConfiguration,
) -> AppResult<()> {
    let mut removal_cookie = Cookie::build(config.name.clone(), "")
        .path(config.path.clone())
        .secure(config.secure)
        .http_only(config.http_only)
        .same_site(config.same_site)
        .finish();

    if let Some(domain) = &config.domain {
        removal_cookie.set_domain(domain.clone());
    }

    removal_cookie.make_removal();
    append_set_cookie(response, removal_cookie.to_string())
}

fn append_set_cookie(response: &mut ResponseHead, value: String) -> AppResult<()> {
    let value = HeaderValue::from_str(&value).map_err(|err| {
        AppError::InternalServerError(format!("invalid Set-Cookie header value: {err}"))
    })?;
    response.headers_mut().append(SET_COOKIE, value);
    Ok(())
}

pub(crate) fn parse_same_site(value: &str) -> AppResult<SameSite> {
    match value.trim().to_ascii_lowercase().as_str() {
        "strict" => Ok(SameSite::Strict),
        "lax" => Ok(SameSite::Lax),
        "none" => Ok(SameSite::None),
        other => Err(AppError::Config(format!(
            "invalid APP_SESSION_COOKIE_SAME_SITE value: {other}"
        ))),
    }
}

fn duration_from_secs(secs: u64) -> Duration {
    let secs = i64::try_from(secs).unwrap_or(i64::MAX);
    Duration::seconds(secs)
}

fn is_invalid_state_error(err: &AppError) -> bool {
    matches!(err, AppError::Json(_) | AppError::Config(_))
}

fn e500<E>(err: E) -> actix_web::Error
where
    E: fmt::Debug + fmt::Display + 'static,
{
    tracing::error!(error = %err, "session middleware internal error");
    actix_web::error::InternalError::from_response(
        err,
        HttpResponse::InternalServerError().finish(),
    )
    .into()
}