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>
(cherry picked from commit 1a5e02c7a9ca66fb20f6752225903aa7a4dd0f70)
Reviewed-on: https://chromium-review.googlesource.com/861607
Commit-Queue: Shelley Chen <shchen@chromium.org>
Tested-by: Shelley Chen <shchen@chromium.org>

17 files changed, 160 insertions, 114 deletions

diff --git a/firmware/2lib/2common.c b/firmware/2lib/2common.c
index 9e75d0c3..4db9e77a 100644
--- a/firmware/2lib/2common.c
+++ b/firmware/2lib/2common.c
@@ -61,23 +61,12 @@ void vb2_workbuf_init(struct vb2_workbuf *wb, uint8_t *buf, uint32_t size)
 		wb->size = 0;
 }
 
-/**
- * Round up a number to a multiple of VB2_WORKBUF_ALIGN
- *
- * @param v		Number to round up
- * @return The number, rounded up.
- */
-static __inline uint32_t wb_round_up(uint32_t v)
-{
-	return (v + VB2_WORKBUF_ALIGN - 1) & ~(VB2_WORKBUF_ALIGN - 1);
-}
-
 void *vb2_workbuf_alloc(struct vb2_workbuf *wb, uint32_t size)
 {
 	uint8_t *ptr = wb->buf;
 
 	/* Round up size to work buffer alignment */
-	size = wb_round_up(size);
+	size = vb2_wb_round_up(size);
 
 	if (size > wb->size)
 		return NULL;
@@ -104,7 +93,7 @@ void *vb2_workbuf_realloc(struct vb2_workbuf *wb,
 void vb2_workbuf_free(struct vb2_workbuf *wb, uint32_t size)
 {
 	/* Round up size to work buffer alignment */
-	size = wb_round_up(size);
+	size = vb2_wb_round_up(size);
 
 	wb->buf -= size;
 	wb->size += size;
diff --git a/firmware/2lib/2misc.c b/firmware/2lib/2misc.c
index a955f536..9e4d3297 100644
--- a/firmware/2lib/2misc.c
+++ b/firmware/2lib/2misc.c
@@ -31,6 +31,11 @@ void vb2_workbuf_from_ctx(struct vb2_context *ctx, struct vb2_workbuf *wb)
 			 ctx->workbuf_size - ctx->workbuf_used);
 }
 
+void vb2_set_workbuf_used(struct vb2_context *ctx, uint32_t used)
+{
+	ctx->workbuf_used = vb2_wb_round_up(used);
+}
+
 int vb2_read_gbb_header(struct vb2_context *ctx, struct vb2_gbb_header *gbb)
 {
 	int rv;
@@ -132,7 +137,7 @@ int vb2_init_context(struct vb2_context *ctx)
 
 	/* Initialize the shared data at the start of the work buffer */
 	memset(sd, 0, sizeof(*sd));
-	ctx->workbuf_used = sizeof(*sd);
+	ctx->workbuf_used = vb2_wb_round_up(sizeof(*sd));
 	return VB2_SUCCESS;
 }
 
diff --git a/firmware/2lib/include/2common.h b/firmware/2lib/include/2common.h
index 67032c9b..019a34ac 100644
--- a/firmware/2lib/include/2common.h
+++ b/firmware/2lib/include/2common.h
@@ -75,6 +75,17 @@ struct vb2_public_key;
  * macro for us we'll be safe and use that. */
 #define VB2_WORKBUF_ALIGN __BIGGEST_ALIGNMENT__
 
+/**
+ * Round up a number to a multiple of VB2_WORKBUF_ALIGN
+ *
+ * @param v		Number to round up
+ * @return The number, rounded up.
+ */
+static __inline uint32_t vb2_wb_round_up(uint32_t v)
+{
+	return (v + VB2_WORKBUF_ALIGN - 1) & ~(VB2_WORKBUF_ALIGN - 1);
+}
+
 /* Work buffer */
 struct vb2_workbuf {
 	uint8_t *buf;
diff --git a/firmware/2lib/include/2misc.h b/firmware/2lib/include/2misc.h
index d937b8df..853a1b40 100644
--- a/firmware/2lib/include/2misc.h
+++ b/firmware/2lib/include/2misc.h
@@ -42,6 +42,17 @@ int vb2_validate_gbb_signature(uint8_t *sig);
 void vb2_workbuf_from_ctx(struct vb2_context *ctx, struct vb2_workbuf *wb);
 
 /**
+ * Set the amount of work buffer used in the vboot context.
+ *
+ * This will round up to VB2_WORKBUF_ALIGN, so that the next allocation will
+ * be aligned as expected.
+ *
+ * @param ctx		Vboot context
+ * @param used		Number of bytes used
+ */
+void vb2_set_workbuf_used(struct vb2_context *ctx, uint32_t used);
+
+/**
  * Read the GBB header.
  *
  * @param ctx		Vboot context
diff --git a/firmware/lib20/api.c b/firmware/lib20/api.c
index 327a25c9..606b5ceb 100644
--- a/firmware/lib20/api.c
+++ b/firmware/lib20/api.c
@@ -73,7 +73,7 @@ int vb2api_init_hash(struct vb2_context *ctx, uint32_t tag, uint32_t *size)
 
 		sd->workbuf_hash_offset = vb2_offset_of(ctx->workbuf, dc);
 		sd->workbuf_hash_size = dig_size;
-		ctx->workbuf_used = sd->workbuf_hash_offset + dig_size;
+		vb2_set_workbuf_used(ctx, sd->workbuf_hash_offset + dig_size);
 	}
 
 	/*
diff --git a/firmware/lib20/api_kernel.c b/firmware/lib20/api_kernel.c
index ee5c41a9..dbe47e61 100644
--- a/firmware/lib20/api_kernel.c
+++ b/firmware/lib20/api_kernel.c
@@ -73,6 +73,9 @@ int vb2api_kernel_phase1(struct vb2_context *ctx)
 					 key_data, key_size);
 		if (rv)
 			return rv;
+
+		sd->workbuf_kernel_key_offset =
+				vb2_offset_of(ctx->workbuf, key_data);
 	} else {
 		/* Kernel subkey from firmware preamble */
 		struct vb2_fw_preamble *pre;
@@ -91,7 +94,8 @@ int vb2api_kernel_phase1(struct vb2_context *ctx)
 		 * data key, firmware preamble, or hash data.  So move the
 		 * kernel key from the preamble down after the shared data.
 		 */
-		key_data = (uint8_t *)(sd + 1);
+		sd->workbuf_kernel_key_offset = vb2_wb_round_up(sizeof(*sd));
+		key_data = ctx->workbuf + sd->workbuf_kernel_key_offset;
 		packed_key = (struct vb2_packed_key *)key_data;
 		memmove(packed_key, pre_key, sizeof(*packed_key));
 		packed_key->key_offset = sizeof(*packed_key);
@@ -114,11 +118,9 @@ int vb2api_kernel_phase1(struct vb2_context *ctx)
 	 *   - vb2_shared_data
 	 *   - kernel key
 	 */
-	sd->workbuf_kernel_key_offset =
-		vb2_offset_of(ctx->workbuf, key_data);
 	sd->workbuf_kernel_key_size = key_size;
-	ctx->workbuf_used = sd->workbuf_kernel_key_offset +
-		sd->workbuf_kernel_key_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_kernel_key_offset +
+			     sd->workbuf_kernel_key_size);
 
 	return VB2_SUCCESS;
 }
