firmware: Align workbuf used size

Previously, workbuf used was not rounded up to a multiple of
VB2_WORKBUF_ALIGN.  The next allocation would be aligned, but not
until it was made.

Change this to round up used size when more workbuf is used.  This
provides better predictability of where the next allocation will be
placed.

Uncovered this problem when I added a new member to vb2_shared_data
which changed its size so it wasn't a multiple of VB2_WORKBUF_ALIGN,
and the vb20 and vb21 unit tests which tried to simulate not enough
buffer broke in strange ways.

BUG=chromium:611535
BRANCH=none
TEST=make -j runtests; build bob firmware and boot it

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

7 files changed, 111 insertions, 83 deletions

diff --git a/tests/vb20_api_kernel_tests.c b/tests/vb20_api_kernel_tests.c
index 212ad931..3a2de215 100644
--- a/tests/vb20_api_kernel_tests.c
+++ b/tests/vb20_api_kernel_tests.c
@@ -20,7 +20,6 @@
 /* Common context for tests */
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__ ((aligned (VB2_WORKBUF_ALIGN)));
-static struct vb2_workbuf wb;
 static struct vb2_context cc;
 static struct vb2_shared_data *sd;
 static struct vb2_fw_preamble *fwpre;
@@ -59,7 +58,6 @@ static void reset_common_data(enum reset_type t)
 	memset(&cc, 0, sizeof(cc));
 	cc.workbuf = workbuf;
 	cc.workbuf_size = sizeof(workbuf);
-	vb2_workbuf_from_ctx(&cc, &wb);
 
 	vb2_init_context(&cc);
 	sd = vb2_get_sd(&cc);
@@ -89,6 +87,7 @@ static void reset_common_data(enum reset_type t)
 		mock_gbb.recovery_key.key_offset +
 		mock_gbb.recovery_key.key_size;
 
