From f24fe0efd81ff1466fe9cb7d0df0e8bcfcb4be3b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 20 Aug 2026 07:53:34 +0000 Subject: [PATCH] Fix RSS step-up after daily CVD hot-reload. Compiling a new engine while the previous one stayed live briefly doubled RSS; glibc then kept the old arenas mapped, so memory plateaued around 1.5x after daily.cvd published. Stream CVD verify/unpack from disk, pin Aho-Corasick to a contiguous NFA, switch Linux to jemalloc, and purge unused pages after each swap. Co-authored-by: Eldad A. Fux --- Cargo.lock | 39 ++++++++++ Cargo.toml | 5 ++ Dockerfile | 3 +- README.md | 10 ++- src/alloc.rs | 79 ++++++++++++++++++++ src/cvd/header.rs | 25 +++++++ src/cvd/mod.rs | 2 +- src/cvd/unpack.rs | 77 +++++++++++++++++++- src/cvd/verify.rs | 75 +++++++++++++++++-- src/engine/mod.rs | 74 +++++++++++++++---- src/lib.rs | 1 + src/main.rs | 1 + src/signatures/hash.rs | 13 ++++ src/signatures/mod.rs | 26 +++++++ src/updater/mod.rs | 162 ++++++++++++++++++++++++----------------- tests/memory.rs | 1 + tests/pipeline.rs | 16 ++++ 17 files changed, 510 insertions(+), 99 deletions(-) create mode 100644 src/alloc.rs diff --git a/Cargo.lock b/Cargo.lock index c4d81a0..50b95ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -383,6 +383,8 @@ dependencies = [ "tar", "tempfile", "thiserror", + "tikv-jemalloc-ctl", + "tikv-jemallocator", "tokio", "tokio-stream", "tower", @@ -1001,6 +1003,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1518,6 +1526,37 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "tikv-jemalloc-ctl" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a184c43b8ab2f41df2733b55556e3f5f632f4aeaa205b1bb018f574b7f5f142" +dependencies = [ + "libc", + "paste", + "tikv-jemalloc-sys", +] + +[[package]] +name = "tikv-jemalloc-sys" +version = "0.7.1+5.3.1-0-g81034ce1f1373e37dc865038e1bc8eeecf559ce8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a2825c78386b4ae0314074867860ba9577875de945f05992c38815cbec327f0" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "249f09e49ab1609436f34c776e84231bead18d6a955f119f939bdc1d847561bd" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "tinytemplate" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index f417b65..5dce72c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,11 @@ tower-http = { version = "=0.6.2", features = ["limit", "trace", "timeout"] } tracing = "=0.1.41" tracing-subscriber = { version = "=0.3.19", features = ["env-filter"] } +# jemalloc returns unused pages after a CVD hot-reload; glibc typically does not. +[target.'cfg(target_os = "linux")'.dependencies] +tikv-jemallocator = { version = "=0.7.0", features = ["background_threads", "unprefixed_malloc_on_supported_platforms"] } +tikv-jemalloc-ctl = "=0.7.0" + [dev-dependencies] criterion = { version = "=0.5.1", features = ["html_reports", "async_tokio"] } http-body-util = "=0.1.2" diff --git a/Dockerfile b/Dockerfile index 4c41b8d..6e81e0d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -34,7 +34,8 @@ ENV DEFENDER_LISTEN=0.0.0.0:8080 \ DEFENDER_UPDATE_INTERVAL_SECS=3600 \ DEFENDER_MAX_BYTES=67108864 \ DEFENDER_USER_AGENT="ClamAV/1.4.2 (defender; docker)" \ - RUST_LOG=info + RUST_LOG=info \ + MALLOC_CONF=background_thread:true,dirty_decay_ms:1000,muzzy_decay_ms:1000 USER defender EXPOSE 8080 diff --git a/README.md b/README.md index dc3c4ef..2adbf08 100644 --- a/README.md +++ b/README.md @@ -73,11 +73,13 @@ A background task (same process, independent Tokio task): 1. `Range: bytes=0-511` against configured mirrors to read the remote CVD header 2. Downloads `main.cvd` / `daily.cvd` when the version/MD5 is newer 3. Verifies MD5 + RSA -4. Compiles a **new** `Engine` off the request path -5. Atomically swaps it with `arc-swap` +4. Compiles a **new** `Engine` off the request path (streaming CVD members so the gzip body is not held alongside the compiled signatures) +5. Atomically swaps it with `arc-swap`, then returns unused heap pages to the OS In-flight scans keep the previous `Arc` until they finish. No connection drop, no restart. +A daily `daily.cvd` publish is the usual trigger. RSS will rise while both engines exist, then fall back after the swap (`rss_before` / `rss_compiled` / `rss_after` on the reload log line). + Default mirrors: - `https://database.clamav.net` @@ -162,7 +164,7 @@ Official ClamAV CVD files as of 18 Aug 2026 (from `database.clamav.net`): | `bytecode.cvd` | **0.27 MiB** (281,702 B) | 1.24 MiB | 80 | Not executed (no bytecode VM) | | **Total baked** | **107.56 MiB** | **308.8 MiB** | | Image includes `main` + `daily` | -Loaded into the scanner (main + daily, PUA off): ~540k file hashes, ~102k body signatures, ~307k logical signatures. Resident set with that engine is about **1.4 GiB**. +Loaded into the scanner (main + daily, PUA off): ~540k file hashes, ~102k body signatures, ~307k logical signatures. Resident set with that engine is about **1.4 GiB**. The Aho-Corasick prefilter uses a contiguous NFA (not a DFA) so a database reload cannot balloon the automaton. After a hot-swap, jemalloc purges unused dirty pages so RSS does not stay at the two-engine peak. ## Development @@ -220,7 +222,7 @@ Official `daily.cvd` over the same HTTP path: | `POST /scan` clean 64 KiB | 667 µs | 93.7 MiB/s | | `POST /scan` EICAR ×16 concurrent | 296 µs / batch | **54.0 k req/s** | -Tiny requests are loopback/HTTP-latency bound; large bodies are bounded by MD5+SHA1+SHA256. RSS tests (`tests/memory.rs`) stay stable across 20k scans and 200 engine swaps. +Tiny requests are loopback/HTTP-latency bound; large bodies are bounded by MD5+SHA1+SHA256. RSS tests (`tests/memory.rs`) stay stable across 20k scans and 200 engine swaps. Reloading official CVDs briefly overlaps the previous engine; unused pages are returned after the swap. ## Architecture diff --git a/src/alloc.rs b/src/alloc.rs new file mode 100644 index 0000000..c3016b8 --- /dev/null +++ b/src/alloc.rs @@ -0,0 +1,79 @@ +//! Process allocator helpers. +//! +//! On Linux the binary uses jemalloc so that dropping a compiled [`crate::engine::Engine`] +//! actually returns pages to the OS. glibc `malloc` commonly keeps the old +//! arenas mapped, which is what a CVD hot-reload looks like as a permanent RSS +//! step-up (old engine + new engine, then “stuck” at ~1.5×). + +#[cfg(target_os = "linux")] +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +/// Enable jemalloc’s background purge thread. Safe to call more than once. +pub fn init() { + #[cfg(target_os = "linux")] + { + if let Err(e) = tikv_jemalloc_ctl::background_thread::write(true) { + tracing::debug!(error = ?e, "jemalloc background_thread not enabled"); + } + } +} + +/// Return unused heap pages to the OS. +/// +/// Call after compiling an engine (scratch buffers from unpack/AC construction) +/// and again after swapping so the previous engine’s pages can be unmapped. +pub fn reclaim_unused_pages() { + #[cfg(target_os = "linux")] + { + let _ = tikv_jemalloc_ctl::epoch::advance(); + // `MALLCTL_ARENAS_ALL` is `u32::MAX`: purge every arena, including the + // tokio blocking-pool thread that compiled the engine. + let rc = unsafe { tikv_jemalloc_ctl::raw::write(b"arena.4294967295.purge\0", ()) }; + if let Err(e) = rc { + tracing::debug!(error = ?e, "jemalloc arena purge failed"); + } + } +} + +/// Current resident set in bytes (`VmRSS`), if `/proc` is available. +pub fn rss_bytes() -> Option { + let status = std::fs::read_to_string("/proc/self/status").ok()?; + for line in status.lines() { + if let Some(rest) = line.strip_prefix("VmRSS:") { + let kb: u64 = rest.split_whitespace().next()?.parse().ok()?; + return Some(kb.saturating_mul(1024)); + } + } + None +} + +/// Format a byte count for logs (`123.4 MiB`). +pub fn format_bytes(n: u64) -> String { + const MIB: f64 = 1024.0 * 1024.0; + if n >= 1024 * 1024 { + format!("{:.1} MiB", n as f64 / MIB) + } else if n >= 1024 { + format!("{:.1} KiB", n as f64 / 1024.0) + } else { + format!("{n} B") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn format_bytes_scales() { + assert_eq!(format_bytes(512), "512 B"); + assert_eq!(format_bytes(2048), "2.0 KiB"); + assert_eq!(format_bytes(2 * 1024 * 1024), "2.0 MiB"); + } + + #[test] + fn reclaim_does_not_panic() { + reclaim_unused_pages(); + let _ = rss_bytes(); + } +} diff --git a/src/cvd/header.rs b/src/cvd/header.rs index cce3426..cef4e73 100644 --- a/src/cvd/header.rs +++ b/src/cvd/header.rs @@ -1,5 +1,9 @@ //! 512-byte ClamAV CVD/CLD header parser. +use std::fs::File; +use std::io::Read; +use std::path::Path; + use crate::error::{Error, Result}; /// Size of the fixed CVD header prefix. @@ -101,6 +105,15 @@ impl CvdHeader { }) } + /// Read only the 512-byte header from `path` (does not load the gzip body). + pub fn read_file(path: impl AsRef) -> Result { + let path = path.as_ref(); + let mut file = File::open(path).map_err(|e| Error::io(path, e))?; + let mut buf = [0u8; CVD_HEADER_SIZE]; + file.read_exact(&mut buf).map_err(|e| Error::io(path, e))?; + Self::parse(&buf) + } + /// Serialize back to a 512-byte padded header. pub fn to_bytes(&self) -> [u8; CVD_HEADER_SIZE] { let s = format!( @@ -197,4 +210,16 @@ mod tests { let h2 = CvdHeader::parse(&bytes).unwrap(); assert_eq!(h, h2); } + + #[test] + fn read_file_ignores_body() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("daily.cvd"); + let h = CvdHeader::parse_str(SAMPLE).unwrap(); + let mut bytes = h.to_bytes().to_vec(); + bytes.extend_from_slice(&[0u8; 1024 * 1024]); + std::fs::write(&path, &bytes).unwrap(); + let loaded = CvdHeader::read_file(&path).unwrap(); + assert_eq!(loaded, h); + } } diff --git a/src/cvd/mod.rs b/src/cvd/mod.rs index 2ff1a6d..3be10f5 100644 --- a/src/cvd/mod.rs +++ b/src/cvd/mod.rs @@ -15,7 +15,7 @@ pub mod unpack; pub mod verify; pub use header::CvdHeader; -pub use unpack::{unpack_cvd, UnpackedDb}; +pub use unpack::{for_each_cvd_member, unpack_cvd, UnpackedDb}; pub use verify::{verify_cvd, verify_cvd_bytes, VerifyMode}; use crate::error::Result; diff --git a/src/cvd/unpack.rs b/src/cvd/unpack.rs index 83e1e7a..3e9135c 100644 --- a/src/cvd/unpack.rs +++ b/src/cvd/unpack.rs @@ -1,13 +1,16 @@ //! Unpack the gzip+tar body that follows a 512-byte CVD header. use std::collections::BTreeMap; +use std::fs::File; use std::io::{Cursor, Read}; +use std::path::Path; use flate2::read::GzDecoder; use tar::Archive; use super::header::CVD_HEADER_SIZE; use crate::error::{Error, Result}; +use crate::signatures::is_signature_member; /// Files extracted from a CVD/CLD archive, keyed by file name (no path). #[derive(Debug, Clone, Default)] @@ -39,11 +42,43 @@ pub fn unpack_cvd(data: &[u8]) -> Result { pub fn unpack_body(body: &[u8]) -> Result { let gz = GzDecoder::new(Cursor::new(body)); - let mut archive = Archive::new(gz); + let mut files = BTreeMap::new(); + for_each_archive_file(gz, true, true, |name, data| { + files.insert(name, data); + Ok(()) + })?; + Ok(UnpackedDb { files }) +} + +/// Stream CVD tar members from disk without retaining the compressed file or +/// previously visited members. The 512-byte header is skipped; callers should +/// already have authenticated the file. +/// +/// Non-signature members (bytecode, YARA, …) and PUA files (when `load_pua` is +/// false) are consumed and discarded without keeping their contents. +pub fn for_each_cvd_member( + path: impl AsRef, + load_pua: bool, + mut visit: impl FnMut(&str, &[u8]) -> Result<()>, +) -> Result<()> { + let path = path.as_ref(); + let mut file = File::open(path).map_err(|e| Error::io(path, e))?; + let mut hdr = [0u8; CVD_HEADER_SIZE]; + file.read_exact(&mut hdr).map_err(|e| Error::io(path, e))?; + let gz = GzDecoder::new(file); + for_each_archive_file(gz, false, load_pua, |name, data| visit(&name, &data)) +} + +fn for_each_archive_file( + reader: R, + load_all: bool, + load_pua: bool, + mut visit: impl FnMut(String, Vec) -> Result<()>, +) -> Result<()> { + let mut archive = Archive::new(reader); archive.set_overwrite(false); archive.set_preserve_permissions(false); - let mut files = BTreeMap::new(); let entries = archive .entries() .map_err(|e| Error::CvdUnpack(e.to_string()))?; @@ -59,13 +94,18 @@ pub fn unpack_body(body: &[u8]) -> Result { .and_then(|s| s.to_str()) .ok_or_else(|| Error::CvdUnpack("non-utf8 member name".into()))? .to_string(); + if !load_all && !is_signature_member(&name, load_pua) { + std::io::copy(&mut entry, &mut std::io::sink()) + .map_err(|e| Error::CvdUnpack(e.to_string()))?; + continue; + } let mut buf = Vec::new(); entry .read_to_end(&mut buf) .map_err(|e| Error::CvdUnpack(e.to_string()))?; - files.insert(name, buf); + visit(name, buf)?; } - Ok(UnpackedDb { files }) + Ok(()) } /// Build a synthetic CVD body (gzip tar) from name → contents. Header is not included. @@ -140,4 +180,33 @@ mod tests { assert_eq!(unpacked.get("test.hdb").unwrap(), files[0].1); assert_eq!(unpacked.get("test.ndb").unwrap(), files[1].1); } + + #[test] + fn for_each_member_matches_unpack() { + let files = [ + ( + "test.hdb", + b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa:1:Eicar-Test-File\n".as_slice(), + ), + ("test.ndb", b"Eicar:0:*:585530\n".as_slice()), + ("skip.cbc", b"not-a-signature".as_slice()), + ]; + let cvd = pack_cvd(&files, 1, "unit").unwrap(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t.cvd"); + std::fs::write(&path, &cvd).unwrap(); + + let mut seen = Vec::new(); + for_each_cvd_member(&path, false, |name, data| { + seen.push((name.to_string(), data.to_vec())); + Ok(()) + }) + .unwrap(); + seen.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!(seen.len(), 2); + assert_eq!(seen[0].0, "test.hdb"); + assert_eq!(seen[1].0, "test.ndb"); + assert_eq!(seen[0].1, files[0].1); + assert_eq!(seen[1].1, files[1].1); + } } diff --git a/src/cvd/verify.rs b/src/cvd/verify.rs index 0623ec3..2f2991f 100644 --- a/src/cvd/verify.rs +++ b/src/cvd/verify.rs @@ -69,13 +69,52 @@ pub fn verify_cvd_bytes(data: &[u8], header: &CvdHeader, mode: VerifyMode) -> Re Err(Error::CvdSignature) } -/// Verify a CVD file on disk. +/// Verify a CVD file on disk without buffering the gzip body. pub fn verify_cvd(path: impl AsRef, mode: VerifyMode) -> Result { + use std::io::Read; + let path = path.as_ref(); - let data = std::fs::read(path).map_err(|e| Error::io(path, e))?; - let header = CvdHeader::parse(&data)?; - verify_cvd_bytes(&data, &header, mode)?; - Ok(header) + let mut file = std::fs::File::open(path).map_err(|e| Error::io(path, e))?; + let mut hdr = [0u8; CVD_HEADER_SIZE]; + file.read_exact(&mut hdr).map_err(|e| Error::io(path, e))?; + let header = CvdHeader::parse(&hdr)?; + + let mut md5 = Md5::new(); + let mut sha = Sha256::new(); + let mut buf = [0u8; 128 * 1024]; + loop { + let n = file.read(&mut buf).map_err(|e| Error::io(path, e))?; + if n == 0 { + break; + } + md5.update(&buf[..n]); + if mode == VerifyMode::Official { + sha.update(&buf[..n]); + } + } + let digest = hex::encode(md5.finalize()); + if digest != header.md5 { + return Err(Error::CvdChecksum { + expected: header.md5.clone(), + actual: digest, + }); + } + if mode == VerifyMode::Integrity { + return Ok(header); + } + if verify_legacy_md5(&header.md5, &header.dsig) { + return Ok(header); + } + let sha = sha.finalize(); + if verify_pss(&sha, &header.dsig, CLI_NSTR, CLI_ESTR) { + return Ok(header); + } + for (_name, n, e) in CLAMAV_RSA_KEYS.iter().skip(1) { + if verify_pss(&sha, &header.dsig, n, e) { + return Ok(header); + } + } + Err(Error::CvdSignature) } pub fn md5_hex(data: &[u8]) -> String { @@ -296,4 +335,30 @@ mod tests { data[CVD_HEADER_SIZE..].copy_from_slice(body); assert!(verify_cvd_bytes(&data, &header, VerifyMode::Integrity).is_err()); } + + #[test] + fn verify_cvd_streams_from_disk() { + let body = b"hello-gzip-body"; + let md5 = md5_hex(body); + let mut data = vec![0u8; CVD_HEADER_SIZE + body.len()]; + let header = CvdHeader { + magic: "ClamAV-VDB".into(), + time: "now".into(), + version: 1, + signatures: 0, + flevel: 1, + md5: md5.clone(), + dsig: "placeholder".into(), + builder: "test".into(), + stime: 0, + }; + data[..CVD_HEADER_SIZE].copy_from_slice(&header.to_bytes()); + data[CVD_HEADER_SIZE..].copy_from_slice(body); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("x.cvd"); + std::fs::write(&path, &data).unwrap(); + let loaded = verify_cvd(&path, VerifyMode::Integrity).unwrap(); + assert_eq!(loaded.md5, md5); + assert_eq!(loaded.version, 1); + } } diff --git a/src/engine/mod.rs b/src/engine/mod.rs index 07f7565..6e29be7 100644 --- a/src/engine/mod.rs +++ b/src/engine/mod.rs @@ -6,14 +6,14 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Instant; -use aho_corasick::{AhoCorasick, AhoCorasickBuilder, MatchKind}; +use aho_corasick::{AhoCorasick, AhoCorasickBuilder, AhoCorasickKind, MatchKind}; use arc_swap::ArcSwap; use md5::{Digest, Md5}; use rustc_hash::{FxHashMap, FxHashSet}; use sha1::Sha1; use sha2::Sha256; -use crate::cvd::{load_bytes, CvdHeader, UnpackedDb, VerifyMode}; +use crate::cvd::{for_each_cvd_member, load_bytes, verify_cvd, CvdHeader, UnpackedDb, VerifyMode}; use crate::error::{Error, Result}; use crate::signatures::hash::{FpSet, HashAlgo, HashDb}; use crate::signatures::ldb::{load_ldb, LogicalSig}; @@ -210,19 +210,19 @@ impl Engine { match ext.as_str() { "cvd" | "cld" => { let t1 = Instant::now(); - let bytes = std::fs::read(&path).map_err(|e| Error::io(&path, e))?; let mode = if ext == "cvd" { verify } else { VerifyMode::Integrity }; + let file_bytes = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0); tracing::info!( file = %name, - bytes = bytes.len(), + bytes = file_bytes, verify = ?mode, "reading CVD" ); - let (header, unpacked) = load_bytes(&bytes, mode)?; + let header = verify_cvd(&path, mode)?; tracing::info!( file = %name, version = header.version, @@ -230,25 +230,41 @@ impl Engine { flevel = header.flevel, builder = %header.builder, built = %header.time, - members = unpacked.files.len(), elapsed_ms = t1.elapsed().as_millis() as u64, - "verified and unpacked CVD" + "verified CVD, ingesting members" + ); + builder.add_cvd_header(&name, header); + let mut members = 0usize; + for_each_cvd_member(&path, load_pua, |fname, data| { + builder.add_named_file(fname, data, load_pua); + members += 1; + Ok(()) + })?; + tracing::info!( + file = %name, + members, + elapsed_ms = t1.elapsed().as_millis() as u64, + "ingested CVD members" ); - builder.add_cvd(&name, header, &unpacked, load_pua); cvds += 1; } _ => { - tracing::debug!(file = %name, "loading signature file"); - builder.add_named_file( - &name, - &std::fs::read(&path).unwrap_or_default(), - load_pua, - ); + if crate::signatures::is_signature_member(&name, load_pua) { + tracing::debug!(file = %name, "loading signature file"); + builder.add_named_file( + &name, + &std::fs::read(&path).unwrap_or_default(), + load_pua, + ); + } else { + tracing::debug!(file = %name, "skipping non-signature file"); + } } } } tracing::info!(cvds, "compiling scan engine"); let engine = builder.build(); + crate::alloc::reclaim_unused_pages(); tracing::info!( file_hashes = engine.meta.file_hashes, section_hashes = engine.meta.section_hashes, @@ -665,7 +681,7 @@ impl EngineBuilder { } } - fn add_cvd(&mut self, name: &str, header: CvdHeader, unpacked: &UnpackedDb, load_pua: bool) { + fn add_cvd_header(&mut self, name: &str, header: CvdHeader) { self.databases.push(CvdInfo { name: name.to_string(), version: header.version, @@ -675,6 +691,10 @@ impl EngineBuilder { time: header.time, md5: header.md5, }); + } + + fn add_cvd(&mut self, name: &str, header: CvdHeader, unpacked: &UnpackedDb, load_pua: bool) { + self.add_cvd_header(name, header); for (fname, data) in &unpacked.files { self.add_named_file(fname, data, load_pua); } @@ -736,7 +756,15 @@ impl EngineBuilder { } } - fn build(self) -> Engine { + fn build(mut self) -> Engine { + self.file_hash.shrink_to_fit(); + self.section_hash.shrink_to_fit(); + self.fp.shrink_to_fit(); + self.body.shrink_to_fit(); + self.logical.shrink_to_fit(); + self.ignored.shrink_to_fit(); + self.ignored_prefix.shrink_to_fit(); + let mut needles: Vec> = Vec::new(); let mut needle_idx: FxHashMap, usize> = FxHashMap::default(); @@ -791,12 +819,26 @@ impl EngineBuilder { if ac_sub.len() < n { ac_sub.resize(n, Vec::new()); } + drop(needle_idx); + needles.shrink_to_fit(); + for v in &mut ac_body { + v.shrink_to_fit(); + } + for v in &mut ac_sub { + v.shrink_to_fit(); + } + slow_body.shrink_to_fit(); + slow_logical.shrink_to_fit(); let ac = if needles.is_empty() { None } else { + // ContiguousNFA is the memory-bounded automaton. The default + // heuristic may pick a DFA, which can be many times larger and is + // the main reason a daily CVD reload doubles RSS. AhoCorasickBuilder::new() .match_kind(MatchKind::Standard) + .kind(Some(AhoCorasickKind::ContiguousNFA)) .ascii_case_insensitive(false) .build(&needles) .ok() diff --git a/src/lib.rs b/src/lib.rs index 8bda93d..16c7be7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ //! High-performance streaming virus scanner using ClamAV public databases. +pub mod alloc; pub mod config; pub mod cvd; pub mod engine; diff --git a/src/main.rs b/src/main.rs index 89cc6e2..99b9dfb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ const VERSION: &str = env!("CARGO_PKG_VERSION"); #[tokio::main] async fn main() -> anyhow::Result<()> { init_tracing(); + defender::alloc::init(); let cfg = Config::from_env(); tracing::info!( diff --git a/src/signatures/hash.rs b/src/signatures/hash.rs index 0b8f50e..af0a629 100644 --- a/src/signatures/hash.rs +++ b/src/signatures/hash.rs @@ -186,6 +186,13 @@ impl HashDb { pub fn is_empty(&self) -> bool { self.len() == 0 } + + pub fn shrink_to_fit(&mut self) { + self.md5.shrink_to_fit(); + self.sha1.shrink_to_fit(); + self.sha256.shrink_to_fit(); + self.names.shrink_to_fit(); + } } fn lookup<'a, const N: usize>( @@ -248,6 +255,12 @@ impl FpSet { pub fn contains_sha256(&self, d: &[u8; 32]) -> bool { self.sha256.contains(d) } + + pub fn shrink_to_fit(&mut self) { + self.md5.shrink_to_fit(); + self.sha1.shrink_to_fit(); + self.sha256.shrink_to_fit(); + } } #[cfg(test)] diff --git a/src/signatures/mod.rs b/src/signatures/mod.rs index 4bc097b..bb57b4f 100644 --- a/src/signatures/mod.rs +++ b/src/signatures/mod.rs @@ -9,3 +9,29 @@ pub use hash::{FpSet, HashAlgo, HashDb, HashSig}; pub use hexpat::HexPattern; pub use ldb::{load_ldb, LogicalSig}; pub use ndb::{load_ndb, NdbSig, OffsetKind, TargetType}; + +/// True if `name` is a signature file the engine can ingest. +pub fn is_signature_member(name: &str, load_pua: bool) -> bool { + let lower = name.to_ascii_lowercase(); + let ext = lower.rsplit_once('.').map(|(_, e)| e).unwrap_or(""); + match ext { + "hdb" | "hsb" | "mdb" | "msb" | "ndb" | "ldb" | "fp" | "sfp" | "ign" | "ign2" => true, + "hdu" | "hsu" | "mdu" | "msu" | "ndu" | "ldu" => load_pua, + _ => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classifies_members() { + assert!(is_signature_member("main.mdb", false)); + assert!(is_signature_member("daily.ldb", false)); + assert!(!is_signature_member("main.mdu", false)); + assert!(is_signature_member("main.mdu", true)); + assert!(!is_signature_member("bytecode.cbc", true)); + assert!(!is_signature_member("main.cvd.tmp", false)); + } +} diff --git a/src/updater/mod.rs b/src/updater/mod.rs index 0c113a2..58655d3 100644 --- a/src/updater/mod.rs +++ b/src/updater/mod.rs @@ -1,10 +1,15 @@ //! Background CVD updater: download, verify, hot-swap with zero downtime. +use std::path::Path; use std::time::{Duration, Instant}; +use tokio::io::AsyncWriteExt; +use tokio_stream::StreamExt; + +use crate::alloc; use crate::config::Config; use crate::cvd::header::{CvdHeader, CVD_HEADER_SIZE}; -use crate::cvd::verify::{verify_cvd_bytes, VerifyMode}; +use crate::cvd::verify::{verify_cvd, VerifyMode}; use crate::engine::{Database, Engine}; use crate::error::{Error, Result}; @@ -88,7 +93,7 @@ impl Updater { "remote CVD header" ); if dest.exists() { - if let Ok(local) = CvdHeader::parse(&std::fs::read(&dest).unwrap_or_default()) { + if let Ok(local) = CvdHeader::read_file(&dest) { if local.version >= remote.version && local.md5 == remote.md5 { tracing::info!( db = name, @@ -123,41 +128,7 @@ impl Updater { ); } - let bytes = self.download(name).await?; - tracing::info!( - db = name, - bytes = bytes.len(), - "verifying CVD checksum and digital signature" - ); - let header = CvdHeader::parse(&bytes)?; - let mode = if self.cfg.verify_official { - VerifyMode::Official - } else { - VerifyMode::Integrity - }; - verify_cvd_bytes(&bytes, &header, mode)?; - if header.version != remote.version && dest.exists() { - tracing::debug!( - db = name, - header = header.version, - advertised = remote.version, - "version differs between Range header and full file" - ); - } - - let tmp = dest.with_extension("cvd.tmp"); - std::fs::write(&tmp, &bytes).map_err(|e| Error::io(&tmp, e))?; - std::fs::rename(&tmp, &dest).map_err(|e| Error::io(&dest, e))?; - tracing::info!( - db = name, - version = header.version, - signatures = header.signatures, - builder = %header.builder, - built = %header.time, - bytes = bytes.len(), - path = %dest.display(), - "verified and installed CVD" - ); + self.download_to(name, &dest).await?; Ok(true) } @@ -216,12 +187,12 @@ impl Updater { CvdHeader::parse(&bytes) } - async fn download(&self, name: &str) -> Result> { + async fn download_to(&self, name: &str, dest: &Path) -> Result<()> { let mut last_err = None; for (i, mirror) in self.cfg.mirrors.iter().enumerate() { let url = format!("{mirror}/{name}.cvd"); - match self.download_url(name, &url).await { - Ok(b) => return Ok(b), + match self.download_url_to(name, &url, dest).await { + Ok(()) => return Ok(()), Err(e) => { let remaining = self.cfg.mirrors.len() - i - 1; if remaining > 0 { @@ -242,8 +213,9 @@ impl Updater { Err(last_err.unwrap_or_else(|| Error::Update("no mirrors".into()))) } - async fn download_url(&self, name: &str, url: &str) -> Result> { + async fn download_url_to(&self, name: &str, url: &str, dest: &Path) -> Result<()> { let t0 = Instant::now(); + let tmp = dest.with_extension("cvd.tmp"); let resp = self .client .get(url) @@ -260,30 +232,71 @@ impl Updater { advertised_bytes, "downloading CVD" ); - let bytes = resp - .bytes() - .await - .map_err(|e| Error::Update(e.to_string()))?; - let elapsed_ms = t0.elapsed().as_millis() as u64; - let bytes_per_sec = if elapsed_ms > 0 { - (bytes.len() as u128 * 1000 / elapsed_ms as u128) as u64 - } else { - 0 - }; - tracing::info!( - db = name, - %url, - bytes = bytes.len(), - elapsed_ms, - bytes_per_sec, - "download complete" - ); - Ok(bytes.to_vec()) + + let result = async { + let mut file = tokio::fs::File::create(&tmp) + .await + .map_err(|e| Error::io(&tmp, e))?; + let mut stream = resp.bytes_stream(); + let mut written = 0u64; + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|e| Error::Update(e.to_string()))?; + written += chunk.len() as u64; + file.write_all(&chunk) + .await + .map_err(|e| Error::io(&tmp, e))?; + } + file.flush().await.map_err(|e| Error::io(&tmp, e))?; + drop(file); + + let mode = if self.cfg.verify_official { + VerifyMode::Official + } else { + VerifyMode::Integrity + }; + let tmp_path = tmp.clone(); + let header = tokio::task::spawn_blocking(move || verify_cvd(&tmp_path, mode)) + .await + .map_err(|e| Error::Update(e.to_string()))??; + + tokio::fs::rename(&tmp, dest) + .await + .map_err(|e| Error::io(dest, e))?; + + let elapsed_ms = t0.elapsed().as_millis() as u64; + let bytes_per_sec = if elapsed_ms > 0 { + (written as u128 * 1000 / elapsed_ms as u128) as u64 + } else { + 0 + }; + tracing::info!( + db = name, + %url, + version = header.version, + signatures = header.signatures, + builder = %header.builder, + built = %header.time, + bytes = written, + elapsed_ms, + bytes_per_sec, + path = %dest.display(), + "verified and installed CVD" + ); + Ok(()) + } + .await; + + if result.is_err() { + let _ = tokio::fs::remove_file(&tmp).await; + } + result } async fn reload(&self) -> Result<()> { + let rss_before = alloc::rss_bytes(); tracing::info!( dir = %self.cfg.db_dir.display(), + rss = rss_before.map(alloc::format_bytes), "compiling new scan engine from disk (in-flight scans keep the previous engine)" ); let t0 = Instant::now(); @@ -297,17 +310,31 @@ impl Updater { let engine = tokio::task::spawn_blocking(move || Engine::load_dir(&dir, verify, pua)) .await .map_err(|e| Error::Update(e.to_string()))??; + let rss_compiled = alloc::rss_bytes(); + let file_hashes = engine.meta.file_hashes; + let section_hashes = engine.meta.section_hashes; + let body = engine.meta.body_sigs; + let logical = engine.meta.logical_sigs; + let skipped = engine.meta.skipped_sigs; + let db_count = engine.meta.databases.len(); + let databases = engine.meta.databases.clone(); + self.db.swap(engine); + alloc::reclaim_unused_pages(); + let rss_after = alloc::rss_bytes(); tracing::info!( - file_hashes = engine.meta.file_hashes, - section_hashes = engine.meta.section_hashes, - body = engine.meta.body_sigs, - logical = engine.meta.logical_sigs, - skipped = engine.meta.skipped_sigs, - databases = engine.meta.databases.len(), + file_hashes, + section_hashes, + body, + logical, + skipped, + databases = db_count, elapsed_ms = t0.elapsed().as_millis() as u64, + rss_before = rss_before.map(alloc::format_bytes), + rss_compiled = rss_compiled.map(alloc::format_bytes), + rss_after = rss_after.map(alloc::format_bytes), "scan engine swapped atomically" ); - for d in &engine.meta.databases { + for d in &databases { tracing::info!( db = %d.name, version = d.version, @@ -316,7 +343,6 @@ impl Updater { "active CVD" ); } - self.db.swap(engine); Ok(()) } } diff --git a/tests/memory.rs b/tests/memory.rs index 1a0affc..a902e3c 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -71,6 +71,7 @@ fn rss_stable_after_engine_swaps() { } // Force drop of old engines. db.swap(tiny_engine(9999)); + defender::alloc::reclaim_unused_pages(); for _ in 0..10 { let _ = db.current().scan(EICAR); } diff --git a/tests/pipeline.rs b/tests/pipeline.rs index af66eb0..d8934aa 100644 --- a/tests/pipeline.rs +++ b/tests/pipeline.rs @@ -130,3 +130,19 @@ fn load_dir_roundtrip() { defender::ScanVerdict::Infected { .. } )); } + +#[test] +fn load_dir_skips_tmp_and_unknown_files() { + let dir = tempfile::tempdir().unwrap(); + let md5 = hex::encode(Md5::digest(EICAR)); + let hdb = format!("{md5}:68:Eicar-Test-Signature\n"); + let cvd = pack_cvd(&[("d.hdb", hdb.as_bytes())], 3, "dir").unwrap(); + std::fs::write(dir.path().join("daily.cvd"), &cvd).unwrap(); + std::fs::write(dir.path().join("daily.cvd.tmp"), b"not a cvd").unwrap(); + std::fs::write(dir.path().join("README.txt"), b"ignore me").unwrap(); + let eng = Engine::load_dir(dir.path(), VerifyMode::Integrity, false).unwrap(); + assert!(matches!( + eng.scan(EICAR).verdict, + defender::ScanVerdict::Infected { .. } + )); +}