diff --git a/firmware/lib20/kernel.c b/firmware/lib20/kernel.c
index f5632e58..e06c20eb 100644
--- a/firmware/lib20/kernel.c
+++ b/firmware/lib20/kernel.c
@@ -237,8 +237,8 @@ int vb2_load_kernel_keyblock(struct vb2_context *ctx)
 	 *   - kernel key
 	 *   - packed kernel data key
 	 */
-	ctx->workbuf_used = sd->workbuf_data_key_offset +
-		sd->workbuf_data_key_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_data_key_offset +
+			     sd->workbuf_data_key_size);
 
 	return VB2_SUCCESS;
 }
@@ -445,7 +445,7 @@ int vb2_load_kernel_preamble(struct vb2_context *ctx)
 	 * TODO: we could move the preamble down over the kernel data key
 	 * since we don't need it anymore.
 	 */
-	ctx->workbuf_used = sd->workbuf_preamble_offset + pre_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_preamble_offset + pre_size);
 
 	return VB2_SUCCESS;
 }
diff --git a/firmware/lib20/misc.c b/firmware/lib20/misc.c
index cbe894c9..d7a905a9 100644
--- a/firmware/lib20/misc.c
+++ b/firmware/lib20/misc.c
@@ -185,8 +185,8 @@ int vb2_load_fw_keyblock(struct vb2_context *ctx)
 	 *   - vb2_shared_data
 	 *   - packed firmware data key
 	 */
-	ctx->workbuf_used = sd->workbuf_data_key_offset +
-		sd->workbuf_data_key_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_data_key_offset +
+			     sd->workbuf_data_key_size);
 
 	return VB2_SUCCESS;
 }
@@ -304,7 +304,7 @@ int vb2_load_fw_preamble(struct vb2_context *ctx)
 	 * TODO: we could move the preamble down over the firmware data key
 	 * since we don't need it anymore.
 	 */
-	ctx->workbuf_used = sd->workbuf_preamble_offset + pre_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_preamble_offset + pre_size);
 
 	return VB2_SUCCESS;
 }
diff --git a/firmware/lib21/api.c b/firmware/lib21/api.c
index c9774863..9f14579e 100644
--- a/firmware/lib21/api.c
+++ b/firmware/lib21/api.c
@@ -84,7 +84,7 @@ int vb21api_init_hash(struct vb2_context *ctx,
 
 		sd->workbuf_hash_offset = vb2_offset_of(ctx->workbuf, dc);
 		sd->workbuf_hash_size = dig_size;
-		ctx->workbuf_used = sd->workbuf_hash_offset + dig_size;
+		vb2_set_workbuf_used(ctx, sd->workbuf_hash_offset + dig_size);
 	}
 
 	sd->hash_tag = vb2_offset_of(ctx->workbuf, sig);
diff --git a/firmware/lib21/misc.c b/firmware/lib21/misc.c
index cb6da77e..361d6eda 100644
--- a/firmware/lib21/misc.c
+++ b/firmware/lib21/misc.c
@@ -148,8 +148,8 @@ int vb21_load_fw_keyblock(struct vb2_context *ctx)
 	sd->workbuf_data_key_size = packed_key->c.total_size;
 
 	/* Data key will persist in the workbuf after we return */
-	ctx->workbuf_used = sd->workbuf_data_key_offset +
-		sd->workbuf_data_key_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_data_key_offset +
+			     sd->workbuf_data_key_size);
 
 	return VB2_SUCCESS;
 }
@@ -240,8 +240,8 @@ int vb21_load_fw_preamble(struct vb2_context *ctx)
 	sd->workbuf_preamble_size = pre->c.total_size;
 
 	/* Preamble will persist in work buffer after we return */
-	ctx->workbuf_used = sd->workbuf_preamble_offset +
-		sd->workbuf_preamble_size;
+	vb2_set_workbuf_used(ctx, sd->workbuf_preamble_offset +
+			     sd->workbuf_preamble_size);
 
 	return VB2_SUCCESS;
 }
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 */