Use address space reference counting to fix leak

Previously, thread_destroy() never released address space references,
causing AS structures to leak when threads were destroyed. The shared
counter was incremented when threads shared an AS but never decremented.

Address space is now freed when the last thread holding a reference is
destroyed, rather than leaking indefinitely.

Author: jserv

include/memory.h

@@ -23,9 +23,20 @@ typedef struct {
 
 	struct fpage *mpu_first;	/*! head of MPU fpage list */
 	struct fpage *mpu_stack_first;	/*! head of MPU stack fpage list */
-	uint32_t shared;	/*! shared user number */
+	uint32_t shared;	/*! Reference count: number of threads using this AS */
 } as_t;
 
+/*
+ * Address space reference counting.
+ * as_get(): Increment reference count (thread sharing this AS)
+ * as_put(): Decrement reference count; frees AS when it reaches 0
+ *
+ * Thread must be unlinked from AS before calling as_put().
+ * IRQ protection ensures atomicity of refcount operations.
+ */
+void as_get(as_t *as);
+void as_put(as_t *as);
+
 /**
  * Memory pool represents area of physical address space
  * We set flags to it (kernel & user permissions),

include/platform/irq.h

@@ -25,6 +25,28 @@ static inline void irq_enable(void)
 	__asm__ __volatile__ ("cpsie i" ::: "memory");
 }
 
+/*
+ * Save and restore interrupt state (PRIMASK).
+ * Used for critical sections that may be nested or called from
+ * contexts where IRQ state is unknown.
+ */
+static inline uint32_t irq_save_flags(void)
+{
+	uint32_t flags;
+	__asm__ __volatile__ (
+		"mrs %0, primask\n\t"
+		"cpsid i"
+		: "=r" (flags) :: "memory");
+	return flags;
+}
+
+static inline void irq_restore_flags(uint32_t flags)
+{
+	__asm__ __volatile__ (
+		"msr primask, %0"
+		:: "r" (flags) : "memory");
+}
+
 static inline void irq_svc(void)
 {
 	__asm__ __volatile__ ("svc #0");

include/thread.h

@@ -133,7 +133,6 @@ tcb_t *thread_init(l4_thread_t globalid, utcb_t *utcb);
 tcb_t *thread_create(l4_thread_t globalid, utcb_t *utcb);
 void thread_destroy(tcb_t *thr);
 void thread_space(tcb_t *thr, l4_thread_t spaceid, utcb_t *utcb);
-void thread_free_space(tcb_t *thr);
 void thread_init_ctx(void *sp, void *pc, void *rx, tcb_t *thr);
 void thread_init_kernel_ctx(void *sp, tcb_t *thr);
 void thread_switch(tcb_t *thr);

kernel/memory.c

@@ -12,6 +12,7 @@
 #include <init_hook.h>
 #include <kip.h>
 #include <platform/cortex_m.h>
+#include <platform/irq.h>
 
 /**
  * @file    memory.c
@@ -378,31 +379,79 @@ as_t *as_create(uint32_t as_spaceid)
 
 	as->as_spaceid = as_spaceid;
 	as->first = NULL;
-	as->shared = 0;
+	as->mpu_first = NULL;
+	as->mpu_stack_first = NULL;
+	as->shared = 1;  /* Creator holds initial reference */
 
 	return as;
 }
 
-void as_destroy(as_t *as)
+/*
+ * Increment address space reference count.
+ * Called when a thread shares an existing address space.
+ */
+void as_get(as_t *as)
 {
-	fpage_t *fp, *fpnext;
-	fp = as->first;
+	uint32_t flags;
 
-	if (as->shared > 0) {
-		--as->shared;
+	if (!as)
 		return;
+
+	flags = irq_save_flags();
+	as->shared++;
+	irq_restore_flags(flags);
+}
+
+/*
+ * Decrement address space reference count.
+ * Frees the address space when refcount reaches 0.
+ * IRQs remain disabled through as_destroy() to prevent reentrancy.
+ */
+void as_put(as_t *as)
+{
+	uint32_t flags;
+
+	if (!as)
+		return;
+
+	flags = irq_save_flags();
+
+	/* Guard against refcount underflow (double-put bug) */
+	if (as->shared == 0) {
+		irq_restore_flags(flags);
+		dbg_printf(DL_KDB, "AS: refcount underflow for %p\n",
+		           as->as_spaceid);
+		return;
+	}
+
+	if (--as->shared == 0) {
+		/* Keep IRQs disabled through destroy */
+		as_destroy(as);
 	}
+	irq_restore_flags(flags);
+}
+
+/*
+ * Free all fpages and the address space structure.
+ * Called by as_put() when refcount reaches 0.
+ * Caller must hold IRQs disabled.
+ */
+void as_destroy(as_t *as)
+{
+	fpage_t *fp, *fpnext;
+	fp = as->first;
 
 	/*
 	 * FIXME: What if a CLONED fpage which is MAPPED is to be deleted
 	 */
 	while (fp) {
+		fpnext = fp->as_next;
+
 		if (fp->fpage.flags & FPAGE_CLONE)
 			unmap_fpage(as, fp);
 		else
 			destroy_fpage(fp);
 
-		fpnext = fp->as_next;
 		fp = fpnext;
 	}
 

kernel/syscall.c

@@ -81,7 +81,7 @@ static void sys_thread_control(uint32_t *param1, uint32_t *param2)
 			param1[REG_R0] = 0;
 			return;
 		}
-		thread_free_space(thr);
+		/* thread_destroy() handles AS refcount via as_put() */
 		thread_destroy(thr);
 		param1[REG_R0] = 1;
 	}

kernel/thread.c

@@ -277,6 +277,12 @@ void thread_destroy(tcb_t *thr)
 		caller->t_child = thr->t_child;
 	}
 
+	/* Release address space reference.
+	 * This may free the AS if this was the last reference.
+	 */
+	if (thr->as)
+		as_put(thr->as);
+
 	thread_deinit(thr);
 }
 
@@ -310,10 +316,23 @@ void thread_space(tcb_t *thr, l4_thread_t spaceid, utcb_t *utcb)
 		dbg_printf(DL_THREAD,
 		           "\tNew space: as: %p, utcb: %p \n", thr->as, utcb);
 	} else {
-		tcb_t *space = thread_by_globalid(spaceid);
+		tcb_t *space;
+		as_t *shared_as;
 
-		thr->as = space->as;
-		++(space->as->shared);
+		irq_disable();
+		space = thread_by_globalid(spaceid);
+		if (!space || !space->as) {
+			irq_enable();
+			dbg_printf(DL_KDB,
+			           "THREAD: space thread %t not found\n",
+			           spaceid);
+			return;
+		}
+		shared_as = space->as;
+		as_get(shared_as);
+		irq_enable();
+
+		thr->as = shared_as;
 	}
 
 	/* If no caller, than it is mapping from kernel to root thread
@@ -333,12 +352,6 @@ void thread_space(tcb_t *thr, l4_thread_t spaceid, utcb_t *utcb)
 		           utcb, thr->t_globalid);
 }
 
-void thread_free_space(tcb_t *thr)
-{
-	/* free address space */
-	as_destroy(thr->as);
-}
-
 void thread_init_ctx(void *sp, void *pc, void *regs, tcb_t *thr)
 {
 	int i;