prevent use-after-free caused by overflow in ref_count

crates/libs/core/src/imp/ref_count.rs

@@ -1,19 +1,28 @@
-use std::sync::atomic::{fence, AtomicI32, Ordering};
+use std::sync::atomic::{fence, AtomicI32, Ordering, AtomicU32};
+
+const MAX_REF_COUNT: u32 = i32::MAX as _;
 
 #[doc(hidden)]
 #[repr(transparent)]
 #[derive(Default)]
-pub struct RefCount(pub(crate) AtomicI32);
+pub struct RefCount(AtomicU32);
 
 impl RefCount {
  /// Creates a new `RefCount` with an initial value of `1`.
  pub fn new(count: u32) -> Self {
- Self(AtomicI32::new(count as i32))
+ if count >= MAX_REF_COUNT {
+ panic!("count must be smaller than `MAX_REF_COUNT`");
+ }
+ Self(AtomicU32::new(count))
  }
 
  /// Increments the reference count, returning the new value.
  pub fn add_ref(&self) -> u32 {
- (self.0.fetch_add(1, Ordering::Relaxed) + 1) as u32
+ let new_count = self.0.fetch_add(1, Ordering::Relaxed) + 1;
+ if new_count >= MAX_REF_COUNT { 
+ std::process::abort();
+ }
+ new_count
  }
 
  /// Decrements the reference count, returning the new value.
@@ -23,12 +32,13 @@ impl RefCount {
  pub fn release(&self) -> u32 {
  let remaining = self.0.fetch_sub(1, Ordering::Release) - 1;
 
- match remaining.cmp(&0) {
- std::cmp::Ordering::Equal => fence(Ordering::Acquire),
- std::cmp::Ordering::Less => panic!("Object has been over-released."),
- std::cmp::Ordering::Greater => {}
+ if remaining == 0 {
+ fence(Ordering::Acquire) ;
+ }
+ if remaining == u32::MAX {
+ panic!("Object has been over-released")
  }
 
- remaining as u32
+ remaining
  }
 }