1use std::fmt::Write as _;
14use std::sync::{Arc, Mutex};
15
16use axum::http::{HeaderMap, Method};
17use laterite_core::Text;
18use rand::RngCore;
19use serde::{Deserialize, Serialize};
20
21const VERSION: u8 = 1;
23const MAX_BLOB: usize = 4096;
26const CSRF_HEADER: &str = "x-csrf-token";
27pub(crate) const CSRF_FIELD: &str = "_csrf";
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
33pub(crate) struct Flash {
34 pub(crate) level: FlashLevel,
35 pub(crate) text: Text,
36 #[serde(default)]
38 pub(crate) sticky: bool,
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
43#[serde(rename_all = "snake_case")]
44pub enum FlashLevel {
45 Success,
46 Error,
47 Info,
48}
49
50impl FlashLevel {
51 pub(crate) fn class(self) -> &'static str {
53 match self {
54 FlashLevel::Success => "is-success",
55 FlashLevel::Error => "is-error",
56 FlashLevel::Info => "is-info",
57 }
58 }
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63struct SessionData {
64 #[serde(default = "default_version")]
66 v: u8,
67 #[serde(default)]
69 csrf: String,
70 #[serde(default)]
72 flash: Vec<Flash>,
73}
74
75fn default_version() -> u8 {
76 VERSION
77}
78
79impl Default for SessionData {
80 fn default() -> Self {
81 Self {
82 v: VERSION,
83 csrf: String::new(),
84 flash: Vec::new(),
85 }
86 }
87}
88
89impl SessionData {
90 fn parse(raw: Option<&str>) -> Self {
94 match raw {
95 Some(s) if s.len() <= MAX_BLOB => serde_json::from_str(s).unwrap_or_default(),
96 _ => Self::default(),
97 }
98 }
99}
100
101struct Inner {
102 data: SessionData,
103 dirty: bool,
104}
105
106#[derive(Clone)]
111pub struct SessionHandle {
112 inner: Arc<Mutex<Inner>>,
113}
114
115impl SessionHandle {
116 pub(crate) fn from_blob(raw: Option<&str>) -> Self {
120 let mut data = SessionData::parse(raw);
121 let dirty = data.csrf.is_empty();
122 if dirty {
123 data.csrf = generate_token();
124 }
125 Self {
126 inner: Arc::new(Mutex::new(Inner { data, dirty })),
127 }
128 }
129
130 pub(crate) fn csrf_token(&self) -> String {
132 self.inner.lock().unwrap().data.csrf.clone()
133 }
134
135 pub fn push_flash(&self, level: FlashLevel, text: Text) {
138 let mut g = self.inner.lock().unwrap();
139 g.data.flash.push(Flash {
140 level,
141 text,
142 sticky: false,
143 });
144 g.dirty = true;
145 }
146
147 pub fn push_flash_sticky(&self, level: FlashLevel, text: Text) {
150 let mut g = self.inner.lock().unwrap();
151 g.data.flash.push(Flash {
152 level,
153 text,
154 sticky: true,
155 });
156 g.dirty = true;
157 }
158
159 pub(crate) fn take_flash(&self) -> Vec<Flash> {
161 let mut g = self.inner.lock().unwrap();
162 if g.data.flash.is_empty() {
163 return Vec::new();
164 }
165 g.dirty = true;
166 std::mem::take(&mut g.data.flash)
167 }
168
169 pub(crate) fn dirty_blob(&self) -> Option<String> {
171 let g = self.inner.lock().unwrap();
172 g.dirty
173 .then(|| serde_json::to_string(&g.data).unwrap_or_default())
174 }
175}
176
177pub(crate) fn generate_token() -> String {
179 let mut bytes = [0u8; 32];
180 rand::rngs::OsRng.fill_bytes(&mut bytes);
181 let mut out = String::with_capacity(64);
182 for b in bytes {
183 let _ = write!(out, "{b:02x}");
184 }
185 out
186}
187
188pub(crate) fn is_safe_method(m: &Method) -> bool {
191 matches!(*m, Method::GET | Method::HEAD | Method::OPTIONS) || m.as_str() == "QUERY"
192}
193
194fn ct_eq(a: &[u8], b: &[u8]) -> bool {
196 if a.len() != b.len() {
197 return false;
198 }
199 let mut diff = 0u8;
200 for (x, y) in a.iter().zip(b) {
201 diff |= x ^ y;
202 }
203 diff == 0
204}
205
206pub(crate) fn token_matches(expected: &str, submitted: Option<&str>) -> bool {
209 match submitted {
210 Some(t) if !expected.is_empty() => ct_eq(expected.as_bytes(), t.as_bytes()),
211 _ => false,
212 }
213}
214
215pub(crate) fn origin_ok(headers: &HeaderMap, expected_origin: &str) -> bool {
222 if let Some(site) = headers.get("sec-fetch-site").and_then(|v| v.to_str().ok()) {
223 return matches!(site, "same-origin" | "none");
224 }
225 let origin = match headers.get("origin").and_then(|v| v.to_str().ok()) {
226 Some(o) => o.trim_end_matches('/'),
227 None => return false,
228 };
229 if !expected_origin.is_empty() {
230 return origin == expected_origin.trim_end_matches('/');
231 }
232 match headers.get("host").and_then(|v| v.to_str().ok()) {
233 Some(host) => origin == format!("http://{host}") || origin == format!("https://{host}"),
234 None => false,
235 }
236}
237
238pub(crate) fn submitted_token(headers: &HeaderMap, body: &[u8]) -> Option<String> {
241 if let Some(token) = header_token(headers) {
242 return Some(token);
243 }
244 form_urlencoded::parse(body)
245 .find(|(k, _)| k == CSRF_FIELD)
246 .map(|(_, v)| v.into_owned())
247}
248
249pub(crate) fn header_token(headers: &HeaderMap) -> Option<String> {
251 headers
252 .get(CSRF_HEADER)
253 .and_then(|v| v.to_str().ok())
254 .map(str::to_string)
255}
256
257pub(crate) fn query_token(query: Option<&str>) -> Option<String> {
269 form_urlencoded::parse(query.unwrap_or_default().as_bytes())
270 .find(|(k, _)| k == CSRF_FIELD)
271 .map(|(_, v)| v.into_owned())
272}
273
274#[derive(Clone, Default)]
282pub(crate) struct CsrfPending(std::sync::Arc<std::sync::atomic::AtomicBool>);
283
284impl CsrfPending {
285 pub(crate) fn satisfy(&self) {
287 self.0.store(true, std::sync::atomic::Ordering::Release);
288 }
289
290 pub(crate) fn was_satisfied(&self) -> bool {
291 self.0.load(std::sync::atomic::Ordering::Acquire)
292 }
293}
294
295pub(crate) fn is_multipart(headers: &HeaderMap) -> bool {
297 headers
298 .get(axum::http::header::CONTENT_TYPE)
299 .and_then(|v| v.to_str().ok())
300 .is_some_and(|ct| {
301 ct.trim_start()
302 .to_ascii_lowercase()
303 .starts_with("multipart/form-data")
304 })
305}
306
307#[cfg(test)]
308mod tests {
309 use super::*;
310
311 #[test]
312 fn corrupt_blob_degrades_to_empty_session() {
313 assert!(SessionData::parse(Some("not json")).csrf.is_empty());
314 assert!(SessionData::parse(Some(&"x".repeat(MAX_BLOB + 1)))
315 .csrf
316 .is_empty());
317 assert!(SessionData::parse(None).flash.is_empty());
318 }
319
320 #[test]
321 fn a_fresh_session_mints_a_token_and_is_dirty() {
322 let h = SessionHandle::from_blob(None);
323 assert_eq!(h.csrf_token().len(), 64);
324 assert!(h.dirty_blob().is_some());
325 }
326
327 #[test]
328 fn an_existing_token_is_kept_and_clean() {
329 let blob = format!(r#"{{"v":1,"csrf":"{}"}}"#, "a".repeat(64));
330 let h = SessionHandle::from_blob(Some(&blob));
331 assert_eq!(h.csrf_token(), "a".repeat(64));
332 assert!(h.dirty_blob().is_none());
333 }
334
335 #[test]
336 fn flash_round_trips_and_marks_dirty() {
337 let h = SessionHandle::from_blob(Some(&format!(r#"{{"csrf":"{}"}}"#, "a".repeat(64))));
338 h.push_flash(FlashLevel::Success, Text::dynamic("Saved"));
339 assert!(h.dirty_blob().is_some());
340 let taken = h.take_flash();
341 assert_eq!(taken.len(), 1);
342 assert_eq!(taken[0].text.source(), "Saved");
343 assert!(h.take_flash().is_empty());
344 }
345
346 #[test]
347 fn token_match_is_exact_and_rejects_missing() {
348 assert!(token_matches("abc", Some("abc")));
349 assert!(!token_matches("abc", Some("abd")));
350 assert!(!token_matches("abc", None));
351 assert!(!token_matches("", Some("")));
352 }
353
354 #[test]
355 fn safe_methods_are_exempt() {
356 assert!(is_safe_method(&Method::GET));
357 assert!(is_safe_method(&Method::from_bytes(b"QUERY").unwrap()));
358 assert!(!is_safe_method(&Method::POST));
359 assert!(!is_safe_method(&Method::DELETE));
360 }
361
362 #[test]
365 fn an_upload_carries_its_token_beside_the_body() {
366 let mut h = HeaderMap::new();
367 h.insert(
368 "content-type",
369 "multipart/form-data; boundary=x".parse().unwrap(),
370 );
371 assert!(
372 is_multipart(&h),
373 "an upload must be recognised before buffering"
374 );
375
376 let mut with_header = h.clone();
378 with_header.insert(CSRF_HEADER, "tok".parse().unwrap());
379 assert_eq!(header_token(&with_header).as_deref(), Some("tok"));
380
381 assert_eq!(
383 query_token(Some("_csrf=tok&other=1")).as_deref(),
384 Some("tok")
385 );
386 assert_eq!(query_token(Some("other=1")), None);
387 assert_eq!(query_token(None), None);
388 }
389
390 #[test]
392 fn only_a_multipart_body_skips_buffering() {
393 for (value, expected) in [
394 ("multipart/form-data; boundary=abc", true),
395 ("MULTIPART/FORM-DATA; boundary=abc", true),
396 ("application/x-www-form-urlencoded", false),
397 ("application/json", false),
398 ("text/plain", false),
399 ] {
400 let mut h = HeaderMap::new();
401 h.insert("content-type", value.parse().unwrap());
402 assert_eq!(is_multipart(&h), expected, "{value}");
403 }
404 assert!(
405 !is_multipart(&HeaderMap::new()),
406 "no content type is not an upload"
407 );
408 }
409
410 #[test]
411 fn origin_gate_accepts_same_origin_and_rejects_cross() {
412 let mut h = HeaderMap::new();
413 h.insert("sec-fetch-site", "same-origin".parse().unwrap());
414 assert!(origin_ok(&h, "https://acme.test"));
415
416 let mut h = HeaderMap::new();
417 h.insert("sec-fetch-site", "cross-site".parse().unwrap());
418 assert!(!origin_ok(&h, "https://acme.test"));
419
420 let mut h = HeaderMap::new();
421 h.insert("origin", "https://acme.test".parse().unwrap());
422 assert!(origin_ok(&h, "https://acme.test"));
423 assert!(!origin_ok(&h, "https://evil.test"));
424
425 assert!(!origin_ok(&HeaderMap::new(), "https://acme.test"));
427 }
428
429 #[test]
430 fn submitted_token_reads_header_then_field() {
431 let mut h = HeaderMap::new();
432 h.insert(CSRF_HEADER, "from-header".parse().unwrap());
433 assert_eq!(submitted_token(&h, b"").as_deref(), Some("from-header"));
434
435 let body = b"username=root&_csrf=from-field&x=1";
436 assert_eq!(
437 submitted_token(&HeaderMap::new(), body).as_deref(),
438 Some("from-field")
439 );
440 assert_eq!(submitted_token(&HeaderMap::new(), b"a=b"), None);
441 }
442}