sparse: Refactor chunk parsing function

The chunk parsing code was duplicating a lot of code among the various
chunk types, while all of them could be covered by generic and simple
functions.

Refactor the current code to reuse as much code as possible and hopefully
make the chunk parsing loop more readable and concise.

Signed-off-by: Maxime Ripard <maxime.ripard@free-electrons.com>
Reviewed-by: Tom Rini <trini@konsulko.com>

1 files changed, 239 insertions, 133 deletions

diff --git a/common/aboot.c b/common/aboot.c
index cd1b6a6..243b330 100644
--- a/common/aboot.c
+++ b/common/aboot.c
@@ -37,10 +37,36 @@
 #include <config.h>
 #include <common.h>
 #include <aboot.h>
+#include <errno.h>
 #include <malloc.h>
 #include <part.h>
 #include <sparse_format.h>
 
+typedef struct sparse_buffer {
+	void	*data;
+	u32	length;
+	u32	repeat;
+	u16	type;
+} sparse_buffer_t;
+
+static unsigned int sparse_get_chunk_data_size(sparse_header_t *sparse,
+					       chunk_header_t *chunk)
+{
+	return chunk->total_sz - sparse->chunk_hdr_sz;
+}
+
+static bool sparse_chunk_has_buffer(chunk_header_t *chunk)
+{
+	switch (chunk->chunk_type) {
+	case CHUNK_TYPE_RAW:
+	case CHUNK_TYPE_FILL:
+		return true;
+
+	default:
+		return false;
+	}
+}
+
 static sparse_header_t *sparse_parse_header(void **data)
 {
 	/* Read and skip over sparse image header */
@@ -61,21 +87,185 @@ static sparse_header_t *sparse_parse_header(void **data)
 	return sparse_header;
 }
 
