vboot: vb2_shared_data offsets should hang off parentstabilize-12428.B

vb2_shared_data struct has many offsets to other objects in the
workbuf.  They are all prefixed with `workbuf_`, e.g.:
  uint32_t workbuf_data_key_offset;
  uint32_t workbuf_data_key_size;

In order to adhere to a hierarchical structure on the workbuf,
remove the workbuf_ prefix from these symbols to reflect the
relationship between vb2_shared_data and its children more
accurately.

Create a helper function vb2_member_of to safely look up a child
of a particular object in the workbuf by offset.  Pointer
arithmetic to find vb2_shared_data children is replaced with
calls to this function.

BUG=b:124141368, chromium:994060
TEST=make clean && make runtests
BRANCH=none

Change-Id: Ia82417a35d2067ee5e4f42fea0396e6325127223
Signed-off-by: Joel Kitching <kitching@google.com>
Cq-Depend: chromium:1753400
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/1718264
Commit-Queue: Joel Kitching <kitching@chromium.org>
Tested-by: Joel Kitching <kitching@chromium.org>
Reviewed-by: Julius Werner <jwerner@chromium.org>

7 files changed, 122 insertions, 133 deletions

diff --git a/tests/vb20_api_kernel_tests.c b/tests/vb20_api_kernel_tests.c
index 3a300ab6..cdb87e88 100644
--- a/tests/vb20_api_kernel_tests.c
+++ b/tests/vb20_api_kernel_tests.c
@@ -92,18 +92,18 @@ static void reset_common_data(enum reset_type t)
 		uint8_t *kdata;
 
 		/* Create mock firmware preamble in the context */
-		sd->workbuf_preamble_offset = ctx.workbuf_used;
+		sd->preamble_offset = ctx.workbuf_used;
 		fwpre = (struct vb2_fw_preamble *)
-			(ctx.workbuf + sd->workbuf_preamble_offset);
+			vb2_member_of(sd, sd->preamble_offset);
 		k = &fwpre->kernel_subkey;
 		kdata = (uint8_t *)fwpre + sizeof(*fwpre);
 		memcpy(kdata, fw_kernel_key_data, sizeof(fw_kernel_key_data));
 		k->algorithm = 7;
 		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;
-		vb2_set_workbuf_used(&ctx, sd->workbuf_preamble_offset +
-				     sd->workbuf_preamble_size);
+		sd->preamble_size = sizeof(*fwpre) + k->key_size;
+		vb2_set_workbuf_used(&ctx, sd->preamble_offset +
+				     sd->preamble_size);
 
 	} else if (t == FOR_PHASE2) {
 		struct vb2_signature *sig;
@@ -111,18 +111,18 @@ static void reset_common_data(enum reset_type t)
 		uint8_t *sdata;
 
 		/* Create mock kernel data key */
-		sd->workbuf_data_key_offset = ctx.workbuf_used;
+		sd->data_key_offset = ctx.workbuf_used;
 		kdkey = (struct vb2_packed_key *)
-			(ctx.workbuf + sd->workbuf_data_key_offset);
+			vb2_member_of(sd, sd->data_key_offset);
 		kdkey->algorithm = VB2_ALG_RSA2048_SHA256;
-		sd->workbuf_data_key_size = sizeof(*kdkey);
-		vb2_set_workbuf_used(&ctx, sd->workbuf_data_key_offset +
-				     sd->workbuf_data_key_size);
+		sd->data_key_size = sizeof(*kdkey);
+		vb2_set_workbuf_used(&ctx, sd->data_key_offset +
+				     sd->data_key_size);
 
 		/* Create mock kernel preamble in the context */
-		sd->workbuf_preamble_offset = ctx.workbuf_used;
+		sd->preamble_offset = ctx.workbuf_used;
 		kpre = (struct vb2_kernel_preamble *)
-			(ctx.workbuf + sd->workbuf_preamble_offset);
+			vb2_member_of(sd, sd->preamble_offset);
 		sdata = (uint8_t *)kpre + sizeof(*kpre);
 
 		sig = &kpre->body_signature;
@@ -135,11 +135,11 @@ static void reset_common_data(enum reset_type t)
 				  sizeof(kernel_data));
 		vb2_digest_finalize(&dc, sdata, sig->sig_size);
 
