From ccc13ef0d5aab327d1b3193b031fa12ceab1779e Mon Sep 17 00:00:00 2001
From: Mark <alnyan@airmail.cc>
Date: Sun, 29 Dec 2019 14:15:09 +0200
Subject: [PATCH] Add thread cleanup after termination

---
 include/sys/amd64/sys/thread.h |  2 ++
 include/sys/thread.h           |  1 +
 sys/amd64/mm/map.c             |  7 ++++++-
 sys/amd64/sys/sched.c          |  2 ++
 sys/amd64/sys/thread.c         | 30 ++++++++++++++++++++++++++++++
 sys/vfs/node.c                 |  7 +++++++
 sys/vfs/vfs.c                  | 20 ++++++++++++++++----
 7 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/include/sys/amd64/sys/thread.h b/include/sys/amd64/sys/thread.h
index c9b5da0..7bd98a7 100644
--- a/include/sys/amd64/sys/thread.h
+++ b/include/sys/amd64/sys/thread.h
@@ -11,6 +11,8 @@ struct amd64_thread {
 
     uintptr_t stack3_base;
     uintptr_t stack3_size;
+
+    uint32_t data_flags;
 };
 
 void amd64_thread_set_ip(struct thread *t, uintptr_t ip);
diff --git a/include/sys/thread.h b/include/sys/thread.h
index 7ac7b68..c388e02 100644
--- a/include/sys/thread.h
+++ b/include/sys/thread.h
@@ -45,3 +45,4 @@ int thread_init(struct thread *t,
                 size_t stack3_size,
                 uint32_t flags,
                 void *arg);
