From 8c75b43fe61cb15ee7b8db5fab68828eafdbb7ec Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 31 Jan 2026 02:04:48 +0000 Subject: [PATCH] Fix 9 issues: NaN safety, atomic ordering, TOFU verification, data safety, and more 1. Replace .unwrap() with .unwrap_or(Ordering::Equal) on partial_cmp for NaN safety (db.rs, tools.rs) 2. Use SeqCst ordering and compare_exchange for rate limiter atomics (server.rs) 3. Add TOFU hash verification for tokenizer.json and config.json (embedder.rs) 4. Reorder expire_item to insert-before-delete for crash safety (db.rs) 5. Add comment documenting intentional format! usage in cleanup_expired (db.rs) 6. Log errors instead of silently discarding in transfer_edges (graph.rs) 7. Cap co-access pairing to first 3 IDs to limit O(n^2) edge creation (graph.rs) 8. Fix replace parameter description to match store-before-delete behavior (tools.rs) 9. Strengthen install.sh warning when checksums.txt unavailable, output to stderr https://claude.ai/code/session_01VoNncGCWckv3qRQaKpGRC5 --- install.sh | 2 +- src/db.rs | 25 +++++++++++++++++-------- src/embedder.rs | 2 ++ src/graph.rs | 10 +++++++++- src/mcp/server.rs | 15 ++++++++++----- src/mcp/tools.rs | 8 ++++++-- 6 files changed, 45 insertions(+), 17 deletions(-) diff --git a/install.sh b/install.sh index cab5909..b690907 100755 --- a/install.sh +++ b/install.sh @@ -70,7 +70,7 @@ if curl -fsSL "$CHECKSUMS_URL" -o "${TMPDIR}/checksums.txt" 2>/dev/null; then fi fi else - echo "No checksums.txt available, skipping verification." + echo "WARNING: No checksums.txt available, binary integrity could not be verified." >&2 fi tar -xzf "${TMPDIR}/${TARBALL}" -C "$TMPDIR" diff --git a/src/db.rs b/src/db.rs index 75cbc4f..2363bb6 100644 --- a/src/db.rs +++ b/src/db.rs @@ -535,7 +535,11 @@ impl Database { // Convert map to sorted vec let mut search_results: Vec = results_map.into_values().map(|(r, _)| r).collect(); - search_results.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap()); + search_results.sort_by(|a, b| { + b.similarity + .partial_cmp(&a.similarity) + .unwrap_or(std::cmp::Ordering::Equal) + }); search_results.truncate(limit); Ok(search_results) @@ -598,7 +602,11 @@ impl Database { } // Sort by similarity descending - conflicts.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap()); + conflicts.sort_by(|a, b| { + b.similarity + .partial_cmp(&a.similarity) + .unwrap_or(std::cmp::Ordering::Equal) + }); Ok(conflicts) } @@ -753,12 +761,7 @@ impl Database { item.embedding = self.embedder.embed(&embedding_text)?; item.expires_at = Some(expires_at); - // Delete then re-insert with new expiration - table - .delete(&format!("id = '{}'", sanitize_sql_string(id))) - .await - .map_err(|e| SedimentError::Database(format!("Delete for expire failed: {}", e)))?; - + // Insert-before-delete to avoid data loss on crash let batch = item_to_batch(&item)?; let batches = RecordBatchIterator::new(vec![Ok(batch)], Arc::new(item_schema())); table @@ -767,6 +770,11 @@ impl Database { .await .map_err(|e| SedimentError::Database(format!("Re-insert for expire failed: {}", e)))?; + table + .delete(&format!("id = '{}'", sanitize_sql_string(id))) + .await + .map_err(|e| SedimentError::Database(format!("Delete for expire failed: {}", e)))?; + Ok(()) } @@ -823,6 +831,7 @@ impl Database { }; let now = Utc::now().timestamp(); + // now is a system-generated i64 timestamp, no string sanitization needed let filter = format!("expires_at IS NOT NULL AND expires_at < {}", now); // Count how many will be deleted diff --git a/src/embedder.rs b/src/embedder.rs index 1981cf6..1c428a9 100644 --- a/src/embedder.rs +++ b/src/embedder.rs @@ -215,6 +215,8 @@ fn download_model(model_id: &str) -> Result<(PathBuf, PathBuf, PathBuf)> { // TOFU: verify model integrity verify_tofu_hash(&model_path, "model.safetensors")?; + verify_tofu_hash(&tokenizer_path, "tokenizer.json")?; + verify_tofu_hash(&config_path, "config.json")?; Ok((model_path, tokenizer_path, config_path)) } diff --git a/src/graph.rs b/src/graph.rs index eef6411..e545c3b 100644 --- a/src/graph.rs +++ b/src/graph.rs @@ -225,6 +225,12 @@ impl GraphStore { return Ok(()); } + let item_ids = if item_ids.len() > 3 { + &item_ids[..3] + } else { + item_ids + }; + let now = chrono::Utc::now().timestamp(); for i in 0..item_ids.len() { @@ -336,7 +342,9 @@ impl GraphStore { // Create edges on the new node for (neighbor, strength, rel_type, _) in &edges { - let _ = self.add_related_edge(to_id, neighbor, *strength, rel_type); + if let Err(e) = self.add_related_edge(to_id, neighbor, *strength, rel_type) { + tracing::warn!("transfer edge to {} failed: {}", neighbor, e); + } } Ok(()) diff --git a/src/mcp/server.rs b/src/mcp/server.rs index 5dadbc4..c6cba1e 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -205,13 +205,18 @@ fn handle_call_tool( .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_millis() as u64; - let window = ctx.rate_limit_window.load(Ordering::Relaxed); + let window = ctx.rate_limit_window.load(Ordering::SeqCst); if now_ms - window > 60_000 { - // Reset window - ctx.rate_limit_window.store(now_ms, Ordering::Relaxed); - ctx.rate_limit_count.store(1, Ordering::Relaxed); + // Reset window (use compare_exchange to avoid racing with another reset) + if ctx + .rate_limit_window + .compare_exchange(window, now_ms, Ordering::SeqCst, Ordering::SeqCst) + .is_ok() + { + ctx.rate_limit_count.store(1, Ordering::SeqCst); + } } else { - let count = ctx.rate_limit_count.fetch_add(1, Ordering::Relaxed) + 1; + let count = ctx.rate_limit_count.fetch_add(1, Ordering::SeqCst) + 1; if count > MAX_CALLS_PER_MINUTE { let result = super::protocol::CallToolResult::error("Rate limit exceeded, try again later"); diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index 0002e7d..cd0cf96 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -62,7 +62,7 @@ pub fn get_tools() -> Vec { }, "replace": { "type": "string", - "description": "ID of an existing item to replace (atomically delete before storing)" + "description": "ID of an existing item to replace (stores new item first, then deletes old)" }, "related": { "type": "array", @@ -594,7 +594,11 @@ pub async fn recall_pipeline( result.similarity = (base_score * trust_bonus).min(1.0); } - results.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap()); + results.sort_by(|a, b| { + b.similarity + .partial_cmp(&a.similarity) + .unwrap_or(std::cmp::Ordering::Equal) + }); } // Record access