feat(iroh-relay): Rate-limit client connections (#2961)

## Description

Add a rate limit to the incoming data from client connections.

## Breaking Changes

If not configured there is now a default rate limit for incoming data
from client connections: 4KiB/s steady-stream and 16MiB burst capacity.

## Notes & open questions

- The choice here is made to rate-limit the incoming bytes, regardless
of what they are. The benefit is that the incoming stream is slowed
down, pushing back to the client over the TCP connection. The downside
is that someone who is rate-limited will get a fairly bad experience
since all DISCO traffic is also delayed.

- Only rate-limiting non-disco traffic is an option, but it would not
push back on the TCP stream, which is worse as then you'd still have to
swallow all the incoming traffic. Also it would be open to abuse fairly
easy as the detection of disco packets is based on a magic number which
could easily be spoofed.

- Maybe the `RateLimitedRelayedStream` should live in `stream.rs` next
to the `RelayedStream`? Not really sure.

### TODO

- [x] Allow rate-limit configuration in the config file.
- [x] Test config file loading.
- [x] Set a sensible default rate-limit.
- [x] Improve tests to more fully test the rate limiting.
- [x] Metrics when rate limits are hit.

## Change checklist

- [x] Self-review.
- [x] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [x] Tests if relevant.
- [x] All breaking changes documented.

iroh-relay/Cargo.toml

@@ -98,6 +98,7 @@ clap = { version = "4", features = ["derive"] }
 crypto_box = { version = "0.9.1", features = ["serde", "chacha20"] }
 proptest = "1.2.0"
 rand_chacha = "0.3.1"
+testresult = "0.4.0"
 tokio = { version = "1", features = [
     "io-util",
     "sync",

iroh-relay/src/client/conn.rs

@@ -510,7 +510,7 @@ pub(crate) async fn send_packet<S: Sink<Frame, Error = std::io::Error> + Unpin>(
     };
     if let Some(rate_limiter) = rate_limiter {
         if rate_limiter.check_n(frame.len()).is_err() {
-            tracing::warn!("dropping send: rate limit reached");
+            tracing::debug!("dropping send: rate limit reached");
             return Ok(());
         }
     }
@@ -521,12 +521,7 @@ pub(crate) async fn send_packet<S: Sink<Frame, Error = std::io::Error> + Unpin>(
 }
 
 pub(crate) struct RateLimiter {
-    inner: governor::RateLimiter<
-        governor::state::direct::NotKeyed,
-        governor::state::InMemoryState,
-        governor::clock::DefaultClock,
-        governor::middleware::NoOpMiddleware,
-    >,
+    inner: governor::DefaultDirectRateLimiter,
 }
 
 impl RateLimiter {

iroh-relay/src/main.rs

@@ -8,11 +8,11 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use anyhow::{anyhow, bail, Context as _, Result};
+use anyhow::{bail, Context as _, Result};
 use clap::Parser;
 use iroh_relay::{
     defaults::{DEFAULT_HTTPS_PORT, DEFAULT_HTTP_PORT, DEFAULT_METRICS_PORT, DEFAULT_STUN_PORT},
-    server as relay,
+    server::{self as relay, ClientConnRateLimit},
 };
 use serde::{Deserialize, Serialize};
 use tokio_rustls_acme::{caches::DirCache, AcmeConfig};
@@ -282,6 +282,29 @@ struct Limits {
     accept_conn_limit: Option<f64>,
     /// Burst limit for accepting new connection. Unlimited if not set.
     accept_conn_burst: Option<usize>,
+    /// Rate limiting configuration per client.
+    client: Option<PerClientRateLimitConfig>,
+}
+
+/// Rate limit configuration for each connected client.
+///
+/// The rate limiting uses a token-bucket style algorithm:
+///
+/// - The base rate limit uses a steady-stream rate of bytes allowed.
+/// - Additionally a burst quota allows sending bytes over this steady-stream rate
+///   limit, as long as the maximum burst quota is not exceeded.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+struct PerClientRateLimitConfig {
+    /// Rate limit configuration for the incoming data from the client.
+    rx: Option<RateLimitConfig>,
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+struct RateLimitConfig {
+    /// Maximum number of bytes per second.
+    bytes_per_second: Option<u32>,
+    /// Maximum number of bytes to read in a single burst.
+    max_burst_bytes: Option<u32>,
 }
 
 impl Config {
@@ -295,41 +318,22 @@ impl Config {
         if config_path.exists() {
             Self::read_from_file(&config_path).await
         } else {
-            let config = Config::default();
-            config.write_to_file(&config_path).await?;
-
-            Ok(config)
+            Ok(Config::default())
         }
     }
 
+    fn from_str(config: &str) -> Result<Self> {
+        toml::from_str(config).context("config must be valid toml")
+    }
+
     async fn read_from_file(path: impl AsRef<Path>) -> Result<Self> {
         if !path.as_ref().is_file() {
             bail!("config-path must be a file");
         }
         let config_ser = tokio::fs::read_to_string(&path)
             .await
             .context("unable to read config")?;
-        let config: Self = toml::from_str(&config_ser).context("config file must be valid toml")?;
-
-        Ok(config)
-    }
-
-    /// Write the content of this configuration to the provided path.
-    async fn write_to_file(&self, path: impl AsRef<Path>) -> Result<()> {
-        let p = path
-            .as_ref()
-            .parent()
-            .ok_or_else(|| anyhow!("invalid config file path, no parent"))?;
-        // TODO: correct permissions (0777 for dir, 0600 for file)
-        tokio::fs::create_dir_all(p)
-            .await
-            .with_context(|| format!("unable to create config-path dir: {}", p.display()))?;
-        let config_ser = toml::to_string(self).context("unable to serialize configuration")?;
-        tokio::fs::write(path, config_ser)
-            .await
-            .context("unable to write config file")?;
-
-        Ok(())
+        Self::from_str(&config_ser)
     }
 }
 
@@ -402,17 +406,37 @@ async fn build_relay_config(cfg: Config) -> Result<relay::ServerConfig<std::io::
         }
         None => None,
     };
-    let limits = relay::Limits {
-        accept_conn_limit: cfg
-            .limits
-            .as_ref()
-            .map(|l| l.accept_conn_limit)
-            .unwrap_or_default(),
-        accept_conn_burst: cfg
-            .limits
-            .as_ref()
-            .map(|l| l.accept_conn_burst)
-            .unwrap_or_default(),
+    let limits = match cfg.limits {
+        Some(ref limits) => {
+            let client_rx = match &limits.client {
+                Some(PerClientRateLimitConfig { rx: Some(rx) }) => {
+                    if rx.bytes_per_second.is_none() && rx.max_burst_bytes.is_some() {
+                        bail!("bytes_per_seconds must be specified to enable the rate-limiter");
+                    }
+                    match rx.bytes_per_second {
+                        Some(bps) => Some(ClientConnRateLimit {
+                            bytes_per_second: bps
+                                .try_into()
+                                .context("bytes_per_second must be non-zero u32")?,
+                            max_burst_bytes: rx
+                                .max_burst_bytes
+                                .map(|v| {
+                                    v.try_into().context("max_burst_bytes must be non-zero u32")
+                                })
+                                .transpose()?,
+                        }),
+                        None => None,
+                    }
+                }
+                Some(PerClientRateLimitConfig { rx: None }) | None => None,
+            };
+            relay::Limits {
+                accept_conn_limit: limits.accept_conn_limit,
+                accept_conn_burst: limits.accept_conn_burst,
+                client_rx,
+            }
+        }
+        None => Default::default(),
     };
     let relay_config = relay::RelayConfig {
         http_bind_addr: cfg.http_bind_addr(),
@@ -477,3 +501,46 @@ mod metrics {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::num::NonZeroU32;
+
+    use testresult::TestResult;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_rate_limit_config() -> TestResult {
+        let config = "
+            [limits.client.rx]
+            bytes_per_second = 400
+            max_burst_bytes = 800
+        ";
+        let config = Config::from_str(config)?;
+        let relay_config = build_relay_config(config).await?;
+
+        let relay = relay_config.relay.expect("no relay config");
+        assert_eq!(
+            relay.limits.client_rx.expect("ratelimit").bytes_per_second,
+            NonZeroU32::try_from(400).unwrap()
+        );
+        assert_eq!(
+            relay.limits.client_rx.expect("ratelimit").max_burst_bytes,
+            Some(NonZeroU32::try_from(800).unwrap())
+        );
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_rate_limit_default() -> TestResult {
+        let config = Config::from_str("")?;
+        let relay_config = build_relay_config(config).await?;
+
+        let relay = relay_config.relay.expect("no relay config");
+        assert!(relay.limits.client_rx.is_none());
+
+        Ok(())
+    }
+}

iroh-relay/src/protos/disco.rs

@@ -10,6 +10,7 @@ pub(crate) const MAGIC_LEN: usize = MAGIC.as_bytes().len();
 pub(crate) const KEY_LEN: usize = 32;
 
 const MESSAGE_HEADER_LEN: usize = MAGIC_LEN + KEY_LEN;
+
 /// Reports whether p looks like it's a packet containing an encrypted disco message.
 pub fn looks_like_disco_wrapper(p: &[u8]) -> bool {
     if p.len() < MESSAGE_HEADER_LEN {

iroh-relay/src/protos/relay.rs

@@ -18,6 +18,9 @@ use tokio_util::codec::{Decoder, Encoder};
 /// including its on-wire framing overhead)
 pub const MAX_PACKET_SIZE: usize = 64 * 1024;
 
+/// The maximum frame size.
+///
+/// This is also the minimum burst size that a rate-limiter has to accept.
 const MAX_FRAME_SIZE: usize = 1024 * 1024;
 
 /// The Relay magic number, sent in the FrameType::ClientInfo frame upon initial connection.
@@ -200,9 +203,14 @@ pub(crate) async fn recv_client_key<S: Stream<Item = anyhow::Result<Frame>> + Un
     }
 }
 
+/// The protocol for the relay server.
+///
+/// This is a framed protocol, using [`tokio_util::codec`] to turn the streams of bytes into
+/// [`Frame`]s.
 #[derive(Debug, Default, Clone)]
 pub(crate) struct DerpCodec;
 
+/// The frames in the [`DerpCodec`].
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) enum Frame {
     ClientInfo {
@@ -279,6 +287,12 @@ impl Frame {
         }
     }
 
+    /// Serialized length with frame header.
+    #[cfg(feature = "server")]
+    pub(crate) fn len_with_header(&self) -> usize {
+        self.len() + HEADER_LEN
+    }
+
     /// Tries to decode a frame received over websockets.
     ///
     /// Specifically, bytes received from a binary websocket message frame.

iroh-relay/src/server.rs

@@ -16,7 +16,7 @@
 //! - HTTPS `/generate_204`: Used for net_report probes.
 //! - STUN: UDP port for STUN requests/responses.
 
-use std::{fmt, future::Future, net::SocketAddr, pin::Pin, sync::Arc};
+use std::{fmt, future::Future, net::SocketAddr, num::NonZeroU32, pin::Pin, sync::Arc};
 
 use anyhow::{anyhow, bail, Context, Result};
 use futures_lite::StreamExt;
@@ -140,12 +140,24 @@ pub struct TlsConfig<EC: fmt::Debug, EA: fmt::Debug = EC> {
 }
 
 /// Rate limits.
+// TODO: accept_conn_limit and accept_conn_burst are not currently implemented.
 #[derive(Debug, Default)]
 pub struct Limits {
     /// Rate limit for accepting new connection. Unlimited if not set.
     pub accept_conn_limit: Option<f64>,
     /// Burst limit for accepting new connection. Unlimited if not set.
     pub accept_conn_burst: Option<usize>,
+    /// Rate limits for incoming traffic from a client connection.
+    pub client_rx: Option<ClientConnRateLimit>,
+}
+
+/// Per-client rate limit configuration.
+#[derive(Debug, Copy, Clone)]
+pub struct ClientConnRateLimit {
+    /// Max number of bytes per second to read from the client connection.
+    pub bytes_per_second: NonZeroU32,
+    /// Max number of bytes to read in a single burst.
+    pub max_burst_bytes: Option<NonZeroU32>,
 }
 
 /// TLS certificate configuration.
@@ -260,6 +272,9 @@ impl Server {
                     .request_handler(Method::GET, "/index.html", Box::new(root_handler))
                     .request_handler(Method::GET, RELAY_PROBE_PATH, Box::new(probe_handler))
                     .request_handler(Method::GET, "/robots.txt", Box::new(robots_handler));
+                if let Some(cfg) = relay_config.limits.client_rx {
+                    builder = builder.client_rx_ratelimit(cfg);
+                }
                 let http_addr = match relay_config.tls {
                     Some(tls_config) => {
                         let server_config = rustls::ServerConfig::builder_with_provider(Arc::new(

iroh-relay/src/server/actor.rs

@@ -183,12 +183,8 @@ impl Actor {
             }
             Message::CreateClient(client_builder) => {
                 inc!(Metrics, accepts);
-
-                trace!(
-                    node_id = client_builder.node_id.fmt_short(),
-                    "create client"
-                );
                 let node_id = client_builder.node_id;
+                trace!(node_id = node_id.fmt_short(), "create client");
 
                 // build and register client, starting up read & write loops for the client
                 // connection
@@ -272,6 +268,7 @@ mod tests {
                 stream: RelayedStream::Derp(Framed::new(MaybeTlsStream::Test(io), DerpCodec)),
                 write_timeout: Duration::from_secs(1),
                 channel_capacity: 10,
+                rate_limit: None,
                 server_channel,
             },
             Framed::new(test_io, DerpCodec),