[U-Boot] [PATCH v2] cmd: fsfitxtract: Extract a part of a FIT multi-image stored on a filesystem

[Alexey Ignatov <lexszero@gmail.com>]

Sometimes we interested only in one single small subimage from big multipart
FIT. This command tries to avoid reading out the whole FIT because of memory
or time considerations.
Since we don't have mmap() in U-Boot, this is done by reading the file in
small chunks and trying to parse FIT, skipping reading the payload data
of all images before the requested one, while preserving FIT structure.

Changelog:

v2:
	* Skip unneeded subimages data instead of reading it out
	* Read whole data in one operation
	* Lot of refactoring
	* Faster, more versatile and robust
	* Improved code readability

v1:
	* First submission to maillist

Signed-off-by: Alexey Ignatov <lexszero@gmail.com>
---
 cmd/Kconfig |   7 ++
 cmd/ximg.c  | 356 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 363 insertions(+)

diff --git a/cmd/Kconfig b/cmd/Kconfig
index d9f7151bac..ff8ade96b4 100644
--- a/cmd/Kconfig
+++ b/cmd/Kconfig
@@ -265,6 +265,13 @@ config CMD_XIMG
 	help
 	  Extract a part of a multi-image.
 
+config CMD_FSFITXIMG
+	bool "fsfitxtract"
+	default y
+	depends on FIT
+	help
+	  Extract a part of a FIT multi-image stored on a filesystem
+
 config CMD_POWEROFF
 	bool
 
diff --git a/cmd/ximg.c b/cmd/ximg.c
index 73a571b52b..62e2cb2ac1 100644
--- a/cmd/ximg.c
+++ b/cmd/ximg.c
@@ -22,6 +22,11 @@
 #endif
 #include <asm/byteorder.h>
 #include <asm/io.h>
+#ifdef CONFIG_CMD_FSFITXIMG
+#include <fs.h>
+#include <div64.h>
+#include <libfdt.h>
+#endif
 
 #ifndef CONFIG_SYS_XIMG_LEN
 /* use 8MByte as default max gunzip size */
@@ -283,3 +288,354 @@ U_BOOT_CMD(
 	imxtract, 4, 1, do_imgextract,
 	"extract a part of a multi-image", imgextract_help_text
 );
