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>

8 files changed, 306 insertions, 218 deletions

diff --git a/tests/ec_sync_tests.c b/tests/ec_sync_tests.c
index 28a6a143..dfdd27bc 100644
--- a/tests/ec_sync_tests.c
+++ b/tests/ec_sync_tests.c
@@ -9,6 +9,10 @@
 #include <stdio.h>
 #include <stdlib.h>
 
+#include "2sysincludes.h"
+#include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
 #include "gbb_header.h"
 #include "host_common.h"
 #include "load_kernel_fw.h"
@@ -16,6 +20,7 @@
 #include "test_common.h"
 #include "vboot_audio.h"
 #include "vboot_common.h"
+#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_nvstorage.h"
 #include "vboot_struct.h"
@@ -47,7 +52,8 @@ static int mock_ec_rw_hash_size;
 static uint8_t want_ec_hash[32];
 static uint8_t update_hash;
 static int want_ec_hash_size;
-static VbNvContext mock_vnc;
+static struct vb2_context ctx;
+static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE];
 
 static uint32_t screens_displayed[8];
 static uint32_t screens_count = 0;
@@ -65,9 +71,11 @@ static void ResetMocks(void)
 	gbb.minor_version = GBB_MINOR_VER;
 	gbb.flags = 0;
 
-	memset(&mock_vnc, 0, sizeof(VbNvContext));
-	VbNvSetup(&mock_vnc);
-	VbNvTeardown(&mock_vnc); /* So CRC gets generated */
+	memset(&ctx, 0, sizeof(ctx));
+	ctx.workbuf = workbuf;
+	ctx.workbuf_size = sizeof(workbuf);
+	vb2_init_context(&ctx);
+	vb2_nv_init(&ctx);
 
 	memset(&shared_data, 0, sizeof(shared_data));
 	VbSharedDataInit(shared, sizeof(shared_data));
@@ -191,8 +199,8 @@ VbError_t VbExEcUpdateImage(int devidx, enum VbSelectFirmware_t select,
 	return update_retval;
 }
 
-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)
 {
 	if (screens_count < ARRAY_SIZE(screens_displayed))
 		screens_displayed[screens_count++] = screen;
@@ -202,11 +210,9 @@ VbError_t VbDisplayScreen(VbCommonParams *cparams, uint32_t screen, int force,
 
 static void test_ssync(VbError_t retval, int recovery_reason, const char *desc)
 {
-	uint32_t u;
-
-	TEST_EQ(VbEcSoftwareSync(0, &cparams, &mock_vnc), retval, desc);
-	VbNvGet(&mock_vnc, VBNV_RECOVERY_REQUEST, &u);
-	TEST_EQ(u, recovery_reason, "  recovery reason");
+	TEST_EQ(VbEcSoftwareSync(&ctx, 0, &cparams), retval, desc);
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST),
+		recovery_reason, "  recovery reason");
 }
 
 /* Tests */
@@ -297,7 +303,7 @@ static void VbSoftwareSyncTest(void)
 
 	ResetMocks();
 	mock_ec_rw_hash[0]++;
-	VbNvSet(&mock_vnc, VBNV_TRY_RO_SYNC, 1);
+	vb2_nv_set(&ctx, VB2_NV_TRY_RO_SYNC, 1);
 	test_ssync(0, 0, "Update rw without reboot");
 	TEST_EQ(ec_rw_protected, 1, "  ec rw protected");
 	TEST_EQ(ec_run_image, 1, "  ec run image");
@@ -308,7 +314,7 @@ static void VbSoftwareSyncTest(void)
 	ResetMocks();
 	mock_ec_rw_hash[0]++;
 	mock_ec_ro_hash[0]++;
-	VbNvSet(&mock_vnc, VBNV_TRY_RO_SYNC, 1);
+	vb2_nv_set(&ctx, VB2_NV_TRY_RO_SYNC, 1);
 	test_ssync(0, 0, "Update rw and ro images without reboot");
 	TEST_EQ(ec_rw_protected, 1, "  ec rw protected");
 	TEST_EQ(ec_run_image, 1, "  ec run image");
@@ -318,7 +324,7 @@ static void VbSoftwareSyncTest(void)
 
 	ResetMocks();
 	shared->flags |= VBSD_BOOT_FIRMWARE_WP_ENABLED;
-	VbNvSet(&mock_vnc, VBNV_TRY_RO_SYNC, 1);
+	vb2_nv_set(&ctx, VB2_NV_TRY_RO_SYNC, 1);
 	mock_ec_rw_hash[0]++;
 	mock_ec_ro_hash[0]++;
 	test_ssync(0, 0, "WP enabled");
@@ -329,7 +335,7 @@ static void VbSoftwareSyncTest(void)
 	TEST_EQ(ec_ro_updated, 0, "  ec ro updated");
 
 	ResetMocks();
-	VbNvSet(&mock_vnc, VBNV_TRY_RO_SYNC, 1);
+	vb2_nv_set(&ctx, VB2_NV_TRY_RO_SYNC, 1);
 	mock_ec_ro_hash[0]++;
 	test_ssync(0, 0, "rw update not needed");
 	TEST_EQ(ec_rw_protected, 1, "  ec rw protected");
diff --git a/tests/vboot_api_devmode_tests.c b/tests/vboot_api_devmode_tests.c
index abd8e858..1d680969 100644
--- a/tests/vboot_api_devmode_tests.c
+++ b/tests/vboot_api_devmode_tests.c
@@ -10,6 +10,8 @@
 #include <stdio.h>
 #include <stdlib.h>
 
+#include "2sysincludes.h"
+#include "2api.h"
 #include "crc32.h"
 #include "gbb_header.h"
 #include "host_common.h"
@@ -18,6 +20,7 @@
 #include "test_common.h"
 #include "vboot_common.h"
 #include "vboot_display.h"
+#include "vboot_kernel.h"
 #include "vboot_nvstorage.h"
 #include "vboot_struct.h"
 
