Skip to content

Commit

Permalink
feat(iroh-relay): Rate-limit client connections
Browse files Browse the repository at this point in the history
This has a hardcoded rate-limiter.  Need to make it configurable.
  • Loading branch information
flub committed Nov 22, 2024
1 parent a5e9283 commit 070931d
Show file tree
Hide file tree
Showing 10 changed files with 277 additions and 45 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions iroh-relay/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
9 changes: 2 additions & 7 deletions iroh-relay/src/client/conn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(());
}
}
Expand All @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions iroh-relay/src/protos/disco.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
13 changes: 13 additions & 0 deletions iroh-relay/src/protos/relay.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down
6 changes: 1 addition & 5 deletions iroh-relay/src/server/actor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit 070931d

Please sign in to comment.