From f71a6f5a4139dffeb3dc1a196de41ef7ce8b5724 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Thu, 26 Feb 2026 20:44:21 -0500 Subject: [PATCH] fix: auth failure metrics, command memory budget, migration progress MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **h1 — auth failure counter**: adds `ember_auth_failures_total{reason}` prometheus counter. reason is `wrongpass` for bad credentials, `noauth` for commands sent before authenticating. incremented in `try_auth` and at the NOAUTH error path in both connection modes. **m1 — per-command memory budget**: adds `max_command_memory` (128MB default) to `ConnectionLimits`. `validate_command_sizes` now sums total bytes across all keys and values in MSET/LPUSH/RPUSH and rejects with ERR if the combined payload exceeds the budget. guards against DoS via million-element commands that individually pass per-item limits. **m3 — migration progress overflow**: replaces `saturating_mul(100) / total` with f64 division in `Migration::progress()`. the old formula silently reported 100% when `keys_migrated * 100` overflowed u64. --- crates/ember-cluster/src/migration.rs | 2 +- crates/ember-server/src/concurrent_handler.rs | 11 +++-- crates/ember-server/src/config.rs | 7 +++ .../ember-server/src/connection/dispatch.rs | 9 ++-- crates/ember-server/src/connection/handler.rs | 9 ++-- crates/ember-server/src/connection/mod.rs | 2 + crates/ember-server/src/connection_common.rs | 48 ++++++++++++++----- crates/ember-server/src/metrics.rs | 12 +++++ 8 files changed, 79 insertions(+), 21 deletions(-) diff --git a/crates/ember-cluster/src/migration.rs b/crates/ember-cluster/src/migration.rs index 4a30cc25..1728993a 100644 --- a/crates/ember-cluster/src/migration.rs +++ b/crates/ember-cluster/src/migration.rs @@ -145,7 +145,7 @@ impl Migration { if total == 0 { 100 } else { - (self.keys_migrated.saturating_mul(100) / total).min(100) as u8 + ((self.keys_migrated as f64 / total as f64) * 100.0).min(100.0) as u8 } }) } diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 968923d5..f4903927 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -29,6 +29,7 @@ use crate::connection_common::{ is_auth_frame, is_monitor_frame, try_auth, validate_command_sizes, MonitorEvent, TransactionState, }; +use crate::metrics::on_auth_failure; use crate::pubsub::PubSubManager; use crate::server::{format_client_list, ServerContext}; use crate::slowlog::SlowLog; @@ -126,6 +127,7 @@ where .await; response.serialize(&mut out); } else { + on_auth_failure("noauth"); Frame::Error("NOAUTH Authentication required.".into()) .serialize(&mut out); } @@ -250,9 +252,12 @@ async fn process( match Command::from_frame(frame) { Ok(cmd) => { // reject oversized keys/values before any further processing - if let Some(err) = - validate_command_sizes(&cmd, ctx.limits.max_key_len, ctx.limits.max_value_len) - { + if let Some(err) = validate_command_sizes( + &cmd, + ctx.limits.max_key_len, + ctx.limits.max_value_len, + ctx.limits.max_command_memory, + ) { return err; } diff --git a/crates/ember-server/src/config.rs b/crates/ember-server/src/config.rs index c155a3c5..f1420a96 100644 --- a/crates/ember-server/src/config.rs +++ b/crates/ember-server/src/config.rs @@ -403,6 +403,7 @@ impl EmberConfig { max_pipeline_depth: self.max_pipeline_depth, max_key_len, max_value_len, + max_command_memory: 128 * 1024 * 1024, // 128MB stats_poll_interval: Duration::from_secs(self.engine.stats_poll_interval_secs), }) } @@ -572,6 +573,12 @@ pub struct ConnectionLimits { pub max_pipeline_depth: usize, pub max_key_len: usize, pub max_value_len: usize, + /// Total bytes allowed across all keys and values in a single command. + /// + /// Guards against DoS via commands like MSET or LPUSH with millions of + /// small elements that individually pass the per-item checks. Defaults + /// to 128MB. + pub max_command_memory: usize, pub stats_poll_interval: Duration, } diff --git a/crates/ember-server/src/connection/dispatch.rs b/crates/ember-server/src/connection/dispatch.rs index f4d77e53..b45eb554 100644 --- a/crates/ember-server/src/connection/dispatch.rs +++ b/crates/ember-server/src/connection/dispatch.rs @@ -152,9 +152,12 @@ pub(super) async fn prepare_command( }; // reject oversized keys/values before any further processing - if let Some(err) = - validate_command_sizes(&cmd, ctx.limits.max_key_len, ctx.limits.max_value_len) - { + if let Some(err) = validate_command_sizes( + &cmd, + ctx.limits.max_key_len, + ctx.limits.max_value_len, + ctx.limits.max_command_memory, + ) { return PreparedDispatch::Immediate(PendingResponse::Immediate(err)); } diff --git a/crates/ember-server/src/connection/handler.rs b/crates/ember-server/src/connection/handler.rs index 014e7fb2..63c0c44c 100644 --- a/crates/ember-server/src/connection/handler.rs +++ b/crates/ember-server/src/connection/handler.rs @@ -229,9 +229,12 @@ pub(super) async fn handle_blocking_pop_cmd( Err(e) => return Frame::Error(format!("ERR {e}")), }; - if let Some(err) = - validate_command_sizes(&cmd, ctx.limits.max_key_len, ctx.limits.max_value_len) - { + if let Some(err) = validate_command_sizes( + &cmd, + ctx.limits.max_key_len, + ctx.limits.max_value_len, + ctx.limits.max_command_memory, + ) { return err; } diff --git a/crates/ember-server/src/connection/mod.rs b/crates/ember-server/src/connection/mod.rs index 92e4f5cc..3878d096 100644 --- a/crates/ember-server/src/connection/mod.rs +++ b/crates/ember-server/src/connection/mod.rs @@ -21,6 +21,7 @@ use crate::connection_common::{ initial_acl_user, is_allowed_before_auth, is_auth_frame, is_monitor_frame, try_auth, TransactionState, }; +use crate::metrics::on_auth_failure; use crate::pubsub::PubSubManager; use crate::server::ServerContext; use crate::slowlog::SlowLog; @@ -347,6 +348,7 @@ where .await; response.serialize(&mut out); } else { + on_auth_failure("noauth"); Frame::Error("NOAUTH Authentication required.".into()).serialize(&mut out); } } diff --git a/crates/ember-server/src/connection_common.rs b/crates/ember-server/src/connection_common.rs index 480af814..c8ce235c 100644 --- a/crates/ember-server/src/connection_common.rs +++ b/crates/ember-server/src/connection_common.rs @@ -11,6 +11,7 @@ use ember_protocol::Command; use subtle::ConstantTimeEq; use crate::acl::AclUser; +use crate::metrics::on_auth_failure; use crate::server::ServerContext; // Default values for connection limits. These serve as documentation @@ -23,6 +24,9 @@ pub const DEFAULT_MAX_KEY_LEN: usize = 512 * 1024; #[cfg(test)] /// Default max value length (512MB). pub const DEFAULT_MAX_VALUE_LEN: usize = 512 * 1024 * 1024; +#[cfg(test)] +/// Default max command memory (128MB total payload per command). +pub const DEFAULT_MAX_COMMAND_MEMORY: usize = 128 * 1024 * 1024; /// Checks if a raw frame is an AUTH command (before full parsing). /// @@ -106,6 +110,7 @@ pub fn try_auth(frame: Frame, ctx: &ServerContext) -> AuthResult { let user = match state.get_user(&username) { Some(u) => u, None => { + on_auth_failure("wrongpass"); return AuthResult::fail( "WRONGPASS invalid username-password pair or user is disabled.", ); @@ -113,12 +118,14 @@ pub fn try_auth(frame: Frame, ctx: &ServerContext) -> AuthResult { }; if !user.enabled { + on_auth_failure("wrongpass"); return AuthResult::fail( "WRONGPASS invalid username-password pair or user is disabled.", ); } if !user.verify_password(&password) { + on_auth_failure("wrongpass"); return AuthResult::fail( "WRONGPASS invalid username-password pair or user is disabled.", ); @@ -145,6 +152,7 @@ pub fn try_auth(frame: Frame, ctx: &ServerContext) -> AuthResult { username, } } else { + on_auth_failure("wrongpass"); AuthResult::fail("WRONGPASS invalid username-password pair or user is disabled.") } } @@ -179,13 +187,15 @@ pub fn initial_acl_user(ctx: &ServerContext) -> (Option>, String) { /// Validates key and value sizes for a parsed command. /// -/// Returns an error frame if any key exceeds `max_key_len` or any value -/// exceeds `max_value_len`. Returns `None` when the command passes validation. +/// Returns an error frame if any key exceeds `max_key_len`, any value exceeds +/// `max_value_len`, or the combined payload of a bulk command exceeds +/// `max_command_memory`. Returns `None` when the command passes validation. /// Called on the RESP path to match the limits already enforced by gRPC. pub fn validate_command_sizes( cmd: &Command, max_key_len: usize, max_value_len: usize, + max_command_memory: usize, ) -> Option { // check primary key length if let Some(key) = cmd.primary_key() { @@ -228,6 +238,7 @@ pub fn validate_command_sizes( } } Command::MSet { pairs } => { + let mut total = 0usize; for (k, v) in pairs { if k.len() > max_key_len { return Some(Frame::Error(format!( @@ -241,9 +252,16 @@ pub fn validate_command_sizes( v.len() ))); } + total = total.saturating_add(k.len()).saturating_add(v.len()); + } + if total > max_command_memory { + return Some(Frame::Error(format!( + "ERR command payload {total} bytes exceeds limit of {max_command_memory} bytes" + ))); } } Command::LPush { values, .. } | Command::RPush { values, .. } => { + let mut total = 0usize; for v in values { if v.len() > max_value_len { return Some(Frame::Error(format!( @@ -251,6 +269,12 @@ pub fn validate_command_sizes( v.len() ))); } + total = total.saturating_add(v.len()); + } + if total > max_command_memory { + return Some(Frame::Error(format!( + "ERR command payload {total} bytes exceeds limit of {max_command_memory} bytes" + ))); } } Command::HSet { fields, .. } => { @@ -448,14 +472,14 @@ mod tests { nx: false, xx: false, }; - assert!(validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).is_none()); + assert!(validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).is_none()); } #[test] fn oversized_key_rejected() { let big_key = "x".repeat(DEFAULT_MAX_KEY_LEN + 1); let cmd = Command::Get { key: big_key }; - let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).unwrap(); + let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).unwrap(); assert!(matches!(err, Frame::Error(ref msg) if msg.contains("key length"))); } @@ -463,7 +487,7 @@ mod tests { fn key_at_limit_passes() { let key = "k".repeat(DEFAULT_MAX_KEY_LEN); let cmd = Command::Get { key }; - assert!(validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).is_none()); + assert!(validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).is_none()); } #[test] @@ -476,7 +500,7 @@ mod tests { nx: false, xx: false, }; - let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).unwrap(); + let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).unwrap(); assert!(matches!(err, Frame::Error(ref msg) if msg.contains("value length"))); } @@ -486,7 +510,7 @@ mod tests { let cmd = Command::MSet { pairs: vec![(big_key, Bytes::from_static(b"v"))], }; - let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).unwrap(); + let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).unwrap(); assert!(matches!(err, Frame::Error(ref msg) if msg.contains("key length"))); } @@ -497,7 +521,7 @@ mod tests { key: "mylist".into(), values: vec![Bytes::from_static(b"ok"), big_val], }; - let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).unwrap(); + let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).unwrap(); assert!(matches!(err, Frame::Error(ref msg) if msg.contains("value length"))); } @@ -507,7 +531,7 @@ mod tests { let cmd = Command::Del { keys: vec!["ok".into(), big_key], }; - let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN).unwrap(); + let err = validate_command_sizes(&cmd, DEFAULT_MAX_KEY_LEN, DEFAULT_MAX_VALUE_LEN, DEFAULT_MAX_COMMAND_MEMORY).unwrap(); assert!(matches!(err, Frame::Error(ref msg) if msg.contains("key length"))); } @@ -516,13 +540,15 @@ mod tests { assert!(validate_command_sizes( &Command::Ping(None), DEFAULT_MAX_KEY_LEN, - DEFAULT_MAX_VALUE_LEN + DEFAULT_MAX_VALUE_LEN, + DEFAULT_MAX_COMMAND_MEMORY ) .is_none()); assert!(validate_command_sizes( &Command::DbSize, DEFAULT_MAX_KEY_LEN, - DEFAULT_MAX_VALUE_LEN + DEFAULT_MAX_VALUE_LEN, + DEFAULT_MAX_COMMAND_MEMORY ) .is_none()); } diff --git a/crates/ember-server/src/metrics.rs b/crates/ember-server/src/metrics.rs index ac1857cd..5c476631 100644 --- a/crates/ember-server/src/metrics.rs +++ b/crates/ember-server/src/metrics.rs @@ -331,6 +331,18 @@ pub fn on_connection_rejected() { counter!("ember_connections_rejected").increment(1); } +/// Records an authentication failure. +/// +/// `reason` is a static label describing why auth failed: +/// - `"wrongpass"` — bad password or unknown user +/// - `"noauth"` — command sent without authenticating first +/// - `"acl_deny"` — authenticated but ACL denied the command +#[inline] +pub fn on_auth_failure(reason: &'static str) { + let labels = [("reason", reason)]; + counter!("ember_auth_failures_total", &labels).increment(1); +} + #[cfg(test)] mod tests { use super::*;