+
 	if (t == FOR_PHASE1) {
 		uint8_t *kdata;
 
@@ -103,7 +102,8 @@ static void reset_common_data(enum reset_type t)
 		k->key_offset = vb2_offset_of(k, kdata);
 		k->key_size = sizeof(fw_kernel_key_data);
 		sd->workbuf_preamble_size = sizeof(*fwpre) + k->key_size;
-		cc.workbuf_used += sd->workbuf_preamble_size;
+		vb2_set_workbuf_used(&cc, sd->workbuf_preamble_offset +
+				     sd->workbuf_preamble_size);
 
 	} else if (t == FOR_PHASE2) {
 		struct vb2_signature *sig;
@@ -116,7 +116,8 @@ static void reset_common_data(enum reset_type t)
 			(cc.workbuf + sd->workbuf_data_key_offset);
 		kdkey->algorithm = VB2_ALG_RSA2048_SHA256;
 		sd->workbuf_data_key_size = sizeof(*kdkey);
-		cc.workbuf_used += sd->workbuf_data_key_size;
+		vb2_set_workbuf_used(&cc, sd->workbuf_data_key_offset +
+				     sd->workbuf_data_key_size);
 
 		/* Create mock kernel preamble in the context */
 		sd->workbuf_preamble_offset = cc.workbuf_used;
@@ -137,7 +138,8 @@ static void reset_common_data(enum reset_type t)
 		sd->workbuf_preamble_size = sizeof(*kpre) + sig->sig_size;
 		sd->vblock_preamble_offset =
 			0x10000 - sd->workbuf_preamble_size;
-		cc.workbuf_used += sd->workbuf_preamble_size;
+		vb2_set_workbuf_used(&cc, sd->workbuf_preamble_offset +
+				     sd->workbuf_preamble_size);
 
 	} else {
 		/* Set flags and versions for roll-forward */
@@ -235,8 +237,10 @@ static void phase1_tests(void)
 		(cc.workbuf + sd->workbuf_kernel_key_offset);
 	TEST_EQ(sd->workbuf_kernel_key_size, k->key_offset + k->key_size,
 		"  workbuf key size");
-	TEST_EQ(cc.workbuf_used, sd->workbuf_kernel_key_offset +
-		sd->workbuf_kernel_key_size, "  workbuf used");
+	TEST_EQ(cc.workbuf_used,
+		vb2_wb_round_up(sd->workbuf_kernel_key_offset +
+				sd->workbuf_kernel_key_size),
+		"  workbuf used");
 	TEST_EQ(k->algorithm, 7, "  key algorithm");
 	TEST_EQ(k->key_size, sizeof(fw_kernel_key_data), "  key_size");
 	TEST_EQ(memcmp((uint8_t *)k + k->key_offset, fw_kernel_key_data,
@@ -258,8 +262,10 @@ static void phase1_tests(void)
 		(cc.workbuf + sd->workbuf_kernel_key_offset);
 	TEST_EQ(sd->workbuf_kernel_key_size, k->key_offset + k->key_size,
 		"  workbuf key size");
-	TEST_EQ(cc.workbuf_used, sd->workbuf_kernel_key_offset +
-		sd->workbuf_kernel_key_size, "  workbuf used");
+	TEST_EQ(cc.workbuf_used,
+		vb2_wb_round_up(sd->workbuf_kernel_key_offset +
+				sd->workbuf_kernel_key_size),
+		"  workbuf used");
 	TEST_EQ(k->algorithm, 11, "  key algorithm");
 	TEST_EQ(k->key_size, sizeof(mock_gbb.recovery_key_data), "  key_size");
 	TEST_EQ(memcmp((uint8_t *)k + k->key_offset,
@@ -282,7 +288,8 @@ static void phase1_tests(void)
 	/* Failures while reading recovery key */
 	reset_common_data(FOR_PHASE1);
 	cc.flags |= VB2_CONTEXT_RECOVERY_MODE;
-	cc.workbuf_used = cc.workbuf_size - sizeof(struct vb2_gbb_header) + 1;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb2_gbb_header));
 	TEST_EQ(vb2api_kernel_phase1(&cc), VB2_ERROR_GBB_WORKBUF,
 		"phase1 rec workbuf gbb header");
 
@@ -295,7 +302,8 @@ static void phase1_tests(void)
 	reset_common_data(FOR_PHASE1);
 	cc.flags |= VB2_CONTEXT_RECOVERY_MODE;
 	mock_gbb.h.recovery_key_size = cc.workbuf_size - 1;
-	TEST_EQ(vb2api_kernel_phase1(&cc), VB2_ERROR_API_KPHASE1_WORKBUF_REC_KEY,
+	TEST_EQ(vb2api_kernel_phase1(&cc),
+		VB2_ERROR_API_KPHASE1_WORKBUF_REC_KEY,
 		"phase1 rec workbuf key");
 
 	reset_common_data(FOR_PHASE1);
@@ -351,46 +359,57 @@ static void get_kernel_size_tests(void)
 static void verify_kernel_data_tests(void)
 {
 	reset_common_data(FOR_PHASE2);
-	TEST_SUCC(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	TEST_SUCC(vb2api_verify_kernel_data(&cc, kernel_data,
+					    sizeof(kernel_data)),
 		  "verify data good");
 
 	reset_common_data(FOR_PHASE2);
 	sd->workbuf_preamble_size = 0;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_API_VERIFY_KDATA_PREAMBLE, "verify no preamble");
 
 	reset_common_data(FOR_PHASE2);
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data) + 1),
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data) + 1),
 		VB2_ERROR_API_VERIFY_KDATA_SIZE, "verify size");
 
 	reset_common_data(FOR_PHASE2);
-	cc.workbuf_used = cc.workbuf_size - sizeof(struct vb2_digest_context) + 1;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb2_digest_context));
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_API_VERIFY_KDATA_WORKBUF, "verify workbuf");
 
 	reset_common_data(FOR_PHASE2);
 	sd->workbuf_data_key_size = 0;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_API_VERIFY_KDATA_KEY, "verify no key");
 
 	reset_common_data(FOR_PHASE2);
 	mock_unpack_key_retval = VB2_ERROR_MOCK;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_MOCK, "verify unpack key");
 
 	reset_common_data(FOR_PHASE2);
 	kdkey->algorithm = VB2_ALG_COUNT;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_SHA_INIT_ALGORITHM, "verify hash init");
 
 	reset_common_data(FOR_PHASE2);
