From be4e82592f80cdc352371e6918680876d093c33b Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 09:32:36 -0500 Subject: [PATCH 1/4] harden: cap VSIM COUNT and EF parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit VSIM's COUNT and EF_SEARCH parameters accepted unbounded u64 values, allowing a single request like `VSIM key 0.1 COUNT 999999999999` to cause immediate OOM by attempting to allocate a massive results vector. - add MAX_VSIM_COUNT (10,000) — generous for any practical similarity search while preventing memory exhaustion - add MAX_VSIM_EF (1,024) — consistent with VADD's MAX_HNSW_PARAM, prevents worst-case O(n) graph traversal --- crates/ember-protocol/src/command.rs | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index b46f7501..eac52a65 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -18,6 +18,14 @@ const MAX_VECTOR_DIMS: usize = 65_536; /// Values above 1024 give no practical benefit and waste memory. const MAX_HNSW_PARAM: u64 = 1024; +/// Maximum number of results for VSIM. 10,000 is generous for any practical +/// similarity search while preventing OOM from unbounded result allocation. +const MAX_VSIM_COUNT: u64 = 10_000; + +/// Maximum search beam width for VSIM. Same cap as MAX_HNSW_PARAM — +/// larger values cause worst-case O(n) graph traversal with no accuracy gain. +const MAX_VSIM_EF: u64 = MAX_HNSW_PARAM; + /// Expiration option for the SET command. #[derive(Debug, Clone, PartialEq, Eq)] pub enum SetExpire { @@ -1993,7 +2001,13 @@ fn parse_vsim(args: &[Frame]) -> Result { "VSIM: COUNT requires a value".into(), )); } - count = Some(parse_u64(&args[idx], "VSIM")? as usize); + let c = parse_u64(&args[idx], "VSIM")?; + if c > MAX_VSIM_COUNT { + return Err(ProtocolError::InvalidCommandFrame(format!( + "VSIM: COUNT {c} exceeds max {MAX_VSIM_COUNT}" + ))); + } + count = Some(c as usize); idx += 1; } "EF" => { @@ -2003,7 +2017,13 @@ fn parse_vsim(args: &[Frame]) -> Result { "VSIM: EF requires a value".into(), )); } - ef_search = parse_u64(&args[idx], "VSIM")? as usize; + let ef = parse_u64(&args[idx], "VSIM")?; + if ef > MAX_VSIM_EF { + return Err(ProtocolError::InvalidCommandFrame(format!( + "VSIM: EF {ef} exceeds max {MAX_VSIM_EF}" + ))); + } + ef_search = ef as usize; idx += 1; } "WITHSCORES" => { From 8248d5f2f2ad5bf7932984e0a11651f71311432b Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 09:32:44 -0500 Subject: [PATCH 2/4] harden: fix panics and overflows in vector operations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - replace unreachable!() in vrem with proper WrongType error return - use saturating arithmetic in vadd memory estimate to prevent overflow from bypassing memory limits - use saturating_mul in VectorSet::memory_usage() to prevent overflow in tracking calculations - improve VectorSet::clone fallback chain — adds intermediate fallback layers before last-resort panic, with tracing::error logging - guard search result consistency with min() on key/distance lengths --- crates/ember-core/src/keyspace.rs | 11 ++-- crates/ember-core/src/types/vector.rs | 73 ++++++++++++++++++++------- 2 files changed, 62 insertions(+), 22 deletions(-) diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index 1771f74c..dfdf1671 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -1884,10 +1884,13 @@ impl Keyspace { return Err(VectorWriteError::WrongType); } - // estimate memory for the new vector + // estimate memory for the new vector (saturating to avoid overflow) let dim = vector.len(); - let per_vector = - dim * quantization.bytes_per_element() + connectivity * 2 * 8 + element.len() + 80; + let per_vector = dim + .saturating_mul(quantization.bytes_per_element()) + .saturating_add(connectivity.saturating_mul(16)) + .saturating_add(element.len()) + .saturating_add(80); let estimated_increase = if is_new { memory::ENTRY_OVERHEAD + key.len() + VectorSet::BASE_OVERHEAD + per_vector } else { @@ -1976,7 +1979,7 @@ impl Keyspace { let removed = match entry.value { Value::Vector(ref mut vs) => vs.remove(element), - _ => unreachable!(), + _ => return Err(WrongType), }; if removed { diff --git a/crates/ember-core/src/types/vector.rs b/crates/ember-core/src/types/vector.rs index 48fdcc51..089ad2ef 100644 --- a/crates/ember-core/src/types/vector.rs +++ b/crates/ember-core/src/types/vector.rs @@ -329,7 +329,10 @@ impl VectorSet { .search(query, count) .map_err(|e| VectorError::Index(e.to_string()))?; - let mut results = Vec::with_capacity(matches.keys.len()); + // usearch should always return matching key/distance counts, but + // guard against library bugs by using the shorter length + let result_count = matches.keys.len().min(matches.distances.len()); + let mut results = Vec::with_capacity(result_count); for (key, distance) in matches.keys.iter().zip(matches.distances.iter()) { if let Some(name) = self.names.get(key) { results.push(SearchResult { @@ -417,8 +420,10 @@ impl VectorSet { let count = self.elements.len(); // usearch internal: vector storage + HNSW graph edges - let vector_bytes = count * self.dim * self.quantization.bytes_per_element(); - let graph_bytes = count * self.connectivity * 2 * 8; // each edge is a u64 key + let vector_bytes = count + .saturating_mul(self.dim) + .saturating_mul(self.quantization.bytes_per_element()); + let graph_bytes = count.saturating_mul(self.connectivity).saturating_mul(16); // each edge is a u64 key // rust-side hashmaps: elements + names let name_bytes: usize = self @@ -476,22 +481,54 @@ impl Clone for VectorSet { // try_clone rebuilds the index from scratch. if it fails (e.g. OOM // inside usearch), return an empty set with the same config rather // than panicking. data loss is preferable to a server crash. - self.try_clone().unwrap_or_else(|_| { - Self::new( - self.dim, - self.metric, - self.quantization, - self.connectivity, - self.expansion_add, - ) - .unwrap_or_else(|_| { - // if we can't even create an empty index with the same - // config, fall back to the smallest possible valid set. - // this only happens under extreme memory pressure. - Self::new(self.dim, self.metric, self.quantization, 2, 2) - .expect("failed to allocate even a minimal index — system is out of memory") + self.try_clone() + .or_else(|_| { + Self::new( + self.dim, + self.metric, + self.quantization, + self.connectivity, + self.expansion_add, + ) + }) + .or_else(|_| Self::new(self.dim, self.metric, self.quantization, 2, 2)) + .unwrap_or_else(|e| { + // absolute last resort — log the error and return a 1-dim + // empty set. this loses data but keeps the server alive. + tracing::error!("VectorSet clone failed under extreme memory pressure: {e}"); + Self { + index: Index::new(&IndexOptions { + dimensions: self.dim, + metric: self.metric.to_metric_kind(), + quantization: self.quantization.to_scalar_kind(), + connectivity: 2, + expansion_add: 2, + expansion_search: 0, + multi: false, + }) + .unwrap_or_else(|_| { + // if even this fails, create the absolute minimum + Index::new(&IndexOptions { + dimensions: 1, + metric: MetricKind::Cos, + quantization: ScalarKind::F32, + connectivity: 2, + expansion_add: 2, + expansion_search: 0, + multi: false, + }) + .expect("cannot allocate 1-dim index — system is critically out of memory") + }), + elements: HashMap::new(), + names: HashMap::new(), + next_key: 0, + dim: self.dim, + metric: self.metric, + quantization: self.quantization, + connectivity: 2, + expansion_add: 2, + } }) - }) } } From bfa02b0005f3991e71feebd87929b45e100b8680 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 09:32:53 -0500 Subject: [PATCH 3/4] harden: validate vector dimensions and counts in persistence a crafted AOF or snapshot file could specify dimension=4 billion, causing the recovery loop to iterate 4B times and exhaust memory despite capped_capacity limiting the initial allocation. - add MAX_PERSISTED_VECTOR_DIMS (65,536) and MAX_PERSISTED_VECTOR_COUNT (10M) constants in format.rs - reject AOF records with dim > MAX_PERSISTED_VECTOR_DIMS - reject snapshot entries with dim or count exceeding limits - validate metric (0-2) and quantization (0-2) enum values in snapshot deserialization to catch corruption early - add FormatError::InvalidData variant for structured error reporting --- crates/ember-persistence/src/aof.rs | 8 +++- crates/ember-persistence/src/format.rs | 12 ++++++ crates/ember-persistence/src/snapshot.rs | 48 +++++++++++++++++++++++- 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 6b11914f..96e67a05 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -613,7 +613,13 @@ impl AofRecord { let key = read_string(&mut cursor, "key")?; let element = read_string(&mut cursor, "element")?; let dim = format::read_u32(&mut cursor)?; - let mut vector = Vec::with_capacity(format::capped_capacity(dim)); + if dim > format::MAX_PERSISTED_VECTOR_DIMS { + return Err(FormatError::InvalidData(format!( + "AOF VADD dimension {dim} exceeds max {}", + format::MAX_PERSISTED_VECTOR_DIMS + ))); + } + let mut vector = Vec::with_capacity(dim as usize); for _ in 0..dim { vector.push(format::read_f32(&mut cursor)?); } diff --git a/crates/ember-persistence/src/format.rs b/crates/ember-persistence/src/format.rs index 7c946f99..2f818dbc 100644 --- a/crates/ember-persistence/src/format.rs +++ b/crates/ember-persistence/src/format.rs @@ -43,6 +43,9 @@ pub enum FormatError { #[error("unknown record tag: {0}")] UnknownTag(u8), + #[error("invalid data: {0}")] + InvalidData(String), + #[error("file is encrypted but no encryption key was provided")] EncryptionRequired, @@ -247,6 +250,15 @@ pub fn capped_capacity(count: u32) -> usize { (count as usize).min(65_536) } +/// Maximum vector dimensions allowed in persistence formats. +/// Matches the protocol-layer cap. Records exceeding this are rejected +/// during deserialization to prevent OOM from corrupt files. +pub const MAX_PERSISTED_VECTOR_DIMS: u32 = 65_536; + +/// Maximum element count per vector set in persistence formats. +/// Prevents corrupt count fields from causing unbounded loops. +pub const MAX_PERSISTED_VECTOR_COUNT: u32 = 10_000_000; + /// Verifies that two CRC32 values match. pub fn verify_crc32_values(computed: u32, stored: u32) -> Result<(), FormatError> { if computed != stored { diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index f606656e..41157349 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -106,15 +106,37 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result { #[cfg(feature = "vector")] TYPE_VECTOR => { let metric = format::read_u8(r)?; + if metric > 2 { + return Err(FormatError::InvalidData(format!( + "unknown vector metric: {metric}" + ))); + } let quantization = format::read_u8(r)?; + if quantization > 2 { + return Err(FormatError::InvalidData(format!( + "unknown vector quantization: {quantization}" + ))); + } let connectivity = format::read_u32(r)?; let expansion_add = format::read_u32(r)?; let dim = format::read_u32(r)?; + if dim > format::MAX_PERSISTED_VECTOR_DIMS { + return Err(FormatError::InvalidData(format!( + "vector dimension {dim} exceeds max {}", + format::MAX_PERSISTED_VECTOR_DIMS + ))); + } let count = format::read_u32(r)?; + if count > format::MAX_PERSISTED_VECTOR_COUNT { + return Err(FormatError::InvalidData(format!( + "vector element count {count} exceeds max {}", + format::MAX_PERSISTED_VECTOR_COUNT + ))); + } let mut elements = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { let name = read_snap_string(r, "vector element name")?; - let mut vector = Vec::with_capacity(format::capped_capacity(dim)); + let mut vector = Vec::with_capacity(dim as usize); for _ in 0..dim { vector.push(format::read_f32(r)?); } @@ -629,16 +651,38 @@ impl SnapshotReader { #[cfg(feature = "vector")] TYPE_VECTOR => { let metric = format::read_u8(&mut self.reader)?; + if metric > 2 { + return Err(FormatError::InvalidData(format!( + "unknown vector metric: {metric}" + ))); + } format::write_u8(&mut buf, metric)?; let quantization = format::read_u8(&mut self.reader)?; + if quantization > 2 { + return Err(FormatError::InvalidData(format!( + "unknown vector quantization: {quantization}" + ))); + } format::write_u8(&mut buf, quantization)?; let connectivity = format::read_u32(&mut self.reader)?; format::write_u32(&mut buf, connectivity)?; let expansion_add = format::read_u32(&mut self.reader)?; format::write_u32(&mut buf, expansion_add)?; let dim = format::read_u32(&mut self.reader)?; + if dim > format::MAX_PERSISTED_VECTOR_DIMS { + return Err(FormatError::InvalidData(format!( + "vector dimension {dim} exceeds max {}", + format::MAX_PERSISTED_VECTOR_DIMS + ))); + } format::write_u32(&mut buf, dim)?; let count = format::read_u32(&mut self.reader)?; + if count > format::MAX_PERSISTED_VECTOR_COUNT { + return Err(FormatError::InvalidData(format!( + "vector element count {count} exceeds max {}", + format::MAX_PERSISTED_VECTOR_COUNT + ))); + } format::write_u32(&mut buf, count)?; let mut elements = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { @@ -650,7 +694,7 @@ impl SnapshotReader { "vector element name is not valid utf-8", )) })?; - let mut vector = Vec::with_capacity(format::capped_capacity(dim)); + let mut vector = Vec::with_capacity(dim as usize); for _ in 0..dim { let v = format::read_f32(&mut self.reader)?; format::write_f32(&mut buf, v)?; From 05820b44c024812154db738fbe4f8eca6a7b01fe Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 09:35:03 -0500 Subject: [PATCH 4/4] test: add vector command parser tests 21 tests covering VADD, VSIM, VREM, VGET, VCARD, VDIM, VINFO parsing including edge cases: wrong arity, exceeding limits, unknown options. --- crates/ember-protocol/src/command.rs | 231 +++++++++++++++++++++++++++ 1 file changed, 231 insertions(+) diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index eac52a65..9fa3fb6e 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -4724,4 +4724,235 @@ mod tests { Command::from_frame(cmd(&["PROTO.DELFIELD", "key", "field", "extra"])).unwrap_err(); assert!(matches!(err, ProtocolError::WrongArity(_))); } + + // --- vector commands --- + + #[test] + fn vadd_basic() { + assert_eq!( + Command::from_frame(cmd(&["VADD", "vecs", "elem1", "0.1", "0.2", "0.3"])).unwrap(), + Command::VAdd { + key: "vecs".into(), + element: "elem1".into(), + vector: vec![0.1, 0.2, 0.3], + metric: 0, + quantization: 0, + connectivity: 16, + expansion_add: 64, + }, + ); + } + + #[test] + fn vadd_with_options() { + assert_eq!( + Command::from_frame(cmd(&[ + "VADD", "vecs", "elem1", "1.0", "2.0", "METRIC", "L2", "QUANT", "F16", "M", "32", + "EF", "128" + ])) + .unwrap(), + Command::VAdd { + key: "vecs".into(), + element: "elem1".into(), + vector: vec![1.0, 2.0], + metric: 1, + quantization: 1, + connectivity: 32, + expansion_add: 128, + }, + ); + } + + #[test] + fn vadd_wrong_arity() { + // no args + let err = Command::from_frame(cmd(&["VADD"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + + // key only + let err = Command::from_frame(cmd(&["VADD", "key"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + + // key + element but no vector — "elem" is not numeric so vector is empty + let err = Command::from_frame(cmd(&["VADD", "key", "elem"])).unwrap_err(); + assert!(matches!( + err, + ProtocolError::WrongArity(_) | ProtocolError::InvalidCommandFrame(_) + )); + } + + #[test] + fn vadd_m_exceeds_max() { + let err = + Command::from_frame(cmd(&["VADD", "key", "elem", "1.0", "M", "9999"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn vadd_ef_exceeds_max() { + let err = + Command::from_frame(cmd(&["VADD", "key", "elem", "1.0", "EF", "9999"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn vadd_unknown_metric() { + let err = Command::from_frame(cmd(&["VADD", "key", "elem", "1.0", "METRIC", "HAMMING"])) + .unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn vadd_unknown_quantization() { + let err = + Command::from_frame(cmd(&["VADD", "key", "elem", "1.0", "QUANT", "F64"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn vsim_basic() { + assert_eq!( + Command::from_frame(cmd(&["VSIM", "vecs", "0.1", "0.2", "0.3", "COUNT", "5"])).unwrap(), + Command::VSim { + key: "vecs".into(), + query: vec![0.1, 0.2, 0.3], + count: 5, + ef_search: 0, + with_scores: false, + }, + ); + } + + #[test] + fn vsim_with_ef_and_scores() { + assert_eq!( + Command::from_frame(cmd(&[ + "VSIM", + "vecs", + "1.0", + "2.0", + "COUNT", + "10", + "EF", + "128", + "WITHSCORES" + ])) + .unwrap(), + Command::VSim { + key: "vecs".into(), + query: vec![1.0, 2.0], + count: 10, + ef_search: 128, + with_scores: true, + }, + ); + } + + #[test] + fn vsim_missing_count() { + // all args are numeric so they're consumed as query — COUNT is never found + let err = Command::from_frame(cmd(&["VSIM", "key", "1.0", "2.0"])).unwrap_err(); + assert!(matches!( + err, + ProtocolError::InvalidCommandFrame(_) | ProtocolError::WrongArity(_) + )); + } + + #[test] + fn vsim_wrong_arity() { + let err = Command::from_frame(cmd(&["VSIM"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn vsim_count_exceeds_max() { + let err = Command::from_frame(cmd(&["VSIM", "key", "1.0", "COUNT", "99999"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn vsim_ef_exceeds_max() { + let err = Command::from_frame(cmd(&["VSIM", "key", "1.0", "COUNT", "5", "EF", "9999"])) + .unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn vrem_basic() { + assert_eq!( + Command::from_frame(cmd(&["VREM", "key", "elem"])).unwrap(), + Command::VRem { + key: "key".into(), + element: "elem".into(), + }, + ); + } + + #[test] + fn vrem_wrong_arity() { + let err = Command::from_frame(cmd(&["VREM"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + + let err = Command::from_frame(cmd(&["VREM", "key"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn vget_basic() { + assert_eq!( + Command::from_frame(cmd(&["VGET", "key", "elem"])).unwrap(), + Command::VGet { + key: "key".into(), + element: "elem".into(), + }, + ); + } + + #[test] + fn vget_wrong_arity() { + let err = Command::from_frame(cmd(&["VGET"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn vcard_basic() { + assert_eq!( + Command::from_frame(cmd(&["VCARD", "key"])).unwrap(), + Command::VCard { key: "key".into() }, + ); + } + + #[test] + fn vcard_wrong_arity() { + let err = Command::from_frame(cmd(&["VCARD"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn vdim_basic() { + assert_eq!( + Command::from_frame(cmd(&["VDIM", "key"])).unwrap(), + Command::VDim { key: "key".into() }, + ); + } + + #[test] + fn vdim_wrong_arity() { + let err = Command::from_frame(cmd(&["VDIM"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn vinfo_basic() { + assert_eq!( + Command::from_frame(cmd(&["VINFO", "key"])).unwrap(), + Command::VInfo { key: "key".into() }, + ); + } + + #[test] + fn vinfo_wrong_arity() { + let err = Command::from_frame(cmd(&["VINFO"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } }