[RFC PATCH v2 00/10] dma-mapping: Add preservation of direct allocations

[RFC PATCH v2 00/10] dma-mapping: Add preservation of direct allocations

[Samiullah Khawaja]

This is an RFC v2 to discuss the preservation of DMA allocations for
devices that are using direct allocation mode. Note preservation here
means preservation of physical memory using KHO only. IOVAs and IOMMU
mappings are out of scope.

The complexity and need of this was discussed in the thread mentioned
below. It was proposed to allow preservation of only a subset of
allocations with minimum support. Offline discussion with Will Deacon
concluded that this series is still needed, because we still want to use
coherent allocations for Arm sMMUv3 ste/cds, so sending this v2.

https://lore.kernel.org/all/aepfYzkI7NsVGCF0@google.com/
https://lore.kernel.org/all/aepRy7Gp7Ng85Zr7@willie-the-truck/

Note DMA allocation allows memory allocation from various allocators in
various configurations. This RFC supports preservation with,

- Direct Allocation Only.
- Pages allocated through alloc_page (Buddy allocator).
- No CMA support.
- No DMA pools.
- No device coherent memory regions.

Note that this series does not include/use the page-split and page
refcount fixes during KHO restore, that are in review here:
https://lore.kernel.org/all/20260703020832.1731864-1-praan@google.com/

The changes are also pushed here:
https://github.com/samikhawaja/linux/tree/dma-alloc-preserve-direct-v2

Testing:

This patch series includes a Kunit test, for testing dma allocation
preservation, by using the KHO stubs added in a separate patch series:
https://lore.kernel.org/all/20260521193202.746810-1-skhawaja@google.com/

The following test runs are executed in qemu.

Run on x86:

   KTAP version 1
   1..1
       KTAP version 1
       # Subtest: dma_direct_liveupdate
       # module: direct_test
       1..5
           KTAP version 1
           # Subtest: test_dma_direct_coherent
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_coherent: pass:5 fail:0 skip:0 total:5
       ok 1 test_dma_direct_coherent
           KTAP version 1
           # Subtest: test_dma_direct_non_coherent
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_non_coherent: pass:5 fail:0 skip:0 total:5
       ok 2 test_dma_direct_non_coherent
           KTAP version 1
           # Subtest: test_dma_direct_coherent_atomic
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_coherent_atomic: pass:5 fail:0 skip:0 total:5
       ok 3 test_dma_direct_coherent_atomic
           KTAP version 1
           # Subtest: test_dma_direct_non_coherent_atomic
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_non_coherent_atomic: pass:5 fail:0 skip:0 total:5
       ok 4 test_dma_direct_non_coherent_atomic
       ok 5 test_dmam_preservation
   # dma_direct_liveupdate: pass:5 fail:0 skip:0 total:5
   # Totals: pass:21 fail:0 skip:0 total:21
   ok 1 dma_direct_liveupdate

Run on arm64:

   KTAP version 1
   1..1
       KTAP version 1
       # Subtest: dma_direct_liveupdate
       # module: direct_test
       1..5
           KTAP version 1
           # Subtest: test_dma_direct_coherent
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_coherent: pass:5 fail:0 skip:0 total:5
       ok 1 test_dma_direct_coherent
           KTAP version 1
           # Subtest: test_dma_direct_non_coherent
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_non_coherent: pass:5 fail:0 skip:0 total:5
       ok 2 test_dma_direct_non_coherent
           KTAP version 1
           # Subtest: test_dma_direct_coherent_atomic
           ok 1 size=128
           ok 2 size=4096
           ok 3 size=8192
           ok 4 size=12288
           ok 5 size=32768
       # test_dma_direct_coherent_atomic: pass:5 fail:0 skip:0 total:5
       ok 3 test_dma_direct_coherent_atomic
           KTAP version 1
           # Subtest: test_dma_direct_non_coherent_atomic
       # test_dma_direct_non_coherent_atomic: DMA allocation using pool, expecting -EOPNOTSUPP
           ok 1 size=128
       # test_dma_direct_non_coherent_atomic: DMA allocation using pool, expecting -EOPNOTSUPP
           ok 2 size=4096
       # test_dma_direct_non_coherent_atomic: DMA allocation using pool, expecting -EOPNOTSUPP
           ok 3 size=8192
       # test_dma_direct_non_coherent_atomic: DMA allocation using pool, expecting -EOPNOTSUPP
           ok 4 size=12288
       # test_dma_direct_non_coherent_atomic: DMA allocation using pool, expecting -EOPNOTSUPP
           ok 5 size=32768
       # test_dma_direct_non_coherent_atomic: pass:5 fail:0 skip:0 total:5
       ok 4 test_dma_direct_non_coherent_atomic
       ok 5 test_dmam_preservation
   # dma_direct_liveupdate: pass:5 fail:0 skip:0 total:5
   # Totals: pass:21 fail:0 skip:0 total:21
   ok 1 dma_direct_liveupdate

Looking forward to your feedback on this.

Changelog:

rfc v2:

  - Add dma_is_from_cma() to check if an allocation is from CMA.
  - Add dma_is_from_dev_coherent() to check if an allocation is from a device coherent region.
  - Add support for preserving device-managed DMA allocations (dmam_alloc_*).
  - Add KUnit test suite for dma-direct liveupdate preservation.