-	cc.workbuf_used = cc.workbuf_size - sizeof(struct vb2_digest_context) - 4;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	cc.workbuf_used = cc.workbuf_size -
+			vb2_wb_round_up(sizeof(struct vb2_digest_context));
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_API_CHECK_HASH_WORKBUF_DIGEST, "verify hash workbuf");
 
 	reset_common_data(FOR_PHASE2);
 	kernel_data[3] ^= 0xd0;
-	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data, sizeof(kernel_data)),
+	TEST_EQ(vb2api_verify_kernel_data(&cc, kernel_data,
+					  sizeof(kernel_data)),
 		VB2_ERROR_VDATA_VERIFY_DIGEST, "verify hash digest");
 	kernel_data[3] ^= 0xd0;
 }
diff --git a/tests/vb20_api_tests.c b/tests/vb20_api_tests.c
index a7efca9b..e9dff9fe 100644
--- a/tests/vb20_api_tests.c
+++ b/tests/vb20_api_tests.c
@@ -76,8 +76,8 @@ static void reset_common_data(enum reset_type t)
 
 	sd->workbuf_preamble_offset = cc.workbuf_used;
 	sd->workbuf_preamble_size = sizeof(*pre);
-	cc.workbuf_used = sd->workbuf_preamble_offset
-		+ sd->workbuf_preamble_size;
+	vb2_set_workbuf_used(&cc, sd->workbuf_preamble_offset
+			     + sd->workbuf_preamble_size);
 	pre = (struct vb2_fw_preamble *)
 		(cc.workbuf + sd->workbuf_preamble_offset);
 	pre->body_signature.data_size = mock_body_size;
@@ -89,8 +89,8 @@ static void reset_common_data(enum reset_type t)
 
 	sd->workbuf_data_key_offset = cc.workbuf_used;
 	sd->workbuf_data_key_size = sizeof(*k) + 8;
-	cc.workbuf_used = sd->workbuf_data_key_offset +
-		sd->workbuf_data_key_size;
+	vb2_set_workbuf_used(&cc, sd->workbuf_data_key_offset +
+			     sd->workbuf_data_key_size);
 	k = (struct vb2_packed_key *)
 		(cc.workbuf + sd->workbuf_data_key_offset);
 	k->algorithm = mock_algorithm;
@@ -259,14 +259,13 @@ static void init_hash_tests(void)
 	wb_used_before = cc.workbuf_used;
 	TEST_SUCC(vb2api_init_hash(&cc, VB2_HASH_TAG_FW_BODY, &size),
 		  "init hash good");
-	TEST_EQ(sd->workbuf_hash_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_hash_offset, wb_used_before,
 		"hash context offset");
 	TEST_EQ(sd->workbuf_hash_size, sizeof(struct vb2_digest_context),
 		"hash context size");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_hash_offset + sd->workbuf_hash_size,
+		vb2_wb_round_up(sd->workbuf_hash_offset +
+				sd->workbuf_hash_size),
 		"hash uses workbuf");
 	TEST_EQ(sd->hash_tag, VB2_HASH_TAG_FW_BODY, "hash tag");
 	TEST_EQ(sd->hash_remaining_size, mock_body_size, "hash remaining");
@@ -290,8 +289,8 @@ static void init_hash_tests(void)
 		VB2_ERROR_API_INIT_HASH_TAG, "init hash unknown tag");
 
 	reset_common_data(FOR_MISC);
-	cc.workbuf_used =
-		cc.workbuf_size - sizeof(struct vb2_digest_context) + 8;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb2_digest_context));
 	TEST_EQ(vb2api_init_hash(&cc, VB2_HASH_TAG_FW_BODY, &size),
 		VB2_ERROR_API_INIT_HASH_WORKBUF, "init hash workbuf");
 
diff --git a/tests/vb20_kernel_tests.c b/tests/vb20_kernel_tests.c
index 2823050e..213b9e68 100644
--- a/tests/vb20_kernel_tests.c
+++ b/tests/vb20_kernel_tests.c
@@ -252,12 +252,11 @@ static void load_kernel_keyblock_tests(void)
 	TEST_EQ(sd->kernel_version, 0x20000, "keyblock version");
 	TEST_EQ(sd->vblock_preamble_offset, sizeof(mock_vblock.k),
 		"preamble offset");
