fix: memory leak in icons.rs

mediar-ai · Nov 6, 2024 · 4a1950f · 4a1950f
1 parent 6f7dc35
commit 4a1950f
Showing 1 changed file with 47 additions and 34 deletions.
diff --git a/screenpipe-app-tauri/src-tauri/src/icons.rs b/screenpipe-app-tauri/src-tauri/src/icons.rs
@@ -1,4 +1,3 @@
-
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -13,50 +12,64 @@ pub async fn get_app_icon(
     app_name: &str,
     app_path: Option<String>,
 ) -> Result<Option<AppIcon>, String> {
-    // info!("getting icon for {}", app_name);
     use base64::{engine::general_purpose::STANDARD, Engine};
-    use cocoa::base::id;
-    use cocoa::foundation::{NSData, NSString};
+    use cocoa::base::{id, nil};
+    use cocoa::foundation::{NSAutoreleasePool, NSData, NSString};
     use objc::{class, msg_send, sel, sel_impl};
+
     unsafe {
-        let workspace: id = msg_send![class!(NSWorkspace), sharedWorkspace];
+        // Create autorelease pool
+        let pool = NSAutoreleasePool::new(nil);
+
+        let result = (|| {
+            let workspace: id = msg_send![class!(NSWorkspace), sharedWorkspace];
+
+            let path = if let Some(path) = app_path {
+                path
+            } else {
+                let ns_app_name = NSString::alloc(nil).init_str(app_name);
+                let path: id = msg_send![workspace, fullPathForApplication: ns_app_name];
+                // Release the NSString we created
+                let _: () = msg_send![ns_app_name, release];
+
+                if path == nil {
+                    return Ok(None);
+                }
+                let path: id = msg_send![path, UTF8String];
+                std::ffi::CStr::from_ptr(path as *const _)
+                    .to_string_lossy()
+                    .into_owned()
+            };
+
+            let ns_path = NSString::alloc(nil).init_str(&path);
+            let icon: id = msg_send![workspace, iconForFile:ns_path];
+            let _: () = msg_send![ns_path, release];
 
-        let path = if let Some(path) = app_path {
-            path
-        } else {
-            let ns_app_name = NSString::alloc(cocoa::base::nil).init_str(app_name);
-            let path: id = msg_send![workspace, fullPathForApplication: ns_app_name];
-            if path == cocoa::base::nil {
+            if icon == nil {
                 return Ok(None);
             }
-            let path: id = msg_send![path, UTF8String];
-            std::ffi::CStr::from_ptr(path as *const _)
-                .to_string_lossy()
-                .into_owned()
-        };
 
-        let icon: id =
-            msg_send![workspace, iconForFile:NSString::alloc(cocoa::base::nil).init_str(&path)];
-        if icon == cocoa::base::nil {
-            return Ok(None);
-        }
+            // Rest of the conversion logic remains the same
+            let tiff_data: id = msg_send![icon, TIFFRepresentation];
+            let image_rep: id = msg_send![class!(NSBitmapImageRep), imageRepWithData: tiff_data];
+            let png_data: id = msg_send![image_rep, representationUsingType:4 properties:nil];
 
-        // Convert to PNG data
-        let tiff_data: id = msg_send![icon, TIFFRepresentation];
-        let image_rep: id = msg_send![class!(NSBitmapImageRep), imageRepWithData: tiff_data];
-        let png_data: id =
-            msg_send![image_rep, representationUsingType:4 properties:cocoa::base::nil];
+            let length = NSData::length(png_data);
+            let bytes = NSData::bytes(png_data);
+            let data = std::slice::from_raw_parts(bytes as *const u8, length as usize);
 
-        let length = NSData::length(png_data);
-        let bytes = NSData::bytes(png_data);
-        let data = std::slice::from_raw_parts(bytes as *const u8, length as usize);
+            let base64 = STANDARD.encode(data);
 
-        let base64 = STANDARD.encode(data);
+            Ok(Some(AppIcon {
+                base64,
+                path: Some(path),
+            }))
+        })();
 
-        Ok(Some(AppIcon {
-            base64,
-            path: Some(path),
-        }))
+        // Drain the autorelease pool
+        let _: () = msg_send![pool, drain];
+
+        result
     }
 }