[U-Boot] [PATCH] disk: Don't assume blksz > legacy_mbr

[Rob Clark <robdclark@gmail.com>]

On some devices, this does not appear to be a valid assumption.  So
figure out the number of blocks we actually need to read.

Signed-off-by: Rob Clark <robdclark@gmail.com>
---
Sorry, took a bit longer to get to a point of testing this.. somehow
himport_r() of the default_environment is clobbering my usb keyboard's
usb_device, which I'm still tracking down.  Good thing that pointers
overwritten with ascii are fairly obvious.

 disk/part_dos.c | 6 ++++--
 disk/part_efi.c | 6 ++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/disk/part_dos.c b/disk/part_dos.c
index 1a36be0446..a9d23e121c 100644
--- a/disk/part_dos.c
+++ b/disk/part_dos.c
@@ -89,9 +89,11 @@ static int test_block_type(unsigned char *buffer)
 
 static int part_test_dos(struct blk_desc *dev_desc)
 {
-	ALLOC_CACHE_ALIGN_BUFFER(legacy_mbr, mbr, dev_desc->blksz);
+	/* blksz *should* be a PoT, but to be safe use DIV_ROUND_UP: */
+	lbaint_t blkcnt = DIV_ROUND_UP(sizeof(legacy_mbr), dev_desc->blksz);
+	ALLOC_CACHE_ALIGN_BUFFER(legacy_mbr, mbr, blkcnt * dev_desc->blksz);
 
-	if (blk_dread(dev_desc, 0, 1, (ulong *)mbr) != 1)
+	if (blk_dread(dev_desc, 0, blkcnt, (ulong *)mbr) != 1)
 		return -1;
 
 	if (test_block_type((unsigned char *)mbr) != DOS_MBR)
diff --git a/disk/part_efi.c b/disk/part_efi.c
index 208bb14ee8..e00f6c9d24 100644
--- a/disk/part_efi.c
+++ b/disk/part_efi.c
@@ -923,7 +923,9 @@ static int is_pmbr_valid(legacy_mbr * mbr)
 static int is_gpt_valid(struct blk_desc *dev_desc, u64 lba,
 			gpt_header *pgpt_head, gpt_entry **pgpt_pte)
 {
-	ALLOC_CACHE_ALIGN_BUFFER(legacy_mbr, mbr, dev_desc->blksz);
+	/* blksz *should* be a PoT, but to be safe use DIV_ROUND_UP: */
+	lbaint_t blkcnt = DIV_ROUND_UP(sizeof(legacy_mbr), dev_desc->blksz);
+	ALLOC_CACHE_ALIGN_BUFFER(legacy_mbr, mbr, blkcnt * dev_desc->blksz);
 
 	if (!dev_desc || !pgpt_head) {
 		printf("%s: Invalid Argument(s)\n", __func__);
@@ -931,7 +933,7 @@ static int is_gpt_valid(struct blk_desc *dev_desc, u64 lba,
 	}
 
 	/* Read MBR Header from device */
-	if (blk_dread(dev_desc, 0, 1, (ulong *)mbr) != 1) {
+	if (blk_dread(dev_desc, 0, blkcnt, (ulong *)mbr) != 1) {
 		printf("*** ERROR: Can't read MBR header ***\n");
 		return 0;
 	}
-- 
2.13.5