Fix memory leak of Task structures

framework/aster-frame/src/task/processor.rs

@@ -17,6 +17,9 @@ use crate::{cpu_local, CpuLocal};
 
 pub struct Processor {
     current: Option<Arc<Task>>,
+    /// A temporary variable used in [`switch_to_task`] to avoid dropping `current` while running
+    /// as `current`.
+    prev_task: Option<Arc<Task>>,
     idle_task_ctx: TaskContext,
 }
 
@@ -24,6 +27,7 @@ impl Processor {
     pub const fn new() -> Self {
         Self {
             current: None,
+            prev_task: None,
             idle_task_ctx: TaskContext::new(),
         }
     }
@@ -122,22 +126,35 @@ fn switch_to_task(next_task: Arc<Task>) {
 
     let next_task_ctx_ptr = next_task.ctx().get().cast_const();
 
-    // Change the current task to the next task.
-    CpuLocal::borrow_with(&PROCESSOR, |processor| {
-        processor.borrow_mut().current = Some(next_task.clone());
-    });
-
     if let Some(next_user_space) = next_task.user_space() {
         next_user_space.vm_space().activate();
     }
 
+    // Change the current task to the next task.
+    CpuLocal::borrow_with(&PROCESSOR, |processor| {
+        let mut processor = processor.borrow_mut();
+
+        // We cannot directly overwrite `current` at this point. Since we are running as `current`,
+        // we must avoid dropping `current`. Otherwise, the kernel stack may be unmapped, leading
+        // to soundness problems.
+        let old_current = processor.current.replace(next_task);
+        processor.prev_task = old_current;
+    });
+
     // SAFETY:
     // 1. `ctx` is only used in `schedule()`. We have exclusive access to both the current task
     //    context and the next task context.
     // 2. The next task context is a valid task context.
     unsafe {
+        // This function may not return, for example, when the current task exits. So make sure
+        // that all variables on the stack can be forgotten without causing resource leakage.
         context_switch(current_task_ctx_ptr, next_task_ctx_ptr);
     }
+
+    // Now it's fine to drop `prev_task`. However, we choose not to do this because it is not
+    // always possible. For example, `context_switch` can switch directly to the entry point of the
+    // next task. Not dropping is just fine because the only consequence is that we delay the drop
+    // to the next task switching.
 }
 
 cpu_local! {

framework/aster-frame/src/task/task.rs

@@ -178,8 +178,19 @@ impl Task {
         }
     }
 
-    pub fn exit(&self) -> ! {
+    /// Exits the current task.
+    ///
+    /// The task `self` must be the task that is currently running.
+    ///
+    /// **NOTE:** If there is anything left on the stack, it will be forgotten. This behavior may
+    /// lead to resource leakage.
+    fn exit(self: Arc<Self>) -> ! {
         self.inner_exclusive_access().task_status = TaskStatus::Exited;
+
+        // `current_task()` still holds a strong reference, so nothing is destroyed at this point,
+        // neither is the kernel stack.
+        drop(self);
+
         schedule();
         unreachable!()
     }

kernel/aster-nix/src/process/program_loader/elf/load_elf.rs

@@ -7,7 +7,7 @@
 //! When create a process from elf file, we will use the elf_load_info to construct the VmSpace
 
 use align_ext::AlignExt;
-use aster_frame::{mm::VmIo, task::Task};
+use aster_frame::mm::VmIo;
 use aster_rights::{Full, Rights};
 use xmas_elf::program::{self, ProgramHeader64};
 
@@ -75,7 +75,7 @@ pub fn load_elf_to_vm(
                 user_stack_top,
             })
         }
-        Err(_) => {
+        Err(err) => {
             // Since the process_vm is in invalid state,
             // the process cannot return to user space again,
             // so `Vmar::clear` and `do_exit_group` are called here.
@@ -87,7 +87,9 @@ pub fn load_elf_to_vm(
             // the macro will panic. This corner case should be handled later.
             // FIXME: how to set the correct exit status?
             do_exit_group(TermStatus::Exited(1));
-            Task::current().exit();
+
+            // The process will exit and the error code will be ignored.
+            Err(err)
         }
     }
 }

kernel/aster-nix/src/process/signal/mod.rs

@@ -16,7 +16,7 @@ pub mod signals;
 use core::{mem, sync::atomic::Ordering};
 
 use align_ext::AlignExt;
-use aster_frame::{cpu::UserContext, task::Task, user::UserContextApi};
+use aster_frame::{cpu::UserContext, user::UserContextApi};
 use c_types::{siginfo_t, ucontext_t};
 pub use events::{SigEvents, SigEventsFilter};
 pub use pauser::Pauser;
@@ -90,9 +90,8 @@ pub fn handle_pending_signal(
                         current.executable_path(),
                         sig_num.sig_name()
                     );
-                    do_exit_group(TermStatus::Killed(sig_num));
                     // We should exit current here, since we cannot restore a valid status from trap now.
-                    Task::current().exit();
+                    do_exit_group(TermStatus::Killed(sig_num));
                 }
                 SigDefaultAction::Ign => {}
                 SigDefaultAction::Stop => {

kernel/aster-nix/src/thread/task.rs

@@ -57,22 +57,20 @@ pub fn create_new_user_task(user_space: Arc<UserSpace>, thread_ref: Weak<Thread>
                 break;
             }
             handle_pending_signal(context, &current_thread).unwrap();
-            if current_thread.status().is_exited() {
-                debug!("exit due to signal");
-                break;
-            }
             // If current is suspended, wait for a signal to wake up self
             while current_thread.status().is_stopped() {
                 Thread::yield_now();
                 debug!("{} is suspended.", current_thread.tid());
                 handle_pending_signal(context, &current_thread).unwrap();
             }
+            if current_thread.status().is_exited() {
+                debug!("exit due to signal");
+                break;
+            }
             // a preemption point after handling user event.
             preempt(current_task);
         }
         debug!("exit user loop");
-        // FIXME: This is a work around: exit in kernel task entry may be not called. Why this will happen?
-        current_task.exit();
     }
 
     TaskOptions::new(user_task_entry)