rockchip: mkimage: Add support for idb header V2

Rockchip BootRom supports new idb header v2 instead of legacy version.
Add support for it so that we can generate image for new SoCs.

Signed-off-by: Yi Liu <liuyi@rock-chips.com>
Signed-off-by: Kever Yang <kever.yang@rock-chips.com>
diff --git a/tools/rkcommon.c b/tools/rkcommon.c
index b1d867f..1b37ff3 100644
--- a/tools/rkcommon.c
+++ b/tools/rkcommon.c
@@ -10,18 +10,69 @@
 
 #include "imagetool.h"
 #include <image.h>
+#include <u-boot/sha256.h>
 #include <rc4.h>
 #include "mkimage.h"
 #include "rkcommon.h"
 
 enum {
 	RK_MAGIC		= 0x0ff0aa55,
+	RK_MAGIC_V2		= 0x534E4B52,
 };
 
 enum {
 	RK_HEADER_V1	= 1,
+	RK_HEADER_V2	= 2,
 };
 
+enum hash_type {
+	HASH_NONE	= 0,
+	HASH_SHA256	= 1,
+	HASH_SHA512	= 2,
+};
+
+/**
+ * struct image_entry
+ *
+ * @size_and_off:	[31:16]image size;[15:0]image offset
+ * @address:	default as 0xFFFFFFFF
+ * @flag:	no use
+ * @counter:	no use
+ * @hash:	hash of image
+ *
+ */
+struct image_entry {
+	uint32_t size_and_off;
+	uint32_t address;
+	uint32_t flag;
+	uint32_t counter;
+	uint8_t reserved[8];
+	uint8_t hash[64];
+};
+
+/**
+ * struct header0_info_v2 - v2 header block for rockchip BootRom
+ *
+ * This is stored at SD card block 64 (where each block is 512 bytes)
+ *
+ * @magic:	Magic (must be RK_MAGIC_V2)
+ * @size_and_nimage:	[31:16]number of images;[15:0]
+ *			offset to hash field of header(unit as 4Byte)
+ * @boot_flag:	[3:0]hash type(0:none,1:sha256,2:sha512)
+ * @signature:	hash or signature for header info
+ *
+ */
+struct header0_info_v2 {
+	uint32_t magic;
+	uint8_t reserved[4];
+	uint32_t size_and_nimage;
+	uint32_t boot_flag;
+	uint8_t reserved1[104];
+	struct image_entry images[4];
+	uint8_t reserved2[1064];
+	uint8_t hash[512];
+};
+
 /**
  * struct header0_info - header block for boot ROM
  *
@@ -202,7 +253,6 @@
 	return info->spl_hdr;
 }
 
-
 int rkcommon_get_spl_size(struct image_tool_params *params)
 {
 	struct spl_info *info = rkcommon_get_spl_info(params->imagename);
@@ -223,6 +273,22 @@
 	return info->spl_rc4;
 }
 
+bool rkcommon_is_header_v2(struct image_tool_params *params)
+{
+	struct spl_info *info = rkcommon_get_spl_info(params->imagename);
+
+	return (info->header_ver == RK_HEADER_V2);
+}
+
+static void do_sha256_hash(uint8_t *buf, uint32_t size, uint8_t *out)
+{
+	sha256_context ctx;
+
+	sha256_starts(&ctx);
+	sha256_update(&ctx, buf, size);
+	sha256_finish(&ctx, out);
+}
+
 static void rkcommon_set_header0(void *buf, struct image_tool_params *params)
 {
 	struct header0_info *hdr = buf;
@@ -251,29 +317,67 @@
 	rc4_encode(buf, RK_BLK_SIZE, rc4_key);
 }
 
+static void rkcommon_set_header0_v2(void *buf, struct image_tool_params *params)
+{
+	struct header0_info_v2 *hdr = buf;
+	uint32_t sector_offset, image_sector_count;
+	uint32_t image_size_array[2];
+	uint8_t *image_ptr = NULL;
+	int i;
+
+	printf("Image Type:   Rockchip %s boot image\n",
+		rkcommon_get_spl_hdr(params));
+	memset(buf, '\0', RK_INIT_OFFSET * RK_BLK_SIZE);
+	hdr->magic   = cpu_to_le32(RK_MAGIC_V2);
+	hdr->size_and_nimage = cpu_to_le32((2 << 16) + 384);
+	hdr->boot_flag = cpu_to_le32(HASH_SHA256);
+	sector_offset = 4;
+	image_size_array[0] = spl_params.init_size;
+	image_size_array[1] = spl_params.boot_size;
+
+	for (i = 0; i < 2; i++) {
+		image_sector_count = image_size_array[i] / RK_BLK_SIZE;
+		hdr->images[i].size_and_off = cpu_to_le32((image_sector_count
+							<< 16) + sector_offset);
+		hdr->images[i].address = 0xFFFFFFFF;
+		hdr->images[i].counter = cpu_to_le32(i + 1);
+		image_ptr = buf + sector_offset * RK_BLK_SIZE;
+		do_sha256_hash(image_ptr, image_size_array[i],
+			       hdr->images[i].hash);
+		sector_offset = sector_offset + image_sector_count;
+	}
+
+	do_sha256_hash(buf, (void *)hdr->hash - buf, hdr->hash);
+}
+
 void rkcommon_set_header(void *buf,  struct stat *sbuf,  int ifd,
 			 struct image_tool_params *params)
 {
 	struct header1_info *hdr = buf + RK_SPL_HDR_START;
 
-	rkcommon_set_header0(buf, params);
+	if (rkcommon_is_header_v2(params)) {
+		rkcommon_set_header0_v2(buf, params);
+	} else {
+		rkcommon_set_header0(buf, params);
 
-	/* Set up the SPL name (i.e. copy spl_hdr over) */
-	memcpy(&hdr->magic, rkcommon_get_spl_hdr(params), RK_SPL_HDR_SIZE);
+		/* Set up the SPL name (i.e. copy spl_hdr over) */
+		if (memcmp(&hdr->magic, "RSAK", 4))
+			memcpy(&hdr->magic, rkcommon_get_spl_hdr(params), RK_SPL_HDR_SIZE);
 
