From 3cfc52773782143afeb2b0fa93245ea3242d5746 Mon Sep 17 00:00:00 2001
From: Andreas <aweinlein@cloudflare.com>
Date: Thu, 15 Feb 2024 03:45:59 +0100
Subject: [PATCH] Fixed memory leak on Apple platforms (#17)

---
 src/apple.rs | 62 +++++++++++++++++++++++++++++++---------------------
 1 file changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/apple.rs b/src/apple.rs
index 8d71b54..fc1889a 100644
--- a/src/apple.rs
+++ b/src/apple.rs
@@ -2,42 +2,54 @@ extern crate libc;
 
 use std::num::NonZeroUsize;
 
-use self::libc::{kern_return_t, mach_msg_type_number_t, mach_port_t, thread_t};
-
-// This constant is from
-// /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk/
-// usr/include/mach/machine/thread_state.h.
-//
-// It has not been updated since Apple devices started to support 64-bit ARM (iOS), so it
-// should be very stable.
-const THREAD_STATE_MAX: i32 = 1296;
-#[allow(non_camel_case_types)]
-// https://github.com/apple/darwin-xnu/blob/a1babec6b135d1f35b2590a1990af3c5c5393479/osfmk/mach/mach_types.defs#L155
-type task_inspect_t = mach_port_t;
+use self::libc::{
+    kern_return_t, mach_port_t, natural_t, task_threads, thread_act_array_t, vm_size_t,
+};
+
 #[allow(non_camel_case_types)]
-// https://github.com/apple/darwin-xnu/blob/a1babec6b135d1f35b2590a1990af3c5c5393479/osfmk/mach/mach_types.defs#L238
-type thread_array_t = [thread_t; THREAD_STATE_MAX as usize];
+// https://developer.apple.com/documentation/kernel/mach_port_name_t
+type mach_port_name_t = natural_t;
 
 extern "C" {
-    // https://developer.apple.com/documentation/kernel/1537751-task_threads/
-    fn task_threads(
-        target_task: task_inspect_t,
-        act_list: *mut thread_array_t,
-        act_listCnt: *mut mach_msg_type_number_t,
+    // https://developer.apple.com/documentation/kernel/1402285-mach_vm_deallocate
+    fn mach_vm_deallocate(
+        target_task: mach_port_t,
+        address: *mut u32,
+        size: vm_size_t,
     ) -> kern_return_t;
+
+    // https://developer.apple.com/documentation/kernel/1578777-mach_port_deallocate
+    fn mach_port_deallocate(task: mach_port_t, name: mach_port_name_t) -> kern_return_t;
 }
 
 pub(crate) fn num_threads() -> Option<NonZeroUsize> {
-    // http://web.mit.edu/darwin/src/modules/xnu/osfmk/man/task_threads.html
-    let mut thread_state = [0u32; THREAD_STATE_MAX as usize];
+    // https://developer.apple.com/documentation/kernel/1537751-task_threads
+    let mut thread_list: thread_act_array_t = std::ptr::null_mut();
     let mut thread_count = 0;
 
-    // Safety: `mach_task_self` always returns a valid value, `thread_state` is large enough, and
-    // both it and `thread_count` are writable.
-    let result =
-        unsafe { task_threads(libc::mach_task_self(), &mut thread_state, &mut thread_count) };
+    // Safety:
+    //  - `mach_task_self` always returns a valid value,
+    //  - `thread_list` is a pointer that will point to kernel allocated memory that needs to be
+    //    deallocated if the call succeeds
+    let task = unsafe { libc::mach_task_self() };
+    let result = unsafe { task_threads(task, &mut thread_list, &mut thread_count) };
 
     if result == libc::KERN_SUCCESS {
+        // Deallocate the mach port rights for the threads
+        for thread in 0..thread_count {
+            unsafe {
+                mach_port_deallocate(task, *(thread_list.offset(thread as isize)) as u32);
+            }
+        }
+        // Deallocate the thread list
+        unsafe {
+            mach_vm_deallocate(
+                task,
+                thread_list,
+                std::mem::size_of::<mach_port_t>() * thread_count as usize,
+            );
+        }
+
         NonZeroUsize::new(thread_count as usize)
     } else {
         None