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

session_guard/store/
valkey.rs
