vboot: Pass vb2 context and use vboot2 NV routines

Passing the vb2 context around allows using more of the vb2 functions in
future changes, and prepares for a future where we directly use the
context as it was set up in firmware verification.

BUG=chromium:611535
BRANCH=none
TEST=make runtests; emerge-kevin coreboot depthcharge

Change-Id: I8efa606dbdec5d195b66eb899e76fdc84337ad36
Signed-off-by: Randall Spangler <rspangler@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/404997
Reviewed-by: Shelley Chen <shchen@chromium.org>

7 files changed, 321 insertions, 256 deletions

diff --git a/firmware/lib/ec_sync.c b/firmware/lib/ec_sync.c
index e989bbda..a32c296f 100644
--- a/firmware/lib/ec_sync.c
+++ b/firmware/lib/ec_sync.c
@@ -7,6 +7,7 @@
 
 #include "2sysincludes.h"
 #include "2common.h"
+#include "2nvstorage.h"
 #include "sysincludes.h"
 #include "utility.h"
 #include "vb2_common.h"
@@ -14,20 +15,19 @@
 #include "vboot_common.h"
 #include "vboot_display.h"
 #include "vboot_kernel.h"
-#include "vboot_nvstorage.h"
 
-static void request_recovery(VbNvContext *vnc, uint32_t recovery_request)
+static void request_recovery(struct vb2_context *ctx, uint32_t recovery_request)
 {
-	VB2_DEBUG("request_recovery(%d)\n", (int)recovery_request);
+	VB2_DEBUG("request_recovery(%u)\n", recovery_request);
 
-	VbNvSet(vnc, VBNV_RECOVERY_REQUEST, recovery_request);
+	vb2_nv_set(ctx, VB2_NV_RECOVERY_REQUEST, recovery_request);
 }
 
 /**
  * Wrapper around VbExEcProtect() which sets recovery reason on error.
  */
