rustpad/src/security.rs

/*
 * Copyright (C) 2026 Mateusz Gruszczyński @linuxiarz.pl
 * Source-Available Code / Dual-Licensed.
 *
 * Free for non-commercial and evaluation use under terms of BSL/GPLv3.
 * Commercial or production use requires a valid paid license.
 * See LICENSE file in repository root for details.
 */

use axum::{
    Json,
    http::{HeaderMap, HeaderValue, Uri, header},
    response::{IntoResponse, Response},
};
use rand_core::{OsRng, RngCore};
use serde::Serialize;
use sha2::{Digest, Sha256};

pub const SESSION_COOKIE: &str = "__Host-rustpad_session";
pub const CSRF_COOKIE: &str = "__Host-rustpad_csrf";
pub const CSRF_HEADER: &str = "x-rustpad-csrf";

const CSRF_TOKEN_BYTES: usize = 32;
const CSRF_TTL_SECONDS: i64 = 24 * 60 * 60;
const RESOURCE_ACCESS_COOKIE_PREFIX: &str = "__Host-rustpad_access_";
const RESOURCE_ACCESS_COOKIE_SUFFIX_LENGTH: usize = 24;
const RESOURCE_ACCESS_TOKEN_LENGTH: usize = 64;

#[derive(Serialize)]
pub struct CsrfResponse {
    token: String,
}

pub fn session_token(headers: &HeaderMap) -> Option<&str> {
    session_cookie_token(headers)
}

pub fn session_cookie_token(headers: &HeaderMap) -> Option<&str> {
    cookie_value(headers, SESSION_COOKIE)
}

pub fn bearer_token(headers: &HeaderMap) -> Option<&str> {
    headers
        .get(header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer "))
        .map(str::trim)
        .filter(|value| !value.is_empty())
}

pub fn resource_token<'a>(headers: &'a HeaderMap, kind: &str, slug: &str) -> Option<&'a str> {
    let name = resource_cookie_name(kind, slug);
    cookie_value(headers, &name)
}

pub(crate) fn resource_access_cookies(headers: &HeaderMap) -> Vec<(String, Option<String>)> {
    headers
        .get_all(header::COOKIE)
        .iter()
        .filter_map(|value| value.to_str().ok())
        .flat_map(|cookies| cookies.split(';'))
        .filter_map(|part| {
            let (name, value) = part.trim().split_once('=')?;
            if !valid_resource_access_cookie_name(name) {
                return None;
            }
            let value = value.trim();
            let token = (value.len() == RESOURCE_ACCESS_TOKEN_LENGTH
                && value.bytes().all(|byte| byte.is_ascii_hexdigit()))
            .then(|| value.to_owned());
            Some((name.to_owned(), token))
        })
        .collect()
}

pub fn share_session_token<'a>(headers: &'a HeaderMap, kind: &str, slug: &str) -> Option<&'a str> {
    let name = share_session_cookie_name(kind, slug);
    cookie_value(headers, &name)
}

pub fn session_cookie(token: &str, ttl_days: i64) -> HeaderValue {
    secure_cookie(SESSION_COOKIE, token, ttl_days.saturating_mul(86_400))
}

pub fn clear_session_cookie() -> HeaderValue {
    clear_cookie(SESSION_COOKIE)
}

pub(crate) fn clear_resource_access_cookie(name: &str) -> Option<HeaderValue> {
    valid_resource_access_cookie_name(name).then(|| clear_cookie(name))
}

pub async fn csrf_token_endpoint(headers: HeaderMap) -> Response {
    let token = csrf_cookie_token(&headers)
        .filter(|value| valid_csrf_token(value))
        .map(str::to_owned)
        .unwrap_or_else(random_csrf_token);
    let mut response = Json(CsrfResponse {
        token: token.clone(),
    })
    .into_response();
    response
        .headers_mut()
        .insert(header::SET_COOKIE, csrf_cookie(&token));
    response.headers_mut().insert(
        header::CACHE_CONTROL,
        HeaderValue::from_static("no-cache, no-store, max-age=0"),
    );
    response
}

pub fn csrf_request_is_valid(headers: &HeaderMap) -> bool {
    let Some(cookie) = csrf_cookie_token(headers).filter(|value| valid_csrf_token(value)) else {
        return false;
    };
    let Some(provided) = headers
        .get(CSRF_HEADER)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| valid_csrf_token(value))
    else {
        return false;
    };
    constant_time_eq(cookie.as_bytes(), provided.as_bytes())
}

pub fn csrf_cookie_token(headers: &HeaderMap) -> Option<&str> {
    cookie_value(headers, CSRF_COOKIE)
}

pub fn resource_cookie(kind: &str, slug: &str, token: &str, ttl_days: i64) -> HeaderValue {
    secure_cookie(
        &resource_cookie_name(kind, slug),
        token,
        ttl_days.saturating_mul(86_400),
    )
}

pub fn share_session_cookie(
    kind: &str,
    slug: &str,
    token: &str,
    max_age_seconds: i64,
) -> HeaderValue {
    secure_cookie(
        &share_session_cookie_name(kind, slug),
        token,
        max_age_seconds,
    )
}

