From fa715193c083bbde4985c4ff1e49e288ec29763a Mon Sep 17 00:00:00 2001
From: Lokesh Vutla <lokeshvutla@ti.com>
Date: Tue, 24 May 2016 10:34:44 +0530
Subject: [PATCH] spl: Add an option to load a FIT containing U-Boot from UART

This provides a way to load a FIT containing U-Boot and a selection of device
tree files from UART.

Reviewed-by: Tom Rini <trini@konsulko.com>
Signed-off-by: Lokesh Vutla <lokeshvutla@ti.com>
---
 common/spl/spl_ymodem.c | 98 +++++++++++++++++++++++++++++++++++------
 1 file changed, 84 insertions(+), 14 deletions(-)

diff --git a/common/spl/spl_ymodem.c b/common/spl/spl_ymodem.c
index 4f26ea5d21..5402301c78 100644
--- a/common/spl/spl_ymodem.c
+++ b/common/spl/spl_ymodem.c
@@ -14,15 +14,60 @@
 #include <xyzModem.h>
 #include <asm/u-boot.h>
 #include <asm/utils.h>
+#include <libfdt.h>
 
 #define BUF_SIZE 1024
 
+/*
+ * Information required to load image using ymodem.
+ *
+ * @image_read: Now of bytes read from the image.
+ * @buf: pointer to the previous read block.
+ */
+struct ymodem_fit_info {
+	int image_read;
+	char *buf;
+};
+
 static int getcymodem(void) {
 	if (tstc())
 		return (getc());
 	return -1;
 }
 
+static ulong ymodem_read_fit(struct spl_load_info *load, ulong offset,
+			     ulong size, void *addr)
+{
+	int res, err;
+	struct ymodem_fit_info *info = load->priv;
+	char *buf = info->buf;
+
+	while (info->image_read < offset) {
+		res = xyzModem_stream_read(buf, BUF_SIZE, &err);
+		if (res <= 0)
+			return res;
+		info->image_read += res;
+	}
+
+	if (info->image_read > offset) {
+		res = info->image_read - offset;
+		memcpy(addr, &buf[BUF_SIZE - res], res);
+		addr = addr + res;
+	}
+
+	while (info->image_read < offset + size) {
+		res = xyzModem_stream_read(buf, BUF_SIZE, &err);
+		if (res <= 0)
+			return res;
+
+		memcpy(addr, buf, res);
+		info->image_read += res;
+		addr += res;
+	}
+
+	return size;
+}
+
 int spl_ymodem_load_image(void)
 {
 	int size = 0;
@@ -31,30 +76,55 @@ int spl_ymodem_load_image(void)
 	int ret;
 	connection_info_t info;
 	char buf[BUF_SIZE];
-	ulong store_addr = ~0;
 	ulong addr = 0;
 
 	info.mode = xyzModem_ymodem;
 	ret = xyzModem_stream_open(&info, &err);
+	if (ret) {
+		printf("spl: ymodem err - %s\n", xyzModem_error(err));
+		return ret;
+	}
+
+	res = xyzModem_stream_read(buf, BUF_SIZE, &err);
+	if (res <= 0)
+		goto end_stream;
+
+	if (IS_ENABLED(CONFIG_SPL_LOAD_FIT) &&
+	    image_get_magic((struct image_header *)buf) == FDT_MAGIC) {
+		struct spl_load_info load;
+		struct ymodem_fit_info info;
+
+		debug("Found FIT\n");
+		load.dev = NULL;
+		load.priv = (void *)&info;
+		load.filename = NULL;
+		load.bl_len = 1;
+		info.buf = buf;
+		info.image_read = BUF_SIZE;
+		load.read = ymodem_read_fit;
+		ret =  spl_load_simple_fit(&load, 0, (void *)buf);
+		size = info.image_read;
 
-	if (!ret) {
-		while ((res =
-			xyzModem_stream_read(buf, BUF_SIZE, &err)) > 0) {
-			if (addr == 0) {
-				ret = spl_parse_image_header((struct image_header *)buf);
-				if (ret)
-					return ret;
-			}
-			store_addr = addr + spl_image.load_addr;
+		while ((res = xyzModem_stream_read(buf, BUF_SIZE, &err)) > 0)
+			size += res;
+	} else {
+		spl_parse_image_header((struct image_header *)buf);
+		ret = spl_parse_image_header((struct image_header *)buf);
+		if (ret)
+			return ret;
+		addr = spl_image.load_addr;
+		memcpy((void *)addr, buf, res);
+		size += res;
+		addr += res;
+
+		while ((res = xyzModem_stream_read(buf, BUF_SIZE, &err)) > 0) {
+			memcpy((void *)addr, buf, res);
 			size += res;
 			addr += res;
-			memcpy((char *)(store_addr), buf, res);
 		}
-	} else {
-		printf("spl: ymodem err - %s\n", xyzModem_error(err));
-		return ret;
 	}
 
+end_stream:
 	xyzModem_stream_close(&err);
 	xyzModem_stream_terminate(false, &getcymodem);
 
-- 
2.39.5