From 474d4c3249c190b3261bc128ec1847e0a4f7b176 Mon Sep 17 00:00:00 2001
From: Adam Kondraciuk <adam.kondraciuk@nordicsemi.no>
Date: Thu, 17 Oct 2024 12:28:03 +0200
Subject: [PATCH] arch: arm: cortex_m: pm_s2ram: Rework S2RAM mark functions

The S2RAM procedure requires marker checking after reset.
Such checking is performed on very early stage of the system initialization
and must ensure that the stack is not used due to the TLS pointer which is
not initialized yet.

Signed-off-by: Adam Kondraciuk <adam.kondraciuk@nordicsemi.no>
---
 arch/arm/core/cortex_m/pm_s2ram.S     | 23 +++++++--------
 arch/arm/core/cortex_m/pm_s2ram.c     | 40 +++++++++++++++++++++------
 include/zephyr/arch/common/pm_s2ram.h | 10 +++++++
 3 files changed, 54 insertions(+), 19 deletions(-)

diff --git a/arch/arm/core/cortex_m/pm_s2ram.S b/arch/arm/core/cortex_m/pm_s2ram.S
index f9c82b4069b..27c2a1e96a7 100644
--- a/arch/arm/core/cortex_m/pm_s2ram.S
+++ b/arch/arm/core/cortex_m/pm_s2ram.S
@@ -70,7 +70,9 @@ SECTION_FUNC(TEXT, arch_pm_s2ram_suspend)
 	/*
 	 * Mark entering suspend to RAM.
 	 */
-	bl pm_s2ram_mark_set
+	mov 	r1, lr
+	bl 	pm_s2ram_mark_set
+	mov	lr, r1
 
 	/*
 	 * Call the system_off function passed as parameter. This should never
@@ -86,7 +88,9 @@ SECTION_FUNC(TEXT, arch_pm_s2ram_suspend)
 	/*
 	 * Reset the marking of suspend to RAM, return is ignored.
 	 */
-	bl pm_s2ram_mark_check_and_clear
+	mov 	r1, lr
+	bl	pm_s2ram_mark_check_and_clear
+	mov 	lr, r1
 
 	/* Move system_off back to r0 as return value */
 	mov	r0, r4
@@ -99,16 +103,13 @@ GTEXT(arch_pm_s2ram_resume)
 SECTION_FUNC(TEXT, arch_pm_s2ram_resume)
 	/*
 	 * Check if reset occurred after suspending to RAM.
-	 * Store LR to ensure we can continue boot when we are not suspended
-	 * to RAM. In addition to LR, R0 is pushed too, to ensure "SP mod 8 = 0",
-	 * as stated by ARM rule 6.2.1.2 for AAPCS32.
 	 */
-	push    {r0, lr}
-	bl      pm_s2ram_mark_check_and_clear
-	cmp	r0, #0x1
-	pop     {r0, lr}
-	beq	resume
-	bx	lr
+	mov 	r1, lr
+	bl 	pm_s2ram_mark_check_and_clear
+	mov 	lr, r1
+	cmp     r0, #0x1
+	beq     resume
+	bx      lr
 
 resume:
 	/*
diff --git a/arch/arm/core/cortex_m/pm_s2ram.c b/arch/arm/core/cortex_m/pm_s2ram.c
index 2657d48dc32..b7fe5d9b626 100644
--- a/arch/arm/core/cortex_m/pm_s2ram.c
+++ b/arch/arm/core/cortex_m/pm_s2ram.c
@@ -22,20 +22,44 @@ __noinit _cpu_context_t _cpu_context;
  */
 static __noinit uint32_t marker;
 
-void pm_s2ram_mark_set(void)
+void __attribute__((naked)) pm_s2ram_mark_set(void)
 {
-	marker = MAGIC;
+	__asm__ volatile(
+		/* Set the marker to MAGIC value */
+		"str	%[_magic_val], [%[_marker]]\n"
+
+		"bx	lr\n"
+		:
+		: [_magic_val] "r"(MAGIC), [_marker] "r"(&marker)
+		: "r1", "r4", "memory");
 }
 
-bool pm_s2ram_mark_check_and_clear(void)
+bool __attribute__((naked)) pm_s2ram_mark_check_and_clear(void)
 {
-	if (marker == MAGIC) {
-		marker = 0;
+	__asm__ volatile(
+		/* Set return value to 0 */
+		"mov	r0, #0\n"
 
-		return true;
-	}
+		/* Check the marker */
+		"ldr	r3, [%[_marker]]\n"
+		"cmp	r3, %[_magic_val]\n"
+		"bne	exit\n"
 
-	return false;
+		/*
+		 * Reset the marker
+		 */
+		"str	r0, [%[_marker]]\n"
+
+		/*
+		 * Set return value to 1
+		 */
+		"mov	r0, #1\n"
+
+		"exit:\n"
+		"bx lr\n"
+		:
+		: [_magic_val] "r"(MAGIC), [_marker] "r"(&marker)
+		: "r0", "r1", "r3", "r4", "memory");
 }
 
 #endif /* CONFIG_PM_S2RAM_CUSTOM_MARKING */
diff --git a/include/zephyr/arch/common/pm_s2ram.h b/include/zephyr/arch/common/pm_s2ram.h
index 34c544c769b..ad9ab8ad7a4 100644
--- a/include/zephyr/arch/common/pm_s2ram.h
+++ b/include/zephyr/arch/common/pm_s2ram.h
@@ -65,6 +65,11 @@ int arch_pm_s2ram_suspend(pm_s2ram_system_off_fn_t system_off);
  *
  * Default implementation is setting a magic word in RAM. CONFIG_PM_S2RAM_CUSTOM_MARKING
  * allows custom implementation.
+ * The following requirements must be fulfilled:
+ * - the function cannot use stack (asm function or function with 'naked' attribute)
+ * - the content of the R1 and R4 registers must remain unchanged
+ * - returning from the function should be performed with the `bx lr` instruction
+ *
  */
 void pm_s2ram_mark_set(void);
 
@@ -76,6 +81,11 @@ void pm_s2ram_mark_set(void);
  *
  * Default implementation is checking a magic word in RAM. CONFIG_PM_S2RAM_CUSTOM_MARKING
  * allows custom implementation.
+ * The following requirements must be fulfilled:
+ * - the function cannot use stack (most likely asm function)
+ * - the content of the R1 and R4 registers must remain unchanged
+ * - the function's return value is passed by R0 register
+ * - returning from the function should be performed with the `bx lr` instruction
  *
  * @retval true if marking is found which indicates resuming after suspend-to-RAM.
  * @retval false if marking is not found which indicates standard boot.