Samiullah Khawaja (10):
  dma: Add DMA allocation preservation KHO ABI
  dma/pool: Add an API to check if DMA allocation is from pool
  dma: contiguous: Add API to check if an allocation is from CMA
  dma-coherent: Allow checking if allocation is from dev coherent region
  dma-direct: Add API to preserve/restore allocations
  dma-mapping: Add API to preserve/restore DMA allocation
  dma-mapping: Add support of preserving dmam allocations
  dma: contiguous: Export is_from_cma helper for kunit
  dma: pool: Export the is_from_pool helper for kunit
  dma-direct: Add KUnit test for liveupdate preservation

 include/linux/dma-direct.h        |  29 +++++
 include/linux/dma-map-ops.h       |  10 ++
 include/linux/dma-mapping.h       | 100 ++++++++++++++
 include/linux/kho/abi/dma_alloc.h |  32 +++++
 kernel/dma/Kconfig                |  12 ++
 kernel/dma/Makefile               |   1 +
 kernel/dma/coherent.c             |  11 ++
 kernel/dma/contiguous.c           |  27 ++++
 kernel/dma/direct.c               | 189 +++++++++++++++++++++++++++
 kernel/dma/direct_test.c          | 208 ++++++++++++++++++++++++++++++
 kernel/dma/mapping.c              | 132 +++++++++++++++++++
 kernel/dma/pool.c                 |  17 +++
 12 files changed, 768 insertions(+)
 create mode 100644 include/linux/kho/abi/dma_alloc.h
 create mode 100644 kernel/dma/direct_test.c


base-commit: ec4084bc445027a52f600e30a976928be1ba1950
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 01/10] dma: Add DMA allocation preservation KHO ABI

[Samiullah Khawaja]

The DMA allocations can be backed by a variety of allocators. Add KHO
ABI for the preservation of contiguous allocations that are done through
dma-direct.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/kho/abi/dma_alloc.h | 32 +++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)
 create mode 100644 include/linux/kho/abi/dma_alloc.h

diff --git a/include/linux/kho/abi/dma_alloc.h b/include/linux/kho/abi/dma_alloc.h
new file mode 100644
index 000000000000..ae38cf1a1fc9
--- /dev/null
+++ b/include/linux/kho/abi/dma_alloc.h
@@ -0,0 +1,32 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef _LINUX_KHO_ABI_DMA_ALLOC_H
+#define _LINUX_KHO_ABI_DMA_ALLOC_H
+
+#include <linux/types.h>
+
+/**
+ * DOC: DMA Alloc ABI
+ *
+ * This header defines the structures used to serialize the state of DMA
+ * allocations, done by device driver, across a Live Update.
+ *
+ * Only DMA allocations done through dma-direct that are contiguous and
+ * allocated using alloc_page are supported.
+ */
+
+/**
+ * struct dma_alloc_ser - Serialized state of a single DMA allocation
+ * @page_phys: Physical address of the preserved pages
+ * @size: Size of the DMA allocation
+ * @attrs: DMA allocation attributes
+ * @force_decrypted: Whether the memory is force decrypted in previous kernel
+ */
+struct dma_alloc_ser {
+	u64 page_phys;
+	u64 size;
+	u64 attrs;
+	u8 force_decrypted;
+	u8 padding[7];
+} __packed;
+
+#endif /* _LINUX_KHO_ABI_DMA_ALLOC_H */
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 02/10] dma/pool: Add an API to check if DMA allocation is from pool

[Samiullah Khawaja]

DMA allocations can be done through DMA pools, add an API that can be
used to check if an allocation is done from a pool. This will be used in
the later commit during preservation of DMA allocation.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/dma-map-ops.h |  1 +
 kernel/dma/pool.c           | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/include/linux/dma-map-ops.h b/include/linux/dma-map-ops.h
index 6a1832a73cad..9382785a2faa 100644
--- a/include/linux/dma-map-ops.h
+++ b/include/linux/dma-map-ops.h
@@ -216,6 +216,7 @@ struct page *dma_alloc_from_pool(struct device *dev, size_t size,
 		bool (*phys_addr_ok)(struct device *, phys_addr_t, size_t));
 bool dma_free_from_pool(struct device *dev, void *start, size_t size);
 
+bool dma_is_from_pool(void *start, size_t size);
 int dma_direct_set_offset(struct device *dev, phys_addr_t cpu_start,
 		dma_addr_t dma_start, u64 size);
 
diff --git a/kernel/dma/pool.c b/kernel/dma/pool.c
index 2b2fbb709242..d1a0bb03e478 100644
--- a/kernel/dma/pool.c
+++ b/kernel/dma/pool.c
@@ -307,3 +307,16 @@ bool dma_free_from_pool(struct device *dev, void *start, size_t size)
 
 	return false;
 }
+
+bool dma_is_from_pool(void *start, size_t size)
+{
+	struct gen_pool *pool = NULL;
+
+	while ((pool = dma_guess_pool(pool, 0))) {
+		if (!gen_pool_has_addr(pool, (unsigned long)start, size))
+			continue;
+		return true;
+	}
+
+	return false;
+}
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 03/10] dma: contiguous: Add API to check if an allocation is from CMA

[Samiullah Khawaja]

Add an API that uses cma_for_each() to check whether a DMA allocation is
from a CMA region.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/dma-map-ops.h |  6 ++++++
 kernel/dma/contiguous.c     | 24 ++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/include/linux/dma-map-ops.h b/include/linux/dma-map-ops.h
index 9382785a2faa..9b91c839ceb5 100644
--- a/include/linux/dma-map-ops.h
+++ b/include/linux/dma-map-ops.h
@@ -106,6 +106,7 @@ struct page *dma_alloc_contiguous(struct device *dev, size_t size, gfp_t gfp);
 void dma_free_contiguous(struct device *dev, struct page *page, size_t size);
 
 void dma_contiguous_early_fixup(phys_addr_t base, unsigned long size);