-		sd->workbuf_preamble_size = sizeof(*kpre) + sig->sig_size;
+		sd->preamble_size = sizeof(*kpre) + sig->sig_size;
 		sd->vblock_preamble_offset =
-			0x10000 - sd->workbuf_preamble_size;
-		vb2_set_workbuf_used(&ctx, sd->workbuf_preamble_offset +
-				     sd->workbuf_preamble_size);
+			0x10000 - sd->preamble_size;
+		vb2_set_workbuf_used(&ctx, sd->preamble_offset +
+				     sd->preamble_size);
 
 	} else {
 		/* Set flags and versions for roll-forward */
@@ -224,19 +224,18 @@ static void phase1_tests(void)
 
 	/* Test successful call */
 	reset_common_data(FOR_PHASE1);
-	old_preamble_offset = sd->workbuf_preamble_offset;
+	old_preamble_offset = sd->preamble_offset;
 	TEST_SUCC(vb2api_kernel_phase1(&ctx), "phase1 good");
-	TEST_EQ(sd->workbuf_preamble_size, 0, "  no more fw preamble");
+	TEST_EQ(sd->preamble_size, 0, "  no more fw preamble");
 	/* Make sure normal key was loaded */
-	TEST_EQ(sd->workbuf_kernel_key_offset, old_preamble_offset,
+	TEST_EQ(sd->kernel_key_offset, old_preamble_offset,
 		"  workbuf key offset");
-	k = (struct vb2_packed_key *)
-		(ctx.workbuf + sd->workbuf_kernel_key_offset);
-	TEST_EQ(sd->workbuf_kernel_key_size, k->key_offset + k->key_size,
+	k = vb2_member_of(sd, sd->kernel_key_offset);
+	TEST_EQ(sd->kernel_key_size, k->key_offset + k->key_size,
 		"  workbuf key size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_kernel_key_offset +
-				sd->workbuf_kernel_key_size),
+		vb2_wb_round_up(sd->kernel_key_offset +
+				sd->kernel_key_size),
 		"  workbuf used");
 	TEST_EQ(k->algorithm, 7, "  key algorithm");
 	TEST_EQ(k->key_size, sizeof(fw_kernel_key_data), "  key_size");
@@ -248,20 +247,19 @@ static void phase1_tests(void)
 	reset_common_data(FOR_PHASE1);
 	ctx.flags |= VB2_CONTEXT_RECOVERY_MODE;
 	/* No preamble loaded in recovery mode */
-	ctx.workbuf_used = old_preamble_offset = sd->workbuf_preamble_offset;
-	sd->workbuf_preamble_offset = sd->workbuf_preamble_size = 0;
+	ctx.workbuf_used = old_preamble_offset = sd->preamble_offset;
+	sd->preamble_offset = sd->preamble_size = 0;
 	TEST_SUCC(vb2api_kernel_phase1(&ctx), "phase1 rec good");
-	TEST_EQ(sd->workbuf_preamble_size, 0, "no more fw preamble");
+	TEST_EQ(sd->preamble_size, 0, "no more fw preamble");
 	/* Make sure recovery key was loaded */
-	TEST_EQ(sd->workbuf_kernel_key_offset, old_preamble_offset,
+	TEST_EQ(sd->kernel_key_offset, old_preamble_offset,
 		"  workbuf key offset");
-	k = (struct vb2_packed_key *)
-		(ctx.workbuf + sd->workbuf_kernel_key_offset);
-	TEST_EQ(sd->workbuf_kernel_key_size, k->key_offset + k->key_size,
+	k = vb2_member_of(sd, sd->kernel_key_offset);
+	TEST_EQ(sd->kernel_key_size, k->key_offset + k->key_size,
 		"  workbuf key size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_kernel_key_offset +
-				sd->workbuf_kernel_key_size),
+		vb2_wb_round_up(sd->kernel_key_offset +
+				sd->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");
@@ -311,7 +309,7 @@ static void phase1_tests(void)
 
 	/* Failures while parsing subkey from firmware preamble */
 	reset_common_data(FOR_PHASE1);
-	sd->workbuf_preamble_size = 0;
+	sd->preamble_size = 0;
 	TEST_EQ(vb2api_kernel_phase1(&ctx), VB2_ERROR_API_KPHASE1_PREAMBLE,
 		"phase1 fw preamble");
 }
@@ -347,7 +345,7 @@ static void get_kernel_size_tests(void)
 	TEST_SUCC(vb2api_get_kernel_size(&ctx, NULL, NULL), "get size null");
 
 	reset_common_data(FOR_PHASE2);
-	sd->workbuf_preamble_size = 0;
+	sd->preamble_size = 0;
 	TEST_EQ(vb2api_get_kernel_size(&ctx, &offs, &size),
 		VB2_ERROR_API_GET_KERNEL_SIZE_PREAMBLE,
 		"get size no preamble");
@@ -361,7 +359,7 @@ static void verify_kernel_data_tests(void)
 		  "verify data good");
 
 	reset_common_data(FOR_PHASE2);
-	sd->workbuf_preamble_size = 0;
+	sd->preamble_size = 0;
 	TEST_EQ(vb2api_verify_kernel_data(&ctx, kernel_data,
 					  sizeof(kernel_data)),
 		VB2_ERROR_API_VERIFY_KDATA_PREAMBLE, "verify no preamble");
@@ -379,7 +377,7 @@ static void verify_kernel_data_tests(void)
 		VB2_ERROR_API_VERIFY_KDATA_WORKBUF, "verify workbuf");
 
 	reset_common_data(FOR_PHASE2);
-	sd->workbuf_data_key_size = 0;
+	sd->data_key_size = 0;
 	TEST_EQ(vb2api_verify_kernel_data(&ctx, kernel_data,
 					  sizeof(kernel_data)),
 		VB2_ERROR_API_VERIFY_KDATA_KEY, "verify no key");
diff --git a/tests/vb20_api_tests.c b/tests/vb20_api_tests.c
index ae6a216c..c3b48b78 100644
--- a/tests/vb20_api_tests.c
+++ b/tests/vb20_api_tests.c
@@ -74,12 +74,11 @@ static void reset_common_data(enum reset_type t)
 	retval_vb2_digest_finalize = VB2_SUCCESS;
 	retval_vb2_verify_digest = VB2_SUCCESS;
 
-	sd->workbuf_preamble_offset = ctx.workbuf_used;
-	sd->workbuf_preamble_size = sizeof(*pre);
-	vb2_set_workbuf_used(&ctx, sd->workbuf_preamble_offset
-			     + sd->workbuf_preamble_size);
-	pre = (struct vb2_fw_preamble *)
-		(ctx.workbuf + sd->workbuf_preamble_offset);
+	sd->preamble_offset = ctx.workbuf_used;
+	sd->preamble_size = sizeof(*pre);
+	vb2_set_workbuf_used(&ctx, sd->preamble_offset
+			     + sd->preamble_size);
+	pre = vb2_member_of(sd, sd->preamble_offset);
 	pre->body_signature.data_size = mock_body_size;
 	pre->body_signature.sig_size = mock_sig_size;
 	if (hwcrypto_state == HWCRYPTO_FORBIDDEN)
@@ -87,12 +86,11 @@ static void reset_common_data(enum reset_type t)
 	else
 		pre->flags = 0;
 
-	sd->workbuf_data_key_offset = ctx.workbuf_used;
-	sd->workbuf_data_key_size = sizeof(*k) + 8;
-	vb2_set_workbuf_used(&ctx, sd->workbuf_data_key_offset +
-			     sd->workbuf_data_key_size);
-	k = (struct vb2_packed_key *)
-		(ctx.workbuf + sd->workbuf_data_key_offset);
+	sd->data_key_offset = ctx.workbuf_used;
+	sd->data_key_size = sizeof(*k) + 8;
+	vb2_set_workbuf_used(&ctx, sd->data_key_offset +
+			     sd->data_key_size);
+	k = vb2_member_of(sd, sd->data_key_offset);
 	k->algorithm = mock_algorithm;
 
 	if (t == FOR_EXTEND_HASH || t == FOR_CHECK_HASH)
@@ -255,13 +253,13 @@ static void init_hash_tests(void)
 	wb_used_before = ctx.workbuf_used;
 	TEST_SUCC(vb2api_init_hash(&ctx, VB2_HASH_TAG_FW_BODY, &size),
 		  "init hash good");
-	TEST_EQ(sd->workbuf_hash_offset, wb_used_before,
+	TEST_EQ(sd->hash_offset, wb_used_before,
 		"hash context offset");
-	TEST_EQ(sd->workbuf_hash_size, sizeof(struct vb2_digest_context),
+	TEST_EQ(sd->hash_size, sizeof(struct vb2_digest_context),
 		"hash context size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_hash_offset +
-				sd->workbuf_hash_size),
+		vb2_wb_round_up(sd->hash_offset +
+				sd->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");
@@ -276,7 +274,7 @@ static void init_hash_tests(void)
 		VB2_ERROR_API_INIT_HASH_TAG, "init hash invalid tag");
 
 	reset_common_data(FOR_MISC);
-	sd->workbuf_preamble_size = 0;
+	sd->preamble_size = 0;
 	TEST_EQ(vb2api_init_hash(&ctx, VB2_HASH_TAG_FW_BODY, &size),
 		VB2_ERROR_API_INIT_HASH_PREAMBLE, "init hash preamble");
 
@@ -291,17 +289,17 @@ static void init_hash_tests(void)
 		VB2_ERROR_API_INIT_HASH_WORKBUF, "init hash workbuf");
 
 	reset_common_data(FOR_MISC);
-	sd->workbuf_data_key_size = 0;
+	sd->data_key_size = 0;
 	TEST_EQ(vb2api_init_hash(&ctx, VB2_HASH_TAG_FW_BODY, &size),
 		VB2_ERROR_API_INIT_HASH_DATA_KEY, "init hash data key");
 
 	reset_common_data(FOR_MISC);
-	sd->workbuf_data_key_size--;
+	sd->data_key_size--;
 	TEST_EQ(vb2api_init_hash(&ctx, VB2_HASH_TAG_FW_BODY, &size),
 		VB2_ERROR_UNPACK_KEY_SIZE, "init hash data key size");
 
 	reset_common_data(FOR_MISC);
-	k = (struct vb2_packed_key *)(ctx.workbuf + sd->workbuf_data_key_offset);
+	k = vb2_member_of(sd, sd->data_key_offset);
 	k->algorithm--;
 	TEST_EQ(vb2api_init_hash(&ctx, VB2_HASH_TAG_FW_BODY, &size),
 		VB2_ERROR_SHA_INIT_ALGORITHM, "init hash algorithm");
@@ -321,7 +319,7 @@ static void extend_hash_tests(void)
 	TEST_EQ(sd->hash_remaining_size, 0, "hash extend remaining 2");
 
 	reset_common_data(FOR_EXTEND_HASH);
-	sd->workbuf_hash_size = 0;
+	sd->hash_size = 0;
 	TEST_EQ(vb2api_extend_hash(&ctx, mock_body, mock_body_size),
 		VB2_ERROR_API_EXTEND_HASH_WORKBUF, "hash extend no workbuf");
 
@@ -336,7 +334,7 @@ static void extend_hash_tests(void)
 	if (hwcrypto_state != HWCRYPTO_ENABLED) {
 		reset_common_data(FOR_EXTEND_HASH);
 		dc = (struct vb2_digest_context *)
-			(ctx.workbuf + sd->workbuf_hash_offset);
+			vb2_member_of(sd, sd->hash_offset);
 		dc->hash_alg = mock_hash_alg + 1;
 		TEST_EQ(vb2api_extend_hash(&ctx, mock_body, mock_body_size),
 			VB2_ERROR_SHA_EXTEND_ALGORITHM, "hash extend fail");
@@ -366,12 +364,12 @@ static void check_hash_tests(void)
 		"check digest wrong size");
 
 	reset_common_data(FOR_CHECK_HASH);
-	sd->workbuf_preamble_size = 0;
+	sd->preamble_size = 0;
 	TEST_EQ(vb2api_check_hash(&ctx),
 		VB2_ERROR_API_CHECK_HASH_PREAMBLE, "check hash preamble");
 
 	reset_common_data(FOR_CHECK_HASH);
-	sd->workbuf_hash_size = 0;
+	sd->hash_size = 0;
 	TEST_EQ(vb2api_check_hash(&ctx),
 		VB2_ERROR_API_CHECK_HASH_WORKBUF, "check hash no workbuf");
 
@@ -395,18 +393,17 @@ static void check_hash_tests(void)
 		VB2_ERROR_API_CHECK_HASH_TAG, "check hash tag");
 
 	reset_common_data(FOR_CHECK_HASH);
-	sd->workbuf_data_key_size = 0;
+	sd->data_key_size = 0;
 	TEST_EQ(vb2api_check_hash(&ctx),
 		VB2_ERROR_API_CHECK_HASH_DATA_KEY, "check hash data key");
 
 	reset_common_data(FOR_CHECK_HASH);
-	sd->workbuf_data_key_size--;
+	sd->data_key_size--;
 	TEST_EQ(vb2api_check_hash(&ctx),
 		VB2_ERROR_UNPACK_KEY_SIZE, "check hash data key size");
 
 	reset_common_data(FOR_CHECK_HASH);
-	pre = (struct vb2_fw_preamble *)
-		(ctx.workbuf + sd->workbuf_preamble_offset);
+	pre = vb2_member_of(sd, sd->preamble_offset);
 	pre->body_signature.sig_size++;
 	TEST_EQ(vb2api_check_hash(&ctx),
 		VB2_ERROR_VDATA_SIG_SIZE, "check hash sig size");
diff --git a/tests/vb20_kernel_tests.c b/tests/vb20_kernel_tests.c
index 5490ccce..52899cca 100644
--- a/tests/vb20_kernel_tests.c
+++ b/tests/vb20_kernel_tests.c
@@ -248,26 +248,26 @@ 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,
+	TEST_EQ(sd->data_key_offset, wb_used_before,
 		"keyblock data key offset");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_data_key_offset +
-				sd->workbuf_data_key_size),
+		vb2_wb_round_up(sd->data_key_offset +
+				sd->data_key_size),
 		"workbuf used");
 
 	/* Make sure data key was properly saved */