-	TEST_EQ(sd->workbuf_data_key_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_data_key_offset, wb_used_before,
 		"keyblock data key offset");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_data_key_offset + sd->workbuf_data_key_size,
+		vb2_wb_round_up(sd->workbuf_data_key_offset +
+				sd->workbuf_data_key_size),
 		"workbuf used");
 
 	/* Make sure data key was properly saved */
@@ -271,7 +270,8 @@ static void load_kernel_keyblock_tests(void)
 		       sizeof(mock_vblock.k.data_key_data)),
 		0, "data key data");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_data_key_offset + sd->workbuf_data_key_size,
+		vb2_wb_round_up(sd->workbuf_data_key_offset +
+				sd->workbuf_data_key_size),
 		"workbuf used after");
 
 	/* Test failures */
@@ -281,7 +281,8 @@ static void load_kernel_keyblock_tests(void)
 		VB2_ERROR_MOCK, "Kernel keyblock unpack key");
 
 	reset_common_data(FOR_KEYBLOCK);
-	cc.workbuf_used = cc.workbuf_size - (sizeof(*kb) - 1);
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(*kb));
 	TEST_EQ(vb2_load_kernel_keyblock(&cc),
 		VB2_ERROR_KERNEL_KEYBLOCK_WORKBUF_HEADER,
 		"Kernel keyblock workbuf header");
@@ -292,7 +293,8 @@ static void load_kernel_keyblock_tests(void)
 		VB2_ERROR_MOCK, "Kernel keyblock read header");
 
 	reset_common_data(FOR_KEYBLOCK);
-	cc.workbuf_used = cc.workbuf_size - (kb->keyblock_size - 1);
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(kb->keyblock_size);
 	TEST_EQ(vb2_load_kernel_keyblock(&cc),
 		VB2_ERROR_KERNEL_KEYBLOCK_WORKBUF,
 		"Kernel keyblock workbuf");
@@ -405,13 +407,12 @@ static void load_kernel_preamble_tests(void)
 	wb_used_before = cc.workbuf_used;
 	TEST_SUCC(vb2_load_kernel_preamble(&cc), "preamble good");
 	TEST_EQ(sd->kernel_version, 0x20002, "combined version");
-	TEST_EQ(sd->workbuf_preamble_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_preamble_offset, wb_used_before,
 		"preamble offset");
 	TEST_EQ(sd->workbuf_preamble_size, pre->preamble_size, "preamble size");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_preamble_offset + sd->workbuf_preamble_size,
+		vb2_wb_round_up(sd->workbuf_preamble_offset +
+				sd->workbuf_preamble_size),
 		"workbuf used");
 
 	/* Expected failures */
@@ -428,8 +429,8 @@ static void load_kernel_preamble_tests(void)
 		"preamble unpack data key");
 
 	reset_common_data(FOR_PREAMBLE);
-	cc.workbuf_used = cc.workbuf_size -
-		sizeof(struct vb2_kernel_preamble) + 8;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb2_kernel_preamble));
 	TEST_EQ(vb2_load_kernel_preamble(&cc),
 		VB2_ERROR_KERNEL_PREAMBLE2_WORKBUF_HEADER,
 		"preamble not enough workbuf for header");
@@ -441,7 +442,8 @@ static void load_kernel_preamble_tests(void)
 		"preamble read header");
 
 	reset_common_data(FOR_PREAMBLE);
-	cc.workbuf_used = cc.workbuf_size - sizeof(mock_vblock.p) + 8;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(mock_vblock.p));
 	TEST_EQ(vb2_load_kernel_preamble(&cc),
 		VB2_ERROR_KERNEL_PREAMBLE2_WORKBUF,
 		"preamble not enough workbuf");
diff --git a/tests/vb20_misc_tests.c b/tests/vb20_misc_tests.c
index 2156e919..9442aa0a 100644
--- a/tests/vb20_misc_tests.c
+++ b/tests/vb20_misc_tests.c
@@ -185,12 +185,11 @@ static void verify_keyblock_tests(void)
 	TEST_EQ(sd->fw_version, 0x20000, "keyblock version");
 	TEST_EQ(sd->vblock_preamble_offset, sizeof(mock_vblock.k),
 		"preamble offset");
