protocol/limiter: impl AsTokens for limits

Provide `AsTokens` implementations for the `Outbound` and `Inbound`
newtypes of `RateLimit`.

This allows them to be used for the rate limiting method.
This commit is contained in:
Fintan Halpenny 2025-09-08 16:11:17 +01:00 committed by Lorenz Leutgeb
parent c675683da9
commit 58305cda32
3 changed files with 40 additions and 6 deletions

View File

@ -1313,7 +1313,7 @@ where
Err(e) => error!(target: "service", "Error querying ban status for {ip}: {e}"), Err(e) => error!(target: "service", "Error querying ban status for {ip}: {e}"),
} }
let host: HostName = ip.into(); let host: HostName = ip.into();
let tokens = RateLimit::from(self.config.limits.rate.inbound.clone()); let tokens = self.config.limits.rate.inbound;
if self.limiter.limit(host.clone(), None, &tokens, self.clock) { if self.limiter.limit(host.clone(), None, &tokens, self.clock) {
trace!(target: "service", "Rate limiting inbound connection from {host}.."); trace!(target: "service", "Rate limiting inbound connection from {host}..");
@ -1819,8 +1819,8 @@ where
peer.last_active = self.clock; peer.last_active = self.clock;
let limit: RateLimit = match peer.link { let limit: RateLimit = match peer.link {
Link::Outbound => self.config.limits.rate.outbound.clone().into(), Link::Outbound => self.config.limits.rate.outbound.into(),
Link::Inbound => self.config.limits.rate.inbound.clone().into(), Link::Inbound => self.config.limits.rate.inbound.into(),
}; };
if self if self
.limiter .limiter

View File

@ -74,6 +74,26 @@ impl AsTokens for config::RateLimit {
} }
} }
impl AsTokens for config::LimitRateInbound {
fn capacity(&self) -> usize {
config::RateLimit::from(*self).capacity()
}
fn rate(&self) -> f64 {
config::RateLimit::from(*self).rate()
}
}
impl AsTokens for config::LimitRateOutbound {
fn capacity(&self) -> usize {
config::RateLimit::from(*self).capacity()
}
fn rate(&self) -> f64 {
config::RateLimit::from(*self).rate()
}
}
#[derive(Debug, serde::Serialize)] #[derive(Debug, serde::Serialize)]
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
pub struct TokenBucket { pub struct TokenBucket {

View File

@ -244,7 +244,7 @@ pub struct ConnectionLimits {
} }
/// Rate limits for a single connection. /// Rate limits for a single connection.
#[derive(Debug, Clone, Serialize, Deserialize, Display)] #[derive(Debug, Clone, Copy, Serialize, Deserialize, Display)]
#[display("RateLimit(fill_rate={fill_rate}, capacity={capacity})")] #[display("RateLimit(fill_rate={fill_rate}, capacity={capacity})")]
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))] #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
@ -599,6 +599,18 @@ impl From<LimitGossipMaxAge> for LocalDuration {
} }
} }
/// Create a new type (`$name`) around a given type (`$type`), with a provided
/// default (`$default`).
///
/// The macro will attempt to derive any extra `$derive`s passed.
///
/// Note that the macro will provide the following traits automatically:
/// - `Clone`
/// - `Debug`
/// - `Display`
/// - `Serialize`
/// - `Deserialize`
/// - `From<$name> for $type`, i.e. can convert back into the original type
macro_rules! wrapper { macro_rules! wrapper {
($name:ident, $type:ty, $default:expr $(, $derive:ty)*) => { ($name:ident, $type:ty, $default:expr $(, $derive:ty)*) => {
#[derive(Clone, Debug, Deserialize, Display, Serialize, From $(, $derive)*)] #[derive(Clone, Debug, Deserialize, Display, Serialize, From $(, $derive)*)]
@ -631,7 +643,8 @@ wrapper!(
RateLimit { RateLimit {
fill_rate: 5.0, fill_rate: 5.0,
capacity: 1024, capacity: 1024,
} },
Copy
); );
wrapper!(LimitMaxOpenFiles, usize, 4096, Copy); wrapper!(LimitMaxOpenFiles, usize, 4096, Copy);
wrapper!( wrapper!(
@ -640,7 +653,8 @@ wrapper!(
RateLimit { RateLimit {
fill_rate: 10.0, fill_rate: 10.0,
capacity: 2048, capacity: 2048,
} },
Copy
); );
#[cfg(test)] #[cfg(test)]