motis-proxy: Implement route rate limit

ansible/roles/motis-proxy/files/Rocket.toml

@@ -2,9 +2,11 @@
 #
 # SPDX-License-Identifier: AGPL-3.0-or-later
 
+[release]
+ip_header = "X-Forwarded-For"
+
 [release.proxy]
 # TODO: Add /address and /ppr/route once we have it set up
 allowed_endpoints = [ "/intermodal", "/guesser",
                       "/railviz/get_trains", "/railviz/get_trips", "/railviz/get_station",
                       "/lookup/schedule_info", "/gbfs/info", "/trip_to_connection" ]
-

ansible/roles/nginx/templates/transitous.conf.j2

@@ -41,6 +41,8 @@ server {
     # MOTIS Proxy
     location /api/ {
         proxy_pass http://localhost:8000/;
+
+        proxy_set_header X-Forwarded-For $remote_addr;
     }
 
     # API Documentation

motis-proxy/Cargo.toml

@@ -18,3 +18,4 @@ rocket_cors = { version = "0.6.0", default-features = false }
 serde_default = "0.1"
 log = "0.4"
 rocket_okapi = { version = "0.8", features = [ "rapidoc" ] }
+lru = "0.12"

motis-proxy/Rocket.toml

@@ -15,3 +15,4 @@ allowed_endpoints = [ "/intermodal", "/guesser", "/address",
 [debug.proxy]
 proxy_assets = true
 motis_address = "https://europe.motis-project.de"
+routes_per_minute_limit = 20

motis-proxy/src/main.rs

@@ -5,6 +5,8 @@
 #[macro_use]
 extern crate rocket;
 
+mod rate_limit;
+
 use log::{trace, warn};
 use reqwest::{Client, StatusCode};
 use rocket::{
@@ -18,7 +20,9 @@ use rocket_cors::{AllowedOrigins, CorsOptions};
 use serde::{Deserialize, Serialize};
 use serde_default::DefaultFromSerde;
 
-use std::time::Duration;
+use std::{net::IpAddr, num::NonZeroUsize, time::Duration};
+
+use rate_limit::IpRateLimit;
 
 // For documetation generation
 use rocket_okapi::okapi::schemars;
@@ -43,6 +47,12 @@ fn default_proxy_assets() -> bool {
 fn default_allowed_endpoints() -> Option<Vec<Endpoint>> {
     None
 }
+fn default_routes_per_minute_limit() -> u16 {
+    20
+}
+fn default_lru_rate_limit_entries() -> usize {
+    10_000
+}
 
 #[derive(Deserialize, DefaultFromSerde)]
 struct Config {
@@ -62,6 +72,12 @@ struct Config {
     /// If this option is not set, all known endpoints will be allowed.
     #[serde(default = "default_allowed_endpoints")]
     allowed_endpoints: Option<Vec<Endpoint>>,
+
+    #[serde(default = "default_routes_per_minute_limit")]
+    routes_per_minute_limit: u16,
+
+    #[serde(default = "default_lru_rate_limit_entries")]
+    lru_rate_limit_entries: usize,
 }
 
 #[derive(Deserialize, Serialize, PartialEq, Eq, JsonSchema)]
@@ -417,6 +433,8 @@ async fn proxy_api(
     request: Json<Request>,
     http_client: &State<Client>,
     config: &State<Config>,
+    route_rate_limit: &State<IpRateLimit>,
+    remote_address: IpAddr,
 ) -> ResultResponse<Custom<Json<serde_json::Value>>> {
     let request = request.into_inner();
 
@@ -431,6 +449,17 @@ async fn proxy_api(
         }
     }
 
+    // Check if routing limit was exceeded
+    if matches!(
+        request.content,
+        RequestContent::IntermodalConnectionRequest(_)
+            | RequestContent::IntermodalRoutingRequest(_)
+    ) {
+        if route_rate_limit.should_limit(&remote_address) {
+            return Err(Custom(Status::TooManyRequests, ()))
+        }
+    }
+
     trace!("MOTIS Request <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<");
     trace!("{}", serde_json::to_string_pretty(&request).unwrap());
 
@@ -548,6 +577,11 @@ fn rocket() -> _ {
         )
         .attach(cors.clone())
         .manage(cors)
+        .manage(IpRateLimit::new(
+            NonZeroUsize::new(config.lru_rate_limit_entries)
+                .expect("lru_rate_limit_entries must not be zero"),
+            config.routes_per_minute_limit
+        ))
         .manage(config)
         .mount("/", routes)
         .mount("/", rocket_cors::catch_all_options_routes())

motis-proxy/src/rate_limit.rs

@@ -0,0 +1,55 @@
+// SPDX-FileCopyrightText: 2024 Jonah Brüchert <[email protected]>
+//
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+use lru::LruCache;
+use std::{net::IpAddr, num::NonZeroUsize, time::Duration, sync::{Arc, Mutex}};
+
+use std::time::Instant;
+
+struct Inner {
+    requests: LruCache<IpAddr, u16>,
+    rate_limit: u16,
+    last_cleared: Instant,
+}
+
+pub struct IpRateLimit {
+    inner: Arc<Mutex<Inner>>
+}
+
+impl IpRateLimit {
+    pub fn new(size: NonZeroUsize, rate_limit: u16) -> IpRateLimit {
+        IpRateLimit {
+            inner: Arc::new(Mutex::new(Inner {
+                requests: LruCache::new(size),
+                rate_limit,
+                last_cleared: Instant::now(),
+            }))
+        }
+    }
+
+    pub fn should_limit(&self, ip: &IpAddr) -> bool {
+        let mut inner = self.inner.lock().unwrap();
+
+        // Check if data is older than a minute, then delete it
+        if inner.last_cleared.elapsed() >= Duration::from_secs(60) {
+            inner.requests.clear();
+            inner.last_cleared = Instant::now();
+        }
+
+        // Incerement existing request counter for address or create a new one
+        if let Some(count) = inner.requests.get_mut(ip) {
+            *count += 1;
+        } else {
+            inner.requests.put(*ip, 1);
+        }
+
+
+        // Check if ip has reached rate limit
+        let rate_limit = inner.rate_limit;
+        inner.requests
+            .get(&ip)
+            .map(|count| *count >= rate_limit)
+            .unwrap_or(false)
+    }
+}