improvements to shared memory handling

+ Now each shared region is reference counter (technically each region
  is refernce counted, private regions just have a count of 1).
  Also, syscalls that touch the TLB get flushed, but should probably
  look into where else this flushing might be needed.

12 files changed, 171 insertions, 13 deletions

diff --git a/include/kmi/assert.h b/include/kmi/assert.h
index ecf5272..14e0679 100644
--- a/include/kmi/assert.h
+++ b/include/kmi/assert.h
@@ -18,6 +18,7 @@
  * Catastrophic assertions warn about the assertion and crash the kernel.
  */
 
+#include <kmi/power.h>
 #include <kmi/debug.h>
 #include <kmi/utils.h>
 
@@ -35,6 +36,7 @@
 		if (unlikely(!(x))) {                        \
 			error("assertion failed: " #x "\n"); \
 			while (1) {                          \
+				poweroff(SYS_COLD_REBOOT);   \
 			}                                    \
 		}                                            \
 	} while (0);
diff --git a/include/kmi/regions.h b/include/kmi/regions.h
index ebaa07a..5167d28 100644
--- a/include/kmi/regions.h
+++ b/include/kmi/regions.h
@@ -78,6 +78,8 @@ struct mem_region_root {
  * track of free and used regions, respectively. All regions are chained
  * together with a doubly linked list, so that the next region's start address
  * should be the current region's end address.
+ *
+ * Note that addresses are expressed in pages!
  */
 struct mem_region {
 	/** Sp-tree node slot. */
@@ -98,6 +100,11 @@ struct mem_region {
 	/** In shared regions, mark the other pid that shared the region. */
 	id_t pid;
 
+	union {
+		vm_t shaddr;
+		size_t refcount;
+	};
+
 	/** Memory region flags, both access as well as metadata. \see MR_USED,
 	 * MR_SHARED, MR_OWNED, MR_COW, MR_KEEP. */
 	vmflags_t flags;
diff --git a/src/regions.c b/src/regions.c
index d5f431d..2db1a91 100644
--- a/src/regions.c
+++ b/src/regions.c
@@ -440,6 +440,9 @@ static vm_t __partition_region(struct mem_region_root *r, struct mem_region *m,
 	m->start = start;
 	m->flags = flags;
 	m->pid = pid;
+	if (m->pid == 0)
+		m->refcount = 1;
+
 	mark_region_used(m->flags);
 	__insert_used_region(r, m);
 	return __addr(start);
diff --git a/src/tcb.c b/src/tcb.c
index 3e4c75a..70cad5a 100644
--- a/src/tcb.c
+++ b/src/tcb.c
@@ -199,7 +199,7 @@ struct tcb *create_thread(struct tcb *p)
 
 	stat_t r = OK;
 	if (p) r = __init_owned_thread(p, t);
-	else   r = __init_free_thread(t);
+	else r = __init_free_thread(t);
 
 	if (r) {
 		free_page(MM_O0, bottom);
diff --git a/src/uapi/mem.c b/src/uapi/mem.c
index f7237d7..5c18797 100644
--- a/src/uapi/mem.c
+++ b/src/uapi/mem.c
@@ -29,6 +29,7 @@ SYSCALL_DEFINE2(req_mem)(struct tcb *t, sys_arg_t size, sys_arg_t flags)
 	if (ERR_CODE(start))
 		return_args1(t, start);
 
+	flush_tlb_all();
 	return_args2(t, OK, start);
 }
 
@@ -52,6 +53,7 @@ SYSCALL_DEFINE3(req_fixmem)(struct tcb *t, sys_arg_t fixed, sys_arg_t size,
 	if (ERR_CODE(start))
 		return_args1(t, start);
 
+	flush_tlb_all();
 	return_args2(t, OK, start);
 }
 
@@ -75,6 +77,7 @@ SYSCALL_DEFINE1(free_mem)(struct tcb *t, sys_arg_t start)
 	if (!(status = free_devmem(r, start)))
 		return_args1(t, OK);
 
+	flush_tlb_all();
 	return_args1(t, status);
 }
 
@@ -102,6 +105,7 @@ SYSCALL_DEFINE3(req_pmem)(struct tcb *t, sys_arg_t paddr, sys_arg_t size,
 	if (ERR_CODE(start))
 		return_args1(t, start);
 
+	flush_tlb_all();
 	return_args2(t, OK, start);
 }
 
@@ -131,6 +135,7 @@ SYSCALL_DEFINE2(req_page)(struct tcb *t, sys_arg_t size, sys_arg_t flags)
 	if (ERR_CODE(start))
 		return_args1(t, start);
 
+	flush_tlb_all();
 	return_args4(t, OK, start, addr, asize);
 }
 
@@ -154,6 +159,7 @@ SYSCALL_DEFINE2(req_sharedmem)(struct tcb *t, sys_arg_t size, sys_arg_t flags)
 	if (ERR_CODE(start))
 		return_args1(t, start);
 
+	flush_tlb_all();
 	return_args2(t, OK, start);
 }
 
@@ -184,5 +190,6 @@ SYSCALL_DEFINE3(ref_sharedmem)(struct tcb *t, sys_arg_t tid, sys_arg_t addr,
 	if (ERR_CODE(start))
 		return_args1(t, start);
 
+	flush_tlb_all();
 	return_args2(t, OK, start);
 }
diff --git a/src/vmem.c b/src/vmem.c
index 2feff45..f4f1bf7 100644
--- a/src/vmem.c
+++ b/src/vmem.c
@@ -63,6 +63,37 @@ static stat_t __copy_mapped_region(struct tcb *d, struct tcb *s,
 	return res;
 }
 
+static void reference_mem(struct tcb *d, struct tcb *s, vm_t ref, vm_t orig)
+{
+	struct mem_region *src = find_used_region(&s->uvmem.region, orig);
+	assert(src);
+
+	struct mem_region *dst = find_used_region(&d->uvmem.region, ref);
+	assert(dst);
+
+	assert(dst->pid == s->rid);
+	dst->shaddr = orig;
+	src->refcount++;
+
+	reference_thread(s);
+}
+
+static void __free_mapping(struct tcb *t, struct mem_region *m);
+
+static void unreference_mem(struct tcb *s, vm_t addr)
+{
+	struct mem_region *src = find_used_region(&s->uvmem.region, addr);
+	assert(src);
+
+	assert(src->refcount >= 1);
+	if (--src->refcount == 0) {
+		__free_mapping(s, src);
+		free_known_region(&s->uvmem.region, src);
+	}
+
+	unreference_thread(s);
+}
+
 /**
  * Copy shared regions to new process. In these cases, we want to both allocate
  * a fixed region and map some fixed physical memory.
@@ -80,7 +111,6 @@ static stat_t __copy_shared_region(struct tcb *d, struct mem_region *m)
 	vm_t start = m->start * BASE_PAGE_SIZE;
 	vm_t end = m->end * BASE_PAGE_SIZE;
 
-	reference_thread(s);
 
 	size_t size = end - start;
 	vm_t v = alloc_shared_fixed_region(&d->uvmem.region, start, size, &size,
@@ -91,8 +121,10 @@ static stat_t __copy_shared_region(struct tcb *d, struct mem_region *m)
 	assert(v == start);
 	stat_t res = clone_region(d->uvmem.vmem, s->uvmem.vmem, start, v, size,
 	                          m->flags);
-	if (res == OK)
+	if (res == OK) {
+		reference_mem(d, s, v, v);
 		return OK;
+	}
 
 	/* cleanup on error */
 	free_region(&d->uvmem.region, v);
@@ -115,8 +147,6 @@ static vm_t __clone_shared_region(struct tcb *d, struct tcb *s,
 	vm_t start = m->start * BASE_PAGE_SIZE;
 	vm_t end = m->end * BASE_PAGE_SIZE;
 
-	reference_thread(s);
-
 	size_t size = end - start;
 	vm_t v = alloc_shared_region(&d->uvmem.region, size, &size,
 	                             MR_NONBACKED | m->flags, s->rid);
@@ -125,18 +155,19 @@ static vm_t __clone_shared_region(struct tcb *d, struct tcb *s,
 
 	stat_t res = clone_region(d->uvmem.vmem, s->uvmem.vmem,
 	                          start, v, size, flags);
-	if (res == OK)
+	if (res == OK) {
+		reference_mem(d, s, v, start);
 		return v;
+	}
 
 	/* cleanup on error */
-	unreference_thread(s);
 	free_region(&d->uvmem.region, v);
 	unmap_fixed_region(d->uvmem.vmem, v, size);
 	return res;
 }
 
 /**
- * Unmap and free private memory region.
+ * Unmap and free memory region.
  *
  * @param t Current thread.
  * @param m Memory region to free.
@@ -145,7 +176,7 @@ static void __free_mapping(struct tcb *t, struct mem_region *m)
 {
 	struct tcb *owner = get_tcb(m->pid);
 	if (owner)
-		unreference_thread(owner);
+		unreference_mem(owner, m->shaddr);
 
 	if (is_set(m->flags, MR_NONBACKED))
 		return;
@@ -173,6 +204,17 @@ void clear_uvmem(struct tcb *t)
 		if (!is_set(m->flags, MR_USED))
 			continue;
 
+		if (m->pid == 0 && m->refcount > 1) {
+			/* we own this shared region and it's used by someone
+			 * else so we can't outright free it yet, but remove our
+			 * 'reference' to it so the refcount reaches zero when all
+			 * referees unmap it, instead of waiting for this thread to be
+			 * completely destroyed.
+			 * */
+			m->refcount--;
+			continue;
+		}
+
 		__free_mapping(t, m);
 		free_known_region(&t->uvmem.region, m);
 	}