@@ -132,6 +135,7 @@ test_case_t test[] = {
 /* Mock data */
 static VbCommonParams cparams;
 static LoadKernelParams lkparams;
+static struct vb2_context ctx;
 static VbNvContext vnc;
 static uint8_t shared_data[VB_SHARED_DATA_MIN_SIZE];
 static VbSharedDataHeader* shared = (VbSharedDataHeader*)shared_data;
@@ -156,6 +160,7 @@ static void ResetMocks(void) {
   cparams.gbb = &gbb;
 
   memset(&lkparams, 0, sizeof(lkparams));
+  memset(&ctx, 0, sizeof(ctx));
 
   memset(&vnc, 0, sizeof(vnc));
   VbNvSetup(&vnc);
@@ -299,9 +304,6 @@ VbError_t VbExDisplayScreen(uint32_t screen_type, uint32_t locale) {
 
 /****************************************************************************/
 
-VbError_t VbBootDeveloper(VbCommonParams* cparams, LoadKernelParams* p);
-
-
 static void VbBootDeveloperSoundTest(void) {
   int i;
   int num_tests =  sizeof(test) / sizeof(test_case_t);
@@ -315,7 +317,7 @@ static void VbBootDeveloperSoundTest(void) {
     kbd_fire_at = test[i].keypress_at_count;
     max_events = test[i].num_events;
     expected_event = test[i].notes;
-    (void) VbBootDeveloper(&cparams, &lkparams);
+    (void) VbBootDeveloper(&ctx, &cparams, &lkparams);
     VBDEBUG(("INFO: matched %d total %d expected %d\n",
              matched_events, current_event, test[i].num_events));
     TEST_TRUE(matched_events == test[i].num_events &&
diff --git a/tests/vboot_api_kernel2_tests.c b/tests/vboot_api_kernel2_tests.c
index 433933ba..284d4a1f 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -9,6 +9,10 @@
 #include <stdio.h>
 #include <stdlib.h>
 
+#include "2sysincludes.h"
+#include "2api.h"
+#include "2misc.h"
+#include "2nvstorage.h"
 #include "gbb_header.h"
 #include "host_common.h"
 #include "load_kernel_fw.h"
@@ -16,6 +20,7 @@
 #include "test_common.h"
 #include "vboot_audio.h"
 #include "vboot_common.h"
+#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_nvstorage.h"
 #include "vboot_struct.h"
@@ -26,6 +31,8 @@ static uint8_t shared_data[VB_SHARED_DATA_MIN_SIZE];
 static VbSharedDataHeader *shared = (VbSharedDataHeader *)shared_data;
 static GoogleBinaryBlockHeader gbb;
 static LoadKernelParams lkp;
+static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE];
+static struct vb2_context ctx;
 
 static int shutdown_request_calls_left;
 static int audio_looping_calls_left;
@@ -62,14 +69,6 @@ static void ResetMocks(void)
 	gbb.minor_version = GBB_MINOR_VER;
 	gbb.flags = 0;
 
-	/*
-	 * Only the outermost vboot_api_kernel call sets vboot_api_kernel's
-	 * vnc.  So clear it here too.
-	 */
-	memset(VbApiKernelGetVnc(), 0, sizeof(VbNvContext));
-	VbNvSetup(VbApiKernelGetVnc());
-	VbNvTeardown(VbApiKernelGetVnc()); /* So CRC gets generated */
-
 	memset(VbApiKernelGetFwmp(), 0, sizeof(struct RollbackSpaceFwmp));
 
 	memset(&shared_data, 0, sizeof(shared_data));
@@ -77,6 +76,12 @@ static void ResetMocks(void)
 
 	memset(&lkp, 0, sizeof(lkp));
 
+	memset(&ctx, 0, sizeof(ctx));
+	ctx.workbuf = workbuf;
+	ctx.workbuf_size = sizeof(workbuf);
+	vb2_init_context(&ctx);
+	vb2_nv_init(&ctx);
+
 	shutdown_request_calls_left = -1;
 	audio_looping_calls_left = 30;
 	vbtlk_retval = 1000;
@@ -178,14 +183,14 @@ int VbAudioLooping(VbAudioContext *audio)
 	return 1;
 }
 
-uint32_t VbTryLoadKernel(VbCommonParams *cparams, LoadKernelParams *p,
-                         uint32_t get_info_flags)
+uint32_t VbTryLoadKernel(struct vb2_context *ctx, VbCommonParams *cparams,
+			 LoadKernelParams *p, uint32_t get_info_flags)
 {
 	return vbtlk_retval + get_info_flags;
 }
 
-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)
 {
 	if (screens_count < ARRAY_SIZE(screens_displayed))
 		screens_displayed[screens_count++] = screen;
@@ -207,44 +212,44 @@ static void VbUserConfirmsTest(void)
 
 	ResetMocks();
 	shutdown_request_calls_left = 1;
-	TEST_EQ(VbUserConfirms(&cparams, 0), -1, "Shutdown requested");
+	TEST_EQ(VbUserConfirms(&ctx, &cparams, 0), -1, "Shutdown requested");
 
 	ResetMocks();
 	mock_keypress[0] = '\r';
-	TEST_EQ(VbUserConfirms(&cparams, 0), 1, "Enter");
+	TEST_EQ(VbUserConfirms(&ctx, &cparams, 0), 1, "Enter");
 
 	ResetMocks();
 	mock_keypress[0] = 0x1b;
-	TEST_EQ(VbUserConfirms(&cparams, 0), 0, "Esc");
+	TEST_EQ(VbUserConfirms(&ctx, &cparams, 0), 0, "Esc");
 
 	ResetMocks();
 	mock_keypress[0] = ' ';
 	shutdown_request_calls_left = 1;
-	TEST_EQ(VbUserConfirms(&cparams, VB_CONFIRM_SPACE_MEANS_NO), 0,
+	TEST_EQ(VbUserConfirms(&ctx, &cparams, VB_CONFIRM_SPACE_MEANS_NO), 0,
                 "Space means no");
 
 	ResetMocks();
 	mock_keypress[0] = ' ';
 	shutdown_request_calls_left = 1;
-	TEST_EQ(VbUserConfirms(&cparams, 0), -1, "Space ignored");
+	TEST_EQ(VbUserConfirms(&ctx, &cparams, 0), -1, "Space ignored");
 
 	ResetMocks();
 	mock_keypress[0] = '\r';
 	mock_keyflags[0] = VB_KEY_FLAG_TRUSTED_KEYBOARD;
-	TEST_EQ(VbUserConfirms(&cparams, VB_CONFIRM_MUST_TRUST_KEYBOARD),
+	TEST_EQ(VbUserConfirms(&ctx, &cparams, VB_CONFIRM_MUST_TRUST_KEYBOARD),
 		1, "Enter with trusted keyboard");
 
 	ResetMocks();
 	mock_keypress[0] = '\r';	/* untrusted */
 	mock_keypress[1] = ' ';
-	TEST_EQ(VbUserConfirms(&cparams,
+	TEST_EQ(VbUserConfirms(&ctx, &cparams,
 			       VB_CONFIRM_SPACE_MEANS_NO |
 			       VB_CONFIRM_MUST_TRUST_KEYBOARD),
 		0, "Untrusted keyboard");
 
 	ResetMocks();
 	mock_switches[0] = VB_INIT_FLAG_REC_BUTTON_PRESSED;
-	TEST_EQ(VbUserConfirms(&cparams,
+	TEST_EQ(VbUserConfirms(&ctx, &cparams,
 			       VB_CONFIRM_SPACE_MEANS_NO |
 			       VB_CONFIRM_MUST_TRUST_KEYBOARD),
 		1, "Recovery button");
@@ -256,7 +261,7 @@ static void VbUserConfirmsTest(void)
 	mock_keypress[3] = ' ';
 	mock_switches[0] = VB_INIT_FLAG_REC_BUTTON_PRESSED;
 	mock_switches_are_stuck = 1;
-	TEST_EQ(VbUserConfirms(&cparams,
+	TEST_EQ(VbUserConfirms(&ctx, &cparams,
 			       VB_CONFIRM_SPACE_MEANS_NO |
 			       VB_CONFIRM_MUST_TRUST_KEYBOARD),
 		0, "Recovery button stuck");
@@ -268,99 +273,99 @@ static void VbBootTest(void)
 {
 	ResetMocks();
 	VbExEcEnteringMode(0, VB_EC_NORMAL);
-	TEST_EQ(VbBootNormal(&cparams, &lkp), 1002, "VbBootNormal()");
+	TEST_EQ(VbBootNormal(&ctx, &cparams, &lkp), 1002, "VbBootNormal()");
 	TEST_EQ(VbGetMode(), VB_EC_NORMAL, "vboot_mode normal");
 }
 
 static void VbBootDevTest(void)
 {
-	uint32_t u;
-
 	printf("Testing VbBootDeveloper()...\n");
 
 	/* Proceed after timeout */
 	ResetMocks();
 	VbExEcEnteringMode(0, VB_EC_DEVELOPER);
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Timeout");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Timeout");
 	TEST_EQ(VbGetMode(), VB_EC_DEVELOPER, "vboot_mode developer");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_DEVELOPER_WARNING,
 		"  warning screen");
-	VbNvGet(VbApiKernelGetVnc(), VBNV_RECOVERY_REQUEST, &u);
-	TEST_EQ(u, 0, "  recovery reason");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST), 0,
+		"  recovery reason");
 	TEST_EQ(audio_looping_calls_left, 0, "  used up audio");
 
 	/* Proceed to legacy after timeout if GBB flag set */
 	ResetMocks();
 	gbb.flags |= GBB_FLAG_DEFAULT_DEV_BOOT_LEGACY |
 		     GBB_FLAG_FORCE_DEV_BOOT_LEGACY;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Timeout");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Timeout");
 	TEST_EQ(vbexlegacy_called, 1, "  try legacy");
 
 	/* Proceed to legacy after timeout if boot legacy and default boot
 	 * legacy are set */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_DEFAULT_BOOT,
-		VBNV_DEV_DEFAULT_BOOT_LEGACY);
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_BOOT_LEGACY, 1);
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Timeout");
+	vb2_nv_set(&ctx, VB2_NV_DEV_DEFAULT_BOOT,
+		   VBNV_DEV_DEFAULT_BOOT_LEGACY);
+	vb2_nv_set(&ctx, VB2_NV_DEV_BOOT_LEGACY, 1);
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Timeout");
 	TEST_EQ(vbexlegacy_called, 1, "  try legacy");
 
 	/* Proceed to legacy boot mode only if enabled */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_DEFAULT_BOOT,
-		VBNV_DEV_DEFAULT_BOOT_LEGACY);
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Timeout");
+	vb2_nv_set(&ctx, VB2_NV_DEV_DEFAULT_BOOT,
+		   VBNV_DEV_DEFAULT_BOOT_LEGACY);
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Timeout");
 	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
 
 	/* Proceed to usb after timeout if boot usb and default boot
 	 * usb are set */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_DEFAULT_BOOT,
-		VBNV_DEV_DEFAULT_BOOT_USB);
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_BOOT_USB, 1);
+	vb2_nv_set(&ctx, VB2_NV_DEV_DEFAULT_BOOT,
+		   VBNV_DEV_DEFAULT_BOOT_USB);
+	vb2_nv_set(&ctx, VB2_NV_DEV_BOOT_USB, 1);
 	vbtlk_retval = VBERROR_SUCCESS - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 0, "Ctrl+U USB");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 0, "Ctrl+U USB");
 
 	/* Proceed to usb boot mode only if enabled */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_DEFAULT_BOOT,
-		VBNV_DEV_DEFAULT_BOOT_USB);
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Timeout");
+	vb2_nv_set(&ctx, VB2_NV_DEV_DEFAULT_BOOT,
+		   VBNV_DEV_DEFAULT_BOOT_USB);
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Timeout");
 
 	/* If no USB tries fixed disk */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_BOOT_USB, 1);
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_DEFAULT_BOOT,
-		VBNV_DEV_DEFAULT_BOOT_USB);
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+U enabled");
+	vb2_nv_set(&ctx, VB2_NV_DEV_BOOT_USB, 1);
+	vb2_nv_set(&ctx, VB2_NV_DEV_DEFAULT_BOOT,
+		   VBNV_DEV_DEFAULT_BOOT_USB);
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Ctrl+U enabled");
 	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
 
 	/* Up arrow is uninteresting / passed to VbCheckDisplayKey() */
 	ResetMocks();
 	mock_keypress[0] = VB_KEY_UP;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Up arrow");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Up arrow");
 
 	/* Shutdown requested in loop */
 	ResetMocks();
 	shutdown_request_calls_left = 2;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Shutdown requested");
 	TEST_NEQ(audio_looping_calls_left, 0, "  aborts audio");
 
 	/* Space goes straight to recovery if no virtual dev switch */
 	ResetMocks();
 	mock_keypress[0] = ' ';
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_LOAD_KERNEL_RECOVERY,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp),
+		VBERROR_LOAD_KERNEL_RECOVERY,
 		"Space = recovery");
