rules: rate limiting

.devcontainer/devcontainer.json

@@ -1,5 +1,17 @@
 {
   "dockerFile": "Dockerfile",
+  "customizations": {
+    "vscode": {
+      "settings": {
+        "rust-analyzer.cargo.features": ["test-utils"],
+        "rust-analyzer.check.features": ["test-utils"],
+        "rust-analyzer.runnables.extraArgs": [
+          "--features",
+          "test-utils"
+        ]
+      }
+    }
+  },
   "extensions": [
     "rust-lang.rust-analyzer"
   ],

Cargo.toml

@@ -77,6 +77,9 @@ zeroize = { workspace = true, features = ["simd", "derive"] }
 zstd = { workspace = true }
 
 
+[features]
+test-utils = ["tokio/test-util"]
+
 
 [workspace]
 resolver = "2"

Makefile

@@ -20,6 +20,10 @@ fmt:
 check:
 	cargo check
 
+.PHONY: test
+test:
+	cargo test --features test-utils
+
 .PHONY: clean
 clean:
 	rm -rf $(DIST_DIR) captcha/dist

docs/rules.md

@@ -118,3 +118,41 @@ Valid lists types:
 - `String`
 - `Ip`
 
+## Rate limiting
+
+Algorithm used is the [sliding window](https://blog.cloudflare.com/counting-things-a-lot-of-different-things/)
+that uses request count from both current and previous period.
+
+`max` (u16) number of requests in given `period` (u16) denominated in seconds.
+Rate limiters have finite `capacity` that should be a power of 2 (otherwise
+next power of 2 will be allocated). E.g. `capacity: 1024` can store no more
+than 1024 entries in a `period`.
+
+For a case when `max` threshold is crossed Pingoo responds with [HTTP 429 Too
+Many
+Requests](https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Status/429).
+For a case where `capacity` bucket is full Pingoo responds with [HTTP 503
+Service
+Unavailable](https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Status/503).
+
+**pingoo.yml**
+```yml
+rules:
+  rate_limit_api_routes:
+    expression: http_request.path.starts_with("/api/")
+    actions:
+      - action: limit
+    limit:
+      max: 10
+      period: 60
+      capacity: 1024
+```
+
+In this example Pingoo:
+* protects resources under `/api` route
+* allows no more than 10 requests per ONE minute
+* starts returning HTTP 429 to the specific client, when number of incoming
+  requests from IP address of that client crossed the threshold of 10 in
+  sampling period of ONE minute
+* can count requests for 1024 (2^10) unique IP addresses on every minute
+* starts returning HTTP 503 to new clients if bucket is full and their IP is not in the bucket

pingoo/config/config.rs

@@ -9,13 +9,18 @@ use std::{
 use http::StatusCode;
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
-use tokio::fs;
+use tokio::{
+    fs,
+    sync::mpsc::{self, Sender},
+    task::JoinHandle,
+};
 use tracing::{debug, warn};
 
 use crate::{
     Error,
     config::config_file::{ConfigFile, RuleConfigFile, parse_service},
     lists::ListType,
+    rate_limiter::{Probe, get_rate_limit_handle},
     rules::Rule,
     service_discovery::service_registry::Upstream,
     tls::acme::LETSENCRYPT_PRODUCTION_URL,
@@ -46,6 +51,7 @@ pub struct Config {
     pub service_discovery: ServiceDiscoveryConfig,
     pub lists: HashMap<String, ListConfig>,
     pub child_process: Option<ChildProcess>,
+    pub limiter_workers: Vec<JoinHandle<()>>,
 }
 
 #[derive(Clone, Copy, Debug, Deserialize, Serialize, Eq, PartialEq)]
@@ -252,17 +258,27 @@ pub async fn load_and_validate() -> Result<Config, Error> {
         .collect();
     validate_listeners_config(&listeners, &services)?;
 
+    let mut limiter_workers: Vec<JoinHandle<()>> = vec![];
     let rules: Vec<Rule> = config_file
         .rules
         .into_iter()
         .map(|(rule_name, rule_config)| {
+            let mut limiter_tx: Option<Sender<Probe>> = None;
+            if let Some(limiter_cfg) = rule_config.limit {
+                let buffer = 1024; // todo make configurable
+                let (tx, rx) = mpsc::channel(buffer);
+                limiter_workers.push(get_rate_limit_handle(rx, limiter_cfg));
+                limiter_tx = Some(tx);
+            }
+
             Ok(Rule {
                 name: rule_name,
                 expression: rule_config
                     .expression
                     .map(|expression| rules::compile_expression(&expression))
                     .map_or(Ok(None), |r| r.map(Some))?,
                 actions: rule_config.actions,
+                limiter_tx,
             })
         })
         .collect::<Result<_, rules::Error>>()
@@ -317,6 +333,7 @@ pub async fn load_and_validate() -> Result<Config, Error> {
         service_discovery: config_file.service_discovery.unwrap_or_default(),
         lists,
         child_process: config_file.child_process,
+        limiter_workers,
     };
 
     return Ok(config);

pingoo/config/config_file.rs

@@ -98,6 +98,7 @@ pub struct ServiceConfigFileStaticNotFound {
 pub struct RuleConfigFile {
     pub expression: Option<String>,
     pub actions: Vec<rules::Action>,
+    pub limit: Option<rules::RateLimit>,
 }
 
 impl Default for ServiceConfigFileStaticNotFound {

pingoo/listeners/http_listener.rs

@@ -17,6 +17,7 @@ use crate::{
     config::ListenerConfig,
     geoip::{self, GeoipDB, GeoipRecord},
     listeners::{GRACEFUL_SHUTDOWN_TIMEOUT, Listener, accept_tcp_connection, bind_tcp_socket},
+    rate_limiter::limit_http_request,
     rules,
     services::{
         HttpService,
@@ -117,7 +118,7 @@ impl Listener for HttpListener {
     }
 }
 
-pub(super) async fn serve_http_requests<IO: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static>(
+pub async fn serve_http_requests<IO: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static>(
     tcp_stream: IO,
     services: Arc<Vec<Arc<dyn HttpService>>>,
     client_socket_addr: SocketAddr,
@@ -258,6 +259,14 @@ pub(super) async fn serve_http_requests<IO: hyper::rt::Read + hyper::rt::Write +
                                     return Ok(captcha_manager.serve_captcha());
                                 }
                             }
+                            Action::Limit {} => {
+                                // todo: if "action: limit" then this must be defined - not Option
+                                if let Some(tx) = rule.limiter_tx.clone() {
+                                    if let Some(res) = limit_http_request(client_data.ip, tx).await {
+                                        return Ok(res);
+                                    }
+                                }
+                            }
                         }
                     }
                 }

pingoo/main.rs

@@ -13,6 +13,7 @@ mod error;
 mod geoip;
 mod listeners;
 mod lists;
+mod rate_limiter;
 mod rules;
 mod serde_utils;
 mod service_discovery;