-	k = (struct vb2_packed_key *)(ctx.workbuf + sd->workbuf_data_key_offset);
+	k = vb2_member_of(sd, sd->data_key_offset);
 	TEST_EQ(k->algorithm, 7, "data key algorithm");
 	TEST_EQ(k->key_version, 2, "data key version");
 	TEST_EQ(k->key_size, sizeof(mock_vblock.k.data_key_data),
 		"data key size");
-	TEST_EQ(memcmp(ctx.workbuf + sd->workbuf_data_key_offset +
-		       k->key_offset, mock_vblock.k.data_key_data,
+	TEST_EQ(memcmp(vb2_member_of(k, k->key_offset),
+		       mock_vblock.k.data_key_data,
 		       sizeof(mock_vblock.k.data_key_data)),
 		0, "data key data");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_data_key_offset +
-				sd->workbuf_data_key_size),
+		vb2_wb_round_up(sd->data_key_offset +
+				sd->data_key_size),
 		"workbuf used after");
 
 	/* Test failures */
@@ -403,17 +403,17 @@ static void load_kernel_preamble_tests(void)
 	wb_used_before = ctx.workbuf_used;
 	TEST_SUCC(vb2_load_kernel_preamble(&ctx), "preamble good");
 	TEST_EQ(sd->kernel_version, 0x20002, "combined version");
