Move rate-limiter to its own support crate

Cargo.lock

@@ -1124,6 +1124,7 @@ dependencies = [
  "prost",
  "prost-derive",
  "rand_rccrypto",
+ "rate-limiter",
  "rc_crypto",
  "serde",
  "serde_derive",
@@ -2621,6 +2622,10 @@ dependencies = [
  "rand_core 0.3.1",
 ]
 
+[[package]]
+name = "rate-limiter"
+version = "0.1.0"
+
 [[package]]
 name = "rayon"
 version = "1.5.0"

Cargo.toml

@@ -17,6 +17,7 @@ members = [
     "components/support/interrupt",
     "components/support/jwcrypto",
     "components/support/rand_rccrypto",
+    "components/support/rate-limiter",
     "components/support/restmail-client",
     "components/support/rc_crypto",
     "components/support/rc_crypto/nss",

components/fxa-client/Cargo.toml

@@ -13,6 +13,7 @@ lazy_static = "1.4"
 log = "0.4"
 prost = "0.6"
 prost-derive = "0.6"
+rate-limiter = { path = "../support/rate-limiter" }
 rand_rccrypto = { path = "../support/rand_rccrypto" }
 serde = { version = "1", features = ["rc"] }
 serde_derive = "1"

components/fxa-client/src/oauth.rs

@@ -11,6 +11,7 @@ use crate::{
     util, FirefoxAccount,
 };
 use jwcrypto::{EncryptionAlgorithm, EncryptionParameters};
+use rate_limiter::RateLimiter;
 use rc_crypto::digest;
 use serde_derive::*;
 use std::convert::TryFrom;
@@ -461,56 +462,29 @@ impl FirefoxAccount {
 const AUTH_CIRCUIT_BREAKER_CAPACITY: u8 = 5;
 const AUTH_CIRCUIT_BREAKER_RENEWAL_RATE: f32 = 3.0 / 60.0 / 1000.0; // 3 tokens every minute.
 
-// The auth circuit breaker rate-limits access to the `oauth_introspect_refresh_token`
-// using a fairly naively implemented token bucket algorithm.
 #[derive(Clone, Copy)]
 pub(crate) struct AuthCircuitBreaker {
-    tokens: u8,
-    last_refill: u64, // in ms.
+    rate_limiter: RateLimiter,
 }
 
 impl Default for AuthCircuitBreaker {
     fn default() -> Self {
         AuthCircuitBreaker {
-            tokens: AUTH_CIRCUIT_BREAKER_CAPACITY,
-            last_refill: Self::now(),
+            rate_limiter: RateLimiter::new(
+                AUTH_CIRCUIT_BREAKER_CAPACITY,
+                AUTH_CIRCUIT_BREAKER_RENEWAL_RATE,
+            ),
         }
     }
 }
 
 impl AuthCircuitBreaker {
     pub(crate) fn check(&mut self) -> Result<()> {
-        self.refill();
-        if self.tokens == 0 {
+        if !self.rate_limiter.check() {
             return Err(ErrorKind::AuthCircuitBreakerError.into());
         }
-        self.tokens -= 1;
         Ok(())
     }
-
-    fn refill(&mut self) {
-        let now = Self::now();
-        let new_tokens =
-            ((now - self.last_refill) as f64 * AUTH_CIRCUIT_BREAKER_RENEWAL_RATE as f64) as u8; // `as` is a truncating/saturing cast.
-        if new_tokens > 0 {
-            self.last_refill = now;
-            self.tokens = std::cmp::min(
-                AUTH_CIRCUIT_BREAKER_CAPACITY,
-                self.tokens.saturating_add(new_tokens),
-            );
-        }
-    }
-
-    #[cfg(not(test))]
-    #[inline]
-    fn now() -> u64 {
-        util::now()
-    }
-
-    #[cfg(test)]
-    fn now() -> u64 {
-        1600000000000
-    }
 }
 
 #[derive(Clone)]
@@ -979,23 +953,6 @@ mod tests {
         }
     }
 
-    #[test]
-    fn test_auth_circuit_breaker_unit_recovery() {
-        let mut breaker = AuthCircuitBreaker::default();
-        // AuthCircuitBreaker::now is fixed for tests, let's assert that for sanity.
-        assert_eq!(AuthCircuitBreaker::now(), 1600000000000);
-        for _ in 0..AUTH_CIRCUIT_BREAKER_CAPACITY {
-            assert!(breaker.check().is_ok());
-        }
-        assert!(breaker.check().is_err());
-        // Jump back in time (1 min).
-        breaker.last_refill -= 60 * 1000;
-        let expected_tokens_before_check: u8 =
-            (AUTH_CIRCUIT_BREAKER_RENEWAL_RATE * 60.0 * 1000.0) as u8;
-        assert!(breaker.check().is_ok());
-        assert_eq!(breaker.tokens, expected_tokens_before_check - 1);
-    }
-
     use crate::scopes;
 
     #[test]

components/support/rate-limiter/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "rate-limiter"
+version = "0.1.0"
+authors = ["Edouard Oger <eoger@fastmail.com>"]
+edition = "2018"
+license = "MPL-2.0"
+
+[lib]
+crate-type = ["lib"]

components/support/rate-limiter/src/lib.rs

@@ -0,0 +1,73 @@
+/* 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 http://mozilla.org/MPL/2.0/. */
+
+// Simple token bucket rate-limiter implementation.
+
+#[derive(Clone, Copy)]
+pub struct RateLimiter {
+    capacity: u8,
+    tokens: u8,
+    renewal_rate: f32, // per ms.
+    last_refill: u64,  // in ms.
+}
+
+impl RateLimiter {
+    pub fn new(capacity: u8, renewal_rate: f32) -> Self {
+        Self {
+            capacity,
+            tokens: capacity,
+            renewal_rate,
+            last_refill: Self::now(),
+        }
+    }
+
+    pub fn check(&mut self) -> bool {
+        self.refill();
+        if self.tokens == 0 {
+            return false;
+        }
+        self.tokens -= 1;
+        true
+    }
+
+    fn refill(&mut self) {
+        let now = Self::now();
+        let new_tokens = ((now - self.last_refill) as f64 * self.renewal_rate as f64) as u8; // `as` is a truncating/saturing cast.
+        if new_tokens > 0 {
+            self.last_refill = now;
+            self.tokens = std::cmp::min(self.capacity, self.tokens.saturating_add(new_tokens));
+        }
+    }
+
+    #[inline]
+    fn now() -> u64 {
+        use std::time::{SystemTime, UNIX_EPOCH};
+        let since_epoch = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .expect("Current date before unix epoch.");
+        since_epoch.as_secs() * 1000 + u64::from(since_epoch.subsec_nanos()) / 1_000_000
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::RateLimiter;
+    #[test]
+    fn test_recovery() {
+        let capacity = 10;
+        let renewal_rate = 1.0 / 60.0 / 1000.0; // 1 token per second.
+        let mut breaker = RateLimiter::new(capacity, renewal_rate);
+        for _ in 0..capacity {
+            assert!(breaker.check());
+        }
+        assert!(!breaker.check());
+        assert_eq!(breaker.tokens, 0);
+        // Jump back in time (1 min).
+        let jump_ms = 60 * 1000;
+        breaker.last_refill -= jump_ms;
+        let expected_tokens_before_check: u8 = (renewal_rate * jump_ms as f32) as u8;
+        assert!(breaker.check());
+        assert_eq!(breaker.tokens, expected_tokens_before_check - 1);
+    }
+}