Skip to main content

max / audiofiles

Robustness: reject non-finite .afcl vectors, make VP-tree search iterative F2: a crafted or corrupt .afcl carrying a non-finite feature vector could reach the k-NN scorer, where it poisoned the standardization means (one NaN made every standardized vector NaN) and hit partial_cmp().unwrap() distance sorts. afcl::import now rejects non-finite exemplar vectors at the trust boundary; exemplar::build_index filters every vector through is_usable_vector (covering legacy/synced rows too); and the three distance/score sorts in exemplar.rs and trained_head.rs use total_cmp, which never panics on NaN. F3: search_nearest and search_within were recursive. On a degenerate library (many identical vectors) the tree flattens past MAX_BUILD_DEPTH into a single INFINITY-threshold chain of length ~N, and the recursion overflowed the call stack on the first query. Both searches are now iterative with a heap-allocated work-stack; the prune bound is unchanged and stays sound against the monotonically-tightening tau. Tests: .afcl non-finite rejection, is_usable_vector, score-on-NaN-query, 50k-identical-vector VP-tree query. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Author: Max Johnson <me@maxj.phd> · 2026-07-02 21:54 UTC
Commit: 6d2e93623c26564f70d072b63a30cc28ed5c20e3
Parent: f8e9072
4 files changed, +178 insertions, -53 deletions
M crates/audiofiles-core/src/analysis/afcl.rs +28
@@ -297,6 +297,14 @@ pub fn import(db: &Database, afcl: &Afcl, source: Option<&str>) -> Result<Import
297 297 if ex.vector.len() != NUM_FEATURES {
298 298 continue;
299 299 }
300 + // Reject non-finite vectors at the trust boundary. A crafted or corrupt
301 + // `.afcl` carrying a NaN/Inf component would otherwise flow into the k-NN
302 + // index, where it poisons standardization means (one NaN -> every
303 + // standardized vector NaN) and reaches the distance-sort comparators.
304 + if ex.vector.iter().any(|v| !v.is_finite()) {
305 + tracing::warn!("skipping imported exemplar with non-finite feature vector");
306 + continue;
307 + }
300 308 let vector = serde_json::to_string(&ex.vector)
301 309 .map_err(|e| CoreError::Serialization(e.to_string()))?;
302 310 let tags = serde_json::to_string(&ex.tags)
@@ -570,6 +578,26 @@ mod tests {
570 578 }
571 579
572 580 #[test]
581 + fn import_rejects_non_finite_exemplar_vectors() {
582 + let src = Database::open_in_memory().unwrap();
583 + seed_library(&src);
584 + let mut afcl = build_export(&src, &ExportOptions::default()).unwrap();
585 + assert_eq!(afcl.exemplars.len(), 3);
586 + // Poison one exemplar's vector. A crafted `.afcl` with an overflowing
587 + // exponent (`1e999`) parses to inf; it must be skipped at the trust
588 + // boundary rather than stored and later poisoning the k-NN index.
589 + afcl.exemplars[0].vector = vec![f64::INFINITY; NUM_FEATURES];
590 +
591 + let dst = Database::open_in_memory().unwrap();
592 + let summary = import(&dst, &afcl, Some("crafted.afcl")).unwrap();
593 + assert_eq!(summary.exemplars, 2, "the non-finite exemplar is rejected");
594 +
595 + let imp = enabled_imported_exemplars(&dst).unwrap();
596 + assert_eq!(imp.len(), 2);
597 + assert!(imp.iter().all(|e| e.vector.iter().all(|x| x.is_finite())));
598 + }
599 +
600 + #[test]
573 601 fn import_does_not_clobber_existing_policy() {
574 602 let src = Database::open_in_memory().unwrap();
575 603 seed_library(&src); // kick policy 0.4/0.8
M crates/audiofiles-core/src/analysis/exemplar.rs +60 -4
@@ -181,15 +181,20 @@ pub fn build_index(db: &Database) -> Result<ExemplarIndex> {
181 181 let Some(tags) = tags_by_hash.remove(&hash) else { continue };
182 182 let v: Vec<f64> = serde_json::from_str(&json)
183 183 .map_err(|e| CoreError::Serialization(e.to_string()))?;
184 - if v.len() == NUM_FEATURES {
184 + if is_usable_vector(&v) {
185 185 raw.push((hash, v, tags, 1.0));
186 186 }
187 187 }
188 188 }
189 189
190 190 // Imported `.afcl` layer exemplars (enabled layers), weighted below `local`.
191 + // A single non-finite component would poison the standardization means for the
192 + // whole set, so a malformed vector (crafted `.afcl`, a legacy row, or one
193 + // synced from another device) is dropped here as well as at `.afcl` import.
191 194 for imp in super::afcl::enabled_imported_exemplars(db)? {
192 - raw.push((imp.id, imp.vector, imp.tags, imp.weight));
195 + if is_usable_vector(&imp.vector) {
196 + raw.push((imp.id, imp.vector, imp.tags, imp.weight));
197 + }
193 198 }
194 199
195 200 // Standardization params over the full exemplar set (local + imported).
@@ -207,6 +212,13 @@ pub fn build_index(db: &Database) -> Result<ExemplarIndex> {
207 212 Ok(ExemplarIndex { means, stds, exemplars, feature_version: FEATURE_VERSION })
208 213 }
209 214
215 + /// A feature vector is usable in the k-NN index only if it has the expected
216 + /// dimension and every component is finite. Non-finite components would poison
217 + /// the standardization means and reach the distance-sort comparators.
218 + fn is_usable_vector(v: &[f64]) -> bool {
219 + v.len() == NUM_FEATURES && v.iter().all(|x| x.is_finite())
220 + }
221 +
210 222 fn standardization_params(raw: &[(String, Vec<f64>)]) -> (Vec<f64>, Vec<f64>) {
211 223 if raw.is_empty() {
212 224 return (vec![0.0; NUM_FEATURES], vec![1.0; NUM_FEATURES]);
@@ -281,7 +293,9 @@ impl ExemplarIndex {
281 293 if dists.is_empty() {
282 294 return Vec::new();
283 295 }
284 - dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
296 + // total_cmp never panics on NaN (unlike partial_cmp().unwrap()); the
297 + // build_index finiteness filter should keep NaN out, this is the backstop.
298 + dists.sort_by(|a, b| a.1.total_cmp(&b.1));
285 299 dists.truncate(k);
286 300
287 301 // Adaptive Gaussian bandwidth = mean neighbor distance.
@@ -322,7 +336,7 @@ impl ExemplarIndex {
322 336 neighbors,
323 337 })
324 338 .collect();
325 - scores.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap().then(a.tag.cmp(&b.tag)));
339 + scores.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.tag.cmp(&b.tag)));
326 340 scores
327 341 }
328 342 }
@@ -491,6 +505,48 @@ mod tests {
491 505 }
492 506
493 507 #[test]
508 + fn is_usable_vector_rejects_non_finite_and_wrong_len() {
509 + assert!(is_usable_vector(&vec_at(0.5)));
510 + // Wrong dimension.
511 + assert!(!is_usable_vector(&[0.0; NUM_FEATURES - 1]));
512 + // Any non-finite component disqualifies it (would poison the means).
513 + let mut nan = vec_at(0.1).to_vec();
514 + nan[3] = f64::NAN;
515 + assert!(!is_usable_vector(&nan));
516 + let mut inf = vec_at(0.1).to_vec();
517 + inf[7] = f64::INFINITY;
518 + assert!(!is_usable_vector(&inf));
519 + }
520 +
521 + #[test]
522 + fn standardization_means_stay_finite_when_bad_vectors_are_pre_filtered() {
523 + // build_index gates every vector through is_usable_vector, so the params
524 + // only ever see finite input. Prove the params themselves stay finite and
525 + // that a NaN slipping into the *scorer* still cannot panic the sort.
526 + let raw = vec![
527 + ("a".to_string(), vec_at(0.0).to_vec()),
528 + ("b".to_string(), vec_at(1.0).to_vec()),
529 + ];
530 + let (means, stds) = standardization_params(&raw);
531 + assert!(means.iter().all(|m| m.is_finite()));
532 + assert!(stds.iter().all(|s| s.is_finite() && *s > 0.0));
533 + }
534 +
535 + #[test]
536 + fn score_does_not_panic_on_nan_query() {
537 + let db = Database::open_in_memory().unwrap();
538 + insert(&db, "k1", vec_at(0.0), &["instrument.drum.kick"]);
539 + insert(&db, "k2", vec_at(0.1), &["instrument.drum.kick"]);
540 + let index = build_index(&db).unwrap();
541 +
542 + // A NaN in the query reaches the distance sort; total_cmp must keep it
543 + // from panicking (returns some ordering rather than unwrapping None).
544 + let mut q = vec_at(0.02).to_vec();
545 + q[2] = f64::NAN;
546 + let _ = index.score(&q, DEFAULT_K, None); // must not panic
547 + }
548 +
549 + #[test]
494 550 fn auto_applies_dominant_neighbor_tag() {
495 551 let db = Database::open_in_memory().unwrap();
496 552 // A tight cluster of kicks, plus one far-away snare.
M crates/audiofiles-core/src/analysis/trained_head.rs +3 -1
@@ -83,7 +83,9 @@ impl TrainedHead {
83 83 TagScore { tag: m.tag.clone(), score: sigmoid(z), neighbors: Vec::new() }
84 84 })
85 85 .collect();
86 - scores.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap().then(a.tag.cmp(&b.tag)));
86 + // total_cmp never panics on NaN; the head is trained from the same
87 + // finiteness-filtered exemplar set, so this is a backstop.
88 + scores.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.tag.cmp(&b.tag)));
87 89 scores
88 90 }
89 91 }
M crates/audiofiles-core/src/vp_tree.rs +87 -48
@@ -267,44 +267,55 @@ fn search_nearest<T>(
267 267 heap: &mut BinaryHeap<HeapEntry>,
268 268 tau: &mut f64,
269 269 ) {
270 - let node = &nodes[node_idx];
271 - let d = dist(query, &items[node.item_idx]);
272 -
273 - // Consider this node's vantage point.
274 - if d < *tau || heap.len() < k {
275 - heap.push(HeapEntry {
276 - distance: d,
277 - index: node.item_idx,
278 - });
279 - if heap.len() > k {
280 - heap.pop();
281 - }
282 - if heap.len() == k {
283 - *tau = heap.peek().unwrap().distance;
270 + // Iterative traversal with a heap-allocated stack. A recursive walk overflows
271 + // the call stack on a degenerate tree: many identical vectors get flattened at
272 + // MAX_BUILD_DEPTH into a single INFINITY-threshold left-linked chain of length
273 + // ~N, and an uncapped recursion would then use one call frame per node. The
274 + // explicit stack grows on the heap instead, so the first query on such a
275 + // library no longer crashes.
276 + //
277 + // Pruning uses the triangle-inequality bound against the current `tau`, which
278 + // only ever tightens — so pruning a subtree with a later (smaller) tau is
279 + // always sound (every point in it is provably farther than the k-th nearest).
280 + let mut stack: Vec<usize> = vec![node_idx];
281 + while let Some(idx) = stack.pop() {
282 + let node = &nodes[idx];
283 + let d = dist(query, &items[node.item_idx]);
284 +
285 + // Consider this node's vantage point.
286 + if d < *tau || heap.len() < k {
287 + heap.push(HeapEntry {
288 + distance: d,
289 + index: node.item_idx,
290 + });
291 + if heap.len() > k {
292 + heap.pop();
293 + }
294 + if heap.len() == k {
295 + *tau = heap.peek().unwrap().distance;
296 + }
284 297 }
285 - }
286 298
287 - // Search the closer subtree first for better pruning.
288 - if d <= node.threshold {
289 - // Query is inside — search inside first.
290 - if let Some(left) = node.left
291 - && d - *tau <= node.threshold {
292 - search_nearest(items, nodes, left, query, k, dist, heap, tau);
299 + let go_left = node.left.filter(|_| d - *tau <= node.threshold);
300 + let go_right = node.right.filter(|_| d + *tau > node.threshold);
301 +
302 + // Push so the closer subtree is popped first (tighter pruning): the
303 + // closer child goes on last.
304 + if d <= node.threshold {
305 + if let Some(right) = go_right {
306 + stack.push(right);
293 307 }
294 - if let Some(right) = node.right
295 - && d + *tau > node.threshold {
296 - search_nearest(items, nodes, right, query, k, dist, heap, tau);
308 + if let Some(left) = go_left {
309 + stack.push(left);
297 310 }
298 - } else {
299 - // Query is outside — search outside first.
300 - if let Some(right) = node.right
301 - && d + *tau > node.threshold {
302 - search_nearest(items, nodes, right, query, k, dist, heap, tau);
311 + } else {
312 + if let Some(left) = go_left {
313 + stack.push(left);
303 314 }
304 - if let Some(left) = node.left
305 - && d - *tau <= node.threshold {
306 - search_nearest(items, nodes, left, query, k, dist, heap, tau);
315 + if let Some(right) = go_right {
316 + stack.push(right);
307 317 }
318 + }
308 319 }
309 320 }
310 321
@@ -319,25 +330,33 @@ fn search_within<T>(
319 330 dist: &impl Fn(&T, &T) -> f64,
320 331 results: &mut Vec<VpMatch>,
321 332 ) {
322 - let node = &nodes[node_idx];
323 - let d = dist(query, &items[node.item_idx]);
324 -
325 - if d <= radius {
326 - results.push(VpMatch {
327 - index: node.item_idx,
328 - distance: d,
329 - });
330 - }
333 + // Iterative for the same degenerate-chain stack-overflow reason as
334 + // search_nearest. `radius` is fixed, so the prune bounds are identical to the
335 + // recursive version — only the result order differs, and callers sort.
336 + let mut stack: Vec<usize> = vec![node_idx];
337 + while let Some(idx) = stack.pop() {
338 + let node = &nodes[idx];
339 + let d = dist(query, &items[node.item_idx]);
340 +
341 + if d <= radius {
342 + results.push(VpMatch {
343 + index: node.item_idx,
344 + distance: d,
345 + });
346 + }
331 347
332 - // Prune subtrees using triangle inequality bounds.
333 - if let Some(left) = node.left
334 - && d - radius <= node.threshold {
335 - search_within(items, nodes, left, query, radius, dist, results);
348 + // Prune subtrees using triangle inequality bounds.
349 + if let Some(left) = node.left
350 + && d - radius <= node.threshold
351 + {
352 + stack.push(left);
336 353 }
337 - if let Some(right) = node.right
338 - && d + radius > node.threshold {
339 - search_within(items, nodes, right, query, radius, dist, results);
354 + if let Some(right) = node.right
355 + && d + radius > node.threshold
356 + {
357 + stack.push(right);
340 358 }
359 + }
341 360 }
342 361
343 362 #[cfg(test)]
@@ -497,6 +516,26 @@ mod tests {
497 516 }
498 517
499 518 #[test]
519 + fn degenerate_identical_library_does_not_overflow() {
520 + // Tens of thousands of identical points collapse past MAX_BUILD_DEPTH into
521 + // a single INFINITY-threshold left-linked chain of length ~N. The old
522 + // recursive search used one call frame per chain node and stack-overflowed
523 + // on the first query; the iterative search must handle it and still be
524 + // correct.
525 + let n = 50_000;
526 + let items = vec![1.0f64; n];
527 + let tree = VpTree::build(items, euclidean_1d);
528 +
529 + let nearest = tree.find_nearest(&1.0, 10, euclidean_1d);
530 + assert_eq!(nearest.len(), 10);
531 + assert!(nearest.iter().all(|r| r.distance == 0.0));
532 +
533 + // Every item is within any radius, so range search must return them all.
534 + let within = tree.find_within(&1.0, 0.5, euclidean_1d);
535 + assert_eq!(within.len(), n);
536 + }
537 +
538 + #[test]
500 539 fn two_items() {
501 540 let tree = VpTree::build(vec![0.0, 10.0], euclidean_1d);
502 541 let results = tree.find_nearest(&3.0, 1, euclidean_1d);