minidiag: Add storage and memory diagnostic screens

BRANCH=none
BUG=b:156692539, b:156693348
TEST=emerge-hatch vboot_reference
TEST=unittest passed:
( export CC=x86_64-pc-linux-gnu-clang DEBUG=1 MENU_UI=1 DIAGNOSTIC_UI=1
    MINIMAL=1 TPM2_MODE= MOCK_TPM=; make clean &&
            make -j32 test_setup && make runtests; echo $? )

Cq-Depend: chromium:2322286, chromium:2328704, chromium:2336239
Cq-Depend: chromium:2361823, chromium:2361582
Signed-off-by: Meng-Huan Yu <menghuan@chromium.org>
Change-Id: I8b875b09bd5bcdb65f08c11945b046d2b3c3a113
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2372022
Reviewed-by: Joel Kitching <kitching@chromium.org>

5 files changed, 252 insertions, 6 deletions

diff --git a/firmware/2lib/2stub.c b/firmware/2lib/2stub.c
index 90041cbc..697b0fc5 100644
--- a/firmware/2lib/2stub.c
+++ b/firmware/2lib/2stub.c
@@ -91,3 +91,23 @@ uint32_t vb2ex_prepare_log_screen(const char *str)
 {
 	return 1;
 }
+
+__attribute__((weak))
+const char *vb2ex_get_diagnostic_storage(void)
+{
+	return "mock";
+}
+
+__attribute__((weak))
+vb2_error_t vb2ex_diag_memory_quick_test(int reset, const char **out)
+{
+	*out = "mock";
+	return VB2_SUCCESS;
+}
+
+__attribute__((weak))
+vb2_error_t vb2ex_diag_memory_full_test(int reset, const char **out)
+{
+	*out = "mock";
+	return VB2_SUCCESS;
+}
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index e0a6d599..2e5dbfaf 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -866,12 +866,15 @@ static const struct vb2_menu_item diagnostics_items[] = {
 	LANGUAGE_SELECT_ITEM,
 	{
 		.text = "Storage",
+		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE,
 	},
 	{
 		.text = "Quick memory check",
+		.target = VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK,
 	},
 	{
 		.text = "Full memory check",
+		.target = VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL,
 	},
 	POWER_OFF_ITEM,
 };
@@ -883,6 +886,165 @@ static const struct vb2_screen_info diagnostics_screen = {
 };
 
 /******************************************************************************/