-	TEST_EQ(sd->workbuf_preamble_offset, wb_used_before,
+	TEST_EQ(sd->preamble_offset, wb_used_before,
 		"preamble offset");
-	TEST_EQ(sd->workbuf_preamble_size, pre->preamble_size, "preamble size");
+	TEST_EQ(sd->preamble_size, pre->preamble_size, "preamble size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_preamble_offset +
-				sd->workbuf_preamble_size),
+		vb2_wb_round_up(sd->preamble_offset +
+				sd->preamble_size),
 		"workbuf used");
 
 	/* Expected failures */
 	reset_common_data(FOR_PREAMBLE);
-	sd->workbuf_data_key_size = 0;
+	sd->data_key_size = 0;
 	TEST_EQ(vb2_load_kernel_preamble(&ctx),
 		VB2_ERROR_KERNEL_PREAMBLE2_DATA_KEY,
 		"preamble no data key");
diff --git a/tests/vb20_misc_tests.c b/tests/vb20_misc_tests.c
index 46e3f036..1d23a6d6 100644
--- a/tests/vb20_misc_tests.c
+++ b/tests/vb20_misc_tests.c
@@ -188,26 +188,26 @@ 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,
+	TEST_EQ(sd->data_key_offset, wb_used_before,
 		"keyblock data key offset");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_data_key_offset +
