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}"),
}
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) {
trace!(target: "service", "Rate limiting inbound connection from {host}..");
@ -1819,8 +1819,8 @@ where
peer.last_active = self.clock;
let limit: RateLimit = match peer.link {
Link::Outbound => self.config.limits.rate.outbound.clone().into(),
Link::Inbound => self.config.limits.rate.inbound.clone().into(),
Link::Outbound => self.config.limits.rate.outbound.into(),
Link::Inbound => self.config.limits.rate.inbound.into(),
};
if self
.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)]
#[serde(rename_all = "camelCase")]
pub struct TokenBucket {

View File

@ -244,7 +244,7 @@ pub struct ConnectionLimits {
}
/// 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})")]
#[serde(rename_all = "camelCase")]
#[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 {
($name:ident, $type:ty, $default:expr $(, $derive:ty)*) => {
#[derive(Clone, Debug, Deserialize, Display, Serialize, From $(, $derive)*)]
@ -631,7 +643,8 @@ wrapper!(
RateLimit {
fill_rate: 5.0,
capacity: 1024,
}
},
Copy
);
wrapper!(LimitMaxOpenFiles, usize, 4096, Copy);
wrapper!(
@ -640,7 +653,8 @@ wrapper!(
RateLimit {
fill_rate: 10.0,
capacity: 2048,
}
},
Copy
);
#[cfg(test)]