sched: enforce ref-counting on current task/thread pointers

arch/x86_64/panic.c

@@ -217,16 +217,18 @@ static void print_stack_trace(
 
 void ml_print_stack_trace(uintptr_t ip)
 {
-	struct task *task = current_task();
+	struct task *task = get_current_task();
 	struct address_space *space = task ? task->t_address_space : NULL;
 	uintptr_t bp;
 	asm volatile("mov %%rbp, %0" : "=r"(bp));
 	print_stack_trace(space, ip, bp);
+	put_current_task(task);
 }
 
 void ml_print_stack_trace_irq(struct ml_cpu_context *ctx)
 {
-	struct task *task = current_task();
+	struct task *task = get_current_task();
 	struct address_space *space = task ? task->t_address_space : NULL;
 	print_stack_trace(space, ctx->rip, ctx->rbp);
+	put_current_task(task);
 }

arch/x86_64/pmap.c

@@ -496,11 +496,7 @@ kern_status_t pmap_handle_fault(
 {
 	// log_fault(fault_addr, flags);
 
-	if (flags & PMAP_FAULT_PRESENT) {
-		return KERN_FATAL_ERROR;
-	}
-
-	struct task *task = current_task();
+	struct task *task = get_current_task();
 	if (!task) {
 		return KERN_FATAL_ERROR;
 	}
@@ -511,7 +507,11 @@ kern_status_t pmap_handle_fault(
 	}
 
 	/* this must be called with `space` unlocked. */
-	return address_space_demand_map(space, fault_addr, flags);
+	kern_status_t status
+		= address_space_demand_map(space, fault_addr, flags);
+
+	put_current_task(task);
+	return status;
 }
 
 kern_status_t pmap_add(

arch/x86_64/thread.c

@@ -95,25 +95,30 @@ kern_status_t ml_thread_config_set(
 	const void *ptr,
 	size_t len)
 {
+	struct thread *self = get_current_thread();
+	kern_status_t status = KERN_OK;
+
 	switch (key) {
 	case THREAD_CFG_FSBASE:
 		if (len != sizeof(thread->tr_ml.tr_fsbase)) {
-			return KERN_INVALID_ARGUMENT;
+			status = KERN_INVALID_ARGUMENT;
+			break;
 		}
 
 		thread->tr_ml.tr_fsbase = *(virt_addr_t *)ptr;
-		if (thread == current_thread()) {
+		if (thread == self) {
 			wrmsr(MSR_FS_BASE, thread->tr_ml.tr_fsbase);
 		}
 
 		break;
 	case THREAD_CFG_GSBASE:
 		if (len != sizeof(thread->tr_ml.tr_gsbase)) {
-			return KERN_INVALID_ARGUMENT;
+			status = KERN_INVALID_ARGUMENT;
+			break;
 		}
 
 		thread->tr_ml.tr_gsbase = *(virt_addr_t *)ptr;
-		if (thread == current_thread()) {
+		if (thread == self) {
 			/* we're in the kernel right now, so the user and kernel
 			 * gs-base registers are swapped. when we return to
 			 * usermode, this value will be swapped back into
@@ -123,8 +128,10 @@ kern_status_t ml_thread_config_set(
 
 		break;
 	default:
-		return KERN_INVALID_ARGUMENT;
+		status = KERN_INVALID_ARGUMENT;
+		break;
 	}
 
-	return KERN_OK;
+	put_current_thread(self);
+	return status;
 }