lean_ctx/core/
reference_store.rs

1// SPDX-License-Identifier: Apache-2.0
2use std::collections::HashMap;
3use std::sync::{Arc, Mutex, OnceLock};
4use std::time::{Duration, Instant};
5
6use crate::core::archive::authority::ArchiveAuthority;
7// Pack active, or the global gateway governed/sovereign: one predicate.
8use crate::core::context_admission::stored::source_authority_required;
9use crate::core::policy::{content, runtime};
10
11#[derive(Clone)]
12struct RefEntry {
13    content: Arc<str>,
14    authority: Option<ArchiveAuthority>,
15    created_at: Instant,
16}
17
18const MAX_ENTRIES: usize = 200;
19const MAX_STORE_BYTES: usize = 64 * 1024 * 1024;
20const TTL: Duration = Duration::from_mins(5);
21
22fn store_lock() -> &'static Mutex<HashMap<String, RefEntry>> {
23    static STORE: OnceLock<Mutex<HashMap<String, RefEntry>>> = OnceLock::new();
24    STORE.get_or_init(|| Mutex::new(HashMap::new()))
25}
26
27const RESOLVE_LOCK_WAIT: Duration = Duration::from_millis(50);
28
29fn lock_briefly() -> Option<std::sync::MutexGuard<'static, HashMap<String, RefEntry>>> {
30    let deadline = Instant::now() + RESOLVE_LOCK_WAIT;
31    loop {
32        match store_lock().try_lock() {
33            Ok(map) => return Some(map),
34            Err(std::sync::TryLockError::Poisoned(poisoned)) => return Some(poisoned.into_inner()),
35            Err(std::sync::TryLockError::WouldBlock) if Instant::now() < deadline => {
36                std::thread::sleep(Duration::from_micros(200));
37            }
38            Err(std::sync::TryLockError::WouldBlock) => return None,
39        }
40    }
41}
42
43/// Store unbound Community output. Mandatory policy requires observed origin.
44pub fn store(content: &str) -> Option<String> {
45    store_with_authority(content, None)
46}
47
48pub fn store_with_authority(text: &str, authority: Option<&ArchiveAuthority>) -> Option<String> {
49    runtime::with_source_view(|| {
50        if text.len() > entry_budget()
51            || (source_authority_required() && authority.is_none())
52            || authority.is_some_and(|origin| !origin.admitted())
53        {
54            return None;
55        }
56        let text = content::protect_active(text).ok()?;
57        // A recovery store holds admitted text only, never withheld or
58        // restricted content (G5, E3).
59        let text = crate::core::context_admission::recovery::admit_for_storage(&text)?;
60        if text.len() > entry_budget() {
61            return None;
62        }
63        let id = reference_id(&text, authority)?;
64        // Optional compression must never park a request behind a shared mutex.
65        let mut map = store_lock().try_lock().ok()?;
66        map.retain(|_, entry| entry.created_at.elapsed() < TTL);
67        let mut bytes: usize = map.values().map(|entry| entry.content.len()).sum();
68        if let Some(previous) = map.remove(&id) {
69            bytes -= previous.content.len();
70        }
71        while map.len() >= MAX_ENTRIES || bytes.saturating_add(text.len()) > MAX_STORE_BYTES {
72            let oldest = map
73                .iter()
74                .min_by_key(|(_, entry)| entry.created_at)
75                .map(|(key, _)| key.clone())?;
76            bytes -= map.remove(&oldest)?.content.len();
77        }
78        map.insert(
79            id.clone(),
80            RefEntry {
81                content: Arc::from(text.as_str()),
82                authority: authority.cloned(),
83                created_at: Instant::now(),
84            },
85        );
86        Some(id)
87    })
88    .ok()
89    .flatten()
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub enum ReferenceResolveError {
94    Malformed,
95    Missing,
96    Expired,
97    /// Present, but its source authority, project binding or the current
98    /// content rules do not allow it to be handed back. Content-free.
99    Refused(&'static str),
100    /// Present and still authorized by its source, but the current gateway
101    /// policy refuses this delivery. Carries the gateway's own content-free
102    /// reason, so a carrier can answer with the typed refusal.
103    Gateway(String),
104}
105
106/// Detailed form shared by the MCP resolver and the recovery verifier. The
107/// lookup is cheap and lock-scoped; every source, policy and role check runs
108/// after the store lock is released.
109pub fn resolve_checked(id: &str) -> Result<String, ReferenceResolveError> {
110    let Some(digest) = id.strip_prefix("ref_") else {
111        return Err(ReferenceResolveError::Malformed);
112    };
113    if digest.is_empty() || !digest.bytes().all(|b| b.is_ascii_hexdigit()) {
114        return Err(ReferenceResolveError::Malformed);
115    }
116    let entry = {
117        // An explicit resolve waits out a concurrent store (microseconds),
118        // but a store held longer degrades to "missing" instead of parking.
119        let Some(mut map) = lock_briefly() else {
120            return Err(ReferenceResolveError::Missing);
121        };
122        let Some(entry) = map.get(id) else {
123            return Err(ReferenceResolveError::Missing);
124        };
125        if entry.created_at.elapsed() >= TTL {
126            map.remove(id);
127            return Err(ReferenceResolveError::Expired);
128        }
129        entry.clone()
130    };
131    // Do not infer a protected HTTP request's identity from the proxy's cwd.
132    if entry.authority.is_some()
133        && !runtime::REQUEST_PROJECT
134            .try_with(|project| project.borrow().is_some())
135            .unwrap_or(false)
136    {
137        return Err(ReferenceResolveError::Refused(
138            "reference is bound to a project this request does not name",
139        ));
140    }
141    runtime::with_source_view(|| {
142        if source_authority_required() && entry.authority.is_none() {
143            return Err(ReferenceResolveError::Refused(
144                "stored output has no source authority under the active policy",
145            ));
146        }
147        if entry
148            .authority
149            .as_ref()
150            .is_some_and(|origin| !origin.admitted())
151            || reference_id(&entry.content, entry.authority.as_ref()).as_deref() != Some(id)
152        {
153            return Err(ReferenceResolveError::Refused(
154                "the reference's source is no longer admitted",
155            ));
156        }
157        let checked = content::protect_active(&entry.content).map_err(|_| {
158            ReferenceResolveError::Refused("stored output is withheld by current content rules")
159        })?;
160        // Recovery is a new delivery for every resolver (`ctx_expand`, the
161        // proxy, the recovery verifier): re-admitted under the current
162        // gateway policy here, not only by individual callers.
163        crate::core::context_admission::recovery::admit_recovered(&checked, "reference")
164            .map_err(ReferenceResolveError::Gateway)
165    })
166    .unwrap_or(Err(ReferenceResolveError::Refused(
167        "current policy cannot be verified",
168    )))
169}
170
171/// Back-compatible resolver used by `ctx_expand` and the proxy.
172pub fn resolve(id: &str) -> Option<String> {
173    resolve_checked(id).ok()
174}
175
176pub fn stats() -> (usize, usize) {
177    let Ok(map) = store_lock().try_lock() else {
178        return (0, 0);
179    };
180    let live = map
181        .values()
182        .filter(|entry| entry.created_at.elapsed() < TTL);
183    live.fold((0, 0), |(count, bytes), entry| {
184        (count + 1, bytes + entry.content.len())
185    })
186}
187
188fn entry_budget() -> usize {
189    crate::core::limits::max_read_bytes().min(content::MAX_PROTECTED_CONTENT_BYTES)
190}
191
192fn reference_id(text: &str, authority: Option<&ArchiveAuthority>) -> Option<String> {
193    let digest = if let Some(origin) = authority {
194        origin.id(text)?
195    } else {
196        let mut hash = blake3::Hasher::new();
197        hash.update(b"leanctx-public-reference-v1\0");
198        hash.update(text.as_bytes());
199        hash.finalize().to_hex().to_string()
200    };
201    Some(format!("ref_{digest}"))
202}
203
204#[cfg(test)]
205mod tests;
lean_ctx/core/reference_store.rs

lean_ctx/core/
reference_store.rs
