On Tue, Oct 25, 2022, Chao Peng wrote:

+static int kvm_vm_ioctl_set_mem_attr(struct kvm *kvm, gpa_t gpa, gpa_t size,
+				     bool is_private)
+{
+	gfn_t start, end;
+	unsigned long i;
+	void *entry;
+	int idx;
+	int r = 0;
+
+	if (size == 0 || gpa + size < gpa)
+		return -EINVAL;
+	if (gpa & (PAGE_SIZE - 1) || size & (PAGE_SIZE - 1))
+		return -EINVAL;
+
+	start = gpa >> PAGE_SHIFT;
+	end = (gpa + size - 1 + PAGE_SIZE) >> PAGE_SHIFT;
+
+	/*
+	 * Guest memory defaults to private, kvm->mem_attr_array only stores
+	 * shared memory.
+	 */
+	entry = is_private ? NULL : xa_mk_value(KVM_MEM_ATTR_SHARED);
+
+	idx = srcu_read_lock(&kvm->srcu);
+	KVM_MMU_LOCK(kvm);
+	kvm_mmu_invalidate_begin(kvm, start, end);
+
+	for (i = start; i < end; i++) {
+		r = xa_err(xa_store(&kvm->mem_attr_array, i, entry,
+				    GFP_KERNEL_ACCOUNT));
+		if (r)
+			goto err;
+	}
+
+	kvm_unmap_mem_range(kvm, start, end);
+
+	goto ret;
+err:
+	for (; i > start; i--)
+		xa_erase(&kvm->mem_attr_array, i);

I don't think deleting previous entries is correct.  To unwind, the correct thing
to do is restore the original values.  E.g. if userspace space is mapping a large
range as shared, and some of the previous entries were shared, deleting them would
incorrectly "convert" those entries to private.

Tracking the previous state likely isn't the best approach, e.g. it would require
speculatively allocating extra memory for a rare condition that is likely going to
lead to OOM anyways.

Instead of trying to unwind, what about updating the ioctl() params such that
retrying with the updated addr+size would Just Work?  E.g.

diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 55b07aae67cc..f1de592a1a06 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c

@@ -1015,15 +1015,12 @@ static int kvm_vm_ioctl_set_mem_attr(struct kvm *kvm, gpa_t gpa, gpa_t size,
 
        kvm_unmap_mem_range(kvm, start, end, attr);
 
-       goto ret;
-err:
-       for (; i > start; i--)
-               xa_erase(&kvm->mem_attr_array, i);
-ret:
        kvm_mmu_invalidate_end(kvm, start, end);
        KVM_MMU_UNLOCK(kvm);
        srcu_read_unlock(&kvm->srcu, idx);
 
+       <update gpa and size>
+
        return r;
 }
 #endif /* CONFIG_KVM_GENERIC_PRIVATE_MEM */

@@ -4989,6 +4986,8 @@ static long kvm_vm_ioctl(struct file *filp,
 
                r = kvm_vm_ioctl_set_mem_attr(kvm, region.addr,
                                              region.size, set);
+               if (copy_to_user(argp, &region, sizeof(region)) && !r)
+                       r = -EFAULT
                break;
        }
 #endif