vboot: move vb2_context inside vb2_shared_data (persistent context)

Move vb2_context to live inside of vb2_shared_data, instead of
in a separate memory space allocated by the caller.

See design doc:
http://go/vboot2-persistent-context

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

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

12 files changed, 188 insertions, 116 deletions

diff --git a/firmware/2lib/2api.c b/firmware/2lib/2api.c
index 9b7f9533..a94f0193 100644
--- a/firmware/2lib/2api.c
+++ b/firmware/2lib/2api.c
@@ -20,11 +20,7 @@
 vb2_error_t vb2api_fw_phase1(struct vb2_context *ctx)
 {
 	vb2_error_t rv;
-	struct vb2_shared_data *sd;
-
-	/* Initialize the vboot context if it hasn't been yet */
-	vb2_init_context(ctx);
-	sd = vb2_get_sd(ctx);
+	struct vb2_shared_data *sd = vb2_get_sd(ctx);
 
 	/* Initialize NV context */
 	vb2_nv_init(ctx);
diff --git a/firmware/2lib/2misc.c b/firmware/2lib/2misc.c
index a1635f7a..deb4f885 100644
--- a/firmware/2lib/2misc.c
+++ b/firmware/2lib/2misc.c
@@ -29,19 +29,21 @@ vb2_error_t vb2_validate_gbb_signature(uint8_t *sig)
 test_mockable
 struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *ctx)
 {
-	return (struct vb2_gbb_header *)
-	       ((void *)vb2_get_sd(ctx) + vb2_get_sd(ctx)->gbb_offset);
+	struct vb2_shared_data *sd = vb2_get_sd(ctx);
+	return (struct vb2_gbb_header *)((void *)sd + sd->gbb_offset);
 }
 
 void vb2_workbuf_from_ctx(struct vb2_context *ctx, struct vb2_workbuf *wb)
 {
-	vb2_workbuf_init(wb, ctx->workbuf + ctx->workbuf_used,
-			 ctx->workbuf_size - ctx->workbuf_used);
+	struct vb2_shared_data *sd = vb2_get_sd(ctx);
+	vb2_workbuf_init(wb, (void *)sd + sd->workbuf_used,
+			 sd->workbuf_size - sd->workbuf_used);
 }
 
 void vb2_set_workbuf_used(struct vb2_context *ctx, uint32_t used)
 {
-	ctx->workbuf_used = vb2_wb_round_up(used);
+	struct vb2_shared_data *sd = vb2_get_sd(ctx);
+	sd->workbuf_used = vb2_wb_round_up(used);
 }
 
 vb2_error_t vb2_read_gbb_header(struct vb2_context *ctx,
@@ -79,9 +81,6 @@ vb2_error_t vb2_read_gbb_header(struct vb2_context *ctx,
 
 void vb2api_fail(struct vb2_context *ctx, uint8_t reason, uint8_t subcode)
 {
-	/* Initialize the vboot context if it hasn't been yet */
-	vb2_init_context(ctx);
-
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
 
 	/* If NV data hasn't been initialized, initialize it now */
@@ -129,46 +128,82 @@ void vb2api_fail(struct vb2_context *ctx, uint8_t reason, uint8_t subcode)
 	}
 }
 
+vb2_error_t vb2api_init(void *workbuf, uint32_t size,
+			struct vb2_context **ctxptr)
+{
+	struct vb2_shared_data *sd = workbuf;
+	*ctxptr = NULL;
+
+	if (!vb2_aligned(workbuf, VB2_WORKBUF_ALIGN))
+		return VB2_ERROR_WORKBUF_ALIGN;
+
+	if (size < vb2_wb_round_up(sizeof(*sd)))
+		return VB2_ERROR_WORKBUF_SMALL;
+
+	/* Zero out vb2_shared_data (which includes vb2_context). */
+	memset(sd, 0, sizeof(*sd));
+
+	/* Initialize shared data. */
+	sd->magic = VB2_SHARED_DATA_MAGIC;
+	sd->struct_version_major = VB2_SHARED_DATA_VERSION_MAJOR;
+	sd->struct_version_minor = VB2_SHARED_DATA_VERSION_MINOR;
+	sd->workbuf_size = size;
+	sd->workbuf_used = vb2_wb_round_up(sizeof(*sd));
+
+	*ctxptr = &sd->ctx;
+	return VB2_SUCCESS;
+}
+
 #pragma GCC diagnostic push
 /* Don't warn for the version_minor check even if the checked version is 0. */
 #pragma GCC diagnostic ignored "-Wtype-limits"
-vb2_error_t vb2_init_context(struct vb2_context *ctx)
+vb2_error_t vb2api_relocate(void *new_workbuf, void *cur_workbuf, uint32_t size,
+			    struct vb2_context **ctxptr)
 {
-	struct vb2_shared_data *sd = vb2_get_sd(ctx);
+	struct vb2_shared_data *sd = cur_workbuf;
 
-	/* Don't do anything if context and workbuf have already been
-	 * initialized. */
-	if (ctx->workbuf_used) {
-		if (sd->magic != VB2_SHARED_DATA_MAGIC)
-			return VB2_ERROR_SHARED_DATA_MAGIC;
+	if (!vb2_aligned(new_workbuf, VB2_WORKBUF_ALIGN))
+		return VB2_ERROR_WORKBUF_ALIGN;
 
-		if (sd->struct_version_major != VB2_SHARED_DATA_VERSION_MAJOR ||
-		    sd->struct_version_minor < VB2_SHARED_DATA_VERSION_MINOR)
-			return VB2_ERROR_SHARED_DATA_VERSION;
+	/* Check magic and version. */
+	if (sd->magic != VB2_SHARED_DATA_MAGIC)
+		return VB2_ERROR_SHARED_DATA_MAGIC;
 
-		return VB2_SUCCESS;
-	}
+	if (sd->struct_version_major != VB2_SHARED_DATA_VERSION_MAJOR ||
+	    sd->struct_version_minor < VB2_SHARED_DATA_VERSION_MINOR)
+		return VB2_ERROR_SHARED_DATA_VERSION;
 
-	/*
-	 * Workbuf had better be big enough for our shared data struct and
-	 * aligned.  Not much we can do if it isn't; we'll die before we can
-	 * store a recovery reason.
-	 */
-	if (ctx->workbuf_size < sizeof(*sd))
-		return VB2_ERROR_INITCTX_WORKBUF_SMALL;
-	if (!vb2_aligned(ctx->workbuf, VB2_WORKBUF_ALIGN))
-		return VB2_ERROR_INITCTX_WORKBUF_ALIGN;
+	/* Check workbuf integrity. */
+	if (sd->workbuf_used < vb2_wb_round_up(sizeof(*sd)))
+		return VB2_ERROR_WORKBUF_INVALID;
+
+	if (sd->workbuf_size < sd->workbuf_used)
+		return VB2_ERROR_WORKBUF_INVALID;
+
+	if (sd->workbuf_used > size)
+		return VB2_ERROR_WORKBUF_SMALL;
+
+	/* Relocate if necessary. */
+	if (cur_workbuf != new_workbuf)
+		memmove(new_workbuf, cur_workbuf, sd->workbuf_used);
+
+	/* Set the new size, and return the context pointer. */
+	sd = new_workbuf;
+	sd->workbuf_size = size;
+	*ctxptr = &sd->ctx;
 
-	/* Initialize the shared data at the start of the work buffer */
-	memset(sd, 0, sizeof(*sd));
-	sd->magic = VB2_SHARED_DATA_MAGIC;
-	sd->struct_version_major = VB2_SHARED_DATA_VERSION_MAJOR;
-	sd->struct_version_minor = VB2_SHARED_DATA_VERSION_MINOR;
-	ctx->workbuf_used = vb2_wb_round_up(sizeof(*sd));
 	return VB2_SUCCESS;
 }
 #pragma GCC diagnostic pop
 
+vb2_error_t vb2api_reinit(void *workbuf, struct vb2_context **ctxptr)
+{
+	/* Blindly retrieve workbuf_size.  vb2api_relocate() will
+	   perform workbuf validation checks. */
+	struct vb2_shared_data *sd = workbuf;
+	return vb2api_relocate(workbuf, workbuf, sd->workbuf_size, ctxptr);
+}
+
 void vb2_check_recovery(struct vb2_context *ctx)
 {
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
@@ -232,7 +267,7 @@ vb2_error_t vb2_fw_init_gbb(struct vb2_context *ctx)
 
 	/* Keep on the work buffer permanently */
 	sd->gbb_offset = vb2_offset_of(sd, gbb);
-	ctx->workbuf_used = vb2_offset_of(ctx->workbuf, wb.buf);
+	vb2_set_workbuf_used(ctx, vb2_offset_of(ctx, wb.buf));
 
 	/* Set any context flags based on GBB flags */
 	if (gbb->flags & VB2_GBB_FLAG_DISABLE_FWMP)
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index 9aa031a9..0a713265 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -196,29 +196,33 @@ enum vb2_context_flags {
 	VB2_CONTEXT_NO_SECDATA_FWMP = (1 << 21),
 };
 
+/* Helper for aligning fields in vb2_context. */
+#define VB2_PAD_STRUCT3(size, align, count) \
+	uint8_t _pad##count[align - (((size - 1) % align) + 1)]
+#define VB2_PAD_STRUCT2(size, align, count) VB2_PAD_STRUCT3(size, align, count)
+#define VB2_PAD_STRUCT(size, align) VB2_PAD_STRUCT2(size, align, __COUNTER__)
+
 /*
  * Context for firmware verification.  Pass this to all vboot APIs.
  *
- * Caller may relocate this between calls to vboot APIs.
+ * Context is stored as part of vb2_shared_data, initialized with vb2api_init().
+ * Subsequent retrieval of the context object should be done by calling
+ * vb2api_reinit(), e.g. if switching firmware applications.
+ *
+ * The context struct can be seen as the "publicly accessible" portion of
+ * vb2_shared_data, and thus does not require its own magic and version fields.
  */
 struct vb2_context {
+
 	/**********************************************************************
-	 * Fields which must be initialized by caller.
+	 * Fields caller must initialize before calling any API functions.
 	 */
 
 	/*
 	 * Flags; see vb2_context_flags.  Some flags may only be set by caller
 	 * prior to calling vboot functions.
 	 */
-	uint32_t flags;
-
-	/*
-	 * Work buffer, and length in bytes.  Caller may relocate this between
-	 * calls to vboot APIs; it contains no internal pointers.  Caller must
-	 * not examine the contents of this work buffer directly.
-	 */
-	uint8_t *workbuf;
-	uint32_t workbuf_size;
+	uint64_t flags;
 
 	/*
 	 * Non-volatile data.  Caller must fill this from some non-volatile
@@ -228,6 +232,7 @@ struct vb2_context {
 	 * and then clear the flag.
 	 */
 	uint8_t nvdata[VB2_NVDATA_SIZE_V2];
+	VB2_PAD_STRUCT(VB2_NVDATA_SIZE_V2, 8);
 
 	/*
 	 * Secure data for firmware verification stage.  Caller must fill this
@@ -237,26 +242,7 @@ struct vb2_context {
 	 * secure non-volatile location and then clear the flag.
 	 */
 	uint8_t secdata_firmware[VB2_SECDATA_FIRMWARE_SIZE];
-
-	/*
-	 * Context pointer for use by caller.  Verified boot never looks at
-	 * this.  Put context here if you need it for APIs that verified boot
-	 * may call (vb2ex_...() functions).
-	 */
-	void *non_vboot_context;
-
-	/**********************************************************************
-	 * Fields caller may examine after calling vb2api_fw_phase1().  Caller
-	 * must set these fields to 0 before calling any vboot functions.
-	 */
-
-	/*
-	 * Amount of work buffer used so far.  Verified boot sub-calls use
-	 * this to know where the unused work area starts.  Caller may use
-	 * this between calls to vboot APIs to know how much data must be
-	 * copied when relocating the work buffer.
-	 */
-	uint32_t workbuf_used;
+	VB2_PAD_STRUCT(VB2_SECDATA_FIRMWARE_SIZE, 8);
 
 	/**********************************************************************
 	 * Fields caller must initialize before calling vb2api_kernel_phase1().
@@ -270,6 +256,7 @@ struct vb2_context {
 	 * to the secure non-volatile location and then clear the flag.
 	 */
 	uint8_t secdata_kernel[VB2_SECDATA_KERNEL_SIZE];
+	VB2_PAD_STRUCT(VB2_SECDATA_KERNEL_SIZE, 8);
 
 	/*
 	 * Firmware management parameters (FWMP) secure data.  Caller must fill
@@ -282,6 +269,7 @@ struct vb2_context {
 	 * location and then clear the flag.
 	 */
 	uint8_t secdata_fwmp[VB2_SECDATA_FWMP_MAX_SIZE];
+	VB2_PAD_STRUCT(VB2_SECDATA_FWMP_MAX_SIZE, 8);
 };
 
 /* Resource index for vb2ex_read_resource() */
@@ -426,6 +414,65 @@ enum vb2_pcr_digest {
  */
 
 /**
+ * Initialize verified boot data structures.
+ *
+ * Needs to be called once per boot, before using any API functions that
+ * accept a vb2_context object.  Sets up the vboot work buffer, as well as
+ * vb2_shared_data and vb2_context.  A pointer to the context object is
+ * written to ctxptr.  After transitioning between different firmware
+ * applications, or any time the context pointer is lost, vb2api_reinit()
+ * should be used to restore access to the context and data on the workbuf.
+ *
+ * If the workbuf needs to be relocated, call vb2api_relocate() instead
+ * of copying memory manually.
+ *
+ * @param workbuf	Workbuf memory location to initialize
+ * @param size		Size of workbuf being initialized
+ * @param ctxptr	Pointer to a context pointer to be filled in
+ * @return VB2_SUCCESS, or non-zero error code.
+ */
+vb2_error_t vb2api_init(void *workbuf, uint32_t size,
+			struct vb2_context **ctxptr);
+
+/**
+ * Reinitialize vboot data structures.
+ *
+ * After transitioning between different firmware applications, or any time the
+ * context pointer is lost, this function should be called to restore access to
+ * the workbuf.  A pointer to the context object is written to ctxptr.  Returns
+ * an error if the vboot work buffer is inconsistent.
+ *
+ * If the workbuf needs to be relocated, call vb2api_relocate() instead
+ * of copying memory manually.
+ *
+ * @param workbuf	Workbuf memory location to check
+ * @param ctxptr	Pointer to a context pointer to be filled in
+ * @return VB2_SUCCESS, or non-zero error code.
+ */
+vb2_error_t vb2api_reinit(void *workbuf, struct vb2_context **ctxptr);
+
+/**
+ * Relocate vboot data structures.
+ *
+ * Move the vboot work buffer from one memory location to another, and expand
+ * or contract the workbuf to fit.  The target memory location may be the same
+ * as the original (used for a "resize" operation), and it is safe to call this
+ * function with overlapping memory regions.
+ *
+ * A pointer to the context object is written to ctxptr.  Returns an error if
+ * the vboot work buffer is inconsistent, or if the new memory space is too
+ * small to contain the work buffer.
+ *
+ * @param new_workbuf	Target workbuf memory location
+ * @param cur_workbuf	Original workbuf memory location to relocate
+ * @param size		Target size of relocated workbuf
+ * @param ctxptr	Pointer to a context pointer to be filled in
+ * @return VB2_SUCCESS, or non-zero error code.
+ */
+vb2_error_t vb2api_relocate(void *new_workbuf, void *cur_workbuf, uint32_t size,
+			    struct vb2_context **ctxptr);
+
+/**
  * Check the validity of firmware secure storage context.
  *
  * Checks version and CRC.
diff --git a/firmware/2lib/include/2misc.h b/firmware/2lib/include/2misc.h
index 5ce10705..d9be9730 100644
--- a/firmware/2lib/include/2misc.h
+++ b/firmware/2lib/include/2misc.h
@@ -14,6 +14,10 @@
 struct vb2_gbb_header;
 struct vb2_workbuf;
 
+#define vb2_container_of(ptr, type, member) ({                     \
+	const typeof(((type *)0)->member) *__mptr = (ptr);         \
+	(type *)((uint8_t *)__mptr - offsetof(type, member) );})   \
+
 /**
  * Get the shared data pointer from the vboot context
  *
@@ -22,7 +26,7 @@ struct vb2_workbuf;
  */
 static inline struct vb2_shared_data *vb2_get_sd(struct vb2_context *ctx)
 {
-	return (struct vb2_shared_data *)ctx->workbuf;
+	return vb2_container_of(ctx, struct vb2_shared_data, ctx);
 }
 
 /**
@@ -73,18 +77,6 @@ vb2_error_t vb2_read_gbb_header(struct vb2_context *ctx,
 				struct vb2_gbb_header *gbb);
 
 /**
- * Set up the verified boot context data, if not already set up.
- *
- * This uses ctx->workbuf_used=0 as a flag to indicate that the data has not
- * yet been set up.  Caller must set that before calling any vboot functions;
- * see 2api.h.
- *
- * @param ctx		Vboot context to initialize
- * @return VB2_SUCCESS, or error code on error.
- */
-vb2_error_t vb2_init_context(struct vb2_context *ctx);
-
-/**
  * Check for recovery reasons we can determine early in the boot process.
  *
  * On exit, check ctx->flags for VB2_CONTEXT_RECOVERY_MODE; if present, jump to
diff --git a/firmware/2lib/include/2return_codes.h b/firmware/2lib/include/2return_codes.h
index 49ee50ab..9f6a54cd 100644
--- a/firmware/2lib/include/2return_codes.h
+++ b/firmware/2lib/include/2return_codes.h
@@ -454,11 +454,11 @@ enum vb2_return_code {
 	 */
 	VB2_ERROR_MISC = VB2_ERROR_BASE + 0x080000,
 
-	/* Work buffer too small in vb2_init_context() */
-	VB2_ERROR_INITCTX_WORKBUF_SMALL,
+	/* Work buffer too small (see vb2api_init and vb2api_reinit) */
+	VB2_ERROR_WORKBUF_SMALL,
 
-	/* Work buffer unaligned in vb2_init_context() */
-	VB2_ERROR_INITCTX_WORKBUF_ALIGN,
+	/* Work buffer unaligned (see vb2api_init and vb2api_reinit) */
+	VB2_ERROR_WORKBUF_ALIGN,
 
 	/* Work buffer too small in GBB-related function */
 	VB2_ERROR_GBB_WORKBUF,
@@ -608,6 +608,9 @@ enum vb2_return_code {
 	/* Invalid parameter */
 	VB2_ERROR_INVALID_PARAMETER,
 
+	/* Problem with workbuf validity (see vb2api_init and vb2api_reinit) */
+	VB2_ERROR_WORKBUF_INVALID,
+
 	/**********************************************************************
 	 * API-level errors
 	 */
diff --git a/firmware/2lib/include/2struct.h b/firmware/2lib/include/2struct.h
index 8fa0254b..a7bd481b 100644
--- a/firmware/2lib/include/2struct.h
+++ b/firmware/2lib/include/2struct.h
@@ -16,17 +16,11 @@
 #ifndef VBOOT_REFERENCE_2STRUCT_H_
 #define VBOOT_REFERENCE_2STRUCT_H_
 
+#include "2api.h"
 #include "2constants.h"
 #include "2crypto.h"
 #include "2sysincludes.h"
 
-/* "V2CT" = vb2_context.magic */
-#define VB2_CONTEXT_MAGIC 0x54433256
-
-/* Current version of vb2_context struct */
-#define VB2_CONTEXT_VERSION_MAJOR 1
-#define VB2_CONTEXT_VERSION_MINOR 0
-
 /* Flags for vb2_shared_data.flags */
 enum vb2_shared_data_flags {
 	/* User has explicitly and physically requested recovery */
@@ -81,9 +75,11 @@ enum vb2_shared_data_status {
 #define VB2_SHARED_DATA_MAGIC 0x44533256
 
 /* Current version of vb2_shared_data struct */
-#define VB2_SHARED_DATA_VERSION_MAJOR 1
+#define VB2_SHARED_DATA_VERSION_MAJOR 2
 #define VB2_SHARED_DATA_VERSION_MINOR 0
 
+#define VB2_CONTEXT_MAX_SIZE 192
+
 /*
  * Data shared between vboot API calls.  Stored at the start of the work
  * buffer.
@@ -96,6 +92,21 @@ struct vb2_shared_data {
 	uint16_t struct_version_major;
 	uint16_t struct_version_minor;
 
+	/* Public fields are stored in the context object */
+	struct vb2_context ctx;
+
+	/* Padding for adding future vb2_context fields */
+	uint8_t padding[VB2_CONTEXT_MAX_SIZE - sizeof(struct vb2_context)];
+
+	/* Work buffer length in bytes. */
+	uint32_t workbuf_size;
+
+	/*
+	 * Amount of work buffer used so far.  Verified boot sub-calls use
+	 * this to know where the unused work area starts.
+	 */
+	uint32_t workbuf_used;
+
 	/* Flags; see enum vb2_shared_data_flags */
 	uint32_t flags;
 
diff --git a/firmware/include/vb2_api.h b/firmware/include/vb2_api.h
index 4774a1a8..d6400c45 100644
--- a/firmware/include/vb2_api.h
+++ b/firmware/include/vb2_api.h
@@ -33,6 +33,7 @@
  * TODO: Make cleaner APIs to this stuff.
  */
 #ifdef NEED_VB20_INTERNALS
+#include "../2lib/include/2misc.h"  /* for vb2_get_sd */
 #include "../2lib/include/2nvstorage.h"
 #include "../2lib/include/2nvstorage_fields.h"
 #include "../2lib/include/2struct.h"
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 68419aa3..52a5dcd0 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -230,16 +230,8 @@ static vb2_error_t vb2_kernel_setup(struct vb2_context *ctx,
 				    VbSharedDataHeader *shared,
 				    VbSelectAndLoadKernelParams *kparams)
 {
-	vb2_error_t rv;
 	uint32_t tpm_rv;
 
-	rv = vb2_init_context(ctx);
-	if (VB2_SUCCESS != rv) {
-		VB2_DEBUG("Can't init vb2_context\n");
-		vb2api_fail(ctx, VB2_RECOVERY_RW_SHARED_DATA, rv);
-		return VBERROR_INIT_SHARED_DATA;
-	}
-
 	/* Translate vboot1 flags back to vboot2 */
 	if (shared->recovery_reason)
 		ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
diff --git a/firmware/lib20/api_kernel.c b/firmware/lib20/api_kernel.c
index 99a4cf7f..3a4da4ce 100644
--- a/firmware/lib20/api_kernel.c
+++ b/firmware/lib20/api_kernel.c
@@ -110,8 +110,7 @@ vb2_error_t vb2api_kernel_phase1(struct vb2_context *ctx)
 	 *   - kernel key
 	 */
 	sd->kernel_key_size = key_size;
-	vb2_set_workbuf_used(ctx, sd->kernel_key_offset +
-			     sd->kernel_key_size);
+	vb2_set_workbuf_used(ctx, sd->kernel_key_offset + sd->kernel_key_size);
 
 	return VB2_SUCCESS;
 }
diff --git a/firmware/lib20/kernel.c b/firmware/lib20/kernel.c
index eff35bc4..3b3e2c54 100644
--- a/firmware/lib20/kernel.c
+++ b/firmware/lib20/kernel.c
@@ -212,7 +212,7 @@ vb2_error_t vb2_load_kernel_keyblock(struct vb2_context *ctx)
 	 * (which we might still need to verify the next kernel, if the
 	 * assoiciated kernel preamble and data don't verify).
 	 */
-	sd->data_key_offset = ctx->workbuf_used;
+	sd->data_key_offset = sd->workbuf_used;
 	key_data = vb2_member_of(sd, sd->data_key_offset);
 	packed_key = (struct vb2_packed_key *)key_data;
 	memmove(packed_key, &kb->data_key, sizeof(*packed_key));
@@ -233,8 +233,7 @@ vb2_error_t vb2_load_kernel_keyblock(struct vb2_context *ctx)
 	 *   - kernel key
 	 *   - packed kernel data key
 	 */
-	vb2_set_workbuf_used(ctx, sd->data_key_offset +
-			     sd->data_key_size);
+	vb2_set_workbuf_used(ctx, sd->data_key_offset + sd->data_key_size);
 
 	return VB2_SUCCESS;
 }
diff --git a/firmware/lib20/misc.c b/firmware/lib20/misc.c
index e6b4d56c..0fa80e1d 100644
--- a/firmware/lib20/misc.c
+++ b/firmware/lib20/misc.c
@@ -186,8 +186,7 @@ vb2_error_t vb2_load_fw_keyblock(struct vb2_context *ctx)
 	 *   - vb2_shared_data
 	 *   - packed firmware data key
 	 */
-	vb2_set_workbuf_used(ctx, sd->data_key_offset +
-			     sd->data_key_size);
+	vb2_set_workbuf_used(ctx, sd->data_key_offset + sd->data_key_size);
 
 	return VB2_SUCCESS;
 }
diff --git a/firmware/lib21/misc.c b/firmware/lib21/misc.c
index df50436d..ca1a012f 100644
--- a/firmware/lib21/misc.c
+++ b/firmware/lib21/misc.c
@@ -149,8 +149,7 @@ vb2_error_t vb21_load_fw_keyblock(struct vb2_context *ctx)
 	sd->data_key_size = packed_key->c.total_size;
 
 	/* Data key will persist in the workbuf after we return */
-	vb2_set_workbuf_used(ctx, sd->data_key_offset +
-			     sd->data_key_size);
+	vb2_set_workbuf_used(ctx, sd->data_key_offset + sd->data_key_size);
 
 	return VB2_SUCCESS;
 }
@@ -240,8 +239,7 @@ vb2_error_t vb21_load_fw_preamble(struct vb2_context *ctx)
 	sd->preamble_size = pre->c.total_size;
 
 	/* Preamble will persist in work buffer after we return */
-	vb2_set_workbuf_used(ctx, sd->preamble_offset +
-			     sd->preamble_size);
+	vb2_set_workbuf_used(ctx, sd->preamble_offset + sd->preamble_size);
 
 	return VB2_SUCCESS;
 }