iota_types/
traffic_control.rs1use rand::distr::Distribution;
6use serde::{Deserialize, Serialize, de::Deserializer};
7use serde_with::serde_as;
8
9#[derive(Clone, Debug, Deserialize, Serialize, Default)]
57#[serde(rename_all = "kebab-case")]
58pub enum ClientIdSource {
59 #[default]
60 SocketAddr,
61 XForwardedFor(usize),
62}
63
64#[derive(Clone, Debug, Deserialize, Serialize)]
65pub struct TrafficControlReconfigParams {
66 pub error_threshold: Option<u64>,
67 pub spam_threshold: Option<u64>,
68 pub dry_run: Option<bool>,
69 pub connection_blocklist_ttl_sec: Option<u64>,
70 pub proxy_blocklist_ttl_sec: Option<u64>,
71}
72
73#[derive(Clone, Debug, Deserialize, Serialize)]
74pub struct Weight(f32);
75
76impl Weight {
77 pub fn new(value: f32) -> Result<Self, &'static str> {
78 if (0.0..=1.0).contains(&value) {
79 Ok(Self(value))
80 } else {
81 Err("Weight must be between 0.0 and 1.0")
82 }
83 }
84
85 pub fn one() -> Self {
86 Self(1.0)
87 }
88
89 pub fn zero() -> Self {
90 Self(0.0)
91 }
92
93 pub fn value(&self) -> f32 {
94 self.0
95 }
96
97 pub fn is_sampled(&self) -> bool {
98 let mut rng = rand::rng();
99 let sample = rand::distr::Uniform::new(0.0, 1.0)
102 .unwrap()
103 .sample(&mut rng);
104 self.accepts(sample)
105 }
106
107 fn accepts(&self, sample: f32) -> bool {
108 sample < self.value()
109 }
110}
111
112fn validate_sample_rate<'de, D>(deserializer: D) -> Result<Weight, D::Error>
113where
114 D: Deserializer<'de>,
115{
116 let value = f32::deserialize(deserializer)?;
117 Weight::new(value)
118 .map_err(|_| serde::de::Error::custom("spam-sample-rate must be between 0.0 and 1.0"))
119}
120
121impl PartialEq for Weight {
122 fn eq(&self, other: &Self) -> bool {
123 self.value() == other.value()
124 }
125}
126
127#[serde_as]
128#[derive(Clone, Debug, Deserialize, Serialize)]
129#[serde(rename_all = "kebab-case")]
130pub struct FreqThresholdConfig {
131 #[serde(default = "default_client_threshold")]
134 pub client_threshold: u64,
135 #[serde(default = "default_proxied_client_threshold")]
137 pub proxied_client_threshold: u64,
138 #[serde(default = "default_burst_secs", alias = "window-size-secs")]
144 pub burst_secs: u64,
145}
146
147impl Default for FreqThresholdConfig {
148 fn default() -> Self {
149 Self {
150 client_threshold: default_client_threshold(),
151 proxied_client_threshold: default_proxied_client_threshold(),
152 burst_secs: default_burst_secs(),
153 }
154 }
155}
156
157fn default_client_threshold() -> u64 {
158 1_000_000
162}
163
164fn default_proxied_client_threshold() -> u64 {
165 10
166}
167
168fn default_burst_secs() -> u64 {
169 30
170}
171
172#[derive(Clone, Serialize, Deserialize, Debug, Default)]
175pub enum PolicyType {
176 #[default]
178 NoOp,
179
180 #[serde(rename = "freq-threshold", alias = "FreqThreshold")]
183 FreqThreshold(FreqThresholdConfig),
184
185 TestNConnIP(u64),
190 TestPanicOnInvocation,
194}
195
196#[serde_as]
197#[derive(Clone, Debug, Deserialize, Serialize)]
198#[serde(rename_all = "kebab-case")]
199pub struct PolicyConfig {
200 #[serde(default = "default_client_id_source")]
201 pub client_id_source: ClientIdSource,
202 #[serde(default = "default_blocklist_ttl_sec")]
203 pub connection_blocklist_ttl_sec: u64,
204 #[serde(default = "default_blocklist_ttl_sec")]
205 pub proxy_blocklist_ttl_sec: u64,
206 #[serde(default)]
207 pub spam_policy_type: PolicyType,
208 #[serde(default)]
209 pub error_policy_type: PolicyType,
210 #[serde(
211 default = "default_spam_sample_rate",
212 deserialize_with = "validate_sample_rate"
213 )]
214 pub spam_sample_rate: Weight,
220 #[serde(default = "default_dry_run")]
221 pub dry_run: bool,
222 #[serde(default)]
226 pub allow_list: Option<Vec<String>>,
227 #[serde(default = "default_ipv6_prefix_len")]
235 pub ipv6_prefix_len: u8,
236}
237
238impl Default for PolicyConfig {
239 fn default() -> Self {
240 Self {
241 client_id_source: default_client_id_source(),
242 connection_blocklist_ttl_sec: default_blocklist_ttl_sec(),
243 proxy_blocklist_ttl_sec: default_blocklist_ttl_sec(),
244 spam_policy_type: PolicyType::NoOp,
245 error_policy_type: PolicyType::NoOp,
246 spam_sample_rate: default_spam_sample_rate(),
247 dry_run: default_dry_run(),
248 allow_list: None,
249 ipv6_prefix_len: default_ipv6_prefix_len(),
250 }
251 }
252}
253
254impl PolicyConfig {
255 pub fn default_dos_protection_policy() -> Self {
256 Self {
257 client_id_source: ClientIdSource::SocketAddr,
258 spam_policy_type: PolicyType::FreqThreshold(FreqThresholdConfig {
259 client_threshold: 1000,
260 burst_secs: 5,
261 ..FreqThresholdConfig::default()
262 }),
263 error_policy_type: PolicyType::FreqThreshold(FreqThresholdConfig {
264 client_threshold: 50,
265 burst_secs: 5,
266 ..FreqThresholdConfig::default()
267 }),
268 spam_sample_rate: Weight::new(1.0).unwrap(),
269 dry_run: true,
270 ..Self::default()
271 }
272 }
273}
274
275pub fn default_client_id_source() -> ClientIdSource {
276 ClientIdSource::SocketAddr
277}
278
279pub fn default_blocklist_ttl_sec() -> u64 {
280 60
281}
282pub fn default_ipv6_prefix_len() -> u8 {
283 64
284}
285
286pub fn default_dry_run() -> bool {
287 true
288}
289
290pub fn default_spam_sample_rate() -> Weight {
291 Weight::new(0.2).unwrap()
292}
293
294#[cfg(test)]
295mod tests {
296 use super::Weight;
297
298 #[test]
299 fn zero_weight_rejects_the_lowest_sample() {
300 assert!(!Weight::zero().accepts(0.0));
301 }
302
303 #[test]
304 fn full_weight_accepts_the_highest_sample() {
305 assert!(Weight::one().accepts(1.0 - f32::EPSILON));
306 }
307
308 #[test]
309 fn the_default_dos_protection_policy_blocks_for_a_nonzero_time() {
310 let policy = super::PolicyConfig::default_dos_protection_policy();
313 assert!(policy.connection_blocklist_ttl_sec > 0);
314 assert!(policy.proxy_blocklist_ttl_sec > 0);
315 }
316
317 #[test]
318 fn the_former_burst_key_still_parses() {
319 let config: super::FreqThresholdConfig =
320 serde_json::from_str(r#"{"window-size-secs": 7}"#).unwrap();
321 assert_eq!(config.burst_secs, 7);
322 }
323}