-	if (rkcommon_need_rc4_spl(params))
-		rkcommon_rc4_encode_spl(buf, RK_SPL_HDR_START,
-					spl_params.init_size);
-
-	if (spl_params.boot_file) {
 		if (rkcommon_need_rc4_spl(params))
-			rkcommon_rc4_encode_spl(buf + RK_SPL_HDR_START,
-						spl_params.init_size,
-						spl_params.boot_size);
+			rkcommon_rc4_encode_spl(buf, RK_SPL_HDR_START,
+						spl_params.init_size);
+
+		if (spl_params.boot_file) {
+			if (rkcommon_need_rc4_spl(params))
+				rkcommon_rc4_encode_spl(buf + RK_SPL_HDR_START,
+							spl_params.init_size,
+							spl_params.boot_size);
+		}
 	}
 }
 
-static inline unsigned rkcommon_offset_to_spi(unsigned offset)
+static inline unsigned int rkcommon_offset_to_spi(unsigned int offset)
 {
 	/*
 	 * While SD/MMC images use a flat addressing, SPI images are padded
@@ -285,7 +389,7 @@
 static int rkcommon_parse_header(const void *buf, struct header0_info *header0,
 				 struct spl_info **spl_info)
 {
-	unsigned hdr1_offset;
+	unsigned int hdr1_offset;
 	struct header1_info *hdr1_sdmmc, *hdr1_spi;
 	int i;
 
@@ -328,6 +432,16 @@
 	return -1;
 }
 
+static int rkcommon_parse_header_v2(const void *buf, struct header0_info_v2 *header)
+{
+	memcpy((void *)header, buf, sizeof(struct header0_info_v2));
+
+	if (le32_to_cpu(header->magic) != RK_MAGIC_V2)
+		return -EPROTO;
+
+	return 0;
+}
+
 int rkcommon_verify_header(unsigned char *buf, int size,
 			   struct image_tool_params *params)
 {
@@ -362,30 +476,46 @@
 void rkcommon_print_header(const void *buf)
 {
 	struct header0_info header0;
+	struct header0_info_v2 header0_v2;
 	struct spl_info *spl_info;
 	uint8_t image_type;
 	int ret, boot_size, init_size;
 
-	ret = rkcommon_parse_header(buf, &header0, &spl_info);
+	if ((*(uint32_t *)buf) == RK_MAGIC_V2) {
+		ret = rkcommon_parse_header_v2(buf, &header0_v2);
 
-	/* If this is the (unimplemented) RC4 case, then fail silently */
-	if (ret == -ENOSYS)
-		return;
+		if (ret < 0) {
+			fprintf(stderr, "Error: image verification failed\n");
+			return;
+		}
 
-	if (ret < 0) {
-		fprintf(stderr, "Error: image verification failed\n");
-		return;
-	}
+		init_size = header0_v2.images[0].size_and_off >> 16;
+		init_size = init_size * RK_BLK_SIZE;
+		boot_size = header0_v2.images[1].size_and_off >> 16;
+		boot_size = boot_size * RK_BLK_SIZE;
+	} else {
+		ret = rkcommon_parse_header(buf, &header0, &spl_info);
+
+		/* If this is the (unimplemented) RC4 case, then fail silently */
+		if (ret == -ENOSYS)
+			return;
+
+		if (ret < 0) {
+			fprintf(stderr, "Error: image verification failed\n");
+			return;
+		}
+
+		image_type = ret;
+		init_size = header0.init_size * RK_BLK_SIZE;
+		boot_size = header0.init_boot_size * RK_BLK_SIZE - init_size;
 
-	image_type = ret;
+		printf("Image Type:   Rockchip %s (%s) boot image\n",
+		       spl_info->spl_hdr,
+		       (image_type == IH_TYPE_RKSD) ? "SD/MMC" : "SPI");
+	}
 
-	printf("Image Type:   Rockchip %s (%s) boot image\n",
-	       spl_info->spl_hdr,
-	       (image_type == IH_TYPE_RKSD) ? "SD/MMC" : "SPI");
-	init_size = le16_to_cpu(header0.init_size) * RK_BLK_SIZE;
 	printf("Init Data Size: %d bytes\n", init_size);
 
-	boot_size = le16_to_cpu(header0.init_boot_size) * RK_BLK_SIZE - init_size;
 	if (boot_size != RK_MAX_BOOT_SIZE)
 		printf("Boot Data Size: %d bytes\n", boot_size);
 }