diff --git a/todored/src/routes/board/limit.rs b/todored/src/routes/board/limit.rs new file mode 100644 index 0000000..a71caa5 --- /dev/null +++ b/todored/src/routes/board/limit.rs @@ -0,0 +1,34 @@ +//! Rate limiting for board websocket. + +use axum::extract::ws::CloseCode; +use redis::Connection; +use crate::outcome::LoggableOutcome; + +pub fn rate_limit_by_key( + mut rconn: Connection, + key: String, + increment: usize, + limit: usize, + expiration_s: usize +) -> Result<(), CloseCode> { + log::trace!("Incrementing rate limit counter for {key:?}..."); + let response: usize = rconn.cmd("INCRBY") + .arg(&key) + .arg(increment) + .log_err_to_error("Could not increase rate limit counter") + .map_err(|_| 1011u16)?; + + log::trace!("Refreshing rate limit counter expiration for {key:?}..."); + rconn.cmd("EXPIRE") + .arg(&key) + .arg(expiration_s) + .log_err_to_warn("Could not set expiration for rate limit counter"); + + log::trace!("Checking rate limit of {limit} / {expiration_s} s for {key:?}..."); + if response > limit { + log::warn!("Hit rate limit of {limit} / {expiration_s} s for {key:?}: counter is at {response}!"); + return Err(1008u16); + } + + Ok(()) +} diff --git a/todored/src/routes/board/mod.rs b/todored/src/routes/board/mod.rs index 90c34d2..f4d30cf 100644 --- a/todored/src/routes/board/mod.rs +++ b/todored/src/routes/board/mod.rs @@ -7,5 +7,6 @@ pub(self) mod ws_receive; pub(self) mod redis_xadd; pub(self) mod redis_xread; pub(self) mod ws_send; +pub(self) mod limit; pub(crate) use self::axum::handler as board_websocket;