Allow only u32 ref counts (#721)

Author: fafhrd91

Cargo.toml

@@ -40,8 +40,8 @@ ntex-macros = { path = "ntex-macros" }
 ntex-util = { path = "ntex-util" }
 
 [workspace.dependencies]
-ntex = "3.0.0-pre.11"
-ntex-bytes = "1.4.0"
+ntex = "3.0.0-pre.13"
+ntex-bytes = "1.4.1"
 ntex-codec = "1.1.0"
 ntex-io = "3.5.0"
 ntex-net = "3.5.1"

ntex-bytes/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [1.4.1] (2026-01-22)
+
+* Allow only u32 ref counts
+
 ## [1.4.0] (2026-01-21)
 
 * Add BytesMut::reserve_capacity() method

ntex-bytes/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex-bytes"
-version = "1.4.0"
+version = "1.4.1"
 authors = ["Nikolay Kim <fafhrd91@gmail.com>"]
 description = "Types and traits for working with bytes (bytes crate fork)"
 documentation = "https://docs.rs/ntex-bytes"

ntex-bytes/src/buf/buf_mut.rs

@@ -1026,11 +1026,11 @@ mod tests {
         let mut buf = BytesMut::new();
         buf.put_i8(0x01);
         assert_eq!(buf, b"\x01"[..]);
-        assert_eq!((&mut buf).remaining_mut(), 103);
+        assert_eq!((&mut buf).remaining_mut(), 107);
         let chunk = (&mut buf).chunk_mut();
         chunk.write_byte(0, b'9');
         unsafe { (&mut buf).advance_mut(1) };
-        assert_eq!((&mut buf).remaining_mut(), 102);
+        assert_eq!((&mut buf).remaining_mut(), 106);
 
         let mut buf = vec![];
         buf.put_i16(0x0809);

ntex-bytes/src/bvec.rs

@@ -112,7 +112,7 @@ impl BytesMut {
     #[inline]
     pub fn new() -> BytesMut {
         BytesMut {
-            storage: StorageVec::with_capacity(104),
+            storage: StorageVec::with_capacity(crate::storage::MIN_CAPACITY),
         }
     }
 

ntex-bytes/src/lib.rs

@@ -80,7 +80,7 @@ pub mod info {
     #[derive(Copy, Clone, Debug, Eq, PartialEq)]
     pub struct Info {
         pub id: usize,
-        pub refs: usize,
+        pub refs: u32,
         pub kind: Kind,
         pub capacity: usize,
     }

ntex-bytes/src/storage.rs

@@ -96,7 +96,7 @@
 use crate::alloc::alloc::{self, Layout, LayoutError};
 
 use std::sync::atomic::Ordering::{Acquire, Relaxed, Release};
-use std::sync::atomic::{self, AtomicUsize};
+use std::sync::atomic::{self, AtomicU32};
 use std::{cmp, mem, num::NonZeroUsize, ptr, ptr::NonNull, slice};
 
 use crate::{info::Info, info::Kind};
@@ -117,12 +117,14 @@ pub(crate) struct Storage {
     offset: NonZeroUsize,
 }
 
+#[derive(Debug)]
 /// Thread-safe reference-counted container for the shared storage.
 struct SharedVec {
-    len: u32,
     offset: u32,
+    len: u32,
+    capacity: u32,
     remaining: u32,
-    ref_count: AtomicUsize,
+    ref_count: AtomicU32,
 }
 
 pub(crate) struct StorageVec(NonNull<SharedVec>);
@@ -138,6 +140,8 @@ const KIND_OFFSET_BITS: usize = 2;
 pub const METADATA_SIZE: usize = mem::size_of::<SharedVec>();
 const METADATA_SIZE_U32: u32 = METADATA_SIZE as u32;
 
+pub(crate) const MIN_CAPACITY: usize = 128 - crate::METADATA_SIZE;
+
 // Bit op constants for extracting the inline length value from the `ptr` field.
 const INLINE_LEN_MASK: usize = 0b1111_1100;
 
@@ -493,7 +497,7 @@ impl Storage {
             // ptr points to SharedVec
             let shared = self.shared_vec();
             let ref_cnt = (*shared).ref_count.fetch_add(1, Relaxed);
-            if ref_cnt == usize::MAX {
+            if ref_cnt == u32::MAX {
                 abort();
             }
 
@@ -649,8 +653,7 @@ impl StorageVec {
     }
 
     pub(crate) fn capacity(&self) -> usize {
-        let ptr = unsafe { &*self.0.as_ptr() };
-        ptr.len as usize + ptr.remaining as usize
+        unsafe { (*self.0.as_ptr()).capacity as usize }
     }
 
     pub(crate) fn remaining(&self) -> usize {
@@ -686,7 +689,10 @@ impl StorageVec {
                 Storage::from_ptr_inline(ptr, at)
             } else {
                 let inner = self.as_inner();
-                inner.ref_count.fetch_add(1, Relaxed);
+                let ref_cnt = inner.ref_count.fetch_add(1, Relaxed);
+                if ref_cnt == u32::MAX {
+                    abort();
+                }
 
                 let offset = inner.offset as usize;
                 Storage {
@@ -710,12 +716,11 @@ impl StorageVec {
             if len == 0 {
                 let inner = self.as_inner();
                 if inner.is_unique() && inner.offset != METADATA_SIZE_U32 {
-                    let cap = (inner.offset as usize)
-                        + inner.len as usize
-                        + inner.remaining as usize;
+                    let cap = (inner.offset as usize) + inner.capacity as usize;
                     inner.len = 0;
                     inner.offset = METADATA_SIZE_U32;
-                    inner.remaining = (cap - METADATA_SIZE) as u32;
+                    inner.capacity = (cap - METADATA_SIZE) as u32;
+                    inner.remaining = inner.capacity;
                     return;
                 }
             }
@@ -768,16 +773,15 @@ impl StorageVec {
             let new_cap = len + additional;
 
             if inner.is_unique() {
-                let capacity = (inner.offset as usize)
-                    + (inner.len as usize)
-                    + (inner.remaining as usize);
+                let capacity = (inner.offset as usize) + (inner.capacity as usize);
 
                 // try to reclaim the buffer. This is possible if the current
                 // handle is the only outstanding handle pointing to the buffer.
                 if capacity >= (new_cap + METADATA_SIZE) {
                     let offset = inner.offset;
                     inner.offset = METADATA_SIZE_U32;
                     inner.remaining = (capacity - len - METADATA_SIZE) as u32;
+                    inner.capacity = inner.len + inner.remaining;
 
                     // The capacity is sufficient, reclaim the buffer
                     if len != 0 {
@@ -794,24 +798,23 @@ impl StorageVec {
 
     #[inline]
     pub(crate) unsafe fn set_len(&mut self, len: usize) {
-        let cap = self.capacity();
-        assert!(len <= cap);
+        let inner = self.0.as_mut();
+        assert!(len as u32 <= inner.capacity);
 
-        let vec = self.0.as_mut();
-        vec.len = len as u32;
-        vec.remaining = (cap - len) as u32;
+        inner.len = len as u32;
+        inner.remaining = inner.capacity - (len as u32);
     }
 
     pub(crate) unsafe fn set_start(&mut self, start: u32) {
         if start != 0 {
-            let cap = self.capacity();
             let inner = self.as_inner();
 
             assert!(
-                start <= cap as u32,
-                "Cannot set start position offset:{} len:{} remaining:{}, new-len:{start}",
+                start <= inner.capacity,
+                "Cannot set start position offset:{} len:{} cap:{} remaining:{} new-len:{start}",
                 inner.offset,
                 inner.len,
+                inner.capacity,
                 inner.remaining,
             );
 
@@ -825,7 +828,8 @@ impl StorageVec {
             } else {
                 inner.len = 0;
             }
-            inner.remaining = cap as u32 - inner.len - start;
+            inner.remaining = inner.capacity - inner.len - start;
+            inner.capacity = inner.remaining + inner.len;
         }
     }
 }
@@ -861,14 +865,16 @@ impl SharedVec {
             if ptr.is_null() {
                 alloc::handle_alloc_error(layout);
             }
+            let capacity = (layout.size() - METADATA_SIZE) as u32;
 
             ptr::write(
                 ptr as *mut SharedVec,
                 SharedVec {
                     len,
+                    capacity,
+                    remaining: capacity - len,
                     offset: METADATA_SIZE_U32,
-                    remaining: (layout.size() - METADATA_SIZE - len as usize) as u32,
-                    ref_count: AtomicUsize::new(1),
+                    ref_count: AtomicU32::new(1),
                 },
             );
             ptr
@@ -881,7 +887,7 @@ impl SharedVec {
     }
 
     fn capacity(&self) -> usize {
-        self.len as usize + self.remaining as usize
+        self.capacity as usize
     }
 }
 
@@ -912,7 +918,7 @@ fn release_shared_vec(ptr: *mut SharedVec) {
         atomic::fence(Acquire);
 
         // Drop the data
-        let cap = (*ptr).offset as usize + (*ptr).remaining as usize + (*ptr).len as usize;
+        let cap = (*ptr).offset as usize + (*ptr).capacity as usize;
         ptr::drop_in_place(ptr);
         let layout = shared_vec_layout(cap - METADATA_SIZE).unwrap();
         alloc::dealloc(ptr as *mut _, layout);

ntex-bytes/tests/test_bytes.rs

@@ -20,7 +20,7 @@ fn is_send<T: Send>() {}
 fn test_size() {
     assert_eq!(24, std::mem::size_of::<Bytes>());
     assert_eq!(24, std::mem::size_of::<Option<Bytes>>());
-    assert_eq!(24, ntex_bytes::METADATA_SIZE);
+    assert_eq!(20, ntex_bytes::METADATA_SIZE);
 
     let mut t = BytesMut::new();
     t.extend_from_slice(&b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"[..]);
@@ -409,15 +409,15 @@ fn fns_defined_for_bytes() {
     let bytes = Bytes::copy_from_slice(&B[..]);
     assert!(bytes.info().kind == Kind::Vec);
     assert!(bytes.info().refs == 1);
-    assert_eq!(bytes.info().capacity, 76);
+    assert_eq!(bytes.info().capacity, 72);
     let b2 = bytes.clone();
     assert!(b2.info().kind == Kind::Vec);
     assert!(b2.info().refs == 2);
-    assert!(b2.info().capacity == 76);
+    assert!(b2.info().capacity == 72);
     drop(b2);
     assert!(bytes.info().kind == Kind::Vec);
     assert!(bytes.info().refs == 1);
-    assert!(bytes.info().capacity == 76);
+    assert!(bytes.info().capacity == 72);
 
     let mut bytes = Bytes::from(&b"hello world"[..]);