diff --git a/crates/ember-cli/src/format.rs b/crates/ember-cli/src/format.rs index 2b23ce5c..ea64da13 100644 --- a/crates/ember-cli/src/format.rs +++ b/crates/ember-cli/src/format.rs @@ -20,13 +20,40 @@ pub fn format_response(frame: &Frame) -> String { format_frame(frame, 0) } +/// Strips ANSI escape sequences and other control characters from +/// server-supplied strings to prevent terminal manipulation attacks. +/// Retains printable ASCII, tabs, and newlines (CR/LF). +fn sanitize(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + let mut chars = s.chars(); + while let Some(ch) = chars.next() { + if ch == '\x1b' { + // skip the ESC and the rest of the ANSI sequence + if let Some(next) = chars.next() { + if next == '[' { + // CSI sequence — consume until a letter + for c in chars.by_ref() { + if c.is_ascii_alphabetic() { + break; + } + } + } + // else: single-char escape, already consumed + } + } else if ch == '\t' || ch == '\n' || ch == '\r' || !ch.is_control() { + out.push(ch); + } + } + out +} + fn format_frame(frame: &Frame, indent: usize) -> String { let prefix = " ".repeat(indent); match frame { - Frame::Simple(s) => format!("{prefix}{}", s.green()), + Frame::Simple(s) => format!("{prefix}{}", sanitize(s).green()), - Frame::Error(e) => format!("{prefix}{} {}", "(error)".red(), e.red()), + Frame::Error(e) => format!("{prefix}{} {}", "(error)".red(), sanitize(e).red()), Frame::Integer(n) => format!( "{prefix}{} {}", @@ -38,9 +65,9 @@ fn format_frame(frame: &Frame, indent: usize) -> String { match std::str::from_utf8(data) { Ok(s) if s.contains("\r\n") || s.contains('\n') => { // multiline output (like INFO) — print unquoted - format!("{prefix}{}", s.green()) + format!("{prefix}{}", sanitize(s).green()) } - Ok(s) => format!("{prefix}{}", format!("\"{}\"", s).green()), + Ok(s) => format!("{prefix}{}", format!("\"{}\"", sanitize(s)).green()), Err(_) => { // binary data — show as hex let hex: String = data.iter().map(|b| format!("{b:02x}")).collect(); @@ -187,4 +214,25 @@ mod tests { }); assert_eq!(out, "1) key => (integer) 1"); } + + #[test] + fn sanitize_strips_ansi_escapes() { + assert_eq!(sanitize("hello\x1b[31mworld\x1b[0m"), "helloworld"); + } + + #[test] + fn sanitize_strips_control_chars() { + assert_eq!(sanitize("hello\x07\x08world"), "helloworld"); + } + + #[test] + fn sanitize_preserves_tabs_and_newlines() { + assert_eq!(sanitize("line1\nline2\ttab"), "line1\nline2\ttab"); + } + + #[test] + fn sanitize_server_response_with_escape() { + let out = no_color(|| format_response(&Frame::Simple("\x1b[31mfake-error\x1b[0m".into()))); + assert_eq!(out, "fake-error"); + } } diff --git a/crates/ember-cli/src/repl.rs b/crates/ember-cli/src/repl.rs index 1b6f0f4c..01edc786 100644 --- a/crates/ember-cli/src/repl.rs +++ b/crates/ember-cli/src/repl.rs @@ -87,7 +87,14 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC continue; } - let _ = rl.add_history_entry(trimmed); + // don't save AUTH commands to history — they contain passwords + if !trimmed + .split_whitespace() + .next() + .is_some_and(|w| w.eq_ignore_ascii_case("auth")) + { + let _ = rl.add_history_entry(trimmed); + } // handle local commands let first_word = trimmed.split_whitespace().next().unwrap_or(""); @@ -157,6 +164,13 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC if let Some(ref path) = history_path { let _ = rl.save_history(path); + // restrict history file permissions — it may contain key names + // and values from previous sessions + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)); + } } // graceful shutdown — send QUIT and close the TCP stream diff --git a/crates/ember-cluster/src/gossip.rs b/crates/ember-cluster/src/gossip.rs index 859ad71b..827f3fcf 100644 --- a/crates/ember-cluster/src/gossip.rs +++ b/crates/ember-cluster/src/gossip.rs @@ -23,6 +23,11 @@ use tracing::{debug, info, trace, warn}; use crate::message::{GossipMessage, MemberInfo, NodeUpdate}; use crate::{NodeId, SlotRange}; +/// Maximum allowed incarnation value. Rejects gossip updates with +/// incarnation numbers beyond this to prevent a malicious node from +/// sending u64::MAX and permanently disabling suspicion refutation. +const MAX_INCARNATION: u64 = u64::MAX / 2; + /// Configuration for the gossip protocol. #[derive(Debug, Clone)] pub struct GossipConfig { @@ -377,6 +382,13 @@ impl GossipEngine { addr, incarnation, } => { + if *incarnation > MAX_INCARNATION { + warn!( + "rejecting alive update for {} with excessive incarnation {}", + node, incarnation + ); + continue; + } if *node == self.local_id { // Someone thinks we're alive, good continue; @@ -419,6 +431,13 @@ impl GossipEngine { } NodeUpdate::Suspect { node, incarnation } => { + if *incarnation > MAX_INCARNATION { + warn!( + "rejecting suspect update for {} with excessive incarnation {}", + node, incarnation + ); + continue; + } if *node == self.local_id { // Refute suspicion by incrementing our incarnation if *incarnation >= self.incarnation { @@ -445,6 +464,13 @@ impl GossipEngine { } NodeUpdate::Dead { node, incarnation } => { + if *incarnation > MAX_INCARNATION { + warn!( + "rejecting dead update for {} with excessive incarnation {}", + node, incarnation + ); + continue; + } if *node == self.local_id { // Refute death claim self.incarnation = incarnation.saturating_add(1); diff --git a/crates/ember-cluster/src/message.rs b/crates/ember-cluster/src/message.rs index fcd0cc76..f121e488 100644 --- a/crates/ember-cluster/src/message.rs +++ b/crates/ember-cluster/src/message.rs @@ -170,8 +170,9 @@ impl GossipMessage { GossipMessage::Welcome { sender, members } => { buf.put_u8(MSG_WELCOME); encode_node_id(buf, sender); - buf.put_u16_le(members.len() as u16); - for member in members { + let count = members.len().min(MAX_COLLECTION_COUNT); + buf.put_u16_le(count as u16); + for member in &members[..count] { encode_member_info(buf, member); } } @@ -324,8 +325,9 @@ fn decode_socket_addr(buf: &mut &[u8]) -> io::Result { } fn encode_updates(buf: &mut BytesMut, updates: &[NodeUpdate]) { - buf.put_u16_le(updates.len() as u16); - for update in updates { + let count = updates.len().min(MAX_COLLECTION_COUNT); + buf.put_u16_le(count as u16); + for update in &updates[..count] { encode_update(buf, update); } } @@ -413,8 +415,9 @@ fn encode_member_info(buf: &mut BytesMut, member: &MemberInfo) { encode_socket_addr(buf, &member.addr); buf.put_u64_le(member.incarnation); buf.put_u8(if member.is_primary { 1 } else { 0 }); - buf.put_u16_le(member.slots.len() as u16); - for slot in &member.slots { + let slot_count = member.slots.len().min(MAX_COLLECTION_COUNT); + buf.put_u16_le(slot_count as u16); + for slot in &member.slots[..slot_count] { buf.put_u16_le(slot.start); buf.put_u16_le(slot.end); } diff --git a/crates/ember-cluster/src/raft.rs b/crates/ember-cluster/src/raft.rs index 846bb193..1dc6c215 100644 --- a/crates/ember-cluster/src/raft.rs +++ b/crates/ember-cluster/src/raft.rs @@ -18,6 +18,7 @@ use openraft::{ use serde::{Deserialize, Serialize}; use tokio::sync::RwLock; +use crate::slots::SLOT_COUNT; use crate::{NodeId, SlotRange}; /// Type configuration for openraft. @@ -178,6 +179,17 @@ impl Storage { } ClusterCommand::AssignSlots { node_id, slots } => { + // validate all slot ranges before applying + for range in slots { + if range.start > range.end || range.end >= SLOT_COUNT { + return ClusterResponse::Error(format!( + "invalid slot range {}..={} (max {})", + range.start, + range.end, + SLOT_COUNT - 1 + )); + } + } let key = node_id.0.to_string(); if let Some(node) = state.nodes.get_mut(&key) { node.slots = slots.clone(); @@ -203,6 +215,12 @@ impl Storage { } ClusterCommand::BeginMigration { slot, from, to } => { + if *slot >= SLOT_COUNT { + return ClusterResponse::Error(format!( + "slot {slot} out of range (max {})", + SLOT_COUNT - 1 + )); + } state.migrations.insert( *slot, MigrationState { @@ -214,6 +232,11 @@ impl Storage { } ClusterCommand::CompleteMigration { slot, new_owner } => { + if !state.migrations.contains_key(slot) { + return ClusterResponse::Error(format!( + "no migration in progress for slot {slot}" + )); + } state.migrations.remove(slot); let key = new_owner.0.to_string(); state.slots.insert(*slot, key); @@ -564,6 +587,112 @@ mod tests { } } + #[tokio::test] + async fn assign_slots_rejects_invalid_range() { + let storage = Arc::new(Storage::new()); + let mut s = Arc::clone(&storage); + + let node_id = NodeId::new(); + let add = Entry { + log_id: log_id(1, 1), + payload: EntryPayload::Normal(ClusterCommand::AddNode { + node_id, + raft_id: 1, + addr: "127.0.0.1:6379".into(), + is_primary: true, + }), + }; + s.apply_to_state_machine(&[add]).await.unwrap(); + + // craft a SlotRange with start > end (bypassing SlotRange::new) + let bad_range = SlotRange { + start: 100, + end: 50, + }; + let assign = Entry { + log_id: log_id(1, 2), + payload: EntryPayload::Normal(ClusterCommand::AssignSlots { + node_id, + slots: vec![bad_range], + }), + }; + let results = s.apply_to_state_machine(&[assign]).await.unwrap(); + assert!( + matches!(&results[0], ClusterResponse::Error(msg) if msg.contains("invalid slot range")) + ); + } + + #[tokio::test] + async fn assign_slots_rejects_out_of_range() { + let storage = Arc::new(Storage::new()); + let mut s = Arc::clone(&storage); + + let node_id = NodeId::new(); + let add = Entry { + log_id: log_id(1, 1), + payload: EntryPayload::Normal(ClusterCommand::AddNode { + node_id, + raft_id: 1, + addr: "127.0.0.1:6379".into(), + is_primary: true, + }), + }; + s.apply_to_state_machine(&[add]).await.unwrap(); + + // slot end >= SLOT_COUNT + let bad_range = SlotRange { + start: 0, + end: 16384, + }; + let assign = Entry { + log_id: log_id(1, 2), + payload: EntryPayload::Normal(ClusterCommand::AssignSlots { + node_id, + slots: vec![bad_range], + }), + }; + let results = s.apply_to_state_machine(&[assign]).await.unwrap(); + assert!( + matches!(&results[0], ClusterResponse::Error(msg) if msg.contains("invalid slot range")) + ); + } + + #[tokio::test] + async fn complete_migration_without_begin_errors() { + let storage = Arc::new(Storage::new()); + let mut s = Arc::clone(&storage); + + let node_id = NodeId::new(); + let complete = Entry { + log_id: log_id(1, 1), + payload: EntryPayload::Normal(ClusterCommand::CompleteMigration { + slot: 100, + new_owner: node_id, + }), + }; + let results = s.apply_to_state_machine(&[complete]).await.unwrap(); + assert!(matches!(&results[0], ClusterResponse::Error(msg) if msg.contains("no migration"))); + } + + #[tokio::test] + async fn begin_migration_rejects_invalid_slot() { + let storage = Arc::new(Storage::new()); + let mut s = Arc::clone(&storage); + + let node1 = NodeId::new(); + let node2 = NodeId::new(); + let begin = Entry { + log_id: log_id(1, 1), + payload: EntryPayload::Normal(ClusterCommand::BeginMigration { + slot: 16384, + from: node1, + to: node2, + }), + }; + let results = s.apply_to_state_machine(&[begin]).await.unwrap(); + assert!(matches!(&results[0], ClusterResponse::Error(msg) if msg.contains("out of range"))); + } + #[tokio::test] async fn storage_log_operations() { let storage = Arc::new(Storage::new()); diff --git a/crates/ember-cluster/src/slots.rs b/crates/ember-cluster/src/slots.rs index 6b4f995f..34748cd1 100644 --- a/crates/ember-cluster/src/slots.rs +++ b/crates/ember-cluster/src/slots.rs @@ -114,10 +114,10 @@ impl SlotRange { /// /// # Panics /// - /// Debug-panics if `start > end` or if `end >= SLOT_COUNT`. + /// Panics if `start > end` or if `end >= SLOT_COUNT`. pub fn new(start: u16, end: u16) -> Self { - debug_assert!(start <= end, "SlotRange requires start <= end"); - debug_assert!(end < SLOT_COUNT, "slot must be < {SLOT_COUNT}"); + assert!(start <= end, "SlotRange requires start <= end"); + assert!(end < SLOT_COUNT, "slot must be < {SLOT_COUNT}"); Self { start, end } } diff --git a/crates/ember-core/src/engine.rs b/crates/ember-core/src/engine.rs index a550cb86..cdadf8ef 100644 --- a/crates/ember-core/src/engine.rs +++ b/crates/ember-core/src/engine.rs @@ -56,6 +56,10 @@ impl Engine { /// Panics if `shard_count` is zero. pub fn with_config(shard_count: usize, config: EngineConfig) -> Self { assert!(shard_count > 0, "shard count must be at least 1"); + assert!( + shard_count <= u16::MAX as usize, + "shard count must fit in u16" + ); let drop_handle = DropHandle::spawn(); diff --git a/crates/ember-core/src/keyspace.rs b/crates/ember-core/src/keyspace.rs index cbb6e910..a2da8eca 100644 --- a/crates/ember-core/src/keyspace.rs +++ b/crates/ember-core/src/keyspace.rs @@ -157,7 +157,7 @@ pub struct VAddBatchResult { /// Errors from vector write operations. #[cfg(feature = "vector")] -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone)] pub enum VectorWriteError { /// The key holds a different type than expected. WrongType, @@ -165,6 +165,12 @@ pub enum VectorWriteError { OutOfMemory, /// usearch index error (dimension mismatch, capacity, etc). IndexError(String), + /// A batch insert partially succeeded before encountering an error. + /// The applied vectors should still be persisted to the AOF. + PartialBatch { + message: String, + applied: Vec<(String, Vec)>, + }, } /// How the keyspace should handle writes when the memory limit is reached. @@ -396,20 +402,14 @@ impl Keyspace { /// Measures entry size before and after a mutation, adjusting the /// memory tracker for the difference. Touches the entry afterwards. - fn track_size(&mut self, key: &str, f: impl FnOnce(&mut Entry) -> T) -> T { - let entry = self - .entries - .get_mut(key) - .expect("caller verified key exists"); + fn track_size(&mut self, key: &str, f: impl FnOnce(&mut Entry) -> T) -> Option { + let entry = self.entries.get_mut(key)?; let old_size = memory::entry_size(key, &entry.value); let result = f(entry); - let entry = self - .entries - .get(key) - .expect("mutation should not remove key"); + let entry = self.entries.get(key)?; let new_size = memory::entry_size(key, &entry.value); self.memory.adjust(old_size, new_size); - result + Some(result) } /// Adjusts the expiry count when replacing an entry whose TTL status @@ -1042,7 +1042,7 @@ impl Keyspace { return None; } let ttl_ms = match time::remaining_ms(entry.expires_at_ms) { - Some(ms) => ms as i64, + Some(ms) => ms.min(i64::MAX as u64) as i64, None => -1, }; Some((key.as_str(), &entry.value, ttl_ms)) @@ -1183,21 +1183,23 @@ impl Keyspace { self.insert_empty(key, Value::List(VecDeque::new())); } - let len = self.track_size(key, |entry| { - let Value::List(ref mut deque) = entry.value else { - unreachable!("type verified by ensure_collection_type"); - }; - for val in values { - if left { - deque.push_front(val.clone()); - } else { - deque.push_back(val.clone()); + let len = self + .track_size(key, |entry| { + let Value::List(ref mut deque) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + for val in values { + if left { + deque.push_front(val.clone()); + } else { + deque.push_back(val.clone()); + } } - } - let len = deque.len(); - entry.touch(); - len - }); + let len = deque.len(); + entry.touch(); + len + }) + .unwrap_or(0); Ok(len) } @@ -1263,28 +1265,30 @@ impl Keyspace { self.insert_empty(key, Value::SortedSet(SortedSet::new())); } - let (count, applied) = self.track_size(key, |entry| { - let Value::SortedSet(ref mut ss) = entry.value else { - unreachable!("type verified by ensure_collection_type"); - }; - let mut count = 0; - let mut applied = Vec::new(); - for (score, member) in members { - let result = ss.add_with_flags(member.clone(), *score, flags); - if result.added || result.updated { - applied.push((*score, member.clone())); - } - if flags.ch { + let (count, applied) = self + .track_size(key, |entry| { + let Value::SortedSet(ref mut ss) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + let mut count = 0; + let mut applied = Vec::new(); + for (score, member) in members { + let result = ss.add_with_flags(member.clone(), *score, flags); if result.added || result.updated { + applied.push((*score, member.clone())); + } + if flags.ch { + if result.added || result.updated { + count += 1; + } + } else if result.added { count += 1; } - } else if result.added { - count += 1; } - } - entry.touch(); - (count, applied) - }); + entry.touch(); + (count, applied) + }) + .unwrap_or_default(); // clean up if the set is still empty (e.g. XX flag on a new key) if let Some(entry) = self.entries.get(key) { @@ -1452,19 +1456,21 @@ impl Keyspace { self.insert_empty(key, Value::Hash(HashMap::new())); } - let added = self.track_size(key, |entry| { - let Value::Hash(ref mut map) = entry.value else { - unreachable!("type verified by ensure_collection_type"); - }; - let mut added = 0; - for (field, value) in fields { - if map.insert(field.clone(), value.clone()).is_none() { - added += 1; + let added = self + .track_size(key, |entry| { + let Value::Hash(ref mut map) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + let mut added = 0; + for (field, value) in fields { + if map.insert(field.clone(), value.clone()).is_none() { + added += 1; + } } - } - entry.touch(); - added - }); + entry.touch(); + added + }) + .unwrap_or(0); Ok(added) } @@ -1718,19 +1724,21 @@ impl Keyspace { self.insert_empty(key, Value::Set(std::collections::HashSet::new())); } - let added = self.track_size(key, |entry| { - let Value::Set(ref mut set) = entry.value else { - unreachable!("type verified by ensure_collection_type"); - }; - let mut added = 0; - for member in members { - if set.insert(member.clone()) { - added += 1; + let added = self + .track_size(key, |entry| { + let Value::Set(ref mut set) = entry.value else { + unreachable!("type verified by ensure_collection_type"); + }; + let mut added = 0; + for member in members { + if set.insert(member.clone()) { + added += 1; + } } - } - entry.touch(); - added - }); + entry.touch(); + added + }) + .unwrap_or(0); Ok(added) } @@ -2048,16 +2056,19 @@ impl Keyspace { applied.push((element.clone(), vector.clone())); } Err(e) => { - // partial insert: return what we applied so far + error - // caller should persist the applied vectors + // partial insert: return applied vectors so they can + // be persisted to AOF despite the error entry.touch(); let new_entry_size = memory::entry_size(key, &entry.value); self.memory.adjust(old_entry_size, new_entry_size); - return Err(VectorWriteError::IndexError(format!( - "error at element '{}': {e} ({} vectors applied before failure)", - element, - applied.len() - ))); + return Err(VectorWriteError::PartialBatch { + message: format!( + "error at element '{}': {e} ({} vectors applied before failure)", + element, + applied.len() + ), + applied, + }); } } } diff --git a/crates/ember-core/src/memory.rs b/crates/ember-core/src/memory.rs index 683c5e08..aa98f524 100644 --- a/crates/ember-core/src/memory.rs +++ b/crates/ember-core/src/memory.rs @@ -38,7 +38,9 @@ pub const MEMORY_SAFETY_MARGIN_PERCENT: usize = 90; /// Returns the number of bytes at which writes should be rejected or /// eviction should begin — always less than the raw configured limit. pub fn effective_limit(max_bytes: usize) -> usize { - max_bytes * MEMORY_SAFETY_MARGIN_PERCENT / 100 + // use u128 intermediate to avoid overflow on large max_bytes values + // while preserving precision for small values + ((max_bytes as u128) * (MEMORY_SAFETY_MARGIN_PERCENT as u128) / 100) as usize } /// Estimated overhead per entry in the HashMap. diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 142cd655..897faa53 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -1073,7 +1073,8 @@ fn dispatch( }, Err(crate::keyspace::VectorWriteError::WrongType) => ShardResponse::WrongType, Err(crate::keyspace::VectorWriteError::OutOfMemory) => ShardResponse::OutOfMemory, - Err(crate::keyspace::VectorWriteError::IndexError(e)) => { + Err(crate::keyspace::VectorWriteError::IndexError(e)) + | Err(crate::keyspace::VectorWriteError::PartialBatch { message: e, .. }) => { ShardResponse::Err(format!("ERR vector index: {e}")) } } @@ -1106,6 +1107,13 @@ fn dispatch( Err(crate::keyspace::VectorWriteError::IndexError(e)) => { ShardResponse::Err(format!("ERR vector index: {e}")) } + Err(crate::keyspace::VectorWriteError::PartialBatch { applied, .. }) => { + // partial success: return applied vectors for AOF persistence + ShardResponse::VAddBatchResult { + added_count: applied.len(), + applied, + } + } } } #[cfg(feature = "vector")] diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 9c67d23a..6cf14e1c 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -790,12 +790,16 @@ impl AofWriter { } /// Truncates the AOF file back to just the header. - /// Used after a successful snapshot to reset the log. + /// + /// Uses write-to-temp-then-rename for crash safety: the old AOF + /// remains intact until the new file (with only a header) is fully + /// synced and atomically renamed into place. pub fn truncate(&mut self) -> Result<(), FormatError> { - // flush and drop the old writer + // flush the old writer so no data is in the BufWriter self.writer.flush()?; - // reopen the file with truncation, write fresh header + // write a fresh header to a temp file next to the real AOF + let tmp_path = self.path.with_extension("aof.tmp"); let mut opts = OpenOptions::new(); opts.create(true).write(true).truncate(true); #[cfg(unix)] @@ -803,26 +807,31 @@ impl AofWriter { use std::os::unix::fs::OpenOptionsExt; opts.mode(0o600); } - let file = opts.open(&self.path)?; - let mut writer = BufWriter::new(file); + let tmp_file = opts.open(&tmp_path)?; + let mut tmp_writer = BufWriter::new(tmp_file); #[cfg(feature = "encryption")] if self.encryption_key.is_some() { format::write_header_versioned( - &mut writer, + &mut tmp_writer, format::AOF_MAGIC, format::FORMAT_VERSION_ENCRYPTED, )?; } else { - format::write_header(&mut writer, format::AOF_MAGIC)?; + format::write_header(&mut tmp_writer, format::AOF_MAGIC)?; } #[cfg(not(feature = "encryption"))] - format::write_header(&mut writer, format::AOF_MAGIC)?; + format::write_header(&mut tmp_writer, format::AOF_MAGIC)?; + + tmp_writer.flush()?; + tmp_writer.get_ref().sync_all()?; + + // atomic rename: old AOF is replaced only after new file is durable + std::fs::rename(&tmp_path, &self.path)?; - writer.flush()?; - // ensure the fresh header is durable before we start appending - writer.get_ref().sync_all()?; - self.writer = writer; + // reopen for appending + let file = OpenOptions::new().append(true).open(&self.path)?; + self.writer = BufWriter::new(file); Ok(()) } } @@ -1102,6 +1111,12 @@ impl AofReader { let element = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &element)?; let dim = format::read_u32(&mut self.reader)?; + if dim > format::MAX_PERSISTED_VECTOR_DIMS { + return Err(FormatError::InvalidData(format!( + "AOF VADD dimension {dim} exceeds max {}", + format::MAX_PERSISTED_VECTOR_DIMS + ))); + } format::write_u32(&mut payload, dim)?; for _ in 0..dim { let v = format::read_f32(&mut self.reader)?; diff --git a/crates/ember-persistence/src/encryption.rs b/crates/ember-persistence/src/encryption.rs index ff6e895c..afcb7ac9 100644 --- a/crates/ember-persistence/src/encryption.rs +++ b/crates/ember-persistence/src/encryption.rs @@ -25,11 +25,25 @@ pub const TAG_SIZE: usize = 16; /// /// The key is stored inline — no heap allocation. Implements `Clone` /// but not `Debug` to avoid accidentally logging key material. +/// On drop, the key bytes are zeroed to prevent key material from +/// lingering in freed memory. #[derive(Clone)] pub struct EncryptionKey { bytes: [u8; 32], } +impl Drop for EncryptionKey { + fn drop(&mut self) { + // Use a volatile write to prevent the compiler from optimizing + // away the zeroing of key material. + for byte in &mut self.bytes { + // SAFETY: this is a regular mutable reference write, using + // write_volatile to prevent dead-store elimination. + unsafe { std::ptr::write_volatile(byte, 0) }; + } + } +} + impl fmt::Debug for EncryptionKey { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("EncryptionKey") diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index baf0890d..c2ecf909 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -144,7 +144,7 @@ fn recover_shard_impl( // step 1: load snapshot let snap_path = snapshot::snapshot_path(data_dir, shard_id); if snap_path.exists() { - match load_snapshot(&snap_path, encryption_key) { + match load_snapshot(&snap_path, shard_id, encryption_key) { Ok(entries) => { for (key, value, ttl_ms) in entries { map.insert(key, (RecoveredValue::from(value), ttl_ms)); @@ -209,6 +209,7 @@ fn recover_shard_impl( /// Returns (key, value, ttl_ms) where ttl_ms is -1 for no expiry. fn load_snapshot( path: &Path, + expected_shard_id: u16, #[allow(unused_variables)] encryption_key: Option>, ) -> Result, FormatError> { #[cfg(feature = "encryption")] @@ -220,6 +221,13 @@ fn load_snapshot( #[cfg(not(feature = "encryption"))] let mut reader = SnapshotReader::open(path)?; + if reader.shard_id != expected_shard_id { + return Err(FormatError::InvalidData(format!( + "snapshot shard_id {} does not match expected {}", + reader.shard_id, expected_shard_id + ))); + } + let mut entries = Vec::new(); while let Some(entry) = reader.read_entry()? { diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index bb10b0d4..3238e164 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -31,6 +31,10 @@ const MAX_VSIM_EF: u64 = MAX_HNSW_PARAM; /// round-trip overhead for bulk inserts. const MAX_VADD_BATCH_SIZE: usize = 10_000; +/// Maximum value for SCAN COUNT. Prevents clients from requesting a scan +/// hint so large it causes pre-allocation issues. +const MAX_SCAN_COUNT: u64 = 10_000_000; + /// Expiration option for the SET command. #[derive(Debug, Clone, PartialEq, Eq)] pub enum SetExpire { @@ -1120,6 +1124,11 @@ fn parse_scan(args: &[Frame]) -> Result { return Err(ProtocolError::WrongArity("SCAN".into())); } let n = parse_u64(&args[idx], "SCAN")?; + if n > MAX_SCAN_COUNT { + return Err(ProtocolError::InvalidCommandFrame(format!( + "SCAN COUNT {n} exceeds max {MAX_SCAN_COUNT}" + ))); + } count = Some(n as usize); idx += 1; } @@ -1869,6 +1878,11 @@ fn parse_vadd(args: &[Frame]) -> Result { } let s = extract_string(&args[idx])?; if let Ok(v) = s.parse::() { + if v.is_nan() || v.is_infinite() { + return Err(ProtocolError::InvalidCommandFrame( + "VADD: vector components must be finite (no NaN/infinity)".into(), + )); + } vector.push(v); idx += 1; } else { @@ -2013,43 +2027,49 @@ fn parse_vadd_batch(args: &[Frame]) -> Result { } // parse entries: each is element_name followed by exactly `dim` floats. - // stop when we see a known flag or run out of args. + // we detect the end of entries by checking whether enough args remain + // for a full entry (1 name + dim floats). this avoids misinterpreting + // element names like "metric" as flags. let mut idx = 3; let mut entries: Vec<(String, Vec)> = Vec::new(); - let flags = ["METRIC", "QUANT", "M", "EF"]; + let entry_len = 1 + dim; // element name + dim floats while idx < args.len() { - let token = extract_string(&args[idx])?; - if flags.contains(&token.to_ascii_uppercase().as_str()) { + // not enough remaining args for a full entry — must be flags + if idx + entry_len > args.len() { break; } - // this token is an element name - let element = token; - idx += 1; - - // read exactly `dim` floats - if idx + dim > args.len() { - return Err(ProtocolError::InvalidCommandFrame(format!( - "VADD_BATCH: not enough floats for element '{element}' (expected {dim})" - ))); + // peek: if the token after the element name isn't a valid float, + // we've reached the flags section + if dim > 0 { + let peek = extract_string(&args[idx + 1])?; + if peek.parse::().is_err() { + break; + } } + let element = extract_string(&args[idx])?; + idx += 1; + let mut vector = Vec::with_capacity(dim); for _ in 0..dim { let s = extract_string(&args[idx])?; let v = s.parse::().map_err(|_| { - ProtocolError::InvalidCommandFrame(format!( - "VADD_BATCH: expected float, got '{s}'" - )) + ProtocolError::InvalidCommandFrame(format!("VADD_BATCH: expected float, got '{s}'")) })?; + if v.is_nan() || v.is_infinite() { + return Err(ProtocolError::InvalidCommandFrame( + "VADD_BATCH: vector components must be finite (no NaN/infinity)".into(), + )); + } vector.push(v); idx += 1; } entries.push((element, vector)); - if entries.len() > MAX_VADD_BATCH_SIZE { + if entries.len() >= MAX_VADD_BATCH_SIZE { return Err(ProtocolError::InvalidCommandFrame(format!( "VADD_BATCH: batch size exceeds max {MAX_VADD_BATCH_SIZE}" ))); @@ -2176,6 +2196,11 @@ fn parse_vsim(args: &[Frame]) -> Result { } let s = extract_string(&args[idx])?; if let Ok(v) = s.parse::() { + if v.is_nan() || v.is_infinite() { + return Err(ProtocolError::InvalidCommandFrame( + "VSIM: query components must be finite (no NaN/infinity)".into(), + )); + } query.push(v); idx += 1; } else { @@ -5018,7 +5043,17 @@ mod tests { fn vadd_batch_basic() { assert_eq!( Command::from_frame(cmd(&[ - "VADD_BATCH", "vecs", "DIM", "3", "a", "0.1", "0.2", "0.3", "b", "0.4", "0.5", + "VADD_BATCH", + "vecs", + "DIM", + "3", + "a", + "0.1", + "0.2", + "0.3", + "b", + "0.4", + "0.5", "0.6" ])) .unwrap(), @@ -5041,8 +5076,21 @@ mod tests { fn vadd_batch_with_options() { assert_eq!( Command::from_frame(cmd(&[ - "VADD_BATCH", "vecs", "DIM", "2", "a", "1.0", "2.0", "METRIC", "L2", "QUANT", - "F16", "M", "32", "EF", "128" + "VADD_BATCH", + "vecs", + "DIM", + "2", + "a", + "1.0", + "2.0", + "METRIC", + "L2", + "QUANT", + "F16", + "M", + "32", + "EF", + "128" ])) .unwrap(), Command::VAddBatch { @@ -5104,39 +5152,42 @@ mod tests { #[test] fn vadd_batch_missing_dim_keyword() { - let err = - Command::from_frame(cmd(&["VADD_BATCH", "key", "3", "a", "1.0"])).unwrap_err(); + let err = Command::from_frame(cmd(&["VADD_BATCH", "key", "3", "a", "1.0"])).unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); } #[test] fn vadd_batch_dim_zero() { - let err = - Command::from_frame(cmd(&["VADD_BATCH", "key", "DIM", "0"])).unwrap_err(); + let err = Command::from_frame(cmd(&["VADD_BATCH", "key", "DIM", "0"])).unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); } #[test] fn vadd_batch_dim_exceeds_max() { - let err = - Command::from_frame(cmd(&["VADD_BATCH", "key", "DIM", "99999"])).unwrap_err(); + let err = Command::from_frame(cmd(&["VADD_BATCH", "key", "DIM", "99999"])).unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); } #[test] fn vadd_batch_insufficient_floats() { // DIM=3 but only 2 floats for element "a" - let err = Command::from_frame(cmd(&[ - "VADD_BATCH", "key", "DIM", "3", "a", "1.0", "2.0", - ])) - .unwrap_err(); + let err = Command::from_frame(cmd(&["VADD_BATCH", "key", "DIM", "3", "a", "1.0", "2.0"])) + .unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); } #[test] fn vadd_batch_m_exceeds_max() { let err = Command::from_frame(cmd(&[ - "VADD_BATCH", "key", "DIM", "2", "a", "1.0", "2.0", "M", "9999", + "VADD_BATCH", + "key", + "DIM", + "2", + "a", + "1.0", + "2.0", + "M", + "9999", ])) .unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); @@ -5145,7 +5196,15 @@ mod tests { #[test] fn vadd_batch_ef_exceeds_max() { let err = Command::from_frame(cmd(&[ - "VADD_BATCH", "key", "DIM", "2", "a", "1.0", "2.0", "EF", "9999", + "VADD_BATCH", + "key", + "DIM", + "2", + "a", + "1.0", + "2.0", + "EF", + "9999", ])) .unwrap_err(); assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); diff --git a/crates/ember-protocol/src/parse.rs b/crates/ember-protocol/src/parse.rs index d61bfcfe..4f141191 100644 --- a/crates/ember-protocol/src/parse.rs +++ b/crates/ember-protocol/src/parse.rs @@ -34,6 +34,12 @@ const MAX_ARRAY_ELEMENTS: usize = 1_048_576; /// Maximum length of a bulk string in bytes (512 MB, matching Redis). const MAX_BULK_LEN: i64 = 512 * 1024 * 1024; +/// Cap for Vec::with_capacity in array/map parsing. A declared count of +/// 1M elements with capacity pre-allocation costs ~72 MB upfront even +/// before any child data is parsed. This cap limits the initial allocation +/// while still letting the Vec grow organically as elements are parsed. +const PREALLOC_CAP: usize = 1024; + /// Checks whether `buf` contains a complete RESP3 frame and parses it. /// /// Returns `Ok(Some(frame))` if a complete frame was parsed, @@ -128,7 +134,7 @@ fn try_parse(cursor: &mut Cursor<&[u8]>, depth: usize) -> Result, depth: usize) -> Result { + if depth >= MAX_SERIALIZE_DEPTH { + // emit an error frame instead of overflowing the stack + dst.put_slice(b"-ERR nesting too deep\r\n"); + return; + } dst.put_u8(b'*'); write_i64(items.len() as i64, dst); dst.put_slice(b"\r\n"); for item in items { - item.serialize(dst); + item.serialize_inner(dst, depth + 1); } } Frame::Null => { dst.put_slice(wire::NULL); } Frame::Map(pairs) => { + if depth >= MAX_SERIALIZE_DEPTH { + dst.put_slice(b"-ERR nesting too deep\r\n"); + return; + } dst.put_u8(b'%'); write_i64(pairs.len() as i64, dst); dst.put_slice(b"\r\n"); for (key, val) in pairs { - key.serialize(dst); - val.serialize(dst); + key.serialize_inner(dst, depth + 1); + val.serialize_inner(dst, depth + 1); } } } diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index e9921416..97b7d839 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -25,7 +25,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use crate::connection_common::{ is_allowed_before_auth, is_auth_frame, try_auth, BUF_CAPACITY, IDLE_TIMEOUT, MAX_AUTH_FAILURES, - MAX_BUF_SIZE, + MAX_BUF_SIZE, MAX_PIPELINE_DEPTH, }; use crate::pubsub::PubSubManager; use crate::server::ServerContext; @@ -69,10 +69,15 @@ where } out.clear(); + let mut pipeline_count: usize = 0; loop { + if pipeline_count >= MAX_PIPELINE_DEPTH { + break; // process this batch, remaining data stays in buf + } match parse_frame(&buf) { Ok(Some((frame, consumed))) => { let _ = buf.split_to(consumed); + pipeline_count += 1; if !authenticated { if is_auth_frame(&frame) { diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 617cc397..d19bf8e6 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -20,7 +20,7 @@ use tokio::sync::{broadcast, oneshot}; use crate::connection_common::{ is_allowed_before_auth, is_auth_frame, try_auth, BUF_CAPACITY, IDLE_TIMEOUT, MAX_AUTH_FAILURES, - MAX_BUF_SIZE, MAX_PATTERN_LEN, MAX_SUBSCRIPTIONS_PER_CONN, + MAX_BUF_SIZE, MAX_PATTERN_LEN, MAX_PIPELINE_DEPTH, MAX_SUBSCRIPTIONS_PER_CONN, }; use crate::pubsub::{PubMessage, PubSubManager}; use crate::server::ServerContext; @@ -200,6 +200,9 @@ where Ok(Some((frame, consumed))) => { let _ = buf.split_to(consumed); frames.push(frame); + if frames.len() >= MAX_PIPELINE_DEPTH { + break; // process this batch, remaining data stays in buf + } } Ok(None) => break, // need more data Err(e) => { diff --git a/crates/ember-server/src/connection_common.rs b/crates/ember-server/src/connection_common.rs index 31f9e3d1..8ab3427a 100644 --- a/crates/ember-server/src/connection_common.rs +++ b/crates/ember-server/src/connection_common.rs @@ -41,6 +41,12 @@ pub const MAX_SUBSCRIPTIONS_PER_CONN: usize = 10_000; /// reasonable channel naming scheme. pub const MAX_PATTERN_LEN: usize = 256; +/// Maximum number of commands parsed from a single read buffer before +/// flushing responses. Prevents a huge pipeline from consuming unbounded +/// memory for pending responses. 10,000 is generous for legitimate +/// pipelining (Redis typically handles hundreds at a time). +pub const MAX_PIPELINE_DEPTH: usize = 10_000; + /// Checks if a raw frame is an AUTH command (before full parsing). /// /// Peeks at the first bulk element to avoid a full `Command::from_frame` diff --git a/crates/ember-server/src/grpc.rs b/crates/ember-server/src/grpc.rs index cb24cd1a..aa9cb8e0 100644 --- a/crates/ember-server/src/grpc.rs +++ b/crates/ember-server/src/grpc.rs @@ -11,7 +11,9 @@ use std::time::{Duration, Instant}; use bytes::Bytes; use ember_core::{Engine, ShardRequest, ShardResponse, TtlResult, Value}; +use subtle::ConstantTimeEq; use tokio_stream::wrappers::ReceiverStream; +use tonic::service::interceptor::InterceptedService; use tonic::{Request, Response, Status, Streaming}; use crate::pubsub::PubSubManager; @@ -48,11 +50,22 @@ impl EmberService { } } - /// Build this service into a tonic router, optionally with auth. - pub fn into_service(self) -> proto::ember_cache_server::EmberCacheServer { - proto::ember_cache_server::EmberCacheServer::new(self) + /// Build this service into a tonic router with optional authentication. + /// + /// When `requirepass` is configured on the server, every gRPC request + /// must carry a matching `authorization` metadata header. Comparison + /// uses constant-time equality to prevent timing side-channels. + pub fn into_service( + self, + ) -> InterceptedService, AuthInterceptor> + { + let interceptor = AuthInterceptor { + requirepass: self.ctx.requirepass.clone(), + }; + let svc = proto::ember_cache_server::EmberCacheServer::new(self) .max_decoding_message_size(4 * 1024 * 1024) // 4 MB - .max_encoding_message_size(4 * 1024 * 1024) + .max_encoding_message_size(4 * 1024 * 1024); + InterceptedService::new(svc, interceptor) } /// Routes a single-key request through the engine. @@ -82,6 +95,34 @@ impl EmberService { } } +/// gRPC authentication interceptor. +/// +/// When `requirepass` is `Some`, every request must include an +/// `authorization` metadata header whose value matches the password. +/// Uses constant-time comparison to prevent timing side-channels. +/// When `requirepass` is `None`, all requests pass through. +#[derive(Clone)] +pub struct AuthInterceptor { + requirepass: Option, +} + +impl tonic::service::Interceptor for AuthInterceptor { + fn call(&mut self, req: Request<()>) -> Result, Status> { + let password = match &self.requirepass { + Some(pw) => pw, + None => return Ok(req), + }; + let token = req + .metadata() + .get("authorization") + .and_then(|v| v.to_str().ok()); + match token { + Some(t) if bool::from(t.as_bytes().ct_eq(password.as_bytes())) => Ok(req), + _ => Err(Status::unauthenticated("authentication required")), + } + } +} + /// Extracts the bytes from a Value::String, or returns an empty vec for /// non-string types. This is intentionally lenient — callers that need /// strict type checking should match on Value::String directly. @@ -1945,6 +1986,24 @@ impl EmberCache for EmberService { )); } + let total_subs = req.channels.len() + req.patterns.len(); + if total_subs > crate::connection_common::MAX_SUBSCRIPTIONS_PER_CONN { + return Err(Status::invalid_argument(format!( + "too many subscriptions ({total_subs}), max {}", + crate::connection_common::MAX_SUBSCRIPTIONS_PER_CONN + ))); + } + + for pat in &req.patterns { + if pat.len() > crate::connection_common::MAX_PATTERN_LEN { + return Err(Status::invalid_argument(format!( + "pattern too long ({} bytes), max {}", + pat.len(), + crate::connection_common::MAX_PATTERN_LEN + ))); + } + } + let (tx, rx) = tokio::sync::mpsc::channel(256); let pubsub = Arc::clone(&self.pubsub); diff --git a/crates/ember-server/src/server.rs b/crates/ember-server/src/server.rs index 57e9fda5..a4fee6c7 100644 --- a/crates/ember-server/src/server.rs +++ b/crates/ember-server/src/server.rs @@ -139,6 +139,7 @@ pub async fn run( ); info!("gRPC listening on {grpc_addr}"); let server = tonic::transport::Server::builder() + .concurrency_limit_per_connection(256) .add_service(svc.into_service()) .serve(grpc_addr); Some(tokio::spawn(async move { @@ -408,6 +409,7 @@ pub async fn run_concurrent( ); info!("gRPC listening on {grpc_addr}"); let server = tonic::transport::Server::builder() + .concurrency_limit_per_connection(256) .add_service(svc.into_service()) .serve(grpc_addr); Some(tokio::spawn(async move {