Skip to content

Commit

Permalink
Addressed some feedback from Trent
Browse files Browse the repository at this point in the history
  • Loading branch information
lijunwangs committed Apr 24, 2024
1 parent b19a240 commit 07ee465
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 7 deletions.
13 changes: 8 additions & 5 deletions streamer/src/nonblocking/connection_rate_limiter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,17 @@ pub struct ConnectionRateLimiter {
}

impl ConnectionRateLimiter {
/// Create a new rate limiter per IpAddr. The rate is specified as the count per minute to allow for
/// less frequent connections.
pub fn new(limit_per_minute: u32) -> Self {
let quota = Quota::per_minute(NonZeroU32::new(limit_per_minute).unwrap()); // Adjust the rate limit as needed
let quota = Quota::per_minute(NonZeroU32::new(limit_per_minute).unwrap());
Self {
limiter: DefaultKeyedRateLimiter::keyed(quota),
}
}

pub fn check(&self, ip: &IpAddr) -> bool {
/// Check if the connection from the said `ip` is allowed.
pub fn is_allowed(&self, ip: &IpAddr) -> bool {
// Acquire a permit from the rate limiter for the given IP address
if self.limiter.check_key(ip).is_ok() {
debug!("Request from IP {:?} allowed", ip);
Expand All @@ -34,19 +37,19 @@ pub struct TotalConnectionRateLimiter {
}

impl TotalConnectionRateLimiter {
/// Create a new rate limiter. The rate is specified as the count per second.
pub fn new(limit_per_second: u32) -> Self {
let quota = Quota::per_second(NonZeroU32::new(limit_per_second).unwrap()); // Adjust the rate limit as needed
Self {
limiter: RateLimiter::direct(quota),
}
}

pub fn check(&self, ip: &IpAddr) -> bool {
/// Check if a connection is allowed.
pub fn is_allowed(&self) -> bool {
if self.limiter.check().is_ok() {
debug!("Request from IP {:?} allowed", ip);
true // Request allowed
} else {
debug!("Request from IP {:?} blocked", ip);
false // Request blocked
}
}
Expand Down
8 changes: 6 additions & 2 deletions streamer/src/nonblocking/quic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -234,15 +234,19 @@ async fn run_server(
let remote_address = connection.remote_address();

// first check overall connection rate limit:
if !overall_connection_rate_limiter.check(&remote_address.ip()) {
if !overall_connection_rate_limiter.is_allowed() {
debug!(
"Reject connection from {:?} -- total rate limiting exceeded",
remote_address.ip()
);
stats
.connection_throttled_across_all
.fetch_add(1, Ordering::Relaxed);
continue;
}

info!("Got a connection {remote_address:?}");
if !rate_limiter.check(&remote_address.ip()) {
if !rate_limiter.is_allowed(&remote_address.ip()) {
info!(
"Reject connection from {:?} -- rate limiting exceeded",
remote_address
Expand Down
1 change: 1 addition & 0 deletions validator/src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -890,6 +890,7 @@ pub fn app<'a>(version: &'a str, default_args: &'a DefaultArgs) -> App<'a, 'a> {
.takes_value(true)
.default_value(&default_args.tpu_max_connections_per_ipaddr_per_minute)
.validator(is_parsable::<u32>)
.hidden(hidden_unless_forced())
.help("Controls the rate of the clients connections per IpAddr per minute."),
)
.arg(
Expand Down

0 comments on commit 07ee465

Please sign in to comment.