From 39b0c17f12387ef81ee23fcbd56665b569bfba99 Mon Sep 17 00:00:00 2001 From: Diggory Hardy Date: Mon, 17 Apr 2023 09:28:18 +0100 Subject: [PATCH] WeightedIndex::new: trap overflow in release builds only --- src/distributions/weighted_index.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/distributions/weighted_index.rs b/src/distributions/weighted_index.rs index 05b9c6106c..c71842d66b 100644 --- a/src/distributions/weighted_index.rs +++ b/src/distributions/weighted_index.rs @@ -117,7 +117,11 @@ impl WeightedIndex { return Err(WeightedError::InvalidWeight); } weights.push(total_weight.clone()); + total_weight += w.borrow(); + if total_weight < *w.borrow() { + return Err(WeightedError::Overflow); + } } if total_weight == zero { @@ -423,7 +427,7 @@ mod test { #[test] fn overflow() { - assert!(WeightedIndex::new([2, usize::MAX]).is_err()); + assert_eq!(WeightedIndex::new([2, usize::MAX]), Err(WeightedError::Overflow)); } } @@ -443,6 +447,9 @@ pub enum WeightedError { /// Too many weights are provided (length greater than `u32::MAX`) TooMany, + + /// The sum of weights overflows + Overflow, } #[cfg(feature = "std")] @@ -455,6 +462,7 @@ impl fmt::Display for WeightedError { WeightedError::InvalidWeight => "A weight is invalid in distribution", WeightedError::AllWeightsZero => "All weights are zero in distribution", WeightedError::TooMany => "Too many weights (hit u32::MAX) in distribution", + WeightedError::Overflow => "The sum of weights overflowed", }) } }