[RFC PATCH v2 5/5] vfio/pci: Add revocation fence for ZONE_DEVICE DMABUFs

[Pranjal Shrivastava <praan@google.com>]

Implement a synchronization fence to safely revoke ZONE_DEVICE-backed
DMABUFs. Introduce a fence in vfio_pci_dma_buf_set_status(). The fence
waits for all struct page refcounts to drop to 1.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 drivers/vfio/pci/vfio_pci_dmabuf.c | 55 ++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/drivers/vfio/pci/vfio_pci_dmabuf.c b/drivers/vfio/pci/vfio_pci_dmabuf.c
index b582e856ba7a..b4284b5cad03 100644
--- a/drivers/vfio/pci/vfio_pci_dmabuf.c
+++ b/drivers/vfio/pci/vfio_pci_dmabuf.c
@@ -4,6 +4,7 @@
 #include <linux/dma-buf-mapping.h>
 #include <linux/pci-p2pdma.h>
 #include <linux/dma-resv.h>
+#include <linux/iopoll.h>
 #include <linux/sched.h>
 #include <uapi/linux/dma-buf.h>
 
@@ -745,6 +746,44 @@ int vfio_pci_core_mmap_prep_dmabuf(struct vfio_pci_core_device *vdev,
 	return ret;
 }
 
+static void vfio_pci_zone_device_wait_fence(struct vfio_pci_dma_buf *priv)
+{
+	unsigned int i;
+
+	if (!priv->zone_device_backed)
+		return;
+
+	/*
+	 * Fence: Wait for any active references to the ZONE_DEVICE
+	 * pages to be dropped. A refcount of 1 represents the base
+	 * ownership.
+	 */
+	for (i = 0; i < priv->nr_ranges; i++) {
+		unsigned long pfn = priv->phys_vec[i].paddr >> PAGE_SHIFT;
+		unsigned long npgs = PAGE_ALIGN(priv->phys_vec[i].len) >> PAGE_SHIFT;
+
+		while (npgs--) {
+			struct page *page = pfn_to_page(pfn++);
+			int count, ret;
+
+			/*
+			 * Poll page_count() and block indefinitely until all
+			 * refs drop to avoid DMA-after-free.
+			 */
+			do {
+				ret = read_poll_timeout(page_count, count,
+							(count == 1),
+							1000, 10000000,
+							false, page);
+				if (ret)
+					dev_warn(&priv->vdev->pdev->dev,
+						 "Waiting for GUP pins to drop on PFN 0x%lx... (importer hung?)\n",
+						 pfn - 1);
+			} while (ret);
+		}
+	}
+}
+
 /* Set the DMABUF's revocation status (OK or temporarily/permanently revoked) */
 static void vfio_pci_dma_buf_set_status(struct vfio_pci_dma_buf *priv,
 					enum vfio_pci_dma_buf_status new_status)
@@ -779,8 +818,24 @@ static void vfio_pci_dma_buf_set_status(struct vfio_pci_dma_buf *priv,
 		dma_resv_unlock(priv->dmabuf->resv);
 		kref_put(&priv->kref, vfio_pci_dma_buf_done);
 		wait_for_completion(&priv->comp);
+
+		/*
+		 * Note: Rmap Deadlocks
+		 * unmap_mapping_range() is safe to call here within memory_lock
+		 * despite the VMA being VM_MIXEDMAP. Because our ZONE_DEVICE pages
+		 * are allocated via devm_memremap_pages(), page->mapping is never
+		 * set which makes them invisible to rmap.
+		 *
+		 * If this changes in the future, this call must be factored outside
+		 * the memory_lock to prevent a 3-way circular deadlock:
+		 * (mmap_lock -> memory_lock -> i_mmap_rwsem).
+		 */
 		unmap_mapping_range(priv->dmabuf->file->f_mapping,
 				    0, 0, true);
+
+		/* Wait for all page refs to drop if ZONE_DEVICE registered */
+		vfio_pci_zone_device_wait_fence(priv);
+
 		/*
 		 * Re-arm the registered kref reference and the
 		 * completion so the post-revoke state matches the
-- 
2.55.0.571.g244d577d93-goog