-	VbNvGet(VbApiKernelGetVnc(), VBNV_RECOVERY_REQUEST, &u);
-	TEST_EQ(u, VBNV_RECOVERY_RW_DEV_SCREEN, "  recovery reason");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST),
+		VBNV_RECOVERY_RW_DEV_SCREEN, "  recovery reason");
 
 	/* Space asks to disable virtual dev switch */
 	ResetMocks();
 	shared->flags = VBSD_HONOR_VIRT_DEV_SWITCH | VBSD_BOOT_DEV_SWITCH_ON;
 	mock_keypress[0] = ' ';
 	mock_keypress[1] = '\r';
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_REBOOT_REQUIRED,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), VBERROR_REBOOT_REQUIRED,
 		"Space = tonorm");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_DEVELOPER_WARNING,
 		"  warning screen");
@@ -368,8 +373,8 @@ static void VbBootDevTest(void)
 		"  tonorm screen");
 	TEST_EQ(screens_displayed[2], VB_SCREEN_TO_NORM_CONFIRMED,
 		"  confirm screen");
-	VbNvGet(VbApiKernelGetVnc(), VBNV_DISABLE_DEV_REQUEST, &u);
-	TEST_EQ(u, 1, "  disable dev request");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DISABLE_DEV_REQUEST), 1,
+		"  disable dev request");
 
 	/* Space-space doesn't disable it */
 	ResetMocks();
@@ -377,7 +382,7 @@ static void VbBootDevTest(void)
 	mock_keypress[0] = ' ';
 	mock_keypress[1] = ' ';
 	mock_keypress[2] = 0x1b;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Space-space");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Space-space");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_DEVELOPER_WARNING,
 		"  warning screen");
 	TEST_EQ(screens_displayed[1], VB_SCREEN_DEVELOPER_TO_NORM,
@@ -390,7 +395,7 @@ static void VbBootDevTest(void)
 	shared->flags = VBSD_HONOR_VIRT_DEV_SWITCH | VBSD_BOOT_DEV_SWITCH_ON;
 	mock_keypress[0] = '\r';
 	mock_keypress[1] = '\r';
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Enter ignored");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Enter ignored");
 
 	/* Enter does if GBB flag set */
 	ResetMocks();
@@ -398,7 +403,7 @@ static void VbBootDevTest(void)
 	gbb.flags |= GBB_FLAG_ENTER_TRIGGERS_TONORM;
 	mock_keypress[0] = '\r';
 	mock_keypress[1] = '\r';
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_REBOOT_REQUIRED,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), VBERROR_REBOOT_REQUIRED,
 		"Enter = tonorm");
 
 	/* Tonorm ignored if GBB forces dev switch on */
