In preparation for wp to support mthp, add the kmsan_copy_pages_meta()
function to copy multiple pages of the source page to the target page.

Signed-off-by: Vernon Yang <vernon2gm@gmail.com>
---
 include/linux/kmsan.h | 13 ++++++++++---
 mm/kmsan/shadow.c     | 26 +++++++++++++++++++-------
 mm/memory.c           |  2 +-
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/include/linux/kmsan.h b/include/linux/kmsan.h
index 2b1432cc16d5..a3f227c3947f 100644
--- a/include/linux/kmsan.h
+++ b/include/linux/kmsan.h
@@ -78,15 +78,16 @@ void kmsan_alloc_page(struct page *page, unsigned int order, gfp_t flags);
 void kmsan_free_page(struct page *page, unsigned int order);
 
 /**
- * kmsan_copy_page_meta() - Copy KMSAN metadata between two pages.
+ * kmsan_copy_pages_meta() - Copy KMSAN metadata between two pages.
  * @dst: destination page.
  * @src: source page.
+ * @nr_pages: copy number of page.
  *
  * KMSAN copies the contents of metadata pages for @src into the metadata pages
  * for @dst. If @dst has no associated metadata pages, nothing happens.
  * If @src has no associated metadata pages, @dst metadata pages are unpoisoned.
  */
-void kmsan_copy_page_meta(struct page *dst, struct page *src);
+void kmsan_copy_pages_meta(struct page *dst, struct page *src, int nr_pages);
 
 /**
  * kmsan_slab_alloc() - Notify KMSAN about a slab allocation.
@@ -324,7 +325,8 @@ static inline void kmsan_free_page(struct page *page, unsigned int order)
 {
 }
 
-static inline void kmsan_copy_page_meta(struct page *dst, struct page *src)
+static inline void kmsan_copy_pages_meta(struct page *dst, struct page *src,
+					int nr_pages)
 {
 }
 
@@ -407,4 +409,9 @@ static inline void *memset_no_sanitize_memory(void *s, int c, size_t n)
 
 #endif
 
+static inline void kmsan_copy_page_meta(struct page *dst, struct page *src)
+{
+	kmsan_copy_pages_meta(dst, src, 1);
+}
+
 #endif /* _LINUX_KMSAN_H */
diff --git a/mm/kmsan/shadow.c b/mm/kmsan/shadow.c
index 54f3c3c962f0..1dd0f7a1eb5f 100644
--- a/mm/kmsan/shadow.c
+++ b/mm/kmsan/shadow.c
@@ -148,24 +148,36 @@ void *kmsan_get_metadata(void *address, bool is_origin)
 	return (is_origin ? origin_ptr_for(page) : shadow_ptr_for(page)) + off;
 }
 
-void kmsan_copy_page_meta(struct page *dst, struct page *src)
+
+void kmsan_copy_pages_meta(struct page *dst, struct page *src, int nr_pages)
 {
+	int i;
+
 	if (!kmsan_enabled || kmsan_in_runtime())
 		return;
-	if (!dst || !page_has_metadata(dst))
+
+	for (i = 0; i < nr_pages; i++) {
+		if (!dst || !page_has_metadata(dst))
+			break;
+		if (!src || !page_has_metadata(src))
+			break;
+	}
+
+	if (i == 0 && !dst) {
 		return;
-	if (!src || !page_has_metadata(src)) {
-		kmsan_internal_unpoison_memory(page_address(dst), PAGE_SIZE,
+	} else if (i < nr_pages) {
+		kmsan_internal_unpoison_memory(page_address(dst),
+					       nr_pages * PAGE_SIZE,
 					       /*checked*/ false);
 		return;
 	}
 
 	kmsan_enter_runtime();
-	__memcpy(shadow_ptr_for(dst), shadow_ptr_for(src), PAGE_SIZE);
-	__memcpy(origin_ptr_for(dst), origin_ptr_for(src), PAGE_SIZE);
+	__memcpy(shadow_ptr_for(dst), shadow_ptr_for(src), nr_pages * PAGE_SIZE);
+	__memcpy(origin_ptr_for(dst), origin_ptr_for(src), nr_pages * PAGE_SIZE);
 	kmsan_leave_runtime();
 }
-EXPORT_SYMBOL(kmsan_copy_page_meta);
+EXPORT_SYMBOL(kmsan_copy_pages_meta);
 
 void kmsan_alloc_page(struct page *page, unsigned int order, gfp_t flags)
 {
diff --git a/mm/memory.c b/mm/memory.c
index 90cbed5ad150..7b8c7d0f9ff4 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -3589,7 +3589,7 @@ static vm_fault_t wp_page_copy(struct vm_fault *vmf)
 			delayacct_wpcopy_end();
 			return err == -EHWPOISON ? VM_FAULT_HWPOISON : 0;
 		}
-		kmsan_copy_page_meta(&new_folio->page, vmf->page);
+		kmsan_copy_pages_meta(&new_folio->page, vmf->page, 1);
 	}
 
 	__folio_mark_uptodate(new_folio);
-- 
2.50.1