[RFC: DMA_PMD 14/22] iommu/dma: Support decrypted and pinned DMA_PMD pages

@@ -53,7 +53,7 @@ static void test_meta_invalid_phys(struct kunit *test)
 
 static void dma_pmd_meta_set_pooled(struct dma_pmd_meta *m, bool pooled)
 {
-	WRITE_ONCE(m->pooled, pooled);
+	WRITE_ONCE(m->flags, pooled ? DMA_PMD_POOLED : 0);
 }
 
 static void test_meta_pooled_toggle(struct kunit *test)

@@ -8,6 +8,7 @@
 #include <linux/bitmap.h>
 #include <linux/cache.h>
 #include <linux/cacheflush.h>
+#include <linux/cc_platform.h>
 #include <linux/dma-pmd.h>
 #include <linux/export.h>
 #include <linux/gfp.h>

@@ -64,10 +65,22 @@ static __always_inline struct dma_pmd_meta *__dma_pmd_meta_of_pfn(unsigned long
 
 bool __dma_is_pmd_page(unsigned long pfn)
 {
-	return READ_ONCE(__dma_pmd_meta_of_pfn(pfn)->pooled);
+	return READ_ONCE(__dma_pmd_meta_of_pfn(pfn)->flags) & DMA_PMD_POOLED;
 }
 EXPORT_SYMBOL(__dma_is_pmd_page);
 
+bool __dma_is_pmd_direct(unsigned long pfn)
+{
+	u8 flags = READ_ONCE(__dma_pmd_meta_of_pfn(pfn)->flags);
+
+	if (cc_platform_has(CC_ATTR_GUEST_MEM_ENCRYPT) &&
+	    !(flags & DMA_PMD_DECRYPTED))
+		return false;
+	return (flags & DMA_PMD_POOLED) &&
+	       (flags & (DMA_PMD_DECRYPTED | DMA_PMD_PINNED));
+}
+EXPORT_SYMBOL(__dma_is_pmd_direct);
+
 /**
  * dma_pmd_meta_of_pfn - Metadata for a PFN known to be used by DMA_PMD.
  * @pfn: PFN the caller already holds a struct page for

@@ -9,6 +9,7 @@
 #include <linux/bitmap.h>
 #include <linux/bitops.h>
 #include <linux/cache.h>
+#include <linux/cc_platform.h>
 #include <linux/debugfs.h>
 #include <linux/dma-mapping.h>
 #include <linux/dma-pmd.h>

@@ -22,6 +23,7 @@
 #include <linux/mutex.h>
 #include <linux/rcupdate.h>
 #include <linux/seq_file.h>
+#include <linux/set_memory.h>
 #include <linux/shrinker.h>
 #include <linux/slab.h>
 #include <linux/spinlock.h>

@@ -39,6 +41,14 @@ static LLIST_HEAD(dma_pmd_free_list);
 /* Dedicated WQ_MEM_RECLAIM workqueue, can run without allocations. */
 static struct workqueue_struct *dma_pmd_wq __ro_after_init;
 
+/* Mark pooled PMD pages decrypted for DMA (auto-enabled on CoCo guests). */
+static bool dma_pmd_decrypt __read_mostly;
+module_param_named(decrypt, dma_pmd_decrypt, bool, 0644);
+
+/* Pin pooled PMD pages in the host hypervisor on allocation and unpin on release. */
+static bool dma_pmd_pin __read_mostly;
+module_param_named(pin, dma_pmd_pin, bool, 0644);
+
 /*
  * Running total and global ceiling on DMA_PMD pages pinned across all pools.
  * If zero, it is set at boot at 1/8 of total RAM.

@@ -65,6 +75,70 @@ static void dma_pmd_pool_free_kref(struct kref *kref)
 	dma_pmd_schedule_reclaim();
 }
 
+static int dma_pmd_page_prepare(struct page *page, struct dma_pmd_meta *meta,
+				bool can_block)
+{
+	bool want_decrypt = READ_ONCE(dma_pmd_decrypt);
+	bool want_pin = READ_ONCE(dma_pmd_pin);
+	void *vaddr = page_address(page);
+	int ret;
+
+	if (!want_decrypt && !want_pin)
+		return 0;
+
+	/*
+	 * Decryption and host-pinning hypercalls may sleep. In atomic context,
+	 * leave DMA_PMD_DECRYPTED/PINNED unset so the page is still pooled but
+	 * uses normal swiotlb bounce buffering when required.
+	 */
+	if (!can_block) {
+		pr_warn_ratelimited("%s: cannot block but want %s %s\n", __func__,
+				    want_decrypt ? "decrypt" : "", want_pin ? "pin" : "");
+		return 0;
+	}
+
+	/*
+	 * Explicit page decryption is for CoCo guests (where force_dma_unencrypted()
+	 * is true), not bare-metal SME hosts where 64-bit devices and the IOMMU
+	 * set the encryption bit in DMA addresses and IOMMU PTEs.
+	 */
+	if (want_decrypt && !WARN_ON_ONCE(cc_platform_has(CC_ATTR_HOST_MEM_ENCRYPT))) {
+		ret = set_memory_decrypted((unsigned long)vaddr, 1 << PMD_ORDER);
+		if (ret)
+			return ret;
+		/* Initialize cleartext cachelines after flipping encryption state. */
+		memset(vaddr, 0, PMD_SIZE);
+		meta->flags |= DMA_PMD_DECRYPTED;
+	}
+
+	if (want_pin) {
+		/* Hook point for hypervisor GPA range pinning. */
+		meta->flags |= DMA_PMD_PINNED;
+	}
+
+	return 0;
+}
+
+static int dma_pmd_page_unprepare(struct dma_pmd_meta *meta)
+{
+	struct page *page = pfn_to_page(dma_pmd_meta_to_pfn(meta));
+
+	if (meta->flags & DMA_PMD_PINNED) {
+		/* Hook point for hypervisor GPA range unpinning. */
+		WRITE_ONCE(meta->flags, meta->flags & ~DMA_PMD_PINNED);
+	}
+
+	if (meta->flags & DMA_PMD_DECRYPTED) {
+		if (set_memory_encrypted((unsigned long)page_address(page),
+					 1 << PMD_ORDER)) {
+			pr_warn_ratelimited("dma_pmd: leaking PMD page that could not be re-encrypted\n");
+			return -EIO;
+		}
+		WRITE_ONCE(meta->flags, meta->flags & ~DMA_PMD_DECRYPTED);
+	}
+	return 0;
+}
+
 /*
  * Releasing an idle PMD page proceeds in three stages:
  *

@@ -74,7 +148,7 @@ static void dma_pmd_pool_free_kref(struct kref *kref)
  *    pushed locklessly onto dma_pmd_free_list and dma_pmd_reclaim_work is
  *    scheduled on dma_pmd_wq.
  * 2. In process context, dma_pmd_release_page() unmaps all IOMMU domains
- *    (dma_pmd_unmap_all()), clears meta->pooled so new lockless readers stop
+ *    (dma_pmd_unmap_all()), clears DMA_PMD_POOLED so new lockless readers stop
  *    entering @meta, and queues dma_pmd_release_page_rcu() via call_rcu().
  * 3. After an RCU grace period (once no concurrent dma_pmd_free_page() reader
  *    can still be dereferencing @meta), dma_pmd_release_page_rcu() unfreezes

@@ -97,22 +171,21 @@ static void dma_pmd_release_page_rcu(struct rcu_head *head)
 	 * block is freed, the PMD frame can be reallocated and @meta overwritten.
 	 */
 	struct dma_pmd_meta *meta = container_of(head, struct dma_pmd_meta, rcu);
