Skip to main content

dioxus_cloudflare/
session.rs

1//! Session middleware backed by Workers KV or D1.
2//!
3//! Provides ergonomic session management for `#[server]` functions with
4//! automatic cookie handling and backend persistence.
5//!
6//! ## Usage
7//!
8//! Configure sessions on the [`Handler`](crate::Handler) builder:
9//!
10//! ```rust,ignore
11//! Handler::new()
12//!     .session(SessionConfig::kv("SESSIONS"))
13//!     .handle(req, env)
14//!     .await
15//! ```
16//!
17//! Then access the session from any server function:
18//!
19//! ```rust,ignore
20//! #[server]
21//! pub async fn login(user: String) -> Result<(), ServerFnError> {
22//!     let session = cf::session().await?;
23//!     session.set("user_id", &user)?;
24//!     Ok(())
25//! }
26//! ```
27//!
28//! ## Design
29//!
30//! - `cf::session()` is **async** — loads data from KV/D1 on first call, cached after
31//! - `Session` methods are **sync** — operate on in-memory cache
32//! - Session data is flushed to the backend before the response is finalized
33//! - Thread-local state follows the same `RefCell<Option<T>>` pattern as context/cookies
34
35use std::cell::RefCell;
36use std::collections::HashMap;
37
38use serde::de::DeserializeOwned;
39use serde::Serialize;
40use wasm_bindgen::prelude::*;
41
42use crate::context::push_cookie;
43use crate::error::{CfError, CfResultExt};
44
45// ---------------------------------------------------------------------------
46// UUID generation via Workers runtime
47// ---------------------------------------------------------------------------
48
49#[wasm_bindgen]
50extern "C" {
51    #[wasm_bindgen(js_namespace = crypto, js_name = randomUUID)]
52    fn random_uuid() -> String;
53}
54
55// ---------------------------------------------------------------------------
56// Configuration types
57// ---------------------------------------------------------------------------
58
59/// Session backend configuration.
60///
61/// Created via [`SessionConfig::kv()`] or [`SessionConfig::d1()`] and passed
62/// to [`Handler::session()`](crate::Handler::session).
63#[derive(Clone)]
64pub struct SessionConfig {
65    pub(crate) backend: SessionBackend,
66    pub(crate) cookie_name: String,
67    pub(crate) max_age_secs: u64,
68}
69
70#[derive(Clone)]
71pub(crate) enum SessionBackend {
72    Kv { binding: String },
73    D1 { binding: String, table: String },
74}
75
76impl SessionConfig {
77    /// Create a KV-backed session configuration.
78    ///
79    /// `binding` is the KV namespace binding name from `wrangler.toml`.
80    /// Sessions are automatically expired using KV's `expiration_ttl`.
81    ///
82    /// # Example
83    ///
84    /// ```rust,ignore
85    /// Handler::new()
86    ///     .session(SessionConfig::kv("SESSIONS"))
87    ///     .handle(req, env)
88    ///     .await
89    /// ```
90    #[must_use]
91    pub fn kv(binding: &str) -> Self {
92        Self {
93            backend: SessionBackend::Kv {
94                binding: binding.to_owned(),
95            },
96            cookie_name: "__session".to_owned(),
97            max_age_secs: 86400,
98        }
99    }
100
101    /// Create a D1-backed session configuration.
102    ///
103    /// `binding` is the D1 database binding name and `table` is the table name.
104    /// The table must exist with schema:
105    ///
106    /// ```sql
107    /// CREATE TABLE sessions (
108    ///     id TEXT PRIMARY KEY,
109    ///     data TEXT NOT NULL,
110    ///     expires_at INTEGER NOT NULL
111    /// );
112    /// ```
113    ///
114    /// Expired sessions are filtered at read time via `expires_at > unixepoch()`.
115    ///
116    /// # Example
117    ///
118    /// ```rust,ignore
119    /// Handler::new()
120    ///     .session(SessionConfig::d1("DB", "sessions"))
121    ///     .handle(req, env)
122    ///     .await
123    /// ```
124    #[must_use]
125    pub fn d1(binding: &str, table: &str) -> Self {
126        Self {
127            backend: SessionBackend::D1 {
128                binding: binding.to_owned(),
129                table: table.to_owned(),
130            },
131            cookie_name: "__session".to_owned(),
132            max_age_secs: 86400,
133        }
134    }
135
136    /// Set the cookie name. Default: `"__session"`.
137    #[must_use]
138    pub fn cookie_name(mut self, name: &str) -> Self {
139        self.cookie_name = name.to_owned();
140        self
141    }
142
143    /// Set the session max age in seconds. Default: `86400` (24 hours).
144    ///
145    /// This controls both the cookie `Max-Age` and the backend expiry
146    /// (KV `expiration_ttl` or D1 `expires_at`).
147    #[must_use]
148    pub fn max_age(mut self, secs: u64) -> Self {
149        self.max_age_secs = secs;
150        self
151    }
152}
153
154// ---------------------------------------------------------------------------
155// Thread-local state
156// ---------------------------------------------------------------------------
157
158struct SessionState {
159    id: String,
160    data: HashMap<String, serde_json::Value>,
161    dirty: bool,
162    destroyed: bool,
163    is_new: bool,
164}
165
166thread_local! {
167    static SESSION_CONFIG: RefCell<Option<SessionConfig>> = const { RefCell::new(None) };
168    static SESSION_STATE: RefCell<Option<SessionState>> = const { RefCell::new(None) };
169}
170
171// ---------------------------------------------------------------------------
172// Public API
173// ---------------------------------------------------------------------------
174
175/// Zero-sized session handle — all operations access thread-local state.
176///
177/// Obtained via [`cf::session()`](session). Methods are sync because they
178/// operate on the in-memory cache loaded by the async `session()` call.
179pub struct Session;
180
181impl Session {
182    /// Get a typed value from the session.
183    ///
184    /// Returns `Ok(None)` if the key is not present.
185    ///
186    /// # Errors
187    ///
188    /// Returns [`CfError`] if deserialization fails.
189    pub fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, CfError> {
190        SESSION_STATE.with(|cell| {
191            let borrow = cell.borrow();
192            let state = borrow
193                .as_ref()
194                .ok_or_else(|| CfError(worker::Error::RustError("session not loaded".into())))?;
195
196            match state.data.get(key) {
197                Some(value) => {
198                    let typed: T = serde_json::from_value(value.clone()).cf()?;
199                    Ok(Some(typed))
200                }
201                None => Ok(None),
202            }
203        })
204    }
205
206    /// Set a value in the session.
207    ///
208    /// The value is stored in the in-memory cache and flushed to the backend
209    /// when the response is finalized.
210    ///
211    /// # Errors
212    ///
213    /// Returns [`CfError`] if serialization fails.
214    pub fn set<T: Serialize>(&self, key: &str, value: &T) -> Result<(), CfError> {
215        SESSION_STATE.with(|cell| {
216            let mut borrow = cell.borrow_mut();
217            let state = borrow
218                .as_mut()
219                .ok_or_else(|| CfError(worker::Error::RustError("session not loaded".into())))?;
220
221            let json_value = serde_json::to_value(value).cf()?;
222            state.data.insert(key.to_owned(), json_value);
223            state.dirty = true;
224            Ok(())
225        })
226    }
227
228    /// Remove a key from the session.
229    pub fn remove(&self, key: &str) {
230        SESSION_STATE.with(|cell| {
231            if let Some(ref mut state) = *cell.borrow_mut() {
232                if state.data.remove(key).is_some() {
233                    state.dirty = true;
234                }
235            }
236        });
237    }
238
239    /// Destroy the session — deletes backend data and clears the cookie.
240    pub fn destroy(&self) {
241        SESSION_STATE.with(|cell| {
242            if let Some(ref mut state) = *cell.borrow_mut() {
243                state.destroyed = true;
244            }
245        });
246    }
247
248    /// Returns `true` if the session has no data.
249    pub fn is_empty(&self) -> bool {
250        SESSION_STATE.with(|cell| {
251            cell.borrow()
252                .as_ref()
253                .is_none_or(|state| state.data.is_empty())
254        })
255    }
256
257    /// Returns `true` if this is a new session (no existing cookie was found).
258    pub fn is_new(&self) -> bool {
259        SESSION_STATE.with(|cell| {
260            cell.borrow()
261                .as_ref()
262                .is_none_or(|state| state.is_new)
263        })
264    }
265}
266
267/// Load the session for the current request.
268///
269/// Async on first call (reads from KV/D1), cached on subsequent calls.
270/// Returns a zero-sized [`Session`] handle for sync get/set/remove operations.
271///
272/// # Errors
273///
274/// Returns [`CfError`] if:
275/// - Session middleware is not configured (no `.session()` on `Handler`)
276/// - The backend read fails (KV/D1 unavailable)
277/// - Session data cannot be deserialized
278///
279/// # Example
280///
281/// ```rust,ignore
282/// use dioxus_cloudflare::prelude::*;
283///
284/// #[server]
285/// pub async fn get_user() -> Result<String, ServerFnError> {
286///     let session = cf::session().await?;
287///     let user: Option<String> = session.get("user_id")?;
288///     Ok(user.unwrap_or_else(|| "not logged in".into()))
289/// }
290/// ```
291pub async fn session() -> Result<Session, CfError> {
292    // Already loaded — return cached handle
293    let already_loaded = SESSION_STATE.with(|cell| cell.borrow().is_some());
294    if already_loaded {
295        return Ok(Session);
296    }
297
298    // Read config
299    let config = SESSION_CONFIG.with(|cell| cell.borrow().clone()).ok_or_else(|| {
300        CfError(worker::Error::RustError(
301            "cf::session() called but session middleware is not configured — \
302             add .session(SessionConfig::kv(\"SESSIONS\")) to your Handler"
303                .into(),
304        ))
305    })?;
306
307    // Check for existing session cookie
308    let existing_id = crate::cookie::cookie(&config.cookie_name)?;
309
310    let (id, data, is_new) = match existing_id {
311        Some(ref session_id) if !session_id.is_empty() => {
312            // Load from backend
313            let data = load_from_backend(&config, session_id).await?;
314            match data {
315                Some(d) => (session_id.clone(), d, false),
316                // Cookie exists but backend data is gone (expired/deleted)
317                None => (random_uuid(), HashMap::new(), true),
318            }
319        }
320        _ => {
321            // No cookie — new session
322            (random_uuid(), HashMap::new(), true)
323        }
324    };
325
326    SESSION_STATE.with(|cell| {
327        *cell.borrow_mut() = Some(SessionState {
328            id,
329            data,
330            dirty: false,
331            destroyed: false,
332            is_new,
333        });
334    });
335
336    Ok(Session)
337}
338
339// ---------------------------------------------------------------------------
340// Internal API (pub(crate))
341// ---------------------------------------------------------------------------
342
343/// Store session config for the current request. Called by `Handler::handle()`.
344pub(crate) fn set_session_config(config: SessionConfig) {
345    SESSION_CONFIG.with(|cell| *cell.borrow_mut() = Some(config));
346}
347
348/// Flush dirty session data to the backend and queue the session cookie.
349/// Called before `finalize()` in `Handler`.
350///
351/// If no session was loaded during this request, this is a no-op.
352pub(crate) async fn flush_session() -> Result<(), worker::Error> {
353    // Take the state — session is consumed at end of request
354    let state = SESSION_STATE.with(|cell| cell.borrow_mut().take());
355    let Some(state) = state else {
356        return Ok(());
357    };
358
359    let config = SESSION_CONFIG.with(|cell| cell.borrow().clone());
360    let Some(config) = config else {
361        return Ok(());
362    };
363
364    if state.destroyed {
365        // Delete from backend
366        delete_from_backend(&config, &state.id)
367            .await
368            .map_err(|e| worker::Error::RustError(format!("session delete failed: {e}")))?;
369        // Clear the cookie
370        push_cookie(format!(
371            "{}=; HttpOnly; Secure; SameSite=Strict; Path=/; Max-Age=0",
372            config.cookie_name
373        ));
374        return Ok(());
375    }
376
377    if state.dirty || state.is_new {
378        // Only write if there's data, or if dirty (even if empty, to persist removal)
379        save_to_backend(&config, &state.id, &state.data)
380            .await
381            .map_err(|e| worker::Error::RustError(format!("session save failed: {e}")))?;
382        // Set the session cookie
383        push_cookie(format!(
384            "{}={}; HttpOnly; Secure; SameSite=Strict; Path=/; Max-Age={}",
385            config.cookie_name, state.id, config.max_age_secs
386        ));
387    }
388
389    // Clear config for next request
390    SESSION_CONFIG.with(|cell| *cell.borrow_mut() = None);
391
392    Ok(())
393}
394
395// ---------------------------------------------------------------------------
396// Backend operations
397// ---------------------------------------------------------------------------
398
399async fn load_from_backend(
400    config: &SessionConfig,
401    session_id: &str,
402) -> Result<Option<HashMap<String, serde_json::Value>>, CfError> {
403    match &config.backend {
404        SessionBackend::Kv { binding } => load_from_kv(binding, session_id).await,
405        SessionBackend::D1 { binding, table } => load_from_d1(binding, table, session_id).await,
406    }
407}
408
409async fn save_to_backend(
410    config: &SessionConfig,
411    session_id: &str,
412    data: &HashMap<String, serde_json::Value>,
413) -> Result<(), CfError> {
414    match &config.backend {
415        SessionBackend::Kv { binding } => {
416            save_to_kv(binding, session_id, data, config.max_age_secs).await
417        }
418        SessionBackend::D1 { binding, table } => {
419            save_to_d1(binding, table, session_id, data, config.max_age_secs).await
420        }
421    }
422}
423
424async fn delete_from_backend(config: &SessionConfig, session_id: &str) -> Result<(), CfError> {
425    match &config.backend {
426        SessionBackend::Kv { binding } => delete_from_kv(binding, session_id).await,
427        SessionBackend::D1 { binding, table } => delete_from_d1(binding, table, session_id).await,
428    }
429}
430
431// ---------------------------------------------------------------------------
432// KV backend
433// ---------------------------------------------------------------------------
434
435fn session_key(session_id: &str) -> String {
436    format!("session:{session_id}")
437}
438
439async fn load_from_kv(
440    binding: &str,
441    session_id: &str,
442) -> Result<Option<HashMap<String, serde_json::Value>>, CfError> {
443    let kv = crate::bindings::kv(binding)?;
444    let key = session_key(session_id);
445
446    let text = kv.get(&key).text().await.cf()?;
447
448    match text {
449        Some(json) => {
450            let data: HashMap<String, serde_json::Value> = serde_json::from_str(&json).cf()?;
451            Ok(Some(data))
452        }
453        None => Ok(None),
454    }
455}
456
457async fn save_to_kv(
458    binding: &str,
459    session_id: &str,
460    data: &HashMap<String, serde_json::Value>,
461    max_age_secs: u64,
462) -> Result<(), CfError> {
463    let kv = crate::bindings::kv(binding)?;
464    let key = session_key(session_id);
465    let json = serde_json::to_string(data).cf()?;
466
467    kv.put(&key, &json)
468        .cf()?
469        .expiration_ttl(max_age_secs)
470        .execute()
471        .await
472        .cf()?;
473
474    Ok(())
475}
476
477async fn delete_from_kv(binding: &str, session_id: &str) -> Result<(), CfError> {
478    let kv = crate::bindings::kv(binding)?;
479    let key = session_key(session_id);
480
481    kv.delete(&key).await.cf()?;
482
483    Ok(())
484}
485
486// ---------------------------------------------------------------------------
487// D1 backend
488// ---------------------------------------------------------------------------
489
490async fn load_from_d1(
491    binding: &str,
492    table: &str,
493    session_id: &str,
494) -> Result<Option<HashMap<String, serde_json::Value>>, CfError> {
495    let db = crate::bindings::d1(binding)?;
496    let key = session_key(session_id);
497
498    let result = db
499        .prepare(format!(
500            "SELECT data FROM {table} WHERE id = ? AND expires_at > unixepoch()"
501        ))
502        .bind(&[key.into()])
503        .map_err(|e| CfError(worker::Error::RustError(format!("D1 bind failed: {e}"))))?
504        .first::<serde_json::Value>(None)
505        .await
506        .cf()?;
507
508    match result {
509        Some(row) => {
510            // D1 first() returns the row as a JSON object — extract the "data" field
511            let data_str = row
512                .get("data")
513                .and_then(|v| v.as_str())
514                .ok_or_else(|| {
515                    CfError(worker::Error::RustError(
516                        "session D1 row missing 'data' column".into(),
517                    ))
518                })?;
519            let data: HashMap<String, serde_json::Value> =
520                serde_json::from_str(data_str).cf()?;
521            Ok(Some(data))
522        }
523        None => Ok(None),
524    }
525}
526
527async fn save_to_d1(
528    binding: &str,
529    table: &str,
530    session_id: &str,
531    data: &HashMap<String, serde_json::Value>,
532    max_age_secs: u64,
533) -> Result<(), CfError> {
534    let db = crate::bindings::d1(binding)?;
535    let key = session_key(session_id);
536    let json = serde_json::to_string(data).cf()?;
537
538    db.prepare(format!(
539        "INSERT OR REPLACE INTO {table} (id, data, expires_at) VALUES (?, ?, unixepoch() + ?)"
540    ))
541    .bind(&[key.into(), json.into(), max_age_secs.into()])
542    .map_err(|e| CfError(worker::Error::RustError(format!("D1 bind failed: {e}"))))?
543    .run()
544    .await
545    .cf()?;
546
547    Ok(())
548}
549
550async fn delete_from_d1(binding: &str, table: &str, session_id: &str) -> Result<(), CfError> {
551    let db = crate::bindings::d1(binding)?;
552    let key = session_key(session_id);
553
554    db.prepare(format!("DELETE FROM {table} WHERE id = ?"))
555        .bind(&[key.into()])
556        .map_err(|e| CfError(worker::Error::RustError(format!("D1 bind failed: {e}"))))?
557        .run()
558        .await
559        .cf()?;
560
561    Ok(())
562}