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;

#[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 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 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-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 clear_share_session_cookie(kind: &str, slug: &str) -> HeaderValue {
    clear_cookie(&share_session_cookie_name(kind, slug))
}

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!("__Host-rustpad_access_{}", hex::encode(&digest[..12]))
}

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)]
mod tests {
    use super::*;

    fn websocket_headers(origin: &'static str, host: &'static str) -> HeaderMap {
        let mut headers = HeaderMap::new();
        headers.insert(header::ORIGIN, HeaderValue::from_static(origin));
        headers.insert(header::HOST, HeaderValue::from_static(host));
        headers
    }

    #[test]
    fn account_sessions_are_cookie_only() {
        let mut headers = HeaderMap::new();
        headers.insert(
            header::AUTHORIZATION,
            HeaderValue::from_static("Bearer legacy-account-token"),
        );
        assert_eq!(session_token(&headers), None);
        assert_eq!(bearer_token(&headers), Some("legacy-account-token"));

        headers.insert(
            header::COOKIE,
            HeaderValue::from_static("__Host-rustpad_session=cookie-token"),
        );
        assert_eq!(session_token(&headers), Some("cookie-token"));
    }

    #[test]
    fn prefers_proxy_controlled_real_ip() {
        let mut headers = HeaderMap::new();
        headers.insert(
            axum::http::HeaderName::from_static("x-forwarded-for"),
            HeaderValue::from_static("203.0.113.10"),
        );
        headers.insert(
            axum::http::HeaderName::from_static("x-real-ip"),
            HeaderValue::from_static("198.51.100.20"),
        );
        assert_eq!(client_key(&headers), "ip:198.51.100.20");
    }

    #[test]
    fn accepts_same_origin_websocket() {
        let headers = websocket_headers("https://pad.example.com", "pad.example.com");
        assert!(websocket_origin_allowed(&headers));
    }

    #[test]
    fn rejects_cross_origin_websocket() {
        let headers = websocket_headers("https://evil.example", "pad.example.com");
        assert!(!websocket_origin_allowed(&headers));
    }

    #[test]
    fn does_not_trust_forwarded_host_for_websocket_origin() {
        let mut headers = websocket_headers("https://evil.example", "pad.example.com");
        headers.insert(
            axum::http::HeaderName::from_static("x-forwarded-host"),
            HeaderValue::from_static("evil.example"),
        );
        assert!(!websocket_origin_allowed(&headers));
    }

    #[test]
    fn rejects_origin_with_path() {
        let headers = websocket_headers("https://pad.example.com/other", "pad.example.com");
        assert!(!websocket_origin_allowed(&headers));
    }

    #[test]
    fn rejects_missing_websocket_origin() {
        let mut headers = HeaderMap::new();
        headers.insert(header::HOST, HeaderValue::from_static("pad.example.com"));
        assert!(!websocket_origin_allowed(&headers));
    }

    #[test]
    fn secure_cookies_are_not_script_readable() {
        let value = session_cookie("abc123", 7).to_str().unwrap();
        assert!(value.contains("HttpOnly"));
        assert!(value.contains("Secure"));
        assert!(value.contains("SameSite=Lax"));
        assert!(value.starts_with("__Host-rustpad_session=abc123;"));
    }

    #[test]
    fn share_sessions_use_separate_scoped_opaque_cookies() {
        let value = share_session_cookie("workspace", "private-space", "opaque", 600)
            .to_str()
            .unwrap();
        assert!(value.starts_with("__Host-rustpad_share_"));
        assert!(value.contains("=opaque;"));
        assert!(value.contains("Max-Age=600"));
        assert!(value.contains("HttpOnly"));
        assert!(value.contains("Secure"));
        assert!(value.contains("SameSite=Lax"));

        let cleared = clear_share_session_cookie("workspace", "private-space")
            .to_str()
            .unwrap();
        assert!(cleared.contains("Max-Age=0"));
    }

    #[test]
    fn csrf_requires_matching_cookie_and_header() {
        let token = "a".repeat(CSRF_TOKEN_BYTES * 2);
        let mut headers = HeaderMap::new();
        headers.insert(
            header::COOKIE,
            HeaderValue::from_str(&format!("{CSRF_COOKIE}={token}")).unwrap(),
        );
        headers.insert(
            axum::http::HeaderName::from_static(CSRF_HEADER),
            HeaderValue::from_str(&token).unwrap(),
        );
        assert!(csrf_request_is_valid(&headers));

        headers.insert(
            axum::http::HeaderName::from_static(CSRF_HEADER),
            HeaderValue::from_static(
                "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
            ),
        );
        assert!(!csrf_request_is_valid(&headers));
    }

    #[test]
    fn csrf_cookie_is_strict_and_script_readable() {
        let token = "a".repeat(CSRF_TOKEN_BYTES * 2);
        let value = csrf_cookie(&token).to_str().unwrap();
        assert!(value.contains("Secure"));
        assert!(value.contains("SameSite=Strict"));
        assert!(!value.contains("HttpOnly"));
    }
}