+void thread_cleanup(struct thread *t);
diff --git a/sys/amd64/mm/map.c b/sys/amd64/mm/map.c
index 34058c1..3c8824c 100644
--- a/sys/amd64/mm/map.c
+++ b/sys/amd64/mm/map.c
@@ -282,7 +282,7 @@ int mm_space_fork(mm_space_t dst_pml4, const mm_space_t src_pml4, uint32_t flags
 }
 
 void mm_space_release(mm_space_t pml4) {
-    for (size_t pml4i = 0; pml4i < 255; ++pml4i) {
+    for (size_t pml4i = 0; pml4i < AMD64_PML4I_USER_END; ++pml4i) {
         if (!(pml4[pml4i] & 1)) {
             continue;
         }
@@ -324,6 +324,11 @@ void mm_space_release(mm_space_t pml4) {
     }
 }
 
+void mm_space_free(mm_space_t pml4) {
+    mm_space_release(pml4);
+    amd64_mm_pool_free(pml4);
+}
+
 static void amd64_mm_describe_range(const mm_space_t pml4, uintptr_t start_addr, uintptr_t end_addr) {
     uintptr_t addr = AMD64_MM_STRIPSX(start_addr);
 
diff --git a/sys/amd64/sys/sched.c b/sys/amd64/sys/sched.c
index a1d27c1..4dc2321 100644
--- a/sys/amd64/sys/sched.c
+++ b/sys/amd64/sys/sched.c
@@ -230,6 +230,8 @@ void sched_remove_from(int cpu, struct thread *thr) {
     thr->next = NULL;
     thr->prev = NULL;
 
+    thread_cleanup(thr);
+
     // TODO: make other cpus schedule a new thread if they're running a stopped one
 }
 
diff --git a/sys/amd64/sys/thread.c b/sys/amd64/sys/thread.c
index b27148b..7f31b0e 100644
--- a/sys/amd64/sys/thread.c
+++ b/sys/amd64/sys/thread.c
@@ -28,12 +28,16 @@ int thread_init(
         size_t rsp3_size,
         uint32_t flags,
         void *arg) {
+    t->data.data_flags = 0;
+
     if (!rsp0_base) {
         void *kstack = kmalloc(THREAD_KSTACK_SIZE);
         _assert(kstack);
 
         rsp0_base = (uintptr_t) kstack;
         rsp0_size = THREAD_KSTACK_SIZE;
+    } else {
+        t->data.data_flags |= 1 /* Don't free kstack */;
     }
 
     if (!(flags & THREAD_KERNEL) && !rsp3_base) {
@@ -44,6 +48,8 @@ int thread_init(
         rsp3_size = 4 * 0x1000;
 
         kdebug("Allocated user stack: %p\n", ustack);
+    } else if (rsp3_base) {
+        t->data.data_flags |= 2 /* Don't free ustack */;
     }
 
     t->data.stack0_base = rsp0_base;
@@ -112,6 +118,29 @@ int thread_init(
     return 0;
 }
 
+void thread_cleanup(struct thread *t) {
+    _assert(t);
+
+    // Release files
+    for (size_t i = 0; i < 4; ++i) {
+        if (t->fds[i].vnode) {
+            vfs_close(&t->ioctx, &t->fds[i]);
+        }
+    }
+
+    // Release stacks
+    if (!(t->data.data_flags & 1)) {
+        kfree((void *) t->data.stack0_base);
+    }
+
+    if (t->space != mm_kernel) {
+        mm_space_free(t->space);
+    }
+
+    memset(t, 0, sizeof(struct thread));
+    kfree(t);
+}
+
 int sys_execve(const char *filename, const char *const argv[], const char *const envp[]) {
     asm volatile ("cli");
     struct thread *thr = get_cpu()->thread;
@@ -211,6 +240,7 @@ int sys_fork(void) {
     struct amd64_thread *dst_data = &thr_dst->data;
     _assert(dst_kstack);
 
+    dst_data->data_flags |= 1 /* Don't free kstack */;
     dst_data->stack0_base = (uintptr_t) dst_kstack;
     dst_data->stack0_size = THREAD_KSTACK_SIZE;
     dst_data->stack3_base = thr_src->data.stack3_base;
diff --git a/sys/vfs/node.c b/sys/vfs/node.c
index 84a2029..817990e 100644
--- a/sys/vfs/node.c
+++ b/sys/vfs/node.c
@@ -67,6 +67,9 @@ void vnode_ref(vnode_t *vn) {
 //    vfs_vnode_path(buf, vn);
 //    printf("++refcount for %s\n", buf);
     _assert(vn);
+    if (vn->type == VN_BLK || vn->type == VN_CHR) {
+        return;
+    }
     _assert(vn->fs);
     _assert(vn->fs->cls);
 
@@ -86,6 +89,10 @@ void vnode_ref(vnode_t *vn) {
 
 void vnode_unref(vnode_t *vn) {
     _assert(vn);
+    if (vn->type == VN_BLK || vn->type == VN_CHR) {
+        return;
+    }
+
     _assert(vn->fs);
     _assert(vn->fs->cls);
 
diff --git a/sys/vfs/vfs.c b/sys/vfs/vfs.c
index bb8fe1d..fa44d17 100644
--- a/sys/vfs/vfs.c
+++ b/sys/vfs/vfs.c
@@ -868,12 +868,24 @@ int vfs_open_node(struct vfs_ioctx *ctx, struct ofile *of, vnode_t *vn, int opt)
 
 void vfs_close(struct vfs_ioctx *ctx, struct ofile *of) {
     _assert(of);
-    kinfo("vfs_close %p\n", of);
     vnode_t *vn = of->vnode;
-    _assert(vn && vn->op);
+    switch (vn->type) {
+    case VN_REG:
+    case VN_DIR:
+        _assert(vn);
 
-    if (vn->op->close) {
-        vn->op->close(of);
+        if (vn->op->close) {
+            vn->op->close(of);
+        }
+        break;
+
+    // TODO: maybe call something special for that
+    case VN_CHR:
+    case VN_BLK:
+        break;
+
+    default:
+        panic("Unhandled vnode type\n");
     }
 
     vnode_unref(of->vnode);