1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
// This Source Code Form is subject to the terms of the Mozilla Public // License, v. 2.0. If a copy of the MPL was not distributed with this // file, You can obtain one at https://mozilla.org/MPL/2.0/. //! A rate limiter using a fixed window counter for arbitrary keys, backed by Redis. //! //! # About //! //! This library provides a fixed window counter for arbitrary keys backed by a Redis instance. The //! period length and per-period maximum limit are configurable on setup. The communication with //! the backend is performing in a non-blocking, asynchronous fashion allowing the library to pair //! with other async frameworks such as Tokio, Actix, etc. //! //! *Note*: Currently pre-*async/await* Futures are used (that is Futures 0.1.x) but that may be //! upgraded in the near future as the Rust 1.39.0 release is near. //! //! # Usage //! //! Add `limitation` to your `Cargo.toml`: //! //! ```toml //! [dependencies] //! limitation = "0.1.1" //! ``` //! //! ## Quick Example //! //! The primary type is the [`Limiter`] which uses a builder to construct itself. Once built, use //! the [`count`] method with a key representing a user, a session, an interaction, etc. Note that //! `count` is a Future and therefore has to be driving to completion with a runtime. For example, //! to run one count on a key of `"10.0.0.5"`: //! //! [`Limiter`]: struct.Limiter.html //! [`count`]: struct.Limiter.html#method.count //! //! ```no_run //! use limitation::Limiter; //! use futures::Future; //! //! // Build a Limiter with a default rate with a local running Redis instance //! let limiter = Limiter::build("redis://127.0.0.1/").finish()?; //! // The key to count and track //! let key = "10.0.0.5"; //! //! // Start and run a Tokio runtime to drive the Future to completion //! tokio::run( //! limiter //! // Count returns a Status if the key is under the limit and an `Error::LimitExceeded` //! // containing a Status if the limit has been exceeded //! .count(key) //! // This example deals with both limit exceeded and any Redis connection issues. In this //! // case we'll print the error to the standard error stream to show the current limit //! // status //! .map_err(|err| eprintln!("err: {}", err)) //! // In this case we are under the limit and can print out the limit status to the //! // standard output stream //! .and_then(|status| { //! println!("ok: {:?}", status); //! Ok(()) //! }), //! ); //! # Ok::<(), limitation::Error>(()) //! ``` //! //! ## The Limiter Builder //! //! The builder for the `Limiter` has 2 settings which can be customized to the use case: //! //! - [`limit`]: The high water mark for number of requests in the period. The default is `5000`. //! - [`period`]: A `Duration` for the period window. The default is 60 minutes. //! //! ```no_run //! use limitation::Limiter; //! use std::time::Duration; //! //! let limiter = Limiter::build("redis://127.0.0.1/") //! .limit(5) //! .period(Duration::from_secs(10)) //! .finish()?; //! # Ok::<(), limitation::Error>(()) //! ``` //! //! [`limit`]: struct.Builder.html#method.limit //! [`period`]: struct.Builder.html#method.period //! //! # Examples //! //! A simple example that uses this library can be found in [limitation-example]. //! //! [limitation-example]: https://github.com/fnichol/limitation/tree/master/limitation-example //! //! # Related Projects and References //! //! The primary inspiration for the implementation was a blog post called [An alternative approach //! to building a simple API Rate limiter using NodeJS and Redis][blog-post] by [Atul //! R](https://twitter.com/masteratul94). //! //! Other research and references used for this library: //! //! - <https://nordicapis.com/everything-you-need-to-know-about-api-rate-limiting/> //! - <https://hechao.li/2018/06/25/Rate-Limiter-Part1/> //! - <https://www.figma.com/blog/an-alternative-approach-to-rate-limiting/> //! - <https://engagor.github.io/blog/2017/05/02/sliding-window-rate-limiter-redis/> //! //! [blog-post]: https://blog.atulr.com/rate-limiter/ //! //! # Ideas and Future Work //! //! These are some ideas and potential future work for this project. If you're reading this then //! maybe you're curious or interesting in helping out? Great! Be sure to check out the //! [Contributing] section and dig in! //! //! - Investigate and offer alternative rate-limiting algorithms, notably a Sliding Window //! solution. //! - Add async Redis connection pooling with the `bb8` and `bb8-redis` crates to reduce //! connection establishment delays. //! - Add a `status` method on `Limiter` which returns they key's `Status` without counting a //! request. //! - Add `RedisServer` support in an integration testing suite, similar to the infrastructure in //! the [redis] crate. //! //! [Contributing]: https://github.com/fnichol/limitation/tree/master/limitation#contributing //! [redis]: https://github.com/mitsuhiko/redis-rs/blob/master/tests/support/mod.rs #![doc(html_root_url = "https://docs.rs/limitation/0.1.1")] #![deny(missing_docs)] use chrono::SubsecRound; use futures::Future; use redis::Client; use std::convert::TryInto; use std::error; use std::fmt; use std::ops::Add; use std::time::Duration; /// The default limit of requests in a period const DEFAULT_LIMIT: usize = 5000; /// The default length of the period in seconds const DEFAULT_PERIOD_SECS: u64 = 60 * 60; /// A rate limiter using a fixed window counter, backed by Redis. /// /// The per-period limit and the period duration are customizable when building an instance. /// /// The [`count`] method is the primary unit of interaction which requires a key representing a /// user, a session, an interaction, etc. This method returns a Future which needs a runtime to /// drive the task to completion asynchronously. /// /// [`count`]: #method.count #[derive(Clone, Debug)] pub struct Limiter { /// The Redis client client: Client, /// The per-period limit limit: usize, /// The period duration period: Duration, } impl Limiter { /// Returns a builder for a `Limiter`. /// /// The [`finish`] method will build the final `Limiter` and perform a **synchronous** /// connection test to the Redis backend. /// /// [`finish`]: struct.Builder.html#method.finish pub fn build(redis_url: &str) -> Builder { Builder { redis_url, limit: DEFAULT_LIMIT, period: Duration::from_secs(DEFAULT_PERIOD_SECS), } } /// Counts a request on a key over a period and returns a [`Status`]. /// /// The `Status` type gives the caller a current state snapshot for the given key. If the limit /// is exceeded a `Error::LimitExceeded` will be returned which also contains a `Status`. /// Critically, the `Status` contains a time when the next period begins and the limit will /// reset. The time is a "unix timestamp" in UTC time. /// /// # Errors /// /// Returns an `Err` if: /// /// - The limit has been exceeded in the current period /// - A client error has occurred /// - A time computation failed /// /// [`Status`]: struct.Status.html pub fn count<K: Into<String>>(&self, key: K) -> impl Future<Item = Status, Error = Error> { let limit = self.limit; self.track(key).and_then(move |(count, reset_epoch_utc)| { let status = build_status(count, limit, reset_epoch_utc); if count > limit { Err(Error::LimitExceeded(status)) } else { Ok(status) } }) } /// Tracks the given key in a period and returns the count and TTL for the key in seconds. fn track<K: Into<String>>(&self, key: K) -> impl Future<Item = (usize, usize), Error = Error> { let key = key.into(); let exipres = self.period.as_secs(); self.client .get_async_connection() .from_err() .and_then(move |con| { // The seed of this approach is outlined Atul R in a blog post about rate limiting // using NodeJS and Redis. For more details, see // https://blog.atulr.com/rate-limiter/ let mut pipe = redis::pipe(); pipe.atomic() .cmd("SET") .arg(&key) .arg(0) .arg("EX") .arg(exipres) .arg("NX") .ignore() .cmd("INCR") .arg(&key) .cmd("TTL") .arg(&key); pipe.query_async(con) .from_err() .and_then(|(_, (count, ttl)): (_, (usize, u64))| { Ok((count, epoch_utc_plus(Duration::from_secs(ttl))?)) }) }) } } /// A report for a given key containing the limit status. /// /// The status contains the following information: /// /// - [`limit`]: the maximum number of requests allowed in the current period /// - [`remaining`]: how many requests are left in the current period /// - [`reset_epoch_utc`]: a UNIX timestamp in UTC approximately when the next period will begin /// /// [`limit`]: #method.limit /// [`remaining`]: #method.remaining /// [`reset_epoch_utc`]: #method.reset_epoch_utc #[derive(Clone, Debug)] pub struct Status { limit: usize, remaining: usize, reset_epoch_utc: usize, } impl Status { /// Returns the maximum number of requests permitted in the current period. pub fn limit(&self) -> usize { self.limit } /// Returns the number of requests remaining in the current period. pub fn remaining(&self) -> usize { self.remaining } /// Returns a UNIX timestamp in UTC approximately when the next period will begin. pub fn reset_epoch_utc(&self) -> usize { self.reset_epoch_utc } } /// A builder for a [`Limiter`]. /// /// [`Limiter`]: struct.Limiter.html pub struct Builder<'a> { redis_url: &'a str, limit: usize, period: Duration, } impl Builder<'_> { /// Sets a new maximum limit for the Limiter. pub fn limit(&mut self, limit: usize) -> &mut Self { self.limit = limit; self } /// Sets a new period duration for the Limiter. pub fn period(&mut self, period: Duration) -> &mut Self { self.period = period; self } /// Finializes and returns a `Limiter`. /// /// Note that this method will connect to the Redis server to test its connection which is a /// **synchronous** operation. /// /// # Errors /// /// Returns an `Err` if the Redis client fails to be created or fails to connect. pub fn finish(&self) -> Result<Limiter, Error> { Ok(Limiter { client: Client::open(self.redis_url)?, limit: self.limit, period: self.period, }) } } /// Error type for this crate. #[derive(Debug)] pub enum Error { /// The Redis client failed to connect or run a query. Client(redis::RedisError), /// The limit is exceeded for a key. LimitExceeded(Status), /// A time conversion failed. Time(time::OutOfRangeError), } impl fmt::Display for Error { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { Error::Client(ref err) => write!(f, "client error ({})", err), Error::LimitExceeded(ref status) => write!(f, "rate limit exceeded ({:?})", status), Error::Time(ref err) => write!(f, "time conversion error ({})", err), } } } impl error::Error for Error { fn source(&self) -> Option<&(dyn error::Error + 'static)> { match self { Error::Client(ref err) => err.source(), Error::LimitExceeded(_) => None, Error::Time(ref err) => err.source(), } } } impl From<redis::RedisError> for Error { fn from(err: redis::RedisError) -> Self { Error::Client(err) } } impl From<time::OutOfRangeError> for Error { fn from(err: time::OutOfRangeError) -> Self { Error::Time(err) } } /// Builds a `Status`. fn build_status(count: usize, limit: usize, reset_epoch_utc: usize) -> Status { let remaining = if count >= limit { 0 } else { limit - count }; Status { limit, remaining, reset_epoch_utc, } } /// Calculates a timestamp for "now plus a duration". fn epoch_utc_plus(duration: Duration) -> Result<usize, time::OutOfRangeError> { Ok(chrono::Utc::now() .add(chrono::Duration::from_std(duration)?) .round_subsecs(0) .timestamp() .try_into() .unwrap_or(0)) }