@@ -407,14 +412,16 @@ static void VbBootDevTest(void)
 	gbb.flags |= GBB_FLAG_FORCE_DEV_SWITCH_ON;
 	mock_keypress[0] = ' ';
 	mock_keypress[1] = '\r';
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Can't tonorm gbb-dev");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002,
+		"Can't tonorm gbb-dev");
 
 	/* Shutdown requested at tonorm screen */
 	ResetMocks();
 	shared->flags = VBSD_HONOR_VIRT_DEV_SWITCH | VBSD_BOOT_DEV_SWITCH_ON;
 	mock_keypress[0] = ' ';
 	shutdown_request_calls_left = 2;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Shutdown requested at tonorm");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_DEVELOPER_WARNING,
 		"  warning screen");
@@ -424,9 +431,9 @@ static void VbBootDevTest(void)
 	/* Ctrl+D dismisses warning */
 	ResetMocks();
 	mock_keypress[0] = 0x04;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+D");
-	VbNvGet(VbApiKernelGetVnc(), VBNV_RECOVERY_REQUEST, &u);
-	TEST_EQ(u, 0, "  recovery reason");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Ctrl+D");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST), 0,
+		"  recovery reason");
 	TEST_NEQ(audio_looping_calls_left, 0, "  aborts audio");
 	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
 
@@ -434,67 +441,70 @@ static void VbBootDevTest(void)
 	ResetMocks();
 	mock_keypress[0] = 0x04;
 	gbb.flags |= GBB_FLAG_DEFAULT_DEV_BOOT_LEGACY;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+D");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Ctrl+D");
 	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
 
 	/* Ctrl+L tries legacy boot mode only if enabled */
 	ResetMocks();
 	mock_keypress[0] = 0x0c;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+L normal");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Ctrl+L normal");
 	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
 
 	ResetMocks();
 	gbb.flags |= GBB_FLAG_FORCE_DEV_BOOT_LEGACY;
 	mock_keypress[0] = 0x0c;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+L force legacy");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002,
+		"Ctrl+L force legacy");
 	TEST_EQ(vbexlegacy_called, 1, "  try legacy");
 
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_BOOT_LEGACY, 1);
+	vb2_nv_set(&ctx, VB2_NV_DEV_BOOT_LEGACY, 1);
 	mock_keypress[0] = 0x0c;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+L nv legacy");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002,
+		"Ctrl+L nv legacy");
 	TEST_EQ(vbexlegacy_called, 1, "  try legacy");
 
 	ResetMocks();
 	VbApiKernelGetFwmp()->flags |= FWMP_DEV_ENABLE_LEGACY;
 	mock_keypress[0] = 0x0c;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+L fwmp legacy");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002,
+		"Ctrl+L fwmp legacy");
 	TEST_EQ(vbexlegacy_called, 1, "  fwmp legacy");
 
 	/* Ctrl+U boots USB only if enabled */
 	ResetMocks();
 	mock_keypress[0] = 0x15;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+U normal");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Ctrl+U normal");
 
 	/* Ctrl+U enabled, with good USB boot */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_BOOT_USB, 1);
+	vb2_nv_set(&ctx, VB2_NV_DEV_BOOT_USB, 1);
 	mock_keypress[0] = 0x15;
 	vbtlk_retval = VBERROR_SUCCESS - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 0, "Ctrl+U USB");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 0, "Ctrl+U USB");
 
 	/* Ctrl+U enabled via GBB */
 	ResetMocks();
 	gbb.flags |= GBB_FLAG_FORCE_DEV_BOOT_USB;
 	mock_keypress[0] = 0x15;
 	vbtlk_retval = VBERROR_SUCCESS - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 0, "Ctrl+U force USB");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 0, "Ctrl+U force USB");
 
 	/* Ctrl+U enabled via FWMP */
 	ResetMocks();
 	VbApiKernelGetFwmp()->flags |= FWMP_DEV_ENABLE_USB;
 	mock_keypress[0] = 0x15;
 	vbtlk_retval = VBERROR_SUCCESS - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 0, "Ctrl+U force USB");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 0, "Ctrl+U force USB");
 
 	/* If no USB, eventually times out and tries fixed disk */
 	ResetMocks();
-	VbNvSet(VbApiKernelGetVnc(), VBNV_DEV_BOOT_USB, 1);
+	vb2_nv_set(&ctx, VB2_NV_DEV_BOOT_USB, 1);
 	mock_keypress[0] = 0x15;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), 1002, "Ctrl+U enabled");
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), 1002, "Ctrl+U enabled");
 	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
-	VbNvGet(VbApiKernelGetVnc(), VBNV_RECOVERY_REQUEST, &u);
-	TEST_EQ(u, 0, "  recovery reason");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST), 0,
+		"  recovery reason");
 	TEST_EQ(audio_looping_calls_left, 0, "  used up audio");
 
 	/* If dev mode is disabled, goes to TONORM screen repeatedly */
@@ -502,7 +512,7 @@ static void VbBootDevTest(void)
 	VbApiKernelGetFwmp()->flags |= FWMP_DEV_DISABLE_BOOT;
 	mock_keypress[0] = '\x1b';  /* Just causes TONORM again */
 	mock_keypress[1] = '\r';
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_REBOOT_REQUIRED,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp), VBERROR_REBOOT_REQUIRED,
 		"FWMP dev disabled");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_DEVELOPER_TO_NORM,
 		"  tonorm screen");
@@ -510,15 +520,16 @@ static void VbBootDevTest(void)
 		"  tonorm screen");
 	TEST_EQ(screens_displayed[2], VB_SCREEN_TO_NORM_CONFIRMED,
 		"  confirm screen");
-	VbNvGet(VbApiKernelGetVnc(), VBNV_DISABLE_DEV_REQUEST, &u);
-	TEST_EQ(u, 1, "  disable dev request");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DISABLE_DEV_REQUEST), 1,
+		"  disable dev request");
 
 	/* Shutdown requested when dev disabled */
 	ResetMocks();
 	shared->flags = VBSD_HONOR_VIRT_DEV_SWITCH | VBSD_BOOT_DEV_SWITCH_ON;
 	VbApiKernelGetFwmp()->flags |= FWMP_DEV_DISABLE_BOOT;
 	shutdown_request_calls_left = 1;
-	TEST_EQ(VbBootDeveloper(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootDeveloper(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Shutdown requested when dev disabled");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_DEVELOPER_TO_NORM,
 		"  tonorm screen");
@@ -528,20 +539,19 @@ static void VbBootDevTest(void)
 
 static void VbBootRecTest(void)
 {
-	uint32_t u;
-
 	printf("Testing VbBootRecovery()...\n");
 
 	/* Shutdown requested in loop */
 	ResetMocks();
 	shutdown_request_calls_left = 10;
 	VbExEcEnteringMode(0, VB_EC_RECOVERY);
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Shutdown requested");
 	TEST_EQ(VbGetMode(), VB_EC_RECOVERY, "vboot_mode recovery");
 
-	VbNvGet(VbApiKernelGetVnc(), VBNV_RECOVERY_REQUEST, &u);
-	TEST_EQ(u, 0, "  recovery reason");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST), 0,
+		"  recovery reason");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_OS_BROKEN,
 		"  broken screen");
 
