platform: msm_shared: Add flashing support for multiple LUNs in UFS.

CRs-Fixed: 573103
Change-Id: I0dc0aacb2ebf32c5150bd02e5215d01a2edafab2
diff --git a/platform/msm_shared/ufs.c b/platform/msm_shared/ufs.c
index 5728bf7..8244576 100644
--- a/platform/msm_shared/ufs.c
+++ b/platform/msm_shared/ufs.c
@@ -91,7 +91,7 @@
 	int                  ret;
 
 	req.data_buffer_base = buffer;
-	req.lun              = 0;
+	req.lun              = dev->current_lun;
 	req.num_blocks       = num_blocks;
 	req.start_lba        = start_lba / dev->block_size;
 
@@ -111,7 +111,7 @@
 	int                  ret;
 
 	req.data_buffer_base = buffer;
-	req.lun              = 0;
+	req.lun              = dev->current_lun;
 	req.num_blocks       = num_blocks;
 	req.start_lba        = start_lba / dev->block_size;
 
@@ -145,12 +145,22 @@
 
 uint64_t ufs_get_dev_capacity(struct ufs_dev* dev)
 {
-	return dev->capacity;
+	uint64_t capacity;
+	uint8_t lun = dev->current_lun;
+
+	capacity = dev->lun_cfg[lun].logical_blk_cnt * dev->block_size;
+
+	return capacity;
 }
 
 uint32_t ufs_get_erase_blk_size(struct ufs_dev* dev)
 {
-	return dev->erase_blk_size;
+	uint32_t erase_blk_size;
+	uint8_t lun = dev->current_lun;
+
+	erase_blk_size = dev->lun_cfg[lun].erase_blk_size;
+
+	return erase_blk_size;
 }
 
 uint32_t ufs_get_serial_num(struct ufs_dev* dev)
@@ -172,10 +182,16 @@
 	return dev->block_size;
 }
 
+uint8_t ufs_get_num_of_luns(struct ufs_dev* dev)
+{
+	return dev->num_lus;
+}
+
 int ufs_init(struct ufs_dev *dev)
 {
 	uint32_t ret = UFS_SUCCESS;
 	uint64_t cap;
+	uint8_t lun = 0;
 
 	dev->block_size = 4096;
 
@@ -220,13 +236,24 @@
 		goto ufs_init_err;
 	}
 
-	ret = dme_read_unit_desc(dev, 0);
+	ret = dme_read_device_desc(dev);
 	if (ret != UFS_SUCCESS)
 	{
-		dprintf(CRITICAL, "UFS dme_read_unit_desc failed\n");
+		dprintf(CRITICAL, "dme_read_dev_desc read failed\n");
 		goto ufs_init_err;
 	}
 
+
+	for(lun=0; lun < dev->num_lus; lun++)
+	{
+		ret = dme_read_unit_desc(dev, lun);
+		if (ret != UFS_SUCCESS)
+		{
+			dprintf(CRITICAL, "UFS dme_read_unit_desc failed\n");
+			goto ufs_init_err;
+		}
+	}
+
 	dprintf(CRITICAL,"UFS init success\n");
 
 ufs_init_err: