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::http::{HeaderMap, HeaderValue, Uri, header};
use sha2::{Digest, Sha256};

pub const SESSION_COOKIE: &str = "__Host-rustpad_session";

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 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 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 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]))
}

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 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;"));
    }
}