use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::sync::Mutex;
use async_trait::async_trait;
use crate::StoreError;
pub const NAR_REF_PREFIX: &str = "nar-refs/";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NarRefKey<'a> {
pub nar_path: &'a str,
pub hash: &'a str,
}
impl fmt::Display for NarRefKey<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{NAR_REF_PREFIX}{}/{}", self.nar_path, self.hash)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NarRefScan<'a> {
pub nar_path: &'a str,
}
impl fmt::Display for NarRefScan<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{NAR_REF_PREFIX}{}/", self.nar_path)
}
}
#[must_use]
pub fn referrer_of<'k>(scan: &NarRefScan<'_>, key: &'k str) -> Option<&'k str> {
let prefix = scan.to_string();
let rest = key.strip_prefix(&prefix)?;
if rest.is_empty() || rest.contains('/') {
return None;
}
Some(rest)
}
#[must_use]
pub fn is_addressable_nar_path(url: &str) -> bool {
!url.is_empty()
&& !url.starts_with('/')
&& !url.contains('\\')
&& !url.chars().any(char::is_control)
&& url.split('/').all(|seg| !seg.is_empty() && seg != "." && seg != "..")
}
#[must_use]
pub fn advertised_nar_url(narinfo: &str) -> Option<String> {
let url = advertised_url_line(narinfo)?;
is_addressable_nar_path(url).then(|| url.to_string())
}
#[must_use]
pub fn advertised_url_line(narinfo: &str) -> Option<&str> {
narinfo.lines().find_map(|line| {
let (key, value) = line.split_once(':')?;
(key.trim() == "URL").then(|| value.trim())
})
}
#[must_use]
pub fn is_servable_narinfo(narinfo: &str) -> bool {
narinfo
.lines()
.any(|line| line.split_once(':').is_some_and(|(k, v)| k.trim() == "StorePath" && !v.trim().is_empty()))
}
#[async_trait]
pub trait NarRefIndex: Send + Sync {
async fn record(&self, nar_path: &str, hash: &str) -> Result<(), StoreError>;
async fn forget(&self, nar_path: &str, hash: &str) -> Result<(), StoreError>;
async fn referrers(&self, nar_path: &str) -> Result<Vec<String>, StoreError>;
}
#[derive(Debug, Default)]
pub struct MemNarRefIndex {
edges: Mutex<BTreeMap<String, BTreeSet<String>>>,
}
impl MemNarRefIndex {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn len(&self) -> usize {
self.edges.lock().unwrap_or_else(std::sync::PoisonError::into_inner).values().map(BTreeSet::len).sum()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[async_trait]
impl NarRefIndex for MemNarRefIndex {
async fn record(&self, nar_path: &str, hash: &str) -> Result<(), StoreError> {
self.edges
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.entry(nar_path.to_string())
.or_default()
.insert(hash.to_string());
Ok(())
}
async fn forget(&self, nar_path: &str, hash: &str) -> Result<(), StoreError> {
let mut edges = self.edges.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(set) = edges.get_mut(nar_path) {
set.remove(hash);
if set.is_empty() {
edges.remove(nar_path);
}
}
Ok(())
}
async fn referrers(&self, nar_path: &str) -> Result<Vec<String>, StoreError> {
Ok(self
.edges
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(nar_path)
.map(|s| s.iter().cloned().collect())
.unwrap_or_default())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn an_edge_key_is_scanned_by_its_own_prefix() {
let key = NarRefKey { nar_path: "nar/abc.nar.xz", hash: "sss" }.to_string();
assert_eq!(key, "nar-refs/nar/abc.nar.xz/sss");
let scan = NarRefScan { nar_path: "nar/abc.nar.xz" };
assert!(key.starts_with(&scan.to_string()));
assert_eq!(referrer_of(&scan, &key), Some("sss"));
}
#[test]
fn a_scan_prefix_does_not_reach_a_longer_neighbour() {
let neighbour = NarRefKey { nar_path: "nar/ab.nar.xz", hash: "sss" }.to_string();
let scan = NarRefScan { nar_path: "nar/ab.nar" };
assert!(!neighbour.starts_with(&scan.to_string()));
assert_eq!(referrer_of(&scan, &neighbour), None);
}
#[test]
fn referrer_of_rejects_a_key_from_another_nar() {
let scan = NarRefScan { nar_path: "nar/a.nar" };
assert_eq!(referrer_of(&scan, "nar-refs/nar/b.nar/sss"), None);
assert_eq!(referrer_of(&scan, "nar-refs/nar/a.nar/"), None);
assert_eq!(referrer_of(&scan, "nar-refs/nar/a.nar/deep/sss"), None);
}
#[test]
fn traversal_and_absolute_urls_are_not_addressable() {
assert!(is_addressable_nar_path("nar/abc.nar.xz"));
assert!(is_addressable_nar_path("nar/deep/abc.nar"));
assert!(!is_addressable_nar_path(""));
assert!(!is_addressable_nar_path("/etc/passwd"));
assert!(!is_addressable_nar_path("../../etc/passwd"));
assert!(!is_addressable_nar_path("nar/../../etc/passwd"));
assert!(!is_addressable_nar_path("nar/./abc.nar"));
assert!(!is_addressable_nar_path("nar//abc.nar"));
assert!(!is_addressable_nar_path("nar\\abc.nar"));
assert!(!is_addressable_nar_path("nar/abc\n.nar"));
}
#[test]
fn an_unaddressable_url_advertises_nothing() {
let good = "StorePath: /nix/store/abc-hello\nURL: nar/abc.nar.xz\nCompression: xz\n\
FileHash: sha256:aaa\nFileSize: 100\nNarHash: sha256:bbb\nNarSize: 200\n\
References: \n";
assert_eq!(advertised_nar_url(good).as_deref(), Some("nar/abc.nar.xz"));
let traversal = "StorePath: /nix/store/abc-hello\nURL: ../../etc/passwd\n\
Compression: xz\nFileHash: sha256:aaa\nFileSize: 100\n\
NarHash: sha256:bbb\nNarSize: 200\nReferences: \n";
assert_eq!(advertised_nar_url(traversal), None);
assert_eq!(advertised_nar_url("not a narinfo at all"), None);
}
#[test]
fn a_narinfo_the_strict_parser_rejects_still_advertises_its_nar() {
let partial = "StorePath: /nix/store/abc-hello\nURL: nar/abc.nar.xz\n\
Compression: xz\nNarHash: sha256:bbb\nNarSize: 200\nReferences: \n";
assert!(
sui_compat::narinfo::NarInfo::parse(partial).is_err(),
"fixture must actually be one the strict parser rejects",
);
assert_eq!(advertised_nar_url(partial).as_deref(), Some("nar/abc.nar.xz"));
}
#[tokio::test]
async fn the_in_memory_index_is_a_set_per_nar() {
let ix = MemNarRefIndex::new();
assert!(ix.is_empty());
ix.record("nar/x.nar", "aaa").await.unwrap();
ix.record("nar/x.nar", "bbb").await.unwrap();
ix.record("nar/x.nar", "aaa").await.unwrap();
assert_eq!(ix.referrers("nar/x.nar").await.unwrap(), vec!["aaa", "bbb"]);
assert_eq!(ix.len(), 2);
ix.forget("nar/x.nar", "aaa").await.unwrap();
assert_eq!(ix.referrers("nar/x.nar").await.unwrap(), vec!["bbb"]);
ix.forget("nar/x.nar", "aaa").await.unwrap();
ix.forget("nar/absent.nar", "zzz").await.unwrap();
ix.forget("nar/x.nar", "bbb").await.unwrap();
assert!(ix.referrers("nar/x.nar").await.unwrap().is_empty());
assert!(ix.is_empty());
}
}