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}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
40#[serde(rename_all = "snake_case")]
41pub enum FlashLevel {
42 Success,
43 Error,
44 Info,
45}
46
47impl FlashLevel {
48 pub(crate) fn class(self) -> &'static str {
50 match self {
51 FlashLevel::Success => "is-success",
52 FlashLevel::Error => "is-error",
53 FlashLevel::Info => "is-info",
54 }
55 }
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
60struct SessionData {
61 #[serde(default = "default_version")]
63 v: u8,
64 #[serde(default)]
66 csrf: String,
67 #[serde(default)]
69 flash: Vec<Flash>,
70}
71
72fn default_version() -> u8 {
73 VERSION
74}
75
76impl Default for SessionData {
77 fn default() -> Self {
78 Self {
79 v: VERSION,
80 csrf: String::new(),
81 flash: Vec::new(),
82 }
83 }
84}
85
86impl SessionData {
87 fn parse(raw: Option<&str>) -> Self {
91 match raw {
92 Some(s) if s.len() <= MAX_BLOB => serde_json::from_str(s).unwrap_or_default(),
93 _ => Self::default(),
94 }
95 }
96}
97
98struct Inner {
99 data: SessionData,
100 dirty: bool,
101}
102
103#[derive(Clone)]
108pub struct SessionHandle {
109 inner: Arc<Mutex<Inner>>,
110}
111
112impl SessionHandle {
113 pub(crate) fn from_blob(raw: Option<&str>) -> Self {
117 let mut data = SessionData::parse(raw);
118 let dirty = data.csrf.is_empty();
119 if dirty {
120 data.csrf = generate_token();
121 }
122 Self {
123 inner: Arc::new(Mutex::new(Inner { data, dirty })),
124 }
125 }
126
127 pub(crate) fn csrf_token(&self) -> String {
129 self.inner.lock().unwrap().data.csrf.clone()
130 }
131
132 pub fn push_flash(&self, level: FlashLevel, text: Text) {
135 let mut g = self.inner.lock().unwrap();
136 g.data.flash.push(Flash { level, text });
137 g.dirty = true;
138 }
139
140 pub(crate) fn take_flash(&self) -> Vec<Flash> {
142 let mut g = self.inner.lock().unwrap();
143 if g.data.flash.is_empty() {
144 return Vec::new();
145 }
146 g.dirty = true;
147 std::mem::take(&mut g.data.flash)
148 }
149
150 pub(crate) fn dirty_blob(&self) -> Option<String> {
152 let g = self.inner.lock().unwrap();
153 g.dirty
154 .then(|| serde_json::to_string(&g.data).unwrap_or_default())
155 }
156}
157
158pub(crate) fn generate_token() -> String {
160 let mut bytes = [0u8; 32];
161 rand::rngs::OsRng.fill_bytes(&mut bytes);
162 let mut out = String::with_capacity(64);
163 for b in bytes {
164 let _ = write!(out, "{b:02x}");
165 }
166 out
167}
168
169pub(crate) fn is_safe_method(m: &Method) -> bool {
172 matches!(*m, Method::GET | Method::HEAD | Method::OPTIONS) || m.as_str() == "QUERY"
173}
174
175fn ct_eq(a: &[u8], b: &[u8]) -> bool {
177 if a.len() != b.len() {
178 return false;
179 }
180 let mut diff = 0u8;
181 for (x, y) in a.iter().zip(b) {
182 diff |= x ^ y;
183 }
184 diff == 0
185}
186
187pub(crate) fn token_matches(expected: &str, submitted: Option<&str>) -> bool {
190 match submitted {
191 Some(t) if !expected.is_empty() => ct_eq(expected.as_bytes(), t.as_bytes()),
192 _ => false,
193 }
194}
195
196pub(crate) fn origin_ok(headers: &HeaderMap, expected_origin: &str) -> bool {
203 if let Some(site) = headers.get("sec-fetch-site").and_then(|v| v.to_str().ok()) {
204 return matches!(site, "same-origin" | "none");
205 }
206 let origin = match headers.get("origin").and_then(|v| v.to_str().ok()) {
207 Some(o) => o.trim_end_matches('/'),
208 None => return false,
209 };
210 if !expected_origin.is_empty() {
211 return origin == expected_origin.trim_end_matches('/');
212 }
213 match headers.get("host").and_then(|v| v.to_str().ok()) {
214 Some(host) => origin == format!("http://{host}") || origin == format!("https://{host}"),
215 None => false,
216 }
217}
218
219pub(crate) fn submitted_token(headers: &HeaderMap, body: &[u8]) -> Option<String> {
222 if let Some(token) = header_token(headers) {
223 return Some(token);
224 }
225 form_urlencoded::parse(body)
226 .find(|(k, _)| k == CSRF_FIELD)
227 .map(|(_, v)| v.into_owned())
228}
229
230pub(crate) fn header_token(headers: &HeaderMap) -> Option<String> {
232 headers
233 .get(CSRF_HEADER)
234 .and_then(|v| v.to_str().ok())
235 .map(str::to_string)
236}
237
238pub(crate) fn query_token(query: Option<&str>) -> Option<String> {
250 form_urlencoded::parse(query.unwrap_or_default().as_bytes())
251 .find(|(k, _)| k == CSRF_FIELD)
252 .map(|(_, v)| v.into_owned())
253}
254
255#[derive(Clone, Default)]
263pub(crate) struct CsrfPending(std::sync::Arc<std::sync::atomic::AtomicBool>);
264
265impl CsrfPending {
266 pub(crate) fn satisfy(&self) {
268 self.0.store(true, std::sync::atomic::Ordering::Release);
269 }
270
271 pub(crate) fn was_satisfied(&self) -> bool {
272 self.0.load(std::sync::atomic::Ordering::Acquire)
273 }
274}
275
276pub(crate) fn is_multipart(headers: &HeaderMap) -> bool {
278 headers
279 .get(axum::http::header::CONTENT_TYPE)
280 .and_then(|v| v.to_str().ok())
281 .is_some_and(|ct| {
282 ct.trim_start()
283 .to_ascii_lowercase()
284 .starts_with("multipart/form-data")
285 })
286}
287
288#[cfg(test)]
289mod tests {
290 use super::*;
291
292 #[test]
293 fn corrupt_blob_degrades_to_empty_session() {
294 assert!(SessionData::parse(Some("not json")).csrf.is_empty());
295 assert!(SessionData::parse(Some(&"x".repeat(MAX_BLOB + 1)))
296 .csrf
297 .is_empty());
298 assert!(SessionData::parse(None).flash.is_empty());
299 }
300
301 #[test]
302 fn a_fresh_session_mints_a_token_and_is_dirty() {
303 let h = SessionHandle::from_blob(None);
304 assert_eq!(h.csrf_token().len(), 64);
305 assert!(h.dirty_blob().is_some());
306 }
307
308 #[test]
309 fn an_existing_token_is_kept_and_clean() {
310 let blob = format!(r#"{{"v":1,"csrf":"{}"}}"#, "a".repeat(64));
311 let h = SessionHandle::from_blob(Some(&blob));
312 assert_eq!(h.csrf_token(), "a".repeat(64));
313 assert!(h.dirty_blob().is_none());
314 }
315
316 #[test]
317 fn flash_round_trips_and_marks_dirty() {
318 let h = SessionHandle::from_blob(Some(&format!(r#"{{"csrf":"{}"}}"#, "a".repeat(64))));
319 h.push_flash(FlashLevel::Success, Text::dynamic("Saved"));
320 assert!(h.dirty_blob().is_some());
321 let taken = h.take_flash();
322 assert_eq!(taken.len(), 1);
323 assert_eq!(taken[0].text.source(), "Saved");
324 assert!(h.take_flash().is_empty());
325 }
326
327 #[test]
328 fn token_match_is_exact_and_rejects_missing() {
329 assert!(token_matches("abc", Some("abc")));
330 assert!(!token_matches("abc", Some("abd")));
331 assert!(!token_matches("abc", None));
332 assert!(!token_matches("", Some("")));
333 }
334
335 #[test]
336 fn safe_methods_are_exempt() {
337 assert!(is_safe_method(&Method::GET));
338 assert!(is_safe_method(&Method::from_bytes(b"QUERY").unwrap()));
339 assert!(!is_safe_method(&Method::POST));
340 assert!(!is_safe_method(&Method::DELETE));
341 }
342
343 #[test]
346 fn an_upload_carries_its_token_beside_the_body() {
347 let mut h = HeaderMap::new();
348 h.insert(
349 "content-type",
350 "multipart/form-data; boundary=x".parse().unwrap(),
351 );
352 assert!(
353 is_multipart(&h),
354 "an upload must be recognised before buffering"
355 );
356
357 let mut with_header = h.clone();
359 with_header.insert(CSRF_HEADER, "tok".parse().unwrap());
360 assert_eq!(header_token(&with_header).as_deref(), Some("tok"));
361
362 assert_eq!(
364 query_token(Some("_csrf=tok&other=1")).as_deref(),
365 Some("tok")
366 );
367 assert_eq!(query_token(Some("other=1")), None);
368 assert_eq!(query_token(None), None);
369 }
370
371 #[test]
373 fn only_a_multipart_body_skips_buffering() {
374 for (value, expected) in [
375 ("multipart/form-data; boundary=abc", true),
376 ("MULTIPART/FORM-DATA; boundary=abc", true),
377 ("application/x-www-form-urlencoded", false),
378 ("application/json", false),
379 ("text/plain", false),
380 ] {
381 let mut h = HeaderMap::new();
382 h.insert("content-type", value.parse().unwrap());
383 assert_eq!(is_multipart(&h), expected, "{value}");
384 }
385 assert!(
386 !is_multipart(&HeaderMap::new()),
387 "no content type is not an upload"
388 );
389 }
390
391 #[test]
392 fn origin_gate_accepts_same_origin_and_rejects_cross() {
393 let mut h = HeaderMap::new();
394 h.insert("sec-fetch-site", "same-origin".parse().unwrap());
395 assert!(origin_ok(&h, "https://acme.test"));
396
397 let mut h = HeaderMap::new();
398 h.insert("sec-fetch-site", "cross-site".parse().unwrap());
399 assert!(!origin_ok(&h, "https://acme.test"));
400
401 let mut h = HeaderMap::new();
402 h.insert("origin", "https://acme.test".parse().unwrap());
403 assert!(origin_ok(&h, "https://acme.test"));
404 assert!(!origin_ok(&h, "https://evil.test"));
405
406 assert!(!origin_ok(&HeaderMap::new(), "https://acme.test"));
408 }
409
410 #[test]
411 fn submitted_token_reads_header_then_field() {
412 let mut h = HeaderMap::new();
413 h.insert(CSRF_HEADER, "from-header".parse().unwrap());
414 assert_eq!(submitted_token(&h, b"").as_deref(), Some("from-header"));
415
416 let body = b"username=root&_csrf=from-field&x=1";
417 assert_eq!(
418 submitted_token(&HeaderMap::new(), body).as_deref(),
419 Some("from-field")
420 );
421 assert_eq!(submitted_token(&HeaderMap::new(), b"a=b"), None);
422 }
423}