4 files changed, 22 insertions, 1 deletions
diff --git a/src/kernel/ao_task.c b/src/kernel/ao_task.c
index 55e423bb..104d1074 100644
--- a/src/kernel/ao_task.c
+++ b/src/kernel/ao_task.c
@@ -355,7 +355,6 @@ ao_yield(void) ao_arch_naked_define
 	 */
 	if (ao_cur_task->wchan == NULL)
 		ao_task_to_run_queue(ao_cur_task);
-	ao_cur_task = NULL;
 	for (;;) {
 		ao_arch_memory_barrier();
 		if (!ao_list_is_empty(&run_queue))
@@ -425,6 +424,7 @@ ao_sleep(__xdata void *wchan)
 void
 ao_wakeup(__xdata void *wchan) __reentrant
 {
+	ao_validate_cur_stack();
 #if HAS_TASK_QUEUE
 	struct ao_task	*sleep, *next;
 	struct ao_list	*sleep_queue;
diff --git a/src/kernel/ao_task.h b/src/kernel/ao_task.h
index c6bec0e3..0e353fe8 100644
--- a/src/kernel/ao_task.h
+++ b/src/kernel/ao_task.h
@@ -56,6 +56,10 @@ extern __data uint8_t ao_num_tasks;
 extern __xdata struct ao_task *__data ao_cur_task;
 extern __data uint8_t ao_task_minimize_latency;	/* Reduce IRQ latency */
 
+#ifndef HAS_ARCH_VALIDATE_CUR_STACK
+#define ao_validate_cur_stack()
+#endif
+
 /*
  ao_task.c
  */
diff --git a/src/stm/ao_arch_funcs.h b/src/stm/ao_arch_funcs.h
index 02b344b1..6b38032c 100644
--- a/src/stm/ao_arch_funcs.h
+++ b/src/stm/ao_arch_funcs.h
@@ -366,6 +366,22 @@ static inline void ao_arch_restore_stack(void) {
 #define HAS_SAMPLE_PROFILE 0
 #endif
 
+#if DEBUG
+#define HAS_ARCH_VALIDATE_CUR_STACK	1
+
+static inline void
+ao_validate_cur_stack(void)
+{
+	uint8_t		*psp;
+
+	asm("mrs %0,psp" : "=&r" (psp));
+	if (ao_cur_task &&
+	    psp <= ao_cur_task->stack &&
+	    psp >= ao_cur_task->stack - 256)
+		ao_panic(AO_PANIC_STACK);
+}
+#endif
+
 #if !HAS_SAMPLE_PROFILE
 #define HAS_ARCH_START_SCHEDULER	1
 
diff --git a/src/stm/ao_timer.c b/src/stm/ao_timer.c
index 8db62e76..91ede84b 100644
--- a/src/stm/ao_timer.c
+++ b/src/stm/ao_timer.c
@@ -41,6 +41,7 @@ volatile __data uint8_t	ao_data_count;
 
 void stm_systick_isr(void)
 {
+	ao_validate_cur_stack();
 	if (stm_systick.csr & (1 << STM_SYSTICK_CSR_COUNTFLAG)) {
 		++ao_tick_count;
 #if HAS_TASK_QUEUE