Initial rate limiting impl based on the leaky-bucket and concread…

… crates (#67)

This implements an initial set of three rate limiting rules, broken up into two groups:

* "multi" rules, where we keep a bounded set of key:bucket pairs
* "single" rules, where we keep a single bucket

Rules can be provided that match these rules, requiring requests to match all provided rules.

---------

Co-authored-by: Brandon Pitman <[email protected]>

source/river/Cargo.toml

@@ -24,17 +24,19 @@ rustdoc-args = ["--cfg", "doc_cfg"]
 
 [dependencies]
 async-trait = "0.1.79"
+cidr = "0.2.3"
+concread = "0.5.3"
+futures-util = "0.3.30"
+http = "1.0.0"
+kdl = "4.6.0"
+leaky-bucket = "1.1.2"
 log = "0.4.21"
+miette = { version = "5.10.0", features = ["fancy"] }
 regex = "1.10.4"
+thiserror = "1.0.61"
 tokio = "1.37.0" # TODO: check for implicit feature usage
 toml = "0.8.12"
 tracing = "0.1.40"
-kdl = "4.6.0"
-miette = { version = "5.10.0", features = ["fancy"] }
-thiserror = "1.0.61"
-http = "1.0.0"
-futures-util = "0.3.30"
-cidr = "0.2.3"
 
 [dependencies.static-files-module]
 version = "0.2"

source/river/assets/test-config.kdl

@@ -31,6 +31,45 @@ services {
             "0.0.0.0:4443" cert-path="./assets/test.crt" key-path="./assets/test.key" offer-h2=true
         }
 
+        // Apply Rate limiting to this service
+        //
+        // Note that ALL rules are applied, and a request must receive a token from all
+        // applicable rules.
+        //
+        // For example:
+        //
+        // A request to URI `/index.html` from IP 1.2.3.4 will only need to get a token from
+        // the `source-ip` rule.
+        //
+        // A request to URI `/static/style.css` from IP 1.2.3.4 will need to get a token from
+        // BOTH the `source-ip` rule (from the `1.2.3.4` bucket), AND the `specific-uri` rule
+        // (from the `/static/style.css` bucket)
+        rate-limiting {
+            // This rate limiting rule is based on the source IP address
+            //
+            // * Up to the last 4000 IP addresses will be remembered
+            // * Each IP address can make a burst of 10 requests
+            // * The bucket for each IP will refill at a rate of 1 request per 10 milliseconds
+            rule kind="source-ip" \
+                max-buckets=4000 tokens-per-bucket=10 refill-qty=1 refill-rate-ms=10
+
+            // This rate limiting is based on the specific URI path
+            //
+            // * Up to the last 2000 URI paths will be remembered
+            // * Each URI path can make a burst of 20 requests
+            // * The bucket for each URI will refill at a rate of 5 requests per 1 millisecond
+            rule kind="specific-uri" pattern="static/.*" \
+                max-buckets=2000 tokens-per-bucket=20 refill-qty=5 refill-rate-ms=1
+
+            // This rate limiting is based on ANY URI paths that match the pattern
+            //
+            // * A single bucket will be used for all URIs that match the pattern
+            // * We allow a burst of up to 50 requests for any MP4 files
+            // * The bucket for all MP4 files will refill at a rate of 2 requests per 3 milliseconds
+            rule kind="any-matching-uri" pattern=r".*\.mp4" \
+                tokens-per-bucket=50 refill-qty=2 refill-rate-ms=3
+        }
+
         // Connectors are the "upstream" interfaces that we connect with. We can name as many
         // as we'd like, at least one is required. By default, connectors are distributed
         // round-robin.
@@ -52,7 +91,7 @@ services {
         // This section is optional.
         path-control {
             request-filters {
-                filter kind="block-cidr-range" addrs="192.168.0.0/16, 10.0.0.0/8, 2001:0db8::0/32, 127.0.0.1"
+                filter kind="block-cidr-range" addrs="192.168.0.0/16, 10.0.0.0/8, 2001:0db8::0/32"
             }
             upstream-request {
                 filter kind="remove-header-key-regex" pattern=".*(secret|SECRET).*"

source/river/src/config/internal.rs

@@ -14,7 +14,10 @@ use pingora::{
 };
 use tracing::warn;
 
-use crate::proxy::request_selector::{null_selector, RequestSelector};
+use crate::proxy::{
+    rate_limiting::AllRateConfig,
+    request_selector::{null_selector, RequestSelector},
+};
 
 /// River's internal configuration
 #[derive(Debug, Clone)]
@@ -109,6 +112,12 @@ impl Config {
     }
 }
 
+///
+#[derive(Debug, Default, Clone, PartialEq)]
+pub struct RateLimitingConfig {
+    pub(crate) rules: Vec<AllRateConfig>,
+}
+
 /// Add Path Control Modifiers
 ///
 /// Note that we use `BTreeMap` and NOT `HashMap`, as we want to maintain the
@@ -141,6 +150,7 @@ pub struct ProxyConfig {
     pub(crate) upstream_options: UpstreamOptions,
     pub(crate) upstreams: Vec<HttpPeer>,
     pub(crate) path_control: PathControl,
+    pub(crate) rate_limiting: RateLimitingConfig,
 }
 
 #[derive(Debug, PartialEq, Clone)]