1 files changed, 79 insertions, 43 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 4d363f68..81243ac2 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -95,11 +95,11 @@ const struct vb2_menu *get_menu(struct vb2_ui_context *ui)
 		.num_items = 0,
 		.items = NULL,
 	};
-	if (ui->state.screen->get_menu) {
-		menu = ui->state.screen->get_menu(ui);
+	if (ui->state->screen->get_menu) {
+		menu = ui->state->screen->get_menu(ui);
 		return menu ? menu : &empty_menu;
 	} else {
-		return &ui->state.screen->menu;
+		return &ui->state->screen->menu;
 	}
 }
 
@@ -125,7 +125,7 @@ vb2_error_t menu_navigation_action(struct vb2_ui_context *ui)
 	case VB_KEY_ENTER:
 		return vb2_ui_menu_select(ui);
 	case VB_KEY_ESC:
-		return vb2_ui_change_root(ui);
+		return vb2_ui_screen_back(ui);
 	default:
 		if (key != 0)
 			VB2_DEBUG("Pressed key %#x, trusted? %d\n",
@@ -142,13 +142,13 @@ vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui)
 	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_UP_SHORT_PRESS)
 		return VB2_REQUEST_UI_CONTINUE;
 
-	item = ui->state.selected_item - 1;
+	item = ui->state->selected_item - 1;
 	while (item >= 0 &&
-	       ((1 << item) & ui->state.disabled_item_mask))
+	       ((1 << item) & ui->state->disabled_item_mask))
 		item--;
 	/* Only update if item is valid */
 	if (item >= 0)
-		ui->state.selected_item = item;
+		ui->state->selected_item = item;
 
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -162,13 +162,13 @@ vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
 		return VB2_REQUEST_UI_CONTINUE;
 
 	menu = get_menu(ui);
-	item = ui->state.selected_item + 1;
+	item = ui->state->selected_item + 1;
 	while (item < menu->num_items &&
-	       ((1 << item) & ui->state.disabled_item_mask))
+	       ((1 << item) & ui->state->disabled_item_mask))
 		item++;
 	/* Only update if item is valid */
 	if (item < menu->num_items)
-		ui->state.selected_item = item;
+		ui->state->selected_item = item;
 
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -185,7 +185,7 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 	if (menu->num_items == 0)
 		return VB2_REQUEST_UI_CONTINUE;
 