@@ -272,7 +314,7 @@ vm_t map_shared_fixed_uvmem(struct tcb *t, pm_t start, size_t size,
 	assert(is_aligned(start, BASE_PAGE_SIZE));
 
 	const vm_t v = alloc_shared_region(&t->uvmem.region, size, &size, flags,
-	                                   get_rproc(t)->tid);
+	                                   0);
 	if (ERR_CODE(v))
 		return v;
 
@@ -354,6 +396,9 @@ stat_t free_uvmem(struct tcb *r, vm_t va)
 	if (!m)
 		return ERR_NF;
 
+	if (m->pid == 0 && m->refcount > 1)
+		return ERR_INVAL;
+
 	__free_mapping(r, m);
 	free_known_region(&r->uvmem.region, m);
 	return OK;
diff --git a/tests/common/arch/riscv64/source.mk b/tests/common/arch/riscv64/source.mk
index 517c1c7..e3c0e55 100644
--- a/tests/common/arch/riscv64/source.mk
+++ b/tests/common/arch/riscv64/source.mk
@@ -1 +1,4 @@
-ARCH_FLAGS += -march=rv64imac -mabi=lp64
+# -mno-relax to disable the compiler from using the gp register as a shorthand
+#  for __global_pointer$, which I currently don't do anything with. Should
+#  probably fix at some point
+ARCH_FLAGS += -march=rv64imac -mabi=lp64 -mno-relax
diff --git a/tests/common/sys.h b/tests/common/sys.h
index 6d5e493..3e98abf 100644
--- a/tests/common/sys.h
+++ b/tests/common/sys.h
@@ -86,9 +86,10 @@ static inline void *sys_ref_sharedmem(id_t tid, uintptr_t addr, vmflags_t flags)
 	return (void *)r.a0;
 }
 