-	TEST_EQ(sd->workbuf_data_key_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_data_key_offset, wb_used_before,
 		"keyblock data key offset");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_data_key_offset + sd->workbuf_data_key_size,
+		vb2_wb_round_up(sd->workbuf_data_key_offset +
+				sd->workbuf_data_key_size),
 		"workbuf used");
 
 	/* Make sure data key was properly saved */
@@ -204,12 +203,14 @@ static void verify_keyblock_tests(void)
 		       sizeof(mock_vblock.k.data_key_data)),
 		0, "data key data");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_data_key_offset + sd->workbuf_data_key_size,
+		vb2_wb_round_up(sd->workbuf_data_key_offset +
+				sd->workbuf_data_key_size),
 		"workbuf used after");
 
 	/* Test failures */
 	reset_common_data(FOR_KEYBLOCK);
-	cc.workbuf_used = cc.workbuf_size - sd->gbb_rootkey_size + 8;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sd->gbb_rootkey_size);
 	TEST_EQ(vb2_load_fw_keyblock(&cc),
 		VB2_ERROR_FW_KEYBLOCK_WORKBUF_ROOT_KEY,
 		"keyblock not enough workbuf for root key");
@@ -227,7 +228,8 @@ static void verify_keyblock_tests(void)
 		"keyblock unpack root key");
 
 	reset_common_data(FOR_KEYBLOCK);
-	cc.workbuf_used = cc.workbuf_size - sd->gbb_rootkey_size - 8;
+	cc.workbuf_used = cc.workbuf_size -
+			vb2_wb_round_up(sd->gbb_rootkey_size);
 	TEST_EQ(vb2_load_fw_keyblock(&cc),
 		VB2_ERROR_FW_KEYBLOCK_WORKBUF_HEADER,
 		"keyblock not enough workbuf for header");
@@ -239,8 +241,9 @@ static void verify_keyblock_tests(void)
 		"keyblock read keyblock header");
 
 	reset_common_data(FOR_KEYBLOCK);
-	cc.workbuf_used = cc.workbuf_size - sd->gbb_rootkey_size
-		- sizeof(struct vb2_keyblock);
+	cc.workbuf_used = cc.workbuf_size -
+			vb2_wb_round_up(sd->gbb_rootkey_size) -
+			vb2_wb_round_up(sizeof(struct vb2_keyblock));
 	TEST_EQ(vb2_load_fw_keyblock(&cc),
 		VB2_ERROR_FW_KEYBLOCK_WORKBUF,
 		"keyblock not enough workbuf for entire keyblock");
@@ -286,13 +289,12 @@ static void verify_preamble_tests(void)
 	wb_used_before = cc.workbuf_used;
 	TEST_SUCC(vb2_load_fw_preamble(&cc), "preamble good");
 	TEST_EQ(sd->fw_version, 0x20002, "combined version");
-	TEST_EQ(sd->workbuf_preamble_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_preamble_offset, wb_used_before,
 		"preamble offset");
 	TEST_EQ(sd->workbuf_preamble_size, pre->preamble_size, "preamble size");
 	TEST_EQ(cc.workbuf_used,
-		sd->workbuf_preamble_offset + sd->workbuf_preamble_size,
+		vb2_wb_round_up(sd->workbuf_preamble_offset +
+				sd->workbuf_preamble_size),
 		"workbuf used");
 
 	/* Expected failures */
@@ -309,7 +311,8 @@ static void verify_preamble_tests(void)
 		"preamble unpack data key");
 
 	reset_common_data(FOR_PREAMBLE);
-	cc.workbuf_used = cc.workbuf_size - sizeof(struct vb2_fw_preamble) + 8;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb2_fw_preamble));
 	TEST_EQ(vb2_load_fw_preamble(&cc),
 		VB2_ERROR_FW_PREAMBLE2_WORKBUF_HEADER,
 		"preamble not enough workbuf for header");
@@ -321,7 +324,8 @@ static void verify_preamble_tests(void)
 		"preamble read header");
 
 	reset_common_data(FOR_PREAMBLE);
-	cc.workbuf_used = cc.workbuf_size - sizeof(mock_vblock.p) + 8;
+	cc.workbuf_used = cc.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(mock_vblock.p));
 	TEST_EQ(vb2_load_fw_preamble(&cc),
 		VB2_ERROR_FW_PREAMBLE2_WORKBUF,
 		"preamble not enough workbuf");
