[PATCH 1/2] nvme: fix PRP pool resize before first use

[Ahmad Fatoum <a.fatoum@barebox.org>]

It's not safe to call dma_free_coherent on a NULL pointer. Implement
dma_realloc_coherent, which correctly handles reallocating a buffer
without leaking memory and use it to ensure only actually allocated
buffers are freed.

Fixes: aedcb568afe4 ("drivers: Import a very basic NVME implementation from Linux")
Reported-by: Ben Pye # Matrix
Signed-off-by: Ahmad Fatoum <a.fatoum@barebox.org>
---
 drivers/dma/map.c       | 21 +++++++++++++++++++++
 drivers/nvme/host/pci.c | 17 +++++++++++------
 include/dma.h           |  5 +++++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/drivers/dma/map.c b/drivers/dma/map.c
index 2980d4c2e790..11671f9f663d 100644
--- a/drivers/dma/map.c
+++ b/drivers/dma/map.c
@@ -68,3 +68,24 @@ void dma_unmap_single(struct device *dev, dma_addr_t dma_addr,
 	debug_dma_unmap(dev, dma_addr, size, dir);
 }
 EXPORT_SYMBOL(dma_unmap_single);
+
+void *dma_realloc_coherent(struct device *dev, void *oldbuf,
+			   size_t oldsize, size_t newsize,
+			   dma_addr_t *dma_handle)
+{
+	dma_addr_t newdma = DMA_ERROR_CODE;
+	void *newbuf = ZERO_SIZE_PTR;
+
+	if (newsize) {
+		newbuf = dma_alloc_coherent(dev, newsize, &newdma);
+		if (!newbuf)
+			return NULL;
+	}
+
+	if (oldbuf)
+		dma_free_coherent(dev, oldbuf, *dma_handle, oldsize);
+
+	*dma_handle = newdma;
+	return newbuf;
+}
+EXPORT_SYMBOL(dma_realloc_coherent);
diff --git a/drivers/nvme/host/pci.c b/drivers/nvme/host/pci.c
index 6343af25026b..29a6db76e0f6 100644
--- a/drivers/nvme/host/pci.c
+++ b/drivers/nvme/host/pci.c
@@ -93,13 +93,18 @@ static int nvme_pci_setup_prps(struct nvme_dev *dev,
 
 	nprps = DIV_ROUND_UP(length, page_size);
 	if (nprps > dev->prp_pool_size) {
-		dma_free_coherent(DMA_DEVICE_BROKEN,
-				  dev->prp_pool, dev->prp_dma,
-				  dev->prp_pool_size * sizeof(u64));
+		__le64 *prp_pool;
+
+		prp_pool = dma_realloc_coherent(DMA_DEVICE_BROKEN,
+						dev->prp_pool,
+						dev->prp_pool_size * sizeof(*prp_pool),
+						nprps * sizeof(*prp_pool),
+						&dev->prp_dma);
+		if (!prp_pool)
+			return -ENOMEM;
+
+		dev->prp_pool = prp_pool;
 		dev->prp_pool_size = nprps;
-		dev->prp_pool = dma_alloc_coherent(DMA_DEVICE_BROKEN,
-						   nprps * sizeof(u64),
-						   &dev->prp_dma);
 	}
 
 	prp_list = dev->prp_pool;
diff --git a/include/dma.h b/include/dma.h
index 60937a69aa55..c285b2095a58 100644
--- a/include/dma.h
+++ b/include/dma.h
@@ -161,4 +161,9 @@ static inline bool dma_map_buf_is_aligned(struct device *dev, const void *buf, s
 	return PTR_IS_ALIGNED(buf, ARCH_DMA_MINALIGN) &&
 		IS_ALIGNED(size, ARCH_DMA_MINALIGN);
 }
+
+void *dma_realloc_coherent(struct device *dev, void *oldbuf,
+			   size_t oldsize, size_t newsize,
+			   dma_addr_t *dma_handle);
+
 #endif /* __DMA_H */
-- 
2.47.3