vboot: use vb2 verification functions for kernel verification

This removes old vboot1 functions in favor of the new vboot2 functions.

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

Change-Id: Idc64f7714bbd9d4fa82d14b6b5d73d71c61de854
Signed-off-by: Randall Spangler <rspangler@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/400900
Reviewed-by: Daisuke Nojiri <dnojiri@chromium.org>

4 files changed, 87 insertions, 637 deletions

diff --git a/tests/vboot_api_kernel5_tests.c b/tests/vboot_api_kernel5_tests.c
index 04ee7669..d752f6b7 100644
--- a/tests/vboot_api_kernel5_tests.c
+++ b/tests/vboot_api_kernel5_tests.c
@@ -10,10 +10,14 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "2sysincludes.h"
+#include "2common.h"
+#include "2rsa.h"
 #include "gbb_header.h"
 #include "host_common.h"
 #include "load_kernel_fw.h"
 #include "test_common.h"
+#include "vb2_common.h"
 #include "vboot_api.h"
 #include "vboot_common.h"
 #include "vboot_kernel.h"
@@ -31,8 +35,7 @@ static uint8_t kernel_buffer[80000];
 static int key_block_verify_fail;  /* 0=ok, 1=sig, 2=hash */
 static int preamble_verify_fail;
 static int verify_data_fail;
-static RSAPublicKey *mock_data_key;
-static int mock_data_key_allocated;
+static int unpack_key_fail;
 
 static VbNvContext vnc;
 static VbKeyBlockHeader kbh;
@@ -69,9 +72,6 @@ static void ResetMocks(void)
 	preamble_verify_fail = 0;
 	verify_data_fail = 0;
 
-	mock_data_key = (RSAPublicKey *)"TestDataKey";
-	mock_data_key_allocated = 0;
-
 	memset(&kbh, 0, sizeof(kbh));
 	kbh.data_key.key_version = 2;
 	kbh.key_block_flags = -1;
@@ -96,53 +96,68 @@ static void copy_kbh(void)
 }
 
 /* Mocks */