@@ -552,7 +562,8 @@ static void VbBootRecTest(void)
 	mock_num_disks[1] = 1;
 	mock_num_disks[2] = 1;
 	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Remove");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_OS_BROKEN,
 		"  broken screen");
@@ -564,7 +575,8 @@ static void VbBootRecTest(void)
 	mock_num_disks[1] = 1;
 	shared->flags |= VBSD_BOOT_DEV_SWITCH_ON;
 	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"No remove in dev");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_INSERT,
 		"  insert screen");
@@ -576,7 +588,8 @@ static void VbBootRecTest(void)
 	mock_num_disks[1] = 1;
 	shared->flags |= VBSD_BOOT_REC_SWITCH_ON;
 	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"No remove in rec");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_INSERT,
 		"  insert screen");
@@ -587,7 +600,8 @@ static void VbBootRecTest(void)
 	mock_num_disks[0] = 0;
 	mock_num_disks[1] = 1;
 	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Remove");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_OS_BROKEN,
 		"  broken screen");
@@ -598,7 +612,8 @@ static void VbBootRecTest(void)
 	mock_num_disks[0] = -1;
 	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
 	shutdown_request_calls_left = 10;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Bad disk count");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_OS_BROKEN,
 		"  broken screen");
@@ -609,7 +624,8 @@ static void VbBootRecTest(void)
 	shutdown_request_calls_left = 100;
 	mock_keypress[0] = 0x04;
 	trust_ec = 0;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Ctrl+D ignored if EC not trusted");
 	TEST_EQ(virtdev_set, 0, "  virtual dev mode off");
 	TEST_NEQ(screens_displayed[1], VB_SCREEN_RECOVERY_TO_DEV,
@@ -621,7 +637,8 @@ static void VbBootRecTest(void)
 	trust_ec = 1;
 	shutdown_request_calls_left = 100;
 	mock_keypress[0] = 0x04;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Ctrl+D ignored if already in dev mode");
 	TEST_EQ(virtdev_set, 0, "  virtual dev mode off");
 	TEST_NEQ(screens_displayed[1], VB_SCREEN_RECOVERY_TO_DEV,
@@ -632,7 +649,8 @@ static void VbBootRecTest(void)
 	trust_ec = 1;
 	shutdown_request_calls_left = 100;
 	mock_keypress[0] = 0x04;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Ctrl+D ignored if recovery not manually triggered");
 	TEST_EQ(virtdev_set, 0, "  virtual dev mode off");
 	TEST_NEQ(screens_displayed[1], VB_SCREEN_RECOVERY_TO_DEV,
@@ -643,7 +661,8 @@ static void VbBootRecTest(void)
 	trust_ec = 1;
 	shutdown_request_calls_left = 100;
 	mock_keypress[0] = 0x04;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Ctrl+D ignored if no virtual dev switch");
 	TEST_EQ(virtdev_set, 0, "  virtual dev mode off");
 	TEST_NEQ(screens_displayed[1], VB_SCREEN_RECOVERY_TO_DEV,
@@ -658,7 +677,8 @@ static void VbBootRecTest(void)
 	shutdown_request_calls_left = 100;
 	mock_keypress[0] = 0x04;
 	mock_switches[0] = VB_INIT_FLAG_REC_BUTTON_PRESSED;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Ctrl+D ignored if phys rec button is still pressed");
 	TEST_NEQ(screens_displayed[1], VB_SCREEN_RECOVERY_TO_DEV,
 		 "  todev screen");
@@ -671,7 +691,8 @@ static void VbBootRecTest(void)
 	trust_ec = 1;
 	mock_keypress[0] = 0x04;
 	mock_keypress[1] = ' ';
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_SHUTDOWN_REQUESTED,
 		"Ctrl+D todev abort");
 	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_INSERT,
 		"  insert screen");
@@ -690,7 +711,7 @@ static void VbBootRecTest(void)
 	mock_keypress[0] = 0x04;
 	mock_keypress[1] = '\r';
 	mock_keyflags[1] = VB_KEY_FLAG_TRUSTED_KEYBOARD;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_REBOOT_REQUIRED,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp), VBERROR_REBOOT_REQUIRED,
 		"Ctrl+D todev confirm");
 	TEST_EQ(virtdev_set, 1, "  virtual dev mode on");
 
@@ -704,7 +725,8 @@ static void VbBootRecTest(void)
 	mock_keypress[1] = '\r';
 	mock_keyflags[1] = VB_KEY_FLAG_TRUSTED_KEYBOARD;
 	virtdev_retval = VBERROR_SIMULATED;
-	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_TPM_SET_BOOT_MODE_STATE,
+	TEST_EQ(VbBootRecovery(&ctx, &cparams, &lkp),
+		VBERROR_TPM_SET_BOOT_MODE_STATE,
 		"Ctrl+D todev failure");
 
 	printf("...done.\n");
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
index 7e4b47dd..decc7137 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -81,8 +81,8 @@ VbError_t VbExNvStorageWrite(const uint8_t *buf)
 	return VBERROR_SUCCESS;
 }
 
-VbError_t VbEcSoftwareSync(int devidx, VbCommonParams *cparams,
-			   VbNvContext *vnc)
+VbError_t VbEcSoftwareSync(struct vb2_context *ctx, int devidx,
+			   VbCommonParams *cparams)
 {
 	return ecsync_retval;
 }
@@ -111,7 +111,8 @@ uint32_t RollbackFwmpRead(struct RollbackSpaceFwmp *fwmp)
 	return rfr_retval;
 }
 
-VbError_t VbBootNormal(VbCommonParams *cparams, LoadKernelParams *p)
+VbError_t VbBootNormal(struct vb2_context *ctx, VbCommonParams *cparams,
+		       LoadKernelParams *p)
 {
 	shared->kernel_version_tpm = new_version;
 
@@ -121,7 +122,8 @@ VbError_t VbBootNormal(VbCommonParams *cparams, LoadKernelParams *p)
 	return vbboot_retval;
 }
 
-VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
+VbError_t VbBootDeveloper(struct vb2_context *ctx, VbCommonParams *cparams,
+			  LoadKernelParams *p)
 {
 	shared->kernel_version_tpm = new_version;
 
@@ -131,7 +133,8 @@ VbError_t VbBootDeveloper(VbCommonParams *cparams, LoadKernelParams *p)
 	return vbboot_retval;
 }
 
-VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
+VbError_t VbBootRecovery(struct vb2_context *ctx, VbCommonParams *cparams,
+			 LoadKernelParams *p)
 {
 	shared->kernel_version_tpm = new_version;
 
diff --git a/tests/vboot_api_kernel_tests.c b/tests/vboot_api_kernel_tests.c
index 0baf4dff..0a59160b 100644
--- a/tests/vboot_api_kernel_tests.c
+++ b/tests/vboot_api_kernel_tests.c
@@ -10,6 +10,9 @@
 #include <stdio.h>
 #include <stdlib.h>
 
+#include "2sysincludes.h"
+#include "2api.h"
+#include "2nvstorage.h"
 #include "gbb_header.h"
 #include "load_kernel_fw.h"
 #include "rollback_index.h"
@@ -187,6 +190,8 @@ static const char *got_find_disk;
 static const char *got_load_disk;
 static uint32_t got_return_val;
 static uint32_t got_external_mismatch;
+static VbNvContext vnc;
+static struct vb2_context ctx;
 
 /**
  * Reset mock data (for use before each test)
@@ -196,6 +201,7 @@ static void ResetMocks(int i)
 	memset(&lkparams, 0, sizeof(lkparams));
 	memset(&mock_disks, 0, sizeof(mock_disks));
 	load_kernel_calls = 0;
+	memset(&vnc, 0, sizeof(vnc));
 
 	got_recovery_request_val = VBNV_RECOVERY_NOT_REQUESTED;
 	got_find_disk = 0;
@@ -275,7 +281,8 @@ VbError_t VbExDiskFreeInfo(VbDiskInfo *infos,
 	return VBERROR_SUCCESS;
 }
 
-VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
+VbError_t LoadKernel(struct vb2_context *ctx, LoadKernelParams *params,
+		     VbCommonParams *cparams)
 {
 	got_find_disk = (const char *)params->disk_handle;
 	VBDEBUG(("%s(%d): got_find_disk = %s\n", __FUNCTION__,
@@ -287,12 +294,13 @@ VbError_t LoadKernel(LoadKernelParams *params, VbCommonParams *cparams)
 	return t->loadkernel_return_val[load_kernel_calls++];
 }
 
-int VbNvSet(VbNvContext *context, VbNvParam param, uint32_t value)
+void vb2_nv_set(struct vb2_context *ctx,
+		enum vb2_nv_param param,
+		uint32_t value)
 {
 	VBDEBUG(("%s(): got_recovery_request_val = %d (0x%x)\n", __FUNCTION__,
 		 value, value));
 	got_recovery_request_val = value;
-	return 0;
 }
 
 /****************************************************************************/
@@ -305,7 +313,7 @@ static void VbTryLoadKernelTest(void)
 	for (i = 0; i < num_tests; i++) {
 		printf("Test case: %s ...\n", test[i].name);
 		ResetMocks(i);
-		TEST_EQ(VbTryLoadKernel(0, &lkparams, test[i].want_flags),
+		TEST_EQ(VbTryLoadKernel(&ctx, 0, &lkparams, test[i].want_flags),
 			t->expected_return_val, "  return value");
 		TEST_EQ(got_recovery_request_val,
 			t->expected_recovery_request_val, "  recovery_request");
diff --git a/tests/vboot_display_tests.c b/tests/vboot_display_tests.c
index ff036d23..33421b5b 100644
--- a/tests/vboot_display_tests.c
+++ b/tests/vboot_display_tests.c
@@ -10,6 +10,10 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "2sysincludes.h"
+#include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
 #include "bmpblk_font.h"
 #include "gbb_header.h"
 #include "host_common.h"
@@ -22,13 +26,14 @@
 
 /* Mock data */
 static VbCommonParams cparams;
-static VbNvContext vnc;
 static uint8_t shared_data[VB_SHARED_DATA_MIN_SIZE];
 static VbSharedDataHeader *shared = (VbSharedDataHeader *)shared_data;
 static char gbb_data[4096 + sizeof(GoogleBinaryBlockHeader)];
 static GoogleBinaryBlockHeader *gbb = (GoogleBinaryBlockHeader *)gbb_data;
 static BmpBlockHeader *bhdr;
 static char debug_info[4096];
+static struct vb2_context ctx;
+static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE];
 
 /* Reset mock data (for use before each test) */
 static void ResetMocks(void)
@@ -75,9 +80,11 @@ static void ResetMocks(void)
 	gbb_used += 64;
 	memcpy(cparams.gbb, gbb, sizeof(*gbb));
 
-	memset(&vnc, 0, sizeof(vnc));
-	VbNvSetup(&vnc);
-	VbNvTeardown(&vnc);                   /* So CRC gets generated */
+	memset(&ctx, 0, sizeof(ctx));
+	ctx.workbuf = workbuf;
+	ctx.workbuf_size = sizeof(workbuf);
+	vb2_init_context(&ctx);
+	vb2_nv_init(&ctx);
 
 	memset(&shared_data, 0, sizeof(shared_data));
 	VbSharedDataInit(shared, sizeof(shared_data));
@@ -136,7 +143,7 @@ static void DebugInfoTest(void)
 
 	/* Display debug info */
 	ResetMocks();
-	VbDisplayDebugInfo(&cparams, &vnc);
+	VbDisplayDebugInfo(&ctx, &cparams);
 	TEST_NEQ(*debug_info, '\0', "Some debug info was displayed");
 	VbApiKernelFree(&cparams);
 }
@@ -169,44 +176,43 @@ static void LocalizationTest(void)
 /* Test display key checking */
 static void DisplayKeyTest(void)
 {
-	uint32_t u;
-
 	ResetMocks();
-	VbCheckDisplayKey(&cparams, 'q', &vnc);
+	VbCheckDisplayKey(&ctx, &cparams, 'q');
 	TEST_EQ(*debug_info, '\0', "DisplayKey q = does nothing");
 	VbApiKernelFree(&cparams);
 
 	ResetMocks();
-	VbCheckDisplayKey(&cparams, '\t', &vnc);
+	VbCheckDisplayKey(&ctx, &cparams, '\t');
 	TEST_NEQ(*debug_info, '\0', "DisplayKey tab = display");
 	VbApiKernelFree(&cparams);
 
 	/* Toggle localization */
 	ResetMocks();
-	VbNvSet(&vnc, VBNV_LOCALIZATION_INDEX, 0);
-	VbNvTeardown(&vnc);
-	VbCheckDisplayKey(&cparams, VB_KEY_DOWN, &vnc);
-	VbNvGet(&vnc, VBNV_LOCALIZATION_INDEX, &u);
-	TEST_EQ(u, 2, "DisplayKey up");
-	VbCheckDisplayKey(&cparams, VB_KEY_LEFT, &vnc);
-	VbNvGet(&vnc, VBNV_LOCALIZATION_INDEX, &u);
-	TEST_EQ(u, 1, "DisplayKey left");
-	VbCheckDisplayKey(&cparams, VB_KEY_RIGHT, &vnc);
-	VbNvGet(&vnc, VBNV_LOCALIZATION_INDEX, &u);
-	TEST_EQ(u, 2, "DisplayKey right");
-	VbCheckDisplayKey(&cparams, VB_KEY_UP, &vnc);
-	VbNvGet(&vnc, VBNV_LOCALIZATION_INDEX, &u);
-	TEST_EQ(u, 0, "DisplayKey up");
+	vb2_nv_set(&ctx, VB2_NV_LOCALIZATION_INDEX, 0);
+	VbCheckDisplayKey(&ctx, &cparams, VB_KEY_DOWN);
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX), 2,
+		"DisplayKey up");
+	VbCheckDisplayKey(&ctx, &cparams, VB_KEY_LEFT);
+	vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX);
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX), 1,
+		"DisplayKey left");
+	VbCheckDisplayKey(&ctx, &cparams, VB_KEY_RIGHT);
+	vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX);
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX), 2,
+		"DisplayKey right");
+	VbCheckDisplayKey(&ctx, &cparams, VB_KEY_UP);
+	vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX);
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX), 0,
+		"DisplayKey up");
 	VbApiKernelFree(&cparams);
 
 	/* Reset localization if localization count is invalid */
 	ResetMocks();
