From f6023d48e52989cb3bdff159ca0c416b926f7810 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcin=20Miko=C5=82ajczyk?= <marcinmikolajcz@gmail.com>
Date: Tue, 2 Apr 2024 20:52:16 +0100
Subject: [PATCH] More proper module reference counting (#798)

---
 src/kernel/src/rtld/mod.rs    |  9 ++++++++-
 src/kernel/src/rtld/module.rs | 10 ++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/kernel/src/rtld/mod.rs b/src/kernel/src/rtld/mod.rs
index 504197e95..68aed61fd 100644
--- a/src/kernel/src/rtld/mod.rs
+++ b/src/kernel/src/rtld/mod.rs
@@ -66,6 +66,7 @@ impl RuntimeLinker {
 
         sys.register(591, &ld, Self::sys_dynlib_dlsym);
         sys.register(592, &ld, Self::sys_dynlib_get_list);
+        sys.register(593, &ld, Self::sys_dynlib_get_info);
         sys.register(594, &ld, Self::sys_dynlib_load_prx);
         sys.register(596, &ld, Self::sys_dynlib_do_copy_relocations);
         sys.register(598, &ld, Self::sys_dynlib_get_proc_param);
@@ -180,6 +181,8 @@ impl RuntimeLinker {
         let loaded = bin.list().skip(1).find(|m| m.path() == path);
 
         if let Some(v) = loaded {
+            *v.ref_count_mut() += 1;
+
             return Ok((v.clone(), bin));
         }
 
@@ -448,6 +451,10 @@ impl RuntimeLinker {
         Ok(SysOut::ZERO)
     }
 
+    fn sys_dynlib_get_info(self: &Arc<Self>, _td: &VThread, _i: &SysIn) -> Result<SysOut, SysErr> {
+        todo!()
+    }
+
     fn sys_dynlib_load_prx(self: &Arc<Self>, td: &VThread, i: &SysIn) -> Result<SysOut, SysErr> {
         // Not sure what is this. Maybe kernel only flags?
         let mut flags: u32 = i.args[1].try_into().unwrap();
@@ -920,7 +927,7 @@ impl RuntimeLinker {
         info.datasize = mem.data_segment().len().try_into().unwrap();
         info.unk4 = 3;
         info.unk6 = 2;
-        info.refcount = Arc::strong_count(md).try_into().unwrap();
+        info.refcount = *md.ref_count();
 
         // Copy module name.
         if flags & 2 == 0 || !md.flags().contains(ModuleFlags::UNK1) {
diff --git a/src/kernel/src/rtld/module.rs b/src/kernel/src/rtld/module.rs
index 1320afb5f..6b31b3609 100644
--- a/src/kernel/src/rtld/module.rs
+++ b/src/kernel/src/rtld/module.rs
@@ -47,6 +47,7 @@ pub struct Module {
     file_type: FileType,
     programs: Vec<Program>,
     symbols: Vec<Symbol>,
+    ref_count: Gutex<u32>,
 }
 
 impl Module {
@@ -187,6 +188,7 @@ impl Module {
             file_type,
             programs,
             symbols,
+            ref_count: gg.spawn(1),
         };
 
         if let Some(info) = file_info {
@@ -311,6 +313,14 @@ impl Module {
         self.symbols.as_ref()
     }
 
+    pub fn ref_count(&self) -> GutexReadGuard<'_, u32> {
+        self.ref_count.read()
+    }
+
+    pub fn ref_count_mut(&self) -> GutexWriteGuard<'_, u32> {
+        self.ref_count.write()
+    }
+
     /// # Safety
     /// `off` must be a valid offset without base adjustment of a function in the memory of this
     /// module.