tools: kwbimage: Add support for SATA images with non-512 byte block size

SATA kwbimage contains offsets in block size unit, not in bytes.

Until now kwbimage expected that SATA disk always have block size of 512
bytes. But there are 4K Native SATA disks with block size of 4096 bytes.

New SATA_BLKSZ command allows to specify different block size than 512
bytes and therefore allows to generate kwbimage for disks with different
block sizes.

This change add support for generating SATA images with different block
size. Also it add support for verifying and dumping such images.

Because block size itself is not stored in SATA kwbimage, image
verification is done by checking every possible block size (it is any
power of two value between 512 and 32 kB).

Signed-off-by: Pali Rohár <pali@kernel.org>
diff --git a/tools/kwbimage.c b/tools/kwbimage.c
index 360fedd..4dce495 100644
--- a/tools/kwbimage.c
+++ b/tools/kwbimage.c
@@ -116,6 +116,7 @@
 	IMAGE_CFG_NAND_BADBLK_LOCATION,
 	IMAGE_CFG_NAND_ECC_MODE,
 	IMAGE_CFG_NAND_PAGESZ,
+	IMAGE_CFG_SATA_BLKSZ,
 	IMAGE_CFG_CPU,
 	IMAGE_CFG_BINARY,
 	IMAGE_CFG_DATA,
@@ -147,6 +148,7 @@
 	[IMAGE_CFG_NAND_BADBLK_LOCATION] = "NAND_BADBLK_LOCATION",
 	[IMAGE_CFG_NAND_ECC_MODE] = "NAND_ECC_MODE",
 	[IMAGE_CFG_NAND_PAGESZ] = "NAND_PAGE_SIZE",
+	[IMAGE_CFG_SATA_BLKSZ] = "SATA_BLKSZ",
 	[IMAGE_CFG_CPU] = "CPU",
 	[IMAGE_CFG_BINARY] = "BINARY",
 	[IMAGE_CFG_DATA] = "DATA",
@@ -185,6 +187,7 @@
 		unsigned int nandbadblklocation;
 		unsigned int nandeccmode;
 		unsigned int nandpagesz;
+		unsigned int satablksz;
 		struct ext_hdr_v0_reg regdata;
 		unsigned int regdata_delay;
 		unsigned int baudrate;
@@ -992,13 +995,21 @@
 	return 1;
 }
 