-	menu_item = &menu->items[ui->state.selected_item];
+	menu_item = &menu->items[ui->state->selected_item];
 
 	if (menu_item->action) {
 		VB2_DEBUG("Menu item <%s> run action\n", menu_item->text);
@@ -193,7 +193,7 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 	} else if (menu_item->target) {
 		VB2_DEBUG("Menu item <%s> to target screen %#x\n",
 			  menu_item->text, menu_item->target);
-		return vb2_ui_change_screen(ui, menu_item->target);
+		return vb2_ui_screen_change(ui, menu_item->target);
 	}
 
 	VB2_DEBUG("Menu item <%s> no action or target screen\n",
@@ -204,28 +204,35 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 /*****************************************************************************/
 /* Screen navigation functions */
 
-vb2_error_t vb2_ui_change_root(struct vb2_ui_context *ui)
+vb2_error_t vb2_ui_screen_back(struct vb2_ui_context *ui)
 {
-	return vb2_ui_change_screen(ui, ui->root_screen->id);
+	struct vb2_screen_state *tmp;
+
+	if (ui->state && ui->state->prev) {
+		tmp = ui->state->prev;
+		free(ui->state);
+		ui->state = tmp;
+	} else {
+		VB2_DEBUG("ERROR: No previous screen; ignoring\n");
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
 }
 
 static vb2_error_t default_screen_init(struct vb2_ui_context *ui)
 {
 	const struct vb2_menu *menu = get_menu(ui);
-	ui->state.selected_item = 0;
+	ui->state->selected_item = 0;
 	if (menu->num_items > 1 && menu->items[0].is_language_select)
-		ui->state.selected_item = 1;
+		ui->state->selected_item = 1;
 	return VB2_REQUEST_UI_CONTINUE;
 }
 
-vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
+vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id)
 {
 	const struct vb2_screen_info *new_screen_info;
-
-	if (ui->state.screen && ui->state.screen->id == id) {
-		VB2_DEBUG("WARNING: Already on screen %#x; ignoring\n", id);
-		return VB2_REQUEST_UI_CONTINUE;
-	}
+	struct vb2_screen_state *cur_state;
+	int state_exists = 0;
 
 	new_screen_info = vb2_get_screen_info(id);
 	if (new_screen_info == NULL) {
@@ -233,13 +240,41 @@ vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 		return VB2_REQUEST_UI_CONTINUE;
 	}
 
-	memset(&ui->state, 0, sizeof(ui->state));
-	ui->state.screen = new_screen_info;
+	/* Check to see if the screen state already exists in our stack. */
+	cur_state = ui->state;
+	while (cur_state != NULL) {
+		if (cur_state->screen->id == id) {
+			state_exists = 1;
+			break;
+		}
+		cur_state = cur_state->prev;
+	}
+
+	if (state_exists) {
+		/* Pop until the requested screen is at the top of stack. */
+		while (ui->state->screen->id != id) {
+			cur_state = ui->state;
+			ui->state = cur_state->prev;
+			free(cur_state);
+		}
+	} else {
+		/* Allocate the requested screen on top of the stack. */
+		cur_state = malloc(sizeof(*ui->state));
+		memset(cur_state, 0, sizeof(*ui->state));
+		if (cur_state == NULL) {
+			VB2_DEBUG("WARNING: malloc failed; ignoring\n");
+			return VB2_REQUEST_UI_CONTINUE;
+		}
+		cur_state->prev = ui->state;
+		cur_state->screen = new_screen_info;
+		ui->state = cur_state;
+		if (ui->state->screen->init)
+			return ui->state->screen->init(ui);
+		else
+			return default_screen_init(ui);
+	}
 
-	if (ui->state.screen->init)
-		return ui->state.screen->init(ui);
-	else
-		return default_screen_init(ui);
+	return VB2_REQUEST_UI_CONTINUE;
 }
 
 /*****************************************************************************/
@@ -252,16 +287,17 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 	struct vb2_screen_state prev_state;
 	enum vb2_ui_error prev_error_code;
 	const struct vb2_menu *menu;
+	const struct vb2_screen_info *root_info;
 	uint32_t key_flags;
 	vb2_error_t rv;
 
 	memset(&ui, 0, sizeof(ui));
 	ui.ctx = ctx;
-	ui.root_screen = vb2_get_screen_info(root_screen_id);
-	if (ui.root_screen == NULL)
+	root_info = vb2_get_screen_info(root_screen_id);
+	if (root_info == NULL)
 		VB2_DIE("Root screen not found.\n");
 	ui.locale_id = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
-	rv = vb2_ui_change_screen(&ui, ui.root_screen->id);
+	rv = vb2_ui_screen_change(&ui, root_screen_id);
 	if (rv != VB2_REQUEST_UI_CONTINUE)
 		return rv;
 	memset(&prev_state, 0, sizeof(prev_state));
@@ -269,19 +305,19 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 
 	while (1) {
 		/* Draw if there are state changes. */
-		if (memcmp(&prev_state, &ui.state, sizeof(ui.state)) ||
+		if (memcmp(&prev_state, ui.state, sizeof(*ui.state)) ||
 		    /* we want to redraw/beep on a transition */
 		    prev_error_code != ui.error_code) {
 
 			menu = get_menu(&ui);
 			VB2_DEBUG("<%s> menu item <%s>\n",
-				  ui.state.screen->name,
+				  ui.state->screen->name,
 				  menu->num_items ?
-				  menu->items[ui.state.selected_item].text :
+				  menu->items[ui.state->selected_item].text :
 				  "null");
-			vb2ex_display_ui(ui.state.screen->id, ui.locale_id,
-					 ui.state.selected_item,
-					 ui.state.disabled_item_mask,
+			vb2ex_display_ui(ui.state->screen->id, ui.locale_id,
+					 ui.state->selected_item,
+					 ui.state->disabled_item_mask,
 					 ui.error_code);
 			/*
 			 * Only beep if we're transitioning from no
@@ -292,7 +328,7 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 				vb2ex_beep(250, 400);
 
 			/* Update prev variables. */
-			memcpy(&prev_state, &ui.state, sizeof(ui.state));
+			memcpy(&prev_state, ui.state, sizeof(*ui.state));
 			prev_error_code = ui.error_code;
 		}
 
@@ -313,8 +349,8 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 			return rv;
 
 		/* Run screen action. */
-		if (ui.state.screen->action) {
-			rv = ui.state.screen->action(&ui);
+		if (ui.state->screen->action) {
+			rv = ui.state->screen->action(&ui);
 			if (rv != VB2_REQUEST_UI_CONTINUE)
 				return rv;
 		}
@@ -350,7 +386,7 @@ vb2_error_t developer_action(struct vb2_ui_context *ui)
 {
 	/* Developer mode keyboard shortcuts */
 	if (ui->key == VB_KEY_CTRL('S'))
-		return vb2_ui_change_screen(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
+		return vb2_ui_screen_change(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
 	if (ui->key == VB_KEY_CTRL('U') ||
 	    (DETACHABLE && ui->key == VB_BUTTON_VOL_UP_LONG_PRESS))
 		return vb2_ui_developer_mode_boot_external_action(ui);
@@ -402,7 +438,7 @@ vb2_error_t manual_recovery_action(struct vb2_ui_context *ui)
 		VB2_DEBUG("Recovery VbTryLoadKernel %#x --> %#x\n",
 			  ui->recovery_rv, rv);
 		ui->recovery_rv = rv;
-		return vb2_ui_change_screen(ui,
+		return vb2_ui_screen_change(ui,
 					    rv == VB2_ERROR_LK_NO_DISK_FOUND ?
 					    VB2_SCREEN_RECOVERY_SELECT :
 					    VB2_SCREEN_RECOVERY_INVALID);
@@ -411,7 +447,7 @@ vb2_error_t manual_recovery_action(struct vb2_ui_context *ui)
 	/* Manual recovery keyboard shortcuts */
 	if (ui->key == VB_KEY_CTRL('D') ||
 	    (DETACHABLE && ui->key == VB_BUTTON_VOL_UP_DOWN_COMBO_PRESS))
-		return vb2_ui_change_screen(ui, VB2_SCREEN_RECOVERY_TO_DEV);
+		return vb2_ui_screen_change(ui, VB2_SCREEN_RECOVERY_TO_DEV);
 
 	/* TODO(b/144969088): Re-implement as debug info screen. */
 	if (ui->key == '\t')