+bool dma_is_from_cma(phys_addr_t phys, size_t size);
 #else /* CONFIG_DMA_CMA */
 static inline struct cma *dev_get_cma_area(struct device *dev)
 {
@@ -145,6 +146,11 @@ static inline void dma_free_contiguous(struct device *dev, struct page *page,
 {
 	__free_pages(page, get_order(size));
 }
+
+static inline bool dma_is_from_cma(phys_addr_t phys, size_t size)
+{
+	return false;
+}
 #endif /* CONFIG_DMA_CMA*/
 
 #ifdef CONFIG_DMA_DECLARE_COHERENT
diff --git a/kernel/dma/contiguous.c b/kernel/dma/contiguous.c
index 03f52bd17120..0b1702bf39ca 100644
--- a/kernel/dma/contiguous.c
+++ b/kernel/dma/contiguous.c
@@ -130,6 +130,30 @@ struct cma *dev_get_cma_area(struct device *dev)
 }
 EXPORT_SYMBOL_GPL(dev_get_cma_area);
 
+struct dma_cma_check_data {
+	phys_addr_t start;
+	phys_addr_t end;
+};
+
+static int dma_cma_check_area(struct cma *cma, void *data)
+{
+	struct dma_cma_check_data *chk = data;
+
+	if (cma_intersects(cma, chk->start, chk->end))
+		return 1;
+	return 0;
+}
+
+bool dma_is_from_cma(phys_addr_t phys, size_t size)
+{
+	struct dma_cma_check_data chk = {
+		.start = phys,
+		.end = phys + size - 1,
+	};
+
+	return cma_for_each_area(dma_cma_check_area, &chk) != 0;
+}
+
 #ifdef CONFIG_DMA_NUMA_CMA
 
 static struct cma *dma_contiguous_numa_area[MAX_NUMNODES];
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 04/10] dma-coherent: Allow checking if allocation is from dev coherent region

[Samiullah Khawaja]

Add an API that allows if a DMA allocation is from device coherent
region.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/dma-map-ops.h |  3 +++
 kernel/dma/coherent.c       | 11 +++++++++++
 2 files changed, 14 insertions(+)