-				sd->workbuf_data_key_size),
+		vb2_wb_round_up(sd->data_key_offset +
+				sd->data_key_size),
 		"workbuf used");
 
 	/* Make sure data key was properly saved */
-	k = (struct vb2_packed_key *)(ctx.workbuf + sd->workbuf_data_key_offset);
+	k = vb2_member_of(sd, sd->data_key_offset);
 	TEST_EQ(k->algorithm, 7, "data key algorithm");
 	TEST_EQ(k->key_version, 2, "data key version");
 	TEST_EQ(k->key_size, sizeof(mock_vblock.k.data_key_data),
 		"data key size");
-	TEST_EQ(memcmp(ctx.workbuf + sd->workbuf_data_key_offset +
-		       k->key_offset, mock_vblock.k.data_key_data,
+	TEST_EQ(memcmp(vb2_member_of(k, k->key_offset),
+		       mock_vblock.k.data_key_data,
 		       sizeof(mock_vblock.k.data_key_data)),
 		0, "data key data");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_data_key_offset +
-				sd->workbuf_data_key_size),
+		vb2_wb_round_up(sd->data_key_offset +
+				sd->data_key_size),
 		"workbuf used after");
 
 	/* Test failures */
@@ -292,17 +292,17 @@ static void verify_preamble_tests(void)
 	wb_used_before = ctx.workbuf_used;
 	TEST_SUCC(vb2_load_fw_preamble(&ctx), "preamble good");
 	TEST_EQ(sd->fw_version, 0x20002, "combined version");
