lean_ctx/core/session/
taint.rs

1// SPDX-License-Identifier: Apache-2.0
2//! The session taint (design pz21 §1.7): which untrusted origins the session
3//! the Instance serves has read.
4//!
5//! - **Writer:** only the origin step of the release point, from the call's
6//!   origin ledger ([`record`]). Admissions outside a tool call (egress
7//!   checks, detached store writes) never taint.
8//! - **Active session:** the MCP server binds the session it serves at the
9//!   start of every call ([`bind`]); [`active_view`] is what
10//!   `origin::session_taint()` returns. A taint recorded before the first
11//!   bind is held without a session id and carried into the bound session.
12//! - **Write-through:** a change is saved at once to a sidecar next to the
13//!   session file (`<id>.taint`), not on the session's save cadence, so a
14//!   crash between the read and the next save cannot untaint the session.
15//!   Loading a session merges the sidecar into `SessionState.taint`.
16//! - **Overflow is sticky:** nothing is dropped silently; a ref that cannot
17//!   be recorded marks the whole session, and so does a change the bounded
18//!   lock could not take in time.
19//! - **Locking:** `ACTIVE` is a bounded lock (LOCK_ORDERING L148), held
20//!   only to merge or copy the in-memory taint; sidecar reads and writes run
21//!   after it is released.
22
23use std::sync::atomic::{AtomicBool, Ordering};
24use std::sync::{Arc, LazyLock};
25use std::time::Duration;
26
27use tokio::sync::{OwnedRwLockWriteGuard, RwLock};
28
29use crate::core::bounded_lock;
30use crate::core::context_admission::origin::{InputTaint, TaintView};
31
32#[derive(Default)]
33struct Active {
34    /// Empty while no session is bound (a CLI run, a test, the first call).
35    session_id: String,
36    taint: InputTaint,
37    /// Bumped on every change, so a sidecar write that raced a newer change
38    /// writes again.
39    generation: u64,
40}
41
42static ACTIVE: LazyLock<Arc<RwLock<Option<Active>>>> =
43    LazyLock::new(|| Arc::new(RwLock::new(None)));
44
45/// A taint change was lost to a lock timeout: every view of the taint reads
46/// as overflowed for the rest of the process (fails closed).
47static UNRECORDED: AtomicBool = AtomicBool::new(false);
48
49/// The lock is held for an in-memory merge only; a wait this long means
50/// something is wrong, and the taint fails closed instead of queueing.
51const LOCK_BUDGET: Duration = Duration::from_secs(2);
52
53fn write_guard() -> Option<OwnedRwLockWriteGuard<Option<Active>>> {
54    let guard = bounded_lock::write_for(&ACTIVE, "session taint", LOCK_BUDGET);
55    if guard.is_none() {
56        UNRECORDED.store(true, Ordering::SeqCst);
57    }
58    guard
59}
60
61fn with_unrecorded(mut taint: InputTaint) -> InputTaint {
62    if UNRECORDED.load(Ordering::SeqCst) {
63        taint.mark_overflow();
64    }
65    taint
66}
67
68/// Bind the session the Instance serves now. A different session replaces
69/// the active taint with its own (loaded with the session); the same
70/// session keeps the union, and a taint recorded before any session was
71/// bound is carried into it and saved, so nothing recorded is ever lost.
72pub fn bind(session_id: &str, taint: &InputTaint) {
73    let save = {
74        let Some(mut active) = write_guard() else {
75            return;
76        };
77        match active.as_mut() {
78            Some(current) if current.session_id == session_id => {
79                if !current.taint.merge(taint) {
80                    return;
81                }
82                current.generation += 1;
83                Some((current.taint.clone(), current.generation))
84            }
85            _ => None,
86        }
87    };
88    if let Some((taint, generation)) = save {
89        persist(session_id, taint, generation);
90        return;
91    }
92    // Another (or no) session was bound: read its sidecar without the lock.
93    let mut stored = load_sidecar(session_id).unwrap_or_default();
94    stored.merge(taint);
95    let save = {
96        let Some(mut active) = write_guard() else {
97            return;
98        };
99        let mut carried = false;
100        let mut generation = 1;
101        if let Some(previous) = active.take()
102            && (previous.session_id.is_empty() || previous.session_id == session_id)
103        {
104            // Recorded before this bind (no session id yet), or a bind of the
105            // same session that won the race: keep it.
106            carried = stored.merge(&previous.taint);
107            generation = previous.generation + 1;
108        }
109        *active = Some(Active {
110            session_id: session_id.to_owned(),
111            taint: stored.clone(),
112            generation,
113        });
114        carried.then_some((stored, generation))
115    };
116    if let Some((taint, generation)) = save {
117        persist(session_id, taint, generation);
118    }
119}
120
121/// Add `delta` to the active session's taint and persist a change at once.
122/// Without a bound session the taint is held process-wide (a CLI run, a
123/// test) and carried into the session bound next.
124pub fn record(delta: &InputTaint) {
125    if !delta.is_tainted() {
126        return;
127    }
128    let save = {
129        let Some(mut active) = write_guard() else {
130            return;
131        };
132        let current = active.get_or_insert_with(Active::default);
133        if !current.taint.merge(delta) {
134            return;
135        }
136        current.generation += 1;
137        (!current.session_id.is_empty()).then(|| {
138            (
139                current.session_id.clone(),
140                current.taint.clone(),
141                current.generation,
142            )
143        })
144    };
145    if let Some((session_id, taint, generation)) = save {
146        persist(&session_id, taint, generation);
147    }
148}
149
150/// Write `taint` (of change `generation`) to the sidecar, outside the lock;
151/// when a newer change of the same session landed meanwhile, write that
152/// too, so the last write on disk is never older than the last change.
153fn persist(session_id: &str, mut taint: InputTaint, mut generation: u64) {
154    loop {
155        save_sidecar(session_id, &taint);
156        let Some(active) = bounded_lock::read_for(&ACTIVE, "session taint", LOCK_BUDGET) else {
157            UNRECORDED.store(true, Ordering::SeqCst);
158            return;
159        };
160        match active.as_ref() {
161            Some(current)
162                if current.session_id == session_id && current.generation > generation =>
163            {
164                taint = current.taint.clone();
165                generation = current.generation;
166            }
167            _ => return,
168        }
169    }
170}
171
172/// The active session's taint, read-only.
173pub fn active_view() -> TaintView {
174    snapshot().view()
175}
176
177/// The active session's taint. A lock timeout reads as overflowed.
178pub fn snapshot() -> InputTaint {
179    let Some(active) = bounded_lock::read_for(&ACTIVE, "session taint", LOCK_BUDGET) else {
180        UNRECORDED.store(true, Ordering::SeqCst);
181        return with_unrecorded(InputTaint::default());
182    };
183    with_unrecorded(
184        active
185            .as_ref()
186            .map(|active| active.taint.clone())
187            .unwrap_or_default(),
188    )
189}
190
191/// The id of the session bound last, if any.
192pub fn active_session_id() -> Option<String> {
193    bounded_lock::read_for(&ACTIVE, "session taint", LOCK_BUDGET)?
194        .as_ref()
195        .map(|active| active.session_id.clone())
196        .filter(|id| !id.is_empty())
197}
198
199/// The taint of `session_id` when it is the active session.
200pub fn snapshot_for(session_id: &str) -> Option<InputTaint> {
201    let taint = bounded_lock::read_for(&ACTIVE, "session taint", LOCK_BUDGET)?
202        .as_ref()
203        .filter(|active| active.session_id == session_id)
204        .map(|active| active.taint.clone())?;
205    Some(with_unrecorded(taint))
206}
207
208/// Forget the active session (tests; a server shutting down). Persisted
209/// taints stay on disk.
210#[cfg(any(test, feature = "test-support"))]
211pub fn reset_active() {
212    if let Some(mut active) = bounded_lock::write_for(&ACTIVE, "session taint", LOCK_BUDGET) {
213        *active = None;
214    }
215    UNRECORDED.store(false, Ordering::SeqCst);
216}
217
218fn sidecar_path(session_id: &str) -> Option<std::path::PathBuf> {
219    let valid = !session_id.is_empty()
220        && !session_id.starts_with('.')
221        && session_id
222            .chars()
223            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_'));
224    if !valid {
225        return None;
226    }
227    super::paths::sessions_dir().map(|dir| dir.join(format!("{session_id}.taint")))
228}
229
230/// The persisted taint of `session_id`, merged into the session on load.
231pub(crate) fn load_sidecar(session_id: &str) -> Option<InputTaint> {
232    let path = sidecar_path(session_id)?;
233    let json = std::fs::read_to_string(path).ok()?;
234    // A sidecar that does not parse fails closed: the session is tainted.
235    Some(serde_json::from_str(&json).unwrap_or_else(|_| InputTaint {
236        overflow: true,
237        ..InputTaint::default()
238    }))
239}
240
241fn save_sidecar(session_id: &str, taint: &InputTaint) {
242    let Some(path) = sidecar_path(session_id) else {
243        return;
244    };
245    let Ok(json) = serde_json::to_string(taint) else {
246        return;
247    };
248    if let Some(dir) = path.parent() {
249        let _ = std::fs::create_dir_all(dir);
250    }
251    if let Err(error) = crate::core::atomic_fs::try_atomic_write(&path, json.as_bytes(), None) {
252        tracing::warn!("session taint not persisted: {error}");
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use crate::core::context_admission::origin::{InputKind, InputRef};
260
261    #[test]
262    fn taint_is_written_through_and_survives_a_rebind() {
263        let _data = crate::core::data_dir::isolated_data_dir();
264        reset_active();
265        bind("20261106-000000u000000p1s0", &InputTaint::default());
266        let mut delta = InputTaint::default();
267        delta.add(InputRef::new(InputKind::Issue, "github:o/r#1"));
268        delta.mark_overflow();
269        record(&delta);
270        assert!(active_view().overflow());
271        // A process restart: the active binding is gone, the sidecar is not.
272        reset_active();
273        assert!(!active_view().is_tainted());
274        bind("20261106-000000u000000p1s0", &InputTaint::default());
275        let view = active_view();
276        assert!(view.overflow());
277        assert!(view.contains(&InputRef::new(InputKind::Issue, "github:o/r#1").unwrap()));
278        // Another session starts untainted.
279        bind("20261106-000000u000000p1s1", &InputTaint::default());
280        assert!(!active_view().is_tainted());
281        reset_active();
282    }
283
284    /// Review 22480 finding 3: a taint recorded before the first `bind` was
285    /// held under an empty session id, and `bind` replaced it with the
286    /// sidecar, so the session forgot what it had read.
287    #[test]
288    fn taint_recorded_before_the_first_bind_is_kept_and_saved() {
289        let _data = crate::core::data_dir::isolated_data_dir();
290        reset_active();
291        let issue = InputRef::new(InputKind::Issue, "github:o/r#7").unwrap();
292        let mut delta = InputTaint::default();
293        delta.add(Some(issue.clone()));
294        record(&delta);
295        bind("20261106-000000u000000p1s7", &InputTaint::default());
296        assert!(active_view().contains(&issue), "lost on bind");
297        // Saved: a restart reads it back from the sidecar.
298        reset_active();
299        bind("20261106-000000u000000p1s7", &InputTaint::default());
300        assert!(active_view().contains(&issue), "not written through");
301        reset_active();
302    }
303}
lean_ctx/core/session/taint.rs

lean_ctx/core/session/
taint.rs
