PoC

src/api/ruma_wrapper/axum.rs

@@ -1,4 +1,8 @@
-use std::{collections::BTreeMap, iter::FromIterator, str};
+use std::{
+    collections::BTreeMap,
+    iter::FromIterator,
+    str::{self},
+};
 
 use axum::{
     async_trait,
@@ -23,7 +27,10 @@ use serde::Deserialize;
 use tracing::{debug, error, warn};
 
 use super::{Ruma, RumaResponse};
-use crate::{service::appservice::RegistrationInfo, services, Error, Result};
+use crate::{
+    service::{appservice::RegistrationInfo, rate_limiting::Target},
+    services, Error, Result,
+};
 
 enum Token {
     Appservice(Box<RegistrationInfo>),
@@ -95,6 +102,44 @@ where
             Token::None
         };
 
+        // doesn't work when Conduit is behind proxy
+        // let remote_addr: ConnectInfo<SocketAddr> = parts.extract().await?;
+
+        let target = match &token {
+            Token::User((user_id, _)) => Some(Target::User(user_id.clone())),
+            Token::None => {
+                let header = parts
+                    .headers
+                    .get("x-forwarded-for")
+                    .ok_or_else(|| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting."))?;
+
+                let s = header
+                    .to_str()
+                    .map_err(|_| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting."))?;
+                Some(
+                    s.parse()
+                        .map(Target::Ip)
+                        .map_err(|_| Error::BadRequest(ErrorKind::Unauthorized, "Rate limiting.")),
+                )
+                .transpose()?
+            }
+            _ => None,
+        };
+
+        if let Err(retry_after_ms) = target.map_or(Ok(()), |t| {
+            let key = (t, (&metadata).into());
+
+            services()
+                .rate_limiting
+                .update_or_reject(&key)
+                .map_err(Some)
+        }) {
+            return Err(Error::BadRequest(
+                ErrorKind::LimitExceeded { retry_after_ms },
+                "Rate limit exceeded.",
+            ));
+        }
+
         let mut json_body = serde_json::from_slice::<CanonicalJsonValue>(&body).ok();
 
         let (sender_user, sender_device, sender_servername, appservice_info) =