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::{
7    RedisFuture, Script, Value,
8    aio::{ConnectionLike, ConnectionManager},
9    cluster_async::ClusterConnection,
10};
11use serde::Serialize;
12use sha2::{Digest, Sha256};
13use uuid::Uuid;
14
15use super::{
16    ClaimScope, SessionRecord, SessionStore, SessionUpdate, Supersede, TokenLookup, TokenRecord, session_bucket,
17    token_bucket,
18};
19use crate::{
20    checks::Check,
21    error::StoreError,
22    event::{ClientSummary, EventKind, SessionEvent, SessionOrigin},
23    fingerprint::Baseline,
24    token::TokenHash,
25};
26
27/// How many sessions one `delete_expired` script invocation removes at most.
28const PURGE_BATCH: usize = 500;
29
30/// How many session ids one `SREM` removes at most.
31const SREM_BATCH: usize = 1000;
32
33/// Lua helpers shared by several scripts. `P` is always a bucket prefix (`prefix{bb}:`), so every
34/// key a script builds shares the hash slot of the keys it was called with.
35macro_rules! lua_helpers {
36    () => {
37        r"
38local function read_session(P, sid)
39  local v = redis.call('HMGET', P .. 's:' .. sid,
40    'subject', 'created', 'auth', 'seen', 'expires', 'step_up', 'baseline', 'origin')
41  if not v[1] then return nil end
42  return {sid, v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8] or ''}
43end
44
45-- Deletes a session and its tokens; returns its subject, or false if it did not exist.
46local function del_session(P, sid)
47  local s = P .. 's:' .. sid
48  local k = P .. 'k:' .. sid
49  local subject = redis.call('HGET', s, 'subject')
50  for _, hex in ipairs(redis.call('ZRANGE', k, 0, -1)) do
51    redis.call('DEL', P .. 't:' .. hex)
52  end
53  redis.call('DEL', k)
54  redis.call('ZREM', P .. 'x', sid)
55  redis.call('ZREM', P .. 'i', sid)
56  if subject then
57    redis.call('DEL', s)
58    return subject
59  end
60  return false
61end
62"
63    };
64}
65
66/// KEYS: session, token, token set, expiry index. ARGV: bucket prefix, sid, token hex, subject,
67/// created, auth, seen, expires, `step_up`, baseline, `exp_ms`, origin or `''`.
68static INSERT: LazyLock<Script> = LazyLock::new(|| {
69    Script::new(
70        r"
71local P = ARGV[1]
72if redis.call('EXISTS', KEYS[1]) == 1 or redis.call('EXISTS', KEYS[2]) == 1 then return 0 end
73redis.call('HSET', KEYS[1], 'subject', ARGV[4], 'created', ARGV[5], 'auth', ARGV[6],
74  'seen', ARGV[7], 'expires', ARGV[8], 'step_up', ARGV[9], 'baseline', ARGV[10],
75  'exp_ms', ARGV[11], 'cur', ARGV[3], 'seq', 1)
76if ARGV[12] ~= '' then redis.call('HSET', KEYS[1], 'origin', ARGV[12]) end
77redis.call('HSET', KEYS[2], 'sid', ARGV[2], 'seq', 1, 'issued', ARGV[5])
78redis.call('ZADD', KEYS[3], 1, ARGV[3])
79redis.call('ZADD', KEYS[4], ARGV[8], ARGV[2])
80redis.call('ZADD', P .. 'i', ARGV[7], ARGV[2])
81redis.call('PEXPIREAT', KEYS[1], ARGV[11])
82redis.call('PEXPIREAT', KEYS[2], ARGV[11])
83redis.call('PEXPIREAT', KEYS[3], ARGV[11])
84return 1
85",
86    )
87});
88
89/// KEYS: subject set. ARGV: sid, `exp_ms`. Adds a session to the set and keeps the set alive
90/// until the latest member expires.
91static INDEX_ADD: LazyLock<Script> = LazyLock::new(|| {
92    Script::new(
93        r"
94redis.call('SADD', KEYS[1], ARGV[1])
95local pt = redis.call('PEXPIRETIME', KEYS[1])
96if pt == -1 or pt < tonumber(ARGV[2]) then redis.call('PEXPIREAT', KEYS[1], ARGV[2]) end
97return 1
98",
99    )
100});
101
102/// KEYS: subject set, revoking set. Moves every member of the subject set into the revoking set
103/// (keeping the later expiry) and returns the revoking set's members, which includes any a
104/// crashed revoke left behind.
105static INDEX_MOVE: LazyLock<Script> = LazyLock::new(|| {
106    Script::new(
107        r"
108local pt = redis.call('PEXPIRETIME', KEYS[1])
109if pt ~= -2 then
110  local cur = redis.call('PEXPIRETIME', KEYS[2])
111  redis.call('SUNIONSTORE', KEYS[2], KEYS[2], KEYS[1])
112  local at = math.max(pt, cur)
113  if at > 0 then redis.call('PEXPIREAT', KEYS[2], at) end
114  redis.call('DEL', KEYS[1])
115end
116return redis.call('SMEMBERS', KEYS[2])
117",
118    )
119});
120
121/// KEYS: token. ARGV: bucket prefix.
122static FIND: LazyLock<Script> = LazyLock::new(|| {
123    Script::new(concat!(
124        lua_helpers!(),
125        r"
126local P = ARGV[1]
127local t = redis.call('HMGET', KEYS[1], 'sid', 'seq', 'issued', 'used', 'sup')
128if not t[1] then return false end
129local rec = read_session(P, t[1])
130if not rec then return false end
131local newer = 0
132for _, hex in ipairs(redis.call('ZRANGEBYSCORE', P .. 'k:' .. t[1], '(' .. t[2], '+inf')) do
133  if redis.call('HEXISTS', P .. 't:' .. hex, 'used') == 1 then
134    newer = 1
135    break
136  end
137end
138rec[10] = t[3]
139rec[11] = t[4] or ''
140rec[12] = t[5] or ''
141rec[13] = tostring(newer)
142return rec
143",
144    ))
145});
146
147/// KEYS: session. ARGV: bucket prefix, sid.
148static READ: LazyLock<Script> = LazyLock::new(|| {
149    Script::new(concat!(
150        lua_helpers!(),
151        "return read_session(ARGV[1], ARGV[2]) or false\n"
152    ))
153});
154
155/// KEYS: session. ARGV: bucket prefix, sid, now, supersede mode (`''` for no rotation), mode argument, new token hex,
156/// prune cutoff, `mark_used` hex or `''`, baseline or `''`, `step_up` or `''`, auth or `''`.
157static APPLY: LazyLock<Script> = LazyLock::new(|| {
158    Script::new(
159        r"
160local s = KEYS[1]
161if redis.call('EXISTS', s) == 0 then return 0 end
162local P, sid, now, mode = ARGV[1], ARGV[2], ARGV[3], ARGV[4]
163local rotated = 0
164if mode ~= '' then
165  local cur = redis.call('HGET', s, 'cur')
166  if cur then
167    local ck = P .. 't:' .. cur
168    local holds = false
169    if mode == 'exact' then
170      holds = cur == ARGV[5]
171    elseif mode == 'before' then
172      local issued = tonumber(redis.call('HGET', ck, 'issued'))
173      holds = issued ~= nil and issued < tonumber(ARGV[5])
174    elseif mode == 'current' then
175      holds = redis.call('EXISTS', ck) == 1
176    end
177    local nk = P .. 't:' .. ARGV[6]
178    if holds and redis.call('EXISTS', nk) == 0 then
179      local k = P .. 'k:' .. sid
180      redis.call('HSET', ck, 'sup', now)
181      local seq = redis.call('HINCRBY', s, 'seq', 1)
182      redis.call('HSET', nk, 'sid', sid, 'seq', seq, 'issued', now)
183      local exp = redis.call('HGET', s, 'exp_ms')
184      if exp then redis.call('PEXPIREAT', nk, exp) end
185      redis.call('ZADD', k, seq, ARGV[6])
186      redis.call('HSET', s, 'cur', ARGV[6])
187      local cutoff = tonumber(ARGV[7])
188      for _, hex in ipairs(redis.call('ZRANGE', k, 0, -1)) do
189        local sup = tonumber(redis.call('HGET', P .. 't:' .. hex, 'sup'))
190        if sup ~= nil and sup < cutoff then
191          redis.call('DEL', P .. 't:' .. hex)
192          redis.call('ZREM', k, hex)
193        end
194      end
195      rotated = 1
196    end
197  end
198end
199if ARGV[8] ~= '' then
200  local mk = P .. 't:' .. ARGV[8]
201  if redis.call('HGET', mk, 'sid') == sid and redis.call('HEXISTS', mk, 'used') == 0 then
202    redis.call('HSET', mk, 'used', now)
203  end
204end
205redis.call('HSET', s, 'seen', now)
206redis.call('ZADD', P .. 'i', now, sid)
207if ARGV[9] ~= '' then redis.call('HSET', s, 'baseline', ARGV[9]) end
208if ARGV[10] ~= '' then redis.call('HSET', s, 'step_up', ARGV[10]) end
209if ARGV[11] ~= '' then redis.call('HSET', s, 'auth', ARGV[11]) end
210return rotated
211",
212    )
213});
214
215/// KEYS: session. ARGV: bucket prefix, sid, expected subject (optional). Returns the deleted
216/// session's subject, or nothing when it did not exist or belongs to another subject.
217static DELETE: LazyLock<Script> = LazyLock::new(|| {
218    Script::new(concat!(
219        lua_helpers!(),
220        r"
221if ARGV[3] and redis.call('HGET', KEYS[1], 'subject') ~= ARGV[3] then return false end
222return del_session(ARGV[1], ARGV[2])
223",
224    ))
225});
226
227/// KEYS: expiry index. ARGV: bucket prefix, now, idle cutoff, batch size.
228static DELETE_EXPIRED: LazyLock<Script> = LazyLock::new(|| {
229    Script::new(concat!(
230        lua_helpers!(),
231        r"
232local P = ARGV[1]
233local found, n = 0, 0
234for _, idx in ipairs({{'x', ARGV[2]}, {'i', ARGV[3]}}) do
235  local ids = redis.call('ZRANGEBYSCORE', P .. idx[1], '-inf', idx[2], 'LIMIT', 0, ARGV[4])
236  found = found + #ids
237  for _, sid in ipairs(ids) do
238    if del_session(P, sid) then n = n + 1 end
239  end
240end
241return {n, found}
242",
243    ))
244});
245
246/// KEYS: events, claim index. ARGV: bucket prefix, cutoff.
247static DELETE_EVENTS: LazyLock<Script> = LazyLock::new(|| {
248    Script::new(
249        r"
250local bound = '(' .. ARGV[2]
251local n = redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', bound)
252for _, m in ipairs(redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', bound)) do
253  redis.call('DEL', ARGV[1] .. 'c:' .. m)
254end
255redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', bound)
256return n
257",
258    )
259});
260
261/// KEYS: claim, claim index. ARGV: now, window (µs), window (ms), claim index member.
262static CLAIM: LazyLock<Script> = LazyLock::new(|| {
263    Script::new(
264        r"
265local at = tonumber(redis.call('GET', KEYS[1]))
266if at ~= nil and tonumber(ARGV[1]) - at < tonumber(ARGV[2]) then return 0 end
267redis.call('SET', KEYS[1], ARGV[1], 'PX', ARGV[3])
268redis.call('ZADD', KEYS[2], ARGV[1], ARGV[4])
269return 1
270",
271    )
272});
273
274/// Standalone or cluster connection behind one type, so the store has a single code path.
275#[derive(Clone)]
276enum Conn {
277    Standalone(ConnectionManager),
278    Cluster(ClusterConnection),
279}
280
281impl ConnectionLike for Conn {
282    fn req_packed_command<'a>(&'a mut self, cmd: &'a redis::Cmd) -> RedisFuture<'a, Value> {
283        match self {
284            Self::Standalone(conn) => conn.req_packed_command(cmd),
285            Self::Cluster(conn) => conn.req_packed_command(cmd),
286        }
287    }
288
289    fn req_packed_commands<'a>(
290        &'a mut self,
291        pipeline: &'a redis::Pipeline,
292        offset: usize,
293        count: usize,
294    ) -> RedisFuture<'a, Vec<Value>> {
295        match self {
296            Self::Standalone(conn) => conn.req_packed_commands(pipeline, offset, count),
297            Self::Cluster(conn) => conn.req_packed_commands(pipeline, offset, count),
298        }
299    }
300
301    fn get_db(&self) -> i64 {
302        match self {
303            Self::Standalone(conn) => conn.get_db(),
304            Self::Cluster(conn) => conn.get_db(),
305        }
306    }
307}
308
309/// A [`SessionStore`] on Valkey or Redis, standalone or Cluster.
310///
311/// Every write that touches more than one key runs as a Lua script, so it is atomic across any
312/// number of processes sharing the server. Requires Valkey 7.2+ or Redis 7.0+.
313///
314/// # Cluster
315///
316/// Connect with [`connect_cluster`](Self::connect_cluster) or
317/// [`from_cluster`](Self::from_cluster). Keys are split into 256 fixed buckets, and every key
318/// carries its bucket as a hash tag (`{bb}`), so a session, its tokens and its indexes always
319/// live in one hash slot and each script runs on one node. The bucket is the last byte of the
320/// session id, the first byte of a token's hash, or the first byte of the SHA-256 of a subject.
321/// [`Guard`](crate::Guard) mines tokens so that a session's tokens share its bucket; the store
322/// rejects writes where they do not.
323///
324/// The subject index is the one structure that spans buckets: a session id sits in its subject's
325/// set while the session lives in its own bucket, so subject-wide operations are eventually
326/// consistent. [`delete_subject_sessions`](SessionStore::delete_subject_sessions) moves the set
327/// aside, deletes each session in its bucket and then drops the moved ids. If the caller
328/// crashes in between, the leftover ids stay pending and the next revoke for the subject
329/// finishes them. A session created while a revoke runs is either revoked or its creator gets an
330/// error and the session is removed. [`list_subject_sessions`](SessionStore::list_subject_sessions)
331/// includes sessions that are mid-revoke.
332///
333/// [`delete_expired`](SessionStore::delete_expired) and
334/// [`delete_events_before`](SessionStore::delete_events_before) visit all 256 buckets.
335///
336/// # Data layout
337///
338/// All keys start with the [key prefix](Self::key_prefix) followed by the bucket tag, written
339/// `B` here (for example `sg:{0a}:`). Timestamps are microseconds since the Unix epoch as
340/// decimal integers.
341///
342/// | Key | Type | Contents |
343/// |-----|------|----------|
344/// | `B 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` |
345/// | `B t:{token hash hex}` | hash | `sid`, `seq` (issue order), `issued`, `used`, `sup` (superseded) |
346/// | `B k:{session id}` | sorted set | the session's token hashes, scored by `seq` |
347/// | `B u:{subject}` | set | session ids of a subject (in the subject's bucket) |
348/// | `B ur:{subject}` | set | ids of a subject being revoked (in the subject's bucket) |
349/// | `B x` | sorted set | session ids scored by `expires_at` |
350/// | `B i` | sorted set | session ids scored by `last_seen_at` |
351/// | `B e` | sorted set | events as JSON, scored by `at` |
352/// | `B c:{session id}:{kind}` | string | time of the last session event claim |
353/// | `B c:s:{subject}:{kind}` | string | time of the last subject alert claim (subject's bucket) |
354/// | `B cz` | sorted set | `{session id}:{kind}` and `s:{subject}:{kind}` claims scored by claim time |
355///
356/// Session, token and per-session token-set keys expire at the session's absolute expiry and
357/// subject sets at their last session's expiry, so Valkey cleans up even if
358/// [`delete_expired`](SessionStore::delete_expired) never runs. Because those TTLs follow the
359/// server's wall clock, sessions whose `expires_at` is already in the past vanish immediately.
360/// Claim keys expire after their window.
361///
362/// Subjects are passed to scripts as data, never as code.
363///
364/// Cheap to clone; clones share the connection.
365#[derive(Clone)]
366pub struct ValkeyStore {
367    conn: Conn,
368    prefix: String,
369}
370
371impl fmt::Debug for ValkeyStore {
372    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
373        f.debug_struct("ValkeyStore")
374            .field("key_prefix", &self.prefix)
375            .finish_non_exhaustive()
376    }
377}
378
379impl ValkeyStore {
380    /// A store on an existing connection, with the default key prefix `sg:`.
381    pub fn new(conn: ConnectionManager) -> Self {
382        Self {
383            conn: Conn::Standalone(conn),
384            prefix: "sg:".to_owned(),
385        }
386    }
387
388    /// A store on an existing Cluster connection, with the default key prefix `sg:`.
389    pub fn from_cluster(conn: ClusterConnection) -> Self {
390        Self {
391            conn: Conn::Cluster(conn),
392            prefix: "sg:".to_owned(),
393        }
394    }
395
396    /// Connects to a Valkey or Redis Cluster through any of its nodes (`redis://`, `rediss://`
397    /// or `valkey://` URLs), with the default key prefix `sg:`.
398    ///
399    /// # Errors
400    ///
401    /// If a URL is invalid or no node is reachable.
402    pub async fn connect_cluster<T: redis::IntoConnectionInfo>(
403        nodes: impl IntoIterator<Item = T>,
404    ) -> Result<Self, redis::RedisError> {
405        let client = redis::cluster::ClusterClient::new(nodes)?;
406        Ok(Self::from_cluster(client.get_async_connection().await?))
407    }
408
409    /// Connects to `url` (`redis://`, `rediss://` or `valkey://`), with the default key prefix
410    /// `sg:`.
411    ///
412    /// # Errors
413    ///
414    /// If the URL is invalid or the server is unreachable.
415    pub async fn connect(url: &str) -> Result<Self, redis::RedisError> {
416        let client = redis::Client::open(url)?;
417        Ok(Self::new(ConnectionManager::new(client).await?))
418    }
419
420    /// Prefixes every key with `prefix`, so several apps can share one server. Defaults to
421    /// `sg:`.
422    ///
423    /// # Panics
424    ///
425    /// If `prefix` contains `{` or `}`: the store puts its own hash tag after the prefix, and a
426    /// brace in the prefix would change which slot keys land in.
427    #[must_use]
428    pub fn key_prefix(mut self, prefix: impl Into<String>) -> Self {
429        let prefix = prefix.into();
430        assert!(
431            !prefix.contains(['{', '}']),
432            "ValkeyStore key prefix must not contain '{{' or '}}'"
433        );
434        self.prefix = prefix;
435        self
436    }
437
438    /// `prefix{bb}:`, the start of every key in `bucket`.
439    fn bucket_prefix(&self, bucket: u8) -> String {
440        format!("{}{{{bucket:02x}}}:", self.prefix)
441    }
442
443    /// A key `kind:id` in `bucket`.
444    fn key(&self, bucket: u8, kind: &str, id: &str) -> String {
445        format!("{}{kind}:{id}", self.bucket_prefix(bucket))
446    }
447
448    /// A singleton key `kind` in `bucket`.
449    fn plain_key(&self, bucket: u8, kind: &str) -> String {
450        format!("{}{kind}", self.bucket_prefix(bucket))
451    }
452
453    fn conn(&self) -> Conn {
454        self.conn.clone()
455    }
456
457    /// Removes `session_id` from its subject's sets. Best effort: the sets also expire, and
458    /// readers tolerate ids of sessions that are gone.
459    async fn unindex(&self, subject: &str, session_id: Uuid) {
460        let bucket = subject_bucket(subject);
461        for kind in ["u", "ur"] {
462            let _removed: Result<i64, _> = redis::cmd("SREM")
463                .arg(self.key(bucket, kind, subject))
464                .arg(session_id.to_string())
465                .query_async(&mut self.conn())
466                .await;
467        }
468    }
469}
470
471/// The bucket of a subject's index sets: the first byte of the SHA-256 of its canonical JSON.
472fn subject_bucket(subject: &str) -> u8 {
473    Sha256::digest(subject.as_bytes())[0]
474}
475
476fn ensure_aligned(session_id: Uuid, token: &TokenHash) -> Result<(), StoreError> {
477    if session_bucket(session_id) == token_bucket(token) {
478        Ok(())
479    } else {
480        Err("valkey store: token and session id are in different buckets".into())
481    }
482}
483
484fn micros(at: DateTime<Utc>) -> String {
485    at.timestamp_micros().to_string()
486}
487
488fn parse_time(value: &str) -> Result<DateTime<Utc>, StoreError> {
489    value
490        .parse::<i64>()
491        .ok()
492        .and_then(DateTime::from_timestamp_micros)
493        .ok_or_else(|| "valkey store: malformed timestamp".into())
494}
495
496fn parse_optional_time(value: &str) -> Result<Option<DateTime<Utc>>, StoreError> {
497    if value.is_empty() {
498        Ok(None)
499    } else {
500        parse_time(value).map(Some)
501    }
502}
503
504/// Parses the nine fields `read_session` returns.
505fn parse_session(fields: &[String]) -> Result<SessionRecord, StoreError> {
506    let [id, subject, created, auth, seen, expires, step_up, baseline, origin] = fields else {
507        return Err("valkey store: malformed session".into());
508    };
509    let baseline: Baseline = serde_json::from_str(baseline)?;
510    let origin: Option<SessionOrigin> = if origin.is_empty() {
511        None
512    } else {
513        Some(serde_json::from_str(origin)?)
514    };
515    Ok(SessionRecord {
516        id: Uuid::parse_str(id)?,
517        subject: subject.clone(),
518        created_at: parse_time(created)?,
519        authenticated_at: parse_time(auth)?,
520        last_seen_at: parse_time(seen)?,
521        expires_at: parse_time(expires)?,
522        step_up_required: step_up == "1",
523        baseline,
524        origin,
525    })
526}
527
528#[derive(Serialize)]
529struct EventRow<'a> {
530    /// Keeps identical events distinct in the sorted set.
531    id: Uuid,
532    session_id: Uuid,
533    subject: &'a str,
534    kind: EventKind,
535    enforced: bool,
536    score: Option<u32>,
537    signals: &'a [Check],
538    ip: IpAddr,
539    user_agent: Option<&'a str>,
540    client: &'a ClientSummary,
541    at: DateTime<Utc>,
542}
543
544impl SessionStore for ValkeyStore {
545    fn sharded(&self) -> bool {
546        true
547    }
548
549    async fn insert_session(&self, session: &SessionRecord, token: &TokenHash) -> Result<(), StoreError> {
550        ensure_aligned(session.id, token)?;
551        let bucket = session_bucket(session.id);
552        let sid = session.id.to_string();
553        let hex = token.to_hex();
554        let exp_ms = session
555            .expires_at
556            .timestamp_micros()
557            .saturating_add(999)
558            .div_euclid(1000);
559        // Index first: a concurrent revoke then either sees the id or makes the check below fail.
560        let sbucket = subject_bucket(&session.subject);
561        let _: i64 = INDEX_ADD
562            .key(self.key(sbucket, "u", &session.subject))
563            .arg(&sid)
564            .arg(exp_ms)
565            .invoke_async(&mut self.conn())
566            .await?;
567        let inserted: i64 = INSERT
568            .key(self.key(bucket, "s", &sid))
569            .key(self.key(bucket, "t", &hex))
570            .key(self.key(bucket, "k", &sid))
571            .key(self.plain_key(bucket, "x"))
572            .arg(self.bucket_prefix(bucket))
573            .arg(&sid)
574            .arg(&hex)
575            .arg(&session.subject)
576            .arg(micros(session.created_at))
577            .arg(micros(session.authenticated_at))
578            .arg(micros(session.last_seen_at))
579            .arg(micros(session.expires_at))
580            .arg(if session.step_up_required { "1" } else { "0" })
581            .arg(serde_json::to_string(&session.baseline)?)
582            .arg(exp_ms)
583            .arg(match &session.origin {
584                Some(origin) => serde_json::to_string(origin)?,
585                None => String::new(),
586            })
587            .invoke_async(&mut self.conn())
588            .await?;
589        if inserted != 1 {
590            return Err("duplicate session id or token".into());
591        }
592        let indexed: bool = redis::cmd("SISMEMBER")
593            .arg(self.key(sbucket, "u", &session.subject))
594            .arg(&sid)
595            .query_async(&mut self.conn())
596            .await?;
597        if !indexed {
598            let exists: bool = redis::cmd("EXISTS")
599                .arg(self.key(bucket, "s", &sid))
600                .query_async(&mut self.conn())
601                .await?;
602            if exists {
603                // A revoke moved the index entry between our two writes; do not leave a session
604                // the revoke can no longer find.
605                self.delete_session(session.id).await?;
606                return Err("valkey store: subject was revoked while the session was created".into());
607            }
608        }
609        Ok(())
610    }
611
612    async fn find_by_token(&self, token: &TokenHash) -> Result<Option<TokenLookup>, StoreError> {
613        let bucket = token_bucket(token);
614        let fields: Option<Vec<String>> = FIND
615            .key(self.key(bucket, "t", &token.to_hex()))
616            .arg(self.bucket_prefix(bucket))
617            .invoke_async(&mut self.conn())
618            .await?;
619        let Some(fields) = fields else {
620            return Ok(None);
621        };
622        let [session @ .., issued, used, superseded, newer] = fields.as_slice() else {
623            return Err("valkey store: malformed token lookup".into());
624        };
625        Ok(Some(TokenLookup {
626            session: parse_session(session)?,
627            token: TokenRecord {
628                hash: *token,
629                issued_at: parse_time(issued)?,
630                first_used_at: parse_optional_time(used)?,
631                superseded_at: parse_optional_time(superseded)?,
632            },
633            newer_token_used: newer == "1",
634        }))
635    }
636
637    async fn apply(&self, update: &SessionUpdate) -> Result<bool, StoreError> {
638        if let Some(rotation) = &update.rotation {
639            ensure_aligned(update.session_id, &rotation.new_token)?;
640        }
641        if let Some(used) = &update.mark_used {
642            ensure_aligned(update.session_id, used)?;
643        }
644        let bucket = session_bucket(update.session_id);
645        let sid = update.session_id.to_string();
646        let (mode, mode_arg, new_token, prune) = match &update.rotation {
647            None => ("", String::new(), String::new(), String::new()),
648            Some(rotation) => {
649                let (mode, arg) = match rotation.supersede {
650                    Supersede::Exact(hash) => ("exact", hash.to_hex()),
651                    Supersede::CurrentIssuedBefore(at) => ("before", micros(at)),
652                    Supersede::Current => ("current", String::new()),
653                };
654                (
655                    mode,
656                    arg,
657                    rotation.new_token.to_hex(),
658                    micros(rotation.prune_superseded_before),
659                )
660            }
661        };
662        let baseline = match &update.baseline {
663            Some(baseline) => serde_json::to_string(baseline)?,
664            None => String::new(),
665        };
666        let step_up = match update.step_up_required {
667            Some(true) => "1",
668            Some(false) => "0",
669            None => "",
670        };
671        let rotated: i64 = APPLY
672            .key(self.key(bucket, "s", &sid))
673            .arg(self.bucket_prefix(bucket))
674            .arg(&sid)
675            .arg(micros(update.now))
676            .arg(mode)
677            .arg(mode_arg)
678            .arg(new_token)
679            .arg(prune)
680            .arg(update.mark_used.map(|h| h.to_hex()).unwrap_or_default())
681            .arg(baseline)
682            .arg(step_up)
683            .arg(update.authenticated_at.map(micros).unwrap_or_default())
684            .invoke_async(&mut self.conn())
685            .await?;
686        Ok(rotated == 1)
687    }
688
689    async fn delete_session(&self, session_id: Uuid) -> Result<bool, StoreError> {
690        let bucket = session_bucket(session_id);
691        let sid = session_id.to_string();
692        let subject: Option<String> = DELETE
693            .key(self.key(bucket, "s", &sid))
694            .arg(self.bucket_prefix(bucket))
695            .arg(&sid)
696            .invoke_async(&mut self.conn())
697            .await?;
698        let Some(subject) = subject else {
699            return Ok(false);
700        };
701        self.unindex(&subject, session_id).await;
702        Ok(true)
703    }
704
705    async fn delete_subject_sessions(&self, subject: &str) -> Result<Vec<Uuid>, StoreError> {
706        let sbucket = subject_bucket(subject);
707        let revoking = self.key(sbucket, "ur", subject);
708        let members: Vec<String> = INDEX_MOVE
709            .key(self.key(sbucket, "u", subject))
710            .key(&revoking)
711            .invoke_async(&mut self.conn())
712            .await?;
713        let mut deleted = Vec::new();
714        for member in &members {
715            // A malformed member can't name a session; the SREM below clears it.
716            let Ok(id) = Uuid::parse_str(member) else { continue };
717            let bucket = session_bucket(id);
718            let removed: Option<String> = DELETE
719                .key(self.key(bucket, "s", member))
720                .arg(self.bucket_prefix(bucket))
721                .arg(member)
722                .arg(subject)
723                .invoke_async(&mut self.conn())
724                .await?;
725            if removed.is_some() {
726                deleted.push(id);
727            }
728        }
729        // Only now are the ids done; a failure above leaves them pending for the next revoke.
730        for chunk in members.chunks(SREM_BATCH) {
731            let _removed: i64 = redis::cmd("SREM")
732                .arg(&revoking)
733                .arg(chunk)
734                .query_async(&mut self.conn())
735                .await?;
736        }
737        Ok(deleted)
738    }
739
740    async fn list_subject_sessions(&self, subject: &str) -> Result<Vec<SessionRecord>, StoreError> {
741        let sbucket = subject_bucket(subject);
742        let members: Vec<String> = redis::cmd("SUNION")
743            .arg(self.key(sbucket, "u", subject))
744            .arg(self.key(sbucket, "ur", subject))
745            .query_async(&mut self.conn())
746            .await?;
747        let mut sessions = Vec::with_capacity(members.len());
748        let mut stale = Vec::new();
749        for member in &members {
750            let Ok(id) = Uuid::parse_str(member) else {
751                stale.push(member);
752                continue;
753            };
754            let bucket = session_bucket(id);
755            let fields: Option<Vec<String>> = READ
756                .key(self.key(bucket, "s", member))
757                .arg(self.bucket_prefix(bucket))
758                .arg(member)
759                .invoke_async(&mut self.conn())
760                .await?;
761            match fields.map(|fields| parse_session(&fields)).transpose()? {
762                Some(record) if record.subject == subject => sessions.push(record),
763                _ => stale.push(member),
764            }
765        }
766        // Expired sessions leave their ids behind. Dropping one that a creator has added but
767        // not yet inserted is safe: the creator's membership check then fails the login.
768        for chunk in stale.chunks(SREM_BATCH) {
769            for kind in ["u", "ur"] {
770                let _removed: Result<i64, _> = redis::cmd("SREM")
771                    .arg(self.key(sbucket, kind, subject))
772                    .arg(chunk)
773                    .query_async(&mut self.conn())
774                    .await;
775            }
776        }
777        Ok(sessions)
778    }
779
780    async fn delete_expired(&self, now: DateTime<Utc>, idle_cutoff: DateTime<Utc>) -> Result<u64, StoreError> {
781        let mut total = 0;
782        for bucket in 0..=u8::MAX {
783            loop {
784                let (deleted, found): (u64, u64) = DELETE_EXPIRED
785                    .key(self.plain_key(bucket, "x"))
786                    .arg(self.bucket_prefix(bucket))
787                    .arg(micros(now))
788                    .arg(micros(idle_cutoff))
789                    .arg(PURGE_BATCH)
790                    .invoke_async(&mut self.conn())
791                    .await?;
792                total += deleted;
793                if found == 0 {
794                    break;
795                }
796            }
797        }
798        Ok(total)
799    }
800
801    async fn record_event(&self, event: &SessionEvent) -> Result<(), StoreError> {
802        let row = EventRow {
803            id: Uuid::new_v4(),
804            session_id: event.session_id,
805            subject: &event.subject,
806            kind: event.kind,
807            enforced: event.enforced,
808            score: event.score,
809            signals: &event.signals,
810            ip: event.ip,
811            user_agent: event.user_agent.as_deref(),
812            client: &event.client,
813            at: event.at,
814        };
815        let _added: i64 = redis::cmd("ZADD")
816            .arg(self.plain_key(session_bucket(event.session_id), "e"))
817            .arg(event.at.timestamp_micros())
818            .arg(serde_json::to_string(&row)?)
819            .query_async(&mut self.conn())
820            .await?;
821        Ok(())
822    }
823
824    async fn delete_events_before(&self, cutoff: DateTime<Utc>) -> Result<u64, StoreError> {
825        let mut total = 0;
826        for bucket in 0..=u8::MAX {
827            let deleted: u64 = DELETE_EVENTS
828                .key(self.plain_key(bucket, "e"))
829                .key(self.plain_key(bucket, "cz"))
830                .arg(self.bucket_prefix(bucket))
831                .arg(micros(cutoff))
832                .invoke_async(&mut self.conn())
833                .await?;
834            total += deleted;
835        }
836        Ok(total)
837    }
838
839    async fn claim_event(
840        &self,
841        scope: ClaimScope<'_>,
842        kind: EventKind,
843        now: DateTime<Utc>,
844        window: TimeDelta,
845    ) -> Result<bool, StoreError> {
846        // A session id never starts with `s:`, and no kind contains `:`, so members are unique.
847        let (bucket, member) = match scope {
848            ClaimScope::Session(session_id) => (session_bucket(session_id), format!("{session_id}:{kind}")),
849            ClaimScope::Subject(subject) => (subject_bucket(subject), format!("s:{subject}:{kind}")),
850        };
851        let micros_window = window.num_microseconds().unwrap_or(i64::MAX);
852        let millis_window = window.num_milliseconds().max(1);
853        let claimed: i64 = CLAIM
854            .key(self.key(bucket, "c", &member))
855            .key(self.plain_key(bucket, "cz"))
856            .arg(micros(now))
857            .arg(micros_window)
858            .arg(millis_window)
859            .arg(&member)
860            .invoke_async(&mut self.conn())
861            .await?;
862        Ok(claimed == 1)
863    }
864}
session_guard/store/valkey.rs

session_guard/store/
valkey.rs