+static int sparse_parse_fill_chunk(sparse_header_t *sparse,
+				   chunk_header_t *chunk)
+{
+	unsigned int chunk_data_sz = sparse_get_chunk_data_size(sparse, chunk);
+
+	if (chunk_data_sz != sizeof(uint32_t))
+		return -EINVAL;
+
+	return 0;
+}
+
+static int sparse_parse_raw_chunk(sparse_header_t *sparse,
+				  chunk_header_t *chunk)
+{
+	unsigned int chunk_data_sz = sparse_get_chunk_data_size(sparse, chunk);
+
+	/* Check if the data size is a multiple of the main block size */
+	if (chunk_data_sz % sparse->blk_sz)
+		return -EINVAL;
+
+	/* Check that the chunk size is consistent */
+	if ((chunk_data_sz / sparse->blk_sz) != chunk->chunk_sz)
+		return -EINVAL;
+
+	return 0;
+}
+
+static chunk_header_t *sparse_parse_chunk(sparse_header_t *sparse,
+					  void **image)
+{
+	chunk_header_t *chunk = (chunk_header_t *) *image;
+	int ret;
+
+	debug("=== Chunk Header ===\n");
+	debug("chunk_type: 0x%x\n", chunk->chunk_type);
+	debug("chunk_data_sz: 0x%x\n", chunk->chunk_sz);
+	debug("total_size: 0x%x\n", chunk->total_sz);
+
+	switch (chunk->chunk_type) {
+	case CHUNK_TYPE_RAW:
+		ret = sparse_parse_raw_chunk(sparse, chunk);
+		if (ret)
+			return NULL;
+		break;
+
+	case CHUNK_TYPE_FILL:
+		ret = sparse_parse_fill_chunk(sparse, chunk);
+		if (ret)
+			return NULL;
+		break;
+
+	case CHUNK_TYPE_DONT_CARE:
+	case CHUNK_TYPE_CRC32:
+		debug("Ignoring chunk\n");
+		break;
+
+	default:
+		printf("%s: Unknown chunk type: %x\n", __func__,
+		       chunk->chunk_type);
+		return NULL;
+	}
+
+	*image += sparse->chunk_hdr_sz;
+
+	return chunk;
+}
+
+static int sparse_get_fill_buffer(sparse_header_t *sparse,
+				  chunk_header_t *chunk,
+				  sparse_buffer_t *buffer,
+				  unsigned int blk_sz,
+				  void *data)
+{
+	int i;
+
+	buffer->type = CHUNK_TYPE_FILL;
+
+	/*
+	 * We create a buffer of one block, and ask it to be
+	 * repeated as many times as needed.
+	 */
+	buffer->length = blk_sz;
+	buffer->repeat = (chunk->chunk_sz * sparse->blk_sz) / blk_sz;
+
+	buffer->data = memalign(ARCH_DMA_MINALIGN,
+				ROUNDUP(blk_sz,
+					ARCH_DMA_MINALIGN));
+	if (!buffer->data)
+		return -ENOMEM;
+
+	for (i = 0; i < (buffer->length / sizeof(uint32_t)); i++)
+		((uint32_t *)buffer->data)[i] = *(uint32_t *)(data);
+
+	return 0;
+}
+
+static int sparse_get_raw_buffer(sparse_header_t *sparse,
+				 chunk_header_t *chunk,
+				 sparse_buffer_t *buffer,
+				 unsigned int blk_sz,
+				 void *data)
+{
+	unsigned int chunk_data_sz = sparse_get_chunk_data_size(sparse, chunk);
+
+	buffer->type = CHUNK_TYPE_RAW;
+	buffer->length = chunk_data_sz;
+	buffer->data = data;
+	buffer->repeat = 1;
+
+	return 0;
+}
+
+static sparse_buffer_t *sparse_get_data_buffer(sparse_header_t *sparse,
+					       chunk_header_t *chunk,
+					       unsigned int blk_sz,
+					       void **image)
+{
+	unsigned int chunk_data_sz = sparse_get_chunk_data_size(sparse, chunk);
+	sparse_buffer_t *buffer;
+	void *data = *image;
+	int ret;
+
+	*image += chunk_data_sz;
+
+	if (!sparse_chunk_has_buffer(chunk))
+		return NULL;
+
+	buffer = calloc(sizeof(sparse_buffer_t), 1);
+	if (!buffer)
+		return NULL;
+
+	switch (chunk->chunk_type) {
+	case CHUNK_TYPE_RAW:
+		ret = sparse_get_raw_buffer(sparse, chunk, buffer, blk_sz,
+					    data);
+		if (ret)
+			return NULL;
+		break;
+
+	case CHUNK_TYPE_FILL:
+		ret = sparse_get_fill_buffer(sparse, chunk, buffer, blk_sz,
+					     data);
+		if (ret)
+			return NULL;
+		break;
+
+	default:
+		return NULL;
+	}
+
+	debug("=== Buffer ===\n");
+	debug("length: 0x%x\n", buffer->length);
+	debug("repeat: 0x%x\n", buffer->repeat);
+	debug("type: 0x%x\n", buffer->type);
+	debug("data: 0x%p\n", buffer->data);
+
+	return buffer;
+}
+
+static void sparse_put_data_buffer(sparse_buffer_t *buffer)
+{
+	if (buffer->type == CHUNK_TYPE_FILL)
+		free(buffer->data);
+
+	free(buffer);
+}
+
 void write_sparse_image(block_dev_desc_t *dev_desc,
 		disk_partition_t *info, const char *part_name,
 		void *data, unsigned sz)
 {
-	lbaint_t blk;
+	lbaint_t start;
 	lbaint_t blkcnt;
-	lbaint_t blks;
-	uint32_t bytes_written = 0;
 	unsigned int chunk;
-	unsigned int chunk_data_sz;
-	uint32_t *fill_buf = NULL;
-	uint32_t fill_val;
 	sparse_header_t *sparse_header;
 	chunk_header_t *chunk_header;
+	sparse_buffer_t *buffer;
 	uint32_t total_blocks = 0;
+	uint32_t skipped = 0;
 	int i;
 
 	sparse_header = sparse_parse_header(&data);
@@ -96,148 +286,64 @@ void write_sparse_image(block_dev_desc_t *dev_desc,
 	puts("Flashing Sparse Image\n");
 
 	/* Start processing chunks */
-	blk = info->start;
-	for (chunk=0; chunk<sparse_header->total_chunks; chunk++)
-	{
-		/* Read and skip over chunk header */
-		chunk_header = (chunk_header_t *) data;
-		data += sizeof(chunk_header_t);
-
-		if (chunk_header->chunk_type != CHUNK_TYPE_RAW) {
-			debug("=== Chunk Header ===\n");
-			debug("chunk_type: 0x%x\n", chunk_header->chunk_type);
-			debug("chunk_data_sz: 0x%x\n", chunk_header->chunk_sz);
-			debug("total_size: 0x%x\n", chunk_header->total_sz);
+	start = info->start;
+	for (chunk = 0; chunk < sparse_header->total_chunks; chunk++) {
+		chunk_header = sparse_parse_chunk(sparse_header, &data);
+		if (!chunk_header) {
+			fastboot_fail("Unknown chunk type");
+			return;
 		}
 
-		if (sparse_header->chunk_hdr_sz > sizeof(chunk_header_t))
-		{
-			/*
-			 * Skip the remaining bytes in a header that is longer
-			 * than we expected.
-			 */
-			data += (sparse_header->chunk_hdr_sz -
-				 sizeof(chunk_header_t));
+		/*
+		 * If we have a DONT_CARE type, just skip the blocks
+		 * and go on parsing the rest of the chunks
+		 */
+		if (chunk_header->chunk_type == CHUNK_TYPE_DONT_CARE) {
+			skipped += chunk_header->chunk_sz;
+			continue;
 		}
 
-		chunk_data_sz = sparse_header->blk_sz * chunk_header->chunk_sz;
-		blkcnt = chunk_data_sz / info->blksz;
-		switch (chunk_header->chunk_type)
-		{
-			case CHUNK_TYPE_RAW:
-			if (chunk_header->total_sz !=
-			    (sparse_header->chunk_hdr_sz + chunk_data_sz))
-			{
-				fastboot_fail(
-					"Bogus chunk size for chunk type Raw");
-				return;
-			}
-
-			if (blk + blkcnt > info->start + info->size) {
-				printf(
-				    "%s: Request would exceed partition size!\n",
-				    __func__);
-				fastboot_fail(
-				    "Request would exceed partition size!");
-				return;
-			}
+		/* Retrieve the buffer we're going to write */
+		buffer = sparse_get_data_buffer(sparse_header, chunk_header,
+						info->blksz, &data);
+		if (!buffer)
+			continue;
 
-			blks = dev_desc->block_write(dev_desc->dev, blk, blkcnt,
-						     data);
-			if (blks != blkcnt) {
-				printf("%s: Write failed " LBAFU "\n",
-				       __func__, blks);
-				fastboot_fail("flash write failure");
-				return;
-			}
-			blk += blkcnt;
-			bytes_written += blkcnt * info->blksz;
-			total_blocks += chunk_header->chunk_sz;
-			data += chunk_data_sz;
-			break;
-
-			case CHUNK_TYPE_FILL:
-			if (chunk_header->total_sz !=
-			    (sparse_header->chunk_hdr_sz + sizeof(uint32_t)))
-			{
-				fastboot_fail(
-					"Bogus chunk size for chunk type FILL");
-				return;
-			}
+		blkcnt = (buffer->length / info->blksz) * buffer->repeat;
 
-			fill_buf = (uint32_t *)
-				   memalign(ARCH_DMA_MINALIGN,
-					    ROUNDUP(info->blksz,
-						    ARCH_DMA_MINALIGN));
-			if (!fill_buf)
-			{
-				fastboot_fail(
-					"Malloc failed for: CHUNK_TYPE_FILL");
-				return;
-			}
-
-			fill_val = *(uint32_t *)data;
-			data = (char *) data + sizeof(uint32_t);
+		if ((start + total_blocks + blkcnt) >
+		    (info->start + info->size)) {
+			printf("%s: Request would exceed partition size!\n",
+			       __func__);
+			fastboot_fail("Request would exceed partition size!");
+			return;
+		}
 
-			for (i = 0; i < (info->blksz / sizeof(fill_val)); i++)
-				fill_buf[i] = fill_val;
+		for (i = 0; i < buffer->repeat; i++) {
+			unsigned long buffer_blk_cnt;
+			unsigned long buffer_blks;
 
-			if (blk + blkcnt > info->start + info->size) {
-				printf(
-				    "%s: Request would exceed partition size!\n",
-				    __func__);
-				fastboot_fail(
-				    "Request would exceed partition size!");
-				return;
-			}
+			buffer_blk_cnt = buffer->length / info->blksz;
 
-			for (i = 0; i < blkcnt; i++) {
-				blks = dev_desc->block_write(dev_desc->dev,
-							     blk, 1, fill_buf);
-				if (blks != 1) {
-					printf(
-					    "%s: Write failed, block # " LBAFU "\n",
-					    __func__, blkcnt);
-					fastboot_fail("flash write failure");
-					free(fill_buf);
-					return;
-				}
-				blk++;
-			}
-			bytes_written += blkcnt * info->blksz;
-			total_blocks += chunk_data_sz / sparse_header->blk_sz;
-
-			free(fill_buf);
-			break;
-
-			case CHUNK_TYPE_DONT_CARE:
-			blk += blkcnt;
-			total_blocks += chunk_header->chunk_sz;
-			break;
-
-			case CHUNK_TYPE_CRC32:
-			if (chunk_header->total_sz !=
-			    sparse_header->chunk_hdr_sz)
-			{
-				fastboot_fail(
-					"Bogus chunk size for chunk type Dont Care");
+			buffer_blks = dev_desc->block_write(dev_desc->dev,
+							    start + total_blocks,
+							    buffer_blk_cnt, buffer->data);
+			if (buffer_blks != buffer_blk_cnt) {
+				printf("%s: Write %d failed " LBAFU "\n",
+				       __func__, i, buffer_blks);
+				fastboot_fail("flash write failure");
 				return;
 			}
-			total_blocks += chunk_header->chunk_sz;
-			data += chunk_data_sz;
-			break;
 
-			default:
-			printf("%s: Unknown chunk type: %x\n", __func__,
-			       chunk_header->chunk_type);
-			fastboot_fail("Unknown chunk type");
-			return;
+			total_blocks += buffer_blk_cnt;
 		}
+
+		sparse_put_data_buffer(buffer);
 	}
 
-	debug("Wrote %d blocks, expected to write %d blocks\n",
-	      total_blocks, sparse_header->total_blks);
-	printf("........ wrote %u bytes to '%s'\n", bytes_written, part_name);
+	debug("Wrote %d blocks, skipped %d, expected to write %d blocks\n",
+	      total_blocks, skipped, sparse_header->total_blks);
+	printf("........ wrote %d blocks to '%s'\n", total_blocks, part_name);
 
 	if (total_blocks != sparse_header->total_blks)
 		fastboot_fail("sparse image write failure");