1 files changed, 11 insertions, 8 deletions

diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index e4d358195e54..9a56ca2fa257 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -822,6 +822,8 @@ int __kvm_set_memory_region(struct kvm *kvm,
                 old_memslots = kvm->memslots;
                 rcu_assign_pointer(kvm->memslots, slots);
                 synchronize_srcu_expedited(&kvm->srcu);
+                /* slot was deleted or moved, clear iommu mapping */
+                kvm_iommu_unmap_pages(kvm, &old);
                 /* From this point no new shadow pages pointing to a deleted,
                  * or moved, memslot will be created.
                  *
@@ -837,20 +839,19 @@ int __kvm_set_memory_region(struct kvm *kvm,
         if (r)
                 goto out_free;
 
-        /* map/unmap the pages in iommu page table */
-        if (npages) {
-                r = kvm_iommu_map_pages(kvm, &new);
-                if (r)
-                        goto out_free;
-        } else
-                kvm_iommu_unmap_pages(kvm, &old);
-
         r = -ENOMEM;
         slots = kmemdup(kvm->memslots, sizeof(struct kvm_memslots),
                         GFP_KERNEL);
         if (!slots)
                 goto out_free;
 
+        /* map new memory slot into the iommu */
+        if (npages) {
+                r = kvm_iommu_map_pages(kvm, &new);
+                if (r)
+                        goto out_slots;
+        }
+
         /* actual memory is freed via old in kvm_free_physmem_slot below */
         if (!npages) {
                 new.dirty_bitmap = NULL;
@@ -869,6 +870,8 @@ int __kvm_set_memory_region(struct kvm *kvm,
 
         return 0;
 
+out_slots:
+        kfree(slots);
 out_free:
         kvm_free_physmem_slot(&new, &old);
 out: