vboot/ui: Add reinit screen hook

Re-initialize the log screen when going back from another 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
     when the firmware log screen is implemented,
     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>
(cherry picked from commit 2cb872fa3056158c35f82a9412e667b69ccf5a49)
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2332200
Reviewed-by: Joel Kitching <kitching@chromium.org>
diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 1bdc273..99cedb0 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -214,6 +214,8 @@
 		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 @@
 			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 e8f5916..a2cd60f 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -266,12 +266,25 @@
 		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 = 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,
@@ -304,6 +317,7 @@
 	.id = VB2_SCREEN_DEBUG_INFO,
 	.name = "Debug info",
 	.init = debug_info_init,
+	.reinit = debug_info_reinit,
 	.menu = MENU_ITEMS(debug_info_items),
 };
 
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 5f2434c..867b0c8 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -39,8 +39,17 @@
 	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. */