1 files changed, 6 insertions, 5 deletions

diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 55a5d4804499..d6351a34b297 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -557,6 +557,10 @@ int __kvm_set_memory_region(struct kvm *kvm,
         base_gfn = mem->guest_phys_addr >> PAGE_SHIFT;
         npages = mem->memory_size >> PAGE_SHIFT;
 
+        r = -EINVAL;
+        if (npages > KVM_MEM_MAX_NR_PAGES)
+                goto out;
+
         if (!npages)
                 mem->flags &= ~KVM_MEM_LOG_DIRTY_PAGES;
 
@@ -1187,13 +1191,10 @@ void mark_page_dirty(struct kvm *kvm, gfn_t gfn)
         memslot = gfn_to_memslot_unaliased(kvm, gfn);
         if (memslot && memslot->dirty_bitmap) {
                 unsigned long rel_gfn = gfn - memslot->base_gfn;
-                unsigned long *p = memslot->dirty_bitmap +
-                                        rel_gfn / BITS_PER_LONG;
-                int offset = rel_gfn % BITS_PER_LONG;
 
                 /* avoid RMW */
-                if (!generic_test_le_bit(offset, p))
-                        generic___set_le_bit(offset, p);
+                if (!generic_test_le_bit(rel_gfn, memslot->dirty_bitmap))
+                        generic___set_le_bit(rel_gfn, memslot->dirty_bitmap);
         }
 }