-	TEST_EQ(sd->workbuf_preamble_offset, wb_used_before,
+	TEST_EQ(sd->preamble_offset, wb_used_before,
 		"preamble offset");
-	TEST_EQ(sd->workbuf_preamble_size, pre->preamble_size, "preamble size");
+	TEST_EQ(sd->preamble_size, pre->preamble_size, "preamble size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_preamble_offset +
-				sd->workbuf_preamble_size),
+		vb2_wb_round_up(sd->preamble_offset +
+				sd->preamble_size),
 		"workbuf used");
 
 	/* Expected failures */
 	reset_common_data(FOR_PREAMBLE);
-	sd->workbuf_data_key_size = 0;
+	sd->data_key_size = 0;
 	TEST_EQ(vb2_load_fw_preamble(&ctx),
 		VB2_ERROR_FW_PREAMBLE2_DATA_KEY,
 		"preamble no data key");
diff --git a/tests/vb21_api_tests.c b/tests/vb21_api_tests.c
index 53177957..8ea9bf3a 100644
--- a/tests/vb21_api_tests.c
+++ b/tests/vb21_api_tests.c
@@ -91,9 +91,8 @@ static void reset_common_data(enum reset_type t)
 
 	vb2_private_key_hash(&hash_key, mock_hash_alg);
 
-	sd->workbuf_preamble_offset = ctx.workbuf_used;
-	pre = (struct vb21_fw_preamble *)
-		(ctx.workbuf + sd->workbuf_preamble_offset);
+	sd->preamble_offset = ctx.workbuf_used;
+	pre = vb2_member_of(sd, sd->preamble_offset);
 	pre->hash_count = 3;
 	pre->hash_offset = sig_offset = sizeof(*pre);
 	if (hwcrypto_state == HWCRYPTO_FORBIDDEN)
@@ -111,9 +110,9 @@ static void reset_common_data(enum reset_type t)
 		free(sig);
 	}
 
-	sd->workbuf_preamble_size = sig_offset;
-	ctx.workbuf_used = vb2_wb_round_up(sd->workbuf_preamble_offset +
-					   sd->workbuf_preamble_size);
+	sd->preamble_size = sig_offset;
+	ctx.workbuf_used = vb2_wb_round_up(sd->preamble_offset +
+					   sd->preamble_size);
 
 	if (t == FOR_EXTEND_HASH || t == FOR_CHECK_HASH)
 		vb21api_init_hash(&ctx, test_id, NULL);
@@ -202,23 +201,22 @@ static void init_hash_tests(void)
 	uint32_t size;
 
 	reset_common_data(FOR_MISC);