+
+#if defined(CONFIG_FIT) && defined(CONFIG_CMD_FSFITXIMG)
+
+/* TODO: maybe make this configurable */
+#define CHUNK_SIZE		(8*1024)
+#define IMAGE_TAIL_SIZE	CHUNK_SIZE
+
+struct fs_fdt_context {
+	const char *dev;
+	const char *part;
+	const char *name;
+	
+	ulong dest;
+	ulong len;
+	ulong src_off;
+	void *fdt;
+};
+
+static int fs_load(struct fs_fdt_context *ctx, loff_t len)
+{
+	int ret = 0;
+	loff_t len_read;
+
+	debug("Reading %lld bytes from offset 0x%lx to 0x%lx\n",
+			len, ctx->src_off, ctx->dest + ctx->len);
+
+	if (fs_set_blk_dev(ctx->dev, ctx->part, FS_TYPE_ANY))
+		return -1;
+
+	ret = fs_read(ctx->name, ctx->dest + ctx->len, ctx->src_off, len, &len_read);
+
+	if (ret < 0) {
+		printf("Read failed: %d", ret);
+		return ret;
+	}
+
+	return len_read;
+}
+
+static int fs_load_next_chunk(struct fs_fdt_context *ctx)
+{
+	int ret;
+	ret = fs_load(ctx, CHUNK_SIZE);
+	if (ret < 0)
+		return ret;
+
+	ctx->len += ret;
+	ctx->src_off += ret;
+
+	return 0;
+}
+
+static int fdt_err_retry(int err)
+{
+	return
+		(err == -FDT_ERR_TRUNCATED) ||
+		(err == -FDT_ERR_BADSTRUCTURE);
+}
+
+static int fs_fdt_wrap(struct fs_fdt_context *ctx,
+		int (*fdt_func)(const void *fdt, int offset), int offset)
+{
+	int ret = 0;
+	while (1) {
+		ret = fdt_func(ctx->fdt, offset);
+	
+		if (ret >= 0)
+			break;
+
+		if (!fdt_err_retry(ret)) {
+			if (ret != -FDT_ERR_NOTFOUND)
+				printf("Error while parsing FIT: %s\n", fdt_strerror(ret));
+			break;
+		}
+	
+		ret = fs_load_next_chunk(ctx);
+		if (ret < 0)
+			break;
+	}
+
+	return ret;
+}
+
+static inline int fs_fit_image_skip_data(struct fs_fdt_context *ctx, int noffset)
+{
+	int plen = 0;
+	int poffset = 0;
+	void *ptr;
+	const void *nptr;
+	int tail;
+
+	for (poffset = fs_fdt_wrap(ctx, fdt_first_property_offset, noffset);
+			poffset >= 0;
+			poffset = fs_fdt_wrap(ctx, fdt_next_property_offset, poffset)) {
+		const struct fdt_property *prop =
+			fdt_get_property_by_offset(ctx->fdt, poffset, &plen);
+		const char *name;
+
+		if (!prop) {
+			printf("Can't get property: %s\n", fdt_strerror(plen));
+			return 1;
+		}
+
+		name = fdt_string(ctx->fdt, fdt32_to_cpu(prop->nameoff));
+		if (!name) {
+			printf("Can't get property name\n");
+			return 1;
+		}
+
+		if (strcmp(name, "data") == 0)
+			break;
+	};
+
+	if (poffset < 0)
+		return poffset;
+
+	debug("data prop offset %d, len %d\n", poffset, plen);
+
+	/* Skip the whole data property which can be quite big using
+	 * some smart pointer arithmetic magic.
+	 */
+	printf("   Skip %d bytes of data\n", plen);
+	plen = ALIGN(plen + 3*4, 4);
+	ptr = fdt_offset_ptr_w(ctx->fdt, poffset, 0);
+	nptr = fdt_offset_ptr(ctx->fdt, poffset + plen, 0);
+	tail = ctx->fdt + ctx->len - nptr;
+	if (tail > 0) {
+		memmove(ptr, nptr, tail);
+		ctx->len -= plen;
+		debug("move tail len %d from %p to %p; ctx->len = 0x%lx\n",
+				tail, ptr, nptr, ctx->len);
+	}
+	else {
+		tail = -tail;
+		ctx->len = ptr - ctx->fdt;
+		ctx->src_off += tail;
+		debug("skip partial data, ctx->len = 0x%lx, ctx->src_off = 0x%lx\n",
+				ctx->len, ctx->src_off);
+
+		if (fs_load_next_chunk(ctx) != 0)
+			return 1;
+	}
+
+	return 0;
+}
+
+static int
+do_fsfitextract(cmd_tbl_t * cmdtp, int flag, int argc, char * const argv[])
+{
+	int		ret;
+	int		verify;
+	struct	fs_fdt_context ctx;
+	const char	*uname;
+	const void	*data;
+	size_t	len;
+	int		noffset = 0, tail;
+	uint8_t		comp;
+	unsigned long time;
+
+	if (argc != 6)
+		return CMD_RET_USAGE;
+
+	verify = getenv_yesno("verify");
+
+	ctx.dev = argv[1];
+	ctx.part = argv[2];
+	ctx.name = argv[3];
+	uname = argv[4];
+	ctx.dest = simple_strtoul(argv[5], NULL, 16);
+
+	time = get_timer(0);
+
+	/* Load FIT from the filesystem incrementally to avoid reading data we
+	 * won't need. If requested subimage contained in the beginning of file,
+	 * this can make a huge speedup.
+	 *
+	 * First, read just the header
+	 */
+	ctx.len = 0;
+	ctx.src_off = 0;
+	printf("## Loading FIT header to 0x%08lx ...\n", ctx.dest);
+	ret = fs_load(&ctx, sizeof(struct fdt_header));
+	if (ret < 0)
+		return 1;
+
+	ctx.fdt = map_sysmem(ctx.dest, sizeof(struct fdt_header));
+	ret = fdt_check_header(ctx.fdt);
+	if (ret) {
+		printf("Bad FIT header: %s\n", fdt_strerror(ret));
+		return 1;
+	}
+	debug("FIT header at %p\n", ctx.fdt);
+
+	ctx.len = sizeof(struct fdt_header);
+
+	/* Read "strings" block to resolve all the node names, put it after the
+	 * header and adjust offset appropriately
+	 */
+	printf("## Loading FIT strings to 0x%08lx ...\n", ctx.dest + ctx.len);
+	ctx.src_off = fdt_off_dt_strings(ctx.fdt);
+	ret = fs_load(&ctx, fdt_size_dt_strings(ctx.fdt));
+	if (ret < 0)
+		return 1;
+
+	fdt_set_off_dt_strings(ctx.fdt, ctx.len);
+	ctx.len += fdt_size_dt_strings(ctx.fdt);
+
+	ctx.len = ALIGN(ctx.len, ARCH_DMA_MINALIGN);
+
+	/* Save offset of "struct" block in file and set the real in-memory offset */
+	ctx.src_off = fdt_off_dt_struct(ctx.fdt);
+	fdt_set_off_dt_struct(ctx.fdt, ctx.len);
+
+	/* Now, start to load struct in CHUNK_SIZE chunks, and parse it@the same
+	 * time. When libfdt notices that the data is truncated or broken, read next
+	 * chunk from the filesystem.
+	 */
+	printf("## Loading FIT struct to 0x%08lx ...\n", ctx.dest + ctx.len);
+	while (1) {
+		if (fs_load_next_chunk(&ctx))
+			return 1;
+
+		ret = fdt_path_offset(ctx.fdt, FIT_IMAGES_PATH);
+		if (ret >= 0) {
+			noffset = ret;
+			break;
+		}
+
+		if (!fdt_err_retry(ret)) {
+			printf("Can't find images parent node: %s",
+					fdt_strerror(ret));
+			return 1;
+		}
+	};
+
+	/* Iterate through all images */
+	for (noffset = fs_fdt_wrap(&ctx, fdt_first_subnode, noffset);
+			noffset >= 0;
+			noffset = fs_fdt_wrap(&ctx, fdt_next_subnode, noffset)) {
+		const char *name;
+
+		name = fdt_get_name(ctx.fdt, noffset, NULL);
+		debug("node '%s' at offset %d\n", name, noffset);
+
+		if (strcmp(name, uname) == 0)
+			break;
+
+		if (fs_fit_image_skip_data(&ctx, noffset) != 0)
+			return 1;
+	}
+
+	if (noffset < 0) {
+		printf("Can't find subimage '%s'\n", uname);
+		return 1;
+	}
+
+	debug("found image node at %08x\n", noffset);
+
+	while (1) {
+		/* get subimage data address and length */
+		ret = fit_image_get_data(ctx.fdt, noffset,
+				&data, &len);
+		if (!ret)
+			break;
+
+		if (!fdt_err_retry(len)) {
+			puts("Could not find subimage data\n");
+			return 1;
+		}
+
+		if (fs_load_next_chunk(&ctx) != 0)
+			return 1;
+	};
+
+	debug("data at %p, len %lu\n", data, len);
+
+	/* Now we know the length of data, so read out the remaining part in
+	 * a single operation. Also read some more than exact data len to
+	 * get other properties such as hashes.
+	 *
+	 * The memory layout@this moment (in terms of offsets from header,
+	 * and declared variables):
+	 *
+	 * offset       pointers
+	 * ---------------------
+	 * 0            ctx.fdt
+	 * ...
+	 *              data
+	 * ...
+	 * ctx.len
+	 * ...
+	 * <data_end>   (data + len)
+	 *
+	 * So, size of leftover is calculated like this:
+	 * data_end = data - ctx.fdt + len
+	 * size_remaining = data_end - ctx.len
+	 */
+	tail = data - ctx.fdt + len - ctx.len + IMAGE_TAIL_SIZE;
+	if (tail > 0) {
+		ret = fs_load(&ctx, tail);
+		if (ret < 0)
+			return 1;
+		ctx.len += ret;
+	}
+
+	time = get_timer(time);
+
+	printf("%lu bytes read in %lu ms", ctx.len, time);
+	if (time > 0) {
+		puts(" (");
+		print_size(lldiv(ctx.len, time) * 1000, "/s");
+		puts(")");
+	}
+	puts("\n");
+
+	fit_image_print(ctx.fdt, noffset, "   ");
+
+	if (fit_image_get_comp(ctx.fdt, noffset, &comp)) {
+		puts("Could not find FIT subimage compression type\n");
+		return 1;
+	}
+
+	/* verify integrity */
+	if (verify) {
+		puts("## Verifying ... ");
+		if (!fit_image_verify(ctx.fdt, noffset)) {
+			puts("Bad data hash\n");
+			return 1;
+		}
+		puts("\n");
+	}
+
+	if (!decompress_data(ctx.dest, (ulong)data, (ulong)len, comp, 0))
+		return 1;
+
+	flush_cache(ctx.dest, len);
+
+	setenv_hex("filesize", len);
+
+	return 0;
+}
+
+U_BOOT_CMD(
+	fsfitxtract, 7, 1, do_fsfitextract,
+	"extract a part of a FIT stored on a filesystem",
+	"<interface> <dev[:part]> <filename> <uname> <addr>\n"
+	"    - Extract subimage <uname> from FIT file <filename> from\n"
+	"       partition <part> on device type <interface> instance <dev>\n"
+	"       to address <addr> in memory."
+);
+#endif
-- 
2.12.2