Add rate limiting to message sending.

src/main.rs

@@ -1,5 +1,6 @@
 #![feature(proc_macro_hygiene, decl_macro)]
 #![feature(async_closure)]
+#![feature(const_option)]
 
 #[macro_use]
 extern crate rocket;

src/routes/channels/message_send.rs

@@ -1,6 +1,7 @@
 use std::collections::HashSet;
 
 use crate::database::*;
+use crate::util::ratelimit::RateLimited;
 use crate::util::result::{Error, Result};
 
 use mongodb::{bson::doc, options::FindOneOptions};
@@ -33,7 +34,7 @@ lazy_static! {
 }
 
 #[post("/<target>/messages", data = "<message>")]
-pub async fn req(user: User, target: Ref, message: Json<Data>) -> Result<Value> {
+pub async fn message_send(_r: RateLimited<'_>, user: User, target: Ref, message: Json<Data>) -> Result<Value> {
     let message = message.into_inner();
     message
         .validate()

src/routes/channels/mod.rs

@@ -28,7 +28,7 @@ pub fn routes() -> Vec<Route> {
         channel_delete::req,
         channel_edit::req,
         invite_create::req,
-        message_send::req,
+        message_send::message_send,
         message_query::req,
         message_search::req,
         message_query_stale::req,

src/util/ratelimit.rs

@@ -1,11 +1,22 @@
-use rocket_governor::{Method, Quota, RocketGovernable};
+use rocket_governor::{Method, Quota, RocketGovernable, RocketGovernor, NonZeroU32};
+use phf::phf_map;
 
 pub struct RateLimitGuard;
 
+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())
+};
+
 impl<'r> RocketGovernable<'r> for RateLimitGuard {
-    fn quota(_method: Method, _route_name: &str) -> Quota {
-        dbg!(_method, _route_name);
-        Quota::per_second(Self::nonzero(1u32))
-            .allow_burst(Self::nonzero(10u32))
+    fn quota(_method: Method, route: &str) -> Quota {
+        if let Some(q) = ROUTE_QUOTAS.get(route) {
+            **q
+        } else {
+            Quota::per_second(Self::nonzero(1u32))
+                .allow_burst(Self::nonzero(5u32))
+        }
     }
 }
+
+pub type RateLimited<'a> = RocketGovernor<'a, RateLimitGuard>;