diff --git a/include/linux/dma-map-ops.h b/include/linux/dma-map-ops.h
index 9b91c839ceb5..b13e5eebe6d3 100644
--- a/include/linux/dma-map-ops.h
+++ b/include/linux/dma-map-ops.h
@@ -162,6 +162,7 @@ int dma_alloc_from_dev_coherent(struct device *dev, ssize_t size,
 int dma_release_from_dev_coherent(struct device *dev, int order, void *vaddr);
 int dma_mmap_from_dev_coherent(struct device *dev, struct vm_area_struct *vma,
 		void *cpu_addr, size_t size, int *ret);
+bool dma_is_from_dev_coherent(struct device *dev, void *vaddr);
 #else
 static inline int dma_declare_coherent_memory(struct device *dev,
 		phys_addr_t phys_addr, dma_addr_t device_addr, size_t size)
@@ -172,6 +173,8 @@ static inline int dma_declare_coherent_memory(struct device *dev,
 #define dma_alloc_from_dev_coherent(dev, size, handle, ret) (0)
 #define dma_release_from_dev_coherent(dev, order, vaddr) (0)
 #define dma_mmap_from_dev_coherent(dev, vma, vaddr, order, ret) (0)
+#define dma_is_from_dev_coherent(dev, vaddr) (false)
+
 static inline void dma_release_coherent_memory(struct device *dev) { }
 #endif /* CONFIG_DMA_DECLARE_COHERENT */
 
diff --git a/kernel/dma/coherent.c b/kernel/dma/coherent.c
index bcdc0f76d2e8..76cbe11e1f90 100644
--- a/kernel/dma/coherent.c
+++ b/kernel/dma/coherent.c
@@ -231,6 +231,17 @@ int dma_release_from_dev_coherent(struct device *dev, int order, void *vaddr)
 	return __dma_release_from_coherent(mem, order, vaddr);
 }
 
+bool dma_is_from_dev_coherent(struct device *dev, void *vaddr)
+{
+	struct dma_coherent_mem *mem = dev_get_coherent_memory(dev);
+
+	if (mem && vaddr >= mem->virt_base && vaddr <
+		   (mem->virt_base + ((dma_addr_t)mem->size << PAGE_SHIFT)))
+		return true;
+
+	return false;
+}
+
 static int __dma_mmap_from_coherent(struct dma_coherent_mem *mem,
 		struct vm_area_struct *vma, void *vaddr, size_t size, int *ret)
 {
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 05/10] dma-direct: Add API to preserve/restore allocations

[Samiullah Khawaja]

Add an API to preserve/restore the DMA direct allocation for liveupdate.
The underlying memory is preserved/restored using KHO. During restore
the memory is setup based on the device configuration, gfp flags and
allocation attributes. Once restored, the driver can use the usual
dma_free* API to deallocate the restored DMA allocation.

This API will be used to add support in dma_alloc* APIs to
preserve/restore the DMA allocations.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 include/linux/dma-direct.h |  29 ++++++
 kernel/dma/Kconfig         |   4 +
 kernel/dma/direct.c        | 189 +++++++++++++++++++++++++++++++++++++
 3 files changed, 222 insertions(+)

diff --git a/include/linux/dma-direct.h b/include/linux/dma-direct.h
index c249912456f9..0efe2bc1a815 100644
--- a/include/linux/dma-direct.h
+++ b/include/linux/dma-direct.h
@@ -141,6 +141,35 @@ static inline bool dma_capable(struct device *dev, dma_addr_t addr, size_t size,
 u64 dma_direct_get_required_mask(struct device *dev);
 void *dma_direct_alloc(struct device *dev, size_t size, dma_addr_t *dma_handle,
 		gfp_t gfp, unsigned long attrs);
+
+#ifdef CONFIG_DMA_LIVEUPDATE
+int dma_direct_preserve_allocation(struct device *dev, void *cpu_addr,
+				   size_t size, dma_addr_t dma_handle,
+				   unsigned long attrs, u64 *state);
+void dma_direct_unpreserve_allocation(struct device *dev, u64 state);
+void *dma_direct_restore_allocation(struct device *dev, size_t size,
+				    dma_addr_t *dma_handle, gfp_t gfp,
+				    unsigned long attrs, u64 state);
+#else
+static inline int dma_direct_preserve_allocation(struct device *dev, void *cpu_addr,
+						 size_t size, dma_addr_t dma_handle,
+						 unsigned long attrs, u64 *state)
+{
+	return -EOPNOTSUPP;
+}
+
+static inline void dma_direct_unpreserve_allocation(struct device *dev, u64 state)
+{
+}
+
+static inline void *dma_direct_restore_allocation(struct device *dev, size_t size,
+						  dma_addr_t *dma_handle, gfp_t gfp,
+						  unsigned long attrs, u64 state)
+{
+	return NULL;
+}
+#endif
+
 void dma_direct_free(struct device *dev, size_t size, void *cpu_addr,
 		dma_addr_t dma_addr, unsigned long attrs);
 struct page *dma_direct_alloc_pages(struct device *dev, size_t size,
diff --git a/kernel/dma/Kconfig b/kernel/dma/Kconfig
index bfef21b4a9ae..d93958931303 100644
--- a/kernel/dma/Kconfig
+++ b/kernel/dma/Kconfig
@@ -265,3 +265,7 @@ config DMA_MAP_BENCHMARK
 	  performance of dma_(un)map_page.
 
 	  See tools/testing/selftests/dma/dma_map_benchmark.c
+
+config DMA_LIVEUPDATE
+	bool "Enable preservation of DMA direct allocations"
+	depends on LIVEUPDATE
diff --git a/kernel/dma/direct.c b/kernel/dma/direct.c
index ec887f443741..117ebae8a81f 100644
--- a/kernel/dma/direct.c
+++ b/kernel/dma/direct.c
@@ -6,6 +6,8 @@
  */
 #include <linux/memblock.h> /* for max_pfn */
 #include <linux/export.h>
+#include <linux/kexec_handover.h>
+#include <linux/kho/abi/dma_alloc.h>
 #include <linux/mm.h>
 #include <linux/dma-map-ops.h>
 #include <linux/scatterlist.h>
@@ -307,6 +309,193 @@ void *dma_direct_alloc(struct device *dev, size_t size,
 	return NULL;
 }
 
+#ifdef CONFIG_DMA_LIVEUPDATE
+int dma_direct_preserve_allocation(struct device *dev, void *cpu_addr,
+				   size_t size, dma_addr_t dma_handle,
+				   unsigned long attrs, u64 *state)
+{
+	struct dma_alloc_ser *ser;
+	int ret;
+
+	if (!kho_is_enabled())
+		return -EOPNOTSUPP;
+
+	size = PAGE_ALIGN(size);
+	if (dma_is_from_cma(dma_to_phys(dev, dma_handle), size))
+		return -EOPNOTSUPP;
+
+	if (is_swiotlb_for_alloc(dev))
+		return -EOPNOTSUPP;
+
+	if ((attrs & DMA_ATTR_NO_KERNEL_MAPPING) &&
+	    !force_dma_unencrypted(dev))
+		return -EOPNOTSUPP;
+
+	if (IS_ENABLED(CONFIG_ARCH_HAS_DMA_ALLOC) &&
+	    !dev_is_dma_coherent(dev))
+		return -EOPNOTSUPP;
+
+	if (IS_ENABLED(CONFIG_DMA_GLOBAL_POOL) &&
+	    !dev_is_dma_coherent(dev))
+		return -EOPNOTSUPP;
+
+	if (IS_ENABLED(CONFIG_DMA_COHERENT_POOL) &&
+	    dma_is_from_pool(cpu_addr, size))
+		return -EOPNOTSUPP;
+
+	ser = kho_alloc_preserve(sizeof(*ser));
+	if (IS_ERR(ser))
+		return PTR_ERR(ser);
+
+	ser->page_phys = dma_to_phys(dev, dma_handle);
+	ser->force_decrypted = force_dma_unencrypted(dev);
+	ser->size = size;
+	ser->attrs = (attrs & ~DMA_ATTR_NO_WARN);
+
+	ret = kho_preserve_pages(phys_to_page(ser->page_phys),
+				 1 << get_order(size));
+	if (ret) {
+		kho_unpreserve_free(ser);
+		return ret;
+	}
+
+	*state = virt_to_phys(ser);
+	return 0;
+}
+
+void dma_direct_unpreserve_allocation(struct device *dev, u64 state)
+{
+	struct dma_alloc_ser *ser;
+
+	if (!kho_is_enabled())
+		return;
+
+	ser = phys_to_virt(state);
+	kho_unpreserve_pages(phys_to_page(ser->page_phys),
+			     1 << get_order(ser->size));
+	kho_unpreserve_free(ser);
+}
+
+void *dma_direct_restore_allocation(struct device *dev, size_t size,
+				    dma_addr_t *dma_handle, gfp_t gfp,
+				    unsigned long attrs, u64 state)
+{
+	bool remap = false, set_uncached = false;
+	struct dma_alloc_ser *ser = NULL;
+	struct page *page;
+	void *cpu_addr = NULL;
+	int i;
+
+	if (!kho_is_enabled())
+		return NULL;
+
+	ser = phys_to_virt(state);
+	page = phys_to_page(ser->page_phys);
+
+	size = PAGE_ALIGN(size);
+	if (size != ser->size)
+		return NULL;
+
+	WARN_ON_ONCE((attrs & ~DMA_ATTR_NO_WARN) != ser->attrs);
+
+	/*
+	 * Cannot free the restored pages on error here as these might be in use
+	 * by a device with direct allocation in the previous kernel.
+	 */
+	if (is_swiotlb_for_alloc(dev))
+		goto err;
+
+	if ((attrs & DMA_ATTR_NO_KERNEL_MAPPING) &&
+	    !force_dma_unencrypted(dev))
+		goto err;
+
+	if (!dev_is_dma_coherent(dev)) {
+		if (IS_ENABLED(CONFIG_ARCH_HAS_DMA_ALLOC))
+			goto err;
+
+		if (IS_ENABLED(CONFIG_DMA_GLOBAL_POOL))
+			goto err;
+
+		set_uncached = IS_ENABLED(CONFIG_ARCH_HAS_DMA_SET_UNCACHED);
+		remap = IS_ENABLED(CONFIG_DMA_DIRECT_REMAP);
+		if (!set_uncached && !remap)
+			goto err;
+	}
+
+	if (PageHighMem(page)) {
+		remap = true;
+		set_uncached = false;
+	}
+
+	/*
+	 * Remapping will be blocking so return error. The preserved memory
+	 * might be already decrypted in the previous kernel, but the decryption
+	 * call is not guaranteed to be non-blocking so return error always if
+	 * decryption is required.
+	 */
+	if ((remap || force_dma_unencrypted(dev)) &&
+	    dma_direct_use_pool(dev, gfp))
+		goto err;
+
+	/*
+	 * Encryption scheme changed between two kernels and this might cause
+	 * issues if device/driver is not handling it properly.
+	 */
+	WARN_ON_ONCE(ser->force_decrypted != force_dma_unencrypted(dev));
+
+	/*
+	 * arch_dma_prep_coherent() should make sure that any cache lines from
+	 * the previous kernel, if the device was coherent previously or cached
+	 * mapping in this kernel during init are not problematic for
+	 * non-coherent allocations.
+	 */
+	if (remap) {
+		pgprot_t prot = dma_pgprot(dev, PAGE_KERNEL, attrs);
+
+		if (force_dma_unencrypted(dev))
+			prot = pgprot_decrypted(prot);
+
+		arch_dma_prep_coherent(page, size);
+
+		cpu_addr = dma_common_contiguous_remap(page, size, prot,
+						       __builtin_return_address(0));
+		if (!cpu_addr)
+			goto err;
+	} else {
+		cpu_addr = page_address(page);
+		if (dma_set_decrypted(dev, cpu_addr, size)) {
+			cpu_addr = NULL;
+			goto err;
+		}
+	}
+
+	if (set_uncached) {
+		arch_dma_prep_coherent(page, size);
+		cpu_addr = arch_dma_set_uncached(cpu_addr, size);
+		if (IS_ERR(cpu_addr)) {
+			cpu_addr = NULL;
+			goto err;
+		}
+	}
+
+	*dma_handle = phys_to_dma_direct(dev, ser->page_phys);
+
+err:
+	WARN_ON(!kho_restore_pages(ser->page_phys,
+				   1 << get_order(ser->size)));
+	kho_restore_free(ser);
+
+	/*
+	 * DMA direct API expects high-order allocations, so set refcount of
+	 * tail pages to zero.
+	 */
+	for (i = (1 << get_order(size)) - 1; i > 0; --i)
+		set_page_count(page + i, 0);
+
+	return cpu_addr;
+}
+#endif
+
 void dma_direct_free(struct device *dev, size_t size,
 		void *cpu_addr, dma_addr_t dma_addr, unsigned long attrs)
 {
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 06/10] dma-mapping: Add API to preserve/restore DMA allocation

[Samiullah Khawaja]

Add new DMA APIs that allow preserving/restoring DMA allocations across
live update.

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

diff --git a/include/linux/dma-mapping.h b/include/linux/dma-mapping.h
index db8ab24a54f4..3756fc15467b 100644
--- a/include/linux/dma-mapping.h
+++ b/include/linux/dma-mapping.h
@@ -210,6 +210,15 @@ void *dma_vmap_noncontiguous(struct device *dev, size_t size,
 void dma_vunmap_noncontiguous(struct device *dev, void *vaddr);
 int dma_mmap_noncontiguous(struct device *dev, struct vm_area_struct *vma,
 		size_t size, struct sg_table *sgt);
+#ifdef CONFIG_DMA_LIVEUPDATE
+int dma_preserve_allocation_attrs(struct device *dev, void *cpu_addr,
+				  size_t size, dma_addr_t dma_handle,
+				  unsigned long attrs, u64 *state);
+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);
+#endif
 #else /* CONFIG_HAS_DMA */
 static inline dma_addr_t dma_map_page_attrs(struct device *dev,
 		struct page *page, size_t offset, size_t size,
@@ -496,6 +505,26 @@ static inline bool dma_need_unmap(struct device *dev)
 }
 #endif /* !CONFIG_HAS_DMA || !CONFIG_DMA_NEED_SYNC */
 
+#if !defined(CONFIG_DMA_LIVEUPDATE) || !defined(CONFIG_HAS_DMA)
+static inline int dma_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 dma_unpreserve_allocation(struct device *dev, u64 state)
+{
+}
+
+static inline void *dma_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,
 		dma_addr_t *dma_handle, enum dma_data_direction dir, gfp_t gfp);
 void dma_free_pages(struct device *dev, size_t size, struct page *page,
@@ -618,6 +647,27 @@ static inline void *dma_alloc_coherent(struct device *dev, size_t size,
 			(gfp & __GFP_NOWARN) ? DMA_ATTR_NO_WARN : 0);
 }
 
+static inline int dma_preserve_coherent_allocation(struct device *dev, void *cpu_addr,
+						   size_t size, dma_addr_t dma_handle, u64 *state)
+{
+	return dma_preserve_allocation_attrs(dev, cpu_addr, size,
+					     dma_handle, 0, state);
+}
+
+static inline void dma_unpreserve_coherent_allocation(struct device *dev, u64 state)
+{
+	dma_unpreserve_allocation(dev, state);
+}
+
+static inline void *dma_restore_coherent_allocation(struct device *dev, size_t size,
+						    dma_addr_t *dma_handle,
+						    gfp_t gfp, u64 state)
+{
+	return dma_restore_allocation_attrs(dev, size, dma_handle, gfp,
+					    (gfp & __GFP_NOWARN) ? DMA_ATTR_NO_WARN : 0,
+					    state);
+}
+
 static inline void dma_free_coherent(struct device *dev, size_t size,
 		void *cpu_addr, dma_addr_t dma_handle)
 {
diff --git a/kernel/dma/mapping.c b/kernel/dma/mapping.c
index 23ed8eb9233e..3b07bf30722c 100644
--- a/kernel/dma/mapping.c
+++ b/kernel/dma/mapping.c
@@ -12,6 +12,8 @@
 #include <linux/gfp.h>
 #include <linux/iommu-dma.h>
 #include <linux/kmsan.h>
+#include <linux/kexec_handover.h>
+#include <linux/kho/abi/dma_alloc.h>
 #include <linux/of_device.h>
 #include <linux/slab.h>
 #include <linux/vmalloc.h>
@@ -628,6 +630,53 @@ u64 dma_get_required_mask(struct device *dev)
 }
 EXPORT_SYMBOL_GPL(dma_get_required_mask);
 
+#ifdef CONFIG_DMA_LIVEUPDATE
+int dma_preserve_allocation_attrs(struct device *dev, void *cpu_addr,
+				  size_t size, dma_addr_t dma_handle,
+				  unsigned long attrs, u64 *state)
+{
+	const struct dma_map_ops *ops = get_dma_ops(dev);
+
+	if (dma_is_from_dev_coherent(dev, cpu_addr))
+		return -EOPNOTSUPP;
+
+	if (dma_alloc_direct(dev, ops))
+		return dma_direct_preserve_allocation(dev, cpu_addr, size,
+						      dma_handle, attrs,
+						      state);
+
+	return -EOPNOTSUPP;
+}
+EXPORT_SYMBOL(dma_preserve_allocation_attrs);
+
+void dma_unpreserve_allocation(struct device *dev, u64 state)
+{
+	const struct dma_map_ops *ops = get_dma_ops(dev);
+
+	if (dma_alloc_direct(dev, ops))
+		dma_direct_unpreserve_allocation(dev, state);
+}
+EXPORT_SYMBOL(dma_unpreserve_allocation);
+
+void *dma_restore_allocation_attrs(struct device *dev, size_t size,
+				   dma_addr_t *dma_handle, gfp_t gfp,
+				   unsigned long attrs, u64 state)
+{
+	const struct dma_map_ops *ops = get_dma_ops(dev);
+	void *cpu_addr = NULL;
+
+	WARN_ON_ONCE(!dev->coherent_dma_mask);
+
+	if (dma_alloc_direct(dev, ops))
+		cpu_addr = dma_direct_restore_allocation(dev, size, dma_handle,
+							 gfp, attrs, state);
+
+	debug_dma_alloc_coherent(dev, size, *dma_handle, cpu_addr, attrs);
+	return cpu_addr;
+}
+EXPORT_SYMBOL(dma_restore_allocation_attrs);
+#endif
+
 void *dma_alloc_attrs(struct device *dev, size_t size, dma_addr_t *dma_handle,
 		gfp_t flag, unsigned long attrs)
 {
-- 
2.55.0.795.g602f6c329a-goog

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

[Samiullah Khawaja]

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

[RFC PATCH v2 08/10] dma: contiguous: Export is_from_cma helper for kunit

[Samiullah Khawaja]

The is_from_cm() can be used to check if a memory allocation is done
from CMA. Export the helper when kunit is enabled so it can be used by
the kunit tests.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 kernel/dma/contiguous.c | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/kernel/dma/contiguous.c b/kernel/dma/contiguous.c
index 0b1702bf39ca..6d0cf9c7528c 100644
--- a/kernel/dma/contiguous.c
+++ b/kernel/dma/contiguous.c
@@ -37,6 +37,8 @@
 
 #define pr_fmt(fmt) "cma: " fmt
 
+#include <kunit/visibility.h>
+
 #include <asm/page.h>
 
 #include <linux/memblock.h>
@@ -153,6 +155,7 @@ bool dma_is_from_cma(phys_addr_t phys, size_t size)
 
 	return cma_for_each_area(dma_cma_check_area, &chk) != 0;
 }
+EXPORT_SYMBOL_IF_KUNIT(dma_is_from_cma);
 
 #ifdef CONFIG_DMA_NUMA_CMA
 
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 09/10] dma: pool: Export the is_from_pool helper for kunit

[Samiullah Khawaja]

When kunit is enabled, export the is_from_pool() helper function so it
can be used to check whether the allocations are from DMA pools.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 kernel/dma/pool.c | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/kernel/dma/pool.c b/kernel/dma/pool.c
index d1a0bb03e478..78496568973f 100644
--- a/kernel/dma/pool.c
+++ b/kernel/dma/pool.c
@@ -3,6 +3,9 @@
  * Copyright (C) 2012 ARM Ltd.
  * Copyright (C) 2020 Google LLC
  */
+
+#include <kunit/visibility.h>
+
 #include <linux/cma.h>
 #include <linux/debugfs.h>
 #include <linux/dma-map-ops.h>
@@ -320,3 +323,4 @@ bool dma_is_from_pool(void *start, size_t size)
 
 	return false;
 }
