s2n_quic_core/packet/interceptor/
loss.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::{havoc, Interceptor};
5use core::ops::Range;
6
7#[derive(Debug)]
8struct Direction {
9    loss: Range<u64>,
10    pass: Range<u64>,
11    mode: Mode,
12}
13
14impl Default for Direction {
15    fn default() -> Self {
16        Self {
17            loss: 0..0,
18            pass: u64::MAX..u64::MAX,
19            mode: Mode::Loss { remaining: 0 },
20        }
21    }
22}
23
24impl Direction {
25    #[inline]
26    fn should_pass<R: havoc::Random>(&mut self, random: &mut R) -> bool {
27        let (remaining, mut is_pass) = match &mut self.mode {
28            Mode::Pass { remaining } => (remaining, true),
29            Mode::Loss { remaining } => (remaining, false),
30        };
31
32        // if we can decrement the remaining, then stay on the current mode
33        if let Some(new_value) = remaining.checked_sub(1) {
34            *remaining = new_value;
35            return is_pass;
36        }
37
38        // try 3 times to generate the next value. It's good to be an odd number so we're at least
39        // alternating if the ranges keep returning 0.
40        for _ in 0..3 {
41            // go to the next mode
42            is_pass = !is_pass;
43
44            let remaining = if is_pass {
45                Self::gen_range(&self.pass, random)
46            } else {
47                Self::gen_range(&self.loss, random)
48            };
49
50            // if this round picked 0 then try again
51            if remaining == 0 {
52                continue;
53            }
54
55            if is_pass {
56                self.mode = Mode::Pass { remaining };
57            } else {
58                self.mode = Mode::Loss { remaining };
59            }
60        }
61
62        is_pass
63    }
64
65    #[inline]
66    fn gen_range<R: havoc::Random>(range: &Range<u64>, random: &mut R) -> u64 {
67        if range.start == range.end {
68            return range.start;
69        }
70
71        random.gen_range(range.clone())
72    }
73}
74
75#[derive(Debug)]
76enum Mode {
77    Loss { remaining: u64 },
78    Pass { remaining: u64 },
79}
80
81#[derive(Debug, Default)]
82pub struct Builder<R> {
83    tx: Direction,
84    rx: Direction,
85    random: R,
86}
87
88impl<R> Builder<R>
89where
90    R: 'static + Send + havoc::Random,
91{
92    pub fn new(random: R) -> Self {
93        Self {
94            tx: Direction::default(),
95            rx: Direction::default(),
96            random,
97        }
98    }
99
100    pub fn with_tx_pass(mut self, range: Range<u64>) -> Self {
101        self.tx.pass = range;
102        self
103    }
104
105    pub fn with_tx_loss(mut self, range: Range<u64>) -> Self {
106        self.tx.loss = range;
107        self
108    }
109
110    pub fn with_rx_pass(mut self, range: Range<u64>) -> Self {
111        self.rx.pass = range;
112        self
113    }
114
115    pub fn with_rx_loss(mut self, range: Range<u64>) -> Self {
116        self.rx.loss = range;
117        self
118    }
119
120    pub fn build(self) -> Loss<R> {
121        Loss {
122            tx: self.tx,
123            rx: self.rx,
124            random: self.random,
125        }
126    }
127}
128
129#[derive(Debug, Default)]
130pub struct Loss<R>
131where
132    R: 'static + Send + havoc::Random,
133{
134    tx: Direction,
135    rx: Direction,
136    random: R,
137}
138
139impl<R> Loss<R>
140where
141    R: 'static + Send + havoc::Random,
142{
143    pub fn builder(random: R) -> Builder<R> {
144        Builder::new(random)
145    }
146}
147
148impl<R> Interceptor for Loss<R>
149where
150    R: 'static + Send + havoc::Random,
151{
152    #[inline]
153    fn intercept_rx_datagram<'a>(
154        &mut self,
155        _subject: &crate::event::api::Subject,
156        _datagram: &super::Datagram,
157        payload: s2n_codec::DecoderBufferMut<'a>,
158    ) -> s2n_codec::DecoderBufferMut<'a> {
159        if !self.rx.should_pass(&mut self.random) {
160            return s2n_codec::DecoderBufferMut::new(&mut payload.into_less_safe_slice()[..0]);
161        }
162
163        payload
164    }
165
166    #[inline]
167    fn intercept_tx_datagram(
168        &mut self,
169        _subject: &crate::event::api::Subject,
170        _datagram: &super::Datagram,
171        payload: &mut s2n_codec::EncoderBuffer,
172    ) {
173        if !self.tx.should_pass(&mut self.random) {
174            payload.set_position(0);
175        }
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::{havoc::testing::RandomSlice, *};
182
183    #[test]
184    fn alternate_test() {
185        static SLICE: &[u8] = &{
186            let mut slice = [0u8; 256];
187            let mut i = 0;
188            while i < slice.len() {
189                slice[i] = i as _;
190                i += 1;
191            }
192            slice
193        };
194
195        let mut rand = RandomSlice::new(SLICE);
196
197        let mut rx = Direction {
198            loss: 0..10,
199            pass: 1..10,
200            ..Default::default()
201        };
202
203        let mut passed = 0;
204        let mut dropped = 0;
205
206        for _ in 0..256 {
207            if rx.should_pass(&mut rand) {
208                passed += 1;
209            } else {
210                dropped += 1;
211            }
212        }
213
214        // these values will always be the same since the Random generator is deterministic
215        assert_eq!(passed, 143);
216        assert_eq!(dropped, 113);
217    }
218}
s2n_quic_core/packet/interceptor/loss.rs

s2n_quic_core/packet/interceptor/
loss.rs
