From f09fde616a22e53e581e94b0bba5d1b7ab282dcb Mon Sep 17 00:00:00 2001
From: Paul <paulmakles@gmail.com>
Date: Fri, 6 Aug 2021 17:22:55 +0100
Subject: [PATCH] Add rate limiting to message sending.

---
 Cargo.lock                          | 51 +++++++++++++++++++++++++++++
 Cargo.toml                          |  1 +
 src/main.rs                         |  1 +
 src/routes/channels/message_send.rs |  3 +-
 src/routes/channels/mod.rs          |  2 +-
 src/util/ratelimit.rs               | 21 +++++++++---
 6 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d7819764..b5fdd49c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2066,6 +2066,50 @@ dependencies = [
  "sha-1 0.8.2",
 ]
 
+[[package]]
+name = "phf"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b2ac8b67553a7ca9457ce0e526948cad581819238f4a9d1ea74545851fa24f37"
+dependencies = [
+ "phf_macros",
+ "phf_shared",
+ "proc-macro-hack",
+]
+
+[[package]]
+name = "phf_generator"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fc1437ada0f3a97d538f0bb608137bf53c53969028cab74c89893e1e9a12f0e"
+dependencies = [
+ "phf_shared",
+ "rand 0.8.3",
+]
+
+[[package]]
+name = "phf_macros"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b706f5936eb50ed880ae3009395b43ed19db5bff2ebd459c95e7bf013a89ab86"
+dependencies = [
+ "phf_generator",
+ "phf_shared",
+ "proc-macro-hack",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "phf_shared"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a68318426de33640f02be62b4ae8eb1261be2efbc337b60c54d845bf4484e0d9"
+dependencies = [
+ "siphasher",
+]
+
 [[package]]
 name = "pin-project"
 version = "0.4.27"
@@ -2635,6 +2679,7 @@ dependencies = [
  "nanoid 0.4.0",
  "num_enum",
  "once_cell",
+ "phf",
  "rauth",
  "regex",
  "reqwest",
@@ -3049,6 +3094,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "siphasher"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "729a25c17d72b06c68cb47955d44fda88ad2d3e7d77e025663fdd69b93dd71a1"
+
 [[package]]
 name = "slab"
 version = "0.4.2"
diff --git a/Cargo.toml b/Cargo.toml
index 77cc09d3..ace75f98 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -24,6 +24,7 @@ regex = "1"
 num_enum = "0.5.1"
 impl_ops = "0.1.1"
 bitfield = "0.13.2"
+phf = { version = "0.9.0", features = ["macros"] }
 
 # ID / key generation
 ulid = "0.4.1"
diff --git a/src/main.rs b/src/main.rs
index f96828e6..1c81376b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,6 @@
 #![feature(proc_macro_hygiene, decl_macro)]
 #![feature(async_closure)]
+#![feature(const_option)]
 
 #[macro_use]
 extern crate rocket;
diff --git a/src/routes/channels/message_send.rs b/src/routes/channels/message_send.rs
index dfc596f8..b49d4698 100644
--- a/src/routes/channels/message_send.rs
+++ b/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()
diff --git a/src/routes/channels/mod.rs b/src/routes/channels/mod.rs
index cc3a4d19..84918219 100644
--- a/src/routes/channels/mod.rs
+++ b/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,
diff --git a/src/util/ratelimit.rs b/src/util/ratelimit.rs
index 28e95956..f951b452 100644
--- a/src/util/ratelimit.rs
+++ b/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>;