[PATCH] drm/gpusvm, pagemap: Do not assume DRM pagemap owns device pages

[Matthew Brost <matthew.brost@intel.com>]

Update drm_pagemap_page_zone_device_data() to derive the pgmap ops from
the page and compare them against the DRM pagemap ops. If the ops do not
match, return NULL.

Also harden two risky call sites by checking for NULL after
hmm_range_fault() or migrate_vma_setup() when migrating to device
memory, as it is possible to encounter device pages that are not owned
by DRM pagemap.

Suggested-by: sashiko.dev
Signed-off-by: Matthew Brost <matthew.brost@intel.com>
---
 drivers/gpu/drm/drm_gpusvm.c  |  5 +++++
 drivers/gpu/drm/drm_pagemap.c | 14 ++++++++++----
 include/drm/drm_pagemap.h     |  5 ++++-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/drivers/gpu/drm/drm_gpusvm.c b/drivers/gpu/drm/drm_gpusvm.c
index 365a9c0b522a..b3cccd047a21 100644
--- a/drivers/gpu/drm/drm_gpusvm.c
+++ b/drivers/gpu/drm/drm_gpusvm.c
@@ -1506,6 +1506,11 @@ int drm_gpusvm_get_pages(struct drm_gpusvm *gpusvm,
 			struct drm_pagemap_zdd *__zdd =
 				drm_pagemap_page_zone_device_data(page);
 
+			if (!__zdd) {
+				err = -EINVAL;
+				goto err_unmap;
+			}
+
 			if (!ctx->allow_mixed &&
 			    zdd != __zdd && i > 0) {
 				err = -EOPNOTSUPP;
diff --git a/drivers/gpu/drm/drm_pagemap.c b/drivers/gpu/drm/drm_pagemap.c
index d82ea7ccb8da..95c951c5b569 100644
--- a/drivers/gpu/drm/drm_pagemap.c
+++ b/drivers/gpu/drm/drm_pagemap.c
@@ -753,10 +753,16 @@ int drm_pagemap_migrate_to_devmem(struct drm_pagemap_devmem *devmem_allocation,
 				own_pages++;
 				goto next;
 			}
-			cur.dpagemap = src_zdd->dpagemap;
-			cur.ops = src_zdd->devmem_allocation->ops;
-			cur.device = cur.dpagemap->drm->dev;
-			pages[i] = src_page;
+			if (src_zdd) {
+				cur.dpagemap = src_zdd->dpagemap;
+				cur.ops = src_zdd->devmem_allocation->ops;
+				cur.device = cur.dpagemap->drm->dev;
+				pages[i] = src_page;
+			} else {
+				npages = i;
+				err = -EINVAL;
+				goto err_finalize;
+			}
 		}
 		if (!pages[i]) {
 			cur.dpagemap = NULL;
diff --git a/include/drm/drm_pagemap.h b/include/drm/drm_pagemap.h
index 95eb4b66b057..9b7c50932db5 100644
--- a/include/drm/drm_pagemap.h
+++ b/include/drm/drm_pagemap.h
@@ -367,12 +367,15 @@ int drm_pagemap_reinit(struct drm_pagemap *dpagemap);
  * drm_pagemap_page_zone_device_data() - Page to zone_device_data
  * @page: Pointer to the page
  *
- * Return: Page's zone_device_data
+ * Return: Page's zone_device_data if owned by DRM pagemap, NULL otherwise
  */
 static inline struct drm_pagemap_zdd *drm_pagemap_page_zone_device_data(struct page *page)
 {
 	struct folio *folio = page_folio(page);
 
+	if (WARN_ON_ONCE(page_pgmap(page)->ops != drm_pagemap_pagemap_ops_get()))
+		return NULL;
+
 	return folio_zone_device_data(folio);
 }
 
-- 
2.34.1