-int KeyBlockVerify(const VbKeyBlockHeader *block, uint64_t size,
-		   const VbPublicKey *key, int hash_only) {
-	hash_only_check = hash_only;
+int vb2_unpack_key(struct vb2_public_key *key,
+		   const uint8_t *buf,
+		   uint32_t size)
+{
+	if (--unpack_key_fail == 0)
+		return VB2_ERROR_MOCK;
+
+	return VB2_SUCCESS;
+}
+
+int vb2_verify_keyblock(struct vb2_keyblock *block,
+			uint32_t size,
+			const struct vb2_public_key *key,
+			const struct vb2_workbuf *wb)
+{
+	hash_only_check = 0;
 
 	if (key_block_verify_fail)
-		return VBERROR_SIMULATED;
+		return VB2_ERROR_MOCK;
 
 	/* Use this as an opportunity to override the key block */
 	memcpy((void *)block, &kbh, sizeof(kbh));
-	return VBERROR_SUCCESS;
+	return VB2_SUCCESS;
 }
 
-RSAPublicKey *PublicKeyToRSA(const VbPublicKey *key)
+int vb2_verify_keyblock_hash(const struct vb2_keyblock *block,
+			     uint32_t size,
+			     const struct vb2_workbuf *wb)
 {
-	TEST_EQ(mock_data_key_allocated, 0, "  mock data key not allocated");
-
-	if (mock_data_key)
-		mock_data_key_allocated++;
+	hash_only_check = 1;
 
-	return mock_data_key;
-}
+	if (key_block_verify_fail)
+		return VB2_ERROR_MOCK;
 
-void RSAPublicKeyFree(RSAPublicKey* key)
-{
-	TEST_EQ(mock_data_key_allocated, 1, "  mock data key allocated");
-	TEST_PTR_EQ(key, mock_data_key, "  data key ptr");
-	mock_data_key_allocated--;
+	/* Use this as an opportunity to override the key block */
+	memcpy((void *)block, &kbh, sizeof(kbh));
+	return VB2_SUCCESS;
 }
 
-int VerifyKernelPreamble(const VbKernelPreambleHeader *preamble,
-			 uint64_t size, const RSAPublicKey *key)
+int vb2_verify_kernel_preamble(struct vb2_kernel_preamble *preamble,
+			       uint32_t size,
+			       const struct vb2_public_key *key,
+			       const struct vb2_workbuf *wb)
 {
 	if (preamble_verify_fail)
-		return VBERROR_SIMULATED;
+		return VB2_ERROR_MOCK;
 
 	/* Use this as an opportunity to override the preamble */
 	memcpy((void *)preamble, &kph, sizeof(kph));
-	return VBERROR_SUCCESS;
+	return VB2_SUCCESS;
 }
 
-int VerifyData(const uint8_t *data, uint64_t size, const VbSignature *sig,
-	       const RSAPublicKey *key)
+int vb2_verify_data(const uint8_t *data,
+		    uint32_t size,
+		    struct vb2_signature *sig,
+		    const struct vb2_public_key *key,
+		    const struct vb2_workbuf *wb)
 {
 	if (verify_data_fail)
-		return VBERROR_SIMULATED;
+		return VB2_ERROR_MOCK;
 
-	return VBERROR_SUCCESS;
+	return VB2_SUCCESS;
 }
 
 VbError_t VbExNvStorageRead(uint8_t *buf)
diff --git a/tests/vboot_common2_tests.c b/tests/vboot_common2_tests.c
deleted file mode 100644
index 386f487f..00000000
--- a/tests/vboot_common2_tests.c
+++ /dev/null
@@ -1,303 +0,0 @@
-/* Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- *
- * Tests for firmware image library.
- */
-
-#include <stdint.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include "2sysincludes.h"
-#include "2common.h"
-#include "2sha.h"
-#include "cryptolib.h"
-#include "file_keys.h"
-#include "host_common.h"
-#include "test_common.h"
-#include "vboot_common.h"
-
-static void VerifyPublicKeyToRSA(const VbPublicKey *orig_key)
-{
-	RSAPublicKey *rsa;
-	VbPublicKey *key =
-		(VbPublicKey *)vb2_alloc_packed_key(orig_key->key_size, 0, 0);
-
-	PublicKeyCopy(key, orig_key);
-	key->algorithm = kNumAlgorithms;
-	TEST_EQ((size_t)PublicKeyToRSA(key), 0,
-		"PublicKeyToRSA() invalid algorithm");
-
-	PublicKeyCopy(key, orig_key);
-	key->key_size -= 1;
-	TEST_EQ((size_t)PublicKeyToRSA(key), 0,
-		"PublicKeyToRSA() invalid size");
-
-	rsa = PublicKeyToRSA(orig_key);
-	TEST_NEQ((size_t)rsa, 0, "PublicKeyToRSA() ok");
-	if (rsa) {
-		TEST_EQ((int)rsa->algorithm, (int)key->algorithm,
-			"PublicKeyToRSA() algorithm");
-		RSAPublicKeyFree(rsa);
-	}
-
-	free(key);
-}
-
-static void VerifyDataTest(const VbPublicKey *public_key,
-                           const struct vb2_private_key *private_key)
-{
-	const uint8_t test_data[] = "This is some test data to sign.";
-	const uint64_t test_size = sizeof(test_data);
-	VbSignature *sig;
-	RSAPublicKey *rsa;
-
-	sig = (VbSignature *)vb2_calculate_signature(test_data, test_size,
-						     private_key);
-	TEST_PTR_NEQ(sig, 0, "VerifyData() calculate signature");
-
-	rsa = PublicKeyToRSA(public_key);
-	TEST_PTR_NEQ(rsa, 0, "VerifyData() calculate rsa");
-
-	if (!sig || !rsa)
-		return;
-
-	TEST_EQ(VerifyData(test_data, test_size, sig, rsa), 0,
-		"VerifyData() ok");
-
-	sig->sig_size -= 16;
-	TEST_EQ(VerifyData(test_data, test_size, sig, rsa), 1,
-		"VerifyData() wrong sig size");
-	sig->sig_size += 16;
-
-	TEST_EQ(VerifyData(test_data, test_size - 1, sig, rsa), 1,
-		"VerifyData() input buffer too small");
-
-	GetSignatureData(sig)[0] ^= 0x5A;
-	TEST_EQ(VerifyData(test_data, test_size, sig, rsa), 1,
-		"VerifyData() wrong sig");
-
-	RSAPublicKeyFree(rsa);
-	free(sig);
-}
-
-static void VerifyDigestTest(const VbPublicKey *public_key,
-                             const struct vb2_private_key *private_key)
-{
-	const uint8_t test_data[] = "This is some other test data to sign.";
-	VbSignature *sig;
-	RSAPublicKey *rsa;
-	uint8_t digest[VB2_MAX_DIGEST_SIZE];
-
-	sig = (VbSignature *)vb2_calculate_signature(test_data,
-						     sizeof(test_data),
-						     private_key);
-	rsa = PublicKeyToRSA(public_key);
-	TEST_SUCC(vb2_digest_buffer(test_data, sizeof(test_data),
-				    vb2_crypto_to_hash(public_key->algorithm),
-				    digest, sizeof(digest)),
-		  "VerifyData() digest");
-
-	TEST_NEQ(sig && rsa, 0, "VerifyData() prerequisites");
-	if (!sig || !rsa)
-		return;
-
-	TEST_EQ(VerifyDigest(digest, sig, rsa), 0, "VerifyDigest() ok");
-
-	GetSignatureData(sig)[0] ^= 0x5A;
-	TEST_EQ(VerifyDigest(digest, sig, rsa), 1, "VerifyDigest() wrong sig");
-
-	sig->sig_size = 1;
-	TEST_EQ(VerifyDigest(digest, sig, rsa), 1, "VerifyDigest() sig size");
-
-	RSAPublicKeyFree(rsa);
-	free(sig);
-}
-
-static void ReSignKernelPreamble(VbKernelPreambleHeader *h,
-                                 const struct vb2_private_key *key)
-{
-	struct vb2_signature *sig = vb2_calculate_signature(
-		(const uint8_t *)h, h->preamble_signature.data_size, key);
-
-	vb2_copy_signature((struct vb2_signature *)&h->preamble_signature, sig);
-	free(sig);
-}
-
-static void VerifyKernelPreambleTest(const VbPublicKey *public_key,
-                                     const struct vb2_private_key *private_key)
-{
-	VbKernelPreambleHeader *hdr;
-	VbKernelPreambleHeader *h;
-	RSAPublicKey *rsa;
-	unsigned hsize;
-
-	/* Create a dummy signature */
-	struct vb2_signature *body_sig = vb2_alloc_signature(56, 78);
-
-	rsa = PublicKeyToRSA(public_key);
-	hdr = (VbKernelPreambleHeader *)
-		vb2_create_kernel_preamble(0x1234, 0x100000, 0x300000, 0x4000,
-					   body_sig, 0, 0, 0, 0, private_key);
-	TEST_NEQ(hdr && rsa, 0, "VerifyKernelPreamble() prerequisites");
-	if (!hdr) {
-		free(body_sig);
-		RSAPublicKeyFree(rsa);
-		return;
-	}
-
-	hsize = (unsigned) hdr->preamble_size;
-	h = (VbKernelPreambleHeader *)malloc(hsize + 16384);
-
-	TEST_EQ(VerifyKernelPreamble(hdr, hsize, rsa), 0,
-		"VerifyKernelPreamble() ok using key");
-	TEST_NEQ(VerifyKernelPreamble(hdr, hsize - 1, rsa), 0,
-		 "VerifyKernelPreamble() size--");
-	TEST_NEQ(VerifyKernelPreamble(hdr, 4, rsa), 0,
-		 "VerifyKernelPreamble() size tiny");
-	TEST_EQ(VerifyKernelPreamble(hdr, hsize + 1, rsa), 0,
-		"VerifyKernelPreamble() size++");
-
-	/* Care about major version but not minor */
-	memcpy(h, hdr, hsize);
-	h->header_version_major++;
-	ReSignKernelPreamble(h, private_key);
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() major++");
-
-	memcpy(h, hdr, hsize);
-	h->header_version_major--;
-	ReSignKernelPreamble(h, private_key);
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() major--");
-
-	memcpy(h, hdr, hsize);
-	h->header_version_minor++;
-	ReSignKernelPreamble(h, private_key);
-	TEST_EQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		"VerifyKernelPreamble() minor++");
-
-	memcpy(h, hdr, hsize);
-	h->header_version_minor--;
-	ReSignKernelPreamble(h, private_key);
-	TEST_EQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		"VerifyKernelPreamble() minor--");
-
-	/* Check signature */
-	memcpy(h, hdr, hsize);
-	h->preamble_signature.sig_offset = hsize;
-	ReSignKernelPreamble(h, private_key);
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() sig off end");
-
-	memcpy(h, hdr, hsize);
-	h->preamble_signature.sig_size--;
-	ReSignKernelPreamble(h, private_key);
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() sig too small");
-
-	memcpy(h, hdr, hsize);
-	GetSignatureData(&h->body_signature)[0] ^= 0x34;
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() sig mismatch");
-
-	/* Check that we signed header and body sig */
-	memcpy(h, hdr, hsize);
-	h->preamble_signature.data_size = 4;
-	h->body_signature.sig_offset = 0;
-	h->body_signature.sig_size = 0;
-	ReSignKernelPreamble(h, private_key);
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() didn't sign header");
-
-	memcpy(h, hdr, hsize);
-	h->body_signature.sig_offset = hsize;
-	ReSignKernelPreamble(h, private_key);
-	TEST_NEQ(VerifyKernelPreamble(h, hsize, rsa), 0,
-		 "VerifyKernelPreamble() body sig off end");
-
-	/* TODO: verify parser can support a bigger header. */
-
-	free(h);
-	RSAPublicKeyFree(rsa);
-	free(hdr);
-	free(body_sig);
-}
-
-int test_algorithm(int key_algorithm, const char *keys_dir)
-{
-	char filename[1024];
-	int rsa_len = siglen_map[key_algorithm] * 8;
-
-	VbPublicKey *public_key = NULL;
-
-	printf("***Testing algorithm: %s\n", algo_strings[key_algorithm]);
-
-	snprintf(filename, sizeof(filename),
-		 "%s/key_rsa%d.pem", keys_dir, rsa_len);
-	struct vb2_private_key *private_key =
-		vb2_read_private_key_pem(filename, key_algorithm);
-	if (!private_key) {
-		fprintf(stderr, "Error reading private_key: %s\n", filename);
-		return 1;
-	}
-
-	snprintf(filename, sizeof(filename),
-		 "%s/key_rsa%d.keyb", keys_dir, rsa_len);
-	public_key = (VbPublicKey *)vb2_read_packed_keyb(filename,
-							 key_algorithm, 1);
-	if (!public_key) {
-		fprintf(stderr, "Error reading public_key: %s\n", filename);
-		free(private_key);
-		return 1;
-	}
-
-	VerifyPublicKeyToRSA(public_key);
-	VerifyDataTest(public_key, private_key);
-	VerifyDigestTest(public_key, private_key);
-	VerifyKernelPreambleTest(public_key, private_key);
-
-	free(public_key);
-	free(private_key);
-
-	return 0;
-}
-
-/*
- * Test only the algorithms we use:
- * 4 (rsa2048 sha256)
- * 7 (rsa4096 sha256)
- * 11 (rsa8192 sha512)
- */
-const int key_algs[] = {4, 7, 11};
-
-int main(int argc, char *argv[]) {
-	if (argc == 2) {
-		int i;
-
-		for (i = 0; i < ARRAY_SIZE(key_algs); i++) {
-			if (test_algorithm(key_algs[i], argv[1]))
-				return 1;
-		}
-
-	} else if (argc == 3 && !strcasecmp(argv[2], "--all")) {
-		/* Test all the algorithms */
-		int alg;
-
-		for (alg = 0; alg < kNumAlgorithms; alg++) {
-			if (test_algorithm(alg, argv[1]))
-				return 1;
-		}
-
-	} else {
-		fprintf(stderr, "Usage: %s <keys_dir> [--all]", argv[0]);
-		return -1;
-	}
-
-	if (vboot_api_stub_check_memory())
-		return 255;
-
-	return gTestSuccess ? 0 : 255;
-}
diff --git a/tests/vboot_common3_tests.c b/tests/vboot_common3_tests.c
deleted file mode 100644
index 390b7f30..00000000
--- a/tests/vboot_common3_tests.c
+++ /dev/null
@@ -1,270 +0,0 @@
-/* Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- *
- * Tests for firmware image library.
- */
-
-#include <stdint.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include "2sysincludes.h"
-#include "2common.h"
-#include "2sha.h"
-#include "cryptolib.h"
-#include "file_keys.h"
-#include "host_common.h"
-#include "test_common.h"
-#include "vboot_common.h"
-
-
-static void ReChecksumKeyBlock(VbKeyBlockHeader *h)
-{
-	vb2_digest_buffer((const uint8_t *)h,
-			  h->key_block_checksum.data_size,
-			  VB2_HASH_SHA512,
-			  GetSignatureData(&h->key_block_checksum),
-			  VB2_SHA512_DIGEST_SIZE);
-}
-
-static void KeyBlockVerifyTest(const VbPublicKey *public_key,
-                               const struct vb2_private_key *private_key,
-                               const struct vb2_packed_key *data_key)
-{
-	VbKeyBlockHeader *hdr;
-	VbKeyBlockHeader *h;
-	unsigned hsize;
-
-	hdr = (VbKeyBlockHeader *)
-		vb2_create_keyblock((struct vb2_packed_key *)data_key,
-				    private_key,
-				    0x1234);
-	TEST_NEQ((size_t)hdr, 0, "KeyBlockVerify() prerequisites");
-	if (!hdr)
-		return;
-	hsize = (unsigned) hdr->key_block_size;
-	h = (VbKeyBlockHeader *)malloc(hsize + 1024);
-
-	TEST_EQ(KeyBlockVerify(hdr, hsize, NULL, 1), 0,
-		"KeyBlockVerify() ok using checksum");
-	TEST_EQ(KeyBlockVerify(hdr, hsize, public_key, 0), 0,
-		"KeyBlockVerify() ok using key");
-	TEST_NEQ(KeyBlockVerify(hdr, hsize, NULL, 0), 0,
-		 "KeyBlockVerify() missing key");
-
-	TEST_NEQ(KeyBlockVerify(hdr, hsize - 1, NULL, 1), 0,
-		 "KeyBlockVerify() size--");
-	TEST_EQ(KeyBlockVerify(hdr, hsize + 1, NULL, 1), 0,
-		"KeyBlockVerify() size++");
-
-	memcpy(h, hdr, hsize);
-	h->magic[0] &= 0x12;
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() magic");
-
-	/* Care about major version but not minor */
-	memcpy(h, hdr, hsize);
-	h->header_version_major++;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() major++");
-
-	memcpy(h, hdr, hsize);
-	h->header_version_major--;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() major--");
-
-	memcpy(h, hdr, hsize);
-	h->header_version_minor++;
-	ReChecksumKeyBlock(h);
-	TEST_EQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		"KeyBlockVerify() minor++");
-
-	memcpy(h, hdr, hsize);
-	h->header_version_minor--;
-	ReChecksumKeyBlock(h);
-	TEST_EQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		"KeyBlockVerify() minor--");
-
-	/* Check hash */
-	memcpy(h, hdr, hsize);
-	h->key_block_checksum.sig_offset = hsize;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() checksum off end");
-
-	memcpy(h, hdr, hsize);
-	h->key_block_checksum.sig_size /= 2;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() checksum too small");
-
-	memcpy(h, hdr, hsize);
-	GetPublicKeyData(&h->data_key)[0] ^= 0x34;
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() checksum mismatch");
-
-	/* Check signature */
-	memcpy(h, hdr, hsize);
-	h->key_block_signature.sig_offset = hsize;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, public_key, 0), 0,
-		 "KeyBlockVerify() sig off end");
-
-	memcpy(h, hdr, hsize);
-	h->key_block_signature.sig_size--;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, public_key, 0), 0,
-		 "KeyBlockVerify() sig too small");
-
-	memcpy(h, hdr, hsize);
-	GetPublicKeyData(&h->data_key)[0] ^= 0x34;
-	TEST_NEQ(KeyBlockVerify(h, hsize, public_key, 0), 0,
-		 "KeyBlockVerify() sig mismatch");
-
-	memcpy(h, hdr, hsize);
-	h->key_block_checksum.data_size = h->key_block_size + 1;
-	TEST_NEQ(KeyBlockVerify(h, hsize, public_key, 1), 0,
-		 "KeyBlockVerify() checksum data past end of block");
-
-	/* Check that we signed header and data key */
-	memcpy(h, hdr, hsize);
-	h->key_block_checksum.data_size = 4;
-	h->data_key.key_offset = 0;
-	h->data_key.key_size = 0;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() didn't sign header");
-
-	memcpy(h, hdr, hsize);
-	h->data_key.key_offset = hsize;
-	ReChecksumKeyBlock(h);
-	TEST_NEQ(KeyBlockVerify(h, hsize, NULL, 1), 0,
-		 "KeyBlockVerify() data key off end");
-
-	/* Corner cases for error checking */
-	TEST_NEQ(KeyBlockVerify(NULL, 4, NULL, 1), 0,
-		 "KeyBlockVerify size too small");
-
-	/*
-	 * TODO: verify parser can support a bigger header (i.e., one where
-	 * data_key.key_offset is bigger than expected).
-	 */
-
-	free(h);
-	free(hdr);
-}
-
-int test_permutation(int signing_key_algorithm, int data_key_algorithm,
-		     const char *keys_dir)
-{
-	char filename[1024];
-	int signing_rsa_len = siglen_map[signing_key_algorithm] * 8;
-	int data_rsa_len = siglen_map[data_key_algorithm] * 8;
-
-	struct vb2_private_key *signing_private_key = NULL;
-	struct vb2_packed_key *data_public_key = NULL;
-	VbPublicKey *signing_public_key = NULL;
-	int retval = 1;
-
-	printf("***Testing signing algorithm: %s\n",
-	       algo_strings[signing_key_algorithm]);
-	printf("***With data key algorithm: %s\n",
-	       algo_strings[data_key_algorithm]);
-
-	snprintf(filename, sizeof(filename),
-		 "%s/key_rsa%d.pem", keys_dir, signing_rsa_len);
-	signing_private_key =
-		vb2_read_private_key_pem(filename, signing_key_algorithm);
-	if (!signing_private_key) {
-		fprintf(stderr, "Error reading signing_private_key: %s\n",
-			filename);
-		goto cleanup;
-	}
-
-	snprintf(filename, sizeof(filename),
-		 "%s/key_rsa%d.keyb", keys_dir, signing_rsa_len);
-	signing_public_key = (VbPublicKey *)
-		vb2_read_packed_keyb(filename, signing_key_algorithm, 1);
-	if (!signing_public_key) {
-		fprintf(stderr, "Error reading signing_public_key: %s\n",
-			filename);
-		goto cleanup;
-	}
-
-	snprintf(filename, sizeof(filename),
-		 "%s/key_rsa%d.keyb", keys_dir, data_rsa_len);
-	data_public_key = vb2_read_packed_keyb(filename, data_key_algorithm, 1);
-	if (!data_public_key) {
-		fprintf(stderr, "Error reading data_public_key: %s\n",
-			filename);
-		goto cleanup;
-	}
-
-	KeyBlockVerifyTest(signing_public_key, signing_private_key,
-			   data_public_key);
-	retval = 0;
-
-cleanup:
-	if (signing_public_key)
-		free(signing_public_key);
-	if (signing_private_key)
-		free(signing_private_key);
-	if (data_public_key)
-		free(data_public_key);
-
-	return retval;
-}
-
-struct test_perm
-{
-	int signing_algorithm;
-	int data_key_algorithm;
-};
-
-/*
- * Permutations of signing and data key algorithms in active use:
- * 7 (rsa4096 sha256) - 4 (rsa2048 sha256)
- * 11 (rsa8192 sha512) - 4 (rsa2048 sha256)
- * 11 (rsa8192 sha512) - 7 (rsa4096 sha256)
- */
-const struct test_perm test_perms[] = {{7, 4}, {11, 4}, {11, 7}};
-
-int main(int argc, char *argv[])
-{
-	if (argc == 2) {
-		/* Test only the algorithms we use */
-		int i;
-
-		for (i = 0; i < ARRAY_SIZE(test_perms); i++) {
-			if (test_permutation(test_perms[i].signing_algorithm,
-					     test_perms[i].data_key_algorithm,
-					     argv[1]))
-				return 1;
-		}
-
-	} else if (argc == 3 && !strcasecmp(argv[2], "--all")) {
-		/* Test all the algorithms */
-		int sign_alg, data_alg;
-
-		for (sign_alg = 0; sign_alg < kNumAlgorithms; sign_alg++) {
-			for (data_alg = 0; data_alg < kNumAlgorithms;
-			     data_alg++) {
-				if (test_permutation(sign_alg, data_alg,
-						     argv[1]))
-					return 1;
-			}
-		}
-	} else {
-		fprintf(stderr,	"Usage: %s <keys_dir> [--all]",	argv[0]);
-		return -1;
-	}
-
-	if (vboot_api_stub_check_memory())
-		return 255;
-
-	return gTestSuccess ? 0 : 255;
-}
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index e3695ab1..994b792c 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -23,6 +23,7 @@
 #include "load_kernel_fw.h"
 #include "rollback_index.h"
 #include "test_common.h"
