use std::fmt::Write as _;
use std::sync::{Arc, Mutex};
use axum::http::{HeaderMap, Method};
use laterite_core::Text;
use rand::RngCore;
use serde::{Deserialize, Serialize};
const VERSION: u8 = 1;
const MAX_BLOB: usize = 4096;
const CSRF_HEADER: &str = "x-csrf-token";
pub(crate) const CSRF_FIELD: &str = "_csrf";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct Flash {
pub(crate) level: FlashLevel,
pub(crate) text: Text,
#[serde(default)]
pub(crate) sticky: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FlashLevel {
Success,
Error,
Info,
}
impl FlashLevel {
pub(crate) fn class(self) -> &'static str {
match self {
FlashLevel::Success => "is-success",
FlashLevel::Error => "is-error",
FlashLevel::Info => "is-info",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct SessionData {
#[serde(default = "default_version")]
v: u8,
#[serde(default)]
csrf: String,
#[serde(default)]
flash: Vec<Flash>,
#[serde(default)]
lists: Vec<(String, String)>,
}
const REMEMBERED_LISTS: usize = 8;
const REMEMBERED_QUERY: usize = 512;
fn default_version() -> u8 {
VERSION
}
impl Default for SessionData {
fn default() -> Self {
Self {
v: VERSION,
csrf: String::new(),
flash: Vec::new(),
lists: Vec::new(),
}
}
}
impl SessionData {
fn parse(raw: Option<&str>) -> Self {
match raw {
Some(s) if s.len() <= MAX_BLOB => serde_json::from_str(s).unwrap_or_default(),
_ => Self::default(),
}
}
}
struct Inner {
data: SessionData,
dirty: bool,
}
#[derive(Clone)]
pub struct SessionHandle {
inner: Arc<Mutex<Inner>>,
}
impl SessionHandle {
pub(crate) fn from_blob(raw: Option<&str>) -> Self {
let mut data = SessionData::parse(raw);
let dirty = data.csrf.is_empty();
if dirty {
data.csrf = generate_token();
}
Self {
inner: Arc::new(Mutex::new(Inner { data, dirty })),
}
}
pub(crate) fn csrf_token(&self) -> String {
self.inner.lock().unwrap().data.csrf.clone()
}
pub fn push_flash(&self, level: FlashLevel, text: Text) {
let mut g = self.inner.lock().unwrap();
g.data.flash.push(Flash {
level,
text,
sticky: false,
});
g.dirty = true;
}
pub fn push_flash_sticky(&self, level: FlashLevel, text: Text) {
let mut g = self.inner.lock().unwrap();
g.data.flash.push(Flash {
level,
text,
sticky: true,
});
g.dirty = true;
}
pub(crate) fn remembered_list(&self, path: &str) -> Option<String> {
let g = self.inner.lock().unwrap();
g.data
.lists
.iter()
.find(|(p, _)| p == path)
.map(|(_, q)| q.clone())
}
pub(crate) fn remember_list(&self, path: &str, query: Option<&str>) {
let mut g = self.inner.lock().unwrap();
let before = g.data.lists.clone();
g.data.lists.retain(|(p, _)| p != path);
if let Some(query) = query.filter(|q| q.len() <= REMEMBERED_QUERY) {
g.data.lists.push((path.to_string(), query.to_string()));
if g.data.lists.len() > REMEMBERED_LISTS {
g.data.lists.remove(0);
}
}
if g.data.lists != before {
g.dirty = true;
}
}
pub(crate) fn take_flash(&self) -> Vec<Flash> {
let mut g = self.inner.lock().unwrap();
if g.data.flash.is_empty() {
return Vec::new();
}
g.dirty = true;
std::mem::take(&mut g.data.flash)
}
pub(crate) fn dirty_blob(&self) -> Option<String> {
let g = self.inner.lock().unwrap();
g.dirty
.then(|| serde_json::to_string(&g.data).unwrap_or_default())
}
}
pub(crate) fn generate_token() -> String {
let mut bytes = [0u8; 32];
rand::rngs::OsRng.fill_bytes(&mut bytes);
let mut out = String::with_capacity(64);
for b in bytes {
let _ = write!(out, "{b:02x}");
}
out
}
pub(crate) fn is_safe_method(m: &Method) -> bool {
matches!(*m, Method::GET | Method::HEAD | Method::OPTIONS) || m.as_str() == "QUERY"
}
fn ct_eq(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
let mut diff = 0u8;
for (x, y) in a.iter().zip(b) {
diff |= x ^ y;
}
diff == 0
}
pub(crate) fn token_matches(expected: &str, submitted: Option<&str>) -> bool {
match submitted {
Some(t) if !expected.is_empty() => ct_eq(expected.as_bytes(), t.as_bytes()),
_ => false,
}
}
pub(crate) fn origin_ok(headers: &HeaderMap, expected_origin: &str) -> bool {
if let Some(site) = headers.get("sec-fetch-site").and_then(|v| v.to_str().ok()) {
return matches!(site, "same-origin" | "none");
}
let origin = match headers.get("origin").and_then(|v| v.to_str().ok()) {
Some(o) => o.trim_end_matches('/'),
None => return false,
};
if !expected_origin.is_empty() {
return origin == expected_origin.trim_end_matches('/');
}
match headers.get("host").and_then(|v| v.to_str().ok()) {
Some(host) => origin == format!("http://{host}") || origin == format!("https://{host}"),
None => false,
}
}
pub(crate) fn submitted_token(headers: &HeaderMap, body: &[u8]) -> Option<String> {
if let Some(token) = header_token(headers) {
return Some(token);
}
form_urlencoded::parse(body)
.find(|(k, _)| k == CSRF_FIELD)
.map(|(_, v)| v.into_owned())
}
pub(crate) fn header_token(headers: &HeaderMap) -> Option<String> {
headers
.get(CSRF_HEADER)
.and_then(|v| v.to_str().ok())
.map(str::to_string)
}
pub(crate) fn query_token(query: Option<&str>) -> Option<String> {
form_urlencoded::parse(query.unwrap_or_default().as_bytes())
.find(|(k, _)| k == CSRF_FIELD)
.map(|(_, v)| v.into_owned())
}
#[derive(Clone, Default)]
pub(crate) struct CsrfPending(std::sync::Arc<std::sync::atomic::AtomicBool>);
impl CsrfPending {
pub(crate) fn satisfy(&self) {
self.0.store(true, std::sync::atomic::Ordering::Release);
}
pub(crate) fn was_satisfied(&self) -> bool {
self.0.load(std::sync::atomic::Ordering::Acquire)
}
}
pub(crate) fn is_multipart(headers: &HeaderMap) -> bool {
headers
.get(axum::http::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.is_some_and(|ct| {
ct.trim_start()
.to_ascii_lowercase()
.starts_with("multipart/form-data")
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn corrupt_blob_degrades_to_empty_session() {
assert!(SessionData::parse(Some("not json")).csrf.is_empty());
assert!(SessionData::parse(Some(&"x".repeat(MAX_BLOB + 1)))
.csrf
.is_empty());
assert!(SessionData::parse(None).flash.is_empty());
}
#[test]
fn a_fresh_session_mints_a_token_and_is_dirty() {
let h = SessionHandle::from_blob(None);
assert_eq!(h.csrf_token().len(), 64);
assert!(h.dirty_blob().is_some());
}
#[test]
fn an_existing_token_is_kept_and_clean() {
let blob = format!(r#"{{"v":1,"csrf":"{}"}}"#, "a".repeat(64));
let h = SessionHandle::from_blob(Some(&blob));
assert_eq!(h.csrf_token(), "a".repeat(64));
assert!(h.dirty_blob().is_none());
}
#[test]
fn flash_round_trips_and_marks_dirty() {
let h = SessionHandle::from_blob(Some(&format!(r#"{{"csrf":"{}"}}"#, "a".repeat(64))));
h.push_flash(FlashLevel::Success, Text::dynamic("Saved"));
assert!(h.dirty_blob().is_some());
let taken = h.take_flash();
assert_eq!(taken.len(), 1);
assert_eq!(taken[0].text.source(), "Saved");
assert!(h.take_flash().is_empty());
}
#[test]
fn token_match_is_exact_and_rejects_missing() {
assert!(token_matches("abc", Some("abc")));
assert!(!token_matches("abc", Some("abd")));
assert!(!token_matches("abc", None));
assert!(!token_matches("", Some("")));
}
#[test]
fn safe_methods_are_exempt() {
assert!(is_safe_method(&Method::GET));
assert!(is_safe_method(&Method::from_bytes(b"QUERY").unwrap()));
assert!(!is_safe_method(&Method::POST));
assert!(!is_safe_method(&Method::DELETE));
}
#[test]
fn an_upload_carries_its_token_beside_the_body() {
let mut h = HeaderMap::new();
h.insert(
"content-type",
"multipart/form-data; boundary=x".parse().unwrap(),
);
assert!(
is_multipart(&h),
"an upload must be recognised before buffering"
);
let mut with_header = h.clone();
with_header.insert(CSRF_HEADER, "tok".parse().unwrap());
assert_eq!(header_token(&with_header).as_deref(), Some("tok"));
assert_eq!(
query_token(Some("_csrf=tok&other=1")).as_deref(),
Some("tok")
);
assert_eq!(query_token(Some("other=1")), None);
assert_eq!(query_token(None), None);
}
#[test]
fn only_a_multipart_body_skips_buffering() {
for (value, expected) in [
("multipart/form-data; boundary=abc", true),
("MULTIPART/FORM-DATA; boundary=abc", true),
("application/x-www-form-urlencoded", false),
("application/json", false),
("text/plain", false),
] {
let mut h = HeaderMap::new();
h.insert("content-type", value.parse().unwrap());
assert_eq!(is_multipart(&h), expected, "{value}");
}
assert!(
!is_multipart(&HeaderMap::new()),
"no content type is not an upload"
);
}
#[test]
fn origin_gate_accepts_same_origin_and_rejects_cross() {
let mut h = HeaderMap::new();
h.insert("sec-fetch-site", "same-origin".parse().unwrap());
assert!(origin_ok(&h, "https://acme.test"));
let mut h = HeaderMap::new();
h.insert("sec-fetch-site", "cross-site".parse().unwrap());
assert!(!origin_ok(&h, "https://acme.test"));
let mut h = HeaderMap::new();
h.insert("origin", "https://acme.test".parse().unwrap());
assert!(origin_ok(&h, "https://acme.test"));
assert!(!origin_ok(&h, "https://evil.test"));
assert!(!origin_ok(&HeaderMap::new(), "https://acme.test"));
}
#[test]
fn submitted_token_reads_header_then_field() {
let mut h = HeaderMap::new();
h.insert(CSRF_HEADER, "from-header".parse().unwrap());
assert_eq!(submitted_token(&h, b"").as_deref(), Some("from-header"));
let body = b"username=root&_csrf=from-field&x=1";
assert_eq!(
submitted_token(&HeaderMap::new(), body).as_deref(),
Some("from-field")
);
assert_eq!(submitted_token(&HeaderMap::new(), b"a=b"), None);
}
}