s2n_quic_core/ack/
ranges.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    ack::Settings,
6    frame::ack,
7    interval_set::{IntervalSet, RangeInclusiveIter},
8    packet::number::{PacketNumber, PacketNumberRange},
9    varint::VarInt,
10};
11use core::{
12    convert::TryInto,
13    num::NonZeroUsize,
14    ops::{Bound, Deref, DerefMut, RangeInclusive},
15};
16
17#[derive(Clone, Debug)]
18pub struct Ranges(IntervalSet<PacketNumber>);
19
20impl Default for Ranges {
21    #[inline]
22    fn default() -> Self {
23        Self::new(Settings::default().ack_ranges_limit as usize)
24    }
25}
26
27impl Ranges {
28    #[inline]
29    pub fn new(limit: usize) -> Self {
30        let limit = NonZeroUsize::new(limit).expect("limit should be nonzero");
31        Self(IntervalSet::with_limit(limit))
32    }
33
34    /// Inserts a packet number; dropping smaller values if needed
35    #[inline]
36    pub fn insert_packet_number_range(&mut self, pn_range: PacketNumberRange) -> Result<(), Error> {
37        let interval = (
38            Bound::Included(pn_range.start()),
39            Bound::Included(pn_range.end()),
40        );
41        if self.0.insert(interval).is_ok() {
42            return Ok(());
43        }
44
45        // attempt to shed the lowest packet number ranges to make room for larger ones
46        match self.0.pop_min() {
47            Some(min) => {
48                if min < pn_range.start() {
49                    let insert_res = self.0.insert(interval);
50                    debug_assert!(
51                        insert_res.is_ok(),
52                        "min range was removed, so it should be possible to insert another range",
53                    );
54                    insert_res.map_err(|_| Error::RangeInsertionFailed {
55                        min: pn_range.start(),
56                        max: pn_range.end(),
57                    })?;
58
59                    Err(Error::LowestRangeDropped {
60                        min: min.start_inclusive(),
61                        max: min.end_inclusive(),
62                    })
63                } else {
64                    // new value is smaller than min so inset it back in the front
65                    let _ = self.0.insert_front(min);
66                    Err(Error::RangeInsertionFailed {
67                        min: pn_range.start(),
68                        max: pn_range.end(),
69                    })
70                }
71            }
72            None => {
73                debug_assert!(
74                    false,
75                    "IntervalSet should have capacity and return lowest entry"
76                );
77                Err(Error::RangeInsertionFailed {
78                    min: pn_range.start(),
79                    max: pn_range.end(),
80                })
81            }
82        }
83    }
84
85    /// Inserts a packet number; dropping smaller values if needed
86    #[inline]
87    pub fn insert_packet_number(&mut self, packet_number: PacketNumber) -> Result<(), Error> {
88        self.insert_packet_number_range(PacketNumberRange::new(packet_number, packet_number))
89    }
90
91    /// Returns the overall range of the ack_ranges
92    #[inline]
93    pub fn spread(&self) -> usize {
94        match (self.min_value(), self.max_value()) {
95            (Some(min), Some(max)) => {
96                let min = PacketNumber::as_varint(min);
97                let max = PacketNumber::as_varint(max);
98                (max - min).try_into().unwrap_or(usize::MAX)
99            }
100            _ => 0,
101        }
102    }
103}
104
105type Iter<'a> = core::iter::Map<
106    core::iter::Rev<RangeInclusiveIter<'a, PacketNumber>>,
107    fn(RangeInclusive<PacketNumber>) -> RangeInclusive<VarInt>,
108>;
109
110impl<'a> ack::AckRanges for &'a Ranges {
111    type Iter = Iter<'a>;
112
113    #[inline]
114    fn ack_ranges(&self) -> Self::Iter {
115        self.0.inclusive_ranges().rev().map(|range| {
116            let (start, end) = range.into_inner();
117            PacketNumber::as_varint(start)..=PacketNumber::as_varint(end)
118        })
119    }
120}
121
122impl Deref for Ranges {
123    type Target = IntervalSet<PacketNumber>;
124
125    #[inline]
126    fn deref(&self) -> &Self::Target {
127        &self.0
128    }
129}
130
131impl DerefMut for Ranges {
132    #[inline]
133    fn deref_mut(&mut self) -> &mut Self::Target {
134        &mut self.0
135    }
136}
137
138#[derive(Debug, Copy, Clone, PartialEq, Eq)]
139pub enum Error {
140    RangeInsertionFailed {
141        min: PacketNumber,
142        max: PacketNumber,
143    },
144    LowestRangeDropped {
145        min: PacketNumber,
146        max: PacketNumber,
147    },
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use crate::{
154        packet::number::{testing::iter as packet_numbers_iter, PacketNumberSpace},
155        varint,
156    };
157    use bolero::check;
158
159    #[test]
160    fn insert_value_test() {
161        let mut ack_ranges = Ranges::new(3);
162        let mut packet_numbers = packet_numbers_iter(PacketNumberSpace::ApplicationData).step_by(2); // skip every other packet number
163
164        // insert gaps up to the limit
165        let pn_a = packet_numbers.next().unwrap();
166        assert!(ack_ranges.insert_packet_number(pn_a).is_ok());
167
168        let pn_b = packet_numbers.next().unwrap();
169        assert!(ack_ranges.insert_packet_number(pn_b).is_ok());
170
171        let pn_c = packet_numbers.next().unwrap();
172        assert!(ack_ranges.insert_packet_number(pn_c).is_ok());
173
174        // ensure everything was inserted
175        assert_eq!(ack_ranges.interval_len(), 3);
176        assert!(ack_ranges.contains(&pn_a));
177        assert!(ack_ranges.contains(&pn_b));
178        assert!(ack_ranges.contains(&pn_c));
179
180        // insert a new packet number gap
181        let pn_d = packet_numbers.next().unwrap();
182        assert_eq!(
183            ack_ranges.insert_packet_number(pn_d).err().unwrap(),
184            Error::LowestRangeDropped {
185                min: pn_a,
186                max: pn_a
187            }
188        );
189
190        // ensure the previous smaller packet number was dropped
191        assert_eq!(ack_ranges.interval_len(), 3);
192        assert!(!ack_ranges.contains(&pn_a));
193        assert!(ack_ranges.contains(&pn_b));
194        assert!(ack_ranges.contains(&pn_c));
195        assert!(ack_ranges.contains(&pn_d));
196
197        // ensure smaller values are not recorded
198        assert_eq!(
199            ack_ranges.insert_packet_number(pn_a).err().unwrap(),
200            Error::RangeInsertionFailed {
201                min: pn_a,
202                max: pn_a
203            }
204        );
205        assert_eq!(ack_ranges.interval_len(), 3);
206        assert!(!ack_ranges.contains(&pn_a));
207        assert!(ack_ranges.contains(&pn_b));
208        assert!(ack_ranges.contains(&pn_c));
209        assert!(ack_ranges.contains(&pn_d));
210    }
211
212    #[test]
213    fn overlapping_range_test() {
214        let mut packet_numbers = packet_numbers_iter(PacketNumberSpace::ApplicationData).step_by(2); // skip every other packet number
215        let mut ack_ranges = Ranges::new(3);
216
217        // |---a---b---c---d---e---f---g---h---i---|
218        //     ^0
219        //         ^-1-^
220        //                 ^--2----^
221        //                             ^-3-^
222        //                                 ^-4-^
223        //             ^--1_2--^
224        //     ^--0_1--^
225        let pn_a = packet_numbers.next().unwrap();
226        let pn_b = packet_numbers.next().unwrap();
227        let pn_c = packet_numbers.next().unwrap();
228        let pn_d = packet_numbers.next().unwrap();
229        let pn_e = packet_numbers.next().unwrap();
230        let pn_f = packet_numbers.next().unwrap();
231        let pn_g = packet_numbers.next().unwrap();
232        let pn_h = packet_numbers.next().unwrap();
233        let pn_i = packet_numbers.next().unwrap();
234        let range_0 = PacketNumberRange::new(pn_a, pn_a);
235        let range_1 = PacketNumberRange::new(pn_b, pn_c);
236        let range_2 = PacketNumberRange::new(pn_d, pn_f);
237        let range_3 = PacketNumberRange::new(pn_g, pn_h);
238        let range_4 = PacketNumberRange::new(pn_h, pn_i);
239        let range_0_1 = PacketNumberRange::new(pn_a, pn_c);
240        let range_1_2 = PacketNumberRange::new(pn_c, pn_e);
241
242        assert!(ack_ranges.insert_packet_number_range(range_1).is_ok());
243        assert!(ack_ranges.insert_packet_number_range(range_2).is_ok());
244        assert!(ack_ranges.insert_packet_number_range(range_3).is_ok());
245        assert_eq!(ack_ranges.interval_len(), 3);
246
247        for pn in range_1 {
248            assert!(ack_ranges.contains(&pn));
249        }
250        for pn in range_2 {
251            assert!(ack_ranges.contains(&pn));
252        }
253        for pn in range_3 {
254            assert!(ack_ranges.contains(&pn));
255        }
256
257        // merge ranges 1 and 2
258        assert!(ack_ranges.insert_packet_number_range(range_1_2).is_ok());
259        assert_eq!(ack_ranges.interval_len(), 2);
260
261        // insert range 0 at low end
262        assert!(ack_ranges.insert_packet_number_range(range_0).is_ok());
263        assert_eq!(ack_ranges.interval_len(), 3);
264
265        // merge range 0_1 at low end
266        assert!(ack_ranges.insert_packet_number_range(range_0_1).is_ok());
267        assert_eq!(ack_ranges.interval_len(), 2);
268
269        // merge new range at high end
270        assert!(ack_ranges.insert_packet_number_range(range_4).is_ok());
271        assert_eq!(ack_ranges.interval_len(), 2);
272    }
273
274    #[test]
275    fn large_range_test() {
276        let pn_a = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u32(1));
277        let pn_b = PacketNumberSpace::ApplicationData
278            .new_packet_number(VarInt::new(varint::MAX_VARINT_VALUE).unwrap());
279        let mut ack_ranges = Ranges::new(3);
280
281        let range_1 = PacketNumberRange::new(pn_a, pn_b);
282
283        assert!(ack_ranges.insert_packet_number_range(range_1).is_ok());
284        assert_eq!(ack_ranges.interval_len(), 1);
285    }
286
287    #[test]
288    #[cfg_attr(miri, ignore)]
289    #[cfg(target_pointer_width = "64")]
290    fn size_of_snapshots() {
291        use core::mem::size_of;
292        use insta::assert_debug_snapshot;
293
294        assert_debug_snapshot!("Ranges", size_of::<Ranges>());
295    }
296
297    #[test]
298    fn insert_packet_number_range_fuzz() {
299        check!()
300            .with_type::<(u32, u32)>()
301            .map(|(a, b)| (a.min(b), a.max(b))) // ensure valid range
302            .for_each(|(a, b)| {
303                let mut ack_ranges = Ranges::new(1);
304
305                let pn_a = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u32(*a));
306                let pn_b = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u32(*b));
307                let range_1 = PacketNumberRange::new(pn_a, pn_b);
308
309                assert!(ack_ranges.insert_packet_number_range(range_1).is_ok());
310            });
311    }
312}
s2n_quic_core/ack/ranges.rs

s2n_quic_core/ack/
ranges.rs