+#include "vb2_struct.h"
 #include "vboot_api.h"
 #include "vboot_common.h"
 #include "vboot_kernel.h"
@@ -54,8 +55,7 @@ static int gpt_init_fail;
 static int key_block_verify_fail;  /* 0=ok, 1=sig, 2=hash */
 static int preamble_verify_fail;
 static int verify_data_fail;
-static RSAPublicKey *mock_data_key;
-static int mock_data_key_allocated;
+static int unpack_key_fail;
 static int gpt_flag_external;
 
 static uint8_t gbb_data[sizeof(GoogleBinaryBlockHeader) + 2048];
@@ -130,9 +130,7 @@ static void ResetMocks(void)
 	key_block_verify_fail = 0;
 	preamble_verify_fail = 0;
 	verify_data_fail = 0;
-
-	mock_data_key = (RSAPublicKey *)"TestDataKey";
-	mock_data_key_allocated = 0;
+	unpack_key_fail = 0;
 
 	gpt_flag_external = 0;
 
@@ -246,54 +244,64 @@ void GetCurrentKernelUniqueGuid(GptData *gpt, void *dest)
 	memcpy(dest, fake_guid, sizeof(fake_guid));
 }
 
-int KeyBlockVerify(const VbKeyBlockHeader *block, uint64_t size,
-		   const VbPublicKey *key, int hash_only) {
-
-	if (hash_only && key_block_verify_fail >= 2)
-		return VBERROR_SIMULATED;
-	else if (!hash_only && key_block_verify_fail >= 1)
-		return VBERROR_SIMULATED;
+int vb2_unpack_key(struct vb2_public_key *key,
+		   const uint8_t *buf,
+		   uint32_t size)
+{
+	if (--unpack_key_fail == 0)
+		return VB2_ERROR_MOCK;
 
-	/* Use this as an opportunity to override the key block */
-	memcpy((void *)block, &kbh, sizeof(kbh));
-	return VBERROR_SUCCESS;
+	return VB2_SUCCESS;
 }
 
-RSAPublicKey *PublicKeyToRSA(const VbPublicKey *key)
+int vb2_verify_keyblock(struct vb2_keyblock *block,
+			uint32_t size,
+			const struct vb2_public_key *key,
+			const struct vb2_workbuf *wb)
 {
-	TEST_EQ(mock_data_key_allocated, 0, "  mock data key not allocated");
+	if (key_block_verify_fail >= 1)
+		return VB2_ERROR_MOCK;
 
-	if (mock_data_key)
-		mock_data_key_allocated++;
-
-	return mock_data_key;
+	/* Use this as an opportunity to override the key block */
+	memcpy((void *)block, &kbh, sizeof(kbh));
+	return VB2_SUCCESS;
 }
 
-void RSAPublicKeyFree(RSAPublicKey* key)
+int vb2_verify_keyblock_hash(const struct vb2_keyblock *block,
+			     uint32_t size,
+			     const struct vb2_workbuf *wb)
 {
-	TEST_EQ(mock_data_key_allocated, 1, "  mock data key allocated");
-	TEST_PTR_EQ(key, mock_data_key, "  data key ptr");
-	mock_data_key_allocated--;
+	if (key_block_verify_fail >= 2)
+		return VB2_ERROR_MOCK;
+
+	/* Use this as an opportunity to override the key block */
+	memcpy((void *)block, &kbh, sizeof(kbh));
+	return VB2_SUCCESS;
 }
 
-int VerifyKernelPreamble(const VbKernelPreambleHeader *preamble,
-			 uint64_t size, const RSAPublicKey *key)
+int vb2_verify_kernel_preamble(struct vb2_kernel_preamble *preamble,
+			       uint32_t size,
+			       const struct vb2_public_key *key,
+			       const struct vb2_workbuf *wb)
 {
 	if (preamble_verify_fail)
-		return VBERROR_SIMULATED;
+		return VB2_ERROR_MOCK;
 
 	/* Use this as an opportunity to override the preamble */
 	memcpy((void *)preamble, &kph, sizeof(kph));
-	return VBERROR_SUCCESS;
+	return VB2_SUCCESS;
 }
 
-int VerifyData(const uint8_t *data, uint64_t size, const VbSignature *sig,
-	       const RSAPublicKey *key)
+int vb2_verify_data(const uint8_t *data,
+		    uint32_t size,
+		    struct vb2_signature *sig,
+		    const struct vb2_public_key *key,
+		    const struct vb2_workbuf *wb)
 {
 	if (verify_data_fail)
-		return VBERROR_SIMULATED;
+		return VB2_ERROR_MOCK;
 
-	return VBERROR_SUCCESS;
+	return VB2_SUCCESS;
 }
 
 int vb2_digest_buffer(const uint8_t *buf,
@@ -727,7 +735,7 @@ static void LoadKernelTest(void)
 	TEST_EQ(LoadKernel(&lkp, &cparams), 0, "Key version ignored in rec mode");
 
 	ResetMocks();
-	mock_data_key = NULL;
+	unpack_key_fail = 2;
 	TEST_EQ(LoadKernel(&lkp, &cparams), VBERROR_INVALID_KERNEL_FOUND,
 		"Bad data key");