-	struct page *page = pfn_to_page(dma_pmd_meta_to_pfn(meta));
+	struct page *block, *page = pfn_to_page(dma_pmd_meta_to_pfn(meta));
 	struct dma_pmd_pool *pool = READ_ONCE(meta->pool);
-	unsigned int nr = DMA_PMD_BLOCKS(pool->order);
-	unsigned int order = pool->order;
+	unsigned int order = pool->order, i, nr = DMA_PMD_BLOCKS(order);
+	bool leak = READ_ONCE(meta->flags) & DMA_PMD_DECRYPTED;
 	unsigned long flags;
-	unsigned int i;
-
-	for (i = 0; i < nr; i++) {
-		struct page *block = page + (i << order);
 
-		page_ref_unfreeze(block, 1);
-		__free_pages(block, order);
+	if (!leak) {
+		for (i = 0; i < nr; i++) {
+			block = page + (i << order);
+			page_ref_unfreeze(block, 1);
+			__free_pages(block, order);
+		}
+		atomic_long_dec(&dma_pmd_nr_pages);
 	}
 
-	atomic_long_dec(&dma_pmd_nr_pages);
-
 	spin_lock_irqsave(&pool->lock, flags);
 	pool->pmd_free_cnt++;
 	spin_unlock_irqrestore(&pool->lock, flags);

@@ -138,7 +211,9 @@ static void dma_pmd_release_page(struct dma_pmd_meta *meta)
 	 * dma_is_pmd_page() just before this may still be reading @meta, so
 	 * the PMD page and its metadata entry must outlive the grace period.
 	 */
-	WRITE_ONCE(meta->pooled, false);
+	WRITE_ONCE(meta->flags, meta->flags & ~DMA_PMD_POOLED);
+	dma_pmd_page_unprepare(meta);
+
 	call_rcu(&meta->rcu, dma_pmd_release_page_rcu);
 }
 