pub fn client_key(headers: &HeaderMap) -> String {
    let forwarded_ip = header_ip(headers, "cf-connecting-ip")
        .or_else(|| header_ip(headers, "x-real-ip"))
        .or_else(|| header_ip(headers, "x-forwarded-for"));
    if let Some(value) = forwarded_ip {
        return format!("ip:{value}");
    }

    let user_agent = first_header_value(headers, header::USER_AGENT.as_str()).unwrap_or("");
    let language = first_header_value(headers, header::ACCEPT_LANGUAGE.as_str()).unwrap_or("");
    if user_agent.is_empty() && language.is_empty() {
        return "unknown".into();
    }
    let digest = Sha256::digest(format!("{user_agent}|{language}").as_bytes());
    format!("browser:{}", hex::encode(&digest[..12]))
}

pub fn websocket_origin_allowed(headers: &HeaderMap) -> bool {
    let Some(origin) = headers
        .get(header::ORIGIN)
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty() && *value != "null")
    else {
        return false;
    };
    let Ok(uri) = origin.parse::<Uri>() else {
        return false;
    };
    let Some(origin_authority) = uri.authority().map(|value| value.as_str()) else {
        return false;
    };
    if uri.path() != "/" || uri.query().is_some() {
        return false;
    }
    let Some(expected_authority) = first_header_value(headers, header::HOST.as_str()) else {
        return false;
    };
    if !origin_authority.eq_ignore_ascii_case(expected_authority) {
        return false;
    }
    if let Some(expected_scheme) = first_header_value(headers, "x-forwarded-proto") {
        let Some(origin_scheme) = uri.scheme_str() else {
            return false;
        };
        if !origin_scheme.eq_ignore_ascii_case(expected_scheme) {
            return false;
        }
    }
    matches!(uri.scheme_str(), Some("http" | "https"))
}

fn resource_cookie_name(kind: &str, slug: &str) -> String {
    let digest = Sha256::digest(format!("{kind}:{slug}").as_bytes());
    format!(
        "{RESOURCE_ACCESS_COOKIE_PREFIX}{}",
        hex::encode(&digest[..12])
    )
}

fn valid_resource_access_cookie_name(name: &str) -> bool {
    name.strip_prefix(RESOURCE_ACCESS_COOKIE_PREFIX)
        .is_some_and(|suffix| {
            suffix.len() == RESOURCE_ACCESS_COOKIE_SUFFIX_LENGTH
                && suffix.bytes().all(|byte| byte.is_ascii_hexdigit())
        })
}

fn share_session_cookie_name(kind: &str, slug: &str) -> String {
    let digest = Sha256::digest(format!("{kind}:{slug}").as_bytes());
    format!("__Host-rustpad_share_{}", hex::encode(&digest[..12]))
}

pub fn cookie_value<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
    headers
        .get(header::COOKIE)
        .and_then(|value| value.to_str().ok())
        .and_then(|cookies| {
            cookies.split(';').find_map(|part| {
                let (cookie_name, cookie_value) = part.trim().split_once('=')?;
                (cookie_name == name && !cookie_value.is_empty()).then_some(cookie_value)
            })
        })
}

fn secure_cookie(name: &str, value: &str, max_age: i64) -> HeaderValue {
    HeaderValue::from_str(&format!(
        "{name}={value}; Path=/; Max-Age={}; HttpOnly; Secure; SameSite=Lax",
        max_age.max(1)
    ))
    .expect("valid secure cookie")
}

fn csrf_cookie(token: &str) -> HeaderValue {
    HeaderValue::from_str(&format!(
        "{CSRF_COOKIE}={token}; Path=/; Max-Age={}; Secure; SameSite=Strict",
        CSRF_TTL_SECONDS
    ))
    .expect("valid csrf cookie")
}

fn random_csrf_token() -> String {
    let mut bytes = [0_u8; CSRF_TOKEN_BYTES];
    let mut rng = OsRng;
    rng.fill_bytes(&mut bytes);
    hex::encode(bytes)
}

fn valid_csrf_token(value: &str) -> bool {
    value.len() == CSRF_TOKEN_BYTES * 2 && value.bytes().all(|byte| byte.is_ascii_hexdigit())
}

fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
    if left.len() != right.len() {
        return false;
    }
    left.iter()
        .zip(right)
        .fold(0_u8, |difference, (left, right)| {
            difference | (*left ^ *right)
        })
        == 0
}

fn clear_cookie(name: &str) -> HeaderValue {
    HeaderValue::from_str(&format!(
        "{name}=; Path=/; Max-Age=0; HttpOnly; Secure; SameSite=Lax"
    ))
    .expect("valid clear cookie")
}

fn header_ip(headers: &HeaderMap, name: &str) -> Option<std::net::IpAddr> {
    first_header_value(headers, name)?.parse().ok()
}

fn first_header_value<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
    headers
        .get(name)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.split(',').next())
        .map(str::trim)
        .filter(|value| !value.is_empty())
}

#[cfg(test)]
#[path = "tests/security.rs"]
mod tests;