+EXPORT_SYMBOL_IF_KUNIT(dma_is_from_pool);
-- 
2.55.0.795.g602f6c329a-goog

[RFC PATCH v2 10/10] dma-direct: Add KUnit test for liveupdate preservation

[Samiullah Khawaja]

Use the KHO kunit stubs to mock the KHO preserve/unpreserve/restore API
and tests the preservation of DMA allocations with,

- Device being coherent or non-coherent.
- Atomic or non-atomic context.
- Various allocation sizes.
- Managed DMA allocations

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 kernel/dma/Kconfig       |   8 ++
 kernel/dma/Makefile      |   1 +
 kernel/dma/direct_test.c | 208 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 217 insertions(+)
 create mode 100644 kernel/dma/direct_test.c

diff --git a/kernel/dma/Kconfig b/kernel/dma/Kconfig
index d93958931303..48bc74af551a 100644
--- a/kernel/dma/Kconfig
+++ b/kernel/dma/Kconfig
@@ -269,3 +269,11 @@ config DMA_MAP_BENCHMARK
 config DMA_LIVEUPDATE
 	bool "Enable preservation of DMA direct allocations"
 	depends on LIVEUPDATE
+
+config DMA_LIVEUPDATE_KUNIT_TEST
+	tristate "KUnit test for DMA direct liveupdate preservation"
+	depends on KUNIT
+	depends on DMA_LIVEUPDATE
+	help
+	  This builds the KUnit tests for the DMA direct liveupdate
+	  preservation and restoration logic.
diff --git a/kernel/dma/Makefile b/kernel/dma/Makefile
index 6977033444a3..03703542ee79 100644
--- a/kernel/dma/Makefile
+++ b/kernel/dma/Makefile
@@ -10,3 +10,4 @@ obj-$(CONFIG_SWIOTLB)			+= swiotlb.o
 obj-$(CONFIG_DMA_COHERENT_POOL)		+= pool.o
 obj-$(CONFIG_MMU)			+= remap.o
 obj-$(CONFIG_DMA_MAP_BENCHMARK)		+= map_benchmark.o
