diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index 44b85ea..5c460a2 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -61,28 +61,6 @@
  * to a correct menu item index.
  */
 
-static vb2_error_t log_page_init(struct vb2_ui_context *ui)
-{
-	const struct vb2_screen_info *screen = ui->state->screen;
-
-	ui->state->current_page = 0;
-
-	/* TODO(b/174127808): Split out enabling/disabling buttons. */
-	if (ui->state->page_count == 1) {
-		VB2_SET_BIT(ui->state->disabled_item_mask,
-			    screen->page_up_item);
-		VB2_SET_BIT(ui->state->disabled_item_mask,
-			    screen->page_down_item);
-		ui->state->selected_item = screen->back_item;
-	} else {
-		VB2_SET_BIT(ui->state->disabled_item_mask,
-			    screen->page_up_item);
-		ui->state->selected_item = screen->page_down_item;
-	}
-
-	return VB2_SUCCESS;
-}
-
 static vb2_error_t log_page_update(struct vb2_ui_context *ui,
 				   const char *new_log_string)
 {
@@ -91,8 +69,10 @@
 	if (new_log_string) {
 		ui->state->page_count = vb2ex_prepare_log_screen(
 			screen->id, ui->locale_id, new_log_string);
-		if (ui->state->page_count == 0)
+		if (ui->state->page_count == 0) {
+			VB2_DEBUG("vb2ex_prepare_log_screen failed");
 			return VB2_ERROR_UI_LOG_INIT;
+		}
 		if (ui->state->current_page >= ui->state->page_count)
 			ui->state->current_page = ui->state->page_count - 1;
 		ui->force_display = 1;
@@ -109,6 +89,36 @@
 	return VB2_SUCCESS;
 }
 
+static vb2_error_t log_page_reset_to_top(struct vb2_ui_context *ui)
+{
+	const struct vb2_screen_info *screen = ui->state->screen;
+
+	ui->state->current_page = 0;
+	ui->state->selected_item = ui->state->page_count > 1
+					   ? screen->page_down_item
+					   : screen->back_item;
+	return log_page_update(ui, NULL);
+}
+
+static vb2_error_t log_page_show_back_or_cancel(struct vb2_ui_context *ui,
+						int is_show_cancel)
+{
+	int back_item = ui->state->screen->back_item;
+	int cancel_item = ui->state->screen->cancel_item;
+	VB2_CLR_BIT(ui->state->hidden_item_mask, back_item);
+	VB2_CLR_BIT(ui->state->hidden_item_mask, cancel_item);
+	if (is_show_cancel) {
+		VB2_SET_BIT(ui->state->hidden_item_mask, back_item);
+		if (ui->state->selected_item == back_item)
+			ui->state->selected_item = cancel_item;
+	} else {
+		VB2_SET_BIT(ui->state->hidden_item_mask, cancel_item);
+		if (ui->state->selected_item == cancel_item)
+			ui->state->selected_item = back_item;
+	}
+	return VB2_SUCCESS;
+}
+
 static vb2_error_t log_page_prev_action(struct vb2_ui_context *ui)
 {
 	/* Validity check. */
@@ -290,7 +300,7 @@
 #define DEBUG_INFO_ITEM_PAGE_DOWN 2
 #define DEBUG_INFO_ITEM_BACK 3
 
-static vb2_error_t debug_info_init(struct vb2_ui_context *ui)
+static vb2_error_t debug_info_set_content(struct vb2_ui_context *ui)
 {
 	const char *log_string = vb2ex_get_debug_info(ui->ctx);
 	if (!log_string) {
@@ -298,34 +308,23 @@
 		ui->error_code = VB2_UI_ERROR_DEBUG_LOG;
 		return vb2_ui_screen_back(ui);
 	}
-	ui->state->page_count = vb2ex_prepare_log_screen(
-		ui->state->screen->id, ui->locale_id, log_string);
-	if (ui->state->page_count == 0) {
-		VB2_DEBUG("ERROR: Failed to prepare debug info screen\n");
+	vb2_error_t rv = log_page_update(ui, log_string);
+	if (rv) {
 		ui->error_code = VB2_UI_ERROR_DEBUG_LOG;
 		return vb2_ui_screen_back(ui);
 	}
+	return VB2_SUCCESS;
+}
 
-	return log_page_init(ui);
+static vb2_error_t debug_info_init(struct vb2_ui_context *ui)
+{
+	VB2_TRY(debug_info_set_content(ui));
+	return log_page_reset_to_top(ui);
 }
 
 static vb2_error_t debug_info_reinit(struct vb2_ui_context *ui)
 {
-	const char *log_string = vb2ex_get_debug_info(ui->ctx);
-	if (!log_string) {
-		VB2_DEBUG("ERROR: Failed to retrieve debug info\n");
-		ui->error_code = VB2_UI_ERROR_DEBUG_LOG;
-		return vb2_ui_screen_back(ui);
-	}
-	ui->state->page_count = vb2ex_prepare_log_screen(
-		ui->state->screen->id, ui->locale_id, log_string);
-	if (ui->state->page_count == 0) {
-		VB2_DEBUG("ERROR: Failed to prepare debug info screen\n");
-		ui->error_code = VB2_UI_ERROR_DEBUG_LOG;
-		return vb2_ui_screen_back(ui);
-	}
-
-	return VB2_SUCCESS;
+	return debug_info_set_content(ui);
 }
 
 static const struct vb2_menu_item debug_info_items[] = {
@@ -354,42 +353,32 @@
 #define FIRMWARE_LOG_ITEM_PAGE_DOWN 2
 #define FIRMWARE_LOG_ITEM_BACK 3
 
-static vb2_error_t firmware_log_init(struct vb2_ui_context *ui)
+static vb2_error_t firmware_log_set_content(struct vb2_ui_context *ui,
+					    int reset)
 {
-	const char *log_string = vb2ex_get_firmware_log(1);
+	const char *log_string = vb2ex_get_firmware_log(reset);
 	if (!log_string) {
 		VB2_DEBUG("ERROR: Failed to retrieve firmware log\n");
 		ui->error_code = VB2_UI_ERROR_FIRMWARE_LOG;
 		return vb2_ui_screen_back(ui);
 	}
-	ui->state->page_count = vb2ex_prepare_log_screen(
-		ui->state->screen->id, ui->locale_id, log_string);
-	if (ui->state->page_count == 0) {
-		VB2_DEBUG("ERROR: Failed to prepare firmware log screen\n");
+	vb2_error_t rv = log_page_update(ui, log_string);
+	if (rv) {
 		ui->error_code = VB2_UI_ERROR_FIRMWARE_LOG;
 		return vb2_ui_screen_back(ui);
 	}
+	return VB2_SUCCESS;
+}
 
-	return log_page_init(ui);
+static vb2_error_t firmware_log_init(struct vb2_ui_context *ui)
+{
+	VB2_TRY(firmware_log_set_content(ui, 1));
+	return log_page_reset_to_top(ui);
 }
 
 static vb2_error_t firmware_log_reinit(struct vb2_ui_context *ui)
 {
-	const char *log_string = vb2ex_get_firmware_log(0);
-	if (!log_string) {
-		VB2_DEBUG("ERROR: Failed to retrieve firmware log\n");
-		ui->error_code = VB2_UI_ERROR_FIRMWARE_LOG;
-		return vb2_ui_screen_back(ui);
-	}
-	ui->state->page_count = vb2ex_prepare_log_screen(
-		ui->state->screen->id, ui->locale_id, log_string);
-	if (ui->state->page_count == 0) {
-		VB2_DEBUG("ERROR: Failed to prepare firmware log screen\n");
-		ui->error_code = VB2_UI_ERROR_FIRMWARE_LOG;
-		return vb2_ui_screen_back(ui);
-	}
-
-	return VB2_SUCCESS;
+	return firmware_log_set_content(ui, 0);
 }
 
 static const struct vb2_menu_item firmware_log_items[] = {
@@ -1079,20 +1068,16 @@
 {
 	const char *log_string;
 	vb2_error_t rv = vb2ex_diag_get_storage_health(&log_string);
-	if (rv != VB2_SUCCESS) {
-		VB2_DEBUG("ERROR: Failed to retrieve storage log message\n");
+	if (rv) {
 		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
 		return vb2_ui_screen_back(ui);
 	}
-
-	ui->state->page_count = vb2ex_prepare_log_screen(
-		ui->state->screen->id, ui->locale_id, log_string);
-	if (ui->state->page_count == 0) {
-		VB2_DEBUG("ERROR: Failed to prepare storage log screen\n");
+	rv = log_page_update(ui, log_string);
+	if (rv) {
 		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
 		return vb2_ui_screen_back(ui);
 	}
-	return log_page_init(ui);
+	return log_page_reset_to_top(ui);
 }
 
 static const struct vb2_menu_item diagnostics_storage_health_items[] = {
@@ -1118,65 +1103,49 @@
 
 #define DIAGNOSTICS_MEMORY_ITEM_PAGE_UP 0
 #define DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN 1
-#define DIAGNOSTICS_MEMORY_ITEM_CANCEL 2
-#define DIAGNOSTICS_MEMORY_ITEM_BACK 3
+#define DIAGNOSTICS_MEMORY_ITEM_BACK 2
+#define DIAGNOSTICS_MEMORY_ITEM_CANCEL 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)
+static vb2_error_t diagnostics_memory_update_screen_impl(
+	struct vb2_ui_context *ui, memory_test_op_t op, int reset)
 {
 	const char *log_string = NULL;
+	vb2_error_t rv;
+	int is_test_running = 0;
 
 	/* Early return if the memory test is done. */
 	if (ui->state->test_finished)
 		return VB2_SUCCESS;
 
-	vb2_error_t rv = op(reset, &log_string);
-
+	rv = op(reset, &log_string);
+	switch (rv) {
 	/* The test is still running but the output buffer was unchanged. */
-	if (rv == VB2_ERROR_EX_DIAG_TEST_RUNNING)
+	case VB2_ERROR_EX_DIAG_TEST_RUNNING:
 		return VB2_SUCCESS;
-
-	if ((rv && rv != VB2_ERROR_EX_DIAG_TEST_UPDATED) || !log_string) {
-		VB2_DEBUG("ERROR: Failed to retrieve memory test status\n");
-		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
-		return vb2_ui_screen_back(ui);
-	}
-
-	ui->state->page_count = vb2ex_prepare_log_screen(
-		ui->state->screen->id, ui->locale_id, 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;
-		return vb2_ui_screen_back(ui);
-	}
-	/* TODO(b/174127808): Integrate this into a new log_page_* function. */
-	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. */
-	VB2_CLR_BIT(ui->state->hidden_item_mask,
-		    DIAGNOSTICS_MEMORY_ITEM_CANCEL);
-	VB2_CLR_BIT(ui->state->hidden_item_mask, DIAGNOSTICS_MEMORY_ITEM_BACK);
-	if (rv == VB2_ERROR_EX_DIAG_TEST_UPDATED) {
-		VB2_SET_BIT(ui->state->hidden_item_mask,
-			    DIAGNOSTICS_MEMORY_ITEM_BACK);
-		if (ui->state->selected_item == DIAGNOSTICS_MEMORY_ITEM_BACK)
-			ui->state->selected_item =
-				DIAGNOSTICS_MEMORY_ITEM_CANCEL;
-	} else {
-		VB2_SET_BIT(ui->state->hidden_item_mask,
-			    DIAGNOSTICS_MEMORY_ITEM_CANCEL);
-		if (ui->state->selected_item == DIAGNOSTICS_MEMORY_ITEM_CANCEL)
-			ui->state->selected_item = DIAGNOSTICS_MEMORY_ITEM_BACK;
+	case VB2_ERROR_EX_DIAG_TEST_UPDATED:
+		is_test_running = 1;
+		break;
+	case VB2_SUCCESS:
 		ui->state->test_finished = 1;
+		break;
+	default:
+		VB2_DEBUG("memory_test_op 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_memory_update_screen(struct vb2_ui_context *ui,
+						    memory_test_op_t op,
+						    int reset)
+{
+	vb2_error_t rv = diagnostics_memory_update_screen_impl(ui, op, reset);
+	if (rv) {
+		ui->error_code = VB2_UI_ERROR_DIAGNOSTICS;
+		return vb2_ui_screen_back(ui);
+	}
 	return VB2_SUCCESS;
 }
 
@@ -1184,14 +1153,14 @@
 {
 	VB2_TRY(diagnostics_memory_update_screen(
 		ui, &vb2ex_diag_memory_quick_test, 1));
-	return log_page_init(ui);
+	return log_page_reset_to_top(ui);
 }
 
 static vb2_error_t diagnostics_memory_init_full(struct vb2_ui_context *ui)
 {
 	VB2_TRY(diagnostics_memory_update_screen(
 		ui, &vb2ex_diag_memory_full_test, 1));
-	return log_page_init(ui);
+	return log_page_reset_to_top(ui);
 }
 
 static vb2_error_t diagnostics_memory_update_quick(struct vb2_ui_context *ui)
@@ -1209,11 +1178,11 @@
 static const struct vb2_menu_item diagnostics_memory_items[] = {
 	[DIAGNOSTICS_MEMORY_ITEM_PAGE_UP] = PAGE_UP_ITEM,
 	[DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN] = PAGE_DOWN_ITEM,
+	[DIAGNOSTICS_MEMORY_ITEM_BACK] = BACK_ITEM,
 	[DIAGNOSTICS_MEMORY_ITEM_CANCEL] = {
 		.text = "Cancel and go back",
 		.action = vb2_ui_screen_back,
 	},
-	[DIAGNOSTICS_MEMORY_ITEM_BACK] = BACK_ITEM,
 	POWER_OFF_ITEM,
 };
 
@@ -1225,7 +1194,8 @@
 	.menu = MENU_ITEMS(diagnostics_memory_items),
 	.page_up_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_UP,
 	.page_down_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
+	.back_item = DIAGNOSTICS_MEMORY_ITEM_BACK,
+	.cancel_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
 };
 
 static const struct vb2_screen_info diagnostics_memory_full_screen = {
@@ -1236,7 +1206,8 @@
 	.menu = MENU_ITEMS(diagnostics_memory_items),
 	.page_up_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_UP,
 	.page_down_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
+	.back_item = DIAGNOSTICS_MEMORY_ITEM_BACK,
+	.cancel_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
 };
 
 /******************************************************************************/
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 7e81541..017bd44 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -66,6 +66,7 @@
 	uint32_t page_up_item;
 	uint32_t page_down_item;
 	uint32_t back_item;
+	uint32_t cancel_item;
 };
 
 struct vb2_screen_state {