1 files changed, 12 insertions, 49 deletions

diff --git a/drivers/virt/acrn/mm.c b/drivers/virt/acrn/mm.c
index 9c75de0656..fa5d9ca6be 100644
--- a/drivers/virt/acrn/mm.c
+++ b/drivers/virt/acrn/mm.c
@@ -155,83 +155,43 @@ int acrn_vm_memseg_unmap(struct acrn_vm *vm, struct acrn_vm_memmap *memmap)
 int acrn_vm_ram_map(struct acrn_vm *vm, struct acrn_vm_memmap *memmap)
 {
 	struct vm_memory_region_batch *regions_info;
-	int nr_pages, i, order, nr_regions = 0;
+	int nr_pages, i = 0, order, nr_regions = 0;
 	struct vm_memory_mapping *region_mapping;
 	struct vm_memory_region_op *vm_region;
 	struct page **pages = NULL, *page;
 	void *remap_vaddr;
 	int ret, pinned;
 	u64 user_vm_pa;
+	unsigned long pfn;
 	struct vm_area_struct *vma;
 
 	if (!vm || !memmap)
 		return -EINVAL;
 
-	/* Get the page number of the map region */
-	nr_pages = memmap->len >> PAGE_SHIFT;
-	if (!nr_pages)
-		return -EINVAL;
-
 	mmap_read_lock(current->mm);
 	vma = vma_lookup(current->mm, memmap->vma_base);
 	if (vma && ((vma->vm_flags & VM_PFNMAP) != 0)) {
-		unsigned long start_pfn, cur_pfn;
-		spinlock_t *ptl;
-		bool writable;
-		pte_t *ptep;
-
 		if ((memmap->vma_base + memmap->len) > vma->vm_end) {
 			mmap_read_unlock(current->mm);
 			return -EINVAL;
 		}
 
-		for (i = 0; i < nr_pages; i++) {
-			ret = follow_pte(vma->vm_mm,
-					 memmap->vma_base + i * PAGE_SIZE,
-					 &ptep, &ptl);
-			if (ret)
-				break;
-
-			cur_pfn = pte_pfn(ptep_get(ptep));
-			if (i == 0)
-				start_pfn = cur_pfn;
-			writable = !!pte_write(ptep_get(ptep));
-			pte_unmap_unlock(ptep, ptl);
-
-			/* Disallow write access if the PTE is not writable. */
-			if (!writable &&
-			    (memmap->attr & ACRN_MEM_ACCESS_WRITE)) {
-				ret = -EFAULT;
-				break;
-			}
-
-			/* Disallow refcounted pages. */
-			if (pfn_valid(cur_pfn) &&
-			    !PageReserved(pfn_to_page(cur_pfn))) {
-				ret = -EFAULT;
-				break;
-			}
-
-			/* Disallow non-contiguous ranges. */
-			if (cur_pfn != start_pfn + i) {
-				ret = -EINVAL;
-				break;
-			}
-		}
+		ret = follow_pfn(vma, memmap->vma_base, &pfn);
 		mmap_read_unlock(current->mm);
-
-		if (ret) {
+		if (ret < 0) {
 			dev_dbg(acrn_dev.this_device,
 				"Failed to lookup PFN at VMA:%pK.\n", (void *)memmap->vma_base);
 			return ret;
 		}
 
 		return acrn_mm_region_add(vm, memmap->user_vm_pa,
-			 PFN_PHYS(start_pfn), memmap->len,
+			 PFN_PHYS(pfn), memmap->len,
 			 ACRN_MEM_TYPE_WB, memmap->attr);
 	}
 	mmap_read_unlock(current->mm);
 
+	/* Get the page number of the map region */
+	nr_pages = memmap->len >> PAGE_SHIFT;
 	pages = vzalloc(array_size(nr_pages, sizeof(*pages)));
 	if (!pages)
 		return -ENOMEM;
@@ -275,11 +235,12 @@ int acrn_vm_ram_map(struct acrn_vm *vm, struct acrn_vm_memmap *memmap)
 	mutex_unlock(&vm->regions_mapping_lock);
 
 	/* Calculate count of vm_memory_region_op */
-	for (i = 0; i < nr_pages; i += 1 << order) {
+	while (i < nr_pages) {
 		page = pages[i];
 		VM_BUG_ON_PAGE(PageTail(page), page);
 		order = compound_order(page);
 		nr_regions++;
+		i += 1 << order;
 	}
 
 	/* Prepare the vm_memory_region_batch */
@@ -296,7 +257,8 @@ int acrn_vm_ram_map(struct acrn_vm *vm, struct acrn_vm_memmap *memmap)
 	regions_info->vmid = vm->vmid;
 	regions_info->regions_gpa = virt_to_phys(vm_region);
 	user_vm_pa = memmap->user_vm_pa;
-	for (i = 0; i < nr_pages; i += 1 << order) {
+	i = 0;
+	while (i < nr_pages) {
 		u32 region_size;
 
 		page = pages[i];
@@ -312,6 +274,7 @@ int acrn_vm_ram_map(struct acrn_vm *vm, struct acrn_vm_memmap *memmap)
 
 		vm_region++;
 		user_vm_pa += region_size;
+		i += 1 << order;
 	}
 
 	/* Inform the ACRN Hypervisor to set up EPT mappings */