fix(ufs): flush the entire PRDT

Previously, if the image being read exceeded 12,800 KB (or 50 PRDT
entries of size 256 KB), the UFS driver would not flush the entire
Physical Region Descriptor Table (PRDT). This would cause the UFS host
controller to read empty PRDT entries, which eventually would make the
system crash. This change updates the UFS driver to flush the entire
PRDT, irrespective of the size of the image being read.

This change also throws an error if the memory allocated for UFS
descriptors is not sufficient to hold the entire Physical Region
Descriptor Table (PRDT).

Signed-off-by: Jorge Troncoso <jatron@google.com>
Change-Id: I291dc62748992481be3cc156ce1474a6e3990ea9
diff --git a/drivers/ufs/ufs.c b/drivers/ufs/ufs.c
index d6c893d..1fe0b0e 100644
--- a/drivers/ufs/ufs.c
+++ b/drivers/ufs/ufs.c
@@ -289,6 +289,8 @@
 	unsigned int ulba;
 	unsigned int lba_cnt;
 	int prdt_size;
+	uintptr_t desc_limit;
+	size_t flush_size;
 
 	hd = (utrd_header_t *)utrd->header;
 	upiu = (cmd_upiu_t *)utrd->upiu;
@@ -342,17 +344,25 @@
 		assert(0);
 		break;
 	}
-	if (hd->dd == DD_IN)
+	if (hd->dd == DD_IN) {
 		flush_dcache_range(buf, length);
-	else if (hd->dd == DD_OUT)
+	} else if (hd->dd == DD_OUT) {
 		inv_dcache_range(buf, length);
+	}
+
+	utrd->size_prdt = 0;
 	if (length) {
 		upiu->exp_data_trans_len = htobe32(length);
 		assert(lba_cnt <= UINT16_MAX);
 		prdt = (prdt_t *)utrd->prdt;
 
+		desc_limit = ufs_params.desc_base + ufs_params.desc_size;
 		prdt_size = 0;
 		while (length > 0) {
+			if ((uintptr_t)prdt + sizeof(prdt_t) > desc_limit) {
+				ERROR("UFS: Exceeded descriptor limit. Image is too large\n");
+				panic();
+			}
 			prdt->dba = (unsigned int)(buf & UINT32_MAX);
 			prdt->dbau = (unsigned int)((buf >> 32) & UINT32_MAX);
 			/* prdt->dbc counts from 0 */
@@ -372,7 +382,8 @@
 		hd->prdto = (utrd->size_upiu + utrd->size_resp_upiu) >> 2;
 	}
 
-	flush_dcache_range((uintptr_t)utrd->header, UFS_DESC_SIZE);
+	flush_size = utrd->prdt + utrd->size_prdt - utrd->header;
+	flush_dcache_range(utrd->header, flush_size);
 	return 0;
 }