diff --git a/tests/vb21_api_tests.c b/tests/vb21_api_tests.c
index cab2d140..e9eab10d 100644
--- a/tests/vb21_api_tests.c
+++ b/tests/vb21_api_tests.c
@@ -112,8 +112,8 @@ static void reset_common_data(enum reset_type t)
 	}
 
 	sd->workbuf_preamble_size = sig_offset;
-	ctx.workbuf_used = sd->workbuf_preamble_offset
-		+ sd->workbuf_preamble_size;
+	ctx.workbuf_used = vb2_wb_round_up(sd->workbuf_preamble_offset +
+					   sd->workbuf_preamble_size);
 
 	if (t == FOR_EXTEND_HASH || t == FOR_CHECK_HASH)
 		vb21api_init_hash(&ctx, test_id, NULL);
@@ -210,14 +210,13 @@ static void init_hash_tests(void)
 	wb_used_before = ctx.workbuf_used;
 	TEST_SUCC(vb21api_init_hash(&ctx, test_id, &size),
 		  "init hash good");
-	TEST_EQ(sd->workbuf_hash_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_hash_offset, wb_used_before,
 		"hash context offset");
 	TEST_EQ(sd->workbuf_hash_size, sizeof(struct vb2_digest_context),
 		"hash context size");
 	TEST_EQ(ctx.workbuf_used,
-		sd->workbuf_hash_offset + sd->workbuf_hash_size,
+		vb2_wb_round_up(sd->workbuf_hash_offset +
+				sd->workbuf_hash_size),
 		"hash uses workbuf");
 	TEST_EQ(sd->hash_tag,
 		sd->workbuf_preamble_offset + pre->hash_offset,
@@ -243,8 +242,8 @@ static void init_hash_tests(void)
 		VB2_ERROR_API_INIT_HASH_PREAMBLE, "init hash preamble");
 
 	reset_common_data(FOR_MISC);
-	ctx.workbuf_used =
-		ctx.workbuf_size - sizeof(struct vb2_digest_context) + 8;
+	ctx.workbuf_used = ctx.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb2_digest_context));
 	TEST_EQ(vb21api_init_hash(&ctx, test_id, &size),
 		VB2_ERROR_API_INIT_HASH_WORKBUF, "init hash workbuf");
 
diff --git a/tests/vb21_misc_tests.c b/tests/vb21_misc_tests.c
index d7cfbd26..8bf2a6e8 100644
--- a/tests/vb21_misc_tests.c
+++ b/tests/vb21_misc_tests.c
@@ -190,12 +190,11 @@ static void load_keyblock_tests(void)
 	TEST_EQ(sd->fw_version, 0x20000, "keyblock version");
 	TEST_EQ(sd->vblock_preamble_offset, sizeof(mock_vblock.k),
 		"preamble offset");
-	TEST_EQ(sd->workbuf_data_key_offset,
-		(wb_used_before + (VB2_WORKBUF_ALIGN - 1)) &
-		~(VB2_WORKBUF_ALIGN - 1),
+	TEST_EQ(sd->workbuf_data_key_offset, wb_used_before,
 		"keyblock data key offset");
 	TEST_EQ(ctx.workbuf_used,
-		sd->workbuf_data_key_offset + sd->workbuf_data_key_size,
+		vb2_wb_round_up(sd->workbuf_data_key_offset +
+				sd->workbuf_data_key_size),
 		"workbuf used");
 
 	/* Make sure data key was properly saved */
@@ -210,12 +209,14 @@ static void load_keyblock_tests(void)
 		       sizeof(mock_vblock.k.data_key_data)),
 		0, "data key data");
 	TEST_EQ(ctx.workbuf_used,
-		sd->workbuf_data_key_offset + sd->workbuf_data_key_size,
+		vb2_wb_round_up(sd->workbuf_data_key_offset +
+				sd->workbuf_data_key_size),
 		"workbuf used after");
 
 	/* Test failures */
 	reset_common_data(FOR_KEYBLOCK);
