minidiag: Add storage self-test screen

Implement the ui for storage self-test.

BUG=b:173364332
TEST=make clean && CC=x86_64-pc-linux-gnu-clang make runtests
BRANCH=none

Cq-Depend: chromium:2641906
Signed-off-by: Chung-Sheng Wu <chungsheng@google.com>
Change-Id: Ieb3a24e0bf19f4e62de41cccd1138719dbc65e7c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2626566
Tested-by: Chung-Sheng Wu <chungsheng@chromium.org>
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
Reviewed-by: Hsuan Ting Chen <roccochen@chromium.org>
Reviewed-by: Meng-Huan Yu <menghuan@chromium.org>
Commit-Queue: Chung-Sheng Wu <chungsheng@chromium.org>
diff --git a/firmware/2lib/2stub.c b/firmware/2lib/2stub.c
index f7b15ce..7cec2e5 100644
--- a/firmware/2lib/2stub.c
+++ b/firmware/2lib/2stub.c
@@ -194,6 +194,19 @@
 }
 
 __attribute__((weak))
+vb2_error_t vb2ex_diag_get_storage_test_log(const char **out)
+{
+	*out = "mock";
+	return VB2_SUCCESS;
+}
+
+__attribute__((weak))
+vb2_error_t vb2ex_diag_storage_test_control(enum vb2_diag_storage_test ops)
+{
+	return VB2_SUCCESS;
+}
+
+__attribute__((weak))
 vb2_error_t vb2ex_diag_memory_quick_test(int reset, const char **out)
 {
 	*out = "mock";
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index ed13b9c..1ac50ed 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -1026,12 +1026,38 @@
 /******************************************************************************/
 /* VB2_SCREEN_DIAGNOSTICS */
 
+#define DIAGNOSTICS_ITEM_STORAGE_HEALTH 1
+#define DIAGNOSTICS_ITEM_STORAGE_TEST_SHORT 2
+#define DIAGNOSTICS_ITEM_STORAGE_TEST_EXTENDED 3
+
+static vb2_error_t diagnostics_init(struct vb2_ui_context *ui)
+{
+	const char *unused_log_string;
+	vb2_error_t rv = vb2ex_diag_get_storage_test_log(&unused_log_string);
+	if (rv == VB2_ERROR_EX_UNIMPLEMENTED) {
+		VB2_SET_BIT(ui->state->disabled_item_mask,
+			    DIAGNOSTICS_ITEM_STORAGE_TEST_SHORT);
+		VB2_SET_BIT(ui->state->disabled_item_mask,
+			    DIAGNOSTICS_ITEM_STORAGE_TEST_EXTENDED);
+	}
+	ui->state->selected_item = DIAGNOSTICS_ITEM_STORAGE_HEALTH;
+	return VB2_SUCCESS;
+}
+
 static const struct vb2_menu_item diagnostics_items[] = {
 	LANGUAGE_SELECT_ITEM,
-	{
+	[DIAGNOSTICS_ITEM_STORAGE_HEALTH] = {
 		.text = "Storage health info",
 		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH,
 	},
+	[DIAGNOSTICS_ITEM_STORAGE_TEST_SHORT] = {
+		.text = "Storage self-test (short)",
+		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT,
+	},
+	[DIAGNOSTICS_ITEM_STORAGE_TEST_EXTENDED] = {
+		.text = "Storage self-test (Extended)",
+		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED,
+	},
 	{
 		.text = "Memory check (quick)",
 		.target = VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK,
@@ -1046,6 +1072,7 @@
 static const struct vb2_screen_info diagnostics_screen = {
 	.id = VB2_SCREEN_DIAGNOSTICS,
 	.name = "Diagnostic tools",
+	.init = diagnostics_init,
 	.menu = MENU_ITEMS(diagnostics_items),
 };
 
@@ -1090,6 +1117,129 @@
 };
 
 /******************************************************************************/
+/* VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST */
+
+#define DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP 0
+#define DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN 1
+#define DIAGNOSTICS_STORAGE_TEST_ITEM_BACK 2
+#define DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL 3
+
+static vb2_error_t diagnostics_storage_test_update_impl(
+	struct vb2_ui_context *ui)
+{
+	const char *log_string;
+	int is_test_running = 0;
+
+	/* Early return if the test is done. */
+	if (ui->state->test_finished)
+		return VB2_SUCCESS;
+
+	vb2_error_t rv = vb2ex_diag_get_storage_test_log(&log_string);
+	switch (rv) {
+	case VB2_ERROR_EX_DIAG_TEST_RUNNING:
+		is_test_running = 1;
+		break;
+	case VB2_SUCCESS:
+		ui->state->test_finished = 1;
+		break;
+	default:
+		VB2_DEBUG("vb2ex_diag_get_storage_test_log returned %#x\n", rv);
+		return rv;
+	}
+	VB2_TRY(log_page_show_back_or_cancel(ui, is_test_running));
+	return log_page_update(ui, log_string);
+}
+
+static vb2_error_t diagnostics_storage_test_update(struct vb2_ui_context *ui)
+{
+	vb2_error_t rv = diagnostics_storage_test_update_impl(ui);
+	if (rv) {
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		return vb2_ui_screen_back(ui);
+	}
+	return VB2_SUCCESS;
+}
+
+static vb2_error_t diagnostics_storage_test_control(
+	struct vb2_ui_context *ui, enum vb2_diag_storage_test op)
+{
+	vb2_error_t rv = vb2ex_diag_storage_test_control(op);
+	if (rv) {
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		return vb2_ui_screen_back(ui);
+	}
+	return VB2_SUCCESS;
+}
+
+static vb2_error_t diagnostics_storage_test_init(struct vb2_ui_context *ui)
+{
+	VB2_TRY(diagnostics_storage_test_update(ui));
+	return log_page_reset_to_top(ui);
+}
+
+static vb2_error_t diagnostics_storage_test_short_init(
+	struct vb2_ui_context *ui)
+{
+	VB2_TRY(diagnostics_storage_test_control(ui,
+						 VB2_DIAG_STORAGE_TEST_STOP));
+	VB2_TRY(diagnostics_storage_test_control(ui,
+						 VB2_DIAG_STORAGE_TEST_SHORT));
+	return diagnostics_storage_test_init(ui);
+}
+
+static vb2_error_t diagnostics_storage_test_extended_init(
+	struct vb2_ui_context *ui)
+{
+	VB2_TRY(diagnostics_storage_test_control(ui,
+						 VB2_DIAG_STORAGE_TEST_STOP));
+	VB2_TRY(diagnostics_storage_test_control(
+		ui, VB2_DIAG_STORAGE_TEST_EXTENDED));
+	return diagnostics_storage_test_init(ui);
+}
+
+static vb2_error_t diagnostics_storage_test_cancel(struct vb2_ui_context *ui)
+{
+	VB2_TRY(diagnostics_storage_test_control(ui,
+						 VB2_DIAG_STORAGE_TEST_STOP));
+	return vb2_ui_screen_back(ui);
+}
+
+static const struct vb2_menu_item diagnostics_storage_test_items[] = {
+	[DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP] = PAGE_UP_ITEM,
+	[DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN] = PAGE_DOWN_ITEM,
+	[DIAGNOSTICS_STORAGE_TEST_ITEM_BACK] = BACK_ITEM,
+	[DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL] = {
+		.text = "Cancel",
+		.action = diagnostics_storage_test_cancel,
+	},
+	POWER_OFF_ITEM,
+};
+
+static const struct vb2_screen_info diagnostics_storage_test_short_screen = {
+	.id = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT,
+	.name = "Storage self-test (short)",
+	.init = diagnostics_storage_test_short_init,
+	.action = diagnostics_storage_test_update,
+	.menu = MENU_ITEMS(diagnostics_storage_test_items),
+	.page_up_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP,
+	.page_down_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN,
+	.back_item = DIAGNOSTICS_STORAGE_TEST_ITEM_BACK,
+	.cancel_item = DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL,
+};
+
+static const struct vb2_screen_info diagnostics_storage_test_extended_screen = {
+	.id = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED,
+	.name = "Storage self-test (extended)",
+	.init = diagnostics_storage_test_extended_init,
+	.action = diagnostics_storage_test_update,
+	.menu = MENU_ITEMS(diagnostics_storage_test_items),
+	.page_up_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP,
+	.page_down_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN,
+	.back_item = DIAGNOSTICS_STORAGE_TEST_ITEM_BACK,
+	.cancel_item = DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL,
+};
+
+/******************************************************************************/
 /* VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK
    VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL */
 
@@ -1231,6 +1381,8 @@
 	&developer_select_bootloader_screen,
 	&diagnostics_screen,
 	&diagnostics_storage_health_screen,
+	&diagnostics_storage_test_short_screen,
+	&diagnostics_storage_test_extended_screen,
 	&diagnostics_memory_quick_screen,
 	&diagnostics_memory_full_screen,
 };
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index 218a416..f7acdc6 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -1343,6 +1343,8 @@
 	VB2_SCREEN_DIAGNOSTICS			= 0x400,
 	/* Storage diagnostic screen */
 	VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH	= 0x410,
+	VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT	= 0x411,
+	VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED	= 0x412,
 	/* Memory diagnostic screens */
 	VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK    	= 0x420,
 	VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL     	= 0x421,
@@ -1524,6 +1526,21 @@
 vb2_error_t vb2ex_diag_get_storage_health(const char **out);
 
 /**
+ * Get the storage self-test log.
+ *
+ * @param out	For returning a read-only pointer of full log string which is
+ *		guaranteed to be null-terminated. The function will manage
+ *		memory internally, so the returned pointer will only be valid
+ *		until next call.
+ * @return The status of storage test. VB2_SUCCESS means the test is finished,
+ * regardless of passing or failing. VB2_ERROR_EX_DIAG_TEST_RUNNING means
+ * the test is still running. VB2_ERROR_EX_UNIMPLEMENTED means the storage
+ * self-test is not supported on this device. Other non-zero codes for internal
+ * errors.
+ */
+vb2_error_t vb2ex_diag_get_storage_test_log(const char **out);
+
+/**
  * Get the memory diagnostic status. When it is called, it will take over the
  * control for a short period of time running memory test, and then return the
  * result of current status. If `reset` is not zero, it will reset the memory
@@ -1545,6 +1562,17 @@
 vb2_error_t vb2ex_diag_memory_full_test(int reset, const char **out);
 
 /*****************************************************************************/
+/* Functions for diagnostics control. */
+
+enum vb2_diag_storage_test {
+	VB2_DIAG_STORAGE_TEST_STOP = 0,
+	VB2_DIAG_STORAGE_TEST_SHORT,
+	VB2_DIAG_STORAGE_TEST_EXTENDED,
+};
+
+vb2_error_t vb2ex_diag_storage_test_control(enum vb2_diag_storage_test ops);
+
+/*****************************************************************************/
 /* Timer. */
 
 /**
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 017bd44..bfb2b8b 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -79,7 +79,7 @@
 	uint32_t page_count;
 	uint32_t current_page;
 
-	/* For memory check screen. */
+	/* For minidiag test screens. */
 	int test_finished;  /* Do not update screen if the content is done */
 
 	struct vb2_screen_state *prev;
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 607087e..3eb0c9f 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -100,6 +100,8 @@
 static int mock_prepare_log_count;
 static uint32_t mock_log_page_count;
 
+static vb2_error_t mock_diag_storage_test_rv;
+
 static void add_mock_key(uint32_t press, int trusted)
 {
 	if (mock_key_total >= ARRAY_SIZE(mock_key) ||
@@ -353,6 +355,8 @@
 	else
 		add_mock_vbtlk(VB2_ERROR_MOCK, 0);
 	add_mock_pp_pressed(0);
+
+	mock_diag_storage_test_rv = VB2_SUCCESS;
 }
 
 /* Mock functions */
@@ -571,6 +575,11 @@
 	return mock_log_page_count;
 }
 
+vb2_error_t vb2ex_diag_get_storage_test_log(const char **log)
+{
+	return mock_diag_storage_test_rv;
+}
+
 /* Tests */
 static void developer_tests(void)
 {
@@ -1794,8 +1803,8 @@
 	reset_common_data(FOR_DIAGNOSTICS);
 	TEST_EQ(vb2_diagnostic_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"diagnostic screen: no disabled or hidden item");
-	DISPLAYED_EQ("diagnostic menu", VB2_SCREEN_DIAGNOSTICS,
-		     MOCK_IGNORE, MOCK_IGNORE, 0x0, 0x0, MOCK_IGNORE);
+	DISPLAYED_EQ("diagnostic menu", VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE,
+		     MOCK_IGNORE, 0x0, 0x0, MOCK_IGNORE);
 
 	/* Diagnostics screen */
 	reset_common_data(FOR_DIAGNOSTICS);
@@ -1803,67 +1812,94 @@
 	/* #0: Language menu */
 	add_mock_keypress(VB_KEY_UP);
 	add_mock_keypress(VB_KEY_ENTER);
-	/* #1: Storage screen */
 	add_mock_keypress(VB_KEY_ESC);
+	/* #1: Storage health screen */
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
-	/* #2: Quick memory test screen */
 	add_mock_keypress(VB_KEY_ESC);
+	/* #2: Short storage self-test screen */
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
-	/* #3: Full memory test screen */
 	add_mock_keypress(VB_KEY_ESC);
+	/* #3: Extended storage self-test screen */
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
-	/* #4: Power off (End of menu) */
 	add_mock_keypress(VB_KEY_ESC);
+	/* #4: Quick memory test screen */
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_ENTER);
+	add_mock_keypress(VB_KEY_ESC);
+	/* #5: Full memory test screen */
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_ENTER);
+	add_mock_keypress(VB_KEY_ESC);
+	/* #6: Power off (End of menu) */
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	mock_calls_until_shutdown = -1;
 	TEST_EQ(vb2_diagnostic_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"diagnostic screen");
 
-	DISPLAYED_EQ("default on first button of menu",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 1, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_EQ("default on first button of menu", VB2_SCREEN_DIAGNOSTICS,
+		     MOCK_IGNORE, 1, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #0: Language menu */
-	DISPLAYED_EQ("language selection",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 0, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_EQ("language selection", VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE,
+		     0, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	DISPLAYED_EQ("#0: language menu", VB2_SCREEN_LANGUAGE_SELECT,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE,
 		     MOCK_IGNORE);
-	/* #1: Storage screen */
 	DISPLAYED_PASS();
-	DISPLAYED_EQ("storage button",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 1, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
+	/* #1: Storage health screen */
+	DISPLAYED_EQ("storage health button", VB2_SCREEN_DIAGNOSTICS,
+		     MOCK_IGNORE, 1, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	DISPLAYED_EQ("#1: storage screen",
 		     VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH, MOCK_IGNORE,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	/* #2: Quick memory test screen */
 	DISPLAYED_PASS();
-	DISPLAYED_EQ("quick memory test button",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 2, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#1: quick memory test screen",
-		     VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK, MOCK_IGNORE,
+	/* #2: Short storage self-test screen */
+	DISPLAYED_EQ("short storage self-test button", VB2_SCREEN_DIAGNOSTICS,
+		     MOCK_IGNORE, 2, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_EQ("#2: short storage self-test screen",
+		     VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT, MOCK_IGNORE,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	/* #3: Full memory test screen */
 	DISPLAYED_PASS();
-	DISPLAYED_EQ("full memory test button",
+	/* #3: Extended storage self-test screen */
+	DISPLAYED_EQ("extended storage self-test button",
 		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 3, MOCK_IGNORE,
 		     MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#3: full memory test screen",
+	DISPLAYED_EQ("#3: extended storage self-test screen",
+		     VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_PASS();
+	/* #4: Quick memory test screen */
+	DISPLAYED_EQ("quick memory test button", VB2_SCREEN_DIAGNOSTICS,
+		     MOCK_IGNORE, 4, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_EQ("#4: quick memory test screen",
+		     VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_PASS();
+	/* #5: Full memory test screen */
+	DISPLAYED_EQ("full memory test button", VB2_SCREEN_DIAGNOSTICS,
+		     MOCK_IGNORE, 5, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_EQ("#5: full memory test screen",
 		     VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL, MOCK_IGNORE,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	/* #4: Power of (End of menu) */
 	DISPLAYED_PASS();
-	DISPLAYED_EQ("power off",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 4, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
+	/* #6: Power of (End of menu) */
+	DISPLAYED_EQ("power off", VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 6,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();
 
+	/* Diagnostics screen: no nvme */
+	reset_common_data(FOR_DIAGNOSTICS);
+	/* Non-nvme storage returns UNIMPLEMENTED. */
+	mock_diag_storage_test_rv = VB2_ERROR_EX_UNIMPLEMENTED;
+	TEST_EQ(vb2_diagnostic_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"diagnostic screen: check disabled item");
+	DISPLAYED_EQ("diagnostic menu: self-test disabled",
+		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, MOCK_IGNORE, 0xc, 0x0,
+		     MOCK_IGNORE);
+
 	VB2_DEBUG("...done.\n");
 }