-	VbNvSet(&vnc, VBNV_LOCALIZATION_INDEX, 1);
-	VbNvTeardown(&vnc);
+	vb2_nv_set(&ctx, VB2_NV_LOCALIZATION_INDEX, 1);
 	bhdr->signature[0] ^= 0x5a;
-	VbCheckDisplayKey(&cparams, VB_KEY_UP, &vnc);
-	VbNvGet(&vnc, VBNV_LOCALIZATION_INDEX, &u);
-	TEST_EQ(u, 0, "DisplayKey invalid");
+	VbCheckDisplayKey(&ctx, &cparams, VB_KEY_UP);
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_LOCALIZATION_INDEX), 0,
+		"DisplayKey invalid");
 	VbApiKernelFree(&cparams);
 }
 
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index d3a25d58..2a472d38 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -11,7 +11,9 @@
 #include <string.h>
 
 #include "2sysincludes.h"
+#include "2api.h"
 #include "2common.h"
+#include "2misc.h"
 #include "2sha.h"
 #include "cgptlib.h"
 #include "cgptlib_internal.h"
@@ -74,6 +76,8 @@ static GptHeader *mock_gpt_primary =
 static GptHeader *mock_gpt_secondary =
 	(GptHeader*)&mock_disk[MOCK_SECTOR_SIZE * (MOCK_SECTOR_COUNT - 1)];
 static uint8_t mock_digest[VB2_SHA256_DIGEST_SIZE] = {12, 34, 56, 78};
+static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE];
+static struct vb2_context ctx;
 
 /**
  * Prepare a valid GPT header that will pass CheckHeader() tests
@@ -182,6 +186,12 @@ static void ResetMocks(void)
 	mock_parts[0].start = 100;
 	mock_parts[0].size = 150;  /* 75 KB */
 	mock_part_next = 0;
+
+	memset(&ctx, 0, sizeof(ctx));
+	memcpy(ctx.nvdata, vnc.raw, VB2_NVDATA_SIZE);
+	ctx.workbuf = workbuf;
+	ctx.workbuf_size = sizeof(workbuf);
+	// TODO: more workbuf fields - flags, secdata, secdatak
 }
 
 /* Mocks */
@@ -564,6 +574,23 @@ static void ReadWriteGptTest(void)
 
 }
 
+static void TestLoadKernel(int expect_retval, char *test_name)
+{
+	memcpy(ctx.nvdata, vnc.raw, VB2_NVDATA_SIZE);
+	if (lkp.boot_flags & BOOT_FLAG_RECOVERY)
+		ctx.flags |= VB2_CONTEXT_RECOVERY_MODE;
+	if (lkp.boot_flags & BOOT_FLAG_DEVELOPER)
+		ctx.flags |= VB2_CONTEXT_DEVELOPER_MODE;
+
+	TEST_EQ(LoadKernel(&ctx, &lkp, &cparams), expect_retval, test_name);
+
+	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;
+	}
+}
+
 /**
  * Trivial invalid calls to LoadKernel()
  */
@@ -571,14 +598,12 @@ static void InvalidParamsTest(void)
 {
 	ResetMocks();
 	gpt_init_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_NO_KERNEL_FOUND,
-		"Bad GPT");
+	TestLoadKernel(VBERROR_NO_KERNEL_FOUND, "Bad GPT");
 
 	/* This causes the stream open call to fail */
 	ResetMocks();
 	lkp.disk_handle = NULL;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Bad disk handle");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Bad disk handle");
 }
 
 static void LoadKernelTest(void)
@@ -587,8 +612,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 
-	u = LoadKernel(&lkp, &cparams);
-	TEST_EQ(u, 0, "First kernel good");
+	TestLoadKernel(0, "First kernel good");
 	TEST_EQ(lkp.partition_number, 1, "  part num");
 	TEST_EQ(lkp.bootloader_address, 0xbeadd008, "  bootloader addr");
 	TEST_EQ(lkp.bootloader_size, 0x1234, "  bootloader size");
@@ -600,46 +624,44 @@ static void LoadKernelTest(void)
 	ResetMocks();
 	mock_parts[1].start = 300;
 	mock_parts[1].size = 150;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Two good kernels");
+	TestLoadKernel(0, "Two good kernels");
 	TEST_EQ(lkp.partition_number, 1, "  part num");
 	TEST_EQ(mock_part_next, 1, "  didn't read second one");
 
 	/* Fail if no kernels found */
 	ResetMocks();
 	mock_parts[0].size = 0;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_NO_KERNEL_FOUND, "No kernels");
+	TestLoadKernel(VBERROR_NO_KERNEL_FOUND, "No kernels");
 	VbNvGet(&vnc, VBNV_RECOVERY_REQUEST, &u);
 	TEST_EQ(u, VBNV_RECOVERY_RW_NO_OS, "  recovery request");
 
 	/* Skip kernels which are too small */
 	ResetMocks();
 	mock_parts[0].size = 10;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND, "Too small");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Too small");
 	VbNvGet(&vnc, VBNV_RECOVERY_REQUEST, &u);
 	TEST_EQ(u, VBNV_RECOVERY_RW_INVALID_OS, "  recovery request");
 
 	ResetMocks();
 	disk_read_to_fail = 100;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail reading kernel start");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Fail reading kernel start");
 
 	ResetMocks();
 	key_block_verify_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail key block sig");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Fail key block sig");
 
 	/* In dev mode, fail if hash is bad too */
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_DEVELOPER;
 	key_block_verify_fail = 2;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail key block dev hash");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Fail key block dev hash");
 
 	/* But just bad sig is ok */
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_DEVELOPER;
 	key_block_verify_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Succeed key block dev sig");
+	TestLoadKernel(0, "Succeed key block dev sig");
 
 	/* In dev mode and requiring signed kernel, fail if sig is bad */
 	ResetMocks();
