From 082ca85b7ce06bf00a460b2efcbab801bfb8ec75 Mon Sep 17 00:00:00 2001
From: Tristan Honscheid <honscheid@google.com>
Date: Thu, 10 Feb 2022 16:41:22 -0700
Subject: [PATCH] ztest: bug: Add friendly name helper function  for all phases

The `handle_signal()` function in the new ztest API (`ztest_new.c`) uses
an array of strings to lookup a friendly name for each test phase, but
the array only has three elements despite there being six test phases.
This can lead to an out-of-bounds issue. Replace the array with a helper
function and switch statement.

Signed-off-by: Tristan Honscheid <honscheid@google.com>
---
 subsys/testsuite/ztest/src/ztest_new.c | 49 +++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 9 deletions(-)

diff --git a/subsys/testsuite/ztest/src/ztest_new.c b/subsys/testsuite/ztest/src/ztest_new.c
index 6a38eb0aeea..b4854900d8c 100644
--- a/subsys/testsuite/ztest/src/ztest_new.c
+++ b/subsys/testsuite/ztest/src/ztest_new.c
@@ -23,14 +23,25 @@ static struct k_thread ztest_thread;
 
 /* ZTEST_DMEM and ZTEST_BMEM are used for the application shared memory test  */
 
-ZTEST_DMEM enum {
+/**
+ * @brief Each enum member represents a distinct phase of execution for the
+ *        test binary. TEST_PHASE_FRAMEWORK is active when internal ztest code
+ *        is executing; the rest refer to corresponding phases of user test
+ *        code.
+ */
+enum ztest_phase {
 	TEST_PHASE_SETUP,
 	TEST_PHASE_BEFORE,
 	TEST_PHASE_TEST,
 	TEST_PHASE_AFTER,
 	TEST_PHASE_TEARDOWN,
 	TEST_PHASE_FRAMEWORK
-} phase = TEST_PHASE_FRAMEWORK;
+};
+
+/**
+ * @brief Tracks the current phase that ztest is operating in.
+ */
+ZTEST_DMEM enum ztest_phase phase = TEST_PHASE_FRAMEWORK;
 
 static ZTEST_BMEM int test_status;
 
@@ -248,14 +259,34 @@ void ztest_test_pass(void)
 	longjmp(test_pass, 1);
 }
 
+/**
+ * @brief Get a friendly name string for a given test phrase.
+ *
+ * @param phase an enum ztest_phase value describing the desired test phase
+ * @returns a string name for `phase`
+ */
+static inline const char *get_friendly_phase_name(enum ztest_phase phase)
+{
+	switch (phase) {
+	case TEST_PHASE_SETUP:
+		return "setup";
+	case TEST_PHASE_BEFORE:
+		return "before";
+	case TEST_PHASE_TEST:
+		return "test";
+	case TEST_PHASE_AFTER:
+		return "after";
+	case TEST_PHASE_TEARDOWN:
+		return "teardown";
+	case TEST_PHASE_FRAMEWORK:
+		return "framework";
+	default:
+		return "(unknown)";
+	}
+}
+
 static void handle_signal(int sig)
 {
-	static const char *const phase_str[] = {
-		"setup",
-		"unit test",
-		"teardown",
-	};
-
 	PRINT("    %s", strsignal(sig));
 	switch (phase) {
 	case TEST_PHASE_SETUP:
@@ -263,7 +294,7 @@ static void handle_signal(int sig)
 	case TEST_PHASE_TEST:
 	case TEST_PHASE_AFTER:
 	case TEST_PHASE_TEARDOWN:
-		PRINT(" at %s function\n", phase_str[phase]);
+		PRINT(" at %s function\n", get_friendly_phase_name(phase));
 		longjmp(test_fail, 1);
 	case TEST_PHASE_FRAMEWORK:
 		PRINT("\n");