Skip to main content

iota_types/
traffic_control.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use 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/// The source that should be used to identify the client's
14/// IP address. To be used to configure cases where a node has
15/// infra running in front of the node that is separate from the
16/// protocol, such as a load balancer. Note that this is not the
17/// same as the client type (e.g a direct client vs a proxy client,
18/// as in the case of a fullnode driving requests from many clients).
19///
20/// For x-forwarded-for, the usize parameter is the number of forwarding
21/// hops between the client and the node for requests going your infra
22/// or infra provider. Example:
23///
24/// ```ignore
25///     (client) -> { (global proxy) -> (regional proxy) -> (node) }
26/// ```
27///
28/// where
29///
30/// ```ignore
31///     { <server>, ... }
32/// ```
33///
34/// are controlled by the Node operator / their cloud provider.
35/// In this case, we set:
36///
37/// ```ignore
38/// policy-config:
39///    client-id-source:
40///      x-forwarded-for: 2
41///    ...
42/// ```
43///
44/// NOTE: x-forwarded-for: 0 is a special case value that can be used by Node
45/// operators to discover the number of hops that should be configured. To use:
46///
47/// 1. Set `x-forwarded-for: 0` for the `client-id-source` in the config.
48/// 2. Run the node and query any endpoint (AuthorityServer for validator, or
49///    json rpc for rpc node) from a known IP address.
50/// 3. Search for lines containing `x-forwarded-for` in the logs. The log lines
51///    should contain the contents of the `x-forwarded-for` header, if present,
52///    or a corresponding error if not.
53/// 4. The value for number of hops is derived from any such log line that
54///    contains your known IP address, and is defined as 1 + the number of IP
55///    addresses in the `x-forwarded-for` that occur **after** the known client
56///    IP address. Example:
57///
58/// ```ignore
59///     [<known client IP>] <--- number of hops is 1
60///     ["1.2.3.4", <known client IP>, "5.6.7.8", "9.10.11.12"] <--- number of hops is 3
61/// ```
62#[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        // `Uniform::new` excludes the upper bound, so a weight of 1.0 accepts every
104        // sample.
105        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    /// Time in secs, after which no registered ingress traffic
164    /// will trigger dead mans switch to drain any firewalls
165    #[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    /// Sustained tallies per second allowed from a single direct client. Zero
182    /// blocks every client; values above 10^9 are clamped.
183    #[serde(default = "default_client_threshold")]
184    pub client_threshold: u64,
185    /// Sustained tallies per second allowed from a single proxied client.
186    #[serde(default = "default_proxied_client_threshold")]
187    pub proxied_client_threshold: u64,
188    /// The burst allowance, in seconds of the sustained rate. A client can
189    /// send `threshold * burst_secs` tallies back to back before the node
190    /// blocks it. The node limits the burst to `u32::MAX` tallies. A value of
191    /// zero stops the node at startup. The former name `window-size-secs` is
192    /// deprecated, but the node continues to accept it.
193    #[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    // By default only block clients with an unreasonably high sustained rate,
209    // as a client could be a single fullnode proxying the traffic of many
210    // behaving clients. In practice this should always be set.
211    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// Serializable representation of policy types, used in config
223// in order to easily change in tests or to killswitch
224#[derive(Clone, Serialize, Deserialize, Debug, Default)]
225pub enum PolicyType {
226    /// Does nothing
227    #[default]
228    NoOp,
229
230    /// Rate limits each client IP to `threshold` tallies per second, tolerating
231    /// a burst of `threshold * burst_secs` tallies.
232    #[serde(rename = "freq-threshold", alias = "FreqThreshold")]
233    FreqThreshold(FreqThresholdConfig),
234
235    // Below this point are test policies, and thus should not be used in production
236    /// Simple policy that adds connection_ip to blocklist when the same
237    /// connection_ip is encountered in tally N times. If used in an error
238    /// policy, this would trigger after N errors
239    TestNConnIP(u64),
240    /// Test policy that panics when invoked. To be used as an error policy in
241    /// tests that do not expect request errors in order to verify that the
242    /// error policy is not invoked
243    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    /// Note that this sample policy is applied on top of the
265    /// endpoint-specific sample policy (not configurable) which
266    /// weighs endpoints by the relative effort required to serve
267    /// them. Therefore a sample rate of N will yield an actual
268    /// sample rate <= N.
269    pub spam_sample_rate: Weight,
270    #[serde(default = "default_dry_run")]
271    pub dry_run: bool,
272    /// List of String which should all parse to type IPAddr.
273    /// If set, only requests from provided IPs will be allowed,
274    /// and any blocklist related configuration will be ignored.
275    #[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}