Use address space reference counting to fix leak

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;