iota_types/
traffic_control.rs1use std::path::PathBuf;
6
7use rand::distributions::Distribution;
8use serde::{Deserialize, Serialize, de::Deserializer};
9use serde_with::serde_as;
10
11const TRAFFIC_SINK_TIMEOUT_SEC: u64 = 300;
12
13#[derive(Clone, Debug, Deserialize, Serialize, Default)]
63#[serde(rename_all = "kebab-case")]
64pub enum ClientIdSource {
65 #[default]
66 SocketAddr,
67 XForwardedFor(usize),
68}
69
70#[derive(Clone, Debug, Deserialize, Serialize)]
71pub struct TrafficControlReconfigParams {
72 pub error_threshold: Option<u64>,
73 pub spam_threshold: Option<u64>,
74 pub dry_run: Option<bool>,
75}
76
77#[derive(Clone, Debug, Deserialize, Serialize)]
78pub struct Weight(f32);
79
80impl Weight {
81 pub fn new(value: f32) -> Result<Self, &'static str> {
82 if (0.0..=1.0).contains(&value) {
83 Ok(Self(value))
84 } else {
85 Err("Weight must be between 0.0 and 1.0")
86 }
87 }
88
89 pub fn one() -> Self {
90 Self(1.0)
91 }
92
93 pub fn zero() -> Self {
94 Self(0.0)
95 }
96
97 pub fn value(&self) -> f32 {
98 self.0
99 }
100
101 pub fn is_sampled(&self) -> bool {
102 let mut rng = rand::thread_rng();
103 let sample = rand::distributions::Uniform::new(0.0, 1.0).sample(&mut rng);
106 self.accepts(sample)
107 }
108
109 fn accepts(&self, sample: f32) -> bool {
110 sample < self.value()
111 }
112}
113
114#[cfg(test)]
115mod tests {
116 use super::Weight;
117
118 #[test]
119 fn zero_weight_rejects_the_lowest_sample() {
120 assert!(!Weight::zero().accepts(0.0));
121 }
122
123 #[test]
124 fn full_weight_accepts_the_highest_sample() {
125 assert!(Weight::one().accepts(1.0 - f32::EPSILON));
126 }
127
128 #[test]
129 fn the_former_burst_key_still_parses() {
130 let config: super::FreqThresholdConfig =
131 serde_json::from_str(r#"{"window-size-secs": 7}"#).unwrap();
132 assert_eq!(config.burst_secs, 7);
133 }
134}
135
136fn validate_sample_rate<'de, D>(deserializer: D) -> Result<Weight, D::Error>
137where
138 D: Deserializer<'de>,
139{
140 let value = f32::deserialize(deserializer)?;
141 Weight::new(value)
142 .map_err(|_| serde::de::Error::custom("spam-sample-rate must be between 0.0 and 1.0"))
143}
144
145impl PartialEq for Weight {
146 fn eq(&self, other: &Self) -> bool {
147 self.value() == other.value()
148 }
149}
150
151#[serde_as]
152#[derive(Clone, Debug, Deserialize, Serialize)]
153#[serde(rename_all = "kebab-case")]
154pub struct RemoteFirewallConfig {
155 pub remote_fw_url: String,
156 pub destination_port: u16,
157 #[serde(default)]
158 pub delegate_spam_blocking: bool,
159 #[serde(default)]
160 pub delegate_error_blocking: bool,
161 #[serde(default = "default_drain_path")]
162 pub drain_path: PathBuf,
163 #[serde(default = "default_drain_timeout")]
166 pub drain_timeout_secs: u64,
167}
168
169fn default_drain_path() -> PathBuf {
170 PathBuf::from("/tmp/drain")
171}
172
173fn default_drain_timeout() -> u64 {
174 TRAFFIC_SINK_TIMEOUT_SEC
175}
176
177#[serde_as]
178#[derive(Clone, Debug, Deserialize, Serialize)]
179#[serde(rename_all = "kebab-case")]
180pub struct FreqThresholdConfig {
181 #[serde(default = "default_client_threshold")]
184 pub client_threshold: u64,
185 #[serde(default = "default_proxied_client_threshold")]
187 pub proxied_client_threshold: u64,
188 #[serde(default = "default_burst_secs", alias = "window-size-secs")]
194 pub burst_secs: u64,
195}
196
197impl Default for FreqThresholdConfig {
198 fn default() -> Self {
199 Self {
200 client_threshold: default_client_threshold(),
201 proxied_client_threshold: default_proxied_client_threshold(),
202 burst_secs: default_burst_secs(),
203 }
204 }
205}
206
207fn default_client_threshold() -> u64 {
208 1_000_000
212}
213
214fn default_proxied_client_threshold() -> u64 {
215 10
216}
217
218fn default_burst_secs() -> u64 {
219 30
220}
221
222#[derive(Clone, Serialize, Deserialize, Debug, Default)]
225pub enum PolicyType {
226 #[default]
228 NoOp,
229
230 #[serde(rename = "freq-threshold", alias = "FreqThreshold")]
233 FreqThreshold(FreqThresholdConfig),
234
235 TestNConnIP(u64),
240 TestPanicOnInvocation,
244}
245
246#[serde_as]
247#[derive(Clone, Debug, Deserialize, Serialize)]
248#[serde(rename_all = "kebab-case")]
249pub struct PolicyConfig {
250 #[serde(default = "default_client_id_source")]
251 pub client_id_source: ClientIdSource,
252 #[serde(default = "default_blocklist_ttl_sec")]
253 pub connection_blocklist_ttl_sec: u64,
254 #[serde(default = "default_blocklist_ttl_sec")]
255 pub proxy_blocklist_ttl_sec: u64,
256 #[serde(default)]
257 pub spam_policy_type: PolicyType,
258 #[serde(default)]
259 pub error_policy_type: PolicyType,
260 #[serde(
261 default = "default_spam_sample_rate",
262 deserialize_with = "validate_sample_rate"
263 )]
264 pub spam_sample_rate: Weight,
270 #[serde(default = "default_dry_run")]
271 pub dry_run: bool,
272 #[serde(default)]
276 pub allow_list: Option<Vec<String>>,
277}
278
279impl Default for PolicyConfig {
280 fn default() -> Self {
281 Self {
282 client_id_source: default_client_id_source(),
283 connection_blocklist_ttl_sec: default_blocklist_ttl_sec(),
284 proxy_blocklist_ttl_sec: default_blocklist_ttl_sec(),
285 spam_policy_type: PolicyType::NoOp,
286 error_policy_type: PolicyType::NoOp,
287 spam_sample_rate: default_spam_sample_rate(),
288 dry_run: default_dry_run(),
289 allow_list: None,
290 }
291 }
292}
293
294impl PolicyConfig {
295 pub fn default_dos_protection_policy() -> Self {
296 Self {
297 client_id_source: ClientIdSource::SocketAddr,
298 spam_policy_type: PolicyType::FreqThreshold(FreqThresholdConfig {
299 client_threshold: 1000,
300 burst_secs: 5,
301 ..FreqThresholdConfig::default()
302 }),
303 error_policy_type: PolicyType::FreqThreshold(FreqThresholdConfig {
304 client_threshold: 50,
305 burst_secs: 5,
306 ..FreqThresholdConfig::default()
307 }),
308 spam_sample_rate: Weight::new(1.0).unwrap(),
309 dry_run: true,
310 ..Self::default()
311 }
312 }
313}
314
315pub fn default_client_id_source() -> ClientIdSource {
316 ClientIdSource::SocketAddr
317}
318
319pub fn default_blocklist_ttl_sec() -> u64 {
320 60
321}
322pub fn default_dry_run() -> bool {
323 true
324}
325
326pub fn default_spam_sample_rate() -> Weight {
327 Weight::new(0.2).unwrap()
328}