+obj-$(CONFIG_DMA_LIVEUPDATE_KUNIT_TEST)	+= direct_test.o
diff --git a/kernel/dma/direct_test.c b/kernel/dma/direct_test.c
new file mode 100644
index 000000000000..f8403ec47807
--- /dev/null
+++ b/kernel/dma/direct_test.c
@@ -0,0 +1,208 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * KUnit test for DMA allocation live update preservation.
+ */
+
+#include <kunit/test.h>
+#include <linux/module.h>
+#include <kunit/static_stub.h>
+#include <linux/dma-map-ops.h>
+#include <linux/dma-mapping.h>
+#include <linux/dma-direct.h>
+#include <linux/kho/abi/dma_alloc.h>
+#include <linux/kexec_handover.h>
+#include <linux/slab.h>
+
+static const size_t dma_test_sizes[] = {
+	128,
+	PAGE_SIZE,
+	2 * PAGE_SIZE,
+	3 * PAGE_SIZE,
+	8 * PAGE_SIZE,
+};
+
+static void dma_size_desc(const size_t *size, char *desc)
+{
+	snprintf(desc, KUNIT_PARAM_DESC_SIZE, "size=%zu", *size);
+}
+
+KUNIT_ARRAY_PARAM(dma_size, dma_test_sizes, dma_size_desc);
+
+static struct page *mock_kho_restore_pages(phys_addr_t phys, unsigned long nr_pages)
+{
+	struct page *page = phys_to_page(phys);
+
+	if (!kho_test_pages_preserved(phys, nr_pages))
+		return NULL;
+
+	kho_unpreserve_pages(page, nr_pages);
+	return page;
+}
+
+static struct folio *mock_kho_restore_folio(phys_addr_t phys)
+{
+	struct folio *folio = page_folio(phys_to_page(phys));
+
+	if (!kho_test_pages_preserved(phys, (1 << folio_order(folio))))
+		return NULL;
+
+	kho_unpreserve_folio(folio);
+	return folio;
+}
+
+static void test_dma_direct_preserve_restore_common(struct kunit *test,
+						    bool coherent,
+						    bool atomic,
+						    size_t size)
+{
+	unsigned long nr_pages = 1 << get_order(size);
+	dma_addr_t handle1, handle2;
+	int ret, expected_ret = 0;
+	struct device dev = {0};
+	gfp_t gfp = GFP_KERNEL;
+	void *addr1, *addr2;
+	u64 state;
+
+	kunit_activate_static_stub(test, kho_restore_pages, mock_kho_restore_pages);
+	kunit_activate_static_stub(test, kho_restore_folio, mock_kho_restore_folio);
+
+	device_initialize(&dev);
+	dev.coherent_dma_mask = DMA_BIT_MASK(64);
+	dev.dma_mask = &dev.coherent_dma_mask;
+#if defined(CONFIG_ARCH_HAS_SYNC_DMA_FOR_DEVICE) || \
+    defined(CONFIG_ARCH_HAS_SYNC_DMA_FOR_CPU) || \
+    defined(CONFIG_ARCH_HAS_SYNC_DMA_FOR_CPU_ALL)
+	dev.dma_coherent = coherent;
+#endif
+
+	if (atomic)
+		gfp = GFP_ATOMIC;
+
+	addr1 = dma_alloc_coherent(&dev, size, &handle1, gfp);
+	if (!addr1) {
+		kunit_skip(test, "DMA allocation failed (unsupported configuration)");
+		return;
+	}
+
+	if (IS_ENABLED(CONFIG_DMA_COHERENT_POOL) && dma_is_from_pool(addr1, size)) {
+		expected_ret = -EOPNOTSUPP;
+		kunit_info(test, "DMA allocation using pool, expecting -EOPNOTSUPP");
+	} else if (dma_is_from_cma(dma_to_phys(&dev, handle1), size)) {
+		expected_ret = -EOPNOTSUPP;
+		kunit_info(test, "DMA allocation using CMA, expecting -EOPNOTSUPP");
+	}
+
+	ret = dma_preserve_coherent_allocation(&dev, addr1, size, handle1, &state);
+	KUNIT_ASSERT_EQ(test, ret, expected_ret);
+
+	if (!expected_ret) {
+		KUNIT_EXPECT_TRUE_MSG(test,
+				      kho_test_pages_preserved(dma_to_phys(&dev, handle1), nr_pages),
+				      "Allocated block not tracked in KHO");
+
+		addr2 = dma_restore_coherent_allocation(&dev, size, &handle2, gfp, state);
+		KUNIT_ASSERT_NOT_NULL(test, addr2);
+
+		KUNIT_EXPECT_EQ(test, handle1, handle2);
+		dma_free_coherent(&dev, size, addr2, handle2);
+		KUNIT_EXPECT_FALSE_MSG(test,
+				       kho_test_pages_preserved(dma_to_phys(&dev, handle1), nr_pages),
+				       "Allocated block still tracked after free");
+	} else {
+		dma_free_coherent(&dev, size, addr1, handle1);
+	}
+}
+
+static void test_dma_direct_coherent(struct kunit *test)
+{
+	const size_t *size = test->param_value;
+
+	test_dma_direct_preserve_restore_common(test, true, false, *size);
+}
+
+static void test_dma_direct_non_coherent(struct kunit *test)
+{
+	const size_t *size = test->param_value;
+
+	test_dma_direct_preserve_restore_common(test, false, false, *size);
+}
+
+static void test_dma_direct_coherent_atomic(struct kunit *test)
+{
+	const size_t *size = test->param_value;
+
+	test_dma_direct_preserve_restore_common(test, true, true, *size);
+}
+
+static void test_dma_direct_non_coherent_atomic(struct kunit *test)
+{
+	const size_t *size = test->param_value;
+
+	test_dma_direct_preserve_restore_common(test, false, true, *size);
+}
+
+static void test_dmam_preservation(struct kunit *test)
+{
+	dma_addr_t handle, non_dmam_handle;
+	void *addr, *non_dmam_addr;
+	struct device dev = {0};
+	size_t size = 128;
+	u64 state;
+	int ret;
+
+	kunit_activate_static_stub(test, kho_restore_pages, mock_kho_restore_pages);
+	kunit_activate_static_stub(test, kho_restore_folio, mock_kho_restore_folio);
+
+	device_initialize(&dev);
+	dev.coherent_dma_mask = DMA_BIT_MASK(64);
+	dev.dma_mask = &dev.coherent_dma_mask;
+#if defined(CONFIG_ARCH_HAS_SYNC_DMA_FOR_DEVICE) || \
+    defined(CONFIG_ARCH_HAS_SYNC_DMA_FOR_CPU) || \
+    defined(CONFIG_ARCH_HAS_SYNC_DMA_FOR_CPU_ALL)
+	dev.dma_coherent = true;
+#endif
+
+	non_dmam_addr = dma_alloc_coherent(&dev, size, &non_dmam_handle, GFP_KERNEL);
+	if (!non_dmam_addr) {
+		kunit_skip(test, "DMA allocation failed");
+		return;
+	}
+
+	ret = dmam_preserve_coherent_allocation(&dev, non_dmam_addr, size, non_dmam_handle, &state);
+	KUNIT_EXPECT_EQ_MSG(test, ret, -EINVAL, "Preserving non-devres allocation should fail");
+	dma_free_coherent(&dev, size, non_dmam_addr, non_dmam_handle);
+
+	addr = dmam_alloc_coherent(&dev, size, &handle, GFP_KERNEL);
+	if (!addr) {
+		kunit_skip(test, "DMA allocation failed");
+		return;
+	}
+
+	ret = dmam_preserve_coherent_allocation(&dev, addr, size, handle, &state);
+	KUNIT_ASSERT_EQ_MSG(test, ret, 0, "Failed to preserve dmam allocation");
+	KUNIT_EXPECT_TRUE(test, kho_test_pages_preserved(dma_to_phys(&dev, handle), 1));
+
+	dmam_unpreserve_coherent_allocation(&dev, addr, size, handle, state);
+	KUNIT_EXPECT_FALSE(test, kho_test_pages_preserved(dma_to_phys(&dev, handle), 1));
+
+	dmam_free_coherent(&dev, size, addr, handle);
+}
+
+static struct kunit_case dma_direct_test_cases[] = {
+	KUNIT_CASE_PARAM(test_dma_direct_coherent, dma_size_gen_params),
+	KUNIT_CASE_PARAM(test_dma_direct_non_coherent, dma_size_gen_params),
+	KUNIT_CASE_PARAM(test_dma_direct_coherent_atomic, dma_size_gen_params),
+	KUNIT_CASE_PARAM(test_dma_direct_non_coherent_atomic, dma_size_gen_params),
+	KUNIT_CASE(test_dmam_preservation),
+	{}
+};
+
+static struct kunit_suite dma_direct_test_suite = {
+	.name = "dma_direct_liveupdate",
+	.test_cases = dma_direct_test_cases,
+};
+kunit_test_suite(dma_direct_test_suite);
+
+MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");
+MODULE_DESCRIPTION("KUnit test for DMA direct live update preservation");
+MODULE_LICENSE("GPL");
-- 
2.55.0.795.g602f6c329a-goog