diff --git a/src/draw_target.rs b/src/draw_target.rs index 378fe68f..137a01fd 100644 --- a/src/draw_target.rs +++ b/src/draw_target.rs @@ -1,3 +1,4 @@ +use std::convert::TryInto; use std::io; use std::sync::{Arc, RwLock, RwLockWriteGuard}; use std::time::{Duration, Instant}; @@ -350,7 +351,15 @@ impl RateLimiter { // We add `new` to `capacity`, subtract one for returning `true` from here, // then make sure it does not exceed a maximum of `MAX_BURST`, then store it. - self.capacity = Ord::min(MAX_BURST, self.capacity + new as u8 - 1); + self.capacity = Ord::min( + MAX_BURST, + u128::from(self.capacity) + .saturating_add(new) + .saturating_sub(1) + .clamp(0, u8::MAX.into()) + .try_into() + .unwrap(), + ); // Store `prev` for the next iteration after subtracting the `remainder`. self.prev = now - Duration::from_nanos(remainder as u64); true diff --git a/src/state.rs b/src/state.rs index f43c3bce..00d203aa 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,4 +1,5 @@ use std::borrow::Cow; +use std::convert::TryInto; use std::fmt; use std::io; use std::sync::atomic::{AtomicU64, AtomicU8, Ordering}; @@ -430,7 +431,15 @@ impl AtomicPosition { let (new, remainder) = ((diff / INTERVAL), (diff % INTERVAL)); // We add `new` to `capacity`, subtract one for returning `true` from here, // then make sure it does not exceed a maximum of `MAX_BURST`. - capacity = Ord::min(MAX_BURST, capacity + new as u8 - 1); + capacity = Ord::min( + MAX_BURST, + u64::from(capacity) + .saturating_add(new) + .saturating_sub(1) + .clamp(0, u8::MAX.into()) + .try_into() + .unwrap(), + ); // Then, we just store `capacity` and `prev` atomically for the next iteration self.capacity.store(capacity, Ordering::Release); @@ -579,4 +588,12 @@ mod tests { // Should not panic. pb.set_position(0); } + + #[test] + fn test_atomic_position_large_time_difference() { + let atomic_position = AtomicPosition::new(); + let later = atomic_position.start + Duration::from_nanos(INTERVAL * u64::from(u8::MAX)); + // Should not panic. + atomic_position.allow(later); + } }