-	pre = (struct vb21_fw_preamble *)
-		(ctx.workbuf + sd->workbuf_preamble_offset);
+	pre = vb2_member_of(sd, sd->preamble_offset);
 	sig = (struct vb21_signature *)((uint8_t *)pre + pre->hash_offset);
 
 	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,
+	TEST_EQ(sd->hash_offset, wb_used_before,
 		"hash context offset");
-	TEST_EQ(sd->workbuf_hash_size, sizeof(struct vb2_digest_context),
+	TEST_EQ(sd->hash_size, sizeof(struct vb2_digest_context),
 		"hash context size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_hash_offset +
-				sd->workbuf_hash_size),
+		vb2_wb_round_up(sd->hash_offset +
+				sd->hash_size),
 		"hash uses workbuf");
 	TEST_EQ(sd->hash_tag,
-		sd->workbuf_preamble_offset + pre->hash_offset,
+		sd->preamble_offset + pre->hash_offset,
 		"hash signature offset");
 	TEST_EQ(sd->hash_remaining_size, mock_body_size, "hash remaining");
 
@@ -227,7 +225,7 @@ static void init_hash_tests(void)
 		  "init hash again");
 	TEST_EQ(ctx.workbuf_used, wb_used_before, "init hash reuses context");
 	TEST_EQ(sd->hash_tag,
-		sd->workbuf_preamble_offset + pre->hash_offset +
+		sd->preamble_offset + pre->hash_offset +
 		2 * mock_sig_size,
 		"hash signature offset 2");
 
@@ -236,7 +234,7 @@ static void init_hash_tests(void)
 		VB2_ERROR_API_INIT_HASH_ID, "init hash invalid id");
 
 	reset_common_data(FOR_MISC);
-	sd->workbuf_preamble_size = 0;
+	sd->preamble_size = 0;
 	TEST_EQ(vb21api_init_hash(&ctx, test_id, &size),
 		VB2_ERROR_API_INIT_HASH_PREAMBLE, "init hash preamble");
 
@@ -273,7 +271,7 @@ static void extend_hash_tests(void)
 	TEST_EQ(sd->hash_remaining_size, 0, "hash extend remaining 2");
 
 	reset_common_data(FOR_EXTEND_HASH);
-	sd->workbuf_hash_size = 0;
+	sd->hash_size = 0;
 	TEST_EQ(vb2api_extend_hash(&ctx, mock_body, mock_body_size),
 		VB2_ERROR_API_EXTEND_HASH_WORKBUF, "hash extend no workbuf");
 
@@ -293,7 +291,7 @@ static void extend_hash_tests(void)
 	} else {
 		reset_common_data(FOR_EXTEND_HASH);
 		dc = (struct vb2_digest_context *)
-			(ctx.workbuf + sd->workbuf_hash_offset);
+			vb2_member_of(sd, sd->hash_offset);
 		dc->hash_alg = VB2_HASH_INVALID;
 		TEST_EQ(vb2api_extend_hash(&ctx, mock_body, mock_body_size),
 			VB2_ERROR_SHA_EXTEND_ALGORITHM, "hash extend fail");
@@ -307,11 +305,9 @@ static void check_hash_tests(void)
 	struct vb2_digest_context *dc;
 
 	reset_common_data(FOR_CHECK_HASH);
-	pre = (struct vb21_fw_preamble *)
-		(ctx.workbuf + sd->workbuf_preamble_offset);
-	sig = (struct vb21_signature *)((uint8_t *)pre + pre->hash_offset);
-	dc = (struct vb2_digest_context *)
-		(ctx.workbuf + sd->workbuf_hash_offset);
+	pre = vb2_member_of(sd, sd->preamble_offset);
+	sig = vb2_member_of(pre, pre->hash_offset);
+	dc = vb2_member_of(sd, sd->hash_offset);
 
 	TEST_SUCC(vb21api_check_hash(&ctx), "check hash good");
 
@@ -321,7 +317,7 @@ static void check_hash_tests(void)
 		VB2_ERROR_API_CHECK_HASH_TAG, "check hash tag");
 
 	reset_common_data(FOR_CHECK_HASH);
-	sd->workbuf_hash_size = 0;
+	sd->hash_size = 0;
 	TEST_EQ(vb21api_check_hash(&ctx),
 		VB2_ERROR_API_CHECK_HASH_WORKBUF, "check hash no workbuf");
 
