av_state/
redis_store.rs

1//! Redis-backed `StateStore` (brief §8 "Redis Cluster" layer).
2//!
3//! Check-and-spend runs as a server-side Lua script so it is atomic across
4//! distributed clients — same contract as `InMemoryStore::try_spend`.
5//! Contract tests live in `tests/redis_contract.rs`, gated on `AV_REDIS_URL`
6//! (skipped loudly when unset — a skipped gate prints, never silently passes).
7
8use crate::store::{Spend, StateError, StateStore};
9use redis::Commands;
10
11/// Atomic check-and-spend using subtraction so `current + amount` never rounds.
12const TRY_SPEND_LUA: &str = r"
13for i, key in ipairs(KEYS) do
14    local current = tonumber(redis.call('GET', key) or '0')
15    local amount = tonumber(ARGV[(i - 1) * 2 + 1])
16    local limit = tonumber(ARGV[(i - 1) * 2 + 2])
17    if current > limit or amount > limit - current then
18        return i
19    end
20end
21for i, key in ipairs(KEYS) do
22    redis.call('INCRBY', key, ARGV[(i - 1) * 2 + 1])
23    redis.call('EXPIRE', key, 86400)
24end
25return 0
26";
27
28const ADD_LUA: &str = r"
29local current = tonumber(redis.call('GET', KEYS[1]) or '0')
30local amount = tonumber(ARGV[1])
31local limit = tonumber(ARGV[2])
32if current > limit or amount > limit - current then
33    return -1
34end
35local result = redis.call('INCRBY', KEYS[1], ARGV[1])
36-- Match TRY_SPEND_LUA's 24 h TTL. Without this, any counter touched
37-- only through `add()` (bookkeeping, telemetry, non-budget spending)
38-- persists forever in Redis; over a long-running deployment that
39-- silently leaks memory until Redis OOMs. The two APIs must be
40-- interchangeable from the persistence perspective.
41redis.call('EXPIRE', KEYS[1], 86400)
42return result
43";
44
45/// Redis-backed store. Connections are pooled internally (r2d2 for both
46/// single-node and cluster; the redis crate implements
47/// `r2d2::ManageConnection` for `ClusterClient` directly).
48pub struct RedisStore {
49    backend: RedisBackend,
50}
51
52enum RedisBackend {
53    Single(r2d2::Pool<RedisConnectionManager>),
54    Cluster(r2d2::Pool<redis::cluster::ClusterClient>),
55}
56
57struct RedisConnectionManager {
58    client: redis::Client,
59    timeout: std::time::Duration,
60}
61
62impl r2d2::ManageConnection for RedisConnectionManager {
63    type Connection = redis::Connection;
64    type Error = redis::RedisError;
65
66    fn connect(&self) -> Result<Self::Connection, Self::Error> {
67        let connection = self.client.get_connection_with_timeout(self.timeout)?;
68        connection.set_read_timeout(Some(self.timeout))?;
69        connection.set_write_timeout(Some(self.timeout))?;
70        Ok(connection)
71    }
72
73    fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), Self::Error> {
74        redis::cmd("PING").query(connection)
75    }
76
77    fn has_broken(&self, _connection: &mut Self::Connection) -> bool {
78        false
79    }
80}
81
82impl RedisStore {
83    /// Connect to `url` (e.g. `redis://127.0.0.1:6379`).
84    /// Comma-separated URLs select Redis Cluster mode.
85    pub fn connect(url: &str) -> Result<Self, StateError> {
86        let nodes: Vec<String> = url
87            .split(',')
88            .map(str::trim)
89            .filter(|node| !node.is_empty())
90            .map(str::to_owned)
91            .collect();
92        if nodes.len() > 1 {
93            // Round-45: previously one `ClusterConnection` behind a mutex —
94            // every quota/budget operation across all sessions serialized on
95            // a single socket, which under 10k-connection load stalled
96            // admission long enough to blow upstream timeouts (observed as
97            // 502 "upstream timed out" in the 10k SLA gate). Pool cluster
98            // connections exactly like the single-node path.
99            let client = redis::cluster::ClusterClientBuilder::new(nodes)
100                .connection_timeout(std::time::Duration::from_secs(2))
101                .response_timeout(std::time::Duration::from_secs(2))
102                .build()
103                .map_err(|e| StateError::Backend(e.to_string()))?;
104            let pool = r2d2::Pool::builder()
105                .max_size(32)
106                .connection_timeout(std::time::Duration::from_secs(2))
107                .build(client)
108                .map_err(|e| StateError::Backend(e.to_string()))?;
109            return Ok(Self {
110                backend: RedisBackend::Cluster(pool),
111            });
112        }
113        let client = redis::Client::open(url).map_err(|e| StateError::Backend(e.to_string()))?;
114        let manager = RedisConnectionManager {
115            client,
116            timeout: std::time::Duration::from_secs(2),
117        };
118        let pool = r2d2::Pool::builder()
119            .max_size(32)
120            .connection_timeout(std::time::Duration::from_secs(2))
121            .build(manager)
122            .map_err(|e| StateError::Backend(e.to_string()))?;
123        Ok(Self {
124            backend: RedisBackend::Single(pool),
125        })
126    }
127}
128
129fn add_on<C: redis::ConnectionLike>(conn: &mut C, key: &str, delta: u64) -> Result<u64, StateError> {
130    if delta > av_core::error::JCS_SAFE_MAX {
131        return Err(StateError::Overflow(key.to_owned()));
132    }
133    let value: i64 = redis::Script::new(ADD_LUA)
134        .key(key)
135        .arg(delta)
136        .arg(av_core::error::JCS_SAFE_MAX)
137        .invoke(conn)
138        .map_err(|e| StateError::Backend(e.to_string()))?;
139    if value < 0 {
140        return Err(StateError::Overflow(key.to_owned()));
141    }
142    u64::try_from(value).map_err(|_| StateError::Overflow(key.to_owned()))
143}
144
145fn get_on<C: redis::ConnectionLike>(conn: &mut C, key: &str) -> Result<u64, StateError> {
146    let value: Option<i64> = conn.get(key).map_err(|e| StateError::Backend(e.to_string()))?;
147    let value = match value {
148        None => 0,
149        Some(v) if v < 0 => return Err(StateError::Overflow(key.to_owned())),
150        Some(v) => u64::try_from(v).map_err(|_| StateError::Overflow(key.to_owned()))?,
151    };
152    if value > av_core::error::JCS_SAFE_MAX {
153        return Err(StateError::Overflow(key.to_owned()));
154    }
155    Ok(value)
156}
157
158fn spend_many_on<C: redis::ConnectionLike>(
159    conn: &mut C,
160    spends: &[Spend],
161) -> Result<Option<usize>, StateError> {
162    // Same duplicate-key guard as `InMemoryStore::try_spend_many`: the Lua
163    // script reads GET(key) once per iteration in the check phase, so two
164    // spends on the same key each see the pre-commit value and pass their
165    // independent limit checks, then the commit phase INCRBYs both.
166    let mut seen = std::collections::HashSet::with_capacity(spends.len());
167    for spend in spends {
168        if !seen.insert(spend.key.as_str()) {
169            return Err(StateError::Backend(format!(
170                "try_spend_many received duplicate key {:?}",
171                spend.key,
172            )));
173        }
174    }
175    for spend in spends {
176        if spend.amount > av_core::error::JCS_SAFE_MAX || spend.limit > av_core::error::JCS_SAFE_MAX {
177            return Err(StateError::Overflow(spend.key.clone()));
178        }
179    }
180    let script = redis::Script::new(TRY_SPEND_LUA);
181    let mut invocation = script.prepare_invoke();
182    for spend in spends {
183        invocation.key(&spend.key).arg(spend.amount).arg(spend.limit);
184    }
185    let failed: i64 = invocation
186        .invoke(conn)
187        .map_err(|e| StateError::Backend(e.to_string()))?;
188    if failed == 0 {
189        Ok(None)
190    } else {
191        usize::try_from(failed - 1)
192            .map(Some)
193            .map_err(|_| StateError::Backend(format!("invalid Lua failure index {failed}")))
194    }
195}
196
197impl StateStore for RedisStore {
198    fn add(&self, key: &str, delta: u64) -> Result<u64, StateError> {
199        match &self.backend {
200            RedisBackend::Single(pool) => add_on(
201                &mut pool.get().map_err(|e| StateError::Backend(e.to_string()))?,
202                key,
203                delta,
204            ),
205            RedisBackend::Cluster(pool) => add_on(
206                &mut *pool.get().map_err(|e| StateError::Backend(e.to_string()))?,
207                key,
208                delta,
209            ),
210        }
211    }
212
213    fn get(&self, key: &str) -> Result<u64, StateError> {
214        match &self.backend {
215            RedisBackend::Single(pool) => get_on(
216                &mut pool.get().map_err(|e| StateError::Backend(e.to_string()))?,
217                key,
218            ),
219            RedisBackend::Cluster(pool) => get_on(
220                &mut *pool.get().map_err(|e| StateError::Backend(e.to_string()))?,
221                key,
222            ),
223        }
224    }
225
226    fn try_spend(&self, key: &str, amount: u64, limit: u64) -> Result<bool, StateError> {
227        Ok(self
228            .try_spend_many(&[Spend {
229                key: key.to_owned(),
230                amount,
231                limit,
232            }])?
233            .is_none())
234    }
235
236    fn try_spend_many(&self, spends: &[Spend]) -> Result<Option<usize>, StateError> {
237        match &self.backend {
238            RedisBackend::Single(pool) => spend_many_on(
239                &mut pool.get().map_err(|e| StateError::Backend(e.to_string()))?,
240                spends,
241            ),
242            RedisBackend::Cluster(pool) => spend_many_on(
243                &mut *pool.get().map_err(|e| StateError::Backend(e.to_string()))?,
244                spends,
245            ),
246        }
247    }
248
249    fn remove(&self, key: &str) {
250        match &self.backend {
251            RedisBackend::Single(pool) => {
252                if let Ok(mut connection) = pool.get() {
253                    let _: Result<(), _> = connection.del(key);
254                }
255            }
256            RedisBackend::Cluster(pool) => {
257                if let Ok(mut connection) = pool.get() {
258                    let _: Result<(), _> = connection.del(key);
259                }
260            }
261        }
262    }
263
264    /// Round-33 F1: saturating refund via `DECRBY` + a MAX(0) clamp.
265    /// Best-effort — errors are silently swallowed so a Redis blip on
266    /// the compensation path can never turn a lost-race response into
267    /// a 5xx.
268    ///
269    /// Round-34 F1: NEVER resurrect a key that was already `DEL`'d
270    /// by a concurrent `remove_prefix`. The prior implementation
271    /// called `DECRBY` on the raw key: Redis initialises a missing
272    /// key to 0 first, so `DECRBY` returned `-amount` and the
273    /// `MAX(0)` clamp branch did `SET key 0` (no `EX`), producing
274    /// a permanent TTL-less key. Under the round-33 lost-claim-
275    /// plus-idle-close ordering (mcp_call sandbox-gate debit
276    /// races with the reconciler's clear_budget_state), the
277    /// refund path leaked one-to-three TTL-less keys per session
278    /// — attacker-choosable memory growth against the exact class
279    /// reconciler.rs:341-343 documents as impossible. Fix: gate
280    /// the whole DECRBY on `EXISTS`. If the session was cleared,
281    /// the budget is already gone and there is nothing to
282    /// compensate; the refund is a silent no-op. If the session
283    /// is alive, we DECRBY-clamp AND refresh the 24 h TTL to
284    /// match `TRY_SPEND_LUA` (a plain DECRBY does not refresh).
285    fn refund(&self, key: &str, amount: u64) {
286        // Redis DECRBY takes i64. Cap at i64::MAX so a caller passing
287        // u64::MAX cannot silently wrap into a negative value.
288        let amount = i64::try_from(amount).unwrap_or(i64::MAX);
289        let clamp_script = r#"
290            if redis.call('EXISTS', KEYS[1]) == 0 then
291                return 0
292            end
293            local new = redis.call('DECRBY', KEYS[1], ARGV[1])
294            if new < 0 then
295                redis.call('SET', KEYS[1], 0, 'EX', 86400)
296                return 0
297            end
298            redis.call('EXPIRE', KEYS[1], 86400)
299            return new
300        "#;
301        match &self.backend {
302            RedisBackend::Single(pool) => {
303                if let Ok(mut connection) = pool.get() {
304                    let _: Result<i64, _> = redis::Script::new(clamp_script)
305                        .key(key)
306                        .arg(amount)
307                        .invoke(&mut *connection);
308                }
309            }
310            RedisBackend::Cluster(pool) => {
311                if let Ok(mut connection) = pool.get() {
312                    let _: Result<i64, _> = redis::Script::new(clamp_script)
313                        .key(key)
314                        .arg(amount)
315                        .invoke(&mut *connection);
316                }
317            }
318        }
319    }
320}
av_state/redis_store.rs

av_state/
redis_store.rs
