From 7cd5623d3623d4fb215444c0b667f3e98bf1dca1 Mon Sep 17 00:00:00 2001
From: kamkow1 <kamkow256@gmail.com>
Date: Wed, 14 Jan 2026 23:11:06 +0100
Subject: [PATCH] Use reference counting to track filetime of process PD

---
 kernel/amd64/mm.h        |  1 +
 kernel/amd64/proc.c      | 44 +++++++++++++++++++++++++++-------------
 kernel/amd64/sched1.c    |  2 +-
 kernel/proc/proc.c       | 12 +++++------
 kernel/proc/proc.h       |  2 +-
 kernel/syscall/syscall.c | 10 ++++-----
 6 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/kernel/amd64/mm.h b/kernel/amd64/mm.h
index 5aa6ceb..f48bc18 100644
--- a/kernel/amd64/mm.h
+++ b/kernel/amd64/mm.h
@@ -9,6 +9,7 @@
 struct pd {
   spin_lock_t lock;
   uintptr_t cr3_paddr;
+  atomic_int refs;
 };
 
 void amd64_load_kernel_cr3 (void);
diff --git a/kernel/amd64/proc.c b/kernel/amd64/proc.c
index 947e3ef..da4342d 100644
--- a/kernel/amd64/proc.c
+++ b/kernel/amd64/proc.c
@@ -27,9 +27,17 @@ struct proc* proc_from_elf (uint8_t* elf_contents) {
   atomic_store (&proc->state, PROC_READY);
   proc->pid = atomic_fetch_add (&pids, 1);
 
-  proc->pd.lock = SPIN_LOCK_INIT;
-  proc->pd.cr3_paddr = mm_alloc_user_pd_phys ();
-  if (proc->pd.cr3_paddr == 0) {
+  proc->pd = malloc (sizeof (*proc->pd));
+  if (proc->pd == NULL) {
+    free (proc);
+    return NULL;
+  }
+
+  proc->pd->lock = SPIN_LOCK_INIT;
+  proc->pd->refs = 1;
+  proc->pd->cr3_paddr = mm_alloc_user_pd_phys ();
+  if (proc->pd->cr3_paddr == 0) {
+    free (proc->pd);
     free (proc);
     return NULL;
   }
@@ -39,6 +47,8 @@ struct proc* proc_from_elf (uint8_t* elf_contents) {
   struct proc_resource* kstk_r =
       proc_create_resource (proc, kstk_rid, PR_MEM, RV_PRIVATE, (void*)&kstk_mem_init);
   if (kstk_r == NULL) {
+    pmm_free (proc->pd->cr3_paddr, 1);
+    free (proc->pd);
     free (proc);
     return NULL;
   }
@@ -52,6 +62,8 @@ struct proc* proc_from_elf (uint8_t* elf_contents) {
   if (ustk_r == NULL) {
     kstk_r->ops.cleanup (proc, kstk_r);
     free (kstk_r);
+    pmm_free (proc->pd->cr3_paddr, 1);
+    free (proc->pd);
     free (proc);
     return NULL;
   }
@@ -78,24 +90,28 @@ void proc_cleanup (struct proc* proc) {
 
   proc_cleanup_resources (proc);
 
-  struct list_node_link *mapping_link, *mapping_link_tmp;
-  spin_lock (&proc->pd.lock, &ctxprpd);
+  if (atomic_fetch_sub (&proc->pd->refs, 1) == 1) {
+    DEBUG ("PID %d Free virtual address space\n", proc->pid);
+    struct list_node_link *mapping_link, *mapping_link_tmp;
+    spin_lock (&proc->pd->lock, &ctxprpd);
 
-  list_foreach (proc->mappings, mapping_link, mapping_link_tmp) {
-    struct proc_mapping* mapping =
-        list_entry (mapping_link, struct proc_mapping, proc_mappings_link);
+    list_foreach (proc->mappings, mapping_link, mapping_link_tmp) {
+      struct proc_mapping* mapping =
+          list_entry (mapping_link, struct proc_mapping, proc_mappings_link);
 
-    list_remove (proc->mappings, mapping_link);
-    free (mapping);
+      list_remove (proc->mappings, mapping_link);
+      free (mapping);
+    }
+
+    pmm_free (proc->pd->cr3_paddr, 1);
+    spin_unlock (&proc->pd->lock, &ctxprpd);
+    free (proc->pd);
   }
 
-  spin_unlock (&proc->pd.lock, &ctxprpd);
-
-  pmm_free (proc->pd.cr3_paddr, 1);
-
   pmm_free (proc->pdata.kernel_stack - (uintptr_t)hhdm->offset - KSTACK_SIZE,
             KSTACK_SIZE / PAGE_SIZE);
   pmm_free (proc->pdata.user_stack, USTACK_SIZE / PAGE_SIZE);
+  DEBUG ("PID %d Free stacks\n", proc->pid);
 
   free (proc);
 }
diff --git a/kernel/amd64/sched1.c b/kernel/amd64/sched1.c
index a55da0f..6290b87 100644
--- a/kernel/amd64/sched1.c
+++ b/kernel/amd64/sched1.c
@@ -12,5 +12,5 @@ void do_sched (struct proc* proc) {
   thiscpu->tss.rsp0 = proc->pdata.kernel_stack;
   thiscpu->syscall_kernel_stack = proc->pdata.kernel_stack;
 
-  amd64_do_sched ((void*)&proc->pdata.regs, (void*)proc->pd.cr3_paddr);
+  amd64_do_sched ((void*)&proc->pdata.regs, (void*)proc->pd->cr3_paddr);
 }
diff --git a/kernel/proc/proc.c b/kernel/proc/proc.c
index 8accd58..62d9839 100644
--- a/kernel/proc/proc.c
+++ b/kernel/proc/proc.c
@@ -60,16 +60,16 @@ bool proc_map (struct proc* proc, uintptr_t start_paddr, uintptr_t start_vaddr,
 
   flags &= ~(MM_PD_LOCK | MM_PD_RELOAD); /* clear LOCK flag if present, because we lock manualy */
 
-  spin_lock (&proc->pd.lock, &ctxprpd);
+  spin_lock (&proc->pd->lock, &ctxprpd);
 
   list_append (proc->mappings, &mapping->proc_mappings_link);
 
   for (uintptr_t vpage = start_vaddr, ppage = start_paddr; vpage < start_vaddr + pages * PAGE_SIZE;
        vpage += PAGE_SIZE, ppage += PAGE_SIZE) {
-    mm_map_page (&proc->pd, ppage, vpage, flags);
+    mm_map_page (proc->pd, ppage, vpage, flags);
   }
 
-  spin_unlock (&proc->pd.lock, &ctxprpd);
+  spin_unlock (&proc->pd->lock, &ctxprpd);
 
   return true;
 }
@@ -85,7 +85,7 @@ bool proc_unmap (struct proc* proc, uintptr_t start_vaddr, size_t pages) {
   if (tail_mapping == NULL)
     return false;
 
-  spin_lock (&proc->pd.lock, &ctxprpd);
+  spin_lock (&proc->pd->lock, &ctxprpd);
 
   list_foreach (proc->mappings, mapping_link, mapping_link_tmp) {
     struct proc_mapping* mapping =
@@ -127,10 +127,10 @@ bool proc_unmap (struct proc* proc, uintptr_t start_vaddr, size_t pages) {
     free (tail_mapping);
 
   for (uintptr_t vpage = start_vaddr; vpage < end_vaddr; vpage += PAGE_SIZE) {
-    mm_unmap_page (&proc->pd, vpage, 0);
+    mm_unmap_page (proc->pd, vpage, 0);
   }
 
-  spin_unlock (&proc->pd.lock, &ctxprpd);
+  spin_unlock (&proc->pd->lock, &ctxprpd);
 
   return true;
 }
diff --git a/kernel/proc/proc.h b/kernel/proc/proc.h
index a451e37..e2b2544 100644
--- a/kernel/proc/proc.h
+++ b/kernel/proc/proc.h
@@ -44,7 +44,7 @@ struct proc {
 
   struct list_node_link* mappings; /* pd.lock implicitly protects this field */
   struct proc_platformdata pdata;
-  struct pd pd;
+  struct pd* pd;
   spin_lock_t lock;
   struct cpu* cpu;
   atomic_int state;
diff --git a/kernel/syscall/syscall.c b/kernel/syscall/syscall.c
index d84803b..3be5ac5 100644
--- a/kernel/syscall/syscall.c
+++ b/kernel/syscall/syscall.c
@@ -66,16 +66,16 @@ DEFINE_SYSCALL (sys_proc_create_resource_mem) {
   int vis = (int)a2;
   uintptr_t* out_paddr_buf = (uintptr_t*)a3;
 
-  spin_lock (&proc->pd.lock, &ctxprpd);
+  spin_lock (&proc->pd->lock, &ctxprpd);
 
-  uintptr_t out_paddr_buf_paddr = mm_v2p (&proc->pd, (uintptr_t)out_paddr_buf, 0);
+  uintptr_t out_paddr_buf_paddr = mm_v2p (proc->pd, (uintptr_t)out_paddr_buf, 0);
 
-  if (!mm_validate_buffer (&proc->pd, (uintptr_t)out_paddr_buf, sizeof (uintptr_t), 0)) {
-    spin_unlock (&proc->pd.lock, &ctxprpd);
+  if (!mm_validate_buffer (proc->pd, (uintptr_t)out_paddr_buf, sizeof (uintptr_t), 0)) {
+    spin_unlock (&proc->pd->lock, &ctxprpd);
     return -SR_BAD_ADDRESS_SPACE;
   }
 
-  spin_unlock (&proc->pd.lock, &ctxprpd);
+  spin_unlock (&proc->pd->lock, &ctxprpd);
 
   uintptr_t* out_paddr_buf_vaddr = (uintptr_t*)((uintptr_t)hhdm->offset + out_paddr_buf_paddr);