Skip to main content

reddb_server/runtime/
execution_context.rs

1//! Per-thread runtime execution context.
2//!
3//! This module owns connection identity, auth/tenant scope, statement-local
4//! config/secret resolvers, and MVCC snapshot visibility helpers.
5
6use std::collections::{HashMap, HashSet};
7use std::sync::Arc;
8
9use crate::api::{RedDBError, RedDBResult};
10use crate::storage::RedDB;
11use reddb_types::Value;
12
13thread_local! {
14    /// Current connection id for the executing statement. Set by the
15    /// per-connection wrapper (stdio/gRPC handlers) before dispatching
16    /// into `execute_query`; falls back to `0` for embedded callers.
17    static CURRENT_CONN_ID: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
18
19    /// Authenticated user + role for the executing statement (Phase 2.5.2
20    /// RLS enforcement). Set by the transport middleware after validating
21    /// credentials (password / cert / oauth); unset means "anonymous" /
22    /// "embedded" — RLS policies degrade to the role-agnostic subset.
23    ///
24    /// `None` skips RLS injection entirely; `Some((username, role))`
25    /// passes `role` to `matching_rls_policies(table, Some(role), action)`.
26    static CURRENT_AUTH_IDENTITY: std::cell::RefCell<Option<(String, crate::auth::Role)>> =
27        const { std::cell::RefCell::new(None) };
28
29    /// MVCC snapshot scoped to the currently-executing statement (Phase
30    /// 2.3.2d PG parity). `execute_query` captures it on entry and drops
31    /// it on exit; every scan consults it via
32    /// `entity_visible_under_current_snapshot` to hide tuples whose xmin
33    /// hasn't committed or whose xmax already has.
34    ///
35    /// `None` means "pre-MVCC semantics" — the read path returns every
36    /// tuple regardless of xmin/xmax. All embedded callers that bypass
37    /// `execute_query` see this default.
38    static CURRENT_SNAPSHOT: std::cell::RefCell<Option<SnapshotContext>> =
39        const { std::cell::RefCell::new(None) };
40
41    /// Cheap presence flag for `CURRENT_SNAPSHOT`. Scan hot paths
42    /// poll this instead of `borrow()`-ing the RefCell on every
43    /// row — the common case (autocommit / no MVCC session) reads
44    /// one atomic `Cell<bool>` and short-circuits, saving ~10ns × N
45    /// rows on aggregate_group / select_range scans.
46    static HAS_SNAPSHOT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
47
48    /// Session-scoped tenant id for the current connection (Phase 2.5.3
49    /// multi-tenancy). Populated by `SET TENANT 'id'` or by transport
50    /// middleware after resolving tenant from auth claims. Read by the
51    /// `CURRENT_TENANT()` scalar function — RLS policies typically
52    /// combine it as `USING (tenant_id = CURRENT_TENANT())` to scope
53    /// every query to one tenant.
54    ///
55    /// `None` means "no tenant bound" — `CURRENT_TENANT()` returns
56    /// NULL, and RLS policies that gate on it hide every row.
57    static CURRENT_TENANT_ID: std::cell::RefCell<Option<String>> =
58        const { std::cell::RefCell::new(None) };
59
60    /// Statement-local config resolver. SQL expressions materialize the
61    /// `red_config` snapshot lazily on the first `$config.*`/`CONFIG()`
62    /// access, keeping ordinary statements on the zero-scan path.
63    static CURRENT_CONFIG_RESOLVER: std::cell::RefCell<Option<ConfigResolver>> =
64        const { std::cell::RefCell::new(None) };
65
66    /// Statement-local secret resolver. SQL expressions materialize the
67    /// vault KV snapshot lazily on first `$secret.*` access, then use
68    /// lock-free map reads for the rest of the statement.
69    static CURRENT_SECRET_RESOLVER: std::cell::RefCell<Option<SecretResolver>> =
70        const { std::cell::RefCell::new(None) };
71
72    /// Statement-local plain-KV resolver (#1602). Mirrors the secret
73    /// resolver but reads the non-encrypted `plain_kv` store and gates on
74    /// `kv:read`. Materialized lazily on the first `$kv.*` access.
75    static CURRENT_KV_RESOLVER: std::cell::RefCell<Option<KvResolver>> =
76        const { std::cell::RefCell::new(None) };
77}
78
79pub use reddb_engine::transaction::read::SnapshotContext;
80use reddb_engine::transaction::read::{record_serializable_read, visibility_check};
81
82/// Install a connection id on the current thread for the duration of a
83/// statement. The Transaction State module is keyed
84/// by this id so different connections can hold independent BEGINs.
85///
86/// Pub so transports (PG wire, gRPC, HTTP per-request spawners) and
87/// tests can emulate per-connection isolation. Call it once when
88/// binding the connection's worker thread; pair with
89/// `clear_current_connection_id` on teardown.
90pub fn set_current_connection_id(id: u64) {
91    CURRENT_CONN_ID.with(|c| c.set(id));
92}
93
94/// Reset the thread's connection id back to `0` (autocommit).
95pub fn clear_current_connection_id() {
96    CURRENT_CONN_ID.with(|c| c.set(0));
97}
98
99/// Read the connection id set by `set_current_connection_id`. Returns
100/// `0` when no wrapper installed one — auto-commit path.
101pub fn current_connection_id() -> u64 {
102    CURRENT_CONN_ID.with(|c| c.get())
103}
104
105/// Install the authenticated identity for the current thread (Phase 2.5.2
106/// RLS enforcement). Transport layers call this right after resolving
107/// auth so the query dispatch can fold RLS policies into the filter.
108pub fn set_current_auth_identity(username: String, role: crate::auth::Role) {
109    CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = Some((username, role)));
110}
111
112/// Clear the thread-local auth identity. Transports call this after the
113/// statement completes so pooled threads don't leak identities across
114/// requests.
115pub fn clear_current_auth_identity() {
116    CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = None);
117}
118
119/// Read the current-thread auth identity. `None` when no transport
120/// installed one (embedded mode / anonymous access).
121pub(crate) fn current_auth_identity() -> Option<(String, crate::auth::Role)> {
122    CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone())
123}
124
125/// Public probe of the thread-local auth identity for callers outside
126/// the `runtime` module (e.g. the AI credential resolver, which audits
127/// who triggered a secret read on behalf of a query).
128pub fn current_auth_identity_for_audit() -> Option<(String, crate::auth::Role)> {
129    current_auth_identity()
130}
131
132/// Install the session tenant id for the current thread (Phase 2.5.3
133/// multi-tenancy). Called by `SET TENANT 'id'` dispatch and by
134/// transport middleware that resolves tenant from auth claims (e.g.
135/// JWT `tenant` claim, HTTP header, subdomain).
136pub fn set_current_tenant(tenant_id: String) {
137    CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = Some(tenant_id));
138}
139
140/// Clear the current-thread tenant — `CURRENT_TENANT()` will then
141/// return NULL and any RLS policy gated on it will hide every row.
142pub fn clear_current_tenant() {
143    CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = None);
144}
145
146/// Read the current-thread tenant id, applying overrides in priority order:
147///   1. `WITHIN TENANT '<id>' …` per-statement override (highest)
148///   2. `SET LOCAL TENANT '<id>'` transaction-local override (consulted
149///      only when the current connection has an open transaction)
150///   3. `SET TENANT '<id>'` session-level thread-local
151///   4. `None` (deny-default for RLS).
152///
153/// The transaction-local layer is read through the runtime; an embedded
154/// helper crate that has no `RedDBRuntime` access still gets correct
155/// behaviour for layers 1, 3, and 4.
156pub fn current_tenant() -> Option<String> {
157    let inherited = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
158    if let Some(over) = current_scope_override() {
159        if over.tenant.is_active() {
160            return over.tenant.resolve(inherited);
161        }
162    }
163    if let Some(tx_local) = current_tx_local_tenant() {
164        return tx_local;
165    }
166    inherited
167}
168
169thread_local! {
170    /// Snapshot of the active connection's transaction-local tenant for
171    /// the current `execute_query` call. Outer `Some(_)` means "a
172    /// transaction-local tenant override is active for this call";
173    /// inner is the override's value (`Some(s)` overrides to `s`,
174    /// `None` overrides to NULL/cleared). Refreshed at the top of every
175    /// `execute_query` invocation and cleared by the RAII guard on
176    /// return so pooled connections cannot leak the override past the
177    /// statement that owns it.
178    static TX_LOCAL_TENANT: std::cell::RefCell<Option<Option<String>>> =
179        const { std::cell::RefCell::new(None) };
180}
181
182fn current_tx_local_tenant() -> Option<Option<String>> {
183    TX_LOCAL_TENANT.with(|cell| cell.borrow().clone())
184}
185
186/// Recognise `SET LOCAL TENANT '<id>'` / `SET LOCAL TENANT NULL` —
187/// returns `Ok(Some(Some(id)))` for an explicit value, `Ok(Some(None))`
188/// for an explicit NULL clear, `Ok(None)` when the input is not a
189/// `SET LOCAL TENANT` statement at all, and `Err` when the prefix
190/// matches but the value is malformed.
191pub(crate) fn parse_set_local_tenant(query: &str) -> RedDBResult<Option<Option<String>>> {
192    let mut tokens = query.split_ascii_whitespace();
193    let Some(w1) = tokens.next() else {
194        return Ok(None);
195    };
196    if !w1.eq_ignore_ascii_case("SET") {
197        return Ok(None);
198    }
199    let Some(w2) = tokens.next() else {
200        return Ok(None);
201    };
202    if !w2.eq_ignore_ascii_case("LOCAL") {
203        return Ok(None);
204    }
205    let Some(w3) = tokens.next() else {
206        return Ok(None);
207    };
208    if !w3.eq_ignore_ascii_case("TENANT") {
209        return Ok(None);
210    }
211    let rest: String = tokens.collect::<Vec<_>>().join(" ");
212    let rest = rest.trim().trim_end_matches(';').trim();
213    let value_str = rest.strip_prefix('=').map(|s| s.trim()).unwrap_or(rest);
214    if value_str.is_empty() {
215        return Err(RedDBError::Query(
216            "SET LOCAL TENANT expects a string literal or NULL".to_string(),
217        ));
218    }
219    if value_str.eq_ignore_ascii_case("NULL") {
220        return Ok(Some(None));
221    }
222    if value_str.starts_with('\'') && value_str.ends_with('\'') && value_str.len() >= 2 {
223        let inner = &value_str[1..value_str.len() - 1];
224        return Ok(Some(Some(inner.to_string())));
225    }
226    Err(RedDBError::Query(format!(
227        "SET LOCAL TENANT expects a string literal or NULL, got `{value_str}`"
228    )))
229}
230
231pub(crate) struct TxLocalTenantGuard;
232
233impl TxLocalTenantGuard {
234    pub fn install(value: Option<Option<String>>) -> Self {
235        TX_LOCAL_TENANT.with(|cell| *cell.borrow_mut() = value);
236        Self
237    }
238}
239
240impl Drop for TxLocalTenantGuard {
241    fn drop(&mut self) {
242        TX_LOCAL_TENANT.with(|cell| *cell.borrow_mut() = None);
243    }
244}
245
246thread_local! {
247    /// Stack of `WITHIN ... <stmt>` overrides active on the current
248    /// thread. Every entry corresponds to one in-flight `execute_query`
249    /// call that started with a `WITHIN` prefix; the entry is pushed
250    /// before dispatch and popped before the call returns. The stack
251    /// shape supports nested invocations (e.g. a view body that itself
252    /// re-enters execute_query).
253    static SCOPE_OVERRIDES: std::cell::RefCell<Vec<crate::runtime::within_clause::ScopeOverride>> =
254        const { std::cell::RefCell::new(Vec::new()) };
255}
256
257pub(crate) fn push_scope_override(over: crate::runtime::within_clause::ScopeOverride) {
258    SCOPE_OVERRIDES.with(|cell| cell.borrow_mut().push(over));
259}
260
261pub(crate) fn pop_scope_override() {
262    SCOPE_OVERRIDES.with(|cell| {
263        cell.borrow_mut().pop();
264    });
265}
266
267pub(crate) fn current_scope_override() -> Option<crate::runtime::within_clause::ScopeOverride> {
268    SCOPE_OVERRIDES.with(|cell| cell.borrow().last().cloned())
269}
270
271/// Cheap probe: is any `WITHIN …` scope override active on this
272/// thread? The fast-path needs to know without paying for the full
273/// `.last().cloned()` allocation — just peek at stack length.
274pub(crate) fn has_scope_override_active() -> bool {
275    SCOPE_OVERRIDES.with(|cell| !cell.borrow().is_empty())
276}
277
278/// RAII guard pairing `push_scope_override` with the matching pop, so
279/// the stack stays balanced even when the inner `execute_query` returns
280/// early via `?`.
281pub(crate) struct ScopeOverrideGuard;
282
283impl ScopeOverrideGuard {
284    pub fn install(over: crate::runtime::within_clause::ScopeOverride) -> Self {
285        push_scope_override(over);
286        Self
287    }
288}
289
290impl Drop for ScopeOverrideGuard {
291    fn drop(&mut self) {
292        pop_scope_override();
293    }
294}
295
296/// Read the current-thread auth identity, honouring per-statement
297/// `WITHIN ... USER '<u>' AS ROLE '<r>'` overrides. The override only
298/// supplies projected strings — it never grants additional privilege —
299/// so callers that need to make authorisation decisions must read from
300/// the underlying `current_auth_identity()` directly.
301pub(crate) fn current_user_projected() -> Option<String> {
302    let inherited = current_auth_identity().map(|(u, _)| u);
303    if let Some(over) = current_scope_override() {
304        if over.user.is_active() {
305            return over.user.resolve(inherited);
306        }
307    }
308    inherited
309}
310
311pub(crate) fn current_role_projected() -> Option<String> {
312    let inherited = current_auth_identity().map(|(_, r)| format!("{r:?}").to_lowercase());
313    if let Some(over) = current_scope_override() {
314        if over.role.is_active() {
315            return over.role.resolve(inherited);
316        }
317    }
318    inherited
319}
320
321pub(crate) fn current_secret_value(path: &str) -> Option<Value> {
322    let key = path.to_ascii_lowercase();
323    CURRENT_SECRET_RESOLVER.with(|cell| {
324        let mut resolver = cell.borrow_mut();
325        let resolver = resolver.as_mut()?;
326        if let Some(value) = resolver.values.get(&key) {
327            return value.clone();
328        }
329        let value = resolver.runtime.query_secret_value(&key).ok().flatten();
330        resolver.values.insert(key, value.clone());
331        value
332    })
333}
334
335/// Decrypt only while evaluating a query. Stored and projected values retain
336/// their Secret representation, so formatters cannot expose plaintext.
337pub(crate) fn secret_query_input(value: Value) -> Option<Value> {
338    if !matches!(value, Value::Secret(_)) {
339        return Some(value);
340    }
341    CURRENT_SECRET_RESOLVER.with(|cell| {
342        let resolver = cell.borrow();
343        resolver.as_ref()?.runtime.secret_query_input(value)
344    })
345}
346
347pub(crate) fn secret_query_output(value: Value) -> crate::RedDBResult<Value> {
348    if value.is_null() || matches!(value, Value::Secret(_)) {
349        return Ok(value);
350    }
351    CURRENT_SECRET_RESOLVER.with(|cell| {
352        let resolver = cell.borrow();
353        let resolver = resolver.as_ref().ok_or_else(|| {
354            RedDBError::Query("secret expression requires an active statement".into())
355        })?;
356        resolver.runtime.secret_query_output(value)
357    })
358}
359
360#[cfg(test)]
361fn secret_value_from_snapshot(values: &HashMap<String, String>, key: &str) -> Option<String> {
362    if key.starts_with("red.secret.") {
363        return None;
364    }
365    values.get(key).cloned()
366}
367
368struct SecretResolver {
369    runtime: super::RedDBRuntime,
370    values: HashMap<String, Option<Value>>,
371}
372
373pub(crate) struct SecretStoreGuard {
374    previous: Option<SecretResolver>,
375}
376
377impl SecretStoreGuard {
378    pub(super) fn install(runtime: &super::RedDBRuntime) -> Self {
379        let previous = CURRENT_SECRET_RESOLVER.with(|cell| {
380            cell.replace(Some(SecretResolver {
381                runtime: runtime.clone(),
382                values: HashMap::new(),
383            }))
384        });
385        Self { previous }
386    }
387}
388
389impl Drop for SecretStoreGuard {
390    fn drop(&mut self) {
391        let previous = self.previous.take();
392        CURRENT_SECRET_RESOLVER.with(|cell| {
393            cell.replace(previous);
394        });
395    }
396}
397
398/// Resolve `$kv.<path>` from the plain (non-encrypted) KV store (#1602).
399///
400/// The parser desugars `$kv.<path>` to `__KV_REF("red.kv/<path>")`; this
401/// strips the `red.kv/` prefix and looks up the bare key in the
402/// statement-local snapshot, gating each hit on `kv:read`. Returns `None`
403/// (→ SQL NULL) when the key is absent or the principal is denied — never
404/// an error, mirroring the secret resolver's deny-to-NULL behaviour.
405pub(crate) fn current_kv_value(path: &str) -> Option<String> {
406    let key = path.to_ascii_lowercase();
407    CURRENT_KV_RESOLVER.with(|cell| {
408        let mut resolver = cell.borrow_mut();
409        let resolver = resolver.as_mut()?;
410        if resolver.values.is_none() {
411            resolver.values = resolver
412                .store
413                .as_ref()
414                .map(|store| store.plain_kv_snapshot());
415        }
416        let values = resolver.values.as_ref()?;
417        // Accept both the namespaced `red.kv/<key>` form (from the
418        // `$kv.*` desugar) and the bare key (from `update_current_kv_value`
419        // after a `SET KV`).
420        let found = values
421            .get(&key)
422            .map(|value| (key.as_str(), value))
423            .or_else(|| {
424                key.strip_prefix("red.kv/")
425                    .and_then(|rest| values.get(rest).map(|value| (rest, value)))
426            })?;
427        if !resolver.can_read(found.0) {
428            return None;
429        }
430        Some(found.1.clone())
431    })
432}
433
434struct KvResolver {
435    store: Option<Arc<crate::auth::store::AuthStore>>,
436    values: Option<HashMap<String, String>>,
437    identity: Option<(String, crate::auth::Role, Option<String>)>,
438}
439
440impl KvResolver {
441    fn can_read(&self, key: &str) -> bool {
442        let Some(store) = &self.store else {
443            return true;
444        };
445        let Some((username, role, tenant)) = &self.identity else {
446            return true;
447        };
448        let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
449        let mut resource =
450            crate::auth::policies::ResourceRef::new("kv".to_string(), key.to_string());
451        if let Some(tenant) = tenant {
452            resource = resource.with_tenant(tenant.clone());
453        }
454        let ctx = crate::auth::policies::EvalContext {
455            principal_tenant: tenant.clone(),
456            current_tenant: tenant.clone(),
457            peer_ip: None,
458            mfa_present: false,
459            now_ms: crate::auth::now_ms(),
460            principal_is_admin_role: *role == crate::auth::Role::Admin,
461            principal_is_platform_scoped: tenant.is_none(),
462        };
463        store.check_policy_authz_with_role(&principal, "kv:read", &resource, &ctx, *role)
464    }
465}
466
467pub(crate) struct KvStoreGuard {
468    // Boxed to keep the stack footprint of nested guards small. Each
469    // StatementFrameGuards is live for the full recursive call depth of a
470    // materialized-view refresh chain; unboxed KvResolver (≈120 B) would
471    // push those deep paths past the 8 MB stack limit.
472    previous: Option<Box<KvResolver>>,
473}
474
475impl KvStoreGuard {
476    pub(super) fn install(store: Option<Arc<crate::auth::store::AuthStore>>) -> Self {
477        let previous = CURRENT_KV_RESOLVER.with(|cell| {
478            cell.replace(Some(KvResolver {
479                store,
480                values: None,
481                identity: current_auth_identity().map(|(username, role)| {
482                    let tenant = current_tenant();
483                    (username, role, tenant)
484                }),
485            }))
486        });
487        Self {
488            previous: previous.map(Box::new),
489        }
490    }
491}
492
493impl Drop for KvStoreGuard {
494    fn drop(&mut self) {
495        let previous = self.previous.take().map(|b| *b);
496        CURRENT_KV_RESOLVER.with(|cell| {
497            cell.replace(previous);
498        });
499    }
500}
501
502/// Update the statement-local KV snapshot after a `SET KV` / `DELETE KV`
503/// so a subsequent `$kv.*` read in the same statement observes the write.
504pub(crate) fn update_current_kv_value(path: &str, value: Option<String>) {
505    let key = path.to_ascii_lowercase();
506    CURRENT_KV_RESOLVER.with(|cell| {
507        if let Some(resolver) = cell.borrow_mut().as_mut() {
508            let Some(values) = resolver.values.as_mut() else {
509                return;
510            };
511            match value {
512                Some(value) => {
513                    values.insert(key, value);
514                }
515                None => {
516                    values.remove(&key);
517                }
518            }
519        }
520    });
521}
522
523pub(crate) fn current_config_value(path: &str) -> Option<Value> {
524    let key = path.to_ascii_lowercase();
525    CURRENT_CONFIG_RESOLVER.with(|cell| {
526        let mut resolver = cell.borrow_mut();
527        let resolver = resolver.as_mut()?;
528        if resolver.values.is_none() {
529            resolver.values = Some(latest_config_snapshot(&resolver.db));
530        }
531        // #1370 — `$config.<path>` desugars to `CONFIG("red.config/<path>")`,
532        // but `SET CONFIG <path>` stores under the bare key. Mirror the secret
533        // resolver: after the namespaced key, fall back to the stripped bare
534        // key, then the dotted `red.config.<rest>` legacy form. Track the key
535        // that matched so the `config:read` gate authorizes the real target.
536        let found = {
537            let values = resolver.values.as_ref()?;
538            values
539                .get(&key)
540                .map(|value| (key.clone(), value.clone()))
541                .or_else(|| {
542                    key.strip_prefix("red.config/").and_then(|rest| {
543                        values
544                            .get(rest)
545                            .map(|value| (rest.to_string(), value.clone()))
546                            .or_else(|| {
547                                let legacy = format!("red.config.{rest}");
548                                values.get(&legacy).map(|value| (legacy, value.clone()))
549                            })
550                    })
551                })
552        }?;
553        // #1743 — hard-block the reserved `red.config.*` namespace and gate the
554        // read on `config:read`, mirroring the `$secret` sibling. Denied reads
555        // resolve to `None` (→ SQL NULL), never an error.
556        if !resolver.can_read(&found.0) {
557            return None;
558        }
559        Some(found.1)
560    })
561}
562
563pub(crate) fn update_current_config_value(path: &str, value: Value) {
564    let key = path.to_ascii_lowercase();
565    CURRENT_CONFIG_RESOLVER.with(|cell| {
566        if let Some(resolver) = cell.borrow_mut().as_mut() {
567            if let Some(values) = resolver.values.as_mut() {
568                values.insert(key, value);
569            }
570        }
571    });
572}
573
574pub(crate) fn update_current_secret_value(_path: &str, _value: Option<String>) {
575    CURRENT_SECRET_RESOLVER.with(|cell| {
576        if let Some(resolver) = cell.borrow_mut().as_mut() {
577            resolver.values.clear();
578        }
579    });
580}
581
582#[cfg(test)]
583mod tests {
584    use super::*;
585
586    /// Companion to `raw_iteration_confined_to_maintenance_allowlist`:
587    /// the snapshot thread-local is an implementation detail of the
588    /// statement frame (issue #2138). It must stay private to this module,
589    /// so the only way any other module can observe ambient snapshot state
590    /// is `capture_current_snapshot`, which hands back an owned `Option`
591    /// the caller then passes explicitly into the scan API.
592    #[test]
593    fn snapshot_thread_local_stays_private_to_the_statement_frame() {
594        let src = include_str!("execution_context.rs");
595
596        for name in ["CURRENT_SNAPSHOT", "HAS_SNAPSHOT"] {
597            assert!(
598                !src.contains(&format!("pub static {name}"))
599                    && !src.contains(&format!("pub(crate) static {name}")),
600                "{name} must stay private to execution_context; expose an explicit \
601                 Option<SnapshotContext> via capture_current_snapshot instead"
602            );
603        }
604
605        assert!(
606            src.contains("the only sanctioned ambient read"),
607            "capture_current_snapshot must keep documenting that it is the single \
608             sanctioned ambient read of the snapshot thread-local"
609        );
610    }
611
612    #[test]
613    fn secret_snapshot_hard_blocks_system_keys() {
614        let values = HashMap::from([
615            ("red.secret.aes_key".into(), "internal".into()),
616            ("acme.key".into(), "user-value".into()),
617        ]);
618        assert_eq!(
619            secret_value_from_snapshot(&values, "red.secret.aes_key"),
620            None
621        );
622        assert_eq!(
623            secret_value_from_snapshot(&values, "acme.key").as_deref(),
624            Some("user-value")
625        );
626    }
627
628    fn with_kv_values<T>(values: HashMap<String, String>, f: impl FnOnce() -> T) -> T {
629        CURRENT_KV_RESOLVER.with(|cell| {
630            cell.replace(Some(KvResolver {
631                store: None,
632                values: Some(values),
633                identity: None,
634            }));
635        });
636        let result = f();
637        CURRENT_KV_RESOLVER.with(|cell| {
638            cell.replace(None);
639        });
640        result
641    }
642
643    #[test]
644    fn current_kv_value_resolves_namespaced_and_bare_keys() {
645        let values = HashMap::from([("acme.key".to_string(), "plain-value".to_string())]);
646
647        with_kv_values(values, || {
648            // `$kv.acme.key` desugars to `red.kv/acme.key`.
649            assert_eq!(
650                current_kv_value("red.kv/acme.key").as_deref(),
651                Some("plain-value")
652            );
653            // Bare-key lookup (post-`SET KV` snapshot update) also resolves.
654            assert_eq!(current_kv_value("acme.key").as_deref(), Some("plain-value"));
655            // Unknown keys resolve to None (→ SQL NULL), never an error.
656            assert_eq!(current_kv_value("red.kv/missing.key"), None);
657        });
658    }
659
660    #[test]
661    fn config_key_read_allowed_hard_blocks_reserved_namespace() {
662        // The reserved `red.config.*` system namespace is denied role- and
663        // store-independently, mirroring the `red.secret.*` block (#1743).
664        assert!(!config_key_read_allowed(
665            None,
666            None,
667            "red.config/red.config.aes_key"
668        ));
669        assert!(!config_key_read_allowed(None, None, "red.config.aes_key"));
670        // Ordinary config keys pass when no store is installed (non-IAM path).
671        assert!(config_key_read_allowed(
672            None,
673            None,
674            "red.config/ai.openrouter.default.key"
675        ));
676        assert!(config_key_read_allowed(None, None, "acme.flags.beta"));
677    }
678
679    #[test]
680    fn is_config_collection_matches_system_config_stores() {
681        assert!(is_config_collection("red_config"));
682        assert!(is_config_collection("RED_CONFIG"));
683        assert!(is_config_collection("red.config"));
684        assert!(!is_config_collection("users"));
685        assert!(!is_config_collection("red_kv"));
686    }
687
688    #[test]
689    fn update_current_kv_value_reflects_in_resolver() {
690        with_kv_values(HashMap::new(), || {
691            assert_eq!(current_kv_value("red.kv/feature.flag"), None);
692            update_current_kv_value("feature.flag", Some("on".to_string()));
693            assert_eq!(
694                current_kv_value("red.kv/feature.flag").as_deref(),
695                Some("on")
696            );
697            update_current_kv_value("feature.flag", None);
698            assert_eq!(current_kv_value("red.kv/feature.flag"), None);
699        });
700    }
701}
702
703fn latest_config_snapshot(db: &RedDB) -> HashMap<String, Value> {
704    let mut latest: HashMap<String, (u64, Value)> = HashMap::new();
705
706    if let Some(manager) = db.store().get_collection("red_config") {
707        manager.for_each_entity(|entity| {
708            let Some(row) = entity.data.as_row() else {
709                return true;
710            };
711            let Some(Value::Text(key)) = row.get_field("key") else {
712                return true;
713            };
714            let value = row.get_field("value").cloned().unwrap_or(Value::Null);
715            let id = entity.id.raw();
716            let key = key.to_ascii_lowercase();
717            insert_latest_config_value(&mut latest, key.clone(), id, value.clone());
718            if let Some(rest) = key.strip_prefix("red.config.") {
719                insert_latest_config_value(&mut latest, format!("red.config/{rest}"), id, value);
720            }
721            true
722        });
723    }
724
725    if let Some(manager) = db.store().get_collection("red.config") {
726        manager.for_each_entity(|entity| {
727            let Some(row) = entity.data.as_row() else {
728                return true;
729            };
730            if matches!(row.get_field("tombstone"), Some(Value::Boolean(true))) {
731                return true;
732            }
733            let Some(Value::Text(key)) = row.get_field("key") else {
734                return true;
735            };
736            let value = row.get_field("value").cloned().unwrap_or(Value::Null);
737            insert_latest_config_value(
738                &mut latest,
739                format!("red.config/{}", key.to_ascii_lowercase()),
740                entity.id.raw(),
741                value,
742            );
743            true
744        });
745    }
746
747    latest
748        .into_iter()
749        .map(|(key, (_, value))| (key, value))
750        .collect()
751}
752
753fn insert_latest_config_value(
754    latest: &mut HashMap<String, (u64, Value)>,
755    key: String,
756    id: u64,
757    value: Value,
758) {
759    match latest.get(&key) {
760        Some((prev_id, _)) if *prev_id > id => {}
761        _ => {
762            latest.insert(key, (id, value));
763        }
764    }
765}
766
767struct ConfigResolver {
768    db: Arc<RedDB>,
769    store: Option<Arc<crate::auth::store::AuthStore>>,
770    identity: Option<(String, crate::auth::Role, Option<String>)>,
771    values: Option<HashMap<String, Value>>,
772}
773
774impl ConfigResolver {
775    fn can_read(&self, key: &str) -> bool {
776        config_key_read_allowed(self.store.as_ref(), self.identity.as_ref(), key)
777    }
778}
779
780/// The reserved `red.config.*` system-config namespace. Never expose it via
781/// the inline `$config`/`CONFIG()`/`KV()` resolver regardless of IAM role,
782/// mirroring the role-independent `red.secret.*` hard-block (#1743).
783fn is_reserved_config_key(key: &str) -> bool {
784    let key = key.strip_prefix("red.config/").unwrap_or(key);
785    key.starts_with("red.config.")
786        || key == "red.config"
787        || key == super::function_catalog::REGISTRY_KEY
788}
789
790/// The system config-store collections (`red_config` / `red.config`) reached
791/// by the inline resolver. A `KV(<collection>, key)` read of one of these is
792/// gated exactly like `$config`; any other collection is unaffected (#1743).
793fn is_config_collection(collection: &str) -> bool {
794    matches!(
795        collection.to_ascii_lowercase().as_str(),
796        "red_config" | "red.config"
797    )
798}
799
800fn config_resource_target(key: &str) -> String {
801    let bare = key.strip_prefix("red.config/").unwrap_or(key);
802    format!("red.config/{bare}")
803}
804
805/// Shared inline-config read gate: a role-independent hard-block on the
806/// reserved `red.config.*` namespace, then a `config:read` capability check
807/// mirroring `SHOW CONFIG` / the `$secret` sibling (#1743). Absent a store or
808/// with IAM authorization disabled, only the hard-block applies so existing
809/// non-IAM `$config` reads keep working.
810fn config_key_read_allowed(
811    store: Option<&Arc<crate::auth::store::AuthStore>>,
812    identity: Option<&(String, crate::auth::Role, Option<String>)>,
813    key: &str,
814) -> bool {
815    if is_reserved_config_key(key) {
816        return false;
817    }
818    let Some(store) = store else {
819        return true;
820    };
821    if !store.iam_authorization_enabled() {
822        return true;
823    }
824    let Some((username, role, tenant)) = identity else {
825        return true;
826    };
827    let principal = crate::auth::UserId::from_parts(tenant.as_deref(), username);
828    let mut resource =
829        crate::auth::policies::ResourceRef::new("config".to_string(), config_resource_target(key));
830    if let Some(tenant) = tenant {
831        resource = resource.with_tenant(tenant.clone());
832    }
833    let ctx = crate::auth::policies::EvalContext {
834        principal_tenant: tenant.clone(),
835        current_tenant: tenant.clone(),
836        peer_ip: None,
837        mfa_present: false,
838        now_ms: crate::auth::now_ms(),
839        principal_is_admin_role: *role == crate::auth::Role::Admin,
840        principal_is_platform_scoped: tenant.is_none(),
841    };
842    store.check_policy_authz_with_role(&principal, "config:read", &resource, &ctx, *role)
843}
844
845/// Whether the current principal may read config `key` through the inline
846/// `$config`/`CONFIG()`/`KV('red_config', …)` resolver. Used by the db-fallback
847/// lookup paths that bypass [`current_config_value`] (#1743).
848pub(crate) fn config_read_permitted(key: &str) -> bool {
849    let key = key.to_ascii_lowercase();
850    CURRENT_CONFIG_RESOLVER.with(|cell| match cell.borrow().as_ref() {
851        Some(resolver) => resolver.can_read(&key),
852        None => !is_reserved_config_key(&key),
853    })
854}
855
856/// Gate a raw `KV(collection, key)` read: config system collections go through
857/// the `$config` gate; every other collection is unaffected (#1743).
858pub(crate) fn kv_read_permitted(collection: &str, key: &str) -> bool {
859    if is_config_collection(collection) {
860        config_read_permitted(key)
861    } else {
862        true
863    }
864}
865
866pub(crate) struct ConfigSnapshotGuard {
867    previous: Option<ConfigResolver>,
868}
869
870impl ConfigSnapshotGuard {
871    pub(super) fn install(
872        db: Arc<RedDB>,
873        store: Option<Arc<crate::auth::store::AuthStore>>,
874    ) -> Self {
875        let identity = current_auth_identity().map(|(username, role)| {
876            let tenant = current_tenant();
877            (username, role, tenant)
878        });
879        let previous = CURRENT_CONFIG_RESOLVER.with(|cell| {
880            cell.replace(Some(ConfigResolver {
881                db,
882                store,
883                identity,
884                values: None,
885            }))
886        });
887        Self { previous }
888    }
889}
890
891impl Drop for ConfigSnapshotGuard {
892    fn drop(&mut self) {
893        let previous = self.previous.take();
894        CURRENT_CONFIG_RESOLVER.with(|cell| {
895            cell.replace(previous);
896        });
897    }
898}
899
900/// Install the MVCC snapshot used by the current thread for the duration
901/// of one statement. Paired with `clear_current_snapshot()` — callers
902/// should prefer the `CurrentSnapshotGuard` RAII wrapper so early returns
903/// still clean up.
904pub fn set_current_snapshot(ctx: SnapshotContext) {
905    CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = Some(ctx));
906    HAS_SNAPSHOT.with(|c| c.set(true));
907}
908
909pub fn clear_current_snapshot() {
910    CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = None);
911    HAS_SNAPSHOT.with(|c| c.set(false));
912}
913
914/// Drop-guard that restores the previous snapshot on scope exit. Safe to
915/// nest — each statement saves the caller's snapshot and puts it back
916/// instead of blindly clearing, so a top-level `execute_query` called
917/// from inside another statement dispatch (e.g. vector source subqueries)
918/// doesn't strip visibility from the outer scan.
919pub(crate) struct CurrentSnapshotGuard {
920    previous: Option<SnapshotContext>,
921}
922
923impl CurrentSnapshotGuard {
924    pub(crate) fn install(ctx: SnapshotContext) -> Self {
925        let previous = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
926        set_current_snapshot(ctx);
927        Self { previous }
928    }
929}
930
931impl Drop for CurrentSnapshotGuard {
932    fn drop(&mut self) {
933        let prev = self.previous.take();
934        let has = prev.is_some();
935        CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = prev);
936        HAS_SNAPSHOT.with(|c| c.set(has));
937    }
938}
939
940/// Is this entity visible under the current thread's MVCC snapshot?
941///
942/// Returns `true` (no filtering) when no snapshot is installed — that
943/// path is used by embedded callers and by operations that intentionally
944/// bypass MVCC (VACUUM, snapshot export, admin introspection).
945///
946/// When a snapshot is installed the result is
947///   `snapshot.sees(xmin, xmax) && !mgr.is_aborted(xmin) && !xmax_half_abort`
948/// where `xmax_half_abort` re-grants visibility for tuples whose
949/// deleting transaction rolled back.
950#[inline]
951pub fn entity_visible_under_current_snapshot(
952    entity: &crate::storage::unified::entity::UnifiedEntity,
953) -> bool {
954    // Moderation visibility gate (#1274, ADR 0057). A row carrying the
955    // moderation status marker — quarantine-pending or rejected-tombstone
956    // — is hidden from every normal read, on top of MVCC visibility. The
957    // marker lives on the row itself, so the check is a single field probe
958    // and rides the existing per-row visibility chokepoint rather than
959    // adding a separate filter pass to each scan call-site.
960    if crate::runtime::ai::moderation::entity_moderation_hidden(entity) {
961        return false;
962    }
963    // Fast path — one `Cell<bool>` read, no RefCell borrow. Autocommit
964    // reads (no active MVCC transaction) still hide superseded physical
965    // versions while avoiding a full snapshot-context lookup.
966    // This runs on every row of every scan; the slow path only fires
967    // inside an explicit transaction.
968    if !HAS_SNAPSHOT.with(|c| c.get()) {
969        return entity.xmax == 0;
970    }
971    CURRENT_SNAPSHOT.with(|cell| {
972        let guard = cell.borrow();
973        let Some(ctx) = guard.as_ref() else {
974            return true;
975        };
976        let visible = visibility_check(ctx, entity.xmin, entity.xmax);
977        if visible {
978            record_serializable_read(ctx, entity);
979        }
980        visible
981    })
982}
983
984/// Direct visibility check from raw `(xmin, xmax)` — bypasses the
985/// entity borrow for callers that already decomposed the tuple (e.g.
986/// pre-materialized scan caches). Same semantics as
987/// `entity_visible_under_current_snapshot`.
988#[inline]
989pub(crate) fn xids_visible_under_current_snapshot(xmin: u64, xmax: u64) -> bool {
990    if !HAS_SNAPSHOT.with(|c| c.get()) {
991        return true;
992    }
993    CURRENT_SNAPSHOT.with(|cell| {
994        let guard = cell.borrow();
995        let Some(ctx) = guard.as_ref() else {
996            return true;
997        };
998        visibility_check(ctx, xmin, xmax)
999    })
1000}
1001
1002/// Clone the current thread's snapshot context. Parallel scan paths
1003/// (`query_all_zoned` with `std::thread::scope`) call this on the main
1004/// thread *before* spawning workers so the captured `SnapshotContext`
1005/// can be moved into every worker closure. Worker threads do not
1006/// inherit thread-locals, so calling `entity_visible_under_current_snapshot`
1007/// from inside a spawned closure would silently skip the filter.
1008///
1009/// This is the only sanctioned ambient read of the snapshot thread-local
1010/// (issue #2138). `CURRENT_SNAPSHOT` and `HAS_SNAPSHOT` are private to this
1011/// module and stay that way: they exist for statement-frame entry and for
1012/// detecting a frameless read, nothing else. Everything downstream must
1013/// take the resulting `Option<SnapshotContext>` as an explicit argument —
1014/// that is what makes `SegmentManager::scan*` thread-independent.
1015///
1016/// Statement-frame build also reads it, so a nested sub-execution inside
1017/// one statement (CTE prelude, materialized-view refresh, inline graph
1018/// TVF) inherits the enclosing snapshot instead of minting a fresh one
1019/// (#2183).
1020///
1021/// The `Option` is load-bearing: `None` means no statement frame is
1022/// installed — the low-level callers that intentionally bypass MVCC
1023/// (VACUUM, snapshot export, admin introspection) — and the scan API
1024/// keeps a documented fallback for that case rather than assuming a frame
1025/// always exists.
1026pub fn capture_current_snapshot() -> Option<SnapshotContext> {
1027    CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone())
1028}
1029
1030/// Whether the active read snapshot may need historical tuple versions
1031/// that the current secondary indexes cannot prove. Index paths can still
1032/// recheck visible candidates, but only a heap scan can discover versions
1033/// whose indexed value was changed or deleted after this snapshot.
1034pub(crate) fn current_snapshot_requires_index_fallback() -> bool {
1035    if !HAS_SNAPSHOT.with(|c| c.get()) {
1036        return false;
1037    }
1038    CURRENT_SNAPSHOT.with(|cell| {
1039        cell.borrow()
1040            .as_ref()
1041            .is_some_and(|ctx| ctx.requires_index_fallback)
1042    })
1043}
1044
1045/// Frozen MVCC + identity context for callers that need to reinstall
1046/// the same view across thread-local boundaries — long-lived cursors,
1047/// background batchers, anything that detaches from the dispatch path
1048/// and re-enters later.
1049///
1050/// The bundle bakes in the three thread-locals every read path
1051/// consults: `SnapshotContext` (MVCC visibility), the auth identity
1052/// (RLS policy gate), and the tenant id (RLS scope). A FETCH that
1053/// reinstalls the bundle sees exactly the same rows as the DECLARE
1054/// would have, regardless of writes that landed in between.
1055///
1056/// Cheap to clone — `SnapshotContext` is a clone of three
1057/// `Arc`-backed fields, identity is a `(String, Role)`, tenant is a
1058/// `String`. None of these contend with the read path.
1059#[derive(Clone, Default)]
1060pub struct SnapshotBundle {
1061    pub snapshot: Option<SnapshotContext>,
1062    pub auth: Option<(String, crate::auth::Role)>,
1063    pub tenant: Option<String>,
1064}
1065
1066/// Capture the three read-path thread-locals into a `SnapshotBundle`.
1067/// Pairs with `with_snapshot_bundle` for re-entry.
1068pub fn snapshot_bundle() -> SnapshotBundle {
1069    SnapshotBundle {
1070        snapshot: capture_current_snapshot(),
1071        auth: current_auth_identity(),
1072        tenant: CURRENT_TENANT_ID.with(|cell| cell.borrow().clone()),
1073    }
1074}
1075
1076/// Reinstall a captured `SnapshotBundle` for the duration of `f`.
1077/// Restores the caller's previous thread-locals on exit (panic-safe via
1078/// the explicit guard struct so a panic in `f` cannot leak the
1079/// installed identity into the worker's next request).
1080pub fn with_snapshot_bundle<R>(bundle: &SnapshotBundle, f: impl FnOnce() -> R) -> R {
1081    struct Guard {
1082        prev_snapshot: Option<SnapshotContext>,
1083        prev_auth: Option<(String, crate::auth::Role)>,
1084        prev_tenant: Option<String>,
1085    }
1086    impl Drop for Guard {
1087        fn drop(&mut self) {
1088            let snap = self.prev_snapshot.take();
1089            let has = snap.is_some();
1090            CURRENT_SNAPSHOT.with(|cell| *cell.borrow_mut() = snap);
1091            HAS_SNAPSHOT.with(|c| c.set(has));
1092            CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = self.prev_auth.take());
1093            CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = self.prev_tenant.take());
1094        }
1095    }
1096
1097    let _guard = {
1098        let prev_snapshot = CURRENT_SNAPSHOT.with(|cell| cell.borrow().clone());
1099        let prev_auth = CURRENT_AUTH_IDENTITY.with(|cell| cell.borrow().clone());
1100        let prev_tenant = CURRENT_TENANT_ID.with(|cell| cell.borrow().clone());
1101
1102        match bundle.snapshot.clone() {
1103            Some(ctx) => set_current_snapshot(ctx),
1104            None => clear_current_snapshot(),
1105        }
1106        CURRENT_AUTH_IDENTITY.with(|cell| *cell.borrow_mut() = bundle.auth.clone());
1107        CURRENT_TENANT_ID.with(|cell| *cell.borrow_mut() = bundle.tenant.clone());
1108
1109        Guard {
1110            prev_snapshot,
1111            prev_auth,
1112            prev_tenant,
1113        }
1114    };
1115    f()
1116}
1117
1118/// Apply the same visibility rules used by the thread-local helpers
1119/// against a caller-provided context. Intended for parallel workers
1120/// that captured the snapshot with `capture_current_snapshot()`.
1121#[inline]
1122pub fn entity_visible_with_context(
1123    ctx: Option<&SnapshotContext>,
1124    entity: &crate::storage::unified::entity::UnifiedEntity,
1125) -> bool {
1126    // Same moderation visibility gate as the thread-local path (#1274):
1127    // parallel scan workers capture the snapshot context but must apply
1128    // the moderation marker check identically.
1129    if crate::runtime::ai::moderation::entity_moderation_hidden(entity) {
1130        return false;
1131    }
1132    match ctx {
1133        Some(ctx) => {
1134            let visible = visibility_check(ctx, entity.xmin, entity.xmax);
1135            if visible {
1136                record_serializable_read(ctx, entity);
1137            }
1138            visible
1139        }
1140        None => true,
1141    }
1142}