Merge pull request #79 from Zomatree/ratelimits

src/util/ratelimit.rs

@@ -1,22 +1,180 @@
-use rocket_governor::{Method, Quota, RocketGovernable, RocketGovernor, NonZeroU32};
-use phf::phf_map;
+use rocket::{async_trait, http::Status, request::{Outcome, FromRequest}, response};
+use std::{collections::{HashMap, hash_map::DefaultHasher}, time};
+use rauth::entities::Session;
+use crate::util::result::Error;
+use std::hash::{Hash, Hasher};
+use std::sync::{Arc, Mutex};
 
-pub struct RateLimitGuard;
+#[inline]
+fn now() -> f64 {
+    // will this ever actually panic?
+    time::SystemTime::now().duration_since(time::SystemTime::UNIX_EPOCH).unwrap().as_secs_f64()
+}
 
-static ROUTE_QUOTAS: phf::Map<&'static str, &'static Quota> = phf_map! {
-    "message_send" => &Quota::per_second(NonZeroU32::new(1u32).unwrap())
-        .allow_burst(NonZeroU32::new(10u32).unwrap())
-};
+struct Ratelimit {
+    pub rate: u8,
+    pub per: u8,
+    window: f64,
+    tokens: u8,
+    pub last: f64,
+}
 
-impl<'r> RocketGovernable<'r> for RateLimitGuard {
-    fn quota(_method: Method, route: &str) -> Quota {
-        if let Some(q) = ROUTE_QUOTAS.get(route) {
-            **q
+impl Ratelimit {
+    pub fn new(rate: u8, per: u8) -> Self {
+        Self {
+            rate,
+            per,
+            window: 0f64,
+            tokens: rate,
+            last: 0.0
+        }
+    }
+
+    fn get_tokens(&self, current: Option<f64>) -> u8 {
+        let current = current.unwrap_or(now());
+        if current > (self.window + self.per as f64) {
+            self.rate
         } else {
-            Quota::per_second(Self::nonzero(1u32))
-                .allow_burst(Self::nonzero(5u32))
+            self.tokens
+        }
+    }
+
+    fn update_ratelimit(&mut self) -> Option<f64> {
+        let current = now();
+        self.last = current;
+
+        self.tokens = self.get_tokens(Some(current));
+
+        if self.tokens == self.rate {
+            self.window = current
+        }
+
+        if self.tokens == 0 {
+            return Some(self.per as f64 - (current - self.window))
+        }
+
+        self.tokens -= 1;
+
+        if self.tokens == 0 {
+            self.window = current
+        }
+
+        None
+    }
+}
+
+impl Clone for Ratelimit {
+    fn clone(&self) -> Ratelimit {
+        Ratelimit::new(self.rate, self.per)
+    }
+}
+
+#[derive(Clone)]
+struct RatelimitMapping {
+    cache: HashMap<u64, Ratelimit>,
+    cooldown: Ratelimit
+}
+
+impl RatelimitMapping {
+    pub fn new(rate: u8, per: u8) -> Self {
+        RatelimitMapping {
+            cache: HashMap::new(),
+            cooldown: Ratelimit::new(rate, per)
+        }
+    }
+
+    fn verify_cache(&mut self) {
+        let current = now();
+        self.cache.retain(|_, v| current < v.last + v.per as f64);
+    }
+
+    pub fn get_bucket(&mut self, key: u64) -> &mut Ratelimit {
+        self.verify_cache();
+        self.cache.entry(key).or_insert(self.cooldown.clone())
+    }
+}
+
+#[derive(Copy, Clone)]
+pub struct Ratelimiter {
+    bucket: u64,
+    limit: u8,
+    remaining: u8,
+    reset: f64
+}
+
+pub struct RatelimitState(Arc<Mutex<HashMap<&'static str, RatelimitMapping>>>);
+
+impl RatelimitState {
+    pub fn new() -> Self {
+        let mut hashmap = HashMap::new();
+        hashmap.insert("message_send", RatelimitMapping::new(10, 10));
+        let arc = Arc::new(Mutex::new(hashmap));
+        RatelimitState(arc)
+    }
+}
+
+#[async_trait]
+impl<'r> FromRequest<'r> for Ratelimiter {
+    type Error = Error;
+
+    async fn from_request(request: &'r rocket::Request<'_>) -> Outcome<Self, Self::Error> {
+        let res = request.local_cache_async(async {
+            if let Some(route) = request.route() {
+                if let Some(route_name) = &route.name {
+                    let session = request.guard::<Session>().await.unwrap();
+                    let state = request.guard::<&rocket::State<RatelimitState>>().await.unwrap().inner();
+                    let arc = Arc::clone(&state.0);
+                    let mut mutex = arc.lock().unwrap();
+                    let mapping = mutex.get_mut(route_name.as_ref()).unwrap();
+
+                    // create a unique key tied to the user id and route they use
+                    let key = format!("{}:{}:{}", session.user_id, route.method.as_str(), route.uri.as_str());
+                    let mut hasher = DefaultHasher::new();
+
+                    key.hash(&mut hasher);
+                    let hashed = hasher.finish();
+
+                    let bucket = mapping.get_bucket(hashed);
+                    let ret = bucket.update_ratelimit();
+
+                    if let Some(retry_after) = ret {
+                        Err(Error::TooManyRequests { retry_after })
+                    } else {
+                        Ok(Ratelimiter {
+                            bucket: hashed,
+                            limit: bucket.rate,
+                            remaining: bucket.get_tokens(None),
+                            reset: bucket.window + (bucket.per as f64)
+                        })
+                    }
+                } else {
+                    unreachable!()
+                }
+            } else {
+                unreachable!()
+            }
+        }).await;
+
+        match res {
+            Ok(rl) => Outcome::Success(*rl),
+            Err(e) => Outcome::Failure((Status::TooManyRequests, e.clone()))
         }
     }
 }
 
-pub type RateLimited<'a> = RocketGovernor<'a, RateLimitGuard>;
+pub struct RatelimitResponse<R>(pub R);
+
+impl<'r, 'o: 'r, R: response::Responder<'r, 'o>> response::Responder<'r, 'o> for RatelimitResponse<R> {
+    fn respond_to(self, request: &'r rocket::Request<'_>) -> response::Result<'o> {
+        let res: &Result<Ratelimiter, Error> = request.local_cache(|| unreachable!());
+        let ratelimiter = res.as_ref().unwrap();
+
+        Ok(response::Builder::new(rocket::Response::new())
+            .raw_header("x-ratelimit-bucket", ratelimiter.bucket.to_string())
+            .raw_header("x-ratelimit-limit", ratelimiter.limit.to_string())
+            .raw_header("x-ratelimit-remaining", ratelimiter.remaining.to_string())
+            .raw_header("x-ratelimit-reset", ratelimiter.reset.to_string())
+            .merge(self.0.respond_to(request)?)
+            .finalize())
+    }
+}

src/util/result.rs

@@ -6,7 +6,7 @@ use serde::Serialize;
 use std::io::Cursor;
 use validator::ValidationErrors;
 
-#[derive(Serialize, Debug)]
+#[derive(Serialize, Debug, Clone)]
 #[serde(tag = "type")]
 pub enum Error {
     LabelMe,
@@ -69,6 +69,9 @@ pub enum Error {
     VosoUnavailable,
     NotFound,
     NoEffect,
+    TooManyRequests {
+        retry_after: f64
+    }
 }
 
 pub struct EmptyResponse;
@@ -131,6 +134,7 @@ impl<'r> Responder<'r, 'static> for Error {
             Error::VosoUnavailable => Status::BadRequest,
             Error::NotFound => Status::NotFound,
             Error::NoEffect => Status::Ok,
+            Error::TooManyRequests { .. } => Status::TooManyRequests,
         };
 
         // Serialize the error data structure into JSON.