arch: fix addrenv_switch changing this_task causing exceptions

After addrenv_switch(), the current running task (this_task) may change due to
deferred work execution. Update all architecture interrupt, syscall, and exit
handlers to re-fetch tcb = this_task() after addrenv_switch(). Ensures scheduler
and context operations use the correct TCB, preventing context corruption and
exceptions across SMP and memory-protected builds.

Signed-off-by: hujun5 <hujun5@xiaomi.com>

arch/arm/src/arm/arm_doirq.c

@@ -98,6 +98,7 @@ uint32_t *arm_doirq(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/arm/src/arm/arm_syscall.c

@@ -95,6 +95,10 @@ uint32_t *arm_syscall(uint32_t *regs)
         nxsched_switch_context(*running_task, tcb);
 
       case SYS_restore_context:
+#ifdef CONFIG_ARCH_ADDRENV
+        addrenv_switch(tcb);
+        tcb = this_task();
+#endif
 
         /* No context switch occurs in SYS_restore_context, or the
          * context switch has been completed, so there is no
@@ -106,9 +110,6 @@ uint32_t *arm_syscall(uint32_t *regs)
         /* Restore the cpu lock */
 
         restore_critical_section(tcb, cpu);
-#ifdef CONFIG_ARCH_ADDRENV
-        addrenv_switch(tcb);
-#endif
         break;
 
       default:

arch/arm/src/armv7-a/arm_doirq.c

@@ -98,6 +98,7 @@ uint32_t *arm_doirq(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/arm/src/armv7-a/arm_syscall.c

@@ -274,6 +274,14 @@ uint32_t *arm_syscall(uint32_t *regs)
 
       case SYS_restore_context:
 
+#ifdef CONFIG_ARCH_ADDRENV
+        /* The addrenv_switch may change this_task, for example addrenv drop
+         * will post sem to hpwork, so we have to call before restore context
+         */
+
+        addrenv_switch(tcb);
+        tcb = this_task();
+#endif
         /* No context switch occurs in SYS_restore_context, or the
          * context switch has been completed, so there is no
          * need to update scheduler parameters.
@@ -285,9 +293,6 @@ uint32_t *arm_syscall(uint32_t *regs)
 
         restore_critical_section(tcb, cpu);
         regs = tcb->xcp.regs;
-#ifdef CONFIG_ARCH_ADDRENV
-        addrenv_switch(tcb);
-#endif
         break;
 
       /* R0=SYS_task_start:  This a user task start

arch/arm/src/armv7-r/arm_syscall.c

@@ -270,7 +270,10 @@ uint32_t *arm_syscall(uint32_t *regs)
         nxsched_switch_context(*running_task, tcb);
 
       case SYS_restore_context:
-
+#ifdef CONFIG_ARCH_ADDRENV
+        addrenv_switch(tcb);
+        tcb = this_task();
+#endif
         /* No context switch occurs in SYS_restore_context, or the
          * context switch has been completed, so there is no
          * need to update scheduler parameters.
@@ -282,9 +285,6 @@ uint32_t *arm_syscall(uint32_t *regs)
 
         restore_critical_section(tcb, cpu);
         regs = tcb->xcp.regs;
-#ifdef CONFIG_ARCH_ADDRENV
-        addrenv_switch(tcb);
-#endif
         break;
 
       /* R0=SYS_task_start:  This a user task start

arch/arm64/src/common/arm64_doirq.c

@@ -128,6 +128,7 @@ uint64_t *arm64_doirq(int irq, uint64_t * regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/arm64/src/common/arm64_syscall.c

@@ -196,11 +196,18 @@ uint64_t *arm64_syscall(uint64_t *regs)
         restore_critical_section(tcb, cpu);
 #ifdef CONFIG_ARCH_ADDRENV
         addrenv_switch(tcb);
+        tcb = this_task();
+        *running_task = tcb;
 #endif
         break;
 
       case SYS_switch_context:
 
+#ifdef CONFIG_ARCH_ADDRENV
+        addrenv_switch(tcb);
+        tcb = this_task();
+#endif
+
         /* Update scheduler parameters */
 
         nxsched_switch_context(*running_task, tcb);
@@ -209,9 +216,6 @@ uint64_t *arm64_syscall(uint64_t *regs)
         /* Restore the cpu lock */
 
         restore_critical_section(tcb, cpu);
-#ifdef CONFIG_ARCH_ADDRENV
-        addrenv_switch(tcb);
-#endif
         break;
 
 #ifdef CONFIG_BUILD_KERNEL

arch/avr/src/avr32/avr_doirq.c

@@ -107,6 +107,7 @@ uint32_t *avr_doirq(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters. */

arch/avr/src/avr32/avr_switchcontext.c

@@ -88,6 +88,7 @@ void up_switch_context(struct tcb_s *tcb, struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
       /* Update scheduler parameters */
 

arch/avr/src/common/avr_exit.c

@@ -67,10 +67,6 @@ void up_exit(int status)
 
   tcb = this_task();
 
-  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
-
-  g_running_tasks[this_cpu()] = tcb;
-
 #ifdef CONFIG_ARCH_ADDRENV
   /* Make sure that the address environment for the previously running
    * task is closed down gracefully (data caches dump, MMU flushed) and
@@ -79,8 +75,13 @@ void up_exit(int status)
    */
 
   addrenv_switch(tcb);
+  tcb = this_task();
 #endif
 
+  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
+
+  g_running_tasks[this_cpu()] = tcb;
+
   /* Then switch contexts */
 
   avr_fullcontextrestore(tcb->xcp.regs);

arch/hc/src/common/hc_doirq.c

@@ -107,6 +107,7 @@ uint8_t *hc_doirq(int irq, uint8_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters. */

arch/hc/src/common/hc_exit.c

@@ -66,10 +66,6 @@ void up_exit(int status)
 
   tcb = this_task();
 
-  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
-
-  g_running_tasks[this_cpu()] = tcb;
-
 #ifdef CONFIG_ARCH_ADDRENV
   /* Make sure that the address environment for the previously running
    * task is closed down gracefully (data caches dump, MMU flushed) and
@@ -78,8 +74,13 @@ void up_exit(int status)
    */
 
   addrenv_switch(tcb);
+  tcb = this_task();
 #endif
 
+  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
+
+  g_running_tasks[this_cpu()] = tcb;
+
   /* Then switch contexts */
 
   hc_fullcontextrestore(tcb->xcp.regs);

arch/hc/src/common/hc_switchcontext.c

@@ -91,6 +91,7 @@ void up_switch_context(struct tcb_s *tcb, struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
       /* Update scheduler parameters */
 

arch/mips/src/mips32/mips_doirq.c

@@ -116,6 +116,7 @@ uint32_t *mips_doirq(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters. */

arch/misoc/src/lm32/lm32_doirq.c

@@ -101,6 +101,7 @@ uint32_t *lm32_doirq(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters. */

arch/misoc/src/minerva/minerva_doirq.c

@@ -100,6 +100,7 @@ uint32_t *minerva_doirq(int irq, uint32_t * regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters. */

arch/or1k/src/common/or1k_exit.c

@@ -68,10 +68,6 @@ void up_exit(int status)
 
   tcb = this_task();
 
-  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
-
-  g_running_tasks[this_cpu()] = tcb;
-
 #ifdef CONFIG_ARCH_ADDRENV
   /* Make sure that the address environment for the previously running
    * task is closed down gracefully (data caches dump, MMU flushed) and
@@ -80,8 +76,13 @@ void up_exit(int status)
    */
 
   addrenv_switch(tcb);
+  tcb = this_task();
 #endif
 
+  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
+
+  g_running_tasks[this_cpu()] = tcb;
+
   /* Then switch contexts */
 
   or1k_fullcontextrestore(tcb->xcp.regs);

arch/or1k/src/common/or1k_switchcontext.c

@@ -95,6 +95,7 @@ void up_switch_context(struct tcb_s *tcb, struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
       /* Update scheduler parameters */
 

arch/renesas/src/common/renesas_doirq.c

@@ -110,6 +110,7 @@ uint32_t *renesas_doirq(int irq, uint32_t * regs)
            */
 
           addrenv_switch(tcb);
+          tcb = this_task();
 #endif
 
           /* Update scheduler parameters. */

arch/renesas/src/common/renesas_exit.c

@@ -66,10 +66,6 @@ void up_exit(int status)
 
   tcb = this_task();
 
-  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
-
-  g_running_tasks[this_cpu()] = tcb;
-
 #ifdef CONFIG_ARCH_ADDRENV
   /* Make sure that the address environment for the previously running
    * task is closed down gracefully (data caches dump, MMU flushed) and
@@ -78,8 +74,13 @@ void up_exit(int status)
    */
 
   addrenv_switch(tcb);
+  tcb = this_task();
 #endif
 
+  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
+
+  g_running_tasks[this_cpu()] = tcb;
+
   /* Then switch contexts */
 
   renesas_fullcontextrestore(tcb->xcp.regs);

arch/renesas/src/common/renesas_switchcontext.c

@@ -91,6 +91,7 @@ void up_switch_context(struct tcb_s *tcb, struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
       /* Update scheduler parameters */
 

arch/risc-v/src/common/riscv_doirq.c

@@ -124,6 +124,7 @@ uintreg_t *riscv_doirq(int irq, uintreg_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/risc-v/src/common/supervisor/riscv_perform_syscall.c

@@ -68,6 +68,7 @@ void *riscv_perform_syscall(uintreg_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
       /* Update scheduler parameters */
 

arch/sparc/src/sparc_v8/sparc_v8_doirq.c

@@ -111,6 +111,7 @@ uint32_t *sparc_doirq(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters. */

arch/tricore/src/common/tricore_doirq.c

@@ -89,6 +89,7 @@ IFX_INTERRUPT_INTERNAL(tricore_doirq, 0, 255)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/x86/src/common/x86_exit.c

@@ -67,12 +67,6 @@ void up_exit(int status)
 
   tcb = this_task();
 
-  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases
-   * NOTE: the API also adjusts the global IRQ control for SMP
-   */
-
-  g_running_tasks[this_cpu()] = tcb;
-
 #ifdef CONFIG_ARCH_ADDRENV
   /* Make sure that the address environment for the previously running
    * task is closed down gracefully (data caches dump, MMU flushed) and
@@ -81,8 +75,15 @@ void up_exit(int status)
    */
 
   addrenv_switch(tcb);
+  tcb = this_task();
 #endif
 
+  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases
+   * NOTE: the API also adjusts the global IRQ control for SMP
+   */
+
+  g_running_tasks[this_cpu()] = tcb;
+
   /* Then switch contexts */
 
   x86_fullcontextrestore(tcb->xcp.regs);

arch/x86/src/common/x86_switchcontext.c

@@ -91,6 +91,7 @@ void up_switch_context(struct tcb_s *tcb, struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
       /* Update scheduler parameters */
 

arch/x86/src/qemu/qemu_handlers.c

@@ -121,6 +121,7 @@ static uint32_t *common_handler(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/x86_64/src/common/x86_64_exit.c

@@ -66,6 +66,17 @@ void up_exit(int status)
 
   tcb = this_task();
 
+#ifdef CONFIG_ARCH_ADDRENV
+  /* Make sure that the address environment for the previously running
+   * task is closed down gracefully (data caches dump, MMU flushed) and
+   * set up the address environment for the new thread at the head of
+   * the ready-to-run list.
+   */
+
+  addrenv_switch(tcb);
+  tcb = this_task();
+#endif
+
   /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases
    * NOTE: the API also adjusts the global IRQ control for SMP
    */
@@ -76,16 +87,6 @@ void up_exit(int status)
 
   x86_64_restore_auxstate(tcb);
 
-#ifdef CONFIG_ARCH_ADDRENV
-  /* Make sure that the address environment for the previously running
-   * task is closed down gracefully (data caches dump, MMU flushed) and
-   * set up the address environment for the new thread at the head of
-   * the ready-to-run list.
-   */
-
-  addrenv_switch(tcb);
-#endif
-
   /* Restore the cpu lock */
 
   restore_critical_section(tcb, this_cpu());

arch/x86_64/src/common/x86_64_switchcontext.c

@@ -97,6 +97,7 @@ void up_switch_context(struct tcb_s *tcb, struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Restore the cpu lock */

arch/x86_64/src/intel64/intel64_handlers.c

@@ -102,6 +102,7 @@ static uint64_t *common_handler(int irq, uint64_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/xtensa/src/common/xtensa_irqdispatch.c

@@ -93,6 +93,7 @@ uint32_t *xtensa_irq_dispatch(int irq, uint32_t *regs)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */

arch/z80/src/common/z80_doirq.c

@@ -93,6 +93,7 @@ FAR chipreg_t *z80_doirq(uint8_t irq, FAR chipreg_t *regs)
            */
 
           addrenv_switch(tcb);
+          tcb = this_task();
 #endif
 
           /* Update scheduler parameters */

arch/z80/src/common/z80_exit.c

@@ -67,11 +67,6 @@ void up_exit(int status)
    */
 
   tcb = this_task();
-  sinfo("New Active Task TCB=%p\n", tcb);
-
-  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
-
-  g_running_tasks[this_cpu()] = tcb;
 
 #ifdef CONFIG_ARCH_ADDRENV
   /* Make sure that the address environment for the previously running
@@ -81,8 +76,15 @@ void up_exit(int status)
    */
 
   addrenv_switch(tcb);
+  tcb = this_task();
 #endif
 
+  sinfo("New Active Task TCB=%p\n", tcb);
+
+  /* Adjusts time slice for SCHED_RR & SCHED_SPORADIC cases */
+
+  g_running_tasks[this_cpu()] = tcb;
+
   /* Then switch contexts */
 
   RESTORE_USERCONTEXT(tcb);

arch/z80/src/common/z80_switchcontext.c

@@ -94,6 +94,7 @@ void up_switch_context(FAR struct tcb_s *tcb, FAR struct tcb_s *rtcb)
        */
 
       addrenv_switch(tcb);
+      tcb = this_task();
 #endif
 
       /* Update scheduler parameters */