@@ -392,20 +467,33 @@ static struct dma_pmd_meta *dma_pmd_add_page(struct dma_pmd_pool *pool,
 		return NULL;
 	}
 
-	/* Fails with -EBUSY if a concurrent PFN walker holds a speculative ref. */
-	if (split_page_compound(page, PMD_ORDER, pool->order)) {
+	meta = dma_pmd_meta_of_pfn(page_to_pfn(page));
+	memset(meta, 0, sizeof(*meta));
+	meta->pool = pool;
+	spin_lock_init(&meta->map_lock);
+
+	if (dma_pmd_page_prepare(page, meta, can_block)) {
 		__free_pages(page, PMD_ORDER);
 		atomic_long_dec(&dma_pmd_nr_pages);
 		return NULL;
 	}
 
-	meta = dma_pmd_meta_of_pfn(page_to_pfn(page));
-	memset(meta, 0, sizeof(*meta));
-	meta->pool = pool;
-	spin_lock_init(&meta->map_lock);
+	/* Fails with -EBUSY if a concurrent PFN walker holds a speculative ref. */
+	if (split_page_compound(page, PMD_ORDER, pool->order)) {
+		if (!dma_pmd_page_unprepare(meta)) {
+			__free_pages(page, PMD_ORDER);
+			atomic_long_dec(&dma_pmd_nr_pages);
+		}
+		return NULL;
+	}
 
-	bitmap_set(meta->dirty_bitmap, 0, nr);
-	meta->nr_dirty = nr;
+	if (meta->flags & DMA_PMD_DECRYPTED) {
+		bitmap_set(meta->free_bitmap, 0, nr);
+		meta->nr_free = nr;
+	} else {
+		bitmap_set(meta->dirty_bitmap, 0, nr);
+		meta->nr_dirty = nr;
+	}
 
 	kref_get(&pool->refcount);
 

@@ -418,7 +506,7 @@ static struct dma_pmd_meta *dma_pmd_add_page(struct dma_pmd_pool *pool,
 	 * function returns and the caller publishes the PMD page under
 	 * @pool->lock, whose release orders both stores against every consumer.
 	 */
-	WRITE_ONCE(meta->pooled, true);
+	WRITE_ONCE(meta->flags, meta->flags | DMA_PMD_POOLED);
 
 	return meta;
 }

@@ -764,10 +852,18 @@ static unsigned long dma_pmd_shrink_scan(struct shrinker *shrink,
 	/*
 	 * Retiring is done outside both locks: dma_pmd_release_page() unmaps
 	 * from the IOMMU, which is far too long to hold pool->lock for.
+	 * Pages that need set_memory_encrypted() cannot be released under
+	 * fs_reclaim (CPA allocates PTE pages on x86); hand them to the
+	 * reclaim workqueue instead.
 	 */
 	list_for_each_entry_safe(meta, tmp, &release_list, list) {
-		list_del(&meta->list);
-		dma_pmd_release_page(meta);
+		list_del_init(&meta->list);
+		if (meta->flags & (DMA_PMD_DECRYPTED | DMA_PMD_PINNED)) {
+			llist_add(&meta->llnode, &dma_pmd_free_list);
+			dma_pmd_schedule_reclaim();
+		} else {
+			dma_pmd_release_page(meta);
+		}
 	}
 	srcu_read_unlock(&dma_pmd_srcu, srcu_idx);
 

@@ -806,6 +902,9 @@ static int __init dma_pmd_init(void)
 	if (!dma_pmd_max_pages)
 		dma_pmd_max_pages = max(totalram_pages() >> (PMD_ORDER + 3), 16UL);
 
+	if (cc_platform_has(CC_ATTR_GUEST_MEM_ENCRYPT))
+		dma_pmd_decrypt = true;
+
 	/*
 	 * Reclaim frees memory, so it must not queue behind arbitrary work on
 	 * system_wq when the machine is already short of it. Failure is not

@@ -107,10 +107,16 @@ enum {
 #define DMA_PMD_META_SHIFT		8
 #define DMA_PMD_META_SIZE		BIT(DMA_PMD_META_SHIFT)
 
+enum {
+	DMA_PMD_POOLED		= BIT(0),
+	DMA_PMD_DECRYPTED	= BIT(1),
+	DMA_PMD_PINNED		= BIT(2),
+};
+
 /**
  * struct dma_pmd_meta - Metadata for a single DMA_PMD page
- * @pooled:	True while this page is owned by an dma_pmd_pool (offset 0)
- *		read locklessly by dma_is_pmd_page().
+ * @flags:	Bitmask of DMA_PMD_POOLED, DMA_PMD_DECRYPTED, DMA_PMD_PINNED
+ *		(offset 0), read locklessly by dma_is_pmd_page() and dma_is_pmd_direct().
  * @nr_free:	Number of zeroed available blocks in @free_bitmap
  * @nr_dirty:	Number of dirty available blocks in @dirty_bitmap
  * @map_lock:	Spinlock protecting slow-path updates to @domains_mapped

@@ -138,7 +144,7 @@ enum {
  * physical pages (128 KB per GB of RAM).
  */
 struct dma_pmd_meta {
-	bool pooled;
+	u8 flags;
 	u16 nr_free;
 	u16 nr_dirty;
 	spinlock_t map_lock;

@@ -160,7 +166,7 @@ struct dma_pmd_meta {
 	DECLARE_BITMAP(dirty_bitmap, 1U << PMD_ORDER) ____cacheline_aligned;
 } __aligned(DMA_PMD_META_SIZE);
 
-static_assert(offsetof(struct dma_pmd_meta, pooled) == 0);
+static_assert(offsetof(struct dma_pmd_meta, flags) == 0);
 static_assert(sizeof(struct dma_pmd_meta) == DMA_PMD_META_SIZE);
 
 static inline unsigned int dma_pmd_meta_avail(const struct dma_pmd_meta *m)

@@ -16,6 +16,7 @@ struct dma_pmd_pool;
 extern void *dma_pmd_meta_array;
 
 bool __dma_is_pmd_page(unsigned long pfn);
+bool __dma_is_pmd_direct(unsigned long pfn);
 
 static inline bool dma_is_pmd_page(unsigned long pfn)
 {

@@ -29,6 +30,15 @@ static inline bool dma_is_pmd_page(unsigned long pfn)
 	return __dma_is_pmd_page(pfn);
 }
 
+static inline bool dma_is_pmd_direct(phys_addr_t phys)
+{
+	/* Pairs with smp_store_release() in dma_pmd_meta_init(). */
+	if (likely(!smp_load_acquire(&dma_pmd_meta_array)))
+		return false;
+
+	return __dma_is_pmd_direct(phys >> PAGE_SHIFT);
+}
+
 bool __dma_pmd_free_page(struct page *page);
 
 static inline bool dma_pmd_free_page(struct page *page)

@@ -117,6 +127,11 @@ static inline bool dma_is_pmd_page(unsigned long pfn)
 	return false;
 }
 
+static inline bool dma_is_pmd_direct(phys_addr_t phys)
+{
+	return false;
+}
+
 static inline bool dma_pmd_free_page(struct page *page)
 {
 	return false;

@@ -7,6 +7,7 @@
 #include <linux/memblock.h> /* for max_pfn */
 #include <linux/export.h>
 #include <linux/mm.h>
+#include <linux/dma-pmd.h>
 #include <linux/dma-map-ops.h>
 #include <linux/scatterlist.h>
 #include <linux/pfn.h>

@@ -679,7 +680,10 @@ dma_addr_t dma_direct_map_phys(struct device *dev, phys_addr_t phys,
 			attrs |= DMA_ATTR_CC_SHARED;
 	}
 
-	if (is_swiotlb_force_bounce(dev)) {
+	if (dma_is_pmd_direct(phys)) {
+		if (force_dma_unencrypted(dev))
+			attrs |= DMA_ATTR_CC_SHARED;
+	} else if (is_swiotlb_force_bounce(dev)) {
 		if (attrs & (DMA_ATTR_MMIO | DMA_ATTR_REQUIRE_COHERENT))
 			return DMA_MAPPING_ERROR;