Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
87039e926c | ||
|
|
fdd8901c37 | ||
|
|
e7e83acf35 | ||
|
|
ca8a3a4a2e | ||
|
|
fdc4572ab7 | ||
|
|
4aee2fa610 | ||
|
|
5ca0b8092e | ||
|
|
4b17bf9101 | ||
|
|
ec6bc80007 | ||
|
|
14db35aa74 |
@@ -22,5 +22,36 @@ jobs:
|
|||||||
run: rustup component add rustfmt clippy
|
run: rustup component add rustfmt clippy
|
||||||
- name: Install thumbv7em-none-eabihf target
|
- name: Install thumbv7em-none-eabihf target
|
||||||
run: rustup target add thumbv7em-none-eabihf
|
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
|
- name: Run CI script
|
||||||
run: bash scripts/ci-test.sh
|
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,3 +31,9 @@ tempfile = "3"
|
|||||||
criterion = { version = "0.5", features = ["html_reports"] }
|
criterion = { version = "0.5", features = ["html_reports"] }
|
||||||
half = "2.7"
|
half = "2.7"
|
||||||
serde = { version = "1", features = ["derive"] }
|
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,6 +23,7 @@ rayon = { version = "1", optional = true }
|
|||||||
matrixmultiply = { version = "0.3", optional = true }
|
matrixmultiply = { version = "0.3", optional = true }
|
||||||
cblas-sys = { version = "0.1", optional = true }
|
cblas-sys = { version = "0.1", optional = true }
|
||||||
tokio = { version = "1", features = ["rt", "sync", "macros", "time"], optional = true }
|
tokio = { version = "1", features = ["rt", "sync", "macros", "time"], optional = true }
|
||||||
|
ring = { version = "0.17", optional = true }
|
||||||
|
|
||||||
[target.'cfg(target_os = "macos")'.dependencies]
|
[target.'cfg(target_os = "macos")'.dependencies]
|
||||||
accelerate-src = { version = "0.3", optional = true }
|
accelerate-src = { version = "0.3", optional = true }
|
||||||
@@ -60,3 +61,5 @@ fast-math = ["matrixmultiply"]
|
|||||||
accelerate = ["accelerate-src", "cblas-sys"]
|
accelerate = ["accelerate-src", "cblas-sys"]
|
||||||
openblas = ["openblas-src", "cblas-sys"]
|
openblas = ["openblas-src", "cblas-sys"]
|
||||||
async = ["tokio"]
|
async = ["tokio"]
|
||||||
|
encryption = ["ring"]
|
||||||
|
signing = ["ring"]
|
||||||
|
|||||||
@@ -0,0 +1,23 @@
|
|||||||
|
[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
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
#![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,6 +262,176 @@ 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
|
// Tests
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -460,4 +630,72 @@ mod tests {
|
|||||||
assert_eq!(det.session_count("sess-b"), 1);
|
assert_eq!(det.session_count("sess-b"), 1);
|
||||||
assert_eq!(det.session_count("unknown"), 0);
|
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?;
|
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||||
//! mem.save(entry).await?; // buffered → background writer
|
//! mem.save(entry).await?; // buffered → background writer
|
||||||
//! mem.save_batch(entries).await?; // also buffered
|
//! mem.save_batch(entries).await?; // also buffered
|
||||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
//! let results = mem.hybrid_search(emb, "query".into(), 0.4, 0.6, 5).await;
|
||||||
//! mem.shutdown().await?; // final flush + stop
|
//! mem.shutdown().await?; // final flush + stop
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
|
|||||||
@@ -218,6 +218,171 @@ 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,
|
/// Tokenize a string: lowercase, split on non-alphanumeric characters,
|
||||||
/// filter empty tokens.
|
/// filter empty tokens.
|
||||||
fn tokenize(text: &str) -> Vec<String> {
|
fn tokenize(text: &str) -> Vec<String> {
|
||||||
@@ -451,4 +616,74 @@ 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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,268 @@
|
|||||||
|
//! 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,6 +439,11 @@ impl KnowledgeCache {
|
|||||||
min_activation: f32,
|
min_activation: f32,
|
||||||
max_steps: usize,
|
max_steps: usize,
|
||||||
) -> Vec<(u64, f32)> {
|
) -> 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();
|
let mut activation: HashMap<u64, f32> = HashMap::new();
|
||||||
|
|
||||||
// Initialise seeds with activation 1.0.
|
// Initialise seeds with activation 1.0.
|
||||||
@@ -1162,4 +1167,63 @@ mod tests {
|
|||||||
assert!(ctx.contains("occupation"));
|
assert!(ctx.contains("occupation"));
|
||||||
assert!(ctx.contains("engineer"));
|
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,6 +20,10 @@ pub mod vector_search;
|
|||||||
|
|
||||||
pub mod agents_md;
|
pub mod agents_md;
|
||||||
pub mod anomaly;
|
pub mod anomaly;
|
||||||
|
#[cfg(feature = "encryption")]
|
||||||
|
pub mod encryption;
|
||||||
|
#[cfg(feature = "signing")]
|
||||||
|
pub mod signing;
|
||||||
pub mod cache;
|
pub mod cache;
|
||||||
pub mod confidence;
|
pub mod confidence;
|
||||||
pub mod consolidation;
|
pub mod consolidation;
|
||||||
@@ -60,6 +64,17 @@ pub fn cosine_similarity_prenorm(
|
|||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
use cache::MemoryCache;
|
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")]
|
#[cfg(feature = "hnsw")]
|
||||||
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
use clawhdf5_ann::{DistanceMetric, HnswIndex};
|
||||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||||
@@ -227,6 +242,10 @@ pub struct HDF5Memory {
|
|||||||
/// search.
|
/// search.
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
hnsw_synced_len: usize,
|
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 {
|
impl std::fmt::Debug for HDF5Memory {
|
||||||
@@ -266,6 +285,7 @@ impl HDF5Memory {
|
|||||||
hnsw_dirty: false,
|
hnsw_dirty: false,
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
hnsw_synced_len: 0,
|
hnsw_synced_len: 0,
|
||||||
|
bm25_cache: None,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -285,6 +305,16 @@ impl HDF5Memory {
|
|||||||
None
|
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 {
|
Ok(Self {
|
||||||
config,
|
config,
|
||||||
cache,
|
cache,
|
||||||
@@ -301,6 +331,7 @@ impl HDF5Memory {
|
|||||||
hnsw_dirty: true,
|
hnsw_dirty: true,
|
||||||
#[cfg(feature = "hnsw")]
|
#[cfg(feature = "hnsw")]
|
||||||
hnsw_synced_len: 0,
|
hnsw_synced_len: 0,
|
||||||
|
bm25_cache,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -320,9 +351,35 @@ impl HDF5Memory {
|
|||||||
if let Some(ref mut w) = self.wal {
|
if let Some(ref mut w) = self.wal {
|
||||||
w.truncate()?;
|
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(())
|
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 --------------------------------------------
|
// ---- HNSW index maintenance --------------------------------------------
|
||||||
//
|
//
|
||||||
// The index mirrors the cache: HNSW node id == cache index, kept aligned by
|
// The index mirrors the cache: HNSW node id == cache index, kept aligned by
|
||||||
@@ -517,6 +574,7 @@ impl HDF5Memory {
|
|||||||
);
|
);
|
||||||
// In-place embedding change: the index node is stale, force rebuild.
|
// In-place embedding change: the index node is stale, force rebuild.
|
||||||
self.hnsw_mark_dirty();
|
self.hnsw_mark_dirty();
|
||||||
|
self.bm25_cache = None;
|
||||||
let needs_flush = self
|
let needs_flush = self
|
||||||
.wal
|
.wal
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -558,6 +616,7 @@ impl AgentMemory for HDF5Memory {
|
|||||||
entry.tags,
|
entry.tags,
|
||||||
);
|
);
|
||||||
self.hnsw_on_insert(idx);
|
self.hnsw_on_insert(idx);
|
||||||
|
self.bm25_cache = None;
|
||||||
let needs_flush = self
|
let needs_flush = self
|
||||||
.wal
|
.wal
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -586,6 +645,7 @@ impl AgentMemory for HDF5Memory {
|
|||||||
}
|
}
|
||||||
// Batch inserts rebuild the index once rather than node-by-node.
|
// Batch inserts rebuild the index once rather than node-by-node.
|
||||||
self.hnsw_mark_dirty();
|
self.hnsw_mark_dirty();
|
||||||
|
self.bm25_cache = None;
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
Ok(indices)
|
Ok(indices)
|
||||||
}
|
}
|
||||||
@@ -597,6 +657,7 @@ impl AgentMemory for HDF5Memory {
|
|||||||
)));
|
)));
|
||||||
}
|
}
|
||||||
self.hnsw_on_delete(id);
|
self.hnsw_on_delete(id);
|
||||||
|
self.bm25_cache = None;
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
|
|
||||||
// Auto-compact if threshold exceeded
|
// Auto-compact if threshold exceeded
|
||||||
@@ -614,6 +675,7 @@ impl AgentMemory for HDF5Memory {
|
|||||||
if removed > 0 {
|
if removed > 0 {
|
||||||
// Compaction renumbers cache indices; rebuild the index to match.
|
// Compaction renumbers cache indices; rebuild the index to match.
|
||||||
self.hnsw_mark_dirty();
|
self.hnsw_mark_dirty();
|
||||||
|
self.bm25_cache = None;
|
||||||
self.flush()?;
|
self.flush()?;
|
||||||
}
|
}
|
||||||
Ok(removed)
|
Ok(removed)
|
||||||
@@ -1586,7 +1648,7 @@ impl HDF5Memory {
|
|||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
// Persistent tier.
|
// Persistent tier.
|
||||||
let persistent = self.hybrid_search(query_embedding, query_text, 0.7, 0.3, k);
|
let persistent = self.hybrid_search(query_embedding, query_text, 0.4, 0.6, k);
|
||||||
const EPHEMERAL_BOOST: f32 = 1.2;
|
const EPHEMERAL_BOOST: f32 = 1.2;
|
||||||
let mut results = persistent;
|
let mut results = persistent;
|
||||||
|
|
||||||
|
|||||||
@@ -133,8 +133,63 @@ impl MediaRef {
|
|||||||
checksum: Some(cs),
|
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)
|
// FNV-1a helper (no external deps)
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -807,4 +862,69 @@ mod tests {
|
|||||||
let r = store.get_record(id).unwrap();
|
let r = store.get_record(id).unwrap();
|
||||||
assert_eq!(r.metadata.get("source").unwrap(), "camera-1");
|
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 candidates = k.saturating_mul(3).max(10);
|
||||||
let raw = self
|
let raw = self
|
||||||
.memory
|
.memory
|
||||||
.hybrid_search(query_embedding, query_text, 0.7, 0.3, candidates);
|
.hybrid_search(query_embedding, query_text, 0.4, 0.6, candidates);
|
||||||
|
|
||||||
if raw.is_empty() {
|
if raw.is_empty() {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
|
|||||||
@@ -90,7 +90,16 @@ impl HDF5Memory {
|
|||||||
keyword_weight: f32,
|
keyword_weight: f32,
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> Vec<SearchResult> {
|
||||||
let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
|
// 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 scored = self.vector_keyword_search(
|
let scored = self.vector_keyword_search(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
@@ -121,6 +130,9 @@ impl HDF5Memory {
|
|||||||
|
|
||||||
let hit_indices: Vec<usize> = results.iter().map(|r| r.index).collect();
|
let hit_indices: Vec<usize> = results.iter().map(|r| r.index).collect();
|
||||||
self.apply_hebbian_boost(&hit_indices);
|
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();
|
self.flush().ok();
|
||||||
|
|
||||||
results
|
results
|
||||||
|
|||||||
@@ -0,0 +1,284 @@
|
|||||||
|
//! 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,13 +3,14 @@
|
|||||||
//! Exposes `extern "C"` functions for use via JNI from Kotlin.
|
//! Exposes `extern "C"` functions for use via JNI from Kotlin.
|
||||||
//! Each HDF5Memory instance is managed via an opaque handle (pointer).
|
//! Each HDF5Memory instance is managed via an opaque handle (pointer).
|
||||||
//!
|
//!
|
||||||
//! Thread safety: the caller (Kotlin side) must synchronize access
|
//! Thread safety: each handle wraps `HDF5Memory` in a `Mutex`, so concurrent
|
||||||
//! to a single handle. Multiple handles are independent.
|
//! calls on the same handle are safe. Multiple handles are fully independent.
|
||||||
|
|
||||||
use std::ffi::{CStr, CString};
|
use std::ffi::{CStr, CString};
|
||||||
use std::os::raw::c_char;
|
use std::os::raw::c_char;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
use std::ptr;
|
use std::ptr;
|
||||||
|
use std::sync::Mutex;
|
||||||
|
|
||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
|
|
||||||
@@ -17,8 +18,12 @@ use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
|||||||
// Handle management
|
// Handle management
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
/// Opaque handle to an HDF5Memory instance.
|
/// Opaque handle to a mutex-protected HDF5Memory instance.
|
||||||
type Handle = *mut HDF5Memory;
|
///
|
||||||
|
/// 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>;
|
||||||
|
|
||||||
/// Create a new HDF5 memory file.
|
/// Create a new HDF5 memory file.
|
||||||
///
|
///
|
||||||
@@ -46,7 +51,7 @@ pub unsafe extern "C" fn edgehdf5_create(
|
|||||||
|
|
||||||
let config = MemoryConfig::new(PathBuf::from(path), &agent_id, embedding_dim as usize);
|
let config = MemoryConfig::new(PathBuf::from(path), &agent_id, embedding_dim as usize);
|
||||||
match HDF5Memory::create(config) {
|
match HDF5Memory::create(config) {
|
||||||
Ok(mem) => Box::into_raw(Box::new(mem)),
|
Ok(mem) => Box::into_raw(Box::new(Mutex::new(mem))),
|
||||||
Err(_) => ptr::null_mut(),
|
Err(_) => ptr::null_mut(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -67,7 +72,7 @@ pub unsafe extern "C" fn edgehdf5_open(path: *const c_char) -> Handle {
|
|||||||
};
|
};
|
||||||
|
|
||||||
match HDF5Memory::open(std::path::Path::new(&path)) {
|
match HDF5Memory::open(std::path::Path::new(&path)) {
|
||||||
Ok(mem) => Box::into_raw(Box::new(mem)),
|
Ok(mem) => Box::into_raw(Box::new(Mutex::new(mem))),
|
||||||
Err(_) => ptr::null_mut(),
|
Err(_) => ptr::null_mut(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -82,7 +87,7 @@ pub unsafe extern "C" fn edgehdf5_open(path: *const c_char) -> Handle {
|
|||||||
pub unsafe extern "C" fn edgehdf5_close(handle: Handle) {
|
pub unsafe extern "C" fn edgehdf5_close(handle: Handle) {
|
||||||
if !handle.is_null() {
|
if !handle.is_null() {
|
||||||
// SAFETY: handle was created by Box::into_raw in edgehdf5_create; this is the final use.
|
// SAFETY: handle was created by Box::into_raw in edgehdf5_create; this is the final use.
|
||||||
unsafe { drop(Box::from_raw(handle)) };
|
unsafe { drop(Box::<Mutex<HDF5Memory>>::from_raw(handle)) };
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -115,11 +120,15 @@ pub unsafe extern "C" fn edgehdf5_save(
|
|||||||
session_id: *const c_char,
|
session_id: *const c_char,
|
||||||
tags: *const c_char,
|
tags: *const c_char,
|
||||||
) -> i64 {
|
) -> i64 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_mut() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return -1,
|
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.
|
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||||
let chunk = match unsafe { cstr_to_string(chunk) } {
|
let chunk = match unsafe { cstr_to_string(chunk) } {
|
||||||
@@ -176,7 +185,7 @@ pub unsafe extern "C" fn edgehdf5_save(
|
|||||||
pub unsafe extern "C" fn edgehdf5_count_active(handle: Handle) -> u64 {
|
pub unsafe extern "C" fn edgehdf5_count_active(handle: Handle) -> u64 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
match unsafe { handle.as_ref() } {
|
match unsafe { handle.as_ref() } {
|
||||||
Some(mem) => mem.count_active() as u64,
|
Some(mtx) => mtx.lock().map(|g| g.count_active() as u64).unwrap_or(0),
|
||||||
None => 0,
|
None => 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -190,7 +199,7 @@ pub unsafe extern "C" fn edgehdf5_count_active(handle: Handle) -> u64 {
|
|||||||
pub unsafe extern "C" fn edgehdf5_count(handle: Handle) -> u64 {
|
pub unsafe extern "C" fn edgehdf5_count(handle: Handle) -> u64 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
match unsafe { handle.as_ref() } {
|
match unsafe { handle.as_ref() } {
|
||||||
Some(mem) => mem.count() as u64,
|
Some(mtx) => mtx.lock().map(|g| g.count() as u64).unwrap_or(0),
|
||||||
None => 0,
|
None => 0,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -202,11 +211,15 @@ pub unsafe extern "C" fn edgehdf5_count(handle: Handle) -> u64 {
|
|||||||
/// `handle` must be a valid, non-null handle.
|
/// `handle` must be a valid, non-null handle.
|
||||||
#[unsafe(no_mangle)]
|
#[unsafe(no_mangle)]
|
||||||
pub unsafe extern "C" fn edgehdf5_delete(handle: Handle, index: u64) -> i32 {
|
pub unsafe extern "C" fn edgehdf5_delete(handle: Handle, index: u64) -> i32 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_mut() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return -1,
|
None => return -1,
|
||||||
};
|
};
|
||||||
|
let mut mem = match mtx.lock() {
|
||||||
|
Ok(g) => g,
|
||||||
|
Err(_) => return -1,
|
||||||
|
};
|
||||||
|
|
||||||
match mem.delete(index as usize) {
|
match mem.delete(index as usize) {
|
||||||
Ok(()) => 0,
|
Ok(()) => 0,
|
||||||
@@ -250,11 +263,15 @@ pub unsafe extern "C" fn edgehdf5_hybrid_search(
|
|||||||
out_scores: *mut f32,
|
out_scores: *mut f32,
|
||||||
out_chunks: *mut *mut c_char,
|
out_chunks: *mut *mut c_char,
|
||||||
) -> u32 {
|
) -> u32 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_mut() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return 0,
|
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.
|
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||||
let query_text = match unsafe { cstr_to_string(query_text) } {
|
let query_text = match unsafe { cstr_to_string(query_text) } {
|
||||||
Some(s) => s,
|
Some(s) => s,
|
||||||
@@ -329,11 +346,15 @@ pub unsafe extern "C" fn edgehdf5_add_session(
|
|||||||
channel: *const c_char,
|
channel: *const c_char,
|
||||||
summary: *const c_char,
|
summary: *const c_char,
|
||||||
) -> i32 {
|
) -> i32 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_mut() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return -1,
|
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.
|
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||||
let id = match unsafe { cstr_to_string(id) } {
|
let id = match unsafe { cstr_to_string(id) } {
|
||||||
Some(s) => s,
|
Some(s) => s,
|
||||||
@@ -375,10 +396,14 @@ pub unsafe extern "C" fn edgehdf5_get_session_summary(
|
|||||||
session_id: *const c_char,
|
session_id: *const c_char,
|
||||||
) -> *mut c_char {
|
) -> *mut c_char {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_ref() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return ptr::null_mut(),
|
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.
|
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||||
let session_id = match unsafe { cstr_to_string(session_id) } {
|
let session_id = match unsafe { cstr_to_string(session_id) } {
|
||||||
Some(s) => s,
|
Some(s) => s,
|
||||||
@@ -411,11 +436,15 @@ pub unsafe extern "C" fn edgehdf5_add_entity(
|
|||||||
entity_type: *const c_char,
|
entity_type: *const c_char,
|
||||||
embedding_idx: i64,
|
embedding_idx: i64,
|
||||||
) -> i64 {
|
) -> i64 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_mut() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return -1,
|
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.
|
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||||
let name = match unsafe { cstr_to_string(name) } {
|
let name = match unsafe { cstr_to_string(name) } {
|
||||||
Some(s) => s,
|
Some(s) => s,
|
||||||
@@ -447,11 +476,15 @@ pub unsafe extern "C" fn edgehdf5_add_relation(
|
|||||||
relation: *const c_char,
|
relation: *const c_char,
|
||||||
weight: f32,
|
weight: f32,
|
||||||
) -> i32 {
|
) -> i32 {
|
||||||
// SAFETY: handle is a valid non-null Handle from edgehdf5_create; caller ensures exclusive access.
|
// SAFETY: handle is a valid non-null Handle from edgehdf5_create.
|
||||||
let mem = match unsafe { handle.as_mut() } {
|
let mtx = match unsafe { handle.as_ref() } {
|
||||||
Some(m) => m,
|
Some(m) => m,
|
||||||
None => return -1,
|
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.
|
// SAFETY: JNI caller guarantees the pointer argument is a valid null-terminated C string.
|
||||||
let relation = match unsafe { cstr_to_string(relation) } {
|
let relation = match unsafe { cstr_to_string(relation) } {
|
||||||
Some(s) => s,
|
Some(s) => s,
|
||||||
@@ -492,7 +525,8 @@ mod tests {
|
|||||||
fn open_handle(dir: &tempfile::TempDir) -> Handle {
|
fn open_handle(dir: &tempfile::TempDir) -> Handle {
|
||||||
let path = CString::new(dir.path().join("mem.h5").to_str().unwrap()).unwrap();
|
let path = CString::new(dir.path().join("mem.h5").to_str().unwrap()).unwrap();
|
||||||
let agent_id = CString::new("test-agent").unwrap();
|
let agent_id = CString::new("test-agent").unwrap();
|
||||||
// SAFETY: both C strings are valid and null-terminated.
|
// SAFETY: both C strings are valid and null-terminated; returned handle
|
||||||
|
// wraps HDF5Memory in a Mutex and is safe to use from multiple threads.
|
||||||
unsafe { edgehdf5_create(path.as_ptr(), agent_id.as_ptr(), EMBEDDING_DIM) }
|
unsafe { edgehdf5_create(path.as_ptr(), agent_id.as_ptr(), EMBEDDING_DIM) }
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -590,4 +624,44 @@ mod tests {
|
|||||||
|
|
||||||
unsafe { edgehdf5_close(handle) };
|
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,12 +739,190 @@ impl HnswIndex {
|
|||||||
pub fn m_max0(&self) -> usize {
|
pub fn m_max0(&self) -> usize {
|
||||||
self.m_max0
|
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
|
// 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`.
|
/// Greedy search: find the single closest node to `query` starting from `ep`.
|
||||||
fn greedy_closest(
|
fn greedy_closest(
|
||||||
vectors: &[Vec<f32>],
|
vectors: &[Vec<f32>],
|
||||||
@@ -1452,4 +1630,75 @@ mod tests {
|
|||||||
assert_eq!(results.len(), 3);
|
assert_eq!(results.len(), 3);
|
||||||
assert_eq!(results[0].0, 0);
|
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
|
harness = false
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["std", "checksum", "deflate", "provenance", "fast-deflate", "system-zlib-decompress"]
|
default = ["std", "checksum", "deflate", "provenance", "system-zlib-decompress"]
|
||||||
std = []
|
std = []
|
||||||
checksum = []
|
checksum = []
|
||||||
deflate = ["flate2"]
|
deflate = ["flate2"]
|
||||||
|
|||||||
@@ -143,10 +143,6 @@ pub fn parse_vds_mappings(
|
|||||||
let source_selection = read_selection(heap_data, &mut pos)?;
|
let source_selection = read_selection(heap_data, &mut pos)?;
|
||||||
let virtual_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 {
|
mappings.push(VdsMapping {
|
||||||
source_file,
|
source_file,
|
||||||
source_dataset,
|
source_dataset,
|
||||||
@@ -158,37 +154,6 @@ pub fn parse_vds_mappings(
|
|||||||
Ok(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`.
|
/// 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> {
|
fn read_null_terminated_string(data: &[u8], pos: &mut usize) -> Result<String, FormatError> {
|
||||||
let start = *pos;
|
let start = *pos;
|
||||||
@@ -897,68 +862,4 @@ mod tests {
|
|||||||
let blob = [0x01u8, 0, 0, 0, 0, 0, 0, 0, 0];
|
let blob = [0x01u8, 0, 0, 0, 0, 0, 0, 0, 0];
|
||||||
assert!(parse_vds_mappings(&blob, 8).unwrap().is_empty());
|
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 | |||||||