1use serde::{Deserialize, Serialize};
13use std::path::Path;
14use std::sync::{Arc, Mutex};
15
16pub(crate) const MAX_SOURCES: usize = 64;
20
21type PermissionKey = (String, String, String);
23
24thread_local! {
25 static PERMITTED: std::cell::RefCell<Option<std::collections::HashMap<PermissionKey, bool>>> =
28 const { std::cell::RefCell::new(None) };
29}
30
31pub fn with_permission_memo<T>(operation: impl FnOnce() -> T) -> T {
36 struct Restore(Option<std::collections::HashMap<PermissionKey, bool>>);
37 impl Drop for Restore {
38 fn drop(&mut self) {
39 let previous = self.0.take();
40 PERMITTED.with(|memo| *memo.borrow_mut() = previous);
41 }
42 }
43 let _restore =
44 Restore(PERMITTED.with(|memo| memo.borrow_mut().replace(std::collections::HashMap::new())));
45 operation()
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49pub struct ArchiveAuthority {
50 version: u8,
51 project: String,
52 path: String,
53 tool: String,
54 admitted_digest: String,
55 #[serde(default, skip_serializing_if = "Vec::is_empty")]
58 more: Vec<AdditionalSource>,
59 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
62 incomplete: bool,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66struct AdditionalSource {
67 path: String,
68 tool: String,
69 admitted_digest: String,
70}
71
72#[derive(Debug, Default)]
74pub struct CaptureSlot {
75 authority: Mutex<Option<ArchiveAuthority>>,
76 unbound: std::sync::atomic::AtomicBool,
79}
80
81const SLOT_WAIT: std::time::Duration = std::time::Duration::from_millis(250);
85
86impl CaptureSlot {
87 fn lock_within(
93 &self,
94 budget: std::time::Duration,
95 ) -> Option<std::sync::MutexGuard<'_, Option<ArchiveAuthority>>> {
96 let deadline = std::time::Instant::now() + budget;
97 loop {
98 match self.authority.try_lock() {
99 Ok(guard) => return Some(guard),
100 Err(std::sync::TryLockError::Poisoned(poisoned)) => {
101 return Some(poisoned.into_inner());
102 }
103 Err(std::sync::TryLockError::WouldBlock)
104 if std::time::Instant::now() < deadline =>
105 {
106 std::thread::sleep(std::time::Duration::from_micros(200));
107 }
108 Err(std::sync::TryLockError::WouldBlock) => {
109 self.unbound
110 .store(true, std::sync::atomic::Ordering::SeqCst);
111 return None;
112 }
113 }
114 }
115 }
116
117 fn publish_within(&self, authority: ArchiveAuthority, budget: std::time::Duration) {
119 if let Some(mut slot) = self.lock_within(budget) {
120 match slot.as_mut() {
121 Some(existing) => existing.join(authority),
122 None => *slot = Some(authority),
123 }
124 }
125 }
126
127 fn finish_within(&self, budget: std::time::Duration) -> Option<ArchiveAuthority> {
131 let mut authority = self.lock_within(budget).and_then(|mut slot| slot.take());
132 if self.unbound.load(std::sync::atomic::Ordering::SeqCst) {
133 match authority.as_mut() {
134 Some(authority) => authority.incomplete = true,
135 None => authority = Some(ArchiveAuthority::unbound()),
136 }
137 }
138 authority
139 }
140}
141
142pub(crate) type Capture = Arc<CaptureSlot>;
143
144tokio::task_local! {
145 pub static CAPTURE: Option<Capture>;
146}
147
148pub fn current_capture() -> Option<Capture> {
149 CAPTURE.try_with(Clone::clone).ok().flatten()
150}
151
152pub fn without_capture<T>(operation: impl FnOnce() -> T) -> T {
156 CAPTURE.sync_scope(None, operation)
157}
158
159pub fn mark_unbound() {
161 if let Some(capture) = current_capture() {
162 capture
163 .unbound
164 .store(true, std::sync::atomic::Ordering::SeqCst);
165 }
166}
167
168pub fn publish_read(project: &Path, canonical_path: &str, tool: &str, admitted: &str) {
180 crate::core::context_admission::capture::note_source(project, canonical_path);
184 if current_capture().is_none() {
185 return;
186 }
187 if let Some(authority) = ArchiveAuthority::file(project, canonical_path, tool, admitted) {
188 authority.publish();
189 }
190}
191
192pub async fn capture<T>(
202 operation: impl std::future::Future<Output = T>,
203) -> (T, Option<ArchiveAuthority>) {
204 let capture = Arc::new(CaptureSlot::default());
205 let result = CAPTURE.scope(Some(capture.clone()), operation).await;
206 (result, capture.finish_within(SLOT_WAIT))
207}
208
209pub fn capture_sync<T>(operation: impl FnOnce() -> T) -> (T, Option<ArchiveAuthority>) {
213 let capture = Arc::new(CaptureSlot::default());
214 let result = CAPTURE.sync_scope(Some(capture.clone()), operation);
215 (result, capture.finish_within(SLOT_WAIT))
216}
217
218impl ArchiveAuthority {
219 pub fn file(project: &Path, path: &str, tool: &str, admitted: &str) -> Option<Self> {
221 if !crate::core::context_admission::stored::source_authority_required() {
222 return None;
223 }
224 let authority = Self::bind(project, path, tool, admitted);
225 if authority.is_none() {
226 mark_unbound();
229 }
230 authority
231 }
232
233 pub(crate) fn admitted_source(
244 project: &Path,
245 path: &str,
246 tool: &str,
247 admitted_digest: &str,
248 ) -> Option<Self> {
249 if !crate::core::context_admission::stored::source_authority_required() {
250 return None;
251 }
252 let authority = Self::bind_digest(project, path, tool, admitted_digest);
253 if authority.is_none() {
254 mark_unbound();
255 }
256 authority
257 }
258
259 fn bind_digest(project: &Path, path: &str, tool: &str, admitted_digest: &str) -> Option<Self> {
260 let well_formed = admitted_digest.len() == 64
261 && admitted_digest
262 .bytes()
263 .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase());
264 if !well_formed {
265 return None;
266 }
267 let project = canonical_archive_path(project)?;
268 let path = canonical_archive_path(&project.join(path))?;
271 if !path.starts_with(&project) {
272 return None;
273 }
274 Some(Self {
275 version: 1,
276 project: project.to_str()?.to_owned(),
277 path: path.to_str()?.to_owned(),
278 tool: tool.to_owned(),
279 admitted_digest: admitted_digest.to_owned(),
280 more: Vec::new(),
281 incomplete: false,
282 })
283 }
284
285 fn unbound() -> Self {
288 Self {
289 version: 1,
290 project: String::new(),
291 path: String::new(),
292 tool: String::new(),
293 admitted_digest: String::new(),
294 more: Vec::new(),
295 incomplete: true,
296 }
297 }
298
299 pub fn is_incomplete(&self) -> bool {
302 self.incomplete
303 }
304
305 fn bind(project: &Path, path: &str, tool: &str, admitted: &str) -> Option<Self> {
306 let project = canonical_archive_path(project)?;
307 let path = canonical_archive_path(Path::new(path))?;
308 if !path.starts_with(&project) {
309 return None;
310 }
311 Some(Self {
312 version: 1,
313 project: project.to_str()?.to_owned(),
314 path: path.to_str()?.to_owned(),
315 tool: tool.to_owned(),
316 admitted_digest: blake3::hash(admitted.as_bytes()).to_hex().to_string(),
317 more: Vec::new(),
318 incomplete: false,
319 })
320 }
321
322 pub fn publish(self) {
326 if let Some(capture) = current_capture() {
327 capture.publish_within(self, SLOT_WAIT);
328 }
329 }
330
331 pub fn join(&mut self, other: Self) {
333 if other.project != self.project {
334 self.incomplete = true;
335 return;
336 }
337 self.incomplete |= other.incomplete;
338 let first = AdditionalSource {
339 path: other.path,
340 tool: other.tool,
341 admitted_digest: other.admitted_digest,
342 };
343 for source in std::iter::once(first).chain(other.more) {
344 let primary = source.path == self.path
345 && source.tool == self.tool
346 && source.admitted_digest == self.admitted_digest;
347 if primary || self.more.contains(&source) {
348 continue;
349 }
350 if self.more.len() + 1 >= MAX_SOURCES {
351 self.incomplete = true;
352 return;
353 }
354 self.more.push(source);
355 }
356 }
357
358 pub(super) fn project(&self) -> &str {
360 &self.project
361 }
362
363 pub(super) fn sources(&self) -> Option<Vec<(&str, &str, &str)>> {
366 if self.incomplete {
367 return None;
368 }
369 Some(
370 std::iter::once((&self.path, &self.tool, &self.admitted_digest))
371 .chain(
372 self.more
373 .iter()
374 .map(|source| (&source.path, &source.tool, &source.admitted_digest)),
375 )
376 .map(|(path, tool, digest)| (path.as_str(), tool.as_str(), digest.as_str()))
377 .collect(),
378 )
379 }
380
381 pub(super) fn absorb(&mut self, later: Self) {
388 if later.project != self.project {
389 self.incomplete = true;
390 return;
391 }
392 self.incomplete |= later.incomplete;
393 let first = AdditionalSource {
394 path: later.path,
395 tool: later.tool,
396 admitted_digest: later.admitted_digest,
397 };
398 for source in std::iter::once(first).chain(later.more) {
399 if source.path == self.path && source.tool == self.tool {
400 self.admitted_digest = source.admitted_digest;
401 continue;
402 }
403 if let Some(known) = self
404 .more
405 .iter_mut()
406 .find(|known| known.path == source.path && known.tool == source.tool)
407 {
408 known.admitted_digest = source.admitted_digest;
409 continue;
410 }
411 if self.more.len() + 1 >= MAX_SOURCES {
412 self.incomplete = true;
413 return;
414 }
415 self.more.push(source);
416 }
417 }
418
419 pub(crate) fn id(&self, content: &str) -> Option<String> {
420 let mut hash = blake3::Hasher::new();
421 hash.update(b"leanctx-source-archive-v1\0");
422 hash.update(&serde_json::to_vec(self).ok()?);
423 hash.update(b"\0");
424 hash.update(content.as_bytes());
425 Some(hash.finalize().to_hex().to_string())
426 }
427
428 pub fn admitted(&self) -> bool {
429 self.admitted_budgeted(&mut crate::core::limits::max_read_bytes())
430 }
431
432 #[must_use]
443 pub(crate) fn still_admitted(&self) -> bool {
444 crate::core::policy::runtime::with_source_view(|| self.admitted()).unwrap_or(false)
445 }
446
447 #[must_use]
466 pub(crate) fn still_permitted(&self) -> bool {
467 crate::core::policy::runtime::with_source_view(|| self.permitted()).unwrap_or(false)
468 }
469
470 fn permitted(&self) -> bool {
471 if self.incomplete || !self.matches_project() {
472 return false;
473 }
474 std::iter::once((&self.path, &self.tool))
475 .chain(self.more.iter().map(|source| (&source.path, &source.tool)))
476 .all(|(path, tool)| {
477 let key = (self.project.clone(), path.clone(), tool.clone());
478 if let Some(known) = PERMITTED.with(|memo| {
479 memo.borrow()
480 .as_ref()
481 .and_then(|memo| memo.get(&key).copied())
482 }) {
483 return known;
484 }
485 let mut remaining = crate::core::limits::max_read_bytes();
486 let permitted = crate::core::rooted_read::read_file_for_tool_rooted_admitted(
487 path,
488 &self.project,
489 tool,
490 &mut remaining,
491 )
492 .is_ok_and(|read| read.complete && read.storable);
493 PERMITTED.with(|memo| {
494 if let Some(memo) = memo.borrow_mut().as_mut() {
495 memo.insert(key, permitted);
496 }
497 });
498 permitted
499 })
500 }
501
502 #[must_use]
507 pub(crate) fn identity_digest(&self) -> Option<String> {
508 if self.incomplete {
509 return None;
510 }
511 let mut sources: Vec<(&str, &str, &str)> =
513 std::iter::once((&self.path, &self.tool, &self.admitted_digest))
514 .chain(
515 self.more
516 .iter()
517 .map(|source| (&source.path, &source.tool, &source.admitted_digest)),
518 )
519 .map(|(path, tool, digest)| (path.as_str(), tool.as_str(), digest.as_str()))
520 .collect();
521 sources.sort_unstable();
522 let mut hash = blake3::Hasher::new();
523 hash.update(b"leanctx-source-lineage-v1\0");
524 hash.update(&serde_json::to_vec(&(self.version, &self.project, sources)).ok()?);
525 Some(hash.finalize().to_hex().to_string())
526 }
527
528 pub(super) fn matches_project(&self) -> bool {
529 if self.version != 1 {
530 return false;
531 }
532 let Some(project) = crate::core::policy::runtime::REQUEST_PROJECT
533 .try_with(|slot| slot.borrow().clone())
534 .ok()
535 .flatten()
536 .and_then(|path| crate::core::pathutil::canonicalize_secure(&path).ok())
537 else {
538 return false;
539 };
540 project == Path::new(&self.project)
541 }
542
543 pub(super) fn admitted_budgeted(&self, remaining: &mut usize) -> bool {
546 if self.incomplete || !self.matches_project() {
547 return false;
548 }
549 let primary = (&self.path, &self.tool, &self.admitted_digest);
550 std::iter::once(primary)
551 .chain(
552 self.more
553 .iter()
554 .map(|source| (&source.path, &source.tool, &source.admitted_digest)),
555 )
556 .all(|(path, tool, digest)| {
557 *remaining != 0
558 && crate::core::rooted_read::read_file_for_tool_rooted_admitted(
559 path,
560 &self.project,
561 tool,
562 remaining,
563 )
564 .is_ok_and(|read| {
565 read.complete
572 && read.storable
573 && blake3::hash(read.content.as_bytes()).to_hex().as_str()
574 == digest.as_str()
575 })
576 })
577 }
578}
579
580fn canonical_archive_path(path: &Path) -> Option<std::path::PathBuf> {
581 if !path.is_absolute() {
582 return None;
583 }
584 crate::core::pathutil::canonicalize_secure(path)
585 .or_else(|error| {
586 let simplified = crate::core::pathutil::strip_verbatim(path.to_path_buf());
587 if simplified.as_path() == path {
588 Err(error)
589 } else {
590 crate::core::pathutil::canonicalize_secure(&simplified)
591 }
592 })
593 .ok()
594}
595
596#[cfg(test)]
597mod tests {
598 use super::*;
599
600 fn source(path: &str) -> ArchiveAuthority {
601 ArchiveAuthority {
602 version: 1,
603 project: "/project".to_owned(),
604 path: path.to_owned(),
605 tool: "ctx_read".to_owned(),
606 admitted_digest: blake3::hash(path.as_bytes()).to_hex().to_string(),
607 more: Vec::new(),
608 incomplete: false,
609 }
610 }
611
612 use std::time::Duration;
613
614 fn hold(slot: &Arc<CaptureSlot>) -> (std::sync::mpsc::Sender<()>, std::thread::JoinHandle<()>) {
616 let (release, wait) = std::sync::mpsc::channel::<()>();
617 let (held, is_held) = std::sync::mpsc::channel::<()>();
618 let slot = slot.clone();
619 let holder = std::thread::spawn(move || {
620 let guard = slot.lock_within(Duration::from_secs(5)).expect("free slot");
621 held.send(()).expect("signal");
622 let _ = wait.recv();
623 drop(guard);
624 });
625 is_held.recv().expect("the holder owns the slot");
626 (release, holder)
627 }
628
629 #[test]
634 fn slot_contention_is_bounded_and_never_loses_a_source_silently() {
635 let slot = Arc::new(CaptureSlot::default());
637 let (release, holder) = hold(&slot);
638 let started = std::time::Instant::now();
639 slot.publish_within(source("/project/a.rs"), Duration::from_millis(20));
640 assert!(started.elapsed() < Duration::from_secs(2), "bounded wait");
641 drop(release);
642 holder.join().expect("holder");
643 let captured = slot
644 .finish_within(Duration::from_secs(5))
645 .expect("an unknown acquisition is reported, not dropped");
646 assert!(captured.is_incomplete(), "exceeded budget → incomplete");
647 assert!(!captured.admitted(), "an incomplete authority never admits");
648
649 let slot = Arc::new(CaptureSlot::default());
651 let (release, holder) = hold(&slot);
652 let publisher = {
653 let slot = slot.clone();
654 std::thread::spawn(move || {
655 slot.publish_within(source("/project/a.rs"), Duration::from_secs(5));
656 slot.publish_within(source("/project/b.rs"), Duration::from_secs(5));
657 })
658 };
659 std::thread::sleep(Duration::from_millis(20));
660 drop(release);
661 holder.join().expect("holder");
662 publisher.join().expect("publisher");
663 let captured = slot
664 .finish_within(Duration::from_secs(5))
665 .expect("captured");
666 assert_eq!(captured.more.len(), 1, "the second source joined the first");
667 assert!(!captured.is_incomplete());
668 }
669}