diff --git a/Cargo.lock b/Cargo.lock index c61a3bacf3..d989793744 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1948,6 +1948,7 @@ dependencies = [ "rustls", "rustls-pki-types", "rustls-platform-verifier", + "rustversion", "slab", "thiserror 2.0.17", "tinyvec", diff --git a/Cargo.toml b/Cargo.toml index f7fc50b435..4b10ebda6a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ rand = "0.9" rcgen = "0.14" ring = "0.17" rustc-hash = "2" +rustversion = "1" rustls = { version = "0.23.5", default-features = false, features = ["std"] } rustls-platform-verifier = "0.6" rustls-pki-types = "1.7" diff --git a/quinn-proto/Cargo.toml b/quinn-proto/Cargo.toml index 9d04b08a1e..bbfa32ff0d 100644 --- a/quinn-proto/Cargo.toml +++ b/quinn-proto/Cargo.toml @@ -57,6 +57,7 @@ slab = { workspace = true } thiserror = { workspace = true } tinyvec = { workspace = true, features = ["alloc"] } tracing = { workspace = true } +rustversion = { workspace = true } # Feature flags & dependencies for wasm # wasm-bindgen is assumed for a wasm*-*-unknown target diff --git a/quinn-proto/src/range_set/btree_range_set.rs b/quinn-proto/src/range_set/btree_range_set.rs index 9b9757ec00..c3da15afee 100644 --- a/quinn-proto/src/range_set/btree_range_set.rs +++ b/quinn-proto/src/range_set/btree_range_set.rs @@ -1,5 +1,6 @@ #[cfg(test)] use std::cmp::Ordering; + use std::{ cmp, collections::{BTreeMap, btree_map}, @@ -58,6 +59,37 @@ impl RangeSet { true } + #[rustversion::since(1.91)] + pub(crate) fn insert(&mut self, mut x: Range) -> bool { + if x.is_empty() { + return false; + } + if let Some((start, end)) = self.pred(x.start) { + if end >= x.end { + // Wholly contained + return false; + } else if end >= x.start { + // Extend overlapping predecessor + self.0.remove(&start); + x.start = start; + } + } + + let existing_end = self + .0 + .extract_if((Excluded(x.start), Included(x.end)), |_, _| true) + .last() + .map(|(_start, end)| end); + + if let Some(existing_end) = existing_end { + x.end = cmp::max(x.end, existing_end); + } + + self.0.insert(x.start, x.end); + true + } + + #[rustversion::before(1.91)] pub(crate) fn insert(&mut self, mut x: Range) -> bool { if x.is_empty() { return false; @@ -83,7 +115,6 @@ impl RangeSet { self.0.insert(x.start, x.end); true } - /// Find closest range to `x` that begins at or before it fn pred(&self, x: u64) -> Option<(u64, u64)> { self.0 @@ -101,6 +132,52 @@ impl RangeSet { } #[cfg(test)] + #[rustversion::since(1.91)] + pub(super) fn remove(&mut self, x: Range) -> bool { + if x.is_empty() { + return false; + } + + let before = match self.pred(x.start) { + Some((start, end)) if end > x.start => { + self.0.remove(&start); + if start < x.start { + self.0.insert(start, x.start); + } + if end > x.end { + self.0.insert(x.end, end); + } + // Short-circuit if we cannot possibly overlap with another range + if end >= x.end { + return true; + } + true + } + Some(_) | None => false, + }; + + let removed_end = self + .0 + .extract_if((Excluded(x.start), Excluded(x.end)), |_, _| true) + .last() + .map(|(_start, end)| end); + + let mut after = false; + + if let Some(removed_end) = removed_end { + after = true; + + let over_removed = removed_end > x.end; + if over_removed { + self.0.insert(x.end, removed_end); + } + } + + before || after + } + + #[cfg(test)] + #[rustversion::before(1.91)] pub(super) fn remove(&mut self, x: Range) -> bool { if x.is_empty() { return false; @@ -392,4 +469,231 @@ mod tests { assert_eq!(set.len(), 1); assert_eq!(set.peek_min().unwrap(), 0..4); } + + #[test] + fn insert_extends_prev_joint_range() { + // given + let mut set = RangeSet::new(); + set.insert(0..10); + + // when + set.insert(10..20); + + // then + let expected_inner = BTreeMap::from_iter([(0, 20)]); + assert_eq!(set.0, expected_inner, "ranges should have merged"); + } + + #[test] + fn insert_merges_touching_range() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + set.insert(5..10); + + // then + let expected_inner = BTreeMap::from_iter([(5, 20)]); + assert_eq!(set.0, expected_inner, "ranges should have merged."); + } + + #[test] + fn insert_merges_existing_crossing_range() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + set.insert(12..25); + + // then + let expected_inner = BTreeMap::from_iter([(10, 25)]); + assert_eq!(set.0, expected_inner, "ranges should have merged."); + } + + #[test] + fn insert_ignored_if_wholly_contained() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + let insert_result = set.insert(14..16); + + // then + let expected_inner = BTreeMap::from_iter([(10, 20)]); + assert_eq!(set.0, expected_inner, "insertion should be ignored"); + assert!( + !insert_result, + "insert should return false when inserting an already contained range" + ); + } + + #[test] + fn insert_bridges_disjoint_ranges() { + // given + let mut set = RangeSet::new(); + set.insert(0..10); + set.insert(20..30); + + // when + set.insert(10..20); + + // then + let expected_inner = BTreeMap::from_iter([(0, 30)]); + assert_eq!(set.0, expected_inner, "existing ranges should have merged"); + } + + #[test] + fn remove_splits_contained_range() { + // given + let mut set = RangeSet::new(); + set.insert(0..10); + + // when + let changed = set.remove(2..8); + + // then + assert!(changed); + let expected_inner = BTreeMap::from_iter([(0, 2), (8, 10)]); + assert_eq!(set.0, expected_inner, "range should have split into two"); + } + + #[test] + fn remove_trims_start_of_range() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + let changed = set.remove(5..15); + + // then + assert!(changed); + let expected_inner = BTreeMap::from_iter([(15, 20)]); + assert_eq!( + set.0, expected_inner, + "start of range should have been trimmed" + ); + } + + #[test] + fn remove_trims_end_of_range() { + // given + let mut set = RangeSet::new(); + set.insert(0..10); + + // when + let changed = set.remove(5..15); + + // then + assert!(changed); + let expected_inner = BTreeMap::from_iter([(0, 5)]); + assert_eq!( + set.0, expected_inner, + "end of range should have been trimmed" + ); + } + + #[test] + fn remove_swallows_multiple_ranges() { + // given + let mut set = RangeSet::new(); + set.insert(0..5); + set.insert(10..15); + set.insert(20..25); + + // when + let changed = set.remove(0..25); + + // then + assert!(changed); + let expected_inner = BTreeMap::new(); + assert_eq!( + set.0, expected_inner, + "all ranges should be gone as the removal swallowed them" + ); + } + + #[test] + fn remove_specific_range_in_set() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + let changed = set.remove(10..20); + + // then + assert!(changed); + let expected_inner = BTreeMap::new(); + assert_eq!(set.0, expected_inner, "set should be empty"); + } + + #[test] + fn removing_range_touching_existing_ranges_is_ignored() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + let changed_start = set.remove(0..10); + let changed_end = set.remove(20..30); + + // then + assert!(!changed_start); + assert!(!changed_end); + let expected_inner = BTreeMap::from_iter([(10, 20)]); + assert_eq!( + set.0, expected_inner, + "touching ranges should not remove anything." + ); + } + + #[test] + fn removing_range_crossing_existing_ranges_that_should_be_trimmed() { + // given + let mut set = RangeSet::new(); + set.insert(0..5); + set.insert(10..15); + + // when + let changed = set.remove(4..11); + + // then + assert!(changed); + let expected_inner = BTreeMap::from_iter([(0, 4), (11, 15)]); + assert_eq!(set.0, expected_inner, "existing ranges should be trimmed"); + } + + #[test] + fn remove_crossing_range_trims_start() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + let changed = set.remove(9..12); + + // then + assert!(changed); + let expected_inner = BTreeMap::from_iter([(12, 20)]); + assert_eq!(set.0, expected_inner); + } + + #[test] + fn remove_no_overlap_before_and_after() { + // given + let mut set = RangeSet::new(); + set.insert(10..20); + + // when + let changed = set.remove(0..5); + + // then + assert!(!changed); + let expected_inner = BTreeMap::from_iter([(10, 20)]); + assert_eq!(set.0, expected_inner, "disjoint removals should be ignored"); + } }