From 36675436cc343fc0e828033278d668020bd897b9 Mon Sep 17 00:00:00 2001
From: venomousmoog <venomousmoog@users.noreply.github.com>
Date: Wed, 5 Mar 2025 01:44:40 -0800
Subject: [PATCH] Fix memory leak in `owned_to_vec` (#773)

---
 src/bytes.rs        | 14 +++++++++-----
 tests/test_bytes.rs | 13 +++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/bytes.rs b/src/bytes.rs
index 96f834f..cdb6ea5 100644
--- a/src/bytes.rs
+++ b/src/bytes.rs
@@ -1141,15 +1141,15 @@ unsafe fn owned_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes
     }
 }
 
-unsafe fn owned_to_vec(_data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8> {
+unsafe fn owned_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8> {
     let slice = slice::from_raw_parts(ptr, len);
-    slice.to_vec()
+    let vec = slice.to_vec();
+    owned_drop_impl(data.load(Ordering::Relaxed));
+    vec
 }
 
 unsafe fn owned_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut {
-    let bytes_mut = BytesMut::from_vec(owned_to_vec(data, ptr, len));
-    owned_drop_impl(data.load(Ordering::Relaxed));
-    bytes_mut
+    BytesMut::from_vec(owned_to_vec(data, ptr, len))
 }
 
 unsafe fn owned_is_unique(_data: &AtomicPtr<()>) -> bool {
@@ -1161,6 +1161,10 @@ unsafe fn owned_drop_impl(owned: *mut ()) {
     let ref_cnt = &(*lifetime).ref_cnt;
 
     let old_cnt = ref_cnt.fetch_sub(1, Ordering::Release);
+    debug_assert!(
+        old_cnt > 0 && old_cnt <= usize::MAX >> 1,
+        "expected non-zero refcount and no underflow"
+    );
     if old_cnt != 1 {
         return;
     }
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index c283a80..613efc8 100644
--- a/tests/test_bytes.rs
+++ b/tests/test_bytes.rs
@@ -1619,6 +1619,19 @@ fn owned_to_vec() {
     assert_eq!(drop_counter.get(), 1);
 }
 
+#[test]
+fn owned_into_vec() {
+    let drop_counter = SharedAtomicCounter::new();
+    let buf: [u8; 10] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    let owner = OwnedTester::new(buf, drop_counter.clone());
+    let b1 = Bytes::from_owner(owner);
+
+    let v1: Vec<u8> = b1.into();
+    assert_eq!(&v1[..], &buf[..]);
+    // into() vec will copy out of the owner and drop it
+    assert_eq!(drop_counter.get(), 1);
+}
+
 #[test]
 #[cfg_attr(not(panic = "unwind"), ignore)]
 fn owned_safe_drop_on_as_ref_panic() {