vboot/ui: Fix the back button in log screen

Re-initialize the log screen when going back from another
log screen.

BRANCH=puff
BUG=b:146399181, b:146105976
TEST=CC=x86_64-pc-linux-gnu-clang;
    make clean && make runtests
TEST=USE="menu_ui" emerge-puff depthcharge
     select "advanced options",
     navigate to firmware log screen,
     press <TAB> to debug info,
     and select back to firmware log screen

Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
Change-Id: Ie44805e3bfb14a4a8b660a18a123a184a4c5ea45
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2329224
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>

3 files changed, 52 insertions, 3 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 8adfab80..3a3308b1 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -214,6 +214,8 @@ vb2_error_t vb2_ui_screen_back(struct vb2_ui_context *ui)
 		tmp = ui->state->prev;
 		free(ui->state);
 		ui->state = tmp;
+		if (ui->state->screen->reinit)
+			return ui->state->screen->reinit(ui);
 	} else {
 		VB2_DEBUG("ERROR: No previous screen; ignoring\n");
 	}
@@ -259,6 +261,8 @@ vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id)
 			ui->state = cur_state->prev;
 			free(cur_state);
 		}
+		if (ui->state->screen->reinit)
+			return ui->state->screen->reinit(ui);
 	} else {
 		/* Allocate the requested screen on top of the stack. */
 		cur_state = malloc(sizeof(*ui->state));
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index 653c279a..1b548cff 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -264,18 +264,35 @@ static const struct vb2_screen_info advanced_options_screen = {
 
 static vb2_error_t debug_info_init(struct vb2_ui_context *ui)
 {
-	const char *log_string = vb2ex_get_debug_info(ui->ctx);
+	const char *log_string;
+
+	log_string = vb2ex_get_debug_info(ui->ctx);
 	ui->state->page_count = vb2ex_prepare_log_screen(log_string);
 	if (ui->state->page_count == 0) {
 		ui->error_code = VB2_UI_ERROR_DEBUG_LOG;
 		return vb2_ui_screen_back(ui);
 	}
+
 	return log_page_init(ui,
 			     DEBUG_INFO_ITEM_PAGE_UP,
 			     DEBUG_INFO_ITEM_PAGE_DOWN,
 			     DEBUG_INFO_ITEM_BACK);
 }
 
+static vb2_error_t debug_info_reinit(struct vb2_ui_context *ui)
+{
+	const char *log_string;
+
+	log_string = vb2ex_get_debug_info(ui->ctx);
+	ui->state->page_count = vb2ex_prepare_log_screen(log_string);
+	if (ui->state->page_count == 0) {
+		ui->error_code = VB2_UI_ERROR_DEBUG_LOG;
+		return vb2_ui_screen_back(ui);
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
 static vb2_error_t debug_info_page_prev_action(struct vb2_ui_context *ui)
 {
 	return log_page_prev(ui,
@@ -308,6 +325,7 @@ static const struct vb2_screen_info debug_info_screen = {
 	.id = VB2_SCREEN_DEBUG_INFO,
 	.name = "Debug info",
 	.init = debug_info_init,
+	.reinit = debug_info_reinit,
 	.menu = MENU_ITEMS(debug_info_items),
 };
 
@@ -320,18 +338,35 @@ static const struct vb2_screen_info debug_info_screen = {
 
 static vb2_error_t firmware_log_init(struct vb2_ui_context *ui)
 {
-	const char *log_string = vb2ex_get_firmware_log();
+	const char *log_string;
+
+	log_string = vb2ex_get_firmware_log();
 	ui->state->page_count = vb2ex_prepare_log_screen(log_string);
 	if (ui->state->page_count == 0) {
 		ui->error_code = VB2_UI_ERROR_FIRMWARE_LOG;
 		return vb2_ui_screen_back(ui);
 	}
+
 	return log_page_init(ui,
 			     FIRMWARE_LOG_ITEM_PAGE_UP,
 			     FIRMWARE_LOG_ITEM_PAGE_DOWN,
 			     FIRMWARE_LOG_ITEM_BACK);
 }
 
+static vb2_error_t firmware_log_reinit(struct vb2_ui_context *ui)
+{
+	const char *log_string;
+
+	log_string = vb2ex_get_firmware_log();
+	ui->state->page_count = vb2ex_prepare_log_screen(log_string);
+	if (ui->state->page_count == 0) {
+		ui->error_code = VB2_UI_ERROR_FIRMWARE_LOG;
+		return vb2_ui_screen_back(ui);
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
 static vb2_error_t firmware_log_page_prev_action(struct vb2_ui_context *ui)
 {
 	return log_page_prev(ui,
@@ -364,6 +399,7 @@ static const struct vb2_screen_info firmware_log_screen = {
 	.id = VB2_SCREEN_FIRMWARE_LOG,
 	.name = "Firmware log",
 	.init = firmware_log_init,
+	.reinit = firmware_log_reinit,
 	.menu = MENU_ITEMS(firmware_log_items),
 };
 
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 5f2434c2..867b0c81 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -39,8 +39,17 @@ struct vb2_screen_info {
 	enum vb2_screen id;
 	/* Screen name for printing to console only */
 	const char *name;
-	/* Init function runs once when changing to the screen. */
+	/*
+	 * Init function runs once when changing to the screen which is not in
+	 * the history stack.
+	 */
 	vb2_error_t (*init)(struct vb2_ui_context *ui);
+	/*
+	 * Re-init function runs once when changing to the screen which is
+	 * already in the history stack, for example, when going back to the
+	 * screen. Exactly one of init() and reinit() will be called.
+	 */
+	vb2_error_t (*reinit)(struct vb2_ui_context *ui);
 	/* Action function runs repeatedly while on the screen. */
 	vb2_error_t (*action)(struct vb2_ui_context *ui);
 	/* Menu items. */