power: fg-memif: fix an OOB access while getting number of partitions

Currently, sram_address is assigned to the start address of next
partition to see if the address is within the next partition.
This can end up accessing the next element of address map causing
an out of bounds access when the last partition is accessed
already. Fix this by simplifying the calculation to find number
of partitions.

CRs-Fixed: 2130104
Change-Id: Idab43471b4e6e964152b5a46f0e2f0a4bfc03471
Signed-off-by: Subbaraman Narayanamurthy <subbaram@codeaurora.org>
diff --git a/drivers/power/supply/qcom/fg-memif.c b/drivers/power/supply/qcom/fg-memif.c
index c1b5adc..279b097 100644
--- a/drivers/power/supply/qcom/fg-memif.c
+++ b/drivers/power/supply/qcom/fg-memif.c
@@ -847,8 +847,8 @@
 static int fg_get_partition_count(struct fg_chip *chip, u16 sram_addr, int len,
 				int *count)
 {
-	int i, num = 0;
-	u16 end_addr, last_addr = 0;
+	int i, start_partn = 0, end_partn = 0;
+	u16 end_addr = 0;
 
 	end_addr = sram_addr + len / BYTES_PER_SRAM_WORD;
 	if (!(len % BYTES_PER_SRAM_WORD))
@@ -860,24 +860,24 @@
 	}
 
 	for (i = 0; i < NUM_PARTITIONS; i++) {
-		pr_debug("address: %d last_addr: %d\n", sram_addr, last_addr);
 		if (sram_addr >= chip->addr_map[i].partition_start
-			&& sram_addr <= chip->addr_map[i].partition_end
-			&& last_addr < end_addr) {
-			num++;
-			last_addr = chip->addr_map[i].partition_end;
-			sram_addr = chip->addr_map[i+1].partition_start;
-		}
+				&& sram_addr <= chip->addr_map[i].partition_end)
+			start_partn = i + 1;
+
+		if (end_addr >= chip->addr_map[i].partition_start
+				&& end_addr <= chip->addr_map[i].partition_end)
+			end_partn = i + 1;
 	}
 
-	if (num > 0) {
-		*count = num;
-		return 0;
+	if (!start_partn || !end_partn) {
+		pr_err("Couldn't find number of partitions for address %d\n",
+			sram_addr);
+		return -ENXIO;
 	}
 
-	pr_err("Couldn't find number of partitions for address %d\n",
-		sram_addr);
-	return -ENXIO;
+	*count = (end_partn - start_partn) + 1;
+
+	return 0;
 }
 
 static int fg_get_partition_avail_bytes(struct fg_chip *chip, u16 sram_addr,