-static inline void sys_free_mem(uintptr_t start)
+static inline enum sys_status sys_free_mem(uintptr_t start)
 {
-	syscall1(SYS_FREE_MEM, start);
+	struct sys_ret r = syscall1(SYS_FREE_MEM, start);
+	return r.s;
 }
 
 static inline uint64_t sys_timebase()
diff --git a/tests/dangling_shmem/init.c b/tests/dangling_shmem/init.c
new file mode 100644
index 0000000..df211a7
--- /dev/null
+++ b/tests/dangling_shmem/init.c
@@ -0,0 +1,50 @@
+#include <common/test.h>
+
+static id_t new_id = 0;
+static char *refmem = NULL;
+
+START(pid, tid, d0, d1, d2, d3)
+{
+	UNUSED(pid);
+	UNUSED(tid);
+	UNUSED(d0);
+	UNUSED(d1);
+	UNUSED(d2);
+	UNUSED(d3);
+
+	if (pid == 0 && d0 == SYS_USER_SPAWNED) {
+		printf("forking\n");
+		id_t our_id = 0;
+		new_id = sys_fork(&our_id);
+		if (new_id == 0) {
+			printf("in child, setting up shared memory\n");
+			void *p = sys_req_sharedmem(1, VM_R | VM_W);
+			check(p, "failed getting shared memory\n");
+
+			printf("referencing shared memory for pid 1\n");
+			void *r = sys_ref_sharedmem(1, (uintptr_t)p, VM_R | VM_W);
+			check(r, "failed referencing shared memory\n");
+
+			sys_ipc_req1(1, (uintptr_t)r);
+			sys_exit(1);
+		}
+
+		printf("swapping to child\n");
+		enum sys_status r = sys_swap(new_id);
+		check(r == OK, "failed swapping to child\n");
+
+		printf("returned from child\n");
+		/* child should now be dead */
+		check(refmem[0] == 'p', "child dying affected our shared memory?\n");
+
+		printf("freeing shared memory\n");
+		r = sys_free_mem((uintptr_t)refmem);
+		check(r == OK, "failed freeing referenced memory\n");
+	}
+	else if (pid == new_id) {
+		refmem = (char *)d0;
+		refmem[0] = 'p';
+		sys_ipc_resp0();
+	}
+	ok();
+}
diff --git a/tests/dangling_shmem/source.mk b/tests/dangling_shmem/source.mk
new file mode 100644
index 0000000..2ecd7a6
--- /dev/null
+++ b/tests/dangling_shmem/source.mk
@@ -0,0 +1,2 @@
+DO	!= ./scripts/gen-prog -n dangling_shmem -p init init.c
+DO	!= ./scripts/gen-simple -n dangling_shmem -p init
diff --git a/tests/shmem/init.c b/tests/shmem/init.c
new file mode 100644
index 0000000..13e1667
--- /dev/null
+++ b/tests/shmem/init.c
@@ -0,0 +1,36 @@
+#include <common/test.h>
+
+START(pid, tid, d0, d1, d2, d3)
+{
+	UNUSED(pid);
+	UNUSED(tid);
+	UNUSED(d0);
+	UNUSED(d1);
+	UNUSED(d2);
+	UNUSED(d3);
+
+	printf("allocating shared memory\n");
+	volatile char *shmem = sys_req_sharedmem(1, VM_R | VM_W);
+	check(shmem, "no shared memory?\n");
+
+	*shmem = 'p';
+
+	printf("sharing memory to ourselves\n");
+	volatile char *refmem = sys_ref_sharedmem(1, (uintptr_t)shmem, VM_R | VM_W);
+	check(refmem[0] == 'p', "wrong mapping?\n");
+
+	printf("freeing shared memory (should fail)\n");
+	enum sys_status r = sys_free_mem((uintptr_t)shmem);
+	check(r != OK, "illegal freeing succeeded?\n");
+
+	printf("freeing referenced memory\n");
+	r = sys_free_mem((uintptr_t)refmem);
+	check(r == OK, "legal ref freeing failed?\n");
+	check(shmem[0] == 'p', "freeing ref freed owned memory?\n");
+
+	printf("freeing recently nonshared memory\n");
+	r = sys_free_mem((uintptr_t)shmem);
+	check(r == OK, "legal freeing failed?\n");
+
+	ok();
+}
diff --git a/tests/shmem/source.mk b/tests/shmem/source.mk
new file mode 100644
index 0000000..da84808
--- /dev/null
+++ b/tests/shmem/source.mk
@@ -0,0 +1,2 @@
+DO	!= ./scripts/gen-prog -n shmem -p init init.c
+DO	!= ./scripts/gen-simple -n shmem -p init