session_guard/store/
valkey.rs

1//! Valkey / Redis store.
2
3use std::{fmt, net::IpAddr, sync::LazyLock};
4
5use chrono::{DateTime, TimeDelta, Utc};
6use redis::{Script, aio::ConnectionManager};
7use serde::Serialize;
8use uuid::Uuid;
9
10use super::{SessionRecord, SessionStore, SessionUpdate, Supersede, TokenLookup, TokenRecord};
11use crate::{
12    checks::Check,
13    error::StoreError,
14    event::{EventKind, SessionEvent},
15    fingerprint::Baseline,
16    token::TokenHash,
17};
18
19/// How many sessions one `delete_expired` script invocation removes at most.
20const PURGE_BATCH: usize = 500;
21
22/// Lua helpers shared by several scripts.
23macro_rules! lua_helpers {
24    () => {
25        r"
26local function read_session(P, sid)
27  local v = redis.call('HMGET', P .. 's:' .. sid,
28    'subject', 'created', 'auth', 'seen', 'expires', 'step_up', 'baseline')
29  if not v[1] then return nil end
30  return {sid, v[1], v[2], v[3], v[4], v[5], v[6], v[7]}
31end
32
33local function del_session(P, sid)
34  local s = P .. 's:' .. sid
35  local k = P .. 'k:' .. sid
36  local subject = redis.call('HGET', s, 'subject')
37  for _, hex in ipairs(redis.call('ZRANGE', k, 0, -1)) do
38    redis.call('DEL', P .. 't:' .. hex)
39  end
40  redis.call('DEL', k)
41  redis.call('ZREM', P .. 'x', sid)
42  redis.call('ZREM', P .. 'i', sid)
43  if subject then
44    redis.call('SREM', P .. 'u:' .. subject, sid)
45    redis.call('DEL', s)
46    return 1
47  end
48  return 0
49end
50"
51    };
52}
53
54/// KEYS: session, token, token set, subject set.
55/// ARGV: prefix, sid, token hex, subject, created, auth, seen, expires, `step_up`, baseline,
56/// `exp_ms`.
57static INSERT: LazyLock<Script> = LazyLock::new(|| {
58    Script::new(
59        r"
60local P = ARGV[1]
61if redis.call('EXISTS', KEYS[1]) == 1 or redis.call('EXISTS', KEYS[2]) == 1 then return 0 end
62redis.call('HSET', KEYS[1], 'subject', ARGV[4], 'created', ARGV[5], 'auth', ARGV[6],
63  'seen', ARGV[7], 'expires', ARGV[8], 'step_up', ARGV[9], 'baseline', ARGV[10],
64  'exp_ms', ARGV[11], 'cur', ARGV[3], 'seq', 1)
65redis.call('HSET', KEYS[2], 'sid', ARGV[2], 'seq', 1, 'issued', ARGV[5])
66redis.call('ZADD', KEYS[3], 1, ARGV[3])
67redis.call('SADD', KEYS[4], ARGV[2])
68redis.call('ZADD', P .. 'x', ARGV[8], ARGV[2])
69redis.call('ZADD', P .. 'i', ARGV[7], ARGV[2])
70redis.call('PEXPIREAT', KEYS[1], ARGV[11])
71redis.call('PEXPIREAT', KEYS[2], ARGV[11])
72redis.call('PEXPIREAT', KEYS[3], ARGV[11])
73local pt = redis.call('PEXPIRETIME', KEYS[4])
74if pt == -1 or pt < tonumber(ARGV[11]) then redis.call('PEXPIREAT', KEYS[4], ARGV[11]) end
75return 1
76",
77    )
78});
79
80/// KEYS: token. ARGV: prefix.
81static FIND: LazyLock<Script> = LazyLock::new(|| {
82    Script::new(concat!(
83        lua_helpers!(),
84        r"
85local P = ARGV[1]
86local t = redis.call('HMGET', KEYS[1], 'sid', 'seq', 'issued', 'used', 'sup')
87if not t[1] then return false end
88local rec = read_session(P, t[1])
89if not rec then return false end
90local newer = 0
91for _, hex in ipairs(redis.call('ZRANGEBYSCORE', P .. 'k:' .. t[1], '(' .. t[2], '+inf')) do
92  if redis.call('HEXISTS', P .. 't:' .. hex, 'used') == 1 then
93    newer = 1
94    break
95  end
96end
97rec[9] = t[3]
98rec[10] = t[4] or ''
99rec[11] = t[5] or ''
100rec[12] = tostring(newer)
101return rec
102",
103    ))
104});
105
106/// KEYS: session.
107/// ARGV: prefix, sid, now, supersede mode (`''` for no rotation), mode argument, new token hex,
108/// prune cutoff, `mark_used` hex or `''`, baseline or `''`, `step_up` or `''`, auth or `''`.
109static APPLY: LazyLock<Script> = LazyLock::new(|| {
110    Script::new(
111        r"
112local s = KEYS[1]
113if redis.call('EXISTS', s) == 0 then return 0 end
114local P, sid, now, mode = ARGV[1], ARGV[2], ARGV[3], ARGV[4]
115local rotated = 0
116if mode ~= '' then
117  local cur = redis.call('HGET', s, 'cur')
118  if cur then
119    local ck = P .. 't:' .. cur
120    local holds = false
121    if mode == 'exact' then
122      holds = cur == ARGV[5]
123    elseif mode == 'before' then
124      local issued = tonumber(redis.call('HGET', ck, 'issued'))
125      holds = issued ~= nil and issued < tonumber(ARGV[5])
126    elseif mode == 'current' then
127      holds = redis.call('EXISTS', ck) == 1
128    end
129    local nk = P .. 't:' .. ARGV[6]
130    if holds and redis.call('EXISTS', nk) == 0 then
131      local k = P .. 'k:' .. sid
132      redis.call('HSET', ck, 'sup', now)
133      local seq = redis.call('HINCRBY', s, 'seq', 1)
134      redis.call('HSET', nk, 'sid', sid, 'seq', seq, 'issued', now)
135      local exp = redis.call('HGET', s, 'exp_ms')
136      if exp then redis.call('PEXPIREAT', nk, exp) end
137      redis.call('ZADD', k, seq, ARGV[6])
138      redis.call('HSET', s, 'cur', ARGV[6])
139      local cutoff = tonumber(ARGV[7])
140      for _, hex in ipairs(redis.call('ZRANGE', k, 0, -1)) do
141        local sup = tonumber(redis.call('HGET', P .. 't:' .. hex, 'sup'))
142        if sup ~= nil and sup < cutoff then
143          redis.call('DEL', P .. 't:' .. hex)
144          redis.call('ZREM', k, hex)
145        end
146      end
147      rotated = 1
148    end
149  end
150end
151if ARGV[8] ~= '' then
152  local mk = P .. 't:' .. ARGV[8]
153  if redis.call('HGET', mk, 'sid') == sid and redis.call('HEXISTS', mk, 'used') == 0 then
154    redis.call('HSET', mk, 'used', now)
155  end
156end
157redis.call('HSET', s, 'seen', now)
158redis.call('ZADD', P .. 'i', now, sid)
159if ARGV[9] ~= '' then redis.call('HSET', s, 'baseline', ARGV[9]) end
160if ARGV[10] ~= '' then redis.call('HSET', s, 'step_up', ARGV[10]) end
161if ARGV[11] ~= '' then redis.call('HSET', s, 'auth', ARGV[11]) end
162return rotated
163",
164    )
165});
166
167/// ARGV: prefix, sid.
168static DELETE: LazyLock<Script> =
169    LazyLock::new(|| Script::new(concat!(lua_helpers!(), "return del_session(ARGV[1], ARGV[2])\n")));
170
171/// KEYS: subject set. ARGV: prefix.
172static DELETE_SUBJECT: LazyLock<Script> = LazyLock::new(|| {
173    Script::new(concat!(
174        lua_helpers!(),
175        r"
176local n = 0
177for _, sid in ipairs(redis.call('SMEMBERS', KEYS[1])) do
178  n = n + del_session(ARGV[1], sid)
179end
180redis.call('DEL', KEYS[1])
181return n
182",
183    ))
184});
185
186/// KEYS: subject set. ARGV: prefix.
187static LIST_SUBJECT: LazyLock<Script> = LazyLock::new(|| {
188    Script::new(concat!(
189        lua_helpers!(),
190        r"
191local out = {}
192for _, sid in ipairs(redis.call('SMEMBERS', KEYS[1])) do
193  local rec = read_session(ARGV[1], sid)
194  if rec then
195    out[#out + 1] = rec
196  else
197    redis.call('SREM', KEYS[1], sid)
198  end
199end
200return out
201",
202    ))
203});
204
205/// ARGV: prefix, now, idle cutoff, batch size.
206static DELETE_EXPIRED: LazyLock<Script> = LazyLock::new(|| {
207    Script::new(concat!(
208        lua_helpers!(),
209        r"
210local P = ARGV[1]
211local found, n = 0, 0
212for _, idx in ipairs({{'x', ARGV[2]}, {'i', ARGV[3]}}) do
213  local ids = redis.call('ZRANGEBYSCORE', P .. idx[1], '-inf', idx[2], 'LIMIT', 0, ARGV[4])
214  found = found + #ids
215  for _, sid in ipairs(ids) do n = n + del_session(P, sid) end
216end
217return {n, found}
218",
219    ))
220});
221
222/// KEYS: events, claim index. ARGV: prefix, cutoff.
223static DELETE_EVENTS: LazyLock<Script> = LazyLock::new(|| {
224    Script::new(
225        r"
226local bound = '(' .. ARGV[2]
227local n = redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', bound)
228for _, m in ipairs(redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', bound)) do
229  redis.call('DEL', ARGV[1] .. 'c:' .. m)
230end
231redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', bound)
232return n
233",
234    )
235});
236
237/// KEYS: claim, claim index. ARGV: now, window (µs), window (ms), claim index member.
238static CLAIM: LazyLock<Script> = LazyLock::new(|| {
239    Script::new(
240        r"
241local at = tonumber(redis.call('GET', KEYS[1]))
242if at ~= nil and tonumber(ARGV[1]) - at < tonumber(ARGV[2]) then return 0 end
243redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[3])
244redis.call('ZADD', KEYS[2], ARGV[1], ARGV[4])
245return 1
246",
247    )
248});
249
250/// A [`SessionStore`] on Valkey or Redis.
251///
252/// Every write that touches more than one key runs as a Lua script, so it is atomic across any
253/// number of processes sharing the server. Requires Valkey 7.2+ or Redis 7.0+ in standalone or
254/// primary/replica mode; Cluster is not supported because the scripts derive key names from
255/// stored data.
256///
257/// # Data layout
258///
259/// All keys start with the [key prefix](Self::key_prefix) (`P`, default `sg:`).
260/// Timestamps are microseconds since the Unix epoch as decimal integers.
261///
262/// | Key | Type | Contents |
263/// |-----|------|----------|
264/// | `P s:{session id}` | hash | `subject`, `created`, `auth`, `seen`, `expires`, `step_up`, `baseline` (JSON), `cur` (current token), `seq` (token counter), `exp_ms` |
265/// | `P t:{token hash hex}` | hash | `sid`, `seq` (issue order), `issued`, `used`, `sup` (superseded) |
266/// | `P k:{session id}` | sorted set | the session's token hashes, scored by `seq` |
267/// | `P u:{subject}` | set | session ids of a subject |
268/// | `P x` | sorted set | session ids scored by `expires_at` |
269/// | `P i` | sorted set | session ids scored by `last_seen_at` |
270/// | `P e` | sorted set | events as JSON, scored by `at` |
271/// | `P c:{session id}:{kind}` | string | time of the last event claim |
272/// | `P cz` | sorted set | `{session id}:{kind}` claims scored by claim time |
273///
274/// Session, token and per-session token-set keys expire at the session's absolute expiry and
275/// subject sets at their last session's expiry, so Valkey cleans up even if
276/// [`delete_expired`](SessionStore::delete_expired) never runs. Because those TTLs follow the
277/// server's wall clock, sessions whose `expires_at` is already in the past vanish immediately.
278/// Claim keys expire after their window.
279///
280/// Subjects are passed to scripts as data, never as code.
281///
282/// Cheap to clone; clones share the connection.
283#[derive(Clone)]
284pub struct ValkeyStore {
285    conn: ConnectionManager,
286    prefix: String,
287}
288
289impl fmt::Debug for ValkeyStore {
290    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
291        f.debug_struct("ValkeyStore")
292            .field("key_prefix", &self.prefix)
293            .finish_non_exhaustive()
294    }
295}
296
297impl ValkeyStore {
298    /// A store on an existing connection, with the default key prefix `sg:`.
299    pub fn new(conn: ConnectionManager) -> Self {
300        Self {
301            conn,
302            prefix: "sg:".to_owned(),
303        }
304    }
305
306    /// Connects to `url` (`redis://`, `rediss://` or `valkey://`), with the default key prefix
307    /// `sg:`.
308    ///
309    /// # Errors
310    ///
311    /// If the URL is invalid or the server is unreachable.
312    pub async fn connect(url: &str) -> Result<Self, redis::RedisError> {
313        let client = redis::Client::open(url)?;
314        Ok(Self::new(ConnectionManager::new(client).await?))
315    }
316
317    /// Prefixes every key with `prefix`, so several apps can share one server. Defaults to
318    /// `sg:`.
319    #[must_use]
320    pub fn key_prefix(mut self, prefix: impl Into<String>) -> Self {
321        self.prefix = prefix.into();
322        self
323    }
324
325    fn key(&self, kind: &str, id: &str) -> String {
326        format!("{}{kind}:{id}", self.prefix)
327    }
328
329    fn conn(&self) -> ConnectionManager {
330        self.conn.clone()
331    }
332}
333
334fn micros(at: DateTime<Utc>) -> String {
335    at.timestamp_micros().to_string()
336}
337
338fn parse_time(value: &str) -> Result<DateTime<Utc>, StoreError> {
339    value
340        .parse::<i64>()
341        .ok()
342        .and_then(DateTime::from_timestamp_micros)
343        .ok_or_else(|| "valkey store: malformed timestamp".into())
344}
345
346fn parse_optional_time(value: &str) -> Result<Option<DateTime<Utc>>, StoreError> {
347    if value.is_empty() {
348        Ok(None)
349    } else {
350        parse_time(value).map(Some)
351    }
352}
353
354/// Parses the eight fields `read_session` returns.
355fn parse_session(fields: &[String]) -> Result<SessionRecord, StoreError> {
356    let [id, subject, created, auth, seen, expires, step_up, baseline] = fields else {
357        return Err("valkey store: malformed session".into());
358    };
359    let baseline: Baseline = serde_json::from_str(baseline)?;
360    Ok(SessionRecord {
361        id: Uuid::parse_str(id)?,
362        subject: subject.clone(),
363        created_at: parse_time(created)?,
364        authenticated_at: parse_time(auth)?,
365        last_seen_at: parse_time(seen)?,
366        expires_at: parse_time(expires)?,
367        step_up_required: step_up == "1",
368        baseline,
369    })
370}
371
372#[derive(Serialize)]
373struct EventRow<'a> {
374    /// Keeps identical events distinct in the sorted set.
375    id: Uuid,
376    session_id: Uuid,
377    subject: &'a str,
378    kind: EventKind,
379    enforced: bool,
380    score: Option<u32>,
381    signals: &'a [Check],
382    ip: IpAddr,
383    user_agent: Option<&'a str>,
384    at: DateTime<Utc>,
385}
386
387impl SessionStore for ValkeyStore {
388    async fn insert_session(&self, session: &SessionRecord, token: &TokenHash) -> Result<(), StoreError> {
389        let sid = session.id.to_string();
390        let hex = token.to_hex();
391        let exp_ms = session
392            .expires_at
393            .timestamp_micros()
394            .saturating_add(999)
395            .div_euclid(1000);
396        let inserted: i64 = INSERT
397            .key(self.key("s", &sid))
398            .key(self.key("t", &hex))
399            .key(self.key("k", &sid))
400            .key(self.key("u", &session.subject))
401            .arg(&self.prefix)
402            .arg(&sid)
403            .arg(&hex)
404            .arg(&session.subject)
405            .arg(micros(session.created_at))
406            .arg(micros(session.authenticated_at))
407            .arg(micros(session.last_seen_at))
408            .arg(micros(session.expires_at))
409            .arg(if session.step_up_required { "1" } else { "0" })
410            .arg(serde_json::to_string(&session.baseline)?)
411            .arg(exp_ms)
412            .invoke_async(&mut self.conn())
413            .await?;
414        if inserted == 1 {
415            Ok(())
416        } else {
417            Err("duplicate session id or token".into())
418        }
419    }
420
421    async fn find_by_token(&self, token: &TokenHash) -> Result<Option<TokenLookup>, StoreError> {
422        let fields: Option<Vec<String>> = FIND
423            .key(self.key("t", &token.to_hex()))
424            .arg(&self.prefix)
425            .invoke_async(&mut self.conn())
426            .await?;
427        let Some(fields) = fields else {
428            return Ok(None);
429        };
430        let [session @ .., issued, used, superseded, newer] = fields.as_slice() else {
431            return Err("valkey store: malformed token lookup".into());
432        };
433        Ok(Some(TokenLookup {
434            session: parse_session(session)?,
435            token: TokenRecord {
436                hash: *token,
437                issued_at: parse_time(issued)?,
438                first_used_at: parse_optional_time(used)?,
439                superseded_at: parse_optional_time(superseded)?,
440            },
441            newer_token_used: newer == "1",
442        }))
443    }
444
445    async fn apply(&self, update: &SessionUpdate) -> Result<bool, StoreError> {
446        let sid = update.session_id.to_string();
447        let (mode, mode_arg, new_token, prune) = match &update.rotation {
448            None => ("", String::new(), String::new(), String::new()),
449            Some(rotation) => {
450                let (mode, arg) = match rotation.supersede {
451                    Supersede::Exact(hash) => ("exact", hash.to_hex()),
452                    Supersede::CurrentIssuedBefore(at) => ("before", micros(at)),
453                    Supersede::Current => ("current", String::new()),
454                };
455                (
456                    mode,
457                    arg,
458                    rotation.new_token.to_hex(),
459                    micros(rotation.prune_superseded_before),
460                )
461            }
462        };
463        let baseline = match &update.baseline {
464            Some(baseline) => serde_json::to_string(baseline)?,
465            None => String::new(),
466        };
467        let step_up = match update.step_up_required {
468            Some(true) => "1",
469            Some(false) => "0",
470            None => "",
471        };
472        let rotated: i64 = APPLY
473            .key(self.key("s", &sid))
474            .arg(&self.prefix)
475            .arg(&sid)
476            .arg(micros(update.now))
477            .arg(mode)
478            .arg(mode_arg)
479            .arg(new_token)
480            .arg(prune)
481            .arg(update.mark_used.map(|h| h.to_hex()).unwrap_or_default())
482            .arg(baseline)
483            .arg(step_up)
484            .arg(update.authenticated_at.map(micros).unwrap_or_default())
485            .invoke_async(&mut self.conn())
486            .await?;
487        Ok(rotated == 1)
488    }
489
490    async fn delete_session(&self, session_id: Uuid) -> Result<bool, StoreError> {
491        let deleted: i64 = DELETE
492            .arg(&self.prefix)
493            .arg(session_id.to_string())
494            .invoke_async(&mut self.conn())
495            .await?;
496        Ok(deleted == 1)
497    }
498
499    async fn delete_subject_sessions(&self, subject: &str) -> Result<u64, StoreError> {
500        let deleted: u64 = DELETE_SUBJECT
501            .key(self.key("u", subject))
502            .arg(&self.prefix)
503            .invoke_async(&mut self.conn())
504            .await?;
505        Ok(deleted)
506    }
507
508    async fn list_subject_sessions(&self, subject: &str) -> Result<Vec<SessionRecord>, StoreError> {
509        let rows: Vec<Vec<String>> = LIST_SUBJECT
510            .key(self.key("u", subject))
511            .arg(&self.prefix)
512            .invoke_async(&mut self.conn())
513            .await?;
514        rows.iter().map(|row| parse_session(row)).collect()
515    }
516
517    async fn delete_expired(&self, now: DateTime<Utc>, idle_cutoff: DateTime<Utc>) -> Result<u64, StoreError> {
518        let mut total = 0;
519        loop {
520            let (deleted, found): (u64, u64) = DELETE_EXPIRED
521                .arg(&self.prefix)
522                .arg(micros(now))
523                .arg(micros(idle_cutoff))
524                .arg(PURGE_BATCH)
525                .invoke_async(&mut self.conn())
526                .await?;
527            total += deleted;
528            if found == 0 {
529                return Ok(total);
530            }
531        }
532    }
533
534    async fn record_event(&self, event: &SessionEvent) -> Result<(), StoreError> {
535        let row = EventRow {
536            id: Uuid::new_v4(),
537            session_id: event.session_id,
538            subject: &event.subject,
539            kind: event.kind,
540            enforced: event.enforced,
541            score: event.score,
542            signals: &event.signals,
543            ip: event.ip,
544            user_agent: event.user_agent.as_deref(),
545            at: event.at,
546        };
547        let _added: i64 = redis::cmd("ZADD")
548            .arg(format!("{}e", self.prefix))
549            .arg(event.at.timestamp_micros())
550            .arg(serde_json::to_string(&row)?)
551            .query_async(&mut self.conn())
552            .await?;
553        Ok(())
554    }
555
556    async fn delete_events_before(&self, cutoff: DateTime<Utc>) -> Result<u64, StoreError> {
557        let deleted: u64 = DELETE_EVENTS
558            .key(format!("{}e", self.prefix))
559            .key(format!("{}cz", self.prefix))
560            .arg(&self.prefix)
561            .arg(micros(cutoff))
562            .invoke_async(&mut self.conn())
563            .await?;
564        Ok(deleted)
565    }
566
567    async fn claim_event(
568        &self,
569        session_id: Uuid,
570        kind: EventKind,
571        now: DateTime<Utc>,
572        window: TimeDelta,
573    ) -> Result<bool, StoreError> {
574        let member = format!("{session_id}:{kind}");
575        let micros_window = window.num_microseconds().unwrap_or(i64::MAX);
576        let millis_window = window.num_milliseconds().max(1);
577        let claimed: i64 = CLAIM
578            .key(self.key("c", &member))
579            .key(format!("{}cz", self.prefix))
580            .arg(micros(now))
581            .arg(micros_window)
582            .arg(millis_window)
583            .arg(&member)
584            .invoke_async(&mut self.conn())
585            .await?;
586        Ok(claimed == 1)
587    }
588}
session_guard/store/valkey.rs

session_guard/store/
valkey.rs
