From 613107e7cf511163cffecc57b29edc6f8f89fadd Mon Sep 17 00:00:00 2001
From: Matthias Ahouansou <matthias@ahouansou.cz>
Date: Thu, 11 Jul 2024 13:14:10 +0100
Subject: [PATCH] add rate limiting for registration token validity

---
 src/config/mod.rs                | 4 ++--
 src/service/rate_limiting/mod.rs | 9 ++++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 2d08d855..105b265a 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -105,13 +105,13 @@ pub struct WellKnownConfig {
 }
 
 #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd)]
-#[serde(rename_all = "lowercase")]
+#[serde(rename_all = "snake_case")]
 pub enum Restriction {
     Registration,
     Login,
+    RegistrationTokenValidity,
 
     #[default]
-    #[serde(rename = "")]
     CatchAll,
 }
 
diff --git a/src/service/rate_limiting/mod.rs b/src/service/rate_limiting/mod.rs
index 3900e843..e4c4e201 100644
--- a/src/service/rate_limiting/mod.rs
+++ b/src/service/rate_limiting/mod.rs
@@ -7,7 +7,10 @@ use std::{
 use dashmap::{mapref::entry::Entry, DashMap};
 use ruma::{
     api::{
-        client::{account::register, session::login},
+        client::{
+            account::{check_registration_token_validity, register},
+            session::login,
+        },
         IncomingRequest, Metadata,
     },
     OwnedServerName, OwnedUserId,
@@ -20,6 +23,10 @@ impl From<Metadata> for Restriction {
         [
             (register::v3::Request::METADATA, Restriction::Registration),
             (login::v3::Request::METADATA, Restriction::Login),
+            (
+                check_registration_token_validity::v1::Request::METADATA,
+                Restriction::RegistrationTokenValidity,
+            ),
         ]
         .into_iter()
         .find(|(other, _)| {