[RFC PATCH v2 07/10] dma-mapping: Add support of preserving dmam allocations

[Samiullah Khawaja <skhawaja@google.com>]

Device drivers use dmam_* allocation APIs for managed lifecycle of DMA
allocations. Add support of preserving managed DMA allocations so these
can be preserved and on restore these are managed again.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/dma-mapping.h | 50 ++++++++++++++++++++++
 kernel/dma/mapping.c        | 83 +++++++++++++++++++++++++++++++++++++
 2 files changed, 133 insertions(+)

diff --git a/include/linux/dma-mapping.h b/include/linux/dma-mapping.h
index 3756fc15467b..4c0c428b6b1c 100644
--- a/include/linux/dma-mapping.h
+++ b/include/linux/dma-mapping.h
@@ -218,6 +218,14 @@ void dma_unpreserve_allocation(struct device *dev, u64 state);
 void *dma_restore_allocation_attrs(struct device *dev, size_t size,
 				   dma_addr_t *dma_handle, gfp_t gfp,
 				   unsigned long attrs, u64 state);
+int dmam_preserve_allocation_attrs(struct device *dev, void *cpu_addr,
+				   size_t size, dma_addr_t dma_handle,
+				   unsigned long attrs, u64 *state);
+void dmam_unpreserve_allocation(struct device *dev, void *cpu_addr,
+				size_t size, dma_addr_t dma_handle, u64 state);
+void *dmam_restore_allocation_attrs(struct device *dev, size_t size,
+				    dma_addr_t *dma_handle, gfp_t gfp,
+				    unsigned long attrs, u64 state);
 #endif
 #else /* CONFIG_HAS_DMA */
 static inline dma_addr_t dma_map_page_attrs(struct device *dev,
@@ -523,6 +531,25 @@ static inline void *dma_restore_allocation_attrs(struct device *dev, size_t size
 {
 	return NULL;
 }
+
+static inline int dmam_preserve_allocation_attrs(struct device *dev, void *cpu_addr,
+						 size_t size, dma_addr_t dma_handle,
+						 unsigned long attrs, u64 *state)
+{
+	return -EOPNOTSUPP;
+}
+
+static inline void dmam_unpreserve_allocation(struct device *dev, void *cpu_addr,
+					      size_t size, dma_addr_t dma_handle, u64 state)
+{
+}
+
+static inline void *dmam_restore_allocation_attrs(struct device *dev, size_t size,
+						  dma_addr_t *dma_handle, gfp_t gfp,
+						  unsigned long attrs, u64 state)
+{
+	return NULL;
+}
 #endif
 
 struct page *dma_alloc_pages(struct device *dev, size_t size,
@@ -797,6 +824,29 @@ static inline void *dmam_alloc_coherent(struct device *dev, size_t size,
 			(gfp & __GFP_NOWARN) ? DMA_ATTR_NO_WARN : 0);
 }
 
+static inline int dmam_preserve_coherent_allocation(struct device *dev, void *cpu_addr,
+						    size_t size, dma_addr_t dma_handle, u64 *state)
+{
+	return dmam_preserve_allocation_attrs(dev, cpu_addr, size,
+					      dma_handle, 0, state);
+}
+
+static inline void dmam_unpreserve_coherent_allocation(struct device *dev, void *cpu_addr,
+						       size_t size, dma_addr_t dma_handle,
+						       u64 state)
+{
+	dmam_unpreserve_allocation(dev, cpu_addr, size, dma_handle, state);
+}
+
+static inline void *dmam_restore_coherent_allocation(struct device *dev, size_t size,
+						     dma_addr_t *dma_handle,
+						     gfp_t gfp, u64 state)
+{
+	return dmam_restore_allocation_attrs(dev, size, dma_handle, gfp,
+					     (gfp & __GFP_NOWARN) ? DMA_ATTR_NO_WARN : 0,
+					     state);
+}
+
 static inline void *dma_alloc_wc(struct device *dev, size_t size,
 				 dma_addr_t *dma_addr, gfp_t gfp)
 {
diff --git a/kernel/dma/mapping.c b/kernel/dma/mapping.c
index 3b07bf30722c..aec1222ab5ae 100644
--- a/kernel/dma/mapping.c
+++ b/kernel/dma/mapping.c
@@ -37,12 +37,25 @@ struct dma_devres {
 	void		*vaddr;
 	dma_addr_t	dma_handle;
 	unsigned long	attrs;
+#ifdef CONFIG_DMA_LIVEUPDATE
+	bool is_preserved;
+#endif
 };
 
 static void dmam_release(struct device *dev, void *res)
 {
 	struct dma_devres *this = res;
 
+#ifdef CONFIG_DMA_LIVEUPDATE
+	/*
+	 * Freeing the preserved memory is dangerous as it can cause UAF in the
+	 * current or next kernel if the memory is still being used by the
+	 * device.
+	 */
+	if (WARN_ON(this->is_preserved))
+		return;
+#endif
+
 	dma_free_attrs(dev, this->size, this->vaddr, this->dma_handle,
 			this->attrs);
 }
@@ -675,6 +688,76 @@ void *dma_restore_allocation_attrs(struct device *dev, size_t size,
 	return cpu_addr;
 }
 EXPORT_SYMBOL(dma_restore_allocation_attrs);
+
+int dmam_preserve_allocation_attrs(struct device *dev, void *cpu_addr,
+				   size_t size, dma_addr_t dma_handle,
+				   unsigned long attrs, u64 *state)
+{
+	struct dma_devres match_data = { size, cpu_addr, dma_handle };
+	struct dma_devres *dr;
+	int ret;
+
+	dr = devres_find(dev, dmam_release, dmam_match, &match_data);
+	if (!dr)
+		return -EINVAL;
+
+	if (dr->is_preserved)
+		return -EINVAL;
+
+	ret = dma_preserve_allocation_attrs(dev, cpu_addr, size, dma_handle, attrs, state);
+	if (ret)
+		return ret;
+
+	dr->is_preserved = true;
+	return 0;
+}
+EXPORT_SYMBOL(dmam_preserve_allocation_attrs);
+
+void dmam_unpreserve_allocation(struct device *dev, void *cpu_addr,
+				size_t size, dma_addr_t dma_handle, u64 state)
+{
+	struct dma_devres match_data = { size, cpu_addr, dma_handle };
+	struct dma_devres *dr;
+
+	dr = devres_find(dev, dmam_release, dmam_match, &match_data);
+	if (!dr)
+		return;
+
+	if (!dr->is_preserved)
+		return;
+
+	dma_unpreserve_allocation(dev, state);
+	dr->is_preserved = false;
+}
+EXPORT_SYMBOL(dmam_unpreserve_allocation);
+
+void *dmam_restore_allocation_attrs(struct device *dev, size_t size,
+				    dma_addr_t *dma_handle, gfp_t gfp,
+				    unsigned long attrs, u64 state)
+{
+	struct dma_devres *dr;
+	void *vaddr;
+
+	dr = devres_alloc(dmam_release, sizeof(*dr), gfp);
+	if (!dr)
+		return NULL;
+
+	vaddr = dma_restore_allocation_attrs(dev, size, dma_handle, gfp, attrs, state);
+	if (!vaddr) {
+		devres_free(dr);
+		return NULL;
+	}
+
+	dr->vaddr = vaddr;
+	dr->dma_handle = *dma_handle;
+	dr->size = size;
+	dr->attrs = attrs;
+
+	devres_add(dev, dr);
+
+	return vaddr;
+}
+EXPORT_SYMBOL(dmam_restore_allocation_attrs);
 #endif
 
 void *dma_alloc_attrs(struct device *dev, size_t size, dma_addr_t *dma_handle,
-- 
2.55.0.795.g602f6c329a-goog