+/* VB2_SCREEN_DIAGNOSTICS_STORAGE */
+
+#define DIAGNOSTICS_STORAGE_ITEM_PAGE_DOWN 1
+#define DIAGNOSTICS_STORAGE_ITEM_BACK 2
+
+static vb2_error_t diagnostics_storage_init(struct vb2_ui_context *ui)
+{
+	const char *log_string = vb2ex_get_diagnostic_storage();
+	if (!log_string) {
+		VB2_DEBUG("ERROR: Failed to retrieve storage log message\n");
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		ui->error_beep = 1;
+		return vb2_ui_screen_back(ui);
+	}
+
+	ui->state->page_count = vb2ex_prepare_log_screen(log_string);
+	if (ui->state->page_count == 0) {
+		VB2_DEBUG("ERROR: Failed to prepare storage log screen\n");
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		ui->error_beep = 1;
+		return vb2_ui_screen_back(ui);
+	}
+	return log_page_init(ui, DIAGNOSTICS_STORAGE_ITEM_PAGE_DOWN,
+			     DIAGNOSTICS_STORAGE_ITEM_BACK);
+}
+
+static const struct vb2_menu_item diagnostics_storage_items[] = {
+	{
+		.text = "Page up",
+		.action = log_page_prev_action,
+	},
+	[DIAGNOSTICS_STORAGE_ITEM_PAGE_DOWN] = {
+		.text = "Page down",
+		.action = log_page_next_action,
+	},
+	[DIAGNOSTICS_STORAGE_ITEM_BACK] = BACK_ITEM,
+	POWER_OFF_ITEM,
+};
+
+static const struct vb2_screen_info diagnostics_storage_screen = {
+	.id = VB2_SCREEN_DIAGNOSTICS_STORAGE,
+	.name = "Storage",
+	.init = diagnostics_storage_init,
+	.menu = MENU_ITEMS(diagnostics_storage_items),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK
+   VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL */
+
+#define DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN 1
+#define DIAGNOSTICS_MEMORY_ITEM_CANCEL 2
+#define DIAGNOSTICS_MEMORY_ITEM_BACK 3
+
+typedef vb2_error_t (*memory_test_op_t)(int reset, const char **out);
+static vb2_error_t diagnostics_memory_update_screen(struct vb2_ui_context *ui,
+						    memory_test_op_t op,
+						    int reset)
+{
+	const char *log_string = NULL;
+	vb2_error_t rv = op(reset, &log_string);
+	if ((rv && rv != VB2_ERROR_EX_DIAG_TEST_RUNNING) || !log_string) {
+		VB2_DEBUG("ERROR: Failed to retrieve memory test status\n");
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		ui->error_beep = 1;
+		return vb2_ui_screen_back(ui);
+	}
+
+	ui->state->page_count = vb2ex_prepare_log_screen(log_string);
+	if (ui->state->page_count == 0) {
+		VB2_DEBUG("ERROR: Failed to prepare memory log screen, error: "
+			  "%#x\n", rv);
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		ui->error_beep = 1;
+		return vb2_ui_screen_back(ui);
+	}
+	if (ui->state->current_page >= ui->state->page_count)
+		ui->state->current_page = ui->state->page_count - 1;
+
+	ui->force_display = 1;
+
+	/* Show cancel button when the test is running, otherwise show the back
+	 * button. VB2_SUCCESS indicates the test is finished. */
+	ui->state->disabled_item_mask &= ~(1 << DIAGNOSTICS_MEMORY_ITEM_CANCEL);
+	ui->state->disabled_item_mask &= ~(1 << DIAGNOSTICS_MEMORY_ITEM_BACK);
+	if (rv == VB2_ERROR_EX_DIAG_TEST_RUNNING) {
+		ui->state->disabled_item_mask |=
+			1 << DIAGNOSTICS_MEMORY_ITEM_BACK;
+		if (ui->state->selected_item == DIAGNOSTICS_MEMORY_ITEM_BACK)
+			ui->state->selected_item =
+				DIAGNOSTICS_MEMORY_ITEM_CANCEL;
+	} else {
+		ui->state->disabled_item_mask |=
+			1 << DIAGNOSTICS_MEMORY_ITEM_CANCEL;
+		if (ui->state->selected_item == DIAGNOSTICS_MEMORY_ITEM_CANCEL)
+			ui->state->selected_item = DIAGNOSTICS_MEMORY_ITEM_BACK;
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static vb2_error_t diagnostics_memory_init_quick(struct vb2_ui_context *ui)
+{
+	return diagnostics_memory_update_screen(
+		ui, &vb2ex_diag_memory_quick_test, 1);
+}
+
+static vb2_error_t diagnostics_memory_init_full(struct vb2_ui_context *ui)
+{
+	return diagnostics_memory_update_screen(
+		ui, &vb2ex_diag_memory_full_test, 1);
+}
+
+static vb2_error_t diagnostics_memory_update_quick(struct vb2_ui_context *ui)
+{
+	return diagnostics_memory_update_screen(
+		ui, &vb2ex_diag_memory_quick_test, 0);
+}
+
+static vb2_error_t diagnostics_memory_update_full(struct vb2_ui_context *ui)
+{
+	return diagnostics_memory_update_screen(
+		ui, &vb2ex_diag_memory_full_test, 0);
+}
+
+static const struct vb2_menu_item diagnostics_memory_items[] = {
+	{
+		.text = "Page up",
+		.action = log_page_prev_action,
+	},
+	[DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN] = {
+		.text = "Page down",
+		.action = log_page_next_action,
+	},
+	[DIAGNOSTICS_MEMORY_ITEM_CANCEL] = {
+		.text = "Cancel and go back",
+		.action = vb2_ui_screen_back,
+	},
+	[DIAGNOSTICS_MEMORY_ITEM_BACK] = BACK_ITEM,
+	POWER_OFF_ITEM,
+};
+
+static const struct vb2_screen_info diagnostics_memory_quick_screen = {
+	.id = VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK,
+	.name = "Quick memory check",
+	.init = diagnostics_memory_init_quick,
+	.action = diagnostics_memory_update_quick,
+	.menu = MENU_ITEMS(diagnostics_memory_items),
+};
+
+static const struct vb2_screen_info diagnostics_memory_full_screen = {
+	.id = VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL,
+	.name = "Full memory check",
+	.init = diagnostics_memory_init_full,
+	.action = diagnostics_memory_update_full,
+	.menu = MENU_ITEMS(diagnostics_memory_items),
+};
+
+/******************************************************************************/
 /*
  * TODO(chromium:1035800): Refactor UI code across vboot and depthcharge.
  * Currently vboot and depthcharge maintain their own copies of menus/screens.
@@ -910,6 +1072,9 @@ static const struct vb2_screen_info *screens[] = {
 	&developer_boot_external_screen,
 	&developer_invalid_disk_screen,
 	&diagnostics_screen,
+	&diagnostics_storage_screen,
+	&diagnostics_memory_quick_screen,
+	&diagnostics_memory_full_screen,
 };
 
 const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen id)
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index a6c76764..ee3bb53e 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -1318,6 +1318,11 @@ enum vb2_screen {
 	VB2_SCREEN_DEVELOPER_INVALID_DISK	= 0x330,
 	/* Diagnostic tools */
 	VB2_SCREEN_DIAGNOSTICS			= 0x400,
+	/* Storage diagnostic screen */
+	VB2_SCREEN_DIAGNOSTICS_STORAGE	      	= 0x410,
+	/* Memory diagnostic screens */
+	VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK    	= 0x420,
+	VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL     	= 0x421,
 };
 
 enum vb2_ui_error {
@@ -1331,6 +1336,8 @@ enum vb2_ui_error {
 	VB2_UI_ERROR_FIRMWARE_LOG,
 	/* Untrusted confirmation */
 	VB2_UI_ERROR_UNTRUSTED_CONFIRMATION,
+	/* Diagnostics internal failure */
+	VB2_UI_ERROR_DIAGNOSTICS,
 };
 
 /**
@@ -1442,6 +1449,37 @@ const char *vb2ex_get_firmware_log(void);
  */
 uint32_t vb2ex_prepare_log_screen(const char *str);
 
+/**
+ * Get the full storage diagnostic log.
+ *
+ * Return a pointer of full log string which is guaranteed to be
+ * null-terminated.  The function implementation should manage string memory
+ * internally.  Subsequent calls may update the string and/or may return a new
+ * pointer.
+ *
+ * @return The pointer to the full debug info string.  NULL on error.
+ */
+const char *vb2ex_get_diagnostic_storage(void);
+
+/**
+ * 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
+ * test state.
+ * *
+ * @param reset	Discard the current memory test result and re-initialize
+ *		a new test.
+ * @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 memory 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. Other non-zero codes for internal errors.
+ */
+vb2_error_t vb2ex_diag_memory_quick_test(int reset, const char **out);
+vb2_error_t vb2ex_diag_memory_full_test(int reset, const char **out);
+
 /*****************************************************************************/
 /* Timer. */
 
diff --git a/firmware/2lib/include/2return_codes.h b/firmware/2lib/include/2return_codes.h
index 336c11e0..3a50f949 100644
--- a/firmware/2lib/include/2return_codes.h
+++ b/firmware/2lib/include/2return_codes.h
@@ -731,6 +731,12 @@ enum vb2_return_code {
 	/* Error setting vendor data (see: VbExSetVendorData). */
 	VB2_ERROR_EX_SET_VENDOR_DATA,
 
+	/* The memory test is running. */
+	VB2_ERROR_EX_DIAG_TEST_RUNNING,
+
+	/* The memory test initialization failed. */
+	VB2_ERROR_EX_DIAG_TEST_INIT_FAILED,
+
 	/**********************************************************************
 	 * LoadKernel errors
 	 *
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index d3b371c9..569a179c 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -1566,14 +1566,20 @@ static void diagnostics_screen_tests(void)
 	/* #0: Language menu */
 	add_mock_keypress(VB_KEY_UP);
 	add_mock_keypress(VB_KEY_ENTER);
-	/* #1: Storage (no-op) */
+	/* #1: Storage screen */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	/* #2: Quick memory test (no-op) */
+	add_mock_keypress(VB_KEY_ENTER);
+	/* #2: Quick memory test screen */
+	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	/* #3: Full memory test (no-op) */
+	add_mock_keypress(VB_KEY_ENTER);
+	/* #3: Full memory test screen */
+	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_ENTER);
 	/* #4: Power off (End of menu) */
+	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	mock_calls_until_shutdown = -1;
@@ -1587,17 +1593,28 @@ static void diagnostics_screen_tests(void)
 		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 0, 0x0, MOCK_IGNORE);
 	DISPLAYED_EQ("#0: language menu", VB2_SCREEN_LANGUAGE_SELECT,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	/* #1: Storage (no-op) */
+	/* #1: Storage screen */
 	DISPLAYED_PASS();
 	DISPLAYED_EQ("storage button",
 		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 1, 0x0, MOCK_IGNORE);
-	/* #2: Quick memory test (no-op) */
+	DISPLAYED_EQ("#1: storage screen", VB2_SCREEN_DIAGNOSTICS_STORAGE,
+		     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, 0x0, MOCK_IGNORE);
-	/* #3: Full memory test (no-op) */
+	DISPLAYED_EQ("#1: quick memory test screen",
+		     VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	/* #3: Full memory test screen */
+	DISPLAYED_PASS();
 	DISPLAYED_EQ("full memory test button",
 		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 3, 0x0, MOCK_IGNORE);
+	DISPLAYED_EQ("#3: full memory test screen",
+		     VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL, 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, 0x0, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();