From aa287ac68cebb397b485dc0d0304cdd8d0b904e1 Mon Sep 17 00:00:00 2001 From: Rigidity Date: Wed, 6 Nov 2024 19:04:14 -0500 Subject: [PATCH] Initial implementation of rate limits --- Cargo.lock | 1 + crates/chia-sdk-client/Cargo.toml | 1 + crates/chia-sdk-client/src/lib.rs | 4 + crates/chia-sdk-client/src/rate_limiter.rs | 119 ++++++++++++ crates/chia-sdk-client/src/rate_limits.rs | 214 +++++++++++++++++++++ 5 files changed, 339 insertions(+) create mode 100644 crates/chia-sdk-client/src/rate_limiter.rs create mode 100644 crates/chia-sdk-client/src/rate_limits.rs diff --git a/Cargo.lock b/Cargo.lock index 5455b79a..1ad06afe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -412,6 +412,7 @@ dependencies = [ "chia-traits 0.15.0", "futures-util", "native-tls", + "once_cell", "rustls", "rustls-pemfile", "thiserror", diff --git a/crates/chia-sdk-client/Cargo.toml b/crates/chia-sdk-client/Cargo.toml index 5fbce905..e8918bad 100644 --- a/crates/chia-sdk-client/Cargo.toml +++ b/crates/chia-sdk-client/Cargo.toml @@ -32,6 +32,7 @@ rustls-pemfile = { workspace = true, optional = true } tracing = { workspace = true } futures-util = { workspace = true } tokio-tungstenite = { workspace = true } +once_cell = { workspace = true } # This is to ensure that the bindgen feature is enabled for the aws-lc-rs crate. # https://aws.github.io/aws-lc-rs/platform_support.html#tested-platforms diff --git a/crates/chia-sdk-client/src/lib.rs b/crates/chia-sdk-client/src/lib.rs index 751521a1..fca64f9d 100644 --- a/crates/chia-sdk-client/src/lib.rs +++ b/crates/chia-sdk-client/src/lib.rs @@ -1,12 +1,16 @@ mod error; mod network; mod peer; +mod rate_limiter; +mod rate_limits; mod request_map; mod tls; pub use error::*; pub use network::*; pub use peer::*; +pub use rate_limiter::*; +pub use rate_limits::*; pub use tls::*; #[cfg(any(feature = "native-tls", feature = "rustls"))] diff --git a/crates/chia-sdk-client/src/rate_limiter.rs b/crates/chia-sdk-client/src/rate_limiter.rs new file mode 100644 index 00000000..51bd51e4 --- /dev/null +++ b/crates/chia-sdk-client/src/rate_limiter.rs @@ -0,0 +1,119 @@ +use std::{ + collections::HashMap, + time::{SystemTime, UNIX_EPOCH}, +}; + +use chia_protocol::{Message, ProtocolMessageTypes}; + +use crate::RateLimits; + +#[derive(Debug, Clone)] +pub struct RateLimiter { + incoming: bool, + reset_seconds: u64, + period: u64, + message_counts: HashMap, + message_cumulative_sizes: HashMap, + limit_factor: f64, + non_tx_count: u32, + non_tx_size: u32, +} + +impl RateLimiter { + pub fn new(incoming: bool, reset_seconds: u64, limit_factor: f64) -> Self { + Self { + incoming, + reset_seconds, + period: time() / reset_seconds, + message_counts: HashMap::new(), + message_cumulative_sizes: HashMap::new(), + limit_factor, + non_tx_count: 0, + non_tx_size: 0, + } + } + + pub fn handle_message(&mut self, message: &Message, settings: &RateLimits) -> bool { + let size: u32 = message.data.len().try_into().expect("Message too large"); + let period = time() / self.reset_seconds; + + if self.period != period { + self.period = period; + self.message_counts.clear(); + self.message_cumulative_sizes.clear(); + self.non_tx_count = 0; + self.non_tx_size = 0; + } + + let new_message_count = self.message_counts.get(&message.msg_type).unwrap_or(&0) + 1; + let new_cumulative_size = self + .message_cumulative_sizes + .get(&message.msg_type) + .unwrap_or(&0) + + size; + let mut new_non_tx_count = self.non_tx_count; + let mut new_non_tx_size = self.non_tx_size; + + let passed = 'checker: { + let mut limits = settings.default_settings; + + if let Some(tx_limits) = settings.tx.get(&message.msg_type) { + limits = *tx_limits; + } else if let Some(other_limits) = settings.other.get(&message.msg_type) { + limits = *other_limits; + + new_non_tx_count += 1; + new_non_tx_size += size; + + if f64::from(new_non_tx_count) + > f64::from(settings.non_tx_frequency) * self.limit_factor + { + break 'checker false; + } + + if f64::from(new_non_tx_size) + > f64::from(settings.non_tx_max_total_size) * self.limit_factor + { + break 'checker false; + } + } + + let max_total_size = limits + .max_total_size + .unwrap_or(limits.frequency * limits.max_size); + + if f64::from(new_message_count) > f64::from(limits.frequency) * self.limit_factor { + break 'checker false; + } + + if size > limits.max_size { + break 'checker false; + } + + if f64::from(new_cumulative_size) > f64::from(max_total_size) * self.limit_factor { + break 'checker false; + } + + true + }; + + if self.incoming || passed { + *self.message_counts.entry(message.msg_type).or_default() = new_message_count; + *self + .message_cumulative_sizes + .entry(message.msg_type) + .or_default() = new_cumulative_size; + self.non_tx_count = new_non_tx_count; + self.non_tx_size = new_non_tx_size; + } + + passed + } +} + +fn time() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Time went backwards") + .as_secs() +} diff --git a/crates/chia-sdk-client/src/rate_limits.rs b/crates/chia-sdk-client/src/rate_limits.rs new file mode 100644 index 00000000..d79c9b23 --- /dev/null +++ b/crates/chia-sdk-client/src/rate_limits.rs @@ -0,0 +1,214 @@ +use std::collections::HashMap; + +use chia_protocol::ProtocolMessageTypes; +use once_cell::sync::Lazy; + +#[derive(Debug, Clone)] +pub struct RateLimits { + pub default_settings: RateLimit, + pub non_tx_frequency: u32, + pub non_tx_max_total_size: u32, + pub tx: HashMap, + pub other: HashMap, +} + +impl RateLimits { + pub fn extend(&mut self, other: &Self) { + self.default_settings = other.default_settings; + self.non_tx_frequency = other.non_tx_frequency; + self.non_tx_max_total_size = other.non_tx_max_total_size; + self.tx.extend(other.tx.clone()); + self.other.extend(other.other.clone()); + } +} + +#[derive(Debug, Clone, Copy)] +pub struct RateLimit { + pub frequency: u32, + pub max_size: u32, + pub max_total_size: Option, +} + +impl RateLimit { + pub fn new(frequency: u32, max_size: u32, max_total_size: Option) -> Self { + Self { + frequency, + max_size, + max_total_size, + } + } +} + +macro_rules! settings { + ($($message:ident => $frequency:expr, $max_size:expr $(, $max_total_size:expr)? ;)*) => { + { + let mut settings = HashMap::new(); + $( + #[allow(unused_mut, unused_assignments)] + let mut max_total_size = None; + $( max_total_size = Some($max_total_size); )? + settings.insert( + ProtocolMessageTypes::$message, + RateLimit::new( + $frequency, + $max_size, + max_total_size, + ) + ); + )* + settings + } + }; +} + +// TODO: Fix commented out rate limits. +pub static V1_RATE_LIMITS: Lazy = Lazy::new(|| RateLimits { + default_settings: RateLimit::new(100, 1024 * 1024, Some(100 * 1024 * 1024)), + non_tx_frequency: 1000, + non_tx_max_total_size: 100 * 1024 * 1024, + tx: settings! { + NewTransaction => 5000, 100, 5000 * 100; + RequestTransaction => 5000, 100, 5000 * 100; + RespondTransaction => 5000, 1024 * 1024, 20 * 1024 * 1024; + SendTransaction => 5000, 1024 * 1024; + TransactionAck => 5000, 2048; + }, + other: settings! { + Handshake => 5, 10 * 1024, 5 * 10 * 1024; + HarvesterHandshake => 5, 1024 * 1024; + NewSignagePointHarvester => 100, 4886; + NewProofOfSpace => 100, 2048; + RequestSignatures => 100, 2048; + RespondSignatures => 100, 2048; + NewSignagePoint => 200, 2048; + DeclareProofOfSpace => 100, 10 * 1024; + RequestSignedValues => 100, 10 * 1024; + FarmingInfo => 100, 1024; + SignedValues => 100, 1024; + NewPeakTimelord => 100, 20 * 1024; + NewUnfinishedBlockTimelord => 100, 10 * 1024; + NewSignagePointVdf => 100, 100 * 1024; + NewInfusionPointVdf => 100, 100 * 1024; + NewEndOfSubSlotVdf => 100, 100 * 1024; + RequestCompactProofOfTime => 100, 10 * 1024; + RespondCompactProofOfTime => 100, 100 * 1024; + NewPeak => 200, 512; + RequestProofOfWeight => 5, 100; + RespondProofOfWeight => 5, 50 * 1024 * 1024, 100 * 1024 * 1024; + RequestBlock => 200, 100; + RejectBlock => 200, 100; + RequestBlocks => 500, 100; + RespondBlocks => 100, 50 * 1024 * 1024, 5 * 50 * 1024 * 1024; + RejectBlocks => 100, 100; + RespondBlock => 200, 2 * 1024 * 1024, 10 * 2 * 1024 * 1024; + NewUnfinishedBlock => 200, 100; + RequestUnfinishedBlock => 200, 100; + NewUnfinishedBlock2 => 200, 100; + RequestUnfinishedBlock2 => 200, 100; + RespondUnfinishedBlock => 200, 2 * 1024 * 1024, 10 * 2 * 1024 * 1024; + NewSignagePointOrEndOfSubSlot => 200, 200; + RequestSignagePointOrEndOfSubSlot => 200, 200; + RespondSignagePoint => 200, 50 * 1024; + RespondEndOfSubSlot => 100, 50 * 1024; + RequestMempoolTransactions => 5, 1024 * 1024; + RequestCompactVDF => 200, 1024; + RespondCompactVDF => 200, 100 * 1024; + NewCompactVDF => 100, 1024; + RequestPeers => 10, 100; + RespondPeers => 10, 1024 * 1024; + RequestPuzzleSolution => 1000, 100; + RespondPuzzleSolution => 1000, 1024 * 1024; + RejectPuzzleSolution => 1000, 100; + NewPeakWallet => 200, 300; + RequestBlockHeader => 500, 100; + RespondBlockHeader => 500, 500 * 1024; + RejectHeaderRequest => 500, 100; + RequestRemovals => 500, 50 * 1024, 10 * 1024 * 1024; + RespondRemovals => 500, 1024 * 1024, 10 * 1024 * 1024; + RejectRemovalsRequest => 500, 100; + RequestAdditions => 500, 1024 * 1024, 10 * 1024 * 1024; + RespondAdditions => 500, 1024 * 1024, 10 * 1024 * 1024; + RejectAdditionsRequest => 500, 100; + RequestHeaderBlocks => 500, 100; + RejectHeaderBlocks => 100, 100; + RespondHeaderBlocks => 500, 2 * 1024 * 1024, 100 * 1024 * 1024; + RequestPeersIntroducer => 100, 100; + RespondPeersIntroducer => 100, 1024 * 1024; + FarmNewBlock => 200, 200; + RequestPlots => 10, 10 * 1024 * 1024; + RespondPlots => 10, 100 * 1024 * 1024; + PlotSyncStart => 1000, 100 * 1024 * 1024; + PlotSyncLoaded => 1000, 100 * 1024 * 1024; + PlotSyncRemoved => 1000, 100 * 1024 * 1024; + PlotSyncInvalid => 1000, 100 * 1024 * 1024; + PlotSyncKeysMissing => 1000, 100 * 1024 * 1024; + PlotSyncDuplicates => 1000, 100 * 1024 * 1024; + PlotSyncDone => 1000, 100 * 1024 * 1024; + PlotSyncResponse => 3000, 100 * 1024 * 1024; + CoinStateUpdate => 1000, 100 * 1024 * 1024; + RegisterForPhUpdates => 1000, 100 * 1024 * 1024; + RespondToPhUpdates => 1000, 100 * 1024 * 1024; + RegisterForCoinUpdates => 1000, 100 * 1024 * 1024; + RespondToCoinUpdates => 1000, 100 * 1024 * 1024; + RequestRemovePuzzleSubscriptions => 1000, 100 * 1024 * 1024; + RespondRemovePuzzleSubscriptions => 1000, 100 * 1024 * 1024; + RequestRemoveCoinSubscriptions => 1000, 100 * 1024 * 1024; + RespondRemoveCoinSubscriptions => 1000, 100 * 1024 * 1024; + RequestPuzzleState => 1000, 100 * 1024 * 1024; + RespondPuzzleState => 1000, 100 * 1024 * 1024; + RejectPuzzleState => 200, 100; + RequestCoinState => 1000, 100 * 1024 * 1024; + RespondCoinState => 1000, 100 * 1024 * 1024; + RejectCoinState => 200, 100; + // MempoolItemsAdded => 1000, 100 * 1024 * 1024; + // MempoolItemsRemoved => 1000, 100 * 1024 * 1024; + // RequestCostInfo => 1000, 100; + // RespondCostInfo => 1000, 1024; + // RequestSesHashes => 2000, 1 * 1024 * 1024; + // RespondSesHashes => 2000, 1 * 1024 * 1024; + RequestChildren => 2000, 1024 * 1024; + RespondChildren => 2000, 1024 * 1024; + }, +}); + +// TODO: Fix commented out rate limits. +// Also, why are these in tx? +static V2_RATE_LIMIT_CHANGES: Lazy = Lazy::new(|| RateLimits { + default_settings: RateLimit::new(100, 1024 * 1024, Some(100 * 1024 * 1024)), + non_tx_frequency: 1000, + non_tx_max_total_size: 100 * 1024 * 1024, + tx: settings! { + RequestBlockHeader => 500, 100; + RespondBlockHeader => 500, 500 * 1024; + RejectHeaderRequest => 500, 100; + RequestRemovals => 5000, 50 * 1024, 10 * 1024 * 1024; + RespondRemovals => 5000, 1024 * 1024, 10 * 1024 * 1024; + RejectRemovalsRequest => 500, 100; + RequestAdditions => 50000, 100 * 1024 * 1024; + RespondAdditions => 50000, 100 * 1024 * 1024; + RejectAdditionsRequest => 500, 100; + RejectHeaderBlocks => 1000, 100; + RespondHeaderBlocks => 5000, 2 * 1024 * 1024; + RequestBlockHeaders => 5000, 100; + RejectBlockHeaders => 1000, 100; + RespondBlockHeaders => 5000, 2 * 1024 * 1024; + // RequestSesHashes => 2000, 1 * 1024 * 1024; + // RespondSesHashes => 2000, 1 * 1024 * 1024; + RequestChildren => 2000, 1024 * 1024; + RespondChildren => 2000, 1024 * 1024; + RequestPuzzleSolution => 5000, 100; + RespondPuzzleSolution => 5000, 1024 * 1024; + RejectPuzzleSolution => 5000, 100; + NoneResponse => 500, 100; + // Error => 50000, 100; + }, + other: settings! { + RequestHeaderBlocks => 5000, 100; + }, +}); + +pub static V2_RATE_LIMITS: Lazy = Lazy::new(|| { + let mut rate_limits = V1_RATE_LIMITS.clone(); + rate_limits.extend(&V2_RATE_LIMIT_CHANGES); + rate_limits +});