-static VbError_t EcProtect(int devidx, enum VbSelectFirmware_t select,
-			   VbNvContext *vnc)
+static VbError_t EcProtect(struct vb2_context *ctx, int devidx,
+			   enum VbSelectFirmware_t select)
 {
 	int rv = VbExEcProtect(devidx, select);
 
@@ -35,14 +35,15 @@ static VbError_t EcProtect(int devidx, enum VbSelectFirmware_t select,
 		VBDEBUG(("VbExEcProtect() needs reboot\n"));
 	} else if (rv != VBERROR_SUCCESS) {
 		VBDEBUG(("VbExEcProtect() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_PROTECT);
+		request_recovery(ctx, VB2_RECOVERY_EC_PROTECT);
 	}
 	return rv;
 }
 
-static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
+static VbError_t EcUpdateImage(struct vb2_context *ctx, int devidx,
+			       VbCommonParams *cparams,
 			       enum VbSelectFirmware_t select,
-			       int *need_update, int in_rw, VbNvContext *vnc)
+			       int *need_update, int in_rw)
 {
 	VbSharedDataHeader *shared =
 		(VbSharedDataHeader *)cparams->shared_data_blob;
@@ -65,7 +66,7 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 	if (rv) {
 		VBDEBUG(("EcUpdateImage() - "
 			 "VbExEcHashImage() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_HASH_FAILED);
+		request_recovery(ctx, VB2_RECOVERY_EC_HASH_FAILED);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 	VBDEBUG(("EC-%s hash: ", rw_request ? "RW" : "RO"));
@@ -78,14 +79,14 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 	if (rv) {
 		VBDEBUG(("EcUpdateImage() - "
 			 "VbExEcGetExpectedImageHash() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_EXPECTED_HASH);
+		request_recovery(ctx, VB2_RECOVERY_EC_EXPECTED_HASH);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 	if (ec_hash_size != hash_size) {
 		VBDEBUG(("EcUpdateImage() - "
 			 "EC uses %d-byte hash, but AP-RW contains %d bytes\n",
 			 ec_hash_size, hash_size));
-		request_recovery(vnc, VBNV_RECOVERY_EC_HASH_SIZE);
+		request_recovery(ctx, VB2_RECOVERY_EC_HASH_SIZE);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 
@@ -103,7 +104,7 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 	if (rv) {
 		VBDEBUG(("EcUpdateImage() - "
 			 "VbExEcGetExpectedImage() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_EXPECTED_IMAGE);
+		request_recovery(ctx, VB2_RECOVERY_EC_EXPECTED_IMAGE);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 	VBDEBUG(("EcUpdateImage() - image len = %d\n", expected_size));
@@ -118,7 +119,7 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 		    !(shared->flags & VBSD_OPROM_LOADED)) {
 			VBDEBUG(("EcUpdateImage() - Reboot to "
 				 "load VGA Option ROM\n"));
-			VbNvSet(vnc, VBNV_OPROM_NEEDED, 1);
+			vb2_nv_set(ctx, VB2_NV_OPROM_NEEDED, 1);
 		}
 
 		/*
@@ -141,11 +142,11 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 		    !(shared->flags & VBSD_OPROM_LOADED)) {
 			VBDEBUG(("EcUpdateImage() - Reboot to "
 				 "load VGA Option ROM\n"));
-			VbNvSet(vnc, VBNV_OPROM_NEEDED, 1);
+			vb2_nv_set(ctx, VB2_NV_OPROM_NEEDED, 1);
 			return VBERROR_VGA_OPROM_MISMATCH;
 		}
 
-		VbDisplayScreen(cparams, VB_SCREEN_WAIT, 0, vnc);
+		VbDisplayScreen(ctx, cparams, VB_SCREEN_WAIT, 0);
 	}
 
 	rv = VbExEcUpdateImage(devidx, select, expected, expected_size);
@@ -163,7 +164,7 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 		 * mode.
 		 */
 		if (rv != VBERROR_EC_REBOOT_TO_RO_REQUIRED)
-			request_recovery(vnc, VBNV_RECOVERY_EC_UPDATE);
+			request_recovery(ctx, VB2_RECOVERY_EC_UPDATE);
 
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
@@ -173,14 +174,14 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 	if (rv) {
 		VBDEBUG(("EcUpdateImage() - "
 			 "VbExEcHashImage() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_HASH_FAILED);
+		request_recovery(ctx, VB2_RECOVERY_EC_HASH_FAILED);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 	if (hash_size != ec_hash_size) {
 		VBDEBUG(("EcUpdateImage() - "
 			 "VbExEcHashImage() says size %d, not %d\n",
 			 ec_hash_size, hash_size));
-		request_recovery(vnc, VBNV_RECOVERY_EC_HASH_SIZE);
+		request_recovery(ctx, VB2_RECOVERY_EC_HASH_SIZE);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 	VBDEBUG(("Updated EC-%s hash: ", rw_request ? "RW" : "RO"));
@@ -192,15 +193,15 @@ static VbError_t EcUpdateImage(int devidx, VbCommonParams *cparams,
 		VBDEBUG(("EcUpdateImage() - "
 			 "Failed to update EC-%s\n", rw_request ?
 			 "RW" : "RO"));
-		request_recovery(vnc, VBNV_RECOVERY_EC_UPDATE);
+		request_recovery(ctx, VB2_RECOVERY_EC_UPDATE);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 
 	return VBERROR_SUCCESS;
 }
 
-VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
-			   VbNvContext *vnc)
+VbError_t VbEcSoftwareSync(struct vb2_context *ctx, int devidx,
+			   VbCommonParams *cparams)
 {
 	VbSharedDataHeader *shared =
 		(VbSharedDataHeader *)cparams->shared_data_blob;
@@ -211,7 +212,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 	int in_rw = 0;
 	int ro_try_count = 2;
 	int num_tries = 0;
-	uint32_t try_ro_sync, recovery_request;
+	uint32_t recovery_request;
 	int rv, updated_rw, updated_ro;
 
 	VBDEBUG(("VbEcSoftwareSync(devidx=%d)\n", devidx));
@@ -234,7 +235,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 			 */
 			VBDEBUG(("VbEcSoftwareSync() - "
 				 "want recovery but got EC-RW\n"));
-			request_recovery(vnc, shared->recovery_reason);
+			request_recovery(ctx, shared->recovery_reason);
 			return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 		}
 
@@ -249,7 +250,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 	if (rv != VBERROR_SUCCESS) {
 		VBDEBUG(("VbEcSoftwareSync() - "
 			 "VbExEcRunningRW() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_UNKNOWN_IMAGE);
+		request_recovery(ctx, VB2_RECOVERY_EC_UNKNOWN_IMAGE);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 
@@ -263,7 +264,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 		}
 
 		/* Protect the RW flash and stay in EC-RO */
-		rv = EcProtect(devidx, select_rw, vnc);
+		rv = EcProtect(ctx, devidx, select_rw);
 		if (rv != VBERROR_SUCCESS)
 			return rv;
 
@@ -271,7 +272,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 		if (rv != VBERROR_SUCCESS) {
 			VBDEBUG(("VbEcSoftwareSync() - "
 				 "VbExEcDisableJump() returned %d\n", rv));
-			request_recovery(vnc, VBNV_RECOVERY_EC_SOFTWARE_SYNC);
+			request_recovery(ctx, VB2_RECOVERY_EC_SOFTWARE_SYNC);
 			return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 		}
 
@@ -282,7 +283,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 	VBDEBUG(("VbEcSoftwareSync() check for RW update.\n"));
 
 	/* Update the RW Image. */
-	rv = EcUpdateImage(devidx, cparams, select_rw, &updated_rw, in_rw, vnc);
+	rv = EcUpdateImage(ctx, devidx, cparams, select_rw, &updated_rw, in_rw);
 
 	if (rv != VBERROR_SUCCESS) {
 		VBDEBUG(("VbEcSoftwareSync() - "
@@ -306,35 +307,34 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 			 * All other errors trigger recovery mode.
 			 */
 			if (rv != VBERROR_EC_REBOOT_TO_RO_REQUIRED)
-				request_recovery(vnc, VBNV_RECOVERY_EC_JUMP_RW);
+				request_recovery(ctx, VB2_RECOVERY_EC_JUMP_RW);
 
 			return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 		}
 	}
 
-	VbNvGet(vnc, VBNV_TRY_RO_SYNC, &try_ro_sync);
-
+	uint32_t try_ro_sync = vb2_nv_get(ctx, VB2_NV_TRY_RO_SYNC);
 	if (!devidx && try_ro_sync &&
 	    !(shared->flags & VBSD_BOOT_FIRMWARE_WP_ENABLED)) {
 		/* Reset RO Software Sync NV flag */
-		VbNvSet(vnc, VBNV_TRY_RO_SYNC, 0);
+		vb2_nv_set(ctx, VB2_NV_TRY_RO_SYNC, 0);
 
-		VbNvGet(vnc, VBNV_RECOVERY_REQUEST, &recovery_request);
+		recovery_request = vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST);
 
 		/* Update the RO Image. */
 		while (num_tries < ro_try_count) {
 			VBDEBUG(("VbEcSoftwareSync() RO Software Sync\n"));
 
 			/* Get expected EC-RO Image. */
-			rv = EcUpdateImage(devidx, cparams, select_ro,
-					   &updated_ro, in_rw, vnc);
+			rv = EcUpdateImage(ctx, devidx, cparams, select_ro,
+					   &updated_ro, in_rw);
 			if (rv == VBERROR_SUCCESS) {
 				/*
 				 * If the RO update had failed, reset the
 				 * recovery request.
 				 */
 				if (num_tries)
-					request_recovery(vnc, recovery_request);
+					request_recovery(ctx, recovery_request);
 				break;
 			} else
 				VBDEBUG(("VbEcSoftwareSync() - "
@@ -347,12 +347,12 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 		return rv;
 
 	/* Protect RO flash */
-	rv = EcProtect(devidx, select_ro, vnc);
+	rv = EcProtect(ctx, devidx, select_ro);
 	if (rv != VBERROR_SUCCESS)
 		return rv;
 
 	/* Protect RW flash */
-	rv = EcProtect(devidx, select_rw, vnc);
+	rv = EcProtect(ctx, devidx, select_rw);
 	if (rv != VBERROR_SUCCESS)
 		return rv;
 
@@ -360,7 +360,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 	if (rv != VBERROR_SUCCESS) {
 		VBDEBUG(("VbEcSoftwareSync() - "
 			"VbExEcDisableJump() returned %d\n", rv));
-		request_recovery(vnc, VBNV_RECOVERY_EC_SOFTWARE_SYNC);
+		request_recovery(ctx, VB2_RECOVERY_EC_SOFTWARE_SYNC);
 		return VBERROR_EC_REBOOT_TO_RO_REQUIRED;
 	}
 
@@ -378,7 +378,7 @@ VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
 	    (shared->flags & VBSD_OPROM_LOADED)) {
 		VBDEBUG(("VbEcSoftwareSync() - Reboot to "
 			 "unload VGA Option ROM\n"));
-		VbNvSet(vnc, VBNV_OPROM_NEEDED, 0);
+		vb2_nv_set(ctx, VB2_NV_OPROM_NEEDED, 0);
 		return VBERROR_VGA_OPROM_MISMATCH;
 	}
 
diff --git a/firmware/lib/include/load_kernel_fw.h b/firmware/lib/include/load_kernel_fw.h
index 957cee13..aeecb448 100644
--- a/firmware/lib/include/load_kernel_fw.h
+++ b/firmware/lib/include/load_kernel_fw.h
@@ -12,6 +12,8 @@
 #include "vboot_api.h"
 #include "vboot_nvstorage.h"
 
+struct vb2_context;
+
 /* Interface provided by verified boot library to BDS */
 
 /* Boot flags for LoadKernel().boot_flags */
@@ -71,9 +73,14 @@ typedef struct LoadKernelParams {
 /**
  * Attempt to load the kernel from the current device.
  *
+ * @param ctx		Vboot context
+ * @param params	Params specific to loading the kernel
+ * @param cparams	Common parameters to vboot1 APIs
+ *
  * Returns VBERROR_SUCCESS if successful.  If unsuccessful, sets a recovery
  * reason via VbNvStorage and returns an error code.
  */
-VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams);
+VbError_t LoadKernel(struct vb2_context *ctx, LoadKernelParams *params,
+		     VbCommonParams *cparams);
 
 #endif  /* VBOOT_REFERENCE_LOAD_KERNEL_FW_H_ */
diff --git a/firmware/lib/include/vboot_display.h b/firmware/lib/include/vboot_display.h
index 0ab93f0f..3ec4eae6 100644
--- a/firmware/lib/include/vboot_display.h
+++ b/firmware/lib/include/vboot_display.h
@@ -12,13 +12,16 @@
 #include "vboot_api.h"
 #include "vboot_nvstorage.h"
 
-VbError_t VbDisplayScreenFromGBB(VbCommonParams *cparams, uint32_t screen,
-                                 VbNvContext *vncptr, uint32_t locale);
-VbError_t VbDisplayScreen(VbCommonParams *cparams, uint32_t screen, int force,
-                          VbNvContext *vncptr);
-VbError_t VbDisplayDebugInfo(VbCommonParams *cparams, VbNvContext *vncptr);
-VbError_t VbCheckDisplayKey(VbCommonParams *cparams, uint32_t key,
-                            VbNvContext *vncptr);
+struct vb2_context;
+
+VbError_t VbDisplayScreenFromGBB(struct vb2_context *ctx,
+				 VbCommonParams *cparams, uint32_t screen,
+                                 uint32_t locale);
+VbError_t VbDisplayScreen(struct vb2_context *ctx, VbCommonParams *cparams,
+			  uint32_t screen, int force);
+VbError_t VbDisplayDebugInfo(struct vb2_context *ctx, VbCommonParams *cparams);
+VbError_t VbCheckDisplayKey(struct vb2_context *ctx, VbCommonParams *cparams,
+			    uint32_t key);
 
 /* Internal functions, for unit testing */
 
diff --git a/firmware/lib/include/vboot_kernel.h b/firmware/lib/include/vboot_kernel.h
index c7134af3..b6a2a858 100644
--- a/firmware/lib/include/vboot_kernel.h
+++ b/firmware/lib/include/vboot_kernel.h
@@ -14,10 +14,7 @@
 #include "load_kernel_fw.h"
 #include "vboot_api.h"
 
-/**
- * Accessors for unit tests only.
- */
-VbNvContext *VbApiKernelGetVnc(void);
+struct vb2_context;
 
 /**
  * Exported for unit tests only - frees memory used by VbSelectAndLoadKernel()
@@ -25,9 +22,20 @@ VbNvContext *VbApiKernelGetVnc(void);
 void VbApiKernelFree(VbCommonParams *cparams);
 
 /**
- * Try to load a kernel.
+ * Attempt loading a kernel from the specified type(s) of disks.
+ *
+ * If successful, sets p->disk_handle to the disk for the kernel and returns
+ * VBERROR_SUCCESS.
+ *
+ * @param ctx			Vboot context
+ * @param cparams		Vboot common params
+ * @param p			Parameters for loading kernel
+ * @param get_info_flags	Flags to pass to VbExDiskGetInfo()
+ * @return VBERROR_SUCCESS, VBERROR_NO_DISK_FOUND if no disks of the specified
+ * type were found, or other non-zero VBERROR_ codes for other failures.
  */
-uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
+uint32_t VbTryLoadKernel(struct vb2_context *ctx, VbCommonParams *cparams,
+			 LoadKernelParams *p,
                          uint32_t get_info_flags);
 
 /* Flags for VbUserConfirms() */
@@ -48,31 +56,35 @@ uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
  *
  * Returns: 1=yes, 0=no, -1 = shutdown.
  */
-int VbUserConfirms(VbCommonParams *cparams, uint32_t confirm_flags);
+int VbUserConfirms(struct vb2_context *ctx, VbCommonParams *cparams,
+		   uint32_t confirm_flags);
 
 /**
  * Handle a normal boot.
  */
-VbError_t VbBootNormal(VbCommonParams *cparams, LoadKernelParams *p);
+VbError_t VbBootNormal(struct vb2_context *ctx, VbCommonParams *cparams,
+		       LoadKernelParams *p);
 
 /**
  * Handle a developer-mode boot.
  */
-VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p);
+VbError_t VbBootDeveloper(struct vb2_context *ctx, VbCommonParams *cparams,
+			  LoadKernelParams *p);
 
 /**
  * Handle a recovery-mode boot.
  */
-VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p);
+VbError_t VbBootRecovery(struct vb2_context *ctx, VbCommonParams *cparams,
+			 LoadKernelParams *p);
 
 /**
  * Sync EC device <devidx> firmware to expected version.
  *
+ * @param ctx		Vboot context
  * @param devidx	EC device index to sync
  * @param cparams	Common vboot params
- * @param vnc		NV storage context
  */
-VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
-			   VbNvContext *vnc);
+VbError_t VbEcSoftwareSync(struct vb2_context *ctx, int devidx,
+			   VbCommonParams *cparams);
 
 #endif  /* VBOOT_REFERENCE_VBOOT_KERNEL_H_ */
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 7cf8d8a4..597dbc18 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -9,6 +9,8 @@
 
 #include "2sysincludes.h"
 #include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
 #include "2rsa.h"
 #include "gbb_access.h"
 #include "gbb_header.h"
@@ -31,11 +33,6 @@ static struct RollbackSpaceFwmp fwmp;
 #ifdef CHROMEOS_ENVIRONMENT
 /* Global variable accessors for unit tests */
 
-VbNvContext *VbApiKernelGetVnc(void)
-{
-	return &vnc;
-}
-
 struct RollbackSpaceFwmp *VbApiKernelGetFwmp(void)
 {
 	return &fwmp;
@@ -45,16 +42,18 @@ struct RollbackSpaceFwmp *VbApiKernelGetFwmp(void)
 /**
  * Set recovery request (called from vboot_api_kernel.c functions only)
  */
-static void VbSetRecoveryRequest(uint32_t recovery_request)
+static void VbSetRecoveryRequest(struct vb2_context *ctx,
+				 uint32_t recovery_request)
 {
 	VBDEBUG(("VbSetRecoveryRequest(%d)\n", (int)recovery_request));
-	VbNvSet(&vnc, VBNV_RECOVERY_REQUEST, recovery_request);
+	vb2_nv_set(ctx, VB2_NV_RECOVERY_REQUEST, recovery_request);
 }
 
-static void VbSetRecoverySubcode(uint32_t recovery_request)
+static void VbSetRecoverySubcode(struct vb2_context *ctx,
+				 uint32_t recovery_request)
 {
 	VBDEBUG(("VbSetRecoverySubcode(%d)\n", (int)recovery_request));
-	VbNvSet(&vnc, VBNV_RECOVERY_SUBCODE, recovery_request);
+	vb2_nv_set(ctx, VB2_NV_RECOVERY_SUBCODE, recovery_request);
 }
 
 static void VbNvLoad(void)
@@ -70,10 +69,10 @@ static void VbNvCommit(void)
 		VbExNvStorageWrite(vnc.raw);
 }
 
-static void VbAllowUsbBoot(void)
+static void VbAllowUsbBoot(struct vb2_context *ctx)
 {
 	VBDEBUG(("%s\n", __func__));
-	VbNvSet(&vnc, VBNV_DEV_BOOT_USB, 1);
+	vb2_nv_set(ctx, VB2_NV_DEV_BOOT_USB, 1);
 }
 
 /**
@@ -114,11 +113,15 @@ static void VbTryLegacy(int allowed)
  * If successful, sets p->disk_handle to the disk for the kernel and returns
  * VBERROR_SUCCESS.
  *
- * Returns VBERROR_NO_DISK_FOUND if no disks of the specified type were found.
- *
- * May return other VBERROR_ codes for other failures.
+ * @param ctx			Vboot context
+ * @param cparams		Vboot common params
+ * @param p			Parameters for loading kernel
+ * @param get_info_flags	Flags to pass to VbExDiskGetInfo()
+ * @return VBERROR_SUCCESS, VBERROR_NO_DISK_FOUND if no disks of the specified
+ * type were found, or other non-zero VBERROR_ codes for other failures.
  */
-uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
+uint32_t VbTryLoadKernel(struct vb2_context *ctx, VbCommonParams *cparams,
+			 LoadKernelParams *p,
                          uint32_t get_info_flags)
 {
 	VbError_t retval = VBERROR_UNKNOWN;
@@ -138,7 +141,7 @@ uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
 
 	VBDEBUG(("VbTryLoadKernel() found %d disks\n", (int)disk_count));
 	if (0 == disk_count) {
-		VbSetRecoveryRequest(VBNV_RECOVERY_RW_NO_DISK);
+		VbSetRecoveryRequest(ctx, VBNV_RECOVERY_RW_NO_DISK);
 		return VBERROR_NO_DISK_FOUND;
 	}
 
@@ -155,7 +158,8 @@ uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
 		 */
 		if (512 != disk_info[i].bytes_per_lba ||
 		    16 > disk_info[i].lba_count ||
-		    get_info_flags != (disk_info[i].flags & ~VB_DISK_FLAG_EXTERNAL_GPT)) {
+		    get_info_flags != (disk_info[i].flags &
+				       ~VB_DISK_FLAG_EXTERNAL_GPT)) {
 			VBDEBUG(("  skipping: bytes_per_lba=%" PRIu64
 				 " lba_count=%" PRIu64 " flags=0x%x\n",
 				 disk_info[i].bytes_per_lba,
@@ -170,7 +174,7 @@ uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
 						?: p->gpt_lba_count;
 		p->boot_flags |= disk_info[i].flags & VB_DISK_FLAG_EXTERNAL_GPT
 				? BOOT_FLAG_EXTERNAL_GPT : 0;
-		retval = LoadKernel(p, cparams);
+		retval = LoadKernel(ctx, p, cparams);
 		VBDEBUG(("VbTryLoadKernel() LoadKernel() = %d\n", retval));
 
 		/*
@@ -186,7 +190,7 @@ uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
 
 	/* If we didn't find any good kernels, don't return a disk handle. */
 	if (VBERROR_SUCCESS != retval) {
-		VbSetRecoveryRequest(VBNV_RECOVERY_RW_NO_KERNEL);
+		VbSetRecoveryRequest(ctx, VBNV_RECOVERY_RW_NO_KERNEL);
 		p->disk_handle = NULL;
 	}
 
@@ -199,9 +203,11 @@ uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
 	return retval;
 }
 
-uint32_t VbTryUsb(VbCommonParams *cparams, LoadKernelParams *p)
+uint32_t VbTryUsb(struct vb2_context *ctx, VbCommonParams *cparams,
+		  LoadKernelParams *p)
 {
-	uint32_t retval = VbTryLoadKernel(cparams, p, VB_DISK_FLAG_REMOVABLE);
+	uint32_t retval = VbTryLoadKernel(ctx, cparams, p,
+					  VB_DISK_FLAG_REMOVABLE);
 	if (VBERROR_SUCCESS == retval) {
 		VBDEBUG(("VbBootDeveloper() - booting USB\n"));
 	} else {
@@ -214,15 +220,15 @@ uint32_t VbTryUsb(VbCommonParams *cparams, LoadKernelParams *p)
 		 * at this point doesn't put us into
 		 * recovery mode.
 		 */
-		VbSetRecoveryRequest(
-			VBNV_RECOVERY_NOT_REQUESTED);
+		VbSetRecoveryRequest(ctx, VBNV_RECOVERY_NOT_REQUESTED);
 	}
 	return retval;
 }
 
 #define CONFIRM_KEY_DELAY 20  /* Check confirm screen keys every 20ms */
 
-int VbUserConfirms(VbCommonParams *cparams, uint32_t confirm_flags)
+int VbUserConfirms(struct vb2_context *ctx, VbCommonParams *cparams,
+		   uint32_t confirm_flags)
 {
 	VbSharedDataHeader *shared =
            (VbSharedDataHeader *)cparams->shared_data_blob;
@@ -279,7 +285,7 @@ int VbUserConfirms(VbCommonParams *cparams, uint32_t confirm_flags)
 					return 1;
 				}
 			}
-			VbCheckDisplayKey(cparams, key, &vnc);
+			VbCheckDisplayKey(ctx, cparams, key);
 		}
 		VbExSleepMs(CONFIRM_KEY_DELAY);
 	}
@@ -289,11 +295,12 @@ int VbUserConfirms(VbCommonParams *cparams, uint32_t confirm_flags)
 }
 
 VbError_t test_mockable
-VbBootNormal(VbCommonParams *cparams, LoadKernelParams *p)
+VbBootNormal(struct vb2_context *ctx, VbCommonParams *cparams,
+	     LoadKernelParams *p)
 {
 	/* Boot from fixed disk only */
 	VBDEBUG(("Entering %s()\n", __func__));
-	return VbTryLoadKernel(cparams, p, VB_DISK_FLAG_FIXED);
+	return VbTryLoadKernel(ctx, cparams, p, VB_DISK_FLAG_FIXED);
 }
 
 static const char dev_disable_msg[] =
@@ -301,18 +308,16 @@ static const char dev_disable_msg[] =
 	"For more information, see http://dev.chromium.org/chromium-os/fwmp\n"
 	"\n";
 
-VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
+VbError_t VbBootDeveloper(struct vb2_context *ctx, VbCommonParams *cparams,
+			  LoadKernelParams *p)
 {
 	GoogleBinaryBlockHeader *gbb = cparams->gbb;
 	VbSharedDataHeader *shared =
 		(VbSharedDataHeader *)cparams->shared_data_blob;
 
-	uint32_t allow_usb = 0;
-	uint32_t allow_legacy = 0;
 	uint32_t disable_dev_boot = 0;
 	uint32_t use_usb = 0;
 	uint32_t use_legacy = 0;
-	uint32_t default_boot = 0;
 	uint32_t ctrl_d_pressed = 0;
 
 	VbAudioContext *audio = 0;
@@ -320,11 +325,11 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 	VBDEBUG(("Entering %s()\n", __func__));
 
 	/* Check if USB booting is allowed */
-	VbNvGet(&vnc, VBNV_DEV_BOOT_USB, &allow_usb);
-	VbNvGet(&vnc, VBNV_DEV_BOOT_LEGACY, &allow_legacy);
+	uint32_t allow_usb = vb2_nv_get(ctx, VB2_NV_DEV_BOOT_USB);
+	uint32_t allow_legacy = vb2_nv_get(ctx, VB2_NV_DEV_BOOT_LEGACY);
 
 	/* Check if the default is to boot using disk, usb, or legacy */
-	VbNvGet(&vnc, VBNV_DEV_DEFAULT_BOOT, &default_boot);
+	uint32_t default_boot = vb2_nv_get(ctx, VB2_NV_DEV_DEFAULT_BOOT);
 
 	if(default_boot == VBNV_DEV_DEFAULT_BOOT_USB)
 		use_usb = 1;
@@ -359,17 +364,17 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 	/* If dev mode is disabled, only allow TONORM */
 	while (disable_dev_boot) {
 		VBDEBUG(("%s() - dev_disable_boot is set.\n", __func__));
-		VbDisplayScreen(cparams, VB_SCREEN_DEVELOPER_TO_NORM, 0, &vnc);
+		VbDisplayScreen(ctx, cparams, VB_SCREEN_DEVELOPER_TO_NORM, 0);
 		VbExDisplayDebugInfo(dev_disable_msg);
 
 		/* Ignore space in VbUserConfirms()... */
-		switch (VbUserConfirms(cparams, 0)) {
+		switch (VbUserConfirms(ctx, cparams, 0)) {
 		case 1:
 			VBDEBUG(("%s() - leaving dev-mode.\n", __func__));
-			VbNvSet(&vnc, VBNV_DISABLE_DEV_REQUEST, 1);
-			VbDisplayScreen(cparams,
+			vb2_nv_set(ctx, VB2_NV_DISABLE_DEV_REQUEST, 1);
+			VbDisplayScreen(ctx, cparams,
 					VB_SCREEN_TO_NORM_CONFIRMED,
-					0, &vnc);
+					0);
 			VbExSleepMs(5000);
 			return VBERROR_REBOOT_REQUIRED;
 		case -1:
@@ -382,7 +387,7 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 	}
 
 	/* Show the dev mode warning screen */
-	VbDisplayScreen(cparams, VB_SCREEN_DEVELOPER_WARNING, 0, &vnc);
+	VbDisplayScreen(ctx, cparams, VB_SCREEN_DEVELOPER_WARNING, 0);
 
 	/* Get audio/delay context */
 	audio = VbAudioOpen(cparams);
@@ -428,20 +433,20 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 					VbExBeep(120, 400);
 					break;
 				}
-				VbDisplayScreen(cparams,
+				VbDisplayScreen(ctx, cparams,
 						VB_SCREEN_DEVELOPER_TO_NORM,
-						0, &vnc);
+						0);
 				/* Ignore space in VbUserConfirms()... */
-				switch (VbUserConfirms(cparams, 0)) {
+				switch (VbUserConfirms(ctx, cparams, 0)) {
 				case 1:
 					VBDEBUG(("%s() - leaving dev-mode.\n",
 						 __func__));
-					VbNvSet(&vnc, VBNV_DISABLE_DEV_REQUEST,
+					vb2_nv_set(ctx, VB2_NV_DISABLE_DEV_REQUEST,
 						1);
-					VbDisplayScreen(
+					VbDisplayScreen(ctx,
 						cparams,
 						VB_SCREEN_TO_NORM_CONFIRMED,
-						0, &vnc);
+						0);
 					VbExSleepMs(5000);
 					return VBERROR_REBOOT_REQUIRED;
 				case -1:
@@ -452,10 +457,10 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 					/* Stay in dev-mode */
 					VBDEBUG(("%s() - stay in dev-mode\n",
 						 __func__));
-					VbDisplayScreen(
+					VbDisplayScreen(ctx,
 						cparams,
 						VB_SCREEN_DEVELOPER_WARNING,
-						0, &vnc);
+						0);
 					/* Start new countdown */
 					audio = VbAudioOpen(cparams);
 				}
@@ -466,7 +471,7 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 				 */
 				VBDEBUG(("%s() - going to recovery\n",
 					 __func__));
-				VbSetRecoveryRequest(
+				VbSetRecoveryRequest(ctx,
 					VBNV_RECOVERY_RW_DEV_SCREEN);
 				VbAudioClose(audio);
 				return VBERROR_LOAD_KERNEL_RECOVERY;
@@ -510,23 +515,24 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 				 * Clear the screen to show we get the Ctrl+U
 				 * key press.
 				 */
-				VbDisplayScreen(cparams, VB_SCREEN_BLANK, 0,
-						&vnc);
-				if (VBERROR_SUCCESS == VbTryUsb(cparams, p)) {
+				VbDisplayScreen(ctx, cparams, VB_SCREEN_BLANK,
+						0);
+				if (VBERROR_SUCCESS ==
+				    VbTryUsb(ctx, cparams, p)) {
 					VbAudioClose(audio);
 					return VBERROR_SUCCESS;
 				} else {
 					/* Show dev mode warning screen again */
-					VbDisplayScreen(
+					VbDisplayScreen(ctx,
 						cparams,
 						VB_SCREEN_DEVELOPER_WARNING,
-						0, &vnc);
+						0);
 				}
 			}
 			break;
 		default:
 			VBDEBUG(("VbBootDeveloper() - pressed key %d\n", key));
-			VbCheckDisplayKey(cparams, key, &vnc);
+			VbCheckDisplayKey(ctx, cparams, key);
 			break;
 		}
 	} while(VbAudioLooping(audio));
@@ -540,7 +546,7 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 	}
 
 	if ((use_usb && !ctrl_d_pressed) && allow_usb) {
-		if (VBERROR_SUCCESS == VbTryUsb(cparams, p)) {
+		if (VBERROR_SUCCESS == VbTryUsb(ctx, cparams, p)) {
 			VbAudioClose(audio);
 			return VBERROR_SUCCESS;
 		}
@@ -549,7 +555,7 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 	/* Timeout or Ctrl+D; attempt loading from fixed disk */
 	VBDEBUG(("VbBootDeveloper() - trying fixed disk\n"));
 	VbAudioClose(audio);
-	return VbTryLoadKernel(cparams, p, VB_DISK_FLAG_FIXED);
+	return VbTryLoadKernel(ctx, cparams, p, VB_DISK_FLAG_FIXED);
 }
 
 /* Delay in recovery mode */
@@ -557,7 +563,8 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 #define REC_KEY_DELAY        20       /* Check keys every 20ms */
 #define REC_MEDIA_INIT_DELAY 500      /* Check removable media every 500ms */
 
-VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
+VbError_t VbBootRecovery(struct vb2_context *ctx, VbCommonParams *cparams,
+			 LoadKernelParams *p)
 {
 	VbSharedDataHeader *shared =
 		(VbSharedDataHeader *)cparams->shared_data_blob;
@@ -585,12 +592,12 @@ VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 		 */
 		VBDEBUG(("VbBootRecovery() saving recovery reason (%#x)\n",
 				shared->recovery_reason));
-		VbSetRecoverySubcode(shared->recovery_reason);
+		VbSetRecoverySubcode(ctx, shared->recovery_reason);
 		VbNvCommit();
-		VbDisplayScreen(cparams, VB_SCREEN_OS_BROKEN, 0, &vnc);
+		VbDisplayScreen(ctx, cparams, VB_SCREEN_OS_BROKEN, 0);
 		VBDEBUG(("VbBootRecovery() waiting for manual recovery\n"));
 		while (1) {
-			VbCheckDisplayKey(cparams, VbExKeyboardRead(), &vnc);
+			VbCheckDisplayKey(ctx, cparams, VbExKeyboardRead());
 			if (VbWantShutdown(cparams->gbb->flags))
 				return VBERROR_SHUTDOWN_REQUESTED;
 			VbExSleepMs(REC_KEY_DELAY);
@@ -601,7 +608,8 @@ VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 	VBDEBUG(("VbBootRecovery() waiting for a recovery image\n"));
 	while (1) {
 		VBDEBUG(("VbBootRecovery() attempting to load kernel2\n"));
-		retval = VbTryLoadKernel(cparams, p, VB_DISK_FLAG_REMOVABLE);
+		retval = VbTryLoadKernel(ctx, cparams, p,
+					 VB_DISK_FLAG_REMOVABLE);
 
 		/*
 		 * Clear recovery requests from failed kernel loading, since
@@ -609,15 +617,15 @@ VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 		 * powering off after inserting an invalid disk doesn't leave
 		 * us stuck in recovery mode.
 		 */
-		VbSetRecoveryRequest(VBNV_RECOVERY_NOT_REQUESTED);
+		VbSetRecoveryRequest(ctx, VBNV_RECOVERY_NOT_REQUESTED);
 
 		if (VBERROR_SUCCESS == retval)
 			break; /* Found a recovery kernel */
 
-		VbDisplayScreen(cparams, VBERROR_NO_DISK_FOUND == retval ?
+		VbDisplayScreen(ctx, cparams, VBERROR_NO_DISK_FOUND == retval ?
 				VB_SCREEN_RECOVERY_INSERT :
 				VB_SCREEN_RECOVERY_NO_GOOD,
-				0, &vnc);
+				0);
 
 		/*
 		 * Scan keyboard more frequently than media, since x86
@@ -655,14 +663,15 @@ VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 				}
 
 				/* Ask the user to confirm entering dev-mode */
-				VbDisplayScreen(cparams,
+				VbDisplayScreen(ctx, cparams,
 						VB_SCREEN_RECOVERY_TO_DEV,
-						0, &vnc);
+						0);
 				/* SPACE means no... */
 				uint32_t vbc_flags =
 					VB_CONFIRM_SPACE_MEANS_NO |
 					VB_CONFIRM_MUST_TRUST_KEYBOARD;
-				switch (VbUserConfirms(cparams, vbc_flags)) {
+				switch (VbUserConfirms(ctx, cparams,
+						       vbc_flags)) {
 				case 1:
 					VBDEBUG(("%s() Enabling dev-mode...\n",
 						 __func__));
@@ -672,7 +681,7 @@ VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 						 "effect\n", __func__));
 					if (VbExGetSwitches
 					    (VB_INIT_FLAG_ALLOW_USB_BOOT))
-						VbAllowUsbBoot();
+						VbAllowUsbBoot(ctx);
 					return VBERROR_REBOOT_REQUIRED;
 				case -1:
 					VBDEBUG(("%s() - Shutdown requested\n",
@@ -689,7 +698,7 @@ VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 					break;
 				}
 			} else {
-				VbCheckDisplayKey(cparams, key, &vnc);
+				VbCheckDisplayKey(ctx, cparams, key);
 			}
 			if (VbWantShutdown(cparams->gbb->flags))
 				return VBERROR_SHUTDOWN_REQUESTED;
@@ -722,7 +731,6 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 	VbError_t retval = VBERROR_SUCCESS;
 	LoadKernelParams p;
 	uint32_t tpm_status = 0;
-	uint32_t battery_cutoff = 0;
 
 	/* Start timer */
 	shared->timer_vb_select_and_load_kernel_enter = VbExGetTimer();
@@ -743,6 +751,53 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 	p.kernel_buffer = kparams->kernel_buffer;
 	p.kernel_buffer_size = kparams->kernel_buffer_size;
 
+	/* Set up boot flags */
+	p.boot_flags = 0;
+	if (shared->flags & VBSD_BOOT_DEV_SWITCH_ON)
+		p.boot_flags |= BOOT_FLAG_DEVELOPER;
+	if (shared->recovery_reason)
+		p.boot_flags |= BOOT_FLAG_RECOVERY;
+
+	/*
+	 * Set up vboot context.
+	 *
+	 * TODO: Propagate this up to higher API levels, and use more of the
+	 * context fields (e.g. secdatak) and flags.
+	 */
+	struct vb2_context ctx;
+	memset(&ctx, 0, sizeof(ctx));
+	memcpy(ctx.nvdata, vnc.raw, VB2_NVDATA_SIZE);
+
+	if (p.boot_flags & BOOT_FLAG_RECOVERY)
+		ctx.flags |= VB2_CONTEXT_RECOVERY_MODE;
+	if (p.boot_flags & BOOT_FLAG_DEVELOPER)
+		ctx.flags |= VB2_CONTEXT_DEVELOPER_MODE;
+
+	ctx.workbuf_size = VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE +
+			VB2_WORKBUF_ALIGN;
+
+	uint8_t *unaligned_workbuf = ctx.workbuf = malloc(ctx.workbuf_size);
+	if (!unaligned_workbuf) {
+		VB2_DEBUG("%s: Can't allocate work buffer\n", __func__);
+		VbSetRecoveryRequest(&ctx, VB2_RECOVERY_RW_SHARED_DATA);
+		return VBERROR_INIT_SHARED_DATA;
+	}
+
+	if (VB2_SUCCESS != vb2_align(&ctx.workbuf, &ctx.workbuf_size,
+				     VB2_WORKBUF_ALIGN,
+				     VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE)) {
+		VB2_DEBUG("%s: Can't align work buffer\n", __func__);
+		VbSetRecoveryRequest(&ctx, VB2_RECOVERY_RW_SHARED_DATA);
+		return VBERROR_INIT_SHARED_DATA;
+	}
+
+	if (VB2_SUCCESS != vb2_init_context(&ctx)) {
+		VB2_DEBUG("%s: Can't init vb2_context\n", __func__);
+		free(unaligned_workbuf);
+		VbSetRecoveryRequest(&ctx, VB2_RECOVERY_RW_SHARED_DATA);
+		return VBERROR_INIT_SHARED_DATA;
+	}
+
 	/* Clear output params in case we fail */
 	kparams->disk_handle = NULL;
 	kparams->partition_number = 0;
@@ -762,7 +817,7 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 	    !(cparams->gbb->flags & GBB_FLAG_DISABLE_EC_SOFTWARE_SYNC)) {
 		int oprom_mismatch = 0;
 
-		retval = VbEcSoftwareSync(0, cparams, &vnc);
+		retval = VbEcSoftwareSync(&ctx, 0, cparams);
 		/* Save reboot requested until after possible PD sync */
 		if (retval == VBERROR_VGA_OPROM_MISMATCH)
 			oprom_mismatch = 1;
@@ -772,7 +827,7 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 #ifdef PD_SYNC
 		if (!(cparams->gbb->flags &
 		      GBB_FLAG_DISABLE_PD_SOFTWARE_SYNC)) {
-			retval = VbEcSoftwareSync(1, cparams, &vnc);
+			retval = VbEcSoftwareSync(&ctx, 1, cparams);
 			if (retval == VBERROR_VGA_OPROM_MISMATCH)
 				oprom_mismatch = 1;
 			else if (retval != VBERROR_SUCCESS)
@@ -794,10 +849,9 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 
 	/* Check if we need to cut-off battery. This must be done after EC
          * firmware updating and before kernel started. */
-	VbNvGet(&vnc, VBNV_BATTERY_CUTOFF_REQUEST, &battery_cutoff);
-	if (battery_cutoff) {
+	if (vb2_nv_get(&ctx, VB2_NV_BATTERY_CUTOFF_REQUEST)) {
 		VBDEBUG(("Request to cut-off battery\n"));
-		VbNvSet(&vnc, VBNV_BATTERY_CUTOFF_REQUEST, 0);
+		vb2_nv_set(&ctx, VB2_NV_BATTERY_CUTOFF_REQUEST, 0);
 		VbExEcBatteryCutOff();
 		retval = VBERROR_SHUTDOWN_REQUESTED;
 		goto VbSelectAndLoadKernel_exit;
@@ -808,7 +862,8 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 	if (0 != tpm_status) {
 		VBDEBUG(("Unable to get kernel versions from TPM\n"));
 		if (!shared->recovery_reason) {
-			VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_R_ERROR);
+			VbSetRecoveryRequest(&ctx,
+					     VBNV_RECOVERY_RW_TPM_R_ERROR);
 			retval = VBERROR_TPM_READ_KERNEL;
 			goto VbSelectAndLoadKernel_exit;
 		}
@@ -825,40 +880,13 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 	if (0 != tpm_status) {
 		VBDEBUG(("Unable to get FWMP from TPM\n"));
 		if (!shared->recovery_reason) {
-			VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_R_ERROR);
+			VbSetRecoveryRequest(&ctx,
+					     VBNV_RECOVERY_RW_TPM_R_ERROR);
 			retval = VBERROR_TPM_READ_FWMP;
 			goto VbSelectAndLoadKernel_exit;
 		}
 	}
 
-	/* Set up boot flags */
-	p.boot_flags = 0;
-	if (shared->flags & VBSD_BOOT_DEV_SWITCH_ON)
-		p.boot_flags |= BOOT_FLAG_DEVELOPER;
-
-	/* Handle separate normal and developer firmware builds. */
-#if defined(VBOOT_FIRMWARE_TYPE_NORMAL)
-	/* Normal-type firmware always acts like the dev switch is off. */
-	p.boot_flags &= ~BOOT_FLAG_DEVELOPER;
-#elif defined(VBOOT_FIRMWARE_TYPE_DEVELOPER)
-	/* Developer-type firmware fails if the dev switch is off. */
-	if (!(p.boot_flags & BOOT_FLAG_DEVELOPER)) {
-		/*
-		 * Dev firmware should be signed with a key that only verifies
-		 * when the dev switch is on, so we should never get here.
-		 */
-		VBDEBUG(("Developer firmware called with dev switch off!\n"));
-		VbSetRecoveryRequest(VBNV_RECOVERY_RW_DEV_MISMATCH);
-		retval = VBERROR_DEV_FIRMWARE_SWITCH_MISMATCH;
-		goto VbSelectAndLoadKernel_exit;
-	}
-#else
-	/*
-	 * Recovery firmware, or merged normal+developer firmware.  No need to
-	 * override flags.
-	 */
-#endif
-
 	/* Select boot path */
 	if (shared->recovery_reason == VBNV_RECOVERY_TRAIN_AND_REBOOT) {
 		/* Reboot requested by user recovery code. */
@@ -867,21 +895,20 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 		retval = VBERROR_REBOOT_REQUIRED;
 	} else if (shared->recovery_reason) {
 		/* Recovery boot */
-		p.boot_flags |= BOOT_FLAG_RECOVERY;
-		retval = VbBootRecovery(cparams, &p);
+		retval = VbBootRecovery(&ctx, cparams, &p);
 		VbExEcEnteringMode(0, VB_EC_RECOVERY);
-		VbDisplayScreen(cparams, VB_SCREEN_BLANK, 0, &vnc);
+		VbDisplayScreen(&ctx, cparams, VB_SCREEN_BLANK, 0);
 
 	} else if (p.boot_flags & BOOT_FLAG_DEVELOPER) {
 		/* Developer boot */
-		retval = VbBootDeveloper(cparams, &p);
+		retval = VbBootDeveloper(&ctx, cparams, &p);
 		VbExEcEnteringMode(0, VB_EC_DEVELOPER);
-		VbDisplayScreen(cparams, VB_SCREEN_BLANK, 0, &vnc);
+		VbDisplayScreen(&ctx, cparams, VB_SCREEN_BLANK, 0);
 
 	} else {
 		/* Normal boot */
 		VbExEcEnteringMode(0, VB_EC_NORMAL);
-		retval = VbBootNormal(cparams, &p);
+		retval = VbBootNormal(&ctx, cparams, &p);
 
 		if ((1 == shared->firmware_index) &&
 		    (shared->flags & VBSD_FWB_TRIED)) {
@@ -908,7 +935,8 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 			if (VBERROR_INVALID_KERNEL_FOUND == retval) {
 				VBDEBUG(("Trying firmware B, "
 					 "and only found invalid kernels.\n"));
-				VbSetRecoveryRequest(VBNV_RECOVERY_NOT_REQUESTED);
+				VbSetRecoveryRequest(&ctx,
+						VBNV_RECOVERY_NOT_REQUESTED);
 				goto VbSelectAndLoadKernel_exit;
 			}
 		} else {
@@ -924,7 +952,8 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 				if (0 != tpm_status) {
 					VBDEBUG(("Error writing kernel "
 						 "versions to TPM.\n"));
-					VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_W_ERROR);
+					VbSetRecoveryRequest(&ctx,
+						VBNV_RECOVERY_RW_TPM_W_ERROR);
 					retval = VBERROR_TPM_WRITE_KERNEL;
 					goto VbSelectAndLoadKernel_exit;
 				}
@@ -949,7 +978,8 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 	if (0 != tpm_status) {
 		VBDEBUG(("Error locking kernel versions.\n"));
 		if (!shared->recovery_reason) {
-			VbSetRecoveryRequest(VBNV_RECOVERY_RW_TPM_L_ERROR);
+			VbSetRecoveryRequest(&ctx,
+					     VBNV_RECOVERY_RW_TPM_L_ERROR);
 			retval = VBERROR_TPM_LOCK_KERNEL;
 			goto VbSelectAndLoadKernel_exit;
 		}
@@ -957,6 +987,20 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 
  VbSelectAndLoadKernel_exit:
 
+	/*
+	 * Clean up vboot context.
+	 *
+	 * TODO: This should propagate up to higher levels
+	 */
+	/* Free buffers */
+	free(unaligned_workbuf);
+	/* Copy nvdata back to old vboot1 nv context if needed */
+	if (ctx.flags & VB2_CONTEXT_NVDATA_CHANGED) {
+		memcpy(vnc.raw, ctx.nvdata, VB2_NVDATA_SIZE);
+		vnc.raw_changed = 1;
+		ctx.flags &= ~VB2_CONTEXT_NVDATA_CHANGED;
+	}
+
 	VbApiKernelFree(cparams);
 
 	VbNvCommit();
@@ -1025,7 +1069,7 @@ VbError_t VbVerifyMemoryBootImage(VbCommonParams *cparams,
 	dev_switch = shared->flags & VBSD_BOOT_DEV_SWITCH_ON;
 
 	VbNvLoad();
-	VbNvGet(&vnc, VBNV_DEV_BOOT_FASTBOOT_FULL_CAP,
+	VbNvGet(&vnc, VB2_NV_DEV_BOOT_FASTBOOT_FULL_CAP,
 		&allow_fastboot_full_cap);
 
 	if (0 == allow_fastboot_full_cap) {
@@ -1166,8 +1210,7 @@ VbError_t VbLockDevice(void)
 
 	VBDEBUG(("%s() - Storing request to leave dev-mode.\n",
 		 __func__));
-	VbNvSet(&vnc, VBNV_DISABLE_DEV_REQUEST,
-		1);
+	VbNvSet(&vnc, VBNV_DISABLE_DEV_REQUEST, 1);
 
 	VbNvCommit();
 
diff --git a/firmware/lib/vboot_display.c b/firmware/lib/vboot_display.c
index 1abed924..39bd2fa7 100644
--- a/firmware/lib/vboot_display.c
+++ b/firmware/lib/vboot_display.c
@@ -9,6 +9,7 @@
 #include "2sysincludes.h"
 
 #include "2common.h"
+#include "2nvstorage.h"
 #include "2sha.h"
 #include "bmpblk_font.h"
 #include "gbb_access.h"
@@ -148,8 +149,9 @@ void VbRenderTextAtPos(const char *text, int right_to_left,
 	}
 }
 
-VbError_t VbDisplayScreenFromGBB(VbCommonParams *cparams, uint32_t screen,
-                                 VbNvContext *vncptr, uint32_t localization)
+VbError_t VbDisplayScreenFromGBB(struct vb2_context *ctx,
+				 VbCommonParams *cparams, uint32_t screen,
+				 uint32_t localization)
 {
 	char *fullimage = NULL;
 	BmpBlockHeader hdr;
@@ -222,8 +224,8 @@ VbError_t VbDisplayScreenFromGBB(VbCommonParams *cparams, uint32_t screen,
 	/* Clip localization to number of localizations present in the GBB */
 	if (localization >= hdr.number_of_localizations) {
 		localization = 0;
-		VbNvSet(vncptr, VBNV_LOCALIZATION_INDEX, localization);
-		VbNvSet(vncptr, VBNV_BACKUP_NVRAM_REQUEST, 1);
+		vb2_nv_set(ctx, VB2_NV_LOCALIZATION_INDEX, localization);
+		vb2_nv_set(ctx, VB2_NV_BACKUP_NVRAM_REQUEST, 1);
 	}
 
 	/* Display all bitmaps for the image */
@@ -319,9 +321,9 @@ VbError_t VbDisplayScreenFromGBB(VbCommonParams *cparams, uint32_t screen,
  * if bmpblk.bin is found in GBB. New devices store graphics data in cbfs
  * and screens are rendered by Depthcharge (chromium:502066).
  */
-static VbError_t VbDisplayScreenLegacy(VbCommonParams *cparams, uint32_t screen,
-				       int force, VbNvContext *vncptr,
-				       uint32_t locale)
+static VbError_t VbDisplayScreenLegacy(struct vb2_context *ctx,
+				       VbCommonParams *cparams, uint32_t screen,
+				       int force, uint32_t locale)
 {
 	VbError_t retval;
 
@@ -336,16 +338,16 @@ static VbError_t VbDisplayScreenLegacy(VbCommonParams *cparams, uint32_t screen,
 	VbExDisplayBacklight(VB_SCREEN_BLANK == screen ? 0 : 1);
 
 	/* Look in the GBB first */
-	if (VBERROR_SUCCESS == VbDisplayScreenFromGBB(cparams, screen,
-						      vncptr, locale))
+	if (VBERROR_SUCCESS == VbDisplayScreenFromGBB(ctx, cparams, screen,
+						      locale))
 		return VBERROR_SUCCESS;
 
 	/* If screen wasn't in the GBB bitmaps, fall back to a default */
 	return VbExDisplayScreen(screen, locale);
 }
 
-VbError_t VbDisplayScreen(VbCommonParams *cparams, uint32_t screen,
-			  int force, VbNvContext *vncptr)
+VbError_t VbDisplayScreen(struct vb2_context *ctx,
+			  VbCommonParams *cparams, uint32_t screen, int force)
 {
 	uint32_t locale;
 	GoogleBinaryBlockHeader *gbb = cparams->gbb;
@@ -356,13 +358,12 @@ VbError_t VbDisplayScreen(VbCommonParams *cparams, uint32_t screen,
 		return VBERROR_SUCCESS;
 
 	/* Read the locale last saved */
-	VbNvGet(vncptr, VBNV_LOCALIZATION_INDEX, &locale);
+	locale = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
 
 	if (gbb->bmpfv_size == 0)
 		rv = VbExDisplayScreen(screen, locale);
 	else
-		rv = VbDisplayScreenLegacy(cparams, screen, force, vncptr,
-					   locale);
+		rv = VbDisplayScreenLegacy(ctx, cparams, screen, force, locale);
 
 	if (rv == VBERROR_SUCCESS)
 		/* Keep track of the currently displayed screen */
@@ -551,7 +552,7 @@ const char *RecoveryReasonString(uint8_t code)
 
 #define DEBUG_INFO_SIZE 512
 
-VbError_t VbDisplayDebugInfo(VbCommonParams *cparams, VbNvContext *vncptr)
+VbError_t VbDisplayDebugInfo(struct vb2_context *ctx, VbCommonParams *cparams)
 {
 	VbSharedDataHeader *shared =
 		(VbSharedDataHeader *)cparams->shared_data_blob;
@@ -565,7 +566,7 @@ VbError_t VbDisplayDebugInfo(VbCommonParams *cparams, VbNvContext *vncptr)
 	uint32_t i;
 
 	/* Redisplay current screen to overwrite any previous debug output */
-	VbDisplayScreen(cparams, disp_current_screen, 1, vncptr);
+	VbDisplayScreen(ctx, cparams, disp_current_screen, 1);
 
 	/* Add hardware ID */
 	VbRegionReadHWID(cparams, hwid, sizeof(hwid));
@@ -573,7 +574,7 @@ VbError_t VbDisplayDebugInfo(VbCommonParams *cparams, VbNvContext *vncptr)
 	used += StrnAppend(buf + used, hwid, DEBUG_INFO_SIZE - used);
 
 	/* Add recovery reason and subcode */
-	VbNvGet(vncptr, VBNV_RECOVERY_SUBCODE, &i);
+	i = vb2_nv_get(ctx, VB2_NV_RECOVERY_SUBCODE);
 	used += StrnAppend(buf + used,
 			"\nrecovery_reason: 0x", DEBUG_INFO_SIZE - used);
 	used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used,
@@ -595,34 +596,34 @@ VbError_t VbDisplayDebugInfo(VbCommonParams *cparams, VbNvContext *vncptr)
 	for (i = 0; i < VBNV_BLOCK_SIZE; i++) {
 		used += StrnAppend(buf + used, " ", DEBUG_INFO_SIZE - used);
 		used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used,
-				       vncptr->raw[i], 16, 2);
+				       ctx->nvdata[i], 16, 2);
 	}
 
 	/* Add dev_boot_usb flag */
-	VbNvGet(vncptr, VBNV_DEV_BOOT_USB, &i);
+	i = vb2_nv_get(ctx, VB2_NV_DEV_BOOT_USB);
 	used += StrnAppend(buf + used, "\ndev_boot_usb: ", DEBUG_INFO_SIZE - used);
 	used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used, i, 10, 0);
 
 	/* Add dev_boot_legacy flag */
-	VbNvGet(vncptr, VBNV_DEV_BOOT_LEGACY, &i);
+	i = vb2_nv_get(ctx, VB2_NV_DEV_BOOT_LEGACY);
 	used += StrnAppend(buf + used,
 			"\ndev_boot_legacy: ", DEBUG_INFO_SIZE - used);
 	used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used, i, 10, 0);
 
 	/* Add dev_default_boot flag */
-	VbNvGet(vncptr, VBNV_DEV_DEFAULT_BOOT, &i);
+	i = vb2_nv_get(ctx, VB2_NV_DEV_DEFAULT_BOOT);
 	used += StrnAppend(buf + used,
 			"\ndev_default_boot: ", DEBUG_INFO_SIZE - used);
 	used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used, i, 10, 0);
 
 	/* Add dev_boot_signed_only flag */
-	VbNvGet(vncptr, VBNV_DEV_BOOT_SIGNED_ONLY, &i);
+	i = vb2_nv_get(ctx, VB2_NV_DEV_BOOT_SIGNED_ONLY);
 	used += StrnAppend(buf + used, "\ndev_boot_signed_only: ",
 			DEBUG_INFO_SIZE - used);
 	used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used, i, 10, 0);
 
 	/* Add dev_boot_fastboot_full_cap flag */
-	VbNvGet(vncptr, VBNV_DEV_BOOT_FASTBOOT_FULL_CAP, &i);
+	i = vb2_nv_get(ctx, VB2_NV_DEV_BOOT_FASTBOOT_FULL_CAP);
 	used += StrnAppend(buf + used, "\ndev_boot_fastboot_full_cap: ",
 			DEBUG_INFO_SIZE - used);
 	used += Uint64ToString(buf + used, DEBUG_INFO_SIZE - used, i, 10, 0);
@@ -688,8 +689,8 @@ VbError_t VbDisplayDebugInfo(VbCommonParams *cparams, VbNvContext *vncptr)
 #define MAGIC_WORD "xyzzy"
 static uint8_t MagicBuffer[MAGIC_WORD_LEN];
 
-VbError_t VbCheckDisplayKey(VbCommonParams *cparams, uint32_t key,
-                            VbNvContext *vncptr)
+VbError_t VbCheckDisplayKey(struct vb2_context *ctx, VbCommonParams *cparams,
+			    uint32_t key)
 {
 	int i;
 
@@ -701,14 +702,14 @@ VbError_t VbCheckDisplayKey(VbCommonParams *cparams, uint32_t key,
 
 	if ('\t' == key) {
 		/* Tab = display debug info */
-		return VbDisplayDebugInfo(cparams, vncptr);
+		return VbDisplayDebugInfo(ctx, cparams);
 	} else if (VB_KEY_LEFT == key || VB_KEY_RIGHT == key ||
 		   VB_KEY_DOWN == key || VB_KEY_UP == key) {
 		/* Arrow keys = change localization */
 		uint32_t loc = 0;
 		uint32_t count = 0;
 
-		VbNvGet(vncptr, VBNV_LOCALIZATION_INDEX, &loc);
+		loc = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
 		if (VBERROR_SUCCESS != VbGetLocalizationCount(cparams, &count))
 			loc = 0;  /* No localization count (bad GBB?) */
 		else if (VB_KEY_RIGHT == key || VB_KEY_UP == key)
@@ -717,23 +718,30 @@ VbError_t VbCheckDisplayKey(VbCommonParams *cparams, uint32_t key,
 			loc = (loc > 0 ? loc - 1 : count - 1);
 		VBDEBUG(("VbCheckDisplayKey() - change localization to %d\n",
 			 (int)loc));
-		VbNvSet(vncptr, VBNV_LOCALIZATION_INDEX, loc);
-		VbNvSet(vncptr, VBNV_BACKUP_NVRAM_REQUEST, 1);
+		vb2_nv_set(ctx, VB2_NV_LOCALIZATION_INDEX, loc);
+		vb2_nv_set(ctx, VB2_NV_BACKUP_NVRAM_REQUEST, 1);
 
 #ifdef SAVE_LOCALE_IMMEDIATELY
-		VbNvTeardown(vncptr);  /* really only computes checksum */
-		if (vncptr->raw_changed)
-			VbExNvStorageWrite(vncptr->raw);
+		/*
+		 * This is a workaround for coreboot on x86, which will power
+		 * off asynchronously without giving us a chance to react.
+		 * This is not an example of the Right Way to do things.  See
+		 * chrome-os-partner:7689.
+		 */
+		if (ctx->flags & VB2_CONTEXT_NVDATA_CHANGED) {
+			VbExNvStorageWrite(ctx.nvdata);
+			ctx.flags &= ~VB2_CONTEXT_NVDATA_CHANGED;
+		}
 #endif
 
 		/* Force redraw of current screen */
-		return VbDisplayScreen(cparams, disp_current_screen, 1, vncptr);
+		return VbDisplayScreen(ctx, cparams, disp_current_screen, 1);
 	}
 
 	if (0 == memcmp(MagicBuffer, MAGIC_WORD, MAGIC_WORD_LEN)) {
 		if (VBEASTEREGG)
-			(void)VbDisplayScreen(cparams, disp_current_screen,
-					      1, vncptr);
+			(void)VbDisplayScreen(ctx, cparams, disp_current_screen,
+					      1);
 	}
 
   return VBERROR_SUCCESS;
diff --git a/firmware/lib/vboot_kernel.c b/firmware/lib/vboot_kernel.c
index c0afd37b..20272448 100644
--- a/firmware/lib/vboot_kernel.c
+++ b/firmware/lib/vboot_kernel.c
@@ -10,6 +10,8 @@
 #include "2sysincludes.h"
 
 #include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
 #include "2rsa.h"
 #include "2sha.h"
 #include "cgptlib.h"
@@ -40,12 +42,12 @@ enum vboot_mode {
  * @param params	Load kernel parameters
  * @return The current boot mode.
  */
-static enum vboot_mode get_kernel_boot_mode(const LoadKernelParams *params)
+static enum vboot_mode get_kernel_boot_mode(struct vb2_context *ctx)
 {
-	if (BOOT_FLAG_RECOVERY & params->boot_flags)
+	if (ctx->flags & VB2_CONTEXT_RECOVERY_MODE)
 		return kBootRecovery;
 
-	if (BOOT_FLAG_DEVELOPER & params->boot_flags)
+	if (ctx->flags & VB2_CONTEXT_DEVELOPER_MODE)
 		return kBootDev;
 
 	return kBootNormal;
@@ -57,10 +59,11 @@ static enum vboot_mode get_kernel_boot_mode(const LoadKernelParams *params)
  * @param params	Load kernel parameters
  * @return 1 if official OS required; 0 if self-signed kernels are ok
  */
-static int require_official_os(const LoadKernelParams *params)
+static int require_official_os(struct vb2_context *ctx,
+			       const LoadKernelParams *params)
 {
 	/* Normal and recovery modes always require official OS */
-	if (get_kernel_boot_mode(params) != kBootDev)
+	if (get_kernel_boot_mode(ctx) != kBootDev)
 		return 1;
 
 	/* FWMP can require developer mode to use official OS */
@@ -69,9 +72,7 @@ static int require_official_os(const LoadKernelParams *params)
 		return 1;
 
 	/* Developer can request official OS via nvstorage */
-	uint32_t signed_only = 1;
-	VbNvGet(params->nv_context, VBNV_DEV_BOOT_SIGNED_ONLY, &signed_only);
-	return signed_only;
+	return vb2_nv_get(ctx, VB2_NV_DEV_BOOT_SIGNED_ONLY);
 }
 
 /**
@@ -128,7 +129,8 @@ static uint32_t get_body_offset(uint8_t *kbuf)
  *			VB2_VERIFY_KERNEL_PREAMBLE_WORKBUF_BYTES bytes.
  * @return VB2_SUCCESS, or non-zero error code.
  */
-int vb2_verify_kernel_vblock(uint8_t *kbuf,
+int vb2_verify_kernel_vblock(struct vb2_context *ctx,
+			     uint8_t *kbuf,
 			     uint32_t kbuf_size,
 			     const struct vb2_packed_key *kernel_subkey,
 			     const LoadKernelParams *params,
@@ -153,7 +155,7 @@ int vb2_verify_kernel_vblock(uint8_t *kbuf,
 		keyblock_valid = 0;
 
 		/* Check if we must have an officially signed kernel */
-		if (require_official_os(params)) {
+		if (require_official_os(ctx, params)) {
 			VB2_DEBUG("Self-signed kernels not enabled.\n");
 			shpart->check_result = VBSD_LKP_CHECK_SELF_SIGNED;
 			return VB2_ERROR_VBLOCK_SELF_SIGNED;
@@ -170,14 +172,14 @@ int vb2_verify_kernel_vblock(uint8_t *kbuf,
 
 	/* Check the key block flags against boot flags. */
 	if (!(keyblock->keyblock_flags &
-	      ((BOOT_FLAG_DEVELOPER & params->boot_flags) ?
+	      ((ctx->flags & VB2_CONTEXT_DEVELOPER_MODE) ?
 	       KEY_BLOCK_FLAG_DEVELOPER_1 : KEY_BLOCK_FLAG_DEVELOPER_0))) {
 		VB2_DEBUG("Key block developer flag mismatch.\n");
 		shpart->check_result = VBSD_LKP_CHECK_DEV_MISMATCH;
 		keyblock_valid = 0;
 	}
 	if (!(keyblock->keyblock_flags &
-	      ((BOOT_FLAG_RECOVERY & params->boot_flags) ?
+	      ((ctx->flags & VB2_CONTEXT_RECOVERY_MODE) ?
 	       KEY_BLOCK_FLAG_RECOVERY_1 : KEY_BLOCK_FLAG_RECOVERY_0))) {
 		VB2_DEBUG("Key block recovery flag mismatch.\n");
 		shpart->check_result = VBSD_LKP_CHECK_REC_MISMATCH;
@@ -185,7 +187,7 @@ int vb2_verify_kernel_vblock(uint8_t *kbuf,
 	}
 
 	/* Check for rollback of key version except in recovery mode. */
-	enum vboot_mode boot_mode = get_kernel_boot_mode(params);
+	enum vboot_mode boot_mode = get_kernel_boot_mode(ctx);
 	uint32_t key_version = keyblock->data_key.key_version;
 	if (kBootRecovery != boot_mode) {
 		if (key_version < (min_version >> 16)) {
@@ -294,31 +296,33 @@ enum vb2_load_partition_flags {
 };
 
 #define KBUF_SIZE 65536  /* Bytes to read at start of kernel partition */
+
+/* Minimum context work buffer size needed for vb2_load_partition() */
 #define VB2_LOAD_PARTITION_WORKBUF_BYTES	\
 	(VB2_VERIFY_KERNEL_PREAMBLE_WORKBUF_BYTES + KBUF_SIZE)
 
 /**
  * Load and verify a partition from the stream.
  *
+ * @param ctx		Vboot context
  * @param stream	Stream to load kernel from
  * @param kernel_subkey	Key to use to verify vblock
  * @param flags		Flags (one or more of vb2_load_partition_flags)
  * @param params	Load-kernel parameters
  * @param min_version	Minimum kernel version from TPM
  * @param shpart	Destination for verification results
- * @param wb		Work buffer.  Must be at least
- *			VB2_LOAD_PARTITION_WORKBUF_BYTES bytes.
  * @return VB2_SUCCESS, or non-zero error code.
  */
-int vb2_load_partition(VbExStream_t stream,
+int vb2_load_partition(struct vb2_context *ctx,
+		       VbExStream_t stream,
 		       const struct vb2_packed_key *kernel_subkey,
 		       uint32_t flags,
 		       LoadKernelParams *params,
 		       uint32_t min_version,
-		       VbSharedDataKernelPart *shpart,
-		       struct vb2_workbuf *wb)
+		       VbSharedDataKernelPart *shpart)
 {
-	struct vb2_workbuf wblocal = *wb;
+	struct vb2_workbuf wblocal;
+	vb2_workbuf_from_ctx(ctx, &wblocal);
 
 	/* Allocate kernel header buffer in workbuf */
 	uint8_t *kbuf = vb2_workbuf_alloc(&wblocal, KBUF_SIZE);
@@ -333,8 +337,8 @@ int vb2_load_partition(VbExStream_t stream,
 	}
 
 	if (VB2_SUCCESS !=
-	    vb2_verify_kernel_vblock(kbuf, KBUF_SIZE, kernel_subkey, params,
-				     min_version, shpart, &wblocal)) {
+	    vb2_verify_kernel_vblock(ctx, kbuf, KBUF_SIZE, kernel_subkey,
+				     params, min_version, shpart, &wblocal)) {
 		return VB2_ERROR_LOAD_PARTITION_VERIFY_VBLOCK;
 	}
 
@@ -424,7 +428,8 @@ int vb2_load_partition(VbExStream_t stream,
 	return VB2_SUCCESS;
 }
 
-VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
+VbError_t LoadKernel(struct vb2_context *ctx, LoadKernelParams *params,
+		     VbCommonParams *cparams)
 {
 	VbSharedDataHeader *shared = cparams->shared_data_blob;
 	VbSharedDataKernelCall *shcall = NULL;
@@ -449,7 +454,7 @@ VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
 			(shared->lk_call_count & (VBSD_MAX_KERNEL_CALLS - 1));
 	memset(shcall, 0, sizeof(*shcall));
 	shcall->boot_flags = (uint32_t)params->boot_flags;
-	shcall->boot_mode = get_kernel_boot_mode(params);
+	shcall->boot_mode = get_kernel_boot_mode(ctx);
 	shcall->sector_size = (uint32_t)params->bytes_per_lba;
 	shcall->sector_count = params->streaming_lba_count;
 	shared->lk_call_count++;
@@ -468,14 +473,6 @@ VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
 		kernel_subkey = (struct vb2_packed_key *)&shared->kernel_subkey;
 	}
 
-	/* Allocate work buffer */
-	uint8_t *workbuf = malloc(VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE);
-	if (!workbuf)
-		goto load_kernel_exit;
-
-	struct vb2_workbuf wb;
-	vb2_workbuf_init(&wb, workbuf, VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE);
-
 	/* Read GPT data */
 	GptData gpt;
 	gpt.sector_bytes = (uint32_t)params->bytes_per_lba;
@@ -546,13 +543,13 @@ VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
 			lpflags |= VB2_LOAD_PARTITION_VBLOCK_ONLY;
 		}
 
-		int rv = vb2_load_partition(stream,
+		int rv = vb2_load_partition(ctx,
+					    stream,
 					    kernel_subkey,
 					    lpflags,
 					    params,
 					    shared->kernel_version_tpm,
-					    shpart,
-					    &wb);
+					    shpart);
 		VbExStreamClose(stream);
 
 		if (rv != VB2_SUCCESS) {
@@ -625,10 +622,6 @@ VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
 	} /* while(GptNextKernelEntry) */
 
 gpt_done:
-	/* Free buffers */
-	if (workbuf)
-		free(workbuf);
-
 	/* Write and free GPT data */
 	WriteAndFreeGptData(params->disk_handle, &gpt);
 
@@ -660,11 +653,10 @@ gpt_done:
 	}
 
 load_kernel_exit:
-
 	/* Store recovery request, if any */
-	VbNvSet(params->nv_context, VBNV_RECOVERY_REQUEST,
-		VBERROR_SUCCESS != retval ?
-		recovery : VBNV_RECOVERY_NOT_REQUESTED);
+	vb2_nv_set(ctx, VB2_NV_RECOVERY_REQUEST,
+		   VBERROR_SUCCESS != retval ?
+		   recovery : VBNV_RECOVERY_NOT_REQUESTED);
 
 	/* Store how much shared data we used, if any */
 	cparams->shared_data_size = shared->data_used;