4 files changed, 29 insertions, 25 deletions

diff --git a/firmware/2lib/2api.c b/firmware/2lib/2api.c
index f10a1e69..008282b6 100644
--- a/firmware/2lib/2api.c
+++ b/firmware/2lib/2api.c
@@ -205,7 +205,7 @@ int vb2api_get_pcr_digest(struct vb2_context *ctx,
 		digest_size = VB2_SHA1_DIGEST_SIZE;
 		break;
 	case HWID_DIGEST_PCR:
-		digest = vb2_get_sd(ctx)->gbb_hwid_digest;
+		digest = vb2_get_gbb(ctx)->hwid_digest;
 		digest_size = VB2_GBB_HWID_DIGEST_SIZE;
 		break;
 	default:
diff --git a/firmware/2lib/2misc.c b/firmware/2lib/2misc.c
index 8b6a19f8..03e0353a 100644
--- a/firmware/2lib/2misc.c
+++ b/firmware/2lib/2misc.c
@@ -14,7 +14,8 @@
 #include "2sha.h"
 #include "2rsa.h"
 
-int vb2_validate_gbb_signature(uint8_t *sig) {
+int vb2_validate_gbb_signature(uint8_t *sig)
+{
 	const static uint8_t sig_xor[VB2_GBB_SIGNATURE_SIZE] =
 			VB2_GBB_XOR_SIGNATURE;
 	int i;
@@ -25,6 +26,13 @@ int vb2_validate_gbb_signature(uint8_t *sig) {
 	return VB2_SUCCESS;
 }
 
+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);
+}
+
 void vb2_workbuf_from_ctx(struct vb2_context *ctx, struct vb2_workbuf *wb)
 {
 	vb2_workbuf_init(wb, ctx->workbuf + ctx->workbuf_used,
@@ -218,11 +226,9 @@ int vb2_fw_parse_gbb(struct vb2_context *ctx)
 	if (rv)
 		return rv;
 
-	/* Extract the only things we care about at firmware time */
-	sd->gbb_flags = gbb->flags;
-	sd->gbb_rootkey_offset = gbb->rootkey_offset;
-	sd->gbb_rootkey_size = gbb->rootkey_size;
-	memcpy(sd->gbb_hwid_digest, gbb->hwid_digest, VB2_GBB_HWID_DIGEST_SIZE);
+	/* Keep on the work buffer permanently */
+	sd->gbb_offset = vb2_offset_of(sd, gbb);
+	ctx->workbuf_used = vb2_offset_of(ctx->workbuf, wb.buf);
 
 	return VB2_SUCCESS;
 }
@@ -230,6 +236,7 @@ int vb2_fw_parse_gbb(struct vb2_context *ctx)
 int vb2_check_dev_switch(struct vb2_context *ctx)
 {
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
+	struct vb2_gbb_header *gbb = vb2_get_gbb(ctx);
 	uint32_t flags = 0;
 	uint32_t old_flags;
 	int is_dev = 0;
@@ -275,7 +282,7 @@ int vb2_check_dev_switch(struct vb2_context *ctx)
 		is_dev = 1;
 
 	/* Check if GBB is forcing dev mode */
-	if (sd->gbb_flags & VB2_GBB_FLAG_FORCE_DEV_SWITCH_ON)
+	if (gbb->flags & VB2_GBB_FLAG_FORCE_DEV_SWITCH_ON)
 		is_dev = 1;
 
 	/* Handle whichever mode we end up in */
diff --git a/firmware/2lib/include/2misc.h b/firmware/2lib/include/2misc.h
index 853a1b40..3e1f1000 100644
--- a/firmware/2lib/include/2misc.h
+++ b/firmware/2lib/include/2misc.h
@@ -9,6 +9,7 @@
 #define VBOOT_REFERENCE_VBOOT_2MISC_H_
 
 #include "2api.h"
+#include "2struct.h"
 
 struct vb2_gbb_header;
 struct vb2_workbuf;
@@ -19,11 +20,20 @@ struct vb2_workbuf;
  * @param ctx		Vboot context
  * @return The shared data pointer.
  */
-static __inline struct vb2_shared_data *vb2_get_sd(struct vb2_context *ctx) {
+static __inline struct vb2_shared_data *vb2_get_sd(struct vb2_context *ctx)
+{
 	return (struct vb2_shared_data *)ctx->workbuf;
 }
 
 /**
+ * Get the GBB header pointer from a vboot context's shared data
+ *
+ * @param ctx		Vboot context
+ * @return The GBB header pointer.
+ */
+struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *ctx);
+
+/**
  * Validate gbb signature (the magic number)
  *
  * @param sig		Pointer to the signature bytes to validate
diff --git a/firmware/2lib/include/2struct.h b/firmware/2lib/include/2struct.h
index a1f11a18..0b94c20a 100644
--- a/firmware/2lib/include/2struct.h
+++ b/firmware/2lib/include/2struct.h
@@ -95,9 +95,6 @@ struct vb2_shared_data {
 	/* Flags; see enum vb2_shared_data_flags */
 	uint32_t flags;
 
-	/* Flags from GBB header */
-	uint32_t gbb_flags;
-
 	/*
 	 * Reason we are in recovery mode this boot (enum vb2_nv_recovery), or
 	 * 0 if we aren't.
@@ -130,6 +127,9 @@ struct vb2_shared_data {
 	 */
 	uint32_t status;
 
+	/* Offset from start of this struct to GBB header */
+	uint32_t gbb_offset;
+
 	/**********************************************************************
 	 * Data from kernel verification stage.
 	 *
@@ -155,13 +155,6 @@ struct vb2_shared_data {
 	 * we can put them.
 	 */
 
-	/* Root key offset and size from GBB header */
-	uint32_t gbb_rootkey_offset;
-	uint32_t gbb_rootkey_size;
-
-	/* HWID digest from GBB header */
-	uint8_t gbb_hwid_digest[VB2_GBB_HWID_DIGEST_SIZE];
-
 	/* Offset of preamble from start of vblock */
 	uint32_t vblock_preamble_offset;
 
@@ -224,12 +217,6 @@ struct vb2_shared_data {
 	 */
 	uint32_t workbuf_kernel_key_offset;
 	uint32_t workbuf_kernel_key_size;
-
-	/* GBB data and size */
-	struct vb2_gbb_header *gbb;
-	uint32_t gbb_size;
-
-
 } __attribute__((packed));
 
 /****************************************************************************/