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>

4 files changed, 30 insertions, 14 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