diff --git a/tests/vb21_misc_tests.c b/tests/vb21_misc_tests.c
index 37c0ea3a..383e3d02 100644
--- a/tests/vb21_misc_tests.c
+++ b/tests/vb21_misc_tests.c
@@ -193,27 +193,26 @@ 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,
+	TEST_EQ(sd->data_key_offset, wb_used_before,
 		"keyblock data key offset");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_data_key_offset +
-				sd->workbuf_data_key_size),
+		vb2_wb_round_up(sd->data_key_offset +
+				sd->data_key_size),
 		"workbuf used");
 
 	/* Make sure data key was properly saved */
-	k = (struct vb21_packed_key *)(ctx.workbuf +
-				       sd->workbuf_data_key_offset);
+	k = vb2_member_of(sd, sd->data_key_offset);
 	TEST_EQ(k->sig_alg, VB2_SIG_RSA4096, "data key algorithm");
 	TEST_EQ(k->key_version, 2, "data key version");
 	TEST_EQ(k->key_size, sizeof(mock_vblock.k.data_key_data),
 		"data key size");
-	TEST_EQ(memcmp(ctx.workbuf + sd->workbuf_data_key_offset +
-		       k->key_offset, mock_vblock.k.data_key_data,
+	TEST_EQ(memcmp(vb2_member_of(k, k->key_offset),
+		       mock_vblock.k.data_key_data,
 		       sizeof(mock_vblock.k.data_key_data)),
 		0, "data key data");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_data_key_offset +
-				sd->workbuf_data_key_size),
+		vb2_wb_round_up(sd->data_key_offset +
+				sd->data_key_size),
 		"workbuf used after");
 
 	/* Test failures */
@@ -295,22 +294,22 @@ static void load_preamble_tests(void)
 
 	/* Test successful call */
 	reset_common_data(FOR_PREAMBLE);
-	data_key_offset_before = sd->workbuf_data_key_offset;
+	data_key_offset_before = sd->data_key_offset;
 	TEST_SUCC(vb21_load_fw_preamble(&ctx), "preamble good");
 	TEST_EQ(sd->fw_version, 0x20002, "combined version");
-	TEST_EQ(sd->workbuf_preamble_offset, data_key_offset_before,
+	TEST_EQ(sd->preamble_offset, data_key_offset_before,
 		"preamble offset");
-	TEST_EQ(sd->workbuf_preamble_size, pre->c.total_size, "preamble size");
+	TEST_EQ(sd->preamble_size, pre->c.total_size, "preamble size");
 	TEST_EQ(ctx.workbuf_used,
-		vb2_wb_round_up(sd->workbuf_preamble_offset +
-				sd->workbuf_preamble_size),
+		vb2_wb_round_up(sd->preamble_offset +
+				sd->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");
+	TEST_EQ(sd->data_key_offset, 0, "data key offset gone");
+	TEST_EQ(sd->data_key_size, 0, "data key size gone");
 
 	/* Expected failures */
 	reset_common_data(FOR_PREAMBLE);
-	sd->workbuf_data_key_size = 0;
+	sd->data_key_size = 0;
 	TEST_EQ(vb21_load_fw_preamble(&ctx),
 		VB2_ERROR_FW_PREAMBLE2_DATA_KEY,
 		"preamble no data key");
diff --git a/tests/vb2_misc_tests.c b/tests/vb2_misc_tests.c
index a9f82531..c2763e4e 100644
--- a/tests/vb2_misc_tests.c
+++ b/tests/vb2_misc_tests.c
@@ -228,8 +228,7 @@ static void gbb_tests(void)
 	TEST_SUCC(vb2_fw_parse_gbb(&ctx), "parse gbb");
 	/* Manually calculate the location of GBB since we have mocked out the
 	   original definition of vb2_get_gbb. */
-	struct vb2_gbb_header *current_gbb =
-		(struct vb2_gbb_header *)((void *)sd + sd->gbb_offset);
+	struct vb2_gbb_header *current_gbb = vb2_member_of(sd, sd->gbb_offset);
 	TEST_SUCC(memcmp(&gbbsrc, current_gbb, sizeof(gbbsrc)),
 		  "copy gbb contents");
 	TEST_EQ(used_before, ctx.workbuf_used - sizeof(gbbsrc),