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 rand::distr::Distribution;
6use serde::{Deserialize, Serialize, de::Deserializer};
7use serde_with::serde_as;
8
9/// The source that should be used to identify the client's
10/// IP address. To be used to configure cases where a node has
11/// infra running in front of the node that is separate from the
12/// protocol, such as a load balancer. Note that this is not the
13/// same as the client type (e.g a direct client vs a proxy client,
14/// as in the case of a fullnode driving requests from many clients).
15///
16/// For x-forwarded-for, the usize parameter is the number of forwarding
17/// hops between the client and the node for requests going your infra
18/// or infra provider. Example:
19///
20/// ```ignore
21///     (client) -> { (global proxy) -> (regional proxy) -> (node) }
22/// ```
23///
24/// where
25///
26/// ```ignore
27///     { <server>, ... }
28/// ```
29///
30/// are controlled by the Node operator / their cloud provider.
31/// In this case, we set:
32///
33/// ```ignore
34/// policy-config:
35///    client-id-source:
36///      x-forwarded-for: 2
37///    ...
38/// ```
39///
40/// NOTE: x-forwarded-for: 0 is a special case value that can be used by Node
41/// operators to discover the number of hops that should be configured. To use:
42///
43/// 1. Set `x-forwarded-for: 0` for the `client-id-source` in the config.
44/// 2. Run the node and query any endpoint (AuthorityServer for validator, or json rpc for rpc node)
45///    from a known IP address.
46/// 3. Search for lines containing `x-forwarded-for` in the logs. The log lines should contain the
47///    contents of the `x-forwarded-for` header, if present, or a corresponding error if not.
48/// 4. The value for number of hops is derived from any such log line that contains your known IP
49///    address, and is defined as 1 + the number of IP addresses in the `x-forwarded-for` that occur
50///    **after** the known client IP address. Example:
51///
52/// ```ignore
53///     [<known client IP>] <--- number of hops is 1
54///     ["1.2.3.4", <known client IP>, "5.6.7.8", "9.10.11.12"] <--- number of hops is 3
55/// ```
56#[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        // `Uniform::new` excludes the upper bound, so a weight of 1.0 accepts every
100        // sample.
101        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    /// Sustained tallies per second allowed from a single direct client. Zero
132    /// blocks every client; values above 10^9 are clamped.
133    #[serde(default = "default_client_threshold")]
134    pub client_threshold: u64,
135    /// Sustained tallies per second allowed from a single proxied client.
136    #[serde(default = "default_proxied_client_threshold")]
137    pub proxied_client_threshold: u64,
138    /// The burst allowance, in seconds of the sustained rate. A client can
139    /// send `threshold * burst_secs` tallies back to back before the node
140    /// blocks it. The node limits the burst to `u32::MAX` tallies. A value of
141    /// zero stops the node at startup. The former name `window-size-secs` is
142    /// deprecated, but the node continues to accept it.
143    #[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    // By default only block clients with an unreasonably high sustained rate,
159    // as a client could be a single fullnode proxying the traffic of many
160    // behaving clients. In practice this should always be set.
161    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// Serializable representation of policy types, used in config
173// in order to easily change in tests or to killswitch
174#[derive(Clone, Serialize, Deserialize, Debug, Default)]
175pub enum PolicyType {
176    /// Does nothing
177    #[default]
178    NoOp,
179
180    /// Rate limits each client IP to `threshold` tallies per second, tolerating
181    /// a burst of `threshold * burst_secs` tallies.
182    #[serde(rename = "freq-threshold", alias = "FreqThreshold")]
183    FreqThreshold(FreqThresholdConfig),
184
185    // Below this point are test policies, and thus should not be used in production
186    /// Simple policy that adds connection_ip to blocklist when the same
187    /// connection_ip is encountered in tally N times. If used in an error
188    /// policy, this would trigger after N errors
189    TestNConnIP(u64),
190    /// Test policy that panics when invoked. To be used as an error policy in
191    /// tests that do not expect request errors in order to verify that the
192    /// error policy is not invoked
193    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    /// Note that this sample policy is applied on top of the
215    /// endpoint-specific sample policy (not configurable) which
216    /// weighs endpoints by the relative effort required to serve
217    /// them. Therefore a sample rate of N will yield an actual
218    /// sample rate <= N.
219    pub spam_sample_rate: Weight,
220    #[serde(default = "default_dry_run")]
221    pub dry_run: bool,
222    /// List of String which should all parse to type IPAddr.
223    /// If set, only requests from provided IPs will be allowed,
224    /// and any blocklist related configuration will be ignored.
225    #[serde(default)]
226    pub allow_list: Option<Vec<String>>,
227    /// How many leading bits of an IPv6 address identify a client, from 32 to
228    /// 128. The node cannot see how long a client's own network prefix is, so
229    /// it assumes this length for every IPv6 client: the policies count, and
230    /// the node blocks, the addresses that share these bits as one client.
231    /// One subscriber usually holds a whole /64 and can switch addresses
232    /// within it. 128 treats every address as a client of its own. An IPv4
233    /// client is always its full address.
234    #[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        // A TTL of zero would expire every block at once, thus the policy would
311        // never keep a client out even with dry run off.
312        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}