-	ctx.workbuf_used = ctx.workbuf_size - sd->gbb_rootkey_size + 8;
+	ctx.workbuf_used = ctx.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sd->gbb_rootkey_size);
 	TEST_EQ(vb21_load_fw_keyblock(&ctx),
 		VB2_ERROR_FW_KEYBLOCK_WORKBUF_ROOT_KEY,
 		"keyblock not enough workbuf for root key");
@@ -233,7 +234,8 @@ static void load_keyblock_tests(void)
 		"keyblock unpack root key");
 
 	reset_common_data(FOR_KEYBLOCK);
-	ctx.workbuf_used = ctx.workbuf_size - sd->gbb_rootkey_size - 8;
+	ctx.workbuf_used = ctx.workbuf_size -
+			vb2_wb_round_up(sd->gbb_rootkey_size);
 	TEST_EQ(vb21_load_fw_keyblock(&ctx),
 		VB2_ERROR_READ_RESOURCE_OBJECT_BUF,
 		"keyblock not enough workbuf for header");
@@ -245,8 +247,9 @@ static void load_keyblock_tests(void)
 		"keyblock read keyblock header");
 
 	reset_common_data(FOR_KEYBLOCK);
-	ctx.workbuf_used = ctx.workbuf_size - sd->gbb_rootkey_size
-		- sizeof(struct vb21_keyblock);
+	ctx.workbuf_used = ctx.workbuf_size -
+			vb2_wb_round_up(sd->gbb_rootkey_size) -
+			vb2_wb_round_up(sizeof(struct vb21_keyblock));
 	TEST_EQ(vb21_load_fw_keyblock(&ctx),
 		VB2_ERROR_READ_RESOURCE_OBJECT_BUF,
 		"keyblock not enough workbuf for entire keyblock");
@@ -296,7 +299,8 @@ static void load_preamble_tests(void)
 		"preamble offset");
 	TEST_EQ(sd->workbuf_preamble_size, pre->c.total_size, "preamble size");
 	TEST_EQ(ctx.workbuf_used,
-		sd->workbuf_preamble_offset + sd->workbuf_preamble_size,
+		vb2_wb_round_up(sd->workbuf_preamble_offset +
+				sd->workbuf_preamble_size),
 		"workbuf used");
 	TEST_EQ(sd->workbuf_data_key_offset, 0, "data key offset gone");
 	TEST_EQ(sd->workbuf_data_key_size, 0, "data key size gone");
@@ -315,8 +319,8 @@ static void load_preamble_tests(void)
 		"preamble unpack data key");
 
 	reset_common_data(FOR_PREAMBLE);
-	ctx.workbuf_used = ctx.workbuf_size
-		- sizeof(struct vb21_fw_preamble) + 8;
+	ctx.workbuf_used = ctx.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(struct vb21_fw_preamble));
 	TEST_EQ(vb21_load_fw_preamble(&ctx),
 		VB2_ERROR_READ_RESOURCE_OBJECT_BUF,
 		"preamble not enough workbuf for header");
@@ -328,7 +332,8 @@ static void load_preamble_tests(void)
 		"preamble read header");
 
 	reset_common_data(FOR_PREAMBLE);
-	ctx.workbuf_used = ctx.workbuf_size - sizeof(mock_vblock.p) + 8;
+	ctx.workbuf_used = ctx.workbuf_size + VB2_WORKBUF_ALIGN -
+			vb2_wb_round_up(sizeof(mock_vblock.p));
 	TEST_EQ(vb21_load_fw_preamble(&ctx),
 		VB2_ERROR_READ_RESOURCE_OBJECT_BUF,
 		"preamble not enough workbuf");
diff --git a/tests/vb2_misc_tests.c b/tests/vb2_misc_tests.c
index aa3b061f..01131b49 100644
--- a/tests/vb2_misc_tests.c
+++ b/tests/vb2_misc_tests.c
@@ -86,7 +86,7 @@ static void init_context_tests(void)
 	reset_common_data();
 
 	TEST_SUCC(vb2_init_context(&c), "Init context good");
-	TEST_EQ(c.workbuf_used, sizeof(struct vb2_shared_data),
+	TEST_EQ(c.workbuf_used, vb2_wb_round_up(sizeof(struct vb2_shared_data)),
 		"Init vbsd");
 
 	/* Don't re-init if used is non-zero */