Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
167671fd79 | ||
|
|
339a5bd06a |
@@ -22,36 +22,5 @@ jobs:
|
||||
run: rustup component add rustfmt clippy
|
||||
- name: Install thumbv7em-none-eabihf target
|
||||
run: rustup target add thumbv7em-none-eabihf
|
||||
- name: Install cargo-audit
|
||||
run: cargo install cargo-audit --locked
|
||||
- name: Install cargo-deny
|
||||
run: cargo install cargo-deny --locked
|
||||
- name: Run CI script
|
||||
run: bash scripts/ci-test.sh
|
||||
benchmark:
|
||||
runs-on: ubuntu-latest
|
||||
container: rust:latest
|
||||
if: github.ref == 'refs/heads/main' || github.event_name == 'pull_request'
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Cache cargo registry/target
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
target
|
||||
key: ${{ runner.os }}-bench-${{ hashFiles('**/Cargo.lock') }}
|
||||
- name: Save baseline on main
|
||||
if: github.ref == 'refs/heads/main'
|
||||
run: |
|
||||
cargo bench -p clawhdf5-agent --bench memory_bench -- --save-baseline main 2>&1 || true
|
||||
- name: Compare against baseline on PRs
|
||||
if: github.event_name == 'pull_request'
|
||||
run: |
|
||||
# Download the saved baseline artifact from the target branch if available
|
||||
cargo bench -p clawhdf5-agent --bench memory_bench -- --load-baseline main --baseline main 2>&1 | tee /tmp/bench_output.txt || true
|
||||
if grep -q "Performance has regressed" /tmp/bench_output.txt; then
|
||||
echo "::error::Benchmark regression detected — see bench output above"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -31,9 +31,3 @@ tempfile = "3"
|
||||
criterion = { version = "0.5", features = ["html_reports"] }
|
||||
half = "2.7"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
|
||||
# Enable overflow checks for the format parser in release mode — this crate
|
||||
# processes untrusted byte offsets where a silent wrapping integer would be a
|
||||
# safety/correctness hazard.
|
||||
[profile.release.package.clawhdf5-format]
|
||||
overflow-checks = true
|
||||
|
||||
@@ -23,7 +23,6 @@ rayon = { version = "1", optional = true }
|
||||
matrixmultiply = { version = "0.3", optional = true }
|
||||
cblas-sys = { version = "0.1", optional = true }
|
||||
tokio = { version = "1", features = ["rt", "sync", "macros", "time"], optional = true }
|
||||
ring = { version = "0.17", optional = true }
|
||||
|
||||
[target.'cfg(target_os = "macos")'.dependencies]
|
||||
accelerate-src = { version = "0.3", optional = true }
|
||||
@@ -61,5 +60,3 @@ fast-math = ["matrixmultiply"]
|
||||
accelerate = ["accelerate-src", "cblas-sys"]
|
||||
openblas = ["openblas-src", "cblas-sys"]
|
||||
async = ["tokio"]
|
||||
encryption = ["ring"]
|
||||
signing = ["ring"]
|
||||
|
||||
@@ -1,23 +0,0 @@
|
||||
[package]
|
||||
name = "clawhdf5-agent-fuzz"
|
||||
version = "0.0.0"
|
||||
publish = false
|
||||
edition = "2024"
|
||||
|
||||
[package.metadata]
|
||||
cargo-fuzz = true
|
||||
|
||||
[dependencies]
|
||||
libfuzzer-sys = "0.4"
|
||||
tempfile = "3"
|
||||
|
||||
[dependencies.clawhdf5-agent]
|
||||
path = ".."
|
||||
|
||||
[workspace]
|
||||
members = ["."]
|
||||
|
||||
[[bin]]
|
||||
name = "fuzz_wal_replay"
|
||||
path = "fuzz_targets/fuzz_wal_replay.rs"
|
||||
doc = false
|
||||
@@ -1,21 +0,0 @@
|
||||
#![no_main]
|
||||
use libfuzzer_sys::fuzz_target;
|
||||
|
||||
use std::io::Write as _;
|
||||
|
||||
fuzz_target!(|data: &[u8]| {
|
||||
// Write the fuzz input to a temporary file, then run it through the WAL
|
||||
// replay path. The goal: verify that no arbitrary byte sequence causes a
|
||||
// panic, OOM, or other safety violation. CRC32 mismatches, truncated
|
||||
// entries, bad magic bytes, and oversized length fields are all expected to
|
||||
// return an error (not crash).
|
||||
let Ok(mut tmp) = tempfile::NamedTempFile::new() else {
|
||||
return;
|
||||
};
|
||||
if tmp.write_all(data).is_err() {
|
||||
return;
|
||||
}
|
||||
// Flush so the reader sees the data.
|
||||
let _ = tmp.flush();
|
||||
let _ = clawhdf5_agent::wal::WalFile::read_entries(tmp.path());
|
||||
});
|
||||
@@ -262,176 +262,6 @@ impl WriteAnomalyDetector {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// EmbeddingAnomalyDetector — embedding-space outlier detection
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Outcome of submitting an embedding to the detector.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum EmbeddingVerdict {
|
||||
/// Embedding is within the learned distribution.
|
||||
Accept,
|
||||
/// Embedding is a statistical outlier. Treat as quarantined until
|
||||
/// explicitly promoted by a trusted code path.
|
||||
Quarantine(String),
|
||||
}
|
||||
|
||||
/// Detects embedding-space outliers via diagonal Mahalanobis distance.
|
||||
///
|
||||
/// The detector learns a running mean and per-dimension variance from
|
||||
/// accepted embeddings using Welford's online algorithm. A new embedding
|
||||
/// whose squared Mahalanobis distance (using the diagonal covariance) exceeds
|
||||
/// `threshold_sigma_sq` standard-deviation-units is flagged as an outlier.
|
||||
///
|
||||
/// The first `warmup` embeddings are always accepted to seed the statistics
|
||||
/// before outlier detection is meaningful.
|
||||
///
|
||||
/// # Embedding-source quarantine
|
||||
///
|
||||
/// When the source is [`MemorySource::Tool`] and the embedding is a spatial
|
||||
/// outlier, the verdict is [`EmbeddingVerdict::Quarantine`]. Callers are
|
||||
/// expected to store the embedding in a quarantine dataset rather than the
|
||||
/// primary memory store, and to require explicit operator promotion before
|
||||
/// the embedding participates in retrieval.
|
||||
#[derive(Debug)]
|
||||
pub struct EmbeddingAnomalyDetector {
|
||||
/// Number of embeddings to absorb before performing outlier checks.
|
||||
warmup: usize,
|
||||
/// Threshold: if the mean squared per-dimension z-score exceeds this
|
||||
/// value the embedding is flagged. A value of `9.0` corresponds roughly
|
||||
/// to 3σ per dimension under a Gaussian model.
|
||||
threshold_sigma_sq: f32,
|
||||
/// Running count of accepted embeddings (used for Welford's update).
|
||||
count: usize,
|
||||
/// Welford's running mean per dimension.
|
||||
mean: Vec<f64>,
|
||||
/// Welford's running M2 (sum of squared deviations) per dimension.
|
||||
m2: Vec<f64>,
|
||||
}
|
||||
|
||||
impl EmbeddingAnomalyDetector {
|
||||
/// Create a detector for embeddings of the given dimensionality.
|
||||
///
|
||||
/// * `dim` — embedding dimension.
|
||||
/// * `warmup` — number of embeddings accepted unconditionally to seed
|
||||
/// the mean/variance statistics. Minimum effective value is 2.
|
||||
/// * `threshold_sigma_sq` — mean squared z-score threshold; 9.0 is a
|
||||
/// reasonable default (≈3σ per dimension).
|
||||
pub fn new(dim: usize, warmup: usize, threshold_sigma_sq: f32) -> Self {
|
||||
Self {
|
||||
warmup: warmup.max(2),
|
||||
threshold_sigma_sq,
|
||||
count: 0,
|
||||
mean: vec![0.0f64; dim],
|
||||
m2: vec![0.0f64; dim],
|
||||
}
|
||||
}
|
||||
|
||||
/// Evaluate `embedding` and update the running statistics.
|
||||
///
|
||||
/// Returns [`EmbeddingVerdict::Accept`] if the embedding is within the
|
||||
/// learned distribution (or the detector is still in warmup), or
|
||||
/// [`EmbeddingVerdict::Quarantine`] if it is a spatial outlier.
|
||||
///
|
||||
/// The statistics are updated unconditionally so that the detector adapts
|
||||
/// to the distribution even when embeddings are quarantined — this prevents
|
||||
/// the mean from drifting away from the true distribution if many outliers
|
||||
/// arrive in a batch.
|
||||
pub fn evaluate(&mut self, embedding: &[f32], source: &MemorySource) -> EmbeddingVerdict {
|
||||
if embedding.len() != self.mean.len() {
|
||||
// Dimension mismatch — reject without updating stats.
|
||||
return EmbeddingVerdict::Quarantine(format!(
|
||||
"embedding dimension {} does not match detector dimension {}",
|
||||
embedding.len(),
|
||||
self.mean.len()
|
||||
));
|
||||
}
|
||||
|
||||
// Snapshot pre-update stats for outlier scoring (so the candidate point
|
||||
// cannot dilute its own z-score by pulling the mean toward itself).
|
||||
let pre_count = self.count;
|
||||
let pre_mean = self.mean.clone();
|
||||
let pre_m2 = self.m2.clone();
|
||||
|
||||
// Welford online update — always runs so stats stay current.
|
||||
self.count += 1;
|
||||
let n = self.count as f64;
|
||||
for (i, &x) in embedding.iter().enumerate() {
|
||||
let x64 = x as f64;
|
||||
let delta = x64 - self.mean[i];
|
||||
self.mean[i] += delta / n;
|
||||
let delta2 = x64 - self.mean[i];
|
||||
self.m2[i] += delta * delta2;
|
||||
}
|
||||
|
||||
// During warmup, always accept.
|
||||
if self.count <= self.warmup {
|
||||
return EmbeddingVerdict::Accept;
|
||||
}
|
||||
|
||||
// Score against pre-update distribution so the candidate cannot move
|
||||
// the mean toward itself and inflate acceptance.
|
||||
let pre_n = pre_count as f64;
|
||||
let mut sum_zsq = 0.0f64;
|
||||
let mut dims_with_variance = 0usize;
|
||||
// Whether any dimension shows a non-trivial deviation from a zero-variance mean.
|
||||
let mut zero_var_outlier = false;
|
||||
for i in 0..pre_mean.len() {
|
||||
// Need at least 2 points to have a variance estimate.
|
||||
if pre_count < 2 {
|
||||
continue;
|
||||
}
|
||||
let var = pre_m2[i] / (pre_n - 1.0);
|
||||
if var > 1e-12 {
|
||||
let z = (embedding[i] as f64 - pre_mean[i]) / var.sqrt();
|
||||
sum_zsq += z * z;
|
||||
dims_with_variance += 1;
|
||||
} else {
|
||||
// Variance is effectively zero: all training points were identical in this
|
||||
// dimension. Any meaningful deviation from the exact mean is an outlier
|
||||
// by definition — flag it so the caller sees Quarantine.
|
||||
let dev = (embedding[i] as f64 - pre_mean[i]).abs();
|
||||
if dev > 1e-6 {
|
||||
zero_var_outlier = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dims_with_variance == 0 {
|
||||
// No estimated variance in any dimension.
|
||||
if zero_var_outlier {
|
||||
return EmbeddingVerdict::Quarantine(format!(
|
||||
"embedding-space outlier (deviation from zero-variance mean, source={:?})",
|
||||
source
|
||||
));
|
||||
}
|
||||
// All dimensions match the mean exactly — accept.
|
||||
return EmbeddingVerdict::Accept;
|
||||
}
|
||||
|
||||
let mean_zsq = (sum_zsq / dims_with_variance as f64) as f32;
|
||||
if mean_zsq > self.threshold_sigma_sq {
|
||||
let reason = format!(
|
||||
"embedding-space outlier (mean z²={:.2}, threshold={:.2}, source={:?})",
|
||||
mean_zsq, self.threshold_sigma_sq, source
|
||||
);
|
||||
EmbeddingVerdict::Quarantine(reason)
|
||||
} else {
|
||||
EmbeddingVerdict::Accept
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of embeddings seen so far (including warmup and quarantined).
|
||||
pub fn count(&self) -> usize {
|
||||
self.count
|
||||
}
|
||||
|
||||
/// Whether the detector has completed its warmup phase.
|
||||
pub fn is_warmed_up(&self) -> bool {
|
||||
self.count > self.warmup
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -630,72 +460,4 @@ mod tests {
|
||||
assert_eq!(det.session_count("sess-b"), 1);
|
||||
assert_eq!(det.session_count("unknown"), 0);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// EmbeddingAnomalyDetector tests
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
fn ebed(v: Vec<f32>) -> Vec<f32> {
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warmup_embeddings_always_accepted() {
|
||||
let mut det = EmbeddingAnomalyDetector::new(3, 5, 9.0);
|
||||
let emb = ebed(vec![1.0, 0.0, 0.0]);
|
||||
for _ in 0..5 {
|
||||
assert_eq!(
|
||||
det.evaluate(&emb, &MemorySource::User),
|
||||
EmbeddingVerdict::Accept
|
||||
);
|
||||
}
|
||||
assert!(!det.is_warmed_up()); // count == warmup, not strictly greater
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn in_distribution_embedding_accepted() {
|
||||
let mut det = EmbeddingAnomalyDetector::new(2, 3, 9.0);
|
||||
// Seed with embeddings near (1.0, 1.0).
|
||||
det.evaluate(&[1.0, 1.0], &MemorySource::User);
|
||||
det.evaluate(&[1.1, 0.9], &MemorySource::User);
|
||||
det.evaluate(&[0.9, 1.1], &MemorySource::User);
|
||||
// A nearby embedding should be accepted.
|
||||
assert_eq!(
|
||||
det.evaluate(&[1.0, 1.0], &MemorySource::User),
|
||||
EmbeddingVerdict::Accept
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn outlier_embedding_quarantined() {
|
||||
let mut det = EmbeddingAnomalyDetector::new(2, 3, 9.0);
|
||||
// Seed: all embeddings near (0.0, 0.0) with very low variance.
|
||||
for _ in 0..3 {
|
||||
det.evaluate(&[0.0, 0.0], &MemorySource::User);
|
||||
}
|
||||
// A far-away embedding should be quarantined.
|
||||
let verdict = det.evaluate(&[100.0, 100.0], &MemorySource::Tool);
|
||||
assert!(
|
||||
matches!(verdict, EmbeddingVerdict::Quarantine(_)),
|
||||
"expected Quarantine, got {:?}",
|
||||
verdict
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dimension_mismatch_quarantined() {
|
||||
let mut det = EmbeddingAnomalyDetector::new(4, 2, 9.0);
|
||||
let verdict = det.evaluate(&[1.0, 2.0], &MemorySource::User);
|
||||
assert!(matches!(verdict, EmbeddingVerdict::Quarantine(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_tracks_all_evaluations() {
|
||||
let mut det = EmbeddingAnomalyDetector::new(2, 2, 9.0);
|
||||
det.evaluate(&[1.0, 0.0], &MemorySource::User);
|
||||
det.evaluate(&[0.0, 1.0], &MemorySource::User);
|
||||
det.evaluate(&[1.0, 1.0], &MemorySource::User);
|
||||
assert_eq!(det.count(), 3);
|
||||
assert!(det.is_warmed_up());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -37,7 +37,7 @@
|
||||
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||
//! mem.save(entry).await?; // buffered → background writer
|
||||
//! mem.save_batch(entries).await?; // also buffered
|
||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.4, 0.6, 5).await;
|
||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
||||
//! mem.shutdown().await?; // final flush + stop
|
||||
//! ```
|
||||
|
||||
|
||||
@@ -218,171 +218,6 @@ impl BM25Index {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sidecar serialization (BM25 persistence — INT-09)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Magic bytes for the `.bm25` sidecar format.
|
||||
const SIDECAR_MAGIC: [u8; 4] = [0x42, 0x4D, 0x32, 0x35]; // "BM25"
|
||||
/// Current sidecar format version.
|
||||
const SIDECAR_VERSION: u8 = 0x01;
|
||||
|
||||
impl BM25Index {
|
||||
/// Serialize the index into a compact binary format suitable for writing to
|
||||
/// the `.bm25` sidecar file.
|
||||
///
|
||||
/// Format:
|
||||
/// ```text
|
||||
/// [4] magic "BM25"
|
||||
/// [1] version byte
|
||||
/// [4] doc_lengths.len() as le u32 (= total chunk count, including tombstones)
|
||||
/// [4] num_docs as le u32
|
||||
/// [4] avg_dl as le f32
|
||||
/// [N*4] doc_lengths as le u32 each
|
||||
/// [4] inverted entry count as le u32
|
||||
/// per inverted entry:
|
||||
/// [4] token byte length as le u32
|
||||
/// [L] UTF-8 token bytes
|
||||
/// [4] posting count as le u32
|
||||
/// per posting: [4] doc_id le u32, [4] term_freq le u32
|
||||
/// [4] idf entry count as le u32
|
||||
/// per idf entry:
|
||||
/// [4] token byte length as le u32
|
||||
/// [L] UTF-8 token bytes
|
||||
/// [4] idf score as le f32
|
||||
/// ```
|
||||
pub fn to_bytes(&self) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(
|
||||
9 + self.doc_lengths.len() * 4 + self.inverted.len() * 16 + self.idf_cache.len() * 16,
|
||||
);
|
||||
|
||||
buf.extend_from_slice(&SIDECAR_MAGIC);
|
||||
buf.push(SIDECAR_VERSION);
|
||||
buf.extend_from_slice(&(self.doc_lengths.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&(self.num_docs as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&self.avg_dl.to_le_bytes());
|
||||
|
||||
for &dl in &self.doc_lengths {
|
||||
buf.extend_from_slice(&dl.to_le_bytes());
|
||||
}
|
||||
|
||||
buf.extend_from_slice(&(self.inverted.len() as u32).to_le_bytes());
|
||||
for (token, postings) in &self.inverted {
|
||||
let tb = token.as_bytes();
|
||||
buf.extend_from_slice(&(tb.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(tb);
|
||||
buf.extend_from_slice(&(postings.len() as u32).to_le_bytes());
|
||||
for &(doc_id, tf) in postings {
|
||||
buf.extend_from_slice(&(doc_id as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&tf.to_le_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
buf.extend_from_slice(&(self.idf_cache.len() as u32).to_le_bytes());
|
||||
for (token, &idf) in &self.idf_cache {
|
||||
let tb = token.as_bytes();
|
||||
buf.extend_from_slice(&(tb.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(tb);
|
||||
buf.extend_from_slice(&idf.to_le_bytes());
|
||||
}
|
||||
|
||||
buf
|
||||
}
|
||||
|
||||
/// Deserialize an index from the bytes produced by [`to_bytes`].
|
||||
///
|
||||
/// Returns `None` if the bytes are malformed (bad magic, wrong version,
|
||||
/// truncated data, or non-UTF-8 tokens). The caller should fall back to
|
||||
/// [`BM25Index::build`] when `None` is returned.
|
||||
///
|
||||
/// `expected_doc_count` is the total number of chunks (including tombstones)
|
||||
/// currently in the cache. If it does not match the serialized
|
||||
/// `doc_lengths.len()`, the sidecar is stale and `None` is returned.
|
||||
pub fn from_bytes(data: &[u8], expected_doc_count: usize) -> Option<Self> {
|
||||
let mut pos = 0usize;
|
||||
|
||||
macro_rules! read_bytes {
|
||||
($n:expr) => {{
|
||||
let end = pos + $n;
|
||||
if end > data.len() {
|
||||
return None;
|
||||
}
|
||||
let slice = &data[pos..end];
|
||||
pos = end;
|
||||
slice
|
||||
}};
|
||||
}
|
||||
macro_rules! read_u32 {
|
||||
() => {{
|
||||
u32::from_le_bytes(read_bytes!(4).try_into().ok()?)
|
||||
}};
|
||||
}
|
||||
macro_rules! read_f32 {
|
||||
() => {{
|
||||
f32::from_le_bytes(read_bytes!(4).try_into().ok()?)
|
||||
}};
|
||||
}
|
||||
|
||||
// Magic + version
|
||||
let magic = read_bytes!(4);
|
||||
if magic != SIDECAR_MAGIC {
|
||||
return None;
|
||||
}
|
||||
let version = read_bytes!(1)[0];
|
||||
if version != SIDECAR_VERSION {
|
||||
return None;
|
||||
}
|
||||
|
||||
// doc_lengths
|
||||
let doc_count = read_u32!() as usize;
|
||||
if doc_count != expected_doc_count {
|
||||
return None; // stale sidecar
|
||||
}
|
||||
let num_docs = read_u32!() as usize;
|
||||
let avg_dl = read_f32!();
|
||||
let mut doc_lengths = Vec::with_capacity(doc_count);
|
||||
for _ in 0..doc_count {
|
||||
doc_lengths.push(read_u32!());
|
||||
}
|
||||
|
||||
// inverted index
|
||||
let inv_count = read_u32!() as usize;
|
||||
let mut inverted: HashMap<String, Vec<(usize, u32)>> = HashMap::with_capacity(inv_count);
|
||||
for _ in 0..inv_count {
|
||||
let tlen = read_u32!() as usize;
|
||||
let token = std::str::from_utf8(read_bytes!(tlen)).ok()?.to_string();
|
||||
let plen = read_u32!() as usize;
|
||||
let mut postings = Vec::with_capacity(plen);
|
||||
for _ in 0..plen {
|
||||
let doc_id = read_u32!() as usize;
|
||||
let tf = read_u32!();
|
||||
postings.push((doc_id, tf));
|
||||
}
|
||||
inverted.insert(token, postings);
|
||||
}
|
||||
|
||||
// idf cache
|
||||
let idf_count = read_u32!() as usize;
|
||||
let mut idf_cache: HashMap<String, f32> = HashMap::with_capacity(idf_count);
|
||||
for _ in 0..idf_count {
|
||||
let tlen = read_u32!() as usize;
|
||||
let token = std::str::from_utf8(read_bytes!(tlen)).ok()?.to_string();
|
||||
let idf = read_f32!();
|
||||
idf_cache.insert(token, idf);
|
||||
}
|
||||
|
||||
Some(Self {
|
||||
inverted,
|
||||
idf_cache,
|
||||
doc_lengths,
|
||||
avg_dl,
|
||||
num_docs,
|
||||
k1: DEFAULT_K1,
|
||||
b: DEFAULT_B,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
||||
/// filter empty tokens.
|
||||
fn tokenize(text: &str) -> Vec<String> {
|
||||
@@ -616,74 +451,4 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Sidecar serialization round-trip (INT-09)
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn sidecar_round_trip_preserves_search_results() {
|
||||
let docs = vec![
|
||||
"the quick brown fox jumps over the lazy dog".to_string(),
|
||||
"rust programming language systems programming".to_string(),
|
||||
"python scripting and data science".to_string(),
|
||||
];
|
||||
let tombstones = vec![0u8, 0, 0];
|
||||
let original = BM25Index::build(&docs, &tombstones);
|
||||
|
||||
// Serialize then deserialize.
|
||||
let bytes = original.to_bytes();
|
||||
let restored =
|
||||
BM25Index::from_bytes(&bytes, docs.len()).expect("round-trip must succeed");
|
||||
|
||||
// Both indexes must return identical results for the same query.
|
||||
let orig_results = original.search("rust programming", 10);
|
||||
let rest_results = restored.search("rust programming", 10);
|
||||
assert_eq!(
|
||||
orig_results.len(),
|
||||
rest_results.len(),
|
||||
"result count mismatch"
|
||||
);
|
||||
for (a, b) in orig_results.iter().zip(rest_results.iter()) {
|
||||
assert_eq!(a.0, b.0, "doc_id mismatch after round-trip");
|
||||
assert!(
|
||||
(a.1 - b.1).abs() < 1e-5,
|
||||
"score mismatch: {} vs {} for doc {}",
|
||||
a.1,
|
||||
b.1,
|
||||
a.0
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_stale_doc_count_rejected() {
|
||||
let docs = vec!["hello world".to_string()];
|
||||
let tombstones = vec![0u8];
|
||||
let idx = BM25Index::build(&docs, &tombstones);
|
||||
let bytes = idx.to_bytes();
|
||||
// Pass wrong expected_doc_count — should return None.
|
||||
assert!(BM25Index::from_bytes(&bytes, 999).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_bad_magic_rejected() {
|
||||
let docs = vec!["hello".to_string()];
|
||||
let tombstones = vec![0u8];
|
||||
let idx = BM25Index::build(&docs, &tombstones);
|
||||
let mut bytes = idx.to_bytes();
|
||||
// Corrupt the magic bytes.
|
||||
bytes[0] = 0xFF;
|
||||
assert!(BM25Index::from_bytes(&bytes, 1).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_empty_index_round_trip() {
|
||||
let docs: Vec<String> = vec![];
|
||||
let tombstones: Vec<u8> = vec![];
|
||||
let idx = BM25Index::build(&docs, &tombstones);
|
||||
let bytes = idx.to_bytes();
|
||||
let restored = BM25Index::from_bytes(&bytes, 0).expect("empty index must round-trip");
|
||||
assert_eq!(restored.search("anything", 5).len(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,268 +0,0 @@
|
||||
//! AES-256-GCM encryption at rest for agent memory files.
|
||||
//!
|
||||
//! # Envelope format
|
||||
//!
|
||||
//! ```text
|
||||
//! [8 bytes magic "CLAWENC\x00"]
|
||||
//! [4 bytes version = 1, little-endian u32]
|
||||
//! [16 bytes PBKDF2 salt]
|
||||
//! [12 bytes AES-GCM nonce]
|
||||
//! [N bytes ciphertext + 16-byte GCM authentication tag]
|
||||
//! ```
|
||||
//!
|
||||
//! Keys are derived from a caller-supplied passphrase using PBKDF2-HMAC-SHA256
|
||||
//! with 200 000 iterations. The same derived key can also be passed directly
|
||||
//! as a raw 32-byte value via [`seal_with_key`] / [`open_with_key`] when the
|
||||
//! caller manages key material externally (e.g. from a hardware key store).
|
||||
|
||||
use std::num::NonZeroU32;
|
||||
|
||||
use ring::aead::{
|
||||
Aad, AES_256_GCM, BoundKey, Nonce, NonceSequence, OpeningKey, SealingKey, UnboundKey,
|
||||
NONCE_LEN,
|
||||
};
|
||||
use ring::error::Unspecified;
|
||||
use ring::pbkdf2;
|
||||
use ring::rand::{SecureRandom, SystemRandom};
|
||||
|
||||
/// Envelope magic bytes.
|
||||
const MAGIC: &[u8; 8] = b"CLAWENC\x00";
|
||||
/// Envelope version.
|
||||
const VERSION: u32 = 1;
|
||||
/// PBKDF2 iteration count (NIST SP 800-132 recommends ≥ 10 000; we use 200 000).
|
||||
const PBKDF2_ITERS: NonZeroU32 = unsafe { NonZeroU32::new_unchecked(200_000) };
|
||||
/// Salt length in bytes.
|
||||
const SALT_LEN: usize = 16;
|
||||
/// Derived key length (AES-256 = 32 bytes).
|
||||
const KEY_LEN: usize = 32;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum EncryptionError {
|
||||
/// Envelope is too short or has incorrect magic/version.
|
||||
MalformedEnvelope,
|
||||
/// AES-GCM authentication tag check failed (wrong key or tampered data).
|
||||
AuthenticationFailed,
|
||||
/// OS random source unavailable.
|
||||
RngFailure,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for EncryptionError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
EncryptionError::MalformedEnvelope => write!(f, "malformed encryption envelope"),
|
||||
EncryptionError::AuthenticationFailed => {
|
||||
write!(f, "AES-GCM authentication failed (wrong key or corrupted data)")
|
||||
}
|
||||
EncryptionError::RngFailure => write!(f, "OS RNG unavailable"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Key derivation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Derive a 32-byte AES-256 key from a passphrase and salt using
|
||||
/// PBKDF2-HMAC-SHA256.
|
||||
pub fn derive_key(passphrase: &[u8], salt: &[u8]) -> [u8; KEY_LEN] {
|
||||
let mut key = [0u8; KEY_LEN];
|
||||
pbkdf2::derive(pbkdf2::PBKDF2_HMAC_SHA256, PBKDF2_ITERS, salt, passphrase, &mut key);
|
||||
key
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Nonce helpers (ring requires a NonceSequence trait)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct FixedNonce([u8; NONCE_LEN]);
|
||||
|
||||
impl NonceSequence for FixedNonce {
|
||||
fn advance(&mut self) -> Result<Nonce, Unspecified> {
|
||||
Ok(Nonce::assume_unique_for_key(self.0))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Core seal / open (raw key)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Encrypt `plaintext` with a raw 32-byte key.
|
||||
///
|
||||
/// Returns the serialized envelope (magic + salt placeholder zeroed +
|
||||
/// nonce + ciphertext). The `salt` field in the envelope is left as zeroes
|
||||
/// because the caller supplies the key directly; use [`seal`] for passphrase-
|
||||
/// based encryption.
|
||||
pub fn seal_with_key(key: &[u8; KEY_LEN], plaintext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
|
||||
let rng = SystemRandom::new();
|
||||
|
||||
let mut nonce_bytes = [0u8; NONCE_LEN];
|
||||
rng.fill(&mut nonce_bytes).map_err(|_| EncryptionError::RngFailure)?;
|
||||
|
||||
let unbound = UnboundKey::new(&AES_256_GCM, key).expect("valid key length");
|
||||
let mut sealing = SealingKey::new(unbound, FixedNonce(nonce_bytes));
|
||||
|
||||
let mut buf: Vec<u8> = plaintext.to_vec();
|
||||
// AES-256-GCM appends a 16-byte authentication tag.
|
||||
buf.extend_from_slice(&[0u8; 16]);
|
||||
let tag = sealing
|
||||
.seal_in_place_separate_tag(Aad::empty(), &mut buf[..plaintext.len()])
|
||||
.map_err(|_| EncryptionError::RngFailure)?;
|
||||
buf[plaintext.len()..].copy_from_slice(tag.as_ref());
|
||||
|
||||
let total = 8 + 4 + SALT_LEN + NONCE_LEN + buf.len();
|
||||
let mut out = Vec::with_capacity(total);
|
||||
out.extend_from_slice(MAGIC);
|
||||
out.extend_from_slice(&VERSION.to_le_bytes());
|
||||
out.extend_from_slice(&[0u8; SALT_LEN]); // salt placeholder
|
||||
out.extend_from_slice(&nonce_bytes);
|
||||
out.extend_from_slice(&buf);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Decrypt an envelope produced by [`seal_with_key`] using the same raw key.
|
||||
pub fn open_with_key(key: &[u8; KEY_LEN], envelope: &[u8]) -> Result<Vec<u8>, EncryptionError> {
|
||||
let header = 8 + 4 + SALT_LEN + NONCE_LEN;
|
||||
if envelope.len() < header + 16 {
|
||||
return Err(EncryptionError::MalformedEnvelope);
|
||||
}
|
||||
if &envelope[..8] != MAGIC {
|
||||
return Err(EncryptionError::MalformedEnvelope);
|
||||
}
|
||||
let ver = u32::from_le_bytes(envelope[8..12].try_into().unwrap());
|
||||
if ver != VERSION {
|
||||
return Err(EncryptionError::MalformedEnvelope);
|
||||
}
|
||||
let nonce_start = 8 + 4 + SALT_LEN;
|
||||
let nonce_bytes: [u8; NONCE_LEN] =
|
||||
envelope[nonce_start..nonce_start + NONCE_LEN].try_into().unwrap();
|
||||
|
||||
let unbound = UnboundKey::new(&AES_256_GCM, key).expect("valid key length");
|
||||
let mut opening = OpeningKey::new(unbound, FixedNonce(nonce_bytes));
|
||||
|
||||
let mut buf: Vec<u8> = envelope[header..].to_vec();
|
||||
let plaintext = opening
|
||||
.open_in_place(Aad::empty(), &mut buf)
|
||||
.map_err(|_| EncryptionError::AuthenticationFailed)?;
|
||||
Ok(plaintext.to_vec())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Passphrase-based seal / open
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Encrypt `plaintext` using a passphrase.
|
||||
///
|
||||
/// A random 16-byte PBKDF2 salt is generated, stored in the envelope header,
|
||||
/// and used to derive the AES-256 key.
|
||||
pub fn seal(passphrase: &[u8], plaintext: &[u8]) -> Result<Vec<u8>, EncryptionError> {
|
||||
let rng = SystemRandom::new();
|
||||
|
||||
let mut salt = [0u8; SALT_LEN];
|
||||
rng.fill(&mut salt).map_err(|_| EncryptionError::RngFailure)?;
|
||||
|
||||
let key = derive_key(passphrase, &salt);
|
||||
|
||||
let mut envelope = seal_with_key(&key, plaintext)?;
|
||||
// Overwrite the zeroed salt placeholder with the real salt.
|
||||
let salt_offset = 8 + 4;
|
||||
envelope[salt_offset..salt_offset + SALT_LEN].copy_from_slice(&salt);
|
||||
Ok(envelope)
|
||||
}
|
||||
|
||||
/// Decrypt an envelope produced by [`seal`].
|
||||
pub fn open(passphrase: &[u8], envelope: &[u8]) -> Result<Vec<u8>, EncryptionError> {
|
||||
let header = 8 + 4 + SALT_LEN + NONCE_LEN;
|
||||
if envelope.len() < header + 16 {
|
||||
return Err(EncryptionError::MalformedEnvelope);
|
||||
}
|
||||
if &envelope[..8] != MAGIC {
|
||||
return Err(EncryptionError::MalformedEnvelope);
|
||||
}
|
||||
let salt_start = 8 + 4;
|
||||
let salt = &envelope[salt_start..salt_start + SALT_LEN];
|
||||
let key = derive_key(passphrase, salt);
|
||||
open_with_key(&key, envelope)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn seal_open_roundtrip_raw_key() {
|
||||
let key = [0xABu8; 32];
|
||||
let plaintext = b"hello, ClawHDF5 AES-256-GCM!";
|
||||
let envelope = seal_with_key(&key, plaintext).unwrap();
|
||||
let recovered = open_with_key(&key, &envelope).unwrap();
|
||||
assert_eq!(recovered, plaintext);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seal_open_roundtrip_passphrase() {
|
||||
let passphrase = b"correct horse battery staple";
|
||||
let plaintext = b"secret agent memory bytes";
|
||||
let envelope = seal(passphrase, plaintext).unwrap();
|
||||
let recovered = open(passphrase, &envelope).unwrap();
|
||||
assert_eq!(recovered, plaintext);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrong_key_fails_authentication() {
|
||||
let key_a = [0x11u8; 32];
|
||||
let key_b = [0x22u8; 32];
|
||||
let envelope = seal_with_key(&key_a, b"sensitive").unwrap();
|
||||
assert!(matches!(open_with_key(&key_b, &envelope), Err(EncryptionError::AuthenticationFailed)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrong_passphrase_fails_authentication() {
|
||||
let envelope = seal(b"right", b"data").unwrap();
|
||||
assert!(matches!(open(b"wrong", &envelope), Err(EncryptionError::AuthenticationFailed)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tampered_ciphertext_fails_authentication() {
|
||||
let key = [0xCCu8; 32];
|
||||
let mut envelope = seal_with_key(&key, b"data").unwrap();
|
||||
let last = envelope.len() - 1;
|
||||
envelope[last] ^= 0xFF;
|
||||
assert!(matches!(open_with_key(&key, &envelope), Err(EncryptionError::AuthenticationFailed)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_envelope_detected() {
|
||||
assert!(matches!(open_with_key(&[0u8; 32], b"too short"), Err(EncryptionError::MalformedEnvelope)));
|
||||
let mut bad_magic = vec![0u8; 64];
|
||||
assert!(matches!(open_with_key(&[0u8; 32], &bad_magic), Err(EncryptionError::MalformedEnvelope)));
|
||||
// correct magic, wrong version
|
||||
bad_magic[..8].copy_from_slice(MAGIC);
|
||||
bad_magic[8..12].copy_from_slice(&99u32.to_le_bytes());
|
||||
assert!(matches!(open_with_key(&[0u8; 32], &bad_magic), Err(EncryptionError::MalformedEnvelope)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn derive_key_is_deterministic() {
|
||||
let k1 = derive_key(b"pass", b"salt1234567890AB");
|
||||
let k2 = derive_key(b"pass", b"salt1234567890AB");
|
||||
assert_eq!(k1, k2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_salts_produce_different_keys() {
|
||||
let k1 = derive_key(b"pass", b"salt1234567890AB");
|
||||
let k2 = derive_key(b"pass", b"SALT1234567890AB");
|
||||
assert_ne!(k1, k2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_plaintext_roundtrip() {
|
||||
let key = [0x77u8; 32];
|
||||
let envelope = seal_with_key(&key, b"").unwrap();
|
||||
let recovered = open_with_key(&key, &envelope).unwrap();
|
||||
assert!(recovered.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -439,11 +439,6 @@ impl KnowledgeCache {
|
||||
min_activation: f32,
|
||||
max_steps: usize,
|
||||
) -> Vec<(u64, f32)> {
|
||||
// decay_factor >= 1.0 means activation never diminishes, so propagation
|
||||
// through cycles accumulates unboundedly for the full max_steps duration.
|
||||
// Clamp to [0.0, 1.0) to guarantee convergence.
|
||||
let decay_factor = decay_factor.clamp(0.0, 1.0 - f32::EPSILON);
|
||||
|
||||
let mut activation: HashMap<u64, f32> = HashMap::new();
|
||||
|
||||
// Initialise seeds with activation 1.0.
|
||||
@@ -1167,63 +1162,4 @@ mod tests {
|
||||
assert!(ctx.contains("occupation"));
|
||||
assert!(ctx.contains("engineer"));
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Cycle safety — BFS and spreading_activation must not loop infinitely
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_bfs_neighbors_cycle_terminates() {
|
||||
let mut cache = KnowledgeCache::new();
|
||||
let a = cache.add_entity("A", "node", -1);
|
||||
let b = cache.add_entity("B", "node", -1);
|
||||
let c = cache.add_entity("C", "node", -1);
|
||||
// A → B → C → A (cycle)
|
||||
cache.add_relation(a, b, "link", 1.0);
|
||||
cache.add_relation(b, c, "link", 1.0);
|
||||
cache.add_relation(c, a, "link", 1.0);
|
||||
|
||||
let result = cache.bfs_neighbors(a, 10);
|
||||
// Should visit b and c exactly once, not loop forever.
|
||||
let ids: HashSet<u64> = result.iter().map(|(e, _)| e.id).collect();
|
||||
assert!(ids.contains(&b), "b must be reachable");
|
||||
assert!(ids.contains(&c), "c must be reachable");
|
||||
assert_eq!(result.len(), 2, "only b and c should appear (no duplicates)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_bfs_neighbors_self_loop_terminates() {
|
||||
let mut cache = KnowledgeCache::new();
|
||||
let a = cache.add_entity("A", "node", -1);
|
||||
// Self-loop: A → A
|
||||
cache.add_relation(a, a, "self", 1.0);
|
||||
|
||||
let result = cache.bfs_neighbors(a, 5);
|
||||
assert!(result.is_empty(), "self-loop seed should not appear in results");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_spreading_activation_cycle_converges() {
|
||||
let mut cache = KnowledgeCache::new();
|
||||
let a = cache.add_entity("A", "node", -1);
|
||||
let b = cache.add_entity("B", "node", -1);
|
||||
let c = cache.add_entity("C", "node", -1);
|
||||
// Cyclic graph A ↔ B ↔ C ↔ A with moderate weights.
|
||||
cache.add_relation(a, b, "link", 0.8);
|
||||
cache.add_relation(b, c, "link", 0.8);
|
||||
cache.add_relation(c, a, "link", 0.8);
|
||||
|
||||
// With decay_factor < 1 the activation decays per step and must
|
||||
// converge within max_steps without panicking or running forever.
|
||||
let result = cache.spreading_activation(&[a], 0.5, 0.001, 20);
|
||||
// At minimum a, b, c should all receive some activation.
|
||||
let activated_ids: HashSet<u64> = result.iter().map(|&(id, _)| id).collect();
|
||||
assert!(activated_ids.contains(&a));
|
||||
assert!(activated_ids.contains(&b));
|
||||
assert!(activated_ids.contains(&c));
|
||||
// Scores must be finite and non-negative.
|
||||
for &(_, score) in &result {
|
||||
assert!(score.is_finite() && score >= 0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,10 +20,6 @@ pub mod vector_search;
|
||||
|
||||
pub mod agents_md;
|
||||
pub mod anomaly;
|
||||
#[cfg(feature = "encryption")]
|
||||
pub mod encryption;
|
||||
#[cfg(feature = "signing")]
|
||||
pub mod signing;
|
||||
pub mod cache;
|
||||
pub mod confidence;
|
||||
pub mod consolidation;
|
||||
@@ -64,17 +60,6 @@ pub fn cosine_similarity_prenorm(
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use cache::MemoryCache;
|
||||
|
||||
/// Returns the path to the BM25 sidecar file for an HDF5 memory file at `h5_path`.
|
||||
///
|
||||
/// The sidecar lives next to the `.h5` file with a `.bm25` extension appended
|
||||
/// (e.g. `memory.h5` → `memory.h5.bm25`). It is loaded on `open()` to skip the
|
||||
/// O(N × terms) rebuild when the cache is large, and written on every `flush()`.
|
||||
fn bm25_sidecar_path(h5_path: &Path) -> PathBuf {
|
||||
let mut p = h5_path.as_os_str().to_owned();
|
||||
p.push(".bm25");
|
||||
PathBuf::from(p)
|
||||
}
|
||||
#[cfg(feature = "hnsw")]
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
@@ -242,10 +227,6 @@ pub struct HDF5Memory {
|
||||
/// search.
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: usize,
|
||||
/// Cached BM25 index. Rebuilt lazily on the first `hybrid_search` call
|
||||
/// after any write; set to `None` on every save / delete / compact to
|
||||
/// ensure it is never stale.
|
||||
bm25_cache: Option<bm25::BM25Index>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HDF5Memory {
|
||||
@@ -285,7 +266,6 @@ impl HDF5Memory {
|
||||
hnsw_dirty: false,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: 0,
|
||||
bm25_cache: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -305,16 +285,6 @@ impl HDF5Memory {
|
||||
None
|
||||
};
|
||||
|
||||
// Try to load the BM25 sidecar so the first hybrid_search after open()
|
||||
// skips the O(N × terms) rebuild. Fall back to None (lazy rebuild) if
|
||||
// the sidecar is absent, malformed, or has a mismatched doc count.
|
||||
let bm25_cache = {
|
||||
let sidecar_path = bm25_sidecar_path(&config.path);
|
||||
std::fs::read(&sidecar_path)
|
||||
.ok()
|
||||
.and_then(|b| bm25::BM25Index::from_bytes(&b, cache.chunks.len()))
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
cache,
|
||||
@@ -331,7 +301,6 @@ impl HDF5Memory {
|
||||
hnsw_dirty: true,
|
||||
#[cfg(feature = "hnsw")]
|
||||
hnsw_synced_len: 0,
|
||||
bm25_cache,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -351,35 +320,9 @@ impl HDF5Memory {
|
||||
if let Some(ref mut w) = self.wal {
|
||||
w.truncate()?;
|
||||
}
|
||||
// Persist the BM25 index alongside the .h5 file so the next open()
|
||||
// can skip the O(N × terms) rebuild. Only write when we have a cached
|
||||
// index; if there is none, leave any existing sidecar in place.
|
||||
if let Some(ref idx) = self.bm25_cache {
|
||||
let sidecar_path = bm25_sidecar_path(&self.config.path);
|
||||
let bytes = idx.to_bytes();
|
||||
// Best-effort: a sidecar write failure is not fatal — the caller
|
||||
// will rebuild from scratch on the next open().
|
||||
let _ = std::fs::write(&sidecar_path, &bytes);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Path to the `.bm25` sidecar file for this memory store.
|
||||
fn bm25_sidecar_path(&self) -> std::path::PathBuf {
|
||||
bm25_sidecar_path(&self.config.path)
|
||||
}
|
||||
|
||||
/// Try to load the BM25 index from the `.bm25` sidecar file.
|
||||
///
|
||||
/// Returns `Some(index)` if the sidecar exists and is valid for the current
|
||||
/// cache state (same total chunk count including tombstones). Returns
|
||||
/// `None` if the sidecar is absent, malformed, or stale.
|
||||
fn load_bm25_sidecar(&self) -> Option<bm25::BM25Index> {
|
||||
let sidecar_path = self.bm25_sidecar_path();
|
||||
let bytes = std::fs::read(&sidecar_path).ok()?;
|
||||
bm25::BM25Index::from_bytes(&bytes, self.cache.chunks.len())
|
||||
}
|
||||
|
||||
// ---- HNSW index maintenance --------------------------------------------
|
||||
//
|
||||
// The index mirrors the cache: HNSW node id == cache index, kept aligned by
|
||||
@@ -574,7 +517,6 @@ impl HDF5Memory {
|
||||
);
|
||||
// In-place embedding change: the index node is stale, force rebuild.
|
||||
self.hnsw_mark_dirty();
|
||||
self.bm25_cache = None;
|
||||
let needs_flush = self
|
||||
.wal
|
||||
.as_ref()
|
||||
@@ -616,7 +558,6 @@ impl AgentMemory for HDF5Memory {
|
||||
entry.tags,
|
||||
);
|
||||
self.hnsw_on_insert(idx);
|
||||
self.bm25_cache = None;
|
||||
let needs_flush = self
|
||||
.wal
|
||||
.as_ref()
|
||||
@@ -645,7 +586,6 @@ impl AgentMemory for HDF5Memory {
|
||||
}
|
||||
// Batch inserts rebuild the index once rather than node-by-node.
|
||||
self.hnsw_mark_dirty();
|
||||
self.bm25_cache = None;
|
||||
self.flush()?;
|
||||
Ok(indices)
|
||||
}
|
||||
@@ -657,7 +597,6 @@ impl AgentMemory for HDF5Memory {
|
||||
)));
|
||||
}
|
||||
self.hnsw_on_delete(id);
|
||||
self.bm25_cache = None;
|
||||
self.flush()?;
|
||||
|
||||
// Auto-compact if threshold exceeded
|
||||
@@ -675,7 +614,6 @@ impl AgentMemory for HDF5Memory {
|
||||
if removed > 0 {
|
||||
// Compaction renumbers cache indices; rebuild the index to match.
|
||||
self.hnsw_mark_dirty();
|
||||
self.bm25_cache = None;
|
||||
self.flush()?;
|
||||
}
|
||||
Ok(removed)
|
||||
@@ -1648,7 +1586,7 @@ impl HDF5Memory {
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
// Persistent tier.
|
||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.4, 0.6, k);
|
||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.7, 0.3, k);
|
||||
const EPHEMERAL_BOOST: f32 = 1.2;
|
||||
let mut results = persistent;
|
||||
|
||||
|
||||
@@ -133,62 +133,7 @@ impl MediaRef {
|
||||
checksum: Some(cs),
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate this reference against a sandbox directory and a URL scheme allowlist.
|
||||
///
|
||||
/// * `Path` references are canonicalized and checked to be within `sandbox`
|
||||
/// (if `sandbox` is `Some`). A path that escapes the sandbox via `..`
|
||||
/// or symlinks is rejected with an error.
|
||||
/// * `Url` references must begin with one of the schemes in
|
||||
/// [`ALLOWED_URL_SCHEMES`]. An empty or scheme-less URL is rejected.
|
||||
/// * `Inline` references are always valid (no external resolution).
|
||||
///
|
||||
/// Returns `Ok(())` when the reference passes all checks, or an `Err`
|
||||
/// with a human-readable reason otherwise.
|
||||
pub fn validate(&self, sandbox: Option<&std::path::Path>) -> Result<(), String> {
|
||||
match &self.ref_type {
|
||||
MediaRefType::Path(raw) => {
|
||||
let candidate = std::path::Path::new(raw);
|
||||
let canonical = candidate
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("path canonicalization failed for {raw:?}: {e}"))?;
|
||||
if let Some(root) = sandbox {
|
||||
let root_canonical = root
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("sandbox canonicalization failed: {e}"))?;
|
||||
if !canonical.starts_with(&root_canonical) {
|
||||
return Err(format!(
|
||||
"path {canonical:?} escapes sandbox {root_canonical:?}"
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
MediaRefType::Url(url) => {
|
||||
let scheme_end = url
|
||||
.find("://")
|
||||
.ok_or_else(|| format!("URL {url:?} has no scheme"))?;
|
||||
let scheme = &url[..scheme_end];
|
||||
if ALLOWED_URL_SCHEMES.contains(&scheme) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!(
|
||||
"URL scheme {scheme:?} is not in the allowlist {:?}",
|
||||
ALLOWED_URL_SCHEMES
|
||||
))
|
||||
}
|
||||
}
|
||||
MediaRefType::Inline(_) => Ok(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// URL schemes that are permitted in `MediaRef::Url` references.
|
||||
///
|
||||
/// Any scheme not in this list is rejected by [`MediaRef::validate`]. Keeping
|
||||
/// the list explicit prevents `file://` or `data:` URIs from being smuggled in
|
||||
/// via adversarial memory content.
|
||||
pub const ALLOWED_URL_SCHEMES: &[&str] = &["https", "http"];
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// FNV-1a helper (no external deps)
|
||||
@@ -862,69 +807,4 @@ mod tests {
|
||||
let r = store.get_record(id).unwrap();
|
||||
assert_eq!(r.metadata.get("source").unwrap(), "camera-1");
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// MediaRef::validate — sandboxing
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn inline_always_valid() {
|
||||
let r = MediaRef::inline(vec![1, 2, 3], "application/octet-stream");
|
||||
assert!(r.validate(None).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_allowed_scheme_https() {
|
||||
let r = MediaRef::url("https://example.com/img.png", "image/png");
|
||||
assert!(r.validate(None).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_allowed_scheme_http() {
|
||||
let r = MediaRef::url("http://example.com/img.png", "image/png");
|
||||
assert!(r.validate(None).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_disallowed_scheme_file() {
|
||||
let r = MediaRef::url("file:///etc/passwd", "text/plain");
|
||||
assert!(r.validate(None).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_disallowed_scheme_data() {
|
||||
let r = MediaRef::url("data:text/html,<script>", "text/html");
|
||||
assert!(r.validate(None).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn url_no_scheme_rejected() {
|
||||
let r = MediaRef::url("not-a-url", "text/plain");
|
||||
assert!(r.validate(None).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn path_within_sandbox_accepted() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file = dir.path().join("audio.mp3");
|
||||
std::fs::write(&file, b"dummy").unwrap();
|
||||
let r = MediaRef::path(file.to_str().unwrap(), "audio/mpeg");
|
||||
assert!(r.validate(Some(dir.path())).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn path_outside_sandbox_rejected() {
|
||||
let sandbox = tempfile::tempdir().unwrap();
|
||||
// /tmp itself exists and is outside the sandbox subdir
|
||||
let r = MediaRef::path("/tmp", "inode/directory");
|
||||
let result = r.validate(Some(sandbox.path()));
|
||||
// May fail at canonicalization or at the starts_with check; either is correct
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn path_nonexistent_rejected_at_canonicalize() {
|
||||
let r = MediaRef::path("/this/path/does/not/exist/abc123", "text/plain");
|
||||
assert!(r.validate(None).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -535,7 +535,7 @@ impl MemoryBackend for ClawhdfBackend {
|
||||
let candidates = k.saturating_mul(3).max(10);
|
||||
let raw = self
|
||||
.memory
|
||||
.hybrid_search(query_embedding, query_text, 0.4, 0.6, candidates);
|
||||
.hybrid_search(query_embedding, query_text, 0.7, 0.3, candidates);
|
||||
|
||||
if raw.is_empty() {
|
||||
return Vec::new();
|
||||
|
||||
@@ -90,16 +90,7 @@ impl HDF5Memory {
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
// Lazily build the BM25 index once and reuse across searches. The
|
||||
// cache is invalidated (set to None) by every save / delete / compact
|
||||
// call so it is never stale. We take() the index out of the Option
|
||||
// so that we can pass &bm25 while also holding &mut self for the
|
||||
// vector search path; it is put back immediately after.
|
||||
if self.bm25_cache.is_none() {
|
||||
self.bm25_cache =
|
||||
Some(bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones));
|
||||
}
|
||||
let bm25 = self.bm25_cache.take().expect("just built");
|
||||
let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
|
||||
let scored = self.vector_keyword_search(
|
||||
query_embedding,
|
||||
query_text,
|
||||
@@ -130,9 +121,6 @@ impl HDF5Memory {
|
||||
|
||||
let hit_indices: Vec<usize> = results.iter().map(|r| r.index).collect();
|
||||
self.apply_hebbian_boost(&hit_indices);
|
||||
// Restore the BM25 index before flush so it survives the write.
|
||||
// flush() does not invalidate bm25_cache; only mutating writes do.
|
||||
self.bm25_cache = Some(bm25);
|
||||
self.flush().ok();
|
||||
|
||||
results
|
||||
|
||||
@@ -1,284 +0,0 @@
|
||||
//! Ed25519 file signing for ClawBrainHub `.brain` files.
|
||||
//!
|
||||
//! # Sidecar format
|
||||
//!
|
||||
//! ```text
|
||||
//! [8 bytes magic "CLAWSIG\x00"]
|
||||
//! [4 bytes version = 1, little-endian u32]
|
||||
//! [1 byte public-key length = 32]
|
||||
//! [32 bytes Ed25519 public key (raw)]
|
||||
//! [1 byte signature length = 64]
|
||||
//! [64 bytes Ed25519 signature over the file's SHA-512 digest]
|
||||
//! ```
|
||||
//!
|
||||
//! The signature covers the **SHA-512 hash** of the file content rather than
|
||||
//! the raw bytes so that large files do not need to be fully loaded into memory
|
||||
//! during verification. Ring's Ed25519 implementation hashes internally, so
|
||||
//! we pass the entire content and let ring handle it.
|
||||
|
||||
use std::io::Read;
|
||||
use std::path::Path;
|
||||
|
||||
use ring::rand::SystemRandom;
|
||||
use ring::signature::{self, Ed25519KeyPair, KeyPair};
|
||||
|
||||
/// Sidecar file magic.
|
||||
const MAGIC: &[u8; 8] = b"CLAWSIG\x00";
|
||||
/// Sidecar format version.
|
||||
const VERSION: u32 = 1;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SigningError {
|
||||
/// Sidecar is too short, has wrong magic, or unsupported version.
|
||||
MalformedSidecar,
|
||||
/// Ed25519 signature did not verify against the file content.
|
||||
InvalidSignature,
|
||||
/// Key generation or signing operation failed.
|
||||
KeyError(String),
|
||||
/// I/O error reading/writing a file.
|
||||
Io(std::io::Error),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SigningError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
SigningError::MalformedSidecar => write!(f, "malformed signing sidecar"),
|
||||
SigningError::InvalidSignature => write!(f, "Ed25519 signature verification failed"),
|
||||
SigningError::KeyError(e) => write!(f, "key error: {e}"),
|
||||
SigningError::Io(e) => write!(f, "I/O error: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for SigningError {
|
||||
fn from(e: std::io::Error) -> Self {
|
||||
SigningError::Io(e)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Key generation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Generate a new Ed25519 key pair.
|
||||
///
|
||||
/// Returns `(pkcs8_document, public_key_bytes)`. The PKCS#8 document should
|
||||
/// be stored securely (it contains the private key). The public key is needed
|
||||
/// for verification and can be distributed freely.
|
||||
pub fn generate_keypair() -> Result<(Vec<u8>, Vec<u8>), SigningError> {
|
||||
let rng = SystemRandom::new();
|
||||
let pkcs8 = Ed25519KeyPair::generate_pkcs8(&rng)
|
||||
.map_err(|_| SigningError::KeyError("key generation failed".into()))?;
|
||||
let pair = Ed25519KeyPair::from_pkcs8(pkcs8.as_ref())
|
||||
.map_err(|_| SigningError::KeyError("pkcs8 decode failed".into()))?;
|
||||
let pubkey = pair.public_key().as_ref().to_vec();
|
||||
Ok((pkcs8.as_ref().to_vec(), pubkey))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sign / verify (in-memory)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Sign `data` with a PKCS#8-encoded Ed25519 private key.
|
||||
///
|
||||
/// Returns the raw 64-byte Ed25519 signature.
|
||||
pub fn sign(pkcs8_key: &[u8], data: &[u8]) -> Result<Vec<u8>, SigningError> {
|
||||
let pair = Ed25519KeyPair::from_pkcs8(pkcs8_key)
|
||||
.map_err(|_| SigningError::KeyError("invalid PKCS#8 key".into()))?;
|
||||
Ok(pair.sign(data).as_ref().to_vec())
|
||||
}
|
||||
|
||||
/// Verify that `signature` is a valid Ed25519 signature of `data` under
|
||||
/// `public_key` (raw 32-byte key).
|
||||
///
|
||||
/// Returns `true` when the signature is valid.
|
||||
pub fn verify(public_key: &[u8], data: &[u8], signature: &[u8]) -> bool {
|
||||
let peer = signature::UnparsedPublicKey::new(&signature::ED25519, public_key);
|
||||
peer.verify(data, signature).is_ok()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sidecar helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Serialize a public key and signature into a sidecar envelope.
|
||||
pub fn encode_sidecar(public_key: &[u8], sig: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(8 + 4 + 1 + public_key.len() + 1 + sig.len());
|
||||
out.extend_from_slice(MAGIC);
|
||||
out.extend_from_slice(&VERSION.to_le_bytes());
|
||||
out.push(public_key.len() as u8);
|
||||
out.extend_from_slice(public_key);
|
||||
out.push(sig.len() as u8);
|
||||
out.extend_from_slice(sig);
|
||||
out
|
||||
}
|
||||
|
||||
/// Parse a sidecar envelope, returning `(public_key, signature)`.
|
||||
pub fn decode_sidecar(sidecar: &[u8]) -> Result<(Vec<u8>, Vec<u8>), SigningError> {
|
||||
if sidecar.len() < 8 + 4 + 1 + 1 {
|
||||
return Err(SigningError::MalformedSidecar);
|
||||
}
|
||||
if &sidecar[..8] != MAGIC {
|
||||
return Err(SigningError::MalformedSidecar);
|
||||
}
|
||||
let ver = u32::from_le_bytes(sidecar[8..12].try_into().unwrap());
|
||||
if ver != VERSION {
|
||||
return Err(SigningError::MalformedSidecar);
|
||||
}
|
||||
let mut pos = 12usize;
|
||||
let pk_len = sidecar[pos] as usize;
|
||||
pos += 1;
|
||||
if pos + pk_len + 1 > sidecar.len() {
|
||||
return Err(SigningError::MalformedSidecar);
|
||||
}
|
||||
let public_key = sidecar[pos..pos + pk_len].to_vec();
|
||||
pos += pk_len;
|
||||
let sig_len = sidecar[pos] as usize;
|
||||
pos += 1;
|
||||
if pos + sig_len > sidecar.len() {
|
||||
return Err(SigningError::MalformedSidecar);
|
||||
}
|
||||
let signature = sidecar[pos..pos + sig_len].to_vec();
|
||||
Ok((public_key, signature))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// File-level helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Returns the path for the sidecar signature file next to `file_path`.
|
||||
///
|
||||
/// Example: `memory.brain` → `memory.brain.sig`
|
||||
pub fn sidecar_path(file_path: &Path) -> std::path::PathBuf {
|
||||
let mut s = file_path.as_os_str().to_owned();
|
||||
s.push(".sig");
|
||||
std::path::PathBuf::from(s)
|
||||
}
|
||||
|
||||
/// Sign `file_path` with `pkcs8_key` and write the sidecar (`.sig` file).
|
||||
pub fn sign_file(file_path: &Path, pkcs8_key: &[u8]) -> Result<(), SigningError> {
|
||||
let data = read_file(file_path)?;
|
||||
let pair = Ed25519KeyPair::from_pkcs8(pkcs8_key)
|
||||
.map_err(|_| SigningError::KeyError("invalid PKCS#8 key".into()))?;
|
||||
let pubkey = pair.public_key().as_ref().to_vec();
|
||||
let sig = pair.sign(&data).as_ref().to_vec();
|
||||
let sidecar = encode_sidecar(&pubkey, &sig);
|
||||
let sidecar_p = sidecar_path(file_path);
|
||||
std::fs::write(&sidecar_p, &sidecar)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Verify the signature sidecar for `file_path`.
|
||||
///
|
||||
/// Reads the `.sig` sidecar next to the file, parses it, and checks the
|
||||
/// signature against `file_path`'s current contents.
|
||||
///
|
||||
/// Returns `Ok(true)` if the signature is valid, `Ok(false)` if the sidecar
|
||||
/// does not exist (not yet signed), and `Err(_)` on parse or I/O failures.
|
||||
pub fn verify_file(file_path: &Path) -> Result<bool, SigningError> {
|
||||
let sidecar_p = sidecar_path(file_path);
|
||||
if !sidecar_p.exists() {
|
||||
return Ok(false);
|
||||
}
|
||||
let sidecar_bytes = read_file(&sidecar_p)?;
|
||||
let (public_key, sig) = decode_sidecar(&sidecar_bytes)?;
|
||||
let data = read_file(file_path)?;
|
||||
if verify(&public_key, &data, &sig) {
|
||||
Ok(true)
|
||||
} else {
|
||||
Err(SigningError::InvalidSignature)
|
||||
}
|
||||
}
|
||||
|
||||
fn read_file(path: &Path) -> Result<Vec<u8>, SigningError> {
|
||||
let mut f = std::fs::File::open(path)?;
|
||||
let mut buf = Vec::new();
|
||||
f.read_to_end(&mut buf)?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
#[test]
|
||||
fn generate_and_sign_verify() {
|
||||
let (pkcs8, pubkey) = generate_keypair().unwrap();
|
||||
let data = b"ClawBrainHub .brain file content";
|
||||
let sig = sign(&pkcs8, data).unwrap();
|
||||
assert_eq!(sig.len(), 64);
|
||||
assert!(verify(&pubkey, data, &sig));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrong_public_key_fails() {
|
||||
let (pkcs8, _) = generate_keypair().unwrap();
|
||||
let (_, other_pubkey) = generate_keypair().unwrap();
|
||||
let sig = sign(&pkcs8, b"data").unwrap();
|
||||
assert!(!verify(&other_pubkey, b"data", &sig));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tampered_data_fails() {
|
||||
let (pkcs8, pubkey) = generate_keypair().unwrap();
|
||||
let sig = sign(&pkcs8, b"original").unwrap();
|
||||
assert!(!verify(&pubkey, b"tampered", &sig));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sidecar_encode_decode_roundtrip() {
|
||||
let pubkey = vec![0xAAu8; 32];
|
||||
let sig = vec![0xBBu8; 64];
|
||||
let sidecar = encode_sidecar(&pubkey, &sig);
|
||||
let (pk2, sig2) = decode_sidecar(&sidecar).unwrap();
|
||||
assert_eq!(pk2, pubkey);
|
||||
assert_eq!(sig2, sig);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_sidecar_detected() {
|
||||
assert!(matches!(decode_sidecar(b"short"), Err(SigningError::MalformedSidecar)));
|
||||
let mut bad = vec![0u8; 20];
|
||||
assert!(matches!(decode_sidecar(&bad), Err(SigningError::MalformedSidecar)));
|
||||
bad[..8].copy_from_slice(MAGIC);
|
||||
bad[8..12].copy_from_slice(&99u32.to_le_bytes()); // wrong version
|
||||
assert!(matches!(decode_sidecar(&bad), Err(SigningError::MalformedSidecar)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sign_and_verify_file() {
|
||||
let (pkcs8, _) = generate_keypair().unwrap();
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
f.write_all(b"brain file content").unwrap();
|
||||
f.flush().unwrap();
|
||||
sign_file(f.path(), &pkcs8).unwrap();
|
||||
// sidecar should exist
|
||||
assert!(sidecar_path(f.path()).exists());
|
||||
// verification should succeed
|
||||
assert!(matches!(verify_file(f.path()), Ok(true)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_file_no_sidecar_returns_false() {
|
||||
let f = NamedTempFile::new().unwrap();
|
||||
assert!(matches!(verify_file(f.path()), Ok(false)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_file_detects_modified_content() {
|
||||
let (pkcs8, _) = generate_keypair().unwrap();
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
f.write_all(b"original content").unwrap();
|
||||
f.flush().unwrap();
|
||||
sign_file(f.path(), &pkcs8).unwrap();
|
||||
// Overwrite the file with different content
|
||||
std::fs::write(f.path(), b"tampered content").unwrap();
|
||||
assert!(matches!(verify_file(f.path()), Err(SigningError::InvalidSignature)));
|
||||
}
|
||||
}
|
||||
@@ -3,14 +3,13 @@
|
||||
//! Exposes `extern "C"` functions for use via JNI from Kotlin.
|
||||
//! Each HDF5Memory instance is managed via an opaque handle (pointer).
|
||||
//!
|
||||
//! Thread safety: each handle wraps `HDF5Memory` in a `Mutex`, so concurrent
|
||||
//! calls on the same handle are safe. Multiple handles are fully independent.
|
||||
//! Thread safety: the caller (Kotlin side) must synchronize access
|
||||
//! to a single handle. Multiple handles are independent.
|
||||
|
||||
use std::ffi::{CStr, CString};
|
||||
use std::os::raw::c_char;
|
||||
use std::path::PathBuf;
|
||||
use std::ptr;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
|
||||
@@ -18,12 +17,8 @@ use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
// Handle management
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Opaque handle to a mutex-protected HDF5Memory instance.
|
||||
///
|
||||
/// Stored on the heap so that the raw pointer (an integer from JNI's
|
||||
/// perspective) is stable across calls. The `Mutex` makes concurrent JNI
|
||||
/// calls on the same handle safe without requiring the caller to synchronize.
|
||||
type Handle = *mut Mutex<HDF5Memory>;
|
||||
/// Opaque handle to an HDF5Memory instance.
|
||||
type Handle = *mut HDF5Memory;
|
||||
|
||||
/// Create a new HDF5 memory file.
|
||||
///
|
||||
@@ -51,7 +46,7 @@ pub unsafe extern "C" fn edgehdf5_create(
|
||||
|
||||
let config = MemoryConfig::new(PathBuf::from(path), &agent_id, embedding_dim as usize);
|
||||
match HDF5Memory::create(config) {
|
||||
Ok(mem) => Box::into_raw(Box::new(Mutex::new(mem))),
|
||||
Ok(mem) => Box::into_raw(Box::new(mem)),
|
||||
Err(_) => ptr::null_mut(),
|
||||
}
|
||||
}
|
||||
@@ -72,7 +67,7 @@ pub unsafe extern "C" fn edgehdf5_open(path: *const c_char) -> Handle {
|
||||
};
|
||||
|
||||
match HDF5Memory::open(std::path::Path::new(&path)) {
|
||||
Ok(mem) => Box::into_raw(Box::new(Mutex::new(mem))),
|
||||
Ok(mem) => Box::into_raw(Box::new(mem)),
|
||||
Err(_) => ptr::null_mut(),
|
||||
}
|
||||
}
|
||||
@@ -87,7 +82,7 @@ pub unsafe extern "C" fn edgehdf5_open(path: *const c_char) -> Handle {
|
||||
pub unsafe extern "C" fn edgehdf5_close(handle: Handle) {
|
||||
if !handle.is_null() {
|
||||
// SAFETY: handle was created by Box::into_raw in edgehdf5_create; this is the final use.
|
||||
unsafe { drop(Box::<Mutex<HDF5Memory>>::from_raw(handle)) };
|
||||
unsafe { drop(Box::from_raw(handle)) };
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,15 +115,11 @@ pub unsafe extern "C" fn edgehdf5_save(
|
||||
session_id: *const c_char,
|
||||
tags: *const c_char,
|
||||
) -> i64 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
||||
let mem = match unsafe { handle.as_mut() } {
|
||||
Some(m) => m,
|
||||
None => return -1,
|
||||
};
|
||||
let mut mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return -1,
|
||||
};
|
||||
|
||||
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||
let chunk = match unsafe { cstr_to_string(chunk) } {
|
||||
@@ -185,7 +176,7 @@ pub unsafe extern "C" fn edgehdf5_save(
|
||||
pub unsafe extern "C" fn edgehdf5_count_active(handle: Handle) -> u64 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
match unsafe { handle.as_ref() } {
|
||||
Some(mtx) => mtx.lock().map(|g| g.count_active() as u64).unwrap_or(0),
|
||||
Some(mem) => mem.count_active() as u64,
|
||||
None => 0,
|
||||
}
|
||||
}
|
||||
@@ -199,7 +190,7 @@ pub unsafe extern "C" fn edgehdf5_count_active(handle: Handle) -> u64 {
|
||||
pub unsafe extern "C" fn edgehdf5_count(handle: Handle) -> u64 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
match unsafe { handle.as_ref() } {
|
||||
Some(mtx) => mtx.lock().map(|g| g.count() as u64).unwrap_or(0),
|
||||
Some(mem) => mem.count() as u64,
|
||||
None => 0,
|
||||
}
|
||||
}
|
||||
@@ -211,15 +202,11 @@ pub unsafe extern "C" fn edgehdf5_count(handle: Handle) -> u64 {
|
||||
/// `handle` must be a valid, non-null handle.
|
||||
#[unsafe(no_mangle)]
|
||||
pub unsafe extern "C" fn edgehdf5_delete(handle: Handle, index: u64) -> i32 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
||||
let mem = match unsafe { handle.as_mut() } {
|
||||
Some(m) => m,
|
||||
None => return -1,
|
||||
};
|
||||
let mut mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return -1,
|
||||
};
|
||||
|
||||
match mem.delete(index as usize) {
|
||||
Ok(()) => 0,
|
||||
@@ -263,15 +250,11 @@ pub unsafe extern "C" fn edgehdf5_hybrid_search(
|
||||
out_scores: *mut f32,
|
||||
out_chunks: *mut *mut c_char,
|
||||
) -> u32 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
||||
let mem = match unsafe { handle.as_mut() } {
|
||||
Some(m) => m,
|
||||
None => return 0,
|
||||
};
|
||||
let mut mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return 0,
|
||||
};
|
||||
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||
let query_text = match unsafe { cstr_to_string(query_text) } {
|
||||
Some(s) => s,
|
||||
@@ -346,15 +329,11 @@ pub unsafe extern "C" fn edgehdf5_add_session(
|
||||
channel: *const c_char,
|
||||
summary: *const c_char,
|
||||
) -> i32 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
||||
let mem = match unsafe { handle.as_mut() } {
|
||||
Some(m) => m,
|
||||
None => return -1,
|
||||
};
|
||||
let mut mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return -1,
|
||||
};
|
||||
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||
let id = match unsafe { cstr_to_string(id) } {
|
||||
Some(s) => s,
|
||||
@@ -396,14 +375,10 @@ pub unsafe extern "C" fn edgehdf5_get_session_summary(
|
||||
session_id: *const c_char,
|
||||
) -> *mut c_char {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
let mem = match unsafe { handle.as_ref() } {
|
||||
Some(m) => m,
|
||||
None => return ptr::null_mut(),
|
||||
};
|
||||
let mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return ptr::null_mut(),
|
||||
};
|
||||
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||
let session_id = match unsafe { cstr_to_string(session_id) } {
|
||||
Some(s) => s,
|
||||
@@ -436,15 +411,11 @@ pub unsafe extern "C" fn edgehdf5_add_entity(
|
||||
entity_type: *const c_char,
|
||||
embedding_idx: i64,
|
||||
) -> i64 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
||||
let mem = match unsafe { handle.as_mut() } {
|
||||
Some(m) => m,
|
||||
None => return -1,
|
||||
};
|
||||
let mut mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return -1,
|
||||
};
|
||||
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||
let name = match unsafe { cstr_to_string(name) } {
|
||||
Some(s) => s,
|
||||
@@ -476,15 +447,11 @@ pub unsafe extern "C" fn edgehdf5_add_relation(
|
||||
relation: *const c_char,
|
||||
weight: f32,
|
||||
) -> i32 {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||
let mtx = match unsafe { handle.as_ref() } {
|
||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
||||
let mem = match unsafe { handle.as_mut() } {
|
||||
Some(m) => m,
|
||||
None => return -1,
|
||||
};
|
||||
let mut mem = match mtx.lock() {
|
||||
Ok(g) => g,
|
||||
Err(_) => return -1,
|
||||
};
|
||||
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||
let relation = match unsafe { cstr_to_string(relation) } {
|
||||
Some(s) => s,
|
||||
@@ -525,8 +492,7 @@ mod tests {
|
||||
fn open_handle(dir: &tempfile::TempDir) -> Handle {
|
||||
let path = CString::new(dir.path().join("mem.h5").to_str().unwrap()).unwrap();
|
||||
let agent_id = CString::new("test-agent").unwrap();
|
||||
// SAFETY: both C strings are valid and null-terminated; returned handle
|
||||
// wraps HDF5Memory in a Mutex and is safe to use from multiple threads.
|
||||
// SAFETY: both C strings are valid and null-terminated.
|
||||
unsafe { edgehdf5_create(path.as_ptr(), agent_id.as_ptr(), EMBEDDING_DIM) }
|
||||
}
|
||||
|
||||
@@ -624,44 +590,4 @@ mod tests {
|
||||
|
||||
unsafe { edgehdf5_close(handle) };
|
||||
}
|
||||
|
||||
/// Verify that concurrent calls on the same handle do not cause data races.
|
||||
///
|
||||
/// Each thread calls `edgehdf5_count_active` on the shared handle. With the
|
||||
/// `Mutex` wrapper in place this must complete without a panic or SIGABRT.
|
||||
/// Without the mutex it would be UB.
|
||||
#[test]
|
||||
fn concurrent_count_active_is_safe() {
|
||||
use std::sync::Arc;
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let handle = open_handle(&dir);
|
||||
assert!(!handle.is_null());
|
||||
|
||||
// Share the raw pointer across threads via a copy-friendly wrapper.
|
||||
// SAFETY: the Mutex inside the handle makes concurrent access sound.
|
||||
#[derive(Clone, Copy)]
|
||||
struct SendableHandle(Handle);
|
||||
unsafe impl Send for SendableHandle {}
|
||||
// SAFETY: the Mutex inside the handle serialises all access,
|
||||
// so sharing the wrapper across threads is sound.
|
||||
unsafe impl Sync for SendableHandle {}
|
||||
|
||||
let shared = Arc::new(SendableHandle(handle));
|
||||
let threads: Vec<_> = (0..8)
|
||||
.map(|_| {
|
||||
let h = Arc::clone(&shared);
|
||||
std::thread::spawn(move || {
|
||||
// SAFETY: handle is valid (not yet closed); Mutex guards access.
|
||||
let count = unsafe { edgehdf5_count_active(h.0) };
|
||||
assert_eq!(count, 0);
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
for t in threads {
|
||||
t.join().expect("thread panicked");
|
||||
}
|
||||
|
||||
unsafe { edgehdf5_close(handle) };
|
||||
}
|
||||
}
|
||||
|
||||
@@ -739,190 +739,12 @@ impl HnswIndex {
|
||||
pub fn m_max0(&self) -> usize {
|
||||
self.m_max0
|
||||
}
|
||||
|
||||
/// Insert a batch of vectors efficiently.
|
||||
///
|
||||
/// With the `parallel` feature enabled, neighbor searches for each new
|
||||
/// vector are executed concurrently against the graph state *before* the
|
||||
/// batch is applied, then edges are wired serially. This trades a small
|
||||
/// reduction in intra-batch connectivity for significant wall-clock
|
||||
/// speedup on large batches.
|
||||
///
|
||||
/// Without the `parallel` feature, this is equivalent to calling
|
||||
/// [`HnswIndex::insert`] for each vector in order.
|
||||
///
|
||||
/// Returns the assigned IDs in insertion order.
|
||||
pub fn batch_insert(&mut self, vectors: Vec<Vec<f32>>) -> Vec<usize> {
|
||||
if vectors.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Empty index: fall through to serial insert so the entry-point
|
||||
// seeding logic in `insert` runs correctly.
|
||||
if self.vectors.is_empty() {
|
||||
return vectors
|
||||
.into_iter()
|
||||
.map(|v| self.insert(v))
|
||||
.collect();
|
||||
}
|
||||
|
||||
let dim = self.vectors[0].len();
|
||||
for v in &vectors {
|
||||
assert_eq!(v.len(), dim, "batch_insert dimension mismatch");
|
||||
}
|
||||
|
||||
let base_id = self.vectors.len();
|
||||
let n = vectors.len();
|
||||
|
||||
// Pre-assign levels to all incoming vectors.
|
||||
let node_levels: Vec<usize> = (0..n)
|
||||
.map(|i| assign_level(base_id + i, self.m))
|
||||
.collect();
|
||||
|
||||
// Phase 1 — neighbor search (read-only on the current graph state).
|
||||
// Returns, for each new vector, the list of (layer, selected_neighbors)
|
||||
// pairs that will become its initial edge set.
|
||||
let per_vector_neighbors: Vec<Vec<(usize, Vec<usize>)>> =
|
||||
self.find_neighbors_batch(&vectors, &node_levels);
|
||||
|
||||
// Phase 2 — extend the vector store (serial).
|
||||
self.vectors.extend(vectors);
|
||||
self.deleted.extend(std::iter::repeat(false).take(n));
|
||||
self.node_levels.extend_from_slice(&node_levels);
|
||||
|
||||
// Grow existing layers to accommodate the new node slots.
|
||||
for layer in self.graph.iter_mut() {
|
||||
layer.resize(self.vectors.len(), Vec::new());
|
||||
}
|
||||
// Add any brand-new top layers introduced by this batch.
|
||||
let new_max_level = node_levels.iter().copied().max().unwrap_or(0);
|
||||
while self.graph.len() <= new_max_level {
|
||||
self.graph.push(vec![Vec::new(); self.vectors.len()]);
|
||||
}
|
||||
|
||||
// Phase 3 — wire edges and track entry-point promotions (serial).
|
||||
for (batch_idx, layer_neighbors) in per_vector_neighbors.into_iter().enumerate() {
|
||||
let id = base_id + batch_idx;
|
||||
for (layer, selected) in layer_neighbors {
|
||||
let max_conn = if layer == 0 { self.m_max0 } else { self.m };
|
||||
self.graph[layer][id] = selected.clone();
|
||||
for &nb in &selected {
|
||||
self.graph[layer][nb].push(id);
|
||||
if self.graph[layer][nb].len() > max_conn {
|
||||
prune_connections(
|
||||
&self.vectors,
|
||||
&mut self.graph[layer][nb],
|
||||
nb,
|
||||
max_conn,
|
||||
self.metric,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Promote entry point if this node sits on a taller layer.
|
||||
let ep_level = self.node_levels[self.entry_point];
|
||||
if node_levels[batch_idx] > ep_level {
|
||||
self.entry_point = id;
|
||||
}
|
||||
}
|
||||
|
||||
(base_id..base_id + n).collect()
|
||||
}
|
||||
|
||||
/// Search for neighbors of each vector in `vectors` against the current
|
||||
/// (read-only) graph. Returns per-vector `(layer_id, neighbor_ids)` pairs.
|
||||
fn find_neighbors_batch(
|
||||
&self,
|
||||
vectors: &[Vec<f32>],
|
||||
node_levels: &[usize],
|
||||
) -> Vec<Vec<(usize, Vec<usize>)>> {
|
||||
let ep_level = self.node_levels[self.entry_point];
|
||||
let entry_point = self.entry_point;
|
||||
|
||||
#[cfg(feature = "parallel")]
|
||||
{
|
||||
use rayon::prelude::*;
|
||||
let existing = &self.vectors;
|
||||
let graph = &self.graph;
|
||||
let metric = self.metric;
|
||||
let m = self.m;
|
||||
let m_max0 = self.m_max0;
|
||||
let ef = self.ef_construction;
|
||||
vectors
|
||||
.par_iter()
|
||||
.zip(node_levels.par_iter())
|
||||
.map(|(v, &nl)| {
|
||||
find_neighbors_for(
|
||||
existing, graph, v, nl, ep_level, entry_point, m, m_max0, ef, metric,
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
#[cfg(not(feature = "parallel"))]
|
||||
{
|
||||
vectors
|
||||
.iter()
|
||||
.zip(node_levels.iter())
|
||||
.map(|(v, &nl)| {
|
||||
find_neighbors_for(
|
||||
&self.vectors,
|
||||
&self.graph,
|
||||
v,
|
||||
nl,
|
||||
ep_level,
|
||||
entry_point,
|
||||
self.m,
|
||||
self.m_max0,
|
||||
self.ef_construction,
|
||||
self.metric,
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Internal HNSW algorithms
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Compute the set of neighbor edges for `new_vec` against a read-only snapshot
|
||||
/// of the existing graph. Used by [`HnswIndex::batch_insert`].
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn find_neighbors_for(
|
||||
existing: &[Vec<f32>],
|
||||
graph: &[Vec<Vec<usize>>],
|
||||
new_vec: &[f32],
|
||||
node_level: usize,
|
||||
ep_level: usize,
|
||||
entry_point: usize,
|
||||
m: usize,
|
||||
m_max0: usize,
|
||||
ef: usize,
|
||||
metric: DistanceMetric,
|
||||
) -> Vec<(usize, Vec<usize>)> {
|
||||
let mut ep = entry_point;
|
||||
|
||||
// Phase 1: greedy descent from the top layer down to node_level + 1.
|
||||
for layer in (node_level + 1..=ep_level).rev() {
|
||||
ep = greedy_closest(existing, &graph[layer], new_vec, ep, metric);
|
||||
}
|
||||
|
||||
// Phase 2: beam search at each layer, collecting selected neighbors.
|
||||
let bottom = node_level.min(ep_level);
|
||||
let mut result = Vec::with_capacity(bottom + 1);
|
||||
for layer in (0..=bottom).rev() {
|
||||
let max_conn = if layer == 0 { m_max0 } else { m };
|
||||
let candidates = search_layer(existing, &graph[layer], new_vec, ep, ef, metric);
|
||||
let selected: Vec<usize> = candidates.iter().take(max_conn).map(|c| c.id).collect();
|
||||
if !selected.is_empty() {
|
||||
ep = selected[0];
|
||||
}
|
||||
result.push((layer, selected));
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Greedy search: find the single closest node to `query` starting from `ep`.
|
||||
fn greedy_closest(
|
||||
vectors: &[Vec<f32>],
|
||||
@@ -1630,75 +1452,4 @@ mod tests {
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(results[0].0, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_insert_ids_are_sequential() {
|
||||
let vectors = make_random_vectors(20, 8, 42);
|
||||
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
|
||||
let ids = index.batch_insert(vectors.clone());
|
||||
assert_eq!(ids, (0..20).collect::<Vec<_>>());
|
||||
assert_eq!(index.len(), 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_insert_into_existing_index() {
|
||||
let first = make_random_vectors(10, 8, 11);
|
||||
let second = make_random_vectors(10, 8, 22);
|
||||
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
|
||||
let ids1 = index.batch_insert(first);
|
||||
assert_eq!(ids1, (0..10).collect::<Vec<_>>());
|
||||
let ids2 = index.batch_insert(second.clone());
|
||||
assert_eq!(ids2, (10..20).collect::<Vec<_>>());
|
||||
assert_eq!(index.len(), 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_insert_search_quality() {
|
||||
// Build index from 50 vectors using serial insert, then build the same
|
||||
// index using batch_insert. The search results should be identical for
|
||||
// the first 50 vectors (which are fully connected in both cases).
|
||||
let vectors = make_random_vectors(50, 16, 99);
|
||||
let mut serial = HnswIndex::new(8, 32, DistanceMetric::Cosine);
|
||||
for v in &vectors {
|
||||
serial.insert(v.clone());
|
||||
}
|
||||
let mut batch = HnswIndex::new(8, 32, DistanceMetric::Cosine);
|
||||
batch.batch_insert(vectors.clone());
|
||||
assert_eq!(batch.len(), serial.len());
|
||||
|
||||
// Both indexes should find the same nearest neighbor for each query.
|
||||
let queries = make_random_vectors(5, 16, 777);
|
||||
for q in &queries {
|
||||
let s = serial.search(q, 1, 32);
|
||||
let b = batch.search(q, 1, 32);
|
||||
assert!(!s.is_empty() && !b.is_empty());
|
||||
// Result must be in the top-3 of the serial index — batch
|
||||
// is slightly less connected due to the read-snapshot approach.
|
||||
let top3_serial: Vec<usize> = serial.search(q, 3, 32).into_iter().map(|(id, _)| id).collect();
|
||||
assert!(top3_serial.contains(&b[0].0), "batch top-1 not in serial top-3");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_insert_empty_is_noop() {
|
||||
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
|
||||
let ids = index.batch_insert(vec![]);
|
||||
assert!(ids.is_empty());
|
||||
assert!(index.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_insert_saves_and_loads() {
|
||||
let vectors = make_random_vectors(30, 6, 55);
|
||||
let mut index = HnswIndex::new(8, 32, DistanceMetric::L2);
|
||||
index.batch_insert(vectors.clone());
|
||||
let bytes = index.to_hdf5_bytes().unwrap();
|
||||
let loaded = HnswIndex::load_from_hdf5(&bytes).unwrap();
|
||||
assert_eq!(loaded.len(), 30);
|
||||
assert_eq!(loaded.metric(), DistanceMetric::L2);
|
||||
// The query's own vector should be the nearest neighbor.
|
||||
let q = &vectors[0];
|
||||
let results = loaded.search(q, 1, 32);
|
||||
assert_eq!(results[0].0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -32,7 +32,7 @@ name = "bench"
|
||||
harness = false
|
||||
|
||||
[features]
|
||||
default = ["std", "checksum", "deflate", "provenance", "system-zlib-decompress"]
|
||||
default = ["std", "checksum", "deflate", "provenance", "fast-deflate", "system-zlib-decompress"]
|
||||
std = []
|
||||
checksum = []
|
||||
deflate = ["flate2"]
|
||||
|
||||
@@ -143,6 +143,10 @@ pub fn parse_vds_mappings(
|
||||
let source_selection = read_selection(heap_data, &mut pos)?;
|
||||
let virtual_selection = read_selection(heap_data, &mut pos)?;
|
||||
|
||||
// Validate external file name to prevent directory traversal attacks
|
||||
// (Dataset paths within files can use absolute HDF5 paths like "/data")
|
||||
validate_vds_file_name(&source_file)?;
|
||||
|
||||
mappings.push(VdsMapping {
|
||||
source_file,
|
||||
source_dataset,
|
||||
@@ -154,6 +158,37 @@ pub fn parse_vds_mappings(
|
||||
Ok(mappings)
|
||||
}
|
||||
|
||||
/// Validate external file names to prevent directory traversal.
|
||||
/// Dataset paths within files can use absolute HDF5 paths (starting with /),
|
||||
/// but external file names must not escape the file tree via .. or absolute paths.
|
||||
fn validate_vds_file_name(filename: &str) -> Result<(), FormatError> {
|
||||
if filename.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// "." means same file - always OK
|
||||
if filename == "." {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Filesystem paths cannot start with / (absolute filesystem path)
|
||||
if filename.starts_with('/') {
|
||||
return Err(FormatError::FilterError(
|
||||
"VDS file name cannot be an absolute filesystem path".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Reject directory traversal (..)
|
||||
if filename.contains("..") {
|
||||
return Err(FormatError::FilterError(
|
||||
"VDS file name contains illegal traversal sequence (..)".into(),
|
||||
));
|
||||
}
|
||||
|
||||
// Relative filesystem paths are OK
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read a null-terminated UTF-8 string from data starting at `pos`.
|
||||
fn read_null_terminated_string(data: &[u8], pos: &mut usize) -> Result<String, FormatError> {
|
||||
let start = *pos;
|
||||
@@ -862,4 +897,68 @@ mod tests {
|
||||
let blob = [0x01u8, 0, 0, 0, 0, 0, 0, 0, 0];
|
||||
assert!(parse_vds_mappings(&blob, 8).unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_vds_mappings_rejects_path_traversal() {
|
||||
// INT-06: Verify that VDS file names containing ".." are rejected
|
||||
let blob = [
|
||||
0x00u8, // version 0 (with explicit file name)
|
||||
0x01, 0, 0, 0, 0, 0, 0, 0, // nused = 1
|
||||
0x2e, 0x2e, 0x2f, 0x65, 0x74, 0x63, 0x2f, 0x70, 0x61, 0x73, 0x73, 0x77, 0x64, 0x00, // "../etc/passwd | ||||