+static unsigned int image_get_satablksz(void)
+{
+	struct image_cfg_element *e;
+	e = image_find_option(IMAGE_CFG_SATA_BLKSZ);
+	return e ? e->satablksz : 512;
+}
+
 static size_t image_headersz_align(size_t headersz, uint8_t blockid)
 {
 	/*
 	 * Header needs to be 4-byte aligned, which is already ensured by code
 	 * above. Moreover UART images must have header aligned to 128 bytes
 	 * (xmodem block size), NAND images to 256 bytes (ECC calculation),
-	 * and SATA and SDIO images to 512 bytes (storage block size).
+	 * SDIO images to 512 bytes (SDHC/SDXC fixed block size) and SATA
+	 * images to specified storage block size (default 512 bytes).
 	 * Note that SPI images do not have to have header size aligned
 	 * to 256 bytes because it is possible to read from SPI storage from
 	 * any offset (read offset does not have to be aligned to block size).
@@ -1007,8 +1018,10 @@
 		return ALIGN(headersz, 128);
 	else if (blockid == IBR_HDR_NAND_ID)
 		return ALIGN(headersz, 256);
-	else if (blockid == IBR_HDR_SATA_ID || blockid == IBR_HDR_SDIO_ID)
+	else if (blockid == IBR_HDR_SDIO_ID)
 		return ALIGN(headersz, 512);
+	else if (blockid == IBR_HDR_SATA_ID)
+		return ALIGN(headersz, image_get_satablksz());
 	else
 		return headersz;
 }
@@ -1076,12 +1089,11 @@
 	if (e)
 		main_hdr->nandbadblklocation = e->nandbadblklocation;
 
-	/*
-	 * For SATA srcaddr is specified in number of sectors.
-	 * This expects the sector size to be 512 bytes.
-	 */
-	if (main_hdr->blockid == IBR_HDR_SATA_ID)
-		main_hdr->srcaddr = cpu_to_le32(le32_to_cpu(main_hdr->srcaddr) / 512);
+	/* For SATA srcaddr is specified in number of sectors. */
+	if (main_hdr->blockid == IBR_HDR_SATA_ID) {
+		params->bl_len = image_get_satablksz();
+		main_hdr->srcaddr = cpu_to_le32(le32_to_cpu(main_hdr->srcaddr) / params->bl_len);
+	}
 
 	/* For PCIe srcaddr is not used and must be set to 0xFFFFFFFF. */
 	if (main_hdr->blockid == IBR_HDR_PEX_ID)
@@ -1533,12 +1545,11 @@
 	if (e)
 		main_hdr->flags = e->debug ? 0x1 : 0;
 
-	/*
-	 * For SATA srcaddr is specified in number of sectors.
-	 * This expects the sector size to be 512 bytes.
-	 */
-	if (main_hdr->blockid == IBR_HDR_SATA_ID)
-		main_hdr->srcaddr = cpu_to_le32(le32_to_cpu(main_hdr->srcaddr) / 512);
+	/* For SATA srcaddr is specified in number of sectors. */
+	if (main_hdr->blockid == IBR_HDR_SATA_ID) {
+		params->bl_len = image_get_satablksz();
+		main_hdr->srcaddr = cpu_to_le32(le32_to_cpu(main_hdr->srcaddr) / params->bl_len);
+	}
 
 	/* For PCIe srcaddr is not used and must be set to 0xFFFFFFFF. */
 	if (main_hdr->blockid == IBR_HDR_PEX_ID)
@@ -1702,6 +1713,13 @@
 	case IMAGE_CFG_NAND_PAGESZ:
 		el->nandpagesz = strtoul(value1, NULL, 16);
 		break;
+	case IMAGE_CFG_SATA_BLKSZ:
+		el->satablksz = strtoul(value1, NULL, 0);
+		if (el->satablksz & (el->satablksz-1)) {
+			fprintf(stderr, "Invalid SATA block size '%s'\n", value1);
+			return -1;
+		}
+		break;
 	case IMAGE_CFG_BINARY:
 		argi = 0;
 
@@ -1893,6 +1911,8 @@
 	struct stat s;
 	int ret;
 
+	params->bl_len = 1;
+
 	/*
 	 * Do not use sbuf->st_size as it contains size with padding.
 	 * We need original image data size, so stat original file.
@@ -2004,10 +2024,11 @@
 	genimg_print_size(le32_to_cpu(mhdr->blocksize) - sizeof(uint32_t));
 	printf("Data Offset:  ");
 	if (mhdr->blockid == IBR_HDR_SATA_ID)
-		printf("%u Sector%s (LBA)\n", le32_to_cpu(mhdr->srcaddr),
+		printf("%u Sector%s (LBA) = ", le32_to_cpu(mhdr->srcaddr),
 		       le32_to_cpu(mhdr->srcaddr) != 1 ? "s" : "");
-	else
-		genimg_print_size(le32_to_cpu(mhdr->srcaddr));
+	genimg_print_size(le32_to_cpu(mhdr->srcaddr) * params->bl_len);
+	if (mhdr->blockid == IBR_HDR_SATA_ID)
+		printf("Sector Size:  %u Bytes\n", params->bl_len);
 	if (mhdr->blockid == IBR_HDR_SPI_ID && le32_to_cpu(mhdr->destaddr) == 0xFFFFFFFF) {
 		printf("Load Address: XIP\n");
 		printf("Execute Offs: %08x\n", le32_to_cpu(mhdr->execaddr));
@@ -2033,6 +2054,7 @@
 	uint32_t offset;
 	uint32_t size;
 	uint8_t csum;
+	int blksz;
 
 	if (header_size > 192*1024)
 		return -FDT_ERR_BADSTRUCTURE;
@@ -2091,12 +2113,28 @@
 		return -FDT_ERR_BADSTRUCTURE;
 	}
 
+	if (size < 4 || size % 4 != 0)
+		return -FDT_ERR_BADSTRUCTURE;
+
 	/*
 	 * For SATA srcaddr is specified in number of sectors.
-	 * This expects that sector size is 512 bytes.
+	 * Try all possible sector sizes which are power of two,
+	 * at least 512 bytes and up to the 32 kB.
 	 */
-	if (blockid == IBR_HDR_SATA_ID)
-		offset *= 512;
+	if (blockid == IBR_HDR_SATA_ID) {
+		for (blksz = 512; blksz < 0x10000; blksz *= 2) {
+			if (offset * blksz > image_size || offset * blksz + size > image_size)
+				break;
+
+			if (image_checksum32(ptr + offset * blksz, size - 4) ==
+			    *(uint32_t *)(ptr + offset * blksz + size - 4)) {
+				params->bl_len = blksz;
+				return 0;
+			}
+		}
+
+		return -FDT_ERR_BADSTRUCTURE;
+	}
 
 	/*
 	 * For PCIe srcaddr is always set to 0xFFFFFFFF.
@@ -2105,16 +2143,14 @@
 	if (blockid == IBR_HDR_PEX_ID && offset == 0xFFFFFFFF)
 		offset = header_size;
 
-	if (offset > image_size || offset % 4 != 0)
-		return -FDT_ERR_BADSTRUCTURE;
-
-	if (size < 4 || offset + size > image_size || size % 4 != 0)
+	if (offset % 4 != 0 || offset > image_size || offset + size > image_size)
 		return -FDT_ERR_BADSTRUCTURE;
 
 	if (image_checksum32(ptr + offset, size - 4) !=
 	    *(uint32_t *)(ptr + offset + size - 4))
 		return -FDT_ERR_BADSTRUCTURE;
 
+	params->bl_len = 1;
 	return 0;
 }
 
@@ -2129,6 +2165,7 @@
 	void *hdr;
 	int ret;
 	int align, size;
+	unsigned int satablksz;
 
 	fcfg = fopen(params->imagename, "r");
 	if (!fcfg) {
@@ -2166,6 +2203,7 @@
 
 	bootfrom = image_get_bootfrom();
 	version = image_get_version();
+	satablksz = image_get_satablksz();
 	switch (version) {
 		/*
 		 * Fallback to version 0 if no version is provided in the
@@ -2211,11 +2249,14 @@
 	tparams->hdr = hdr;
 
 	/*
-	 * Final SATA and SDIO images must be aligned to 512 bytes.
+	 * Final SATA images must be aligned to disk block size.
+	 * Final SDIO images must be aligned to 512 bytes.
 	 * Final SPI and NAND images must be aligned to 256 bytes.
 	 * Final UART image must be aligned to 128 bytes.
 	 */
-	if (bootfrom == IBR_HDR_SATA_ID || bootfrom == IBR_HDR_SDIO_ID)
+	if (bootfrom == IBR_HDR_SATA_ID)
+		align = satablksz;
+	else if (bootfrom == IBR_HDR_SDIO_ID)
 		align = 512;
 	else if (bootfrom == IBR_HDR_SPI_ID || bootfrom == IBR_HDR_NAND_ID)
 		align = 256;
@@ -2306,6 +2347,9 @@
 	if (version == 0 && mhdr->blockid == IBR_HDR_SATA_ID)
 		fprintf(f, "SATA_PIO_MODE %u\n", (unsigned)mhdr0->satapiomode);
 
+	if (mhdr->blockid == IBR_HDR_SATA_ID)
+		fprintf(f, "SATA_BLKSZ %u\n", params->bl_len);
+
 	/*
 	 * Addresses and sizes which are specified by mkimage command line
 	 * arguments and not in kwbimage config file
@@ -2486,7 +2530,7 @@
 		offset = le32_to_cpu(mhdr->srcaddr);
 
 		if (mhdr->blockid == IBR_HDR_SATA_ID)
-			offset *= 512;
+			offset *= params->bl_len;
 
 		if (mhdr->blockid == IBR_HDR_PEX_ID && offset == 0xFFFFFFFF)
 			offset = header_size;