@@ -647,101 +669,97 @@ static void LoadKernelTest(void)
 	VbNvSet(&vnc, VBNV_DEV_BOOT_SIGNED_ONLY, 1);
 	VbNvTeardown(&vnc);
 	key_block_verify_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail key block dev sig");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Fail key block dev sig");
 
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_DEVELOPER;
 	lkp.fwmp = &fwmp;
 	fwmp.flags |= FWMP_DEV_ENABLE_OFFICIAL_ONLY;
 	key_block_verify_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail key block dev sig fwmp");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Fail key block dev sig fwmp");
 
 	/* Check key block flag mismatches */
 	ResetMocks();
 	kbh.key_block_flags =
 		KEY_BLOCK_FLAG_RECOVERY_0 | KEY_BLOCK_FLAG_DEVELOPER_1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Key block dev flag mismatch");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Key block dev flag mismatch");
 
 	ResetMocks();
 	kbh.key_block_flags =
 		KEY_BLOCK_FLAG_RECOVERY_1 | KEY_BLOCK_FLAG_DEVELOPER_0;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Key block rec flag mismatch");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Key block rec flag mismatch");
 
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_RECOVERY;
 	kbh.key_block_flags =
 		KEY_BLOCK_FLAG_RECOVERY_1 | KEY_BLOCK_FLAG_DEVELOPER_1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Key block recdev flag mismatch");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Key block recdev flag mismatch");
 
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_RECOVERY | BOOT_FLAG_DEVELOPER;
 	kbh.key_block_flags =
 		KEY_BLOCK_FLAG_RECOVERY_1 | KEY_BLOCK_FLAG_DEVELOPER_0;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Key block rec!dev flag mismatch");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Key block rec!dev flag mismatch");
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Key block kernel key rollback");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Key block kernel key rollback");
 
 	ResetMocks();
 	kbh.data_key.key_version = 0x10000;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Key block kernel key version too big");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Key block kernel key version too big");
 
 	ResetMocks();
 	kbh.data_key.key_version = 3;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Key block version roll forward");
+	TestLoadKernel(0, "Key block version roll forward");
 	TEST_EQ(shared->kernel_version_tpm, 0x30001, "  shared version");
 
 	ResetMocks();
 	kbh.data_key.key_version = 3;
 	mock_parts[1].start = 300;
 	mock_parts[1].size = 150;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Two kernels roll forward");
+	TestLoadKernel(0, "Two kernels roll forward");
 	TEST_EQ(mock_part_next, 2, "  read both");
 	TEST_EQ(shared->kernel_version_tpm, 0x30001, "  shared version");
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
 	lkp.boot_flags |= BOOT_FLAG_DEVELOPER;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Key version ignored in dev mode");
+	TestLoadKernel(0, "Key version ignored in dev mode");
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
 	lkp.boot_flags |= BOOT_FLAG_RECOVERY;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Key version ignored in rec mode");
+	TestLoadKernel(0, "Key version ignored in rec mode");
 
 	ResetMocks();
 	unpack_key_fail = 2;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Bad data key");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Bad data key");
 
 	ResetMocks();
 	preamble_verify_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Bad preamble");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Bad preamble");
 
 	ResetMocks();
 	kph.kernel_version = 0;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Kernel version rollback");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Kernel version rollback");
 
 	ResetMocks();
 	kph.kernel_version = 0;
 	lkp.boot_flags |= BOOT_FLAG_DEVELOPER;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Kernel version ignored in dev mode");
+	TestLoadKernel(0, "Kernel version ignored in dev mode");
 
 	ResetMocks();
 	kph.kernel_version = 0;
 	lkp.boot_flags |= BOOT_FLAG_RECOVERY;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Kernel version ignored in rec mode");
+	TestLoadKernel(0, "Kernel version ignored in rec mode");
 
 	/* Check developer key hash - bad */
 	ResetMocks();
@@ -749,68 +767,66 @@ static void LoadKernelTest(void)
 	lkp.fwmp = &fwmp;
 	fwmp.flags |= FWMP_DEV_USE_KEY_HASH;
 	fwmp.dev_key_hash[0]++;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail key block dev fwmp hash");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Fail key block dev fwmp hash");
 
 	/* Check developer key hash - good */
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_DEVELOPER;
 	lkp.fwmp = &fwmp;
 	fwmp.flags |= FWMP_DEV_USE_KEY_HASH;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Good key block dev fwmp hash");
+	TestLoadKernel(0, "Good key block dev fwmp hash");
 
 	ResetMocks();
 	kph.preamble_size |= 0x07;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Kernel body offset");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Kernel body offset");
 
 	ResetMocks();
 	kph.preamble_size += 65536;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Kernel body offset huge");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Kernel body offset huge");
 
 	/* Check getting kernel load address from header */
 	ResetMocks();
 	kph.body_load_address = (size_t)kernel_buffer;
 	lkp.kernel_buffer = NULL;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Get load address from preamble");
+	TestLoadKernel(0, "Get load address from preamble");
 	TEST_PTR_EQ(lkp.kernel_buffer, kernel_buffer, "  address");
 	/* Size is rounded up to nearest sector */
 	TEST_EQ(lkp.kernel_buffer_size, 70144, "  size");
 
 	ResetMocks();
 	lkp.kernel_buffer_size = 8192;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Kernel too big for buffer");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Kernel too big for buffer");
 
 	ResetMocks();
 	mock_parts[0].size = 130;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Kernel too big for partition");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Kernel too big for partition");
 
 	ResetMocks();
 	kph.body_signature.data_size = 8192;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Kernel tiny");
+	TestLoadKernel(0, "Kernel tiny");
 
 	ResetMocks();
 	disk_read_to_fail = 228;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
-		"Fail reading kernel data");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND,
+		       "Fail reading kernel data");
 
 	ResetMocks();
 	verify_data_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,	"Bad data");
+	TestLoadKernel(VBERROR_INVALID_KERNEL_FOUND, "Bad data");
 
 	/* Check that EXTERNAL_GPT flag makes it down */
 	ResetMocks();
 	lkp.boot_flags |= BOOT_FLAG_EXTERNAL_GPT;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Succeed external GPT");
+	TestLoadKernel(0, "Succeed external GPT");
 	TEST_EQ(gpt_flag_external, 1, "GPT was external");
 
 	/* Check recovery from unreadble primary GPT */
 	ResetMocks();
 	disk_read_to_fail = 1;
-	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Can't read disk");
+	TestLoadKernel(0, "Can't read disk");
 }
 
 int main(void)
diff --git a/tests/verify_kernel.c b/tests/verify_kernel.c
index 9fc58f33..9a7ee1e9 100644
--- a/tests/verify_kernel.c
+++ b/tests/verify_kernel.c
@@ -9,6 +9,9 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "2sysincludes.h"
+#include "2api.h"
+#include "2misc.h"
 #include "host_common.h"
 #include "util_misc.h"
 #include "vboot_common.h"
@@ -118,8 +121,30 @@ int main(int argc, char *argv[])
 	VbNvSetup(&nvc);
 	params.nv_context = &nvc;
 
+	/*
+	 * Set up vboot context.
+	 *
+	 * TODO: Propagate this up to higher API levels
+	 */
+	struct vb2_context ctx;
+	memset(&ctx, 0, sizeof(ctx));
+	/* No need to initialize ctx->nvdata[]; defaults are fine */
+	/* TODO(chromium:441893): support dev-mode flag and external gpt flag */
+	ctx.workbuf = malloc(VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE);
+	if (!ctx.workbuf) {
+		fprintf(stderr, "Can't allocate workbuf\n");
+		return 1;
+	}
+	ctx.workbuf_size = VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE;
+
+	if (VB2_SUCCESS != vb2_init_context(&ctx)) {
+		free(ctx.workbuf);
+		fprintf(stderr, "Can't init context\n");
+		return 1;
+	}
+
 	/* Try loading kernel */
-	rv = LoadKernel(&params, &cparams);
+	rv = LoadKernel(&ctx, &params, &cparams);
 	if (rv != VBERROR_SUCCESS) {
 		fprintf(stderr, "LoadKernel() failed with code %d\n", rv);
 		return 1;