harden round 3 — overflow, DoS, correctness fixes

crates/ember-cluster/src/gossip.rs

@@ -422,7 +422,7 @@ impl GossipEngine {
                     if *node == self.local_id {
                         // Refute suspicion by incrementing our incarnation
                         if *incarnation >= self.incarnation {
-                            self.incarnation = incarnation + 1;
+                            self.incarnation = incarnation.saturating_add(1);
                             self.queue_update(NodeUpdate::Alive {
                                 node: self.local_id,
                                 addr: self.local_addr,
@@ -447,7 +447,7 @@ impl GossipEngine {
                 NodeUpdate::Dead { node, incarnation } => {
                     if *node == self.local_id {
                         // Refute death claim
-                        self.incarnation = incarnation + 1;
+                        self.incarnation = incarnation.saturating_add(1);
                         self.queue_update(NodeUpdate::Alive {
                             node: self.local_id,
                             addr: self.local_addr,

crates/ember-cluster/src/message.rs

@@ -10,6 +10,33 @@ use bytes::{Buf, BufMut, Bytes, BytesMut};
 
 use crate::{NodeId, SlotRange};
 
+/// Maximum number of members in a Welcome message or updates in a Ping/Ack.
+/// Prevents allocation bombs from crafted messages.
+const MAX_COLLECTION_COUNT: usize = 1024;
+
+// Safe read helpers that return io::Error instead of panicking on truncated input.
+
+fn safe_get_u8(buf: &mut &[u8]) -> io::Result<u8> {
+    if buf.is_empty() {
+        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "need 1 byte"));
+    }
+    Ok(buf.get_u8())
+}
+
+fn safe_get_u16_le(buf: &mut &[u8]) -> io::Result<u16> {
+    if buf.len() < 2 {
+        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "need 2 bytes"));
+    }
+    Ok(buf.get_u16_le())
+}
+
+fn safe_get_u64_le(buf: &mut &[u8]) -> io::Result<u64> {
+    if buf.len() < 8 {
+        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "need 8 bytes"));
+    }
+    Ok(buf.get_u64_le())
+}
+
 /// Message types for the SWIM gossip protocol.
 #[derive(Debug, Clone, PartialEq)]
 pub enum GossipMessage {
@@ -160,10 +187,10 @@ impl GossipMessage {
             ));
         }
 
-        let msg_type = buf.get_u8();
+        let msg_type = safe_get_u8(&mut buf)?;
         match msg_type {
             MSG_PING => {
-                let seq = buf.get_u64_le();
+                let seq = safe_get_u64_le(&mut buf)?;
                 let sender = decode_node_id(&mut buf)?;
                 let updates = decode_updates(&mut buf)?;
                 Ok(GossipMessage::Ping {
@@ -173,7 +200,7 @@ impl GossipMessage {
                 })
             }
             MSG_PING_REQ => {
-                let seq = buf.get_u64_le();
+                let seq = safe_get_u64_le(&mut buf)?;
                 let sender = decode_node_id(&mut buf)?;
                 let target = decode_node_id(&mut buf)?;
                 let target_addr = decode_socket_addr(&mut buf)?;
@@ -185,7 +212,7 @@ impl GossipMessage {
                 })
             }
             MSG_ACK => {
-                let seq = buf.get_u64_le();
+                let seq = safe_get_u64_le(&mut buf)?;
                 let sender = decode_node_id(&mut buf)?;
                 let updates = decode_updates(&mut buf)?;
                 Ok(GossipMessage::Ack {
@@ -204,7 +231,13 @@ impl GossipMessage {
             }
             MSG_WELCOME => {
                 let sender = decode_node_id(&mut buf)?;
-                let count = buf.get_u16_le() as usize;
+                let count = safe_get_u16_le(&mut buf)? as usize;
+                if count > MAX_COLLECTION_COUNT {
+                    return Err(io::Error::new(
+                        io::ErrorKind::InvalidData,
+                        format!("member count {count} exceeds limit"),
+                    ));
+                }
                 let mut members = Vec::with_capacity(count);
                 for _ in 0..count {
                     members.push(decode_member_info(&mut buf)?);
@@ -327,13 +360,13 @@ fn encode_update(buf: &mut BytesMut, update: &NodeUpdate) {
 }
 
 fn decode_updates(buf: &mut &[u8]) -> io::Result<Vec<NodeUpdate>> {
-    if buf.len() < 2 {
+    let count = safe_get_u16_le(buf)? as usize;
+    if count > MAX_COLLECTION_COUNT {
         return Err(io::Error::new(
-            io::ErrorKind::UnexpectedEof,
-            "not enough bytes for update count",
+            io::ErrorKind::InvalidData,
+            format!("update count {count} exceeds limit"),
         ));
     }
-    let count = buf.get_u16_le() as usize;
     let mut updates = Vec::with_capacity(count);
     for _ in 0..count {
         updates.push(decode_update(buf)?);
@@ -342,18 +375,12 @@ fn decode_updates(buf: &mut &[u8]) -> io::Result<Vec<NodeUpdate>> {
 }
 
 fn decode_update(buf: &mut &[u8]) -> io::Result<NodeUpdate> {
-    if buf.is_empty() {
-        return Err(io::Error::new(
-            io::ErrorKind::UnexpectedEof,
-            "not enough bytes for update type",
-        ));
-    }
-    let update_type = buf.get_u8();
+    let update_type = safe_get_u8(buf)?;
     match update_type {
         UPDATE_ALIVE => {
             let node = decode_node_id(buf)?;
             let addr = decode_socket_addr(buf)?;
-            let incarnation = buf.get_u64_le();
+            let incarnation = safe_get_u64_le(buf)?;
             Ok(NodeUpdate::Alive {
                 node,
                 addr,
@@ -362,12 +389,12 @@ fn decode_update(buf: &mut &[u8]) -> io::Result<NodeUpdate> {
         }
         UPDATE_SUSPECT => {
             let node = decode_node_id(buf)?;
-            let incarnation = buf.get_u64_le();
+            let incarnation = safe_get_u64_le(buf)?;
             Ok(NodeUpdate::Suspect { node, incarnation })
         }
         UPDATE_DEAD => {
             let node = decode_node_id(buf)?;
-            let incarnation = buf.get_u64_le();
+            let incarnation = safe_get_u64_le(buf)?;
             Ok(NodeUpdate::Dead { node, incarnation })
         }
         UPDATE_LEFT => {
@@ -396,13 +423,19 @@ fn encode_member_info(buf: &mut BytesMut, member: &MemberInfo) {
 fn decode_member_info(buf: &mut &[u8]) -> io::Result<MemberInfo> {
     let id = decode_node_id(buf)?;
     let addr = decode_socket_addr(buf)?;
-    let incarnation = buf.get_u64_le();
-    let is_primary = buf.get_u8() != 0;
-    let slot_count = buf.get_u16_le() as usize;
+    let incarnation = safe_get_u64_le(buf)?;
+    let is_primary = safe_get_u8(buf)? != 0;
+    let slot_count = safe_get_u16_le(buf)? as usize;
+    if slot_count > MAX_COLLECTION_COUNT {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!("slot range count {slot_count} exceeds limit"),
+        ));
+    }
     let mut slots = Vec::with_capacity(slot_count);
     for _ in 0..slot_count {
-        let start = buf.get_u16_le();
-        let end = buf.get_u16_le();
+        let start = safe_get_u16_le(buf)?;
+        let end = safe_get_u16_le(buf)?;
         slots.push(SlotRange::try_new(start, end)?);
     }
     Ok(MemberInfo {

crates/ember-core/src/concurrent.rs

@@ -169,7 +169,7 @@ impl ConcurrentKeyspace {
             if entry.is_expired() {
                 return false;
             }
-            entry.expires_at_ms = time::now_ms() + seconds * 1000;
+            entry.expires_at_ms = time::now_ms().saturating_add(seconds.saturating_mul(1000));
             true
         } else {
             false

crates/ember-core/src/engine.rs

@@ -185,6 +185,11 @@ impl Engine {
         Ok(results)
     }
 
+    /// Returns true if both keys are owned by the same shard.
+    pub fn same_shard(&self, key1: &str, key2: &str) -> bool {
+        self.shard_for_key(key1) == self.shard_for_key(key2)
+    }
+
     /// Determines which shard owns a given key.
     fn shard_for_key(&self, key: &str) -> usize {
         shard_index(key, self.shards.len())

crates/ember-core/src/keyspace.rs

@@ -513,7 +513,7 @@ impl Keyspace {
                 if entry.expires_at_ms == 0 {
                     self.expiry_count += 1;
                 }
-                entry.expires_at_ms = time::now_ms() + seconds * 1000;
+                entry.expires_at_ms = time::now_ms().saturating_add(seconds.saturating_mul(1000));
                 true
             }
             None => false,
@@ -589,7 +589,7 @@ impl Keyspace {
                 if entry.expires_at_ms == 0 {
                     self.expiry_count += 1;
                 }
-                entry.expires_at_ms = time::now_ms() + millis;
+                entry.expires_at_ms = time::now_ms().saturating_add(millis);
                 true
             }
             None => false,
@@ -1471,8 +1471,12 @@ impl Keyspace {
         };
 
         if is_empty {
+            if let Some(removed_entry) = self.entries.remove(key) {
+                if removed_entry.expires_at_ms != 0 {
+                    self.expiry_count = self.expiry_count.saturating_sub(1);
+                }
+            }
             self.memory.remove_with_size(old_entry_size);
-            self.entries.remove(key);
         } else {
             let new_entry_size = memory::entry_size(key, &self.entries[key].value);
             self.memory.adjust(old_entry_size, new_entry_size);
@@ -1732,8 +1736,12 @@ impl Keyspace {
         };
 
         if is_empty {
+            if let Some(removed_entry) = self.entries.remove(key) {
+                if removed_entry.expires_at_ms != 0 {
+                    self.expiry_count = self.expiry_count.saturating_sub(1);
+                }
+            }
             self.memory.remove_with_size(old_entry_size);
-            self.entries.remove(key);
         } else {
             let new_entry_size = memory::entry_size(key, &self.entries[key].value);
             self.memory.adjust(old_entry_size, new_entry_size);

crates/ember-core/src/shard.rs

@@ -705,11 +705,14 @@ fn dispatch(
             Err(IncrError::OutOfMemory) => ShardResponse::OutOfMemory,
             Err(e) => ShardResponse::Err(e.to_string()),
         },
-        ShardRequest::DecrBy { key, delta } => match ks.incr_by(key, -delta) {
-            Ok(val) => ShardResponse::Integer(val),
-            Err(IncrError::WrongType) => ShardResponse::WrongType,
-            Err(IncrError::OutOfMemory) => ShardResponse::OutOfMemory,
-            Err(e) => ShardResponse::Err(e.to_string()),
+        ShardRequest::DecrBy { key, delta } => match delta.checked_neg() {
+            Some(neg) => match ks.incr_by(key, neg) {
+                Ok(val) => ShardResponse::Integer(val),
+                Err(IncrError::WrongType) => ShardResponse::WrongType,
+                Err(IncrError::OutOfMemory) => ShardResponse::OutOfMemory,
+                Err(e) => ShardResponse::Err(e.to_string()),
+            },
+            None => ShardResponse::Err("ERR increment or decrement would overflow".into()),
         },
         ShardRequest::IncrByFloat { key, delta } => match ks.incr_by_float(key, *delta) {
             Ok(val) => ShardResponse::BulkString(val),

crates/ember-core/src/time.rs

@@ -26,8 +26,11 @@ pub fn is_expired(expires_at_ms: u64) -> bool {
 /// Converts a Duration to an absolute expiry timestamp.
 #[inline]
 pub fn expiry_from_duration(ttl: Option<std::time::Duration>) -> u64 {
-    ttl.map(|d| now_ms() + d.as_millis() as u64)
-        .unwrap_or(NO_EXPIRY)
+    ttl.map(|d| {
+        let ms = d.as_millis().min(u64::MAX as u128) as u64;
+        now_ms().saturating_add(ms)
+    })
+    .unwrap_or(NO_EXPIRY)
 }
 
 /// Returns remaining TTL in seconds, or None if no expiry.

crates/ember-persistence/src/aof.rs

@@ -304,6 +304,13 @@ impl AofRecord {
         Ok(buf)
     }
 
+    /// Cap pre-allocation to avoid huge allocations from corrupt count fields.
+    /// The loop will still iterate `count` times — this just limits the
+    /// up-front reservation so a bogus u32 can't exhaust memory.
+    fn capped_capacity(count: u32) -> usize {
+        (count as usize).min(65_536)
+    }
+
     /// Deserializes a record from a byte slice (tag + payload, no CRC).
     fn from_bytes(data: &[u8]) -> Result<Self, FormatError> {
         let mut cursor = io::Cursor::new(data);
@@ -331,7 +338,7 @@ impl AofRecord {
             TAG_LPUSH | TAG_RPUSH => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut values = Vec::with_capacity(count as usize);
+                let mut values = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     values.push(Bytes::from(format::read_bytes(&mut cursor)?));
                 }
@@ -352,7 +359,7 @@ impl AofRecord {
             TAG_ZADD => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut members = Vec::with_capacity(count as usize);
+                let mut members = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     let score = format::read_f64(&mut cursor)?;
                     let member = read_string(&mut cursor, "member")?;
@@ -363,7 +370,7 @@ impl AofRecord {
             TAG_ZREM => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut members = Vec::with_capacity(count as usize);
+                let mut members = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     members.push(read_string(&mut cursor, "member")?);
                 }
@@ -389,7 +396,7 @@ impl AofRecord {
             TAG_HSET => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut fields = Vec::with_capacity(count as usize);
+                let mut fields = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     let field = read_string(&mut cursor, "field")?;
                     let value = Bytes::from(format::read_bytes(&mut cursor)?);
@@ -400,7 +407,7 @@ impl AofRecord {
             TAG_HDEL => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut fields = Vec::with_capacity(count as usize);
+                let mut fields = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     fields.push(read_string(&mut cursor, "field")?);
                 }
@@ -415,7 +422,7 @@ impl AofRecord {
             TAG_SADD => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut members = Vec::with_capacity(count as usize);
+                let mut members = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     members.push(read_string(&mut cursor, "member")?);
                 }
@@ -424,7 +431,7 @@ impl AofRecord {
             TAG_SREM => {
                 let key = read_string(&mut cursor, "key")?;
                 let count = format::read_u32(&mut cursor)?;
-                let mut members = Vec::with_capacity(count as usize);
+                let mut members = Vec::with_capacity(Self::capped_capacity(count));
                 for _ in 0..count {
                     members.push(read_string(&mut cursor, "member")?);
                 }

crates/ember-persistence/src/recovery.rs

@@ -361,7 +361,7 @@ fn replay_aof(
                 apply_incr(map, key, delta);
             }
             AofRecord::DecrBy { key, delta } => {
-                apply_incr(map, key, -delta);
+                apply_incr(map, key, delta.saturating_neg());
             }
             AofRecord::Append { key, value } => {
                 let entry = map