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

iroh-relay/Cargo.toml

@@ -99,6 +99,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

@@ -5,14 +5,15 @@
 
 use std::{
     net::{Ipv6Addr, SocketAddr},
+    num::NonZeroU32,
     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 +283,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 +319,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 +407,32 @@ 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) }) => {
+                    let mut cfg = ClientConnRateLimit::default();
+                    if let Some(bps) = rx.bytes_per_second {
+                        let v = NonZeroU32::try_from(bps)
+                            .context("bytes_per_second must be non-zero u32")?;
+                        cfg.bytes_per_second = v;
+                    }
+                    if let Some(burst) = rx.max_burst_bytes {
+                        let v = NonZeroU32::try_from(burst)
+                            .context("max_burst_bytes must be non-zero u32")?;
+                        cfg.max_burst_bytes = v;
+                    }
+                    cfg
+                }
+                Some(PerClientRateLimitConfig { rx: None }) | None => Default::default(),
+            };
+            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 +497,33 @@ mod metrics {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    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.bytes_per_second,
+            NonZeroU32::try_from(400).unwrap()
+        );
+        assert_eq!(
+            relay.limits.client_rx.max_burst_bytes,
+            NonZeroU32::try_from(800).unwrap()
+        );
+
+        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,11 @@ impl Frame {
         }
     }
 
+    /// Serialized length with frame header.
+    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,44 @@ 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: 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.
+    ///
+    /// Defaults to 4KiB/s.
+    pub bytes_per_second: NonZeroU32,
+    /// Max number of bytes to read in a single burst.
+    ///
+    /// Defaults to 16 MiB.
+    pub max_burst_bytes: NonZeroU32,
+}
+
+impl ClientConnRateLimit {
+    pub(super) const MAX: ClientConnRateLimit = ClientConnRateLimit {
+        bytes_per_second: NonZeroU32::MAX,
+        max_burst_bytes: NonZeroU32::MAX,
+    };
+}
+
+impl Default for ClientConnRateLimit {
+    fn default() -> Self {
+        Self {
+            bytes_per_second: NonZeroU32::try_from(1024 * 4).expect("nonzero"),
+            max_burst_bytes: NonZeroU32::try_from(1024 * 1024 * 16).expect("nonzero"),
+        }
+    }
 }
 
 /// TLS certificate configuration.
@@ -255,6 +287,7 @@ impl Server {
                     None => relay_config.http_bind_addr,
                 };
                 let mut builder = http_server::ServerBuilder::new(relay_bind_addr)
+                    .client_rx_ratelimit(relay_config.limits.client_rx)
                     .headers(headers)
                     .request_handler(Method::GET, "/", Box::new(root_handler))
                     .request_handler(Method::GET, "/index.html", Box::new(root_handler))

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
@@ -258,6 +254,7 @@ mod tests {
         server::{
             client_conn::ClientConnConfig,
             streams::{MaybeTlsStream, RelayedStream},
+            ClientConnRateLimit,
         },
     };
 
@@ -272,6 +269,7 @@ mod tests {
                 stream: RelayedStream::Derp(Framed::new(MaybeTlsStream::Test(io), DerpCodec)),
                 write_timeout: Duration::from_secs(1),
                 channel_capacity: 10,
+                rate_limit: ClientConnRateLimit::MAX,
                 server_channel,
             },
             Framed::new(test_io, DerpCodec),