Merge branch 'topic/core-cleanup' into for-linus
[safe/jmp/linux-2.6] / arch / x86 / kvm / mmu.c
index fd5579c..19a8906 100644 (file)
@@ -18,6 +18,7 @@
  */
 
 #include "mmu.h"
+#include "x86.h"
 #include "kvm_cache_regs.h"
 
 #include <linux/kvm_host.h>
@@ -29,6 +30,8 @@
 #include <linux/swap.h>
 #include <linux/hugetlb.h>
 #include <linux/compiler.h>
+#include <linux/srcu.h>
+#include <linux/slab.h>
 
 #include <asm/page.h>
 #include <asm/cmpxchg.h>
@@ -108,6 +111,9 @@ module_param(oos_shadow, bool, 0644);
 
 #define PT32_LEVEL_MASK(level) \
                (((1ULL << PT32_LEVEL_BITS) - 1) << PT32_LEVEL_SHIFT(level))
+#define PT32_LVL_OFFSET_MASK(level) \
+       (PT32_BASE_ADDR_MASK & ((1ULL << (PAGE_SHIFT + (((level) - 1) \
+                                               * PT32_LEVEL_BITS))) - 1))
 
 #define PT32_INDEX(address, level)\
        (((address) >> PT32_LEVEL_SHIFT(level)) & ((1 << PT32_LEVEL_BITS) - 1))
@@ -116,23 +122,23 @@ module_param(oos_shadow, bool, 0644);
 #define PT64_BASE_ADDR_MASK (((1ULL << 52) - 1) & ~(u64)(PAGE_SIZE-1))
 #define PT64_DIR_BASE_ADDR_MASK \
        (PT64_BASE_ADDR_MASK & ~((1ULL << (PAGE_SHIFT + PT64_LEVEL_BITS)) - 1))
+#define PT64_LVL_ADDR_MASK(level) \
+       (PT64_BASE_ADDR_MASK & ~((1ULL << (PAGE_SHIFT + (((level) - 1) \
+                                               * PT64_LEVEL_BITS))) - 1))
+#define PT64_LVL_OFFSET_MASK(level) \
+       (PT64_BASE_ADDR_MASK & ((1ULL << (PAGE_SHIFT + (((level) - 1) \
+                                               * PT64_LEVEL_BITS))) - 1))
 
 #define PT32_BASE_ADDR_MASK PAGE_MASK
 #define PT32_DIR_BASE_ADDR_MASK \
        (PAGE_MASK & ~((1ULL << (PAGE_SHIFT + PT32_LEVEL_BITS)) - 1))
+#define PT32_LVL_ADDR_MASK(level) \
+       (PAGE_MASK & ~((1ULL << (PAGE_SHIFT + (((level) - 1) \
+                                           * PT32_LEVEL_BITS))) - 1))
 
 #define PT64_PERM_MASK (PT_PRESENT_MASK | PT_WRITABLE_MASK | PT_USER_MASK \
                        | PT64_NX_MASK)
 
-#define PFERR_PRESENT_MASK (1U << 0)
-#define PFERR_WRITE_MASK (1U << 1)
-#define PFERR_USER_MASK (1U << 2)
-#define PFERR_RSVD_MASK (1U << 3)
-#define PFERR_FETCH_MASK (1U << 4)
-
-#define PT_DIRECTORY_LEVEL 2
-#define PT_PAGE_TABLE_LEVEL 1
-
 #define RMAP_EXT 4
 
 #define ACC_EXEC_MASK    1
@@ -140,6 +146,14 @@ module_param(oos_shadow, bool, 0644);
 #define ACC_USER_MASK    PT_USER_MASK
 #define ACC_ALL          (ACC_EXEC_MASK | ACC_WRITE_MASK | ACC_USER_MASK)
 
+#include <trace/events/kvm.h>
+
+#undef TRACE_INCLUDE_FILE
+#define CREATE_TRACE_POINTS
+#include "mmutrace.h"
+
+#define SPTE_HOST_WRITEABLE (1ULL << PT_FIRST_AVAIL_BITS_SHIFT)
+
 #define SHADOW_PT_INDEX(addr, level) PT64_INDEX(addr, level)
 
 struct kvm_rmap_desc {
@@ -211,7 +225,7 @@ EXPORT_SYMBOL_GPL(kvm_mmu_set_mask_ptes);
 
 static int is_write_protection(struct kvm_vcpu *vcpu)
 {
-       return vcpu->arch.cr0 & X86_CR0_WP;
+       return kvm_read_cr0_bits(vcpu, X86_CR0_WP);
 }
 
 static int is_cpuid_PSE36(void)
@@ -221,7 +235,7 @@ static int is_cpuid_PSE36(void)
 
 static int is_nx(struct kvm_vcpu *vcpu)
 {
-       return vcpu->arch.shadow_efer & EFER_NX;
+       return vcpu->arch.efer & EFER_NX;
 }
 
 static int is_shadow_present_pte(u64 pte)
@@ -235,7 +249,7 @@ static int is_large_pte(u64 pte)
        return pte & PT_PAGE_SIZE_MASK;
 }
 
-static int is_writeble_pte(unsigned long pte)
+static int is_writable_pte(unsigned long pte)
 {
        return pte & PT_WRITABLE_MASK;
 }
@@ -254,7 +268,7 @@ static int is_last_spte(u64 pte, int level)
 {
        if (level == PT_PAGE_TABLE_LEVEL)
                return 1;
-       if (level == PT_DIRECTORY_LEVEL && is_large_pte(pte))
+       if (is_large_pte(pte))
                return 1;
        return 0;
 }
@@ -390,37 +404,52 @@ static void mmu_free_rmap_desc(struct kvm_rmap_desc *rd)
  * Return the pointer to the largepage write count for a given
  * gfn, handling slots that are not large page aligned.
  */
-static int *slot_largepage_idx(gfn_t gfn, struct kvm_memory_slot *slot)
+static int *slot_largepage_idx(gfn_t gfn,
+                              struct kvm_memory_slot *slot,
+                              int level)
 {
        unsigned long idx;
 
-       idx = (gfn / KVM_PAGES_PER_HPAGE) -
-             (slot->base_gfn / KVM_PAGES_PER_HPAGE);
-       return &slot->lpage_info[idx].write_count;
+       idx = (gfn / KVM_PAGES_PER_HPAGE(level)) -
+             (slot->base_gfn / KVM_PAGES_PER_HPAGE(level));
+       return &slot->lpage_info[level - 2][idx].write_count;
 }
 
 static void account_shadowed(struct kvm *kvm, gfn_t gfn)
 {
+       struct kvm_memory_slot *slot;
        int *write_count;
+       int i;
 
        gfn = unalias_gfn(kvm, gfn);
-       write_count = slot_largepage_idx(gfn,
-                                        gfn_to_memslot_unaliased(kvm, gfn));
-       *write_count += 1;
+
+       slot = gfn_to_memslot_unaliased(kvm, gfn);
+       for (i = PT_DIRECTORY_LEVEL;
+            i < PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES; ++i) {
+               write_count   = slot_largepage_idx(gfn, slot, i);
+               *write_count += 1;
+       }
 }
 
 static void unaccount_shadowed(struct kvm *kvm, gfn_t gfn)
 {
+       struct kvm_memory_slot *slot;
        int *write_count;
+       int i;
 
        gfn = unalias_gfn(kvm, gfn);
-       write_count = slot_largepage_idx(gfn,
-                                        gfn_to_memslot_unaliased(kvm, gfn));
-       *write_count -= 1;
-       WARN_ON(*write_count < 0);
+       for (i = PT_DIRECTORY_LEVEL;
+            i < PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES; ++i) {
+               slot          = gfn_to_memslot_unaliased(kvm, gfn);
+               write_count   = slot_largepage_idx(gfn, slot, i);
+               *write_count -= 1;
+               WARN_ON(*write_count < 0);
+       }
 }
 
-static int has_wrprotected_page(struct kvm *kvm, gfn_t gfn)
+static int has_wrprotected_page(struct kvm *kvm,
+                               gfn_t gfn,
+                               int level)
 {
        struct kvm_memory_slot *slot;
        int *largepage_idx;
@@ -428,47 +457,53 @@ static int has_wrprotected_page(struct kvm *kvm, gfn_t gfn)
        gfn = unalias_gfn(kvm, gfn);
        slot = gfn_to_memslot_unaliased(kvm, gfn);
        if (slot) {
-               largepage_idx = slot_largepage_idx(gfn, slot);
+               largepage_idx = slot_largepage_idx(gfn, slot, level);
                return *largepage_idx;
        }
 
        return 1;
 }
 
-static int host_largepage_backed(struct kvm *kvm, gfn_t gfn)
+static int host_mapping_level(struct kvm *kvm, gfn_t gfn)
 {
-       struct vm_area_struct *vma;
-       unsigned long addr;
-       int ret = 0;
+       unsigned long page_size;
+       int i, ret = 0;
 
-       addr = gfn_to_hva(kvm, gfn);
-       if (kvm_is_error_hva(addr))
-               return ret;
+       page_size = kvm_host_page_size(kvm, gfn);
 
-       down_read(&current->mm->mmap_sem);
-       vma = find_vma(current->mm, addr);
-       if (vma && is_vm_hugetlb_page(vma))
-               ret = 1;
-       up_read(&current->mm->mmap_sem);
+       for (i = PT_PAGE_TABLE_LEVEL;
+            i < (PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES); ++i) {
+               if (page_size >= KVM_HPAGE_SIZE(i))
+                       ret = i;
+               else
+                       break;
+       }
 
        return ret;
 }
 
-static int is_largepage_backed(struct kvm_vcpu *vcpu, gfn_t large_gfn)
+static int mapping_level(struct kvm_vcpu *vcpu, gfn_t large_gfn)
 {
        struct kvm_memory_slot *slot;
-
-       if (has_wrprotected_page(vcpu->kvm, large_gfn))
-               return 0;
-
-       if (!host_largepage_backed(vcpu->kvm, large_gfn))
-               return 0;
+       int host_level, level, max_level;
 
        slot = gfn_to_memslot(vcpu->kvm, large_gfn);
        if (slot && slot->dirty_bitmap)
-               return 0;
+               return PT_PAGE_TABLE_LEVEL;
 
-       return 1;
+       host_level = host_mapping_level(vcpu->kvm, large_gfn);
+
+       if (host_level == PT_PAGE_TABLE_LEVEL)
+               return host_level;
+
+       max_level = kvm_x86_ops->get_lpage_level() < host_level ?
+               kvm_x86_ops->get_lpage_level() : host_level;
+
+       for (level = PT_DIRECTORY_LEVEL; level <= max_level; ++level)
+               if (has_wrprotected_page(vcpu->kvm, large_gfn, level))
+                       break;
+
+       return level - 1;
 }
 
 /*
@@ -476,19 +511,19 @@ static int is_largepage_backed(struct kvm_vcpu *vcpu, gfn_t large_gfn)
  * Note: gfn must be unaliased before this function get called
  */
 
-static unsigned long *gfn_to_rmap(struct kvm *kvm, gfn_t gfn, int lpage)
+static unsigned long *gfn_to_rmap(struct kvm *kvm, gfn_t gfn, int level)
 {
        struct kvm_memory_slot *slot;
        unsigned long idx;
 
        slot = gfn_to_memslot(kvm, gfn);
-       if (!lpage)
+       if (likely(level == PT_PAGE_TABLE_LEVEL))
                return &slot->rmap[gfn - slot->base_gfn];
 
-       idx = (gfn / KVM_PAGES_PER_HPAGE) -
-             (slot->base_gfn / KVM_PAGES_PER_HPAGE);
+       idx = (gfn / KVM_PAGES_PER_HPAGE(level)) -
+               (slot->base_gfn / KVM_PAGES_PER_HPAGE(level));
 
-       return &slot->lpage_info[idx].rmap_pde;
+       return &slot->lpage_info[level - 2][idx].rmap_pde;
 }
 
 /*
@@ -504,7 +539,7 @@ static unsigned long *gfn_to_rmap(struct kvm *kvm, gfn_t gfn, int lpage)
  * the spte was not added.
  *
  */
-static int rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn, int lpage)
+static int rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
 {
        struct kvm_mmu_page *sp;
        struct kvm_rmap_desc *desc;
@@ -516,7 +551,7 @@ static int rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn, int lpage)
        gfn = unalias_gfn(vcpu->kvm, gfn);
        sp = page_header(__pa(spte));
        sp->gfns[spte - sp->spt] = gfn;
-       rmapp = gfn_to_rmap(vcpu->kvm, gfn, lpage);
+       rmapp = gfn_to_rmap(vcpu->kvm, gfn, sp->role.level);
        if (!*rmapp) {
                rmap_printk("rmap_add: %p %llx 0->1\n", spte, *spte);
                *rmapp = (unsigned long)spte;
@@ -582,11 +617,9 @@ static void rmap_remove(struct kvm *kvm, u64 *spte)
        pfn = spte_to_pfn(*spte);
        if (*spte & shadow_accessed_mask)
                kvm_set_pfn_accessed(pfn);
-       if (is_writeble_pte(*spte))
-               kvm_release_pfn_dirty(pfn);
-       else
-               kvm_release_pfn_clean(pfn);
-       rmapp = gfn_to_rmap(kvm, sp->gfns[spte - sp->spt], is_large_pte(*spte));
+       if (is_writable_pte(*spte))
+               kvm_set_pfn_dirty(pfn);
+       rmapp = gfn_to_rmap(kvm, sp->gfns[spte - sp->spt], sp->role.level);
        if (!*rmapp) {
                printk(KERN_ERR "rmap_remove: %p %llx 0->BUG\n", spte, *spte);
                BUG();
@@ -613,6 +646,7 @@ static void rmap_remove(struct kvm *kvm, u64 *spte)
                        prev_desc = desc;
                        desc = desc->more;
                }
+               pr_err("rmap_remove: %p %llx many->many\n", spte, *spte);
                BUG();
        }
 }
@@ -649,17 +683,17 @@ static int rmap_write_protect(struct kvm *kvm, u64 gfn)
 {
        unsigned long *rmapp;
        u64 *spte;
-       int write_protected = 0;
+       int i, write_protected = 0;
 
        gfn = unalias_gfn(kvm, gfn);
-       rmapp = gfn_to_rmap(kvm, gfn, 0);
+       rmapp = gfn_to_rmap(kvm, gfn, PT_PAGE_TABLE_LEVEL);
 
        spte = rmap_next(kvm, rmapp, NULL);
        while (spte) {
                BUG_ON(!spte);
                BUG_ON(!(*spte & PT_PRESENT_MASK));
                rmap_printk("rmap_write_protect: spte %p %llx\n", spte, *spte);
-               if (is_writeble_pte(*spte)) {
+               if (is_writable_pte(*spte)) {
                        __set_spte(spte, *spte & ~PT_WRITABLE_MASK);
                        write_protected = 1;
                }
@@ -674,27 +708,31 @@ static int rmap_write_protect(struct kvm *kvm, u64 gfn)
        }
 
        /* check for huge page mappings */
-       rmapp = gfn_to_rmap(kvm, gfn, 1);
-       spte = rmap_next(kvm, rmapp, NULL);
-       while (spte) {
-               BUG_ON(!spte);
-               BUG_ON(!(*spte & PT_PRESENT_MASK));
-               BUG_ON((*spte & (PT_PAGE_SIZE_MASK|PT_PRESENT_MASK)) != (PT_PAGE_SIZE_MASK|PT_PRESENT_MASK));
-               pgprintk("rmap_write_protect(large): spte %p %llx %lld\n", spte, *spte, gfn);
-               if (is_writeble_pte(*spte)) {
-                       rmap_remove(kvm, spte);
-                       --kvm->stat.lpages;
-                       __set_spte(spte, shadow_trap_nonpresent_pte);
-                       spte = NULL;
-                       write_protected = 1;
+       for (i = PT_DIRECTORY_LEVEL;
+            i < PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES; ++i) {
+               rmapp = gfn_to_rmap(kvm, gfn, i);
+               spte = rmap_next(kvm, rmapp, NULL);
+               while (spte) {
+                       BUG_ON(!spte);
+                       BUG_ON(!(*spte & PT_PRESENT_MASK));
+                       BUG_ON((*spte & (PT_PAGE_SIZE_MASK|PT_PRESENT_MASK)) != (PT_PAGE_SIZE_MASK|PT_PRESENT_MASK));
+                       pgprintk("rmap_write_protect(large): spte %p %llx %lld\n", spte, *spte, gfn);
+                       if (is_writable_pte(*spte)) {
+                               rmap_remove(kvm, spte);
+                               --kvm->stat.lpages;
+                               __set_spte(spte, shadow_trap_nonpresent_pte);
+                               spte = NULL;
+                               write_protected = 1;
+                       }
+                       spte = rmap_next(kvm, rmapp, spte);
                }
-               spte = rmap_next(kvm, rmapp, spte);
        }
 
        return write_protected;
 }
 
-static int kvm_unmap_rmapp(struct kvm *kvm, unsigned long *rmapp)
+static int kvm_unmap_rmapp(struct kvm *kvm, unsigned long *rmapp,
+                          unsigned long data)
 {
        u64 *spte;
        int need_tlb_flush = 0;
@@ -709,33 +747,75 @@ static int kvm_unmap_rmapp(struct kvm *kvm, unsigned long *rmapp)
        return need_tlb_flush;
 }
 
+static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
+                            unsigned long data)
+{
+       int need_flush = 0;
+       u64 *spte, new_spte;
+       pte_t *ptep = (pte_t *)data;
+       pfn_t new_pfn;
+
+       WARN_ON(pte_huge(*ptep));
+       new_pfn = pte_pfn(*ptep);
+       spte = rmap_next(kvm, rmapp, NULL);
+       while (spte) {
+               BUG_ON(!is_shadow_present_pte(*spte));
+               rmap_printk("kvm_set_pte_rmapp: spte %p %llx\n", spte, *spte);
+               need_flush = 1;
+               if (pte_write(*ptep)) {
+                       rmap_remove(kvm, spte);
+                       __set_spte(spte, shadow_trap_nonpresent_pte);
+                       spte = rmap_next(kvm, rmapp, NULL);
+               } else {
+                       new_spte = *spte &~ (PT64_BASE_ADDR_MASK);
+                       new_spte |= (u64)new_pfn << PAGE_SHIFT;
+
+                       new_spte &= ~PT_WRITABLE_MASK;
+                       new_spte &= ~SPTE_HOST_WRITEABLE;
+                       if (is_writable_pte(*spte))
+                               kvm_set_pfn_dirty(spte_to_pfn(*spte));
+                       __set_spte(spte, new_spte);
+                       spte = rmap_next(kvm, rmapp, spte);
+               }
+       }
+       if (need_flush)
+               kvm_flush_remote_tlbs(kvm);
+
+       return 0;
+}
+
 static int kvm_handle_hva(struct kvm *kvm, unsigned long hva,
-                         int (*handler)(struct kvm *kvm, unsigned long *rmapp))
+                         unsigned long data,
+                         int (*handler)(struct kvm *kvm, unsigned long *rmapp,
+                                        unsigned long data))
 {
-       int i;
+       int i, j;
+       int ret;
        int retval = 0;
+       struct kvm_memslots *slots;
 
-       /*
-        * If mmap_sem isn't taken, we can look the memslots with only
-        * the mmu_lock by skipping over the slots with userspace_addr == 0.
-        */
-       for (i = 0; i < kvm->nmemslots; i++) {
-               struct kvm_memory_slot *memslot = &kvm->memslots[i];
+       slots = rcu_dereference(kvm->memslots);
+
+       for (i = 0; i < slots->nmemslots; i++) {
+               struct kvm_memory_slot *memslot = &slots->memslots[i];
                unsigned long start = memslot->userspace_addr;
                unsigned long end;
 
-               /* mmu_lock protects userspace_addr */
-               if (!start)
-                       continue;
-
                end = start + (memslot->npages << PAGE_SHIFT);
                if (hva >= start && hva < end) {
                        gfn_t gfn_offset = (hva - start) >> PAGE_SHIFT;
-                       retval |= handler(kvm, &memslot->rmap[gfn_offset]);
-                       retval |= handler(kvm,
-                                         &memslot->lpage_info[
-                                                 gfn_offset /
-                                                 KVM_PAGES_PER_HPAGE].rmap_pde);
+
+                       ret = handler(kvm, &memslot->rmap[gfn_offset], data);
+
+                       for (j = 0; j < KVM_NR_PAGE_SIZES - 1; ++j) {
+                               int idx = gfn_offset;
+                               idx /= KVM_PAGES_PER_HPAGE(PT_DIRECTORY_LEVEL + j);
+                               ret |= handler(kvm,
+                                       &memslot->lpage_info[j][idx].rmap_pde,
+                                       data);
+                       }
+                       trace_kvm_age_page(hva, memslot, ret);
+                       retval |= ret;
                }
        }
 
@@ -744,17 +824,29 @@ static int kvm_handle_hva(struct kvm *kvm, unsigned long hva,
 
 int kvm_unmap_hva(struct kvm *kvm, unsigned long hva)
 {
-       return kvm_handle_hva(kvm, hva, kvm_unmap_rmapp);
+       return kvm_handle_hva(kvm, hva, 0, kvm_unmap_rmapp);
 }
 
-static int kvm_age_rmapp(struct kvm *kvm, unsigned long *rmapp)
+void kvm_set_spte_hva(struct kvm *kvm, unsigned long hva, pte_t pte)
+{
+       kvm_handle_hva(kvm, hva, (unsigned long)&pte, kvm_set_pte_rmapp);
+}
+
+static int kvm_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
+                        unsigned long data)
 {
        u64 *spte;
        int young = 0;
 
-       /* always return old for EPT */
+       /*
+        * Emulate the accessed bit for EPT, by checking if this page has
+        * an EPT mapping, and clearing it if it does. On the next access,
+        * a new EPT mapping will be established.
+        * This has some overhead, but not as much as the cost of swapping
+        * out actively used pages or breaking up actively used hugepages.
+        */
        if (!shadow_accessed_mask)
-               return 0;
+               return kvm_unmap_rmapp(kvm, rmapp, data);
 
        spte = rmap_next(kvm, rmapp, NULL);
        while (spte) {
@@ -773,20 +865,23 @@ static int kvm_age_rmapp(struct kvm *kvm, unsigned long *rmapp)
 
 #define RMAP_RECYCLE_THRESHOLD 1000
 
-static void rmap_recycle(struct kvm_vcpu *vcpu, gfn_t gfn, int lpage)
+static void rmap_recycle(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
 {
        unsigned long *rmapp;
+       struct kvm_mmu_page *sp;
+
+       sp = page_header(__pa(spte));
 
        gfn = unalias_gfn(vcpu->kvm, gfn);
-       rmapp = gfn_to_rmap(vcpu->kvm, gfn, lpage);
+       rmapp = gfn_to_rmap(vcpu->kvm, gfn, sp->role.level);
 
-       kvm_unmap_rmapp(vcpu->kvm, rmapp);
+       kvm_unmap_rmapp(vcpu->kvm, rmapp, 0);
        kvm_flush_remote_tlbs(vcpu->kvm);
 }
 
 int kvm_age_hva(struct kvm *kvm, unsigned long hva)
 {
-       return kvm_handle_hva(kvm, hva, kvm_age_rmapp);
+       return kvm_handle_hva(kvm, hva, 0, kvm_age_rmapp);
 }
 
 #ifdef MMU_DEBUG
@@ -1119,6 +1214,7 @@ static int kvm_sync_page(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp)
                return 1;
        }
 
+       trace_kvm_mmu_sync_page(sp);
        if (rmap_write_protect(vcpu->kvm, sp->gfn))
                kvm_flush_remote_tlbs(vcpu->kvm);
        kvm_unlink_unsync_page(vcpu->kvm, sp);
@@ -1241,8 +1337,6 @@ static struct kvm_mmu_page *kvm_mmu_get_page(struct kvm_vcpu *vcpu,
                quadrant &= (1 << ((PT32_PT_BITS - PT64_PT_BITS) * level)) - 1;
                role.quadrant = quadrant;
        }
-       pgprintk("%s: looking gfn %lx role %x\n", __func__,
-                gfn, role.word);
        index = kvm_page_table_hashfn(gfn);
        bucket = &vcpu->kvm->arch.mmu_page_hash[index];
        hlist_for_each_entry_safe(sp, node, tmp, bucket, hash_link)
@@ -1259,14 +1353,13 @@ static struct kvm_mmu_page *kvm_mmu_get_page(struct kvm_vcpu *vcpu,
                                set_bit(KVM_REQ_MMU_SYNC, &vcpu->requests);
                                kvm_mmu_mark_parents_unsync(vcpu, sp);
                        }
-                       pgprintk("%s: found\n", __func__);
+                       trace_kvm_mmu_get_page(sp, false);
                        return sp;
                }
        ++vcpu->kvm->stat.mmu_cache_miss;
        sp = kvm_mmu_alloc_page(vcpu, parent_pte);
        if (!sp)
                return sp;
-       pgprintk("%s: adding gfn %lx role %x\n", __func__, gfn, role.word);
        sp->gfn = gfn;
        sp->role = role;
        hlist_add_head(&sp->hash_link, bucket);
@@ -1279,6 +1372,7 @@ static struct kvm_mmu_page *kvm_mmu_get_page(struct kvm_vcpu *vcpu,
                vcpu->arch.mmu.prefetch_page(vcpu, sp);
        else
                nonpaging_prefetch_page(vcpu, sp);
+       trace_kvm_mmu_get_page(sp, true);
        return sp;
 }
 
@@ -1302,6 +1396,11 @@ static bool shadow_walk_okay(struct kvm_shadow_walk_iterator *iterator)
 {
        if (iterator->level < PT_PAGE_TABLE_LEVEL)
                return false;
+
+       if (iterator->level == PT_PAGE_TABLE_LEVEL)
+               if (is_large_pte(*iterator->sptep))
+                       return false;
+
        iterator->index = SHADOW_PT_INDEX(iterator->addr, iterator->level);
        iterator->sptep = ((u64 *)__va(iterator->shadow_addr)) + iterator->index;
        return true;
@@ -1391,8 +1490,8 @@ static int mmu_zap_unsync_children(struct kvm *kvm,
                for_each_sp(pages, sp, parents, i) {
                        kvm_mmu_zap_page(kvm, sp);
                        mmu_pages_clear_parents(&parents);
+                       zapped++;
                }
-               zapped += pages.nr;
                kvm_mmu_pages_init(parent, &parents, &pages);
        }
 
@@ -1402,6 +1501,8 @@ static int mmu_zap_unsync_children(struct kvm *kvm,
 static int kvm_mmu_zap_page(struct kvm *kvm, struct kvm_mmu_page *sp)
 {
        int ret;
+
+       trace_kvm_mmu_zap_page(sp);
        ++kvm->stat.mmu_shadow_zapped;
        ret = mmu_zap_unsync_children(kvm, sp);
        kvm_mmu_page_unlink_children(kvm, sp);
@@ -1441,14 +1542,16 @@ void kvm_mmu_change_mmu_pages(struct kvm *kvm, unsigned int kvm_nr_mmu_pages)
         */
 
        if (used_pages > kvm_nr_mmu_pages) {
-               while (used_pages > kvm_nr_mmu_pages) {
+               while (used_pages > kvm_nr_mmu_pages &&
+                       !list_empty(&kvm->arch.active_mmu_pages)) {
                        struct kvm_mmu_page *page;
 
                        page = container_of(kvm->arch.active_mmu_pages.prev,
                                            struct kvm_mmu_page, link);
-                       kvm_mmu_zap_page(kvm, page);
+                       used_pages -= kvm_mmu_zap_page(kvm, page);
                        used_pages--;
                }
+               kvm_nr_mmu_pages = used_pages;
                kvm->arch.n_free_mmu_pages = 0;
        }
        else
@@ -1495,14 +1598,15 @@ static void mmu_unshadow(struct kvm *kvm, gfn_t gfn)
                    && !sp->role.invalid) {
                        pgprintk("%s: zap %lx %x\n",
                                 __func__, gfn, sp->role.word);
-                       kvm_mmu_zap_page(kvm, sp);
+                       if (kvm_mmu_zap_page(kvm, sp))
+                               nn = bucket->first;
                }
        }
 }
 
 static void page_header_update_slot(struct kvm *kvm, void *pte, gfn_t gfn)
 {
-       int slot = memslot_id(kvm, gfn_to_memslot(kvm, gfn));
+       int slot = memslot_id(kvm, gfn);
        struct kvm_mmu_page *sp = page_header(__pa(pte));
 
        __set_bit(slot, sp->slot_bitmap);
@@ -1526,7 +1630,7 @@ struct page *gva_to_page(struct kvm_vcpu *vcpu, gva_t gva)
 {
        struct page *page;
 
-       gpa_t gpa = vcpu->arch.mmu.gva_to_gpa(vcpu, gva);
+       gpa_t gpa = kvm_mmu_gva_to_gpa_read(vcpu, gva, NULL);
 
        if (gpa == UNMAPPED_GVA)
                return NULL;
@@ -1648,6 +1752,7 @@ static int kvm_unsync_page(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp)
        struct kvm_mmu_page *s;
        struct hlist_node *node, *n;
 
+       trace_kvm_mmu_unsync_page(sp);
        index = kvm_page_table_hashfn(sp->gfn);
        bucket = &vcpu->kvm->arch.mmu_page_hash[index];
        /* don't unsync if pagetable is shadowed with multiple roles */
@@ -1686,9 +1791,9 @@ static int mmu_need_write_protect(struct kvm_vcpu *vcpu, gfn_t gfn,
 
 static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                    unsigned pte_access, int user_fault,
-                   int write_fault, int dirty, int largepage,
+                   int write_fault, int dirty, int level,
                    gfn_t gfn, pfn_t pfn, bool speculative,
-                   bool can_unsync)
+                   bool can_unsync, bool reset_host_protection)
 {
        u64 spte;
        int ret = 0;
@@ -1709,18 +1814,22 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                spte |= shadow_nx_mask;
        if (pte_access & ACC_USER_MASK)
                spte |= shadow_user_mask;
-       if (largepage)
+       if (level > PT_PAGE_TABLE_LEVEL)
                spte |= PT_PAGE_SIZE_MASK;
        if (tdp_enabled)
                spte |= kvm_x86_ops->get_mt_mask(vcpu, gfn,
                        kvm_is_mmio_pfn(pfn));
 
+       if (reset_host_protection)
+               spte |= SPTE_HOST_WRITEABLE;
+
        spte |= (u64)pfn << PAGE_SHIFT;
 
        if ((pte_access & ACC_WRITE_MASK)
            || (write_fault && !is_write_protection(vcpu) && !user_fault)) {
 
-               if (largepage && has_wrprotected_page(vcpu->kvm, gfn)) {
+               if (level > PT_PAGE_TABLE_LEVEL &&
+                   has_wrprotected_page(vcpu->kvm, gfn, level)) {
                        ret = 1;
                        spte = shadow_trap_nonpresent_pte;
                        goto set_pte;
@@ -1734,7 +1843,7 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                 * is responsibility of mmu_get_page / kvm_sync_page.
                 * Same reasoning can be applied to dirty page accounting.
                 */
-               if (!can_unsync && is_writeble_pte(*sptep))
+               if (!can_unsync && is_writable_pte(*sptep))
                        goto set_pte;
 
                if (mmu_need_write_protect(vcpu, gfn, can_unsync)) {
@@ -1742,7 +1851,7 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                                 __func__, gfn);
                        ret = 1;
                        pte_access &= ~ACC_WRITE_MASK;
-                       if (is_writeble_pte(spte))
+                       if (is_writable_pte(spte))
                                spte &= ~PT_WRITABLE_MASK;
                }
        }
@@ -1758,11 +1867,12 @@ set_pte:
 static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                         unsigned pt_access, unsigned pte_access,
                         int user_fault, int write_fault, int dirty,
-                        int *ptwrite, int largepage, gfn_t gfn,
-                        pfn_t pfn, bool speculative)
+                        int *ptwrite, int level, gfn_t gfn,
+                        pfn_t pfn, bool speculative,
+                        bool reset_host_protection)
 {
        int was_rmapped = 0;
-       int was_writeble = is_writeble_pte(*sptep);
+       int was_writable = is_writable_pte(*sptep);
        int rmap_count;
 
        pgprintk("%s: spte %llx access %x write_fault %d"
@@ -1775,7 +1885,8 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                 * If we overwrite a PTE page pointer with a 2MB PMD, unlink
                 * the parent of the now unreachable PTE.
                 */
-               if (largepage && !is_large_pte(*sptep)) {
+               if (level > PT_PAGE_TABLE_LEVEL &&
+                   !is_large_pte(*sptep)) {
                        struct kvm_mmu_page *child;
                        u64 pte = *sptep;
 
@@ -1788,8 +1899,10 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
                } else
                        was_rmapped = 1;
        }
+
        if (set_spte(vcpu, sptep, pte_access, user_fault, write_fault,
-                     dirty, largepage, gfn, pfn, speculative, true)) {
+                     dirty, level, gfn, pfn, speculative, true,
+                     reset_host_protection)) {
                if (write_fault)
                        *ptwrite = 1;
                kvm_x86_ops->tlb_flush(vcpu);
@@ -1798,20 +1911,19 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
        pgprintk("%s: setting spte %llx\n", __func__, *sptep);
        pgprintk("instantiating %s PTE (%s) at %ld (%llx) addr %p\n",
                 is_large_pte(*sptep)? "2MB" : "4kB",
-                is_present_pte(*sptep)?"RW":"R", gfn,
-                *shadow_pte, sptep);
+                *sptep & PT_PRESENT_MASK ?"RW":"R", gfn,
+                *sptep, sptep);
        if (!was_rmapped && is_large_pte(*sptep))
                ++vcpu->kvm->stat.lpages;
 
        page_header_update_slot(vcpu->kvm, sptep, gfn);
        if (!was_rmapped) {
-               rmap_count = rmap_add(vcpu, sptep, gfn, largepage);
-               if (!is_rmap_spte(*sptep))
-                       kvm_release_pfn_clean(pfn);
+               rmap_count = rmap_add(vcpu, sptep, gfn);
+               kvm_release_pfn_clean(pfn);
                if (rmap_count > RMAP_RECYCLE_THRESHOLD)
-                       rmap_recycle(vcpu, gfn, largepage);
+                       rmap_recycle(vcpu, sptep, gfn);
        } else {
-               if (was_writeble)
+               if (was_writable)
                        kvm_release_pfn_dirty(pfn);
                else
                        kvm_release_pfn_clean(pfn);
@@ -1827,7 +1939,7 @@ static void nonpaging_new_cr3(struct kvm_vcpu *vcpu)
 }
 
 static int __direct_map(struct kvm_vcpu *vcpu, gpa_t v, int write,
-                       int largepage, gfn_t gfn, pfn_t pfn)
+                       int level, gfn_t gfn, pfn_t pfn)
 {
        struct kvm_shadow_walk_iterator iterator;
        struct kvm_mmu_page *sp;
@@ -1835,11 +1947,10 @@ static int __direct_map(struct kvm_vcpu *vcpu, gpa_t v, int write,
        gfn_t pseudo_gfn;
 
        for_each_shadow_entry(vcpu, (u64)gfn << PAGE_SHIFT, iterator) {
-               if (iterator.level == PT_PAGE_TABLE_LEVEL
-                   || (largepage && iterator.level == PT_DIRECTORY_LEVEL)) {
+               if (iterator.level == level) {
                        mmu_set_spte(vcpu, iterator.sptep, ACC_ALL, ACC_ALL,
                                     0, write, 1, &pt_write,
-                                    largepage, gfn, pfn, false);
+                                    level, gfn, pfn, false, true);
                        ++vcpu->stat.pf_fixed;
                        break;
                }
@@ -1867,14 +1978,20 @@ static int __direct_map(struct kvm_vcpu *vcpu, gpa_t v, int write,
 static int nonpaging_map(struct kvm_vcpu *vcpu, gva_t v, int write, gfn_t gfn)
 {
        int r;
-       int largepage = 0;
+       int level;
        pfn_t pfn;
        unsigned long mmu_seq;
 
-       if (is_largepage_backed(vcpu, gfn & ~(KVM_PAGES_PER_HPAGE-1))) {
-               gfn &= ~(KVM_PAGES_PER_HPAGE-1);
-               largepage = 1;
-       }
+       level = mapping_level(vcpu, gfn);
+
+       /*
+        * This path builds a PAE pagetable - so we can map 2mb pages at
+        * maximum. Therefore check if the level is larger than that.
+        */
+       if (level > PT_DIRECTORY_LEVEL)
+               level = PT_DIRECTORY_LEVEL;
+
+       gfn &= ~(KVM_PAGES_PER_HPAGE(level) - 1);
 
        mmu_seq = vcpu->kvm->mmu_notifier_seq;
        smp_rmb();
@@ -1890,7 +2007,7 @@ static int nonpaging_map(struct kvm_vcpu *vcpu, gva_t v, int write, gfn_t gfn)
        if (mmu_notifier_retry(vcpu, mmu_seq))
                goto out_unlock;
        kvm_mmu_free_some_pages(vcpu);
-       r = __direct_map(vcpu, v, write, largepage, gfn, pfn);
+       r = __direct_map(vcpu, v, write, level, gfn, pfn);
        spin_unlock(&vcpu->kvm->mmu_lock);
 
 
@@ -2036,8 +2153,11 @@ void kvm_mmu_sync_roots(struct kvm_vcpu *vcpu)
        spin_unlock(&vcpu->kvm->mmu_lock);
 }
 
-static gpa_t nonpaging_gva_to_gpa(struct kvm_vcpu *vcpu, gva_t vaddr)
+static gpa_t nonpaging_gva_to_gpa(struct kvm_vcpu *vcpu, gva_t vaddr,
+                                 u32 access, u32 *error)
 {
+       if (error)
+               *error = 0;
        return vaddr;
 }
 
@@ -2066,7 +2186,7 @@ static int tdp_page_fault(struct kvm_vcpu *vcpu, gva_t gpa,
 {
        pfn_t pfn;
        int r;
-       int largepage = 0;
+       int level;
        gfn_t gfn = gpa >> PAGE_SHIFT;
        unsigned long mmu_seq;
 
@@ -2077,10 +2197,10 @@ static int tdp_page_fault(struct kvm_vcpu *vcpu, gva_t gpa,
        if (r)
                return r;
 
-       if (is_largepage_backed(vcpu, gfn & ~(KVM_PAGES_PER_HPAGE-1))) {
-               gfn &= ~(KVM_PAGES_PER_HPAGE-1);
-               largepage = 1;
-       }
+       level = mapping_level(vcpu, gfn);
+
+       gfn &= ~(KVM_PAGES_PER_HPAGE(level) - 1);
+
        mmu_seq = vcpu->kvm->mmu_notifier_seq;
        smp_rmb();
        pfn = gfn_to_pfn(vcpu->kvm, gfn);
@@ -2093,7 +2213,7 @@ static int tdp_page_fault(struct kvm_vcpu *vcpu, gva_t gpa,
                goto out_unlock;
        kvm_mmu_free_some_pages(vcpu);
        r = __direct_map(vcpu, gpa, error_code & PFERR_WRITE_MASK,
-                        largepage, gfn, pfn);
+                        level, gfn, pfn);
        spin_unlock(&vcpu->kvm->mmu_lock);
 
        return r;
@@ -2210,7 +2330,9 @@ static void reset_rsvds_bits_mask(struct kvm_vcpu *vcpu, int level)
                context->rsvd_bits_mask[0][0] = exb_bit_rsvd |
                        rsvd_bits(maxphyaddr, 51);
                context->rsvd_bits_mask[1][3] = context->rsvd_bits_mask[0][3];
-               context->rsvd_bits_mask[1][2] = context->rsvd_bits_mask[0][2];
+               context->rsvd_bits_mask[1][2] = exb_bit_rsvd |
+                       rsvd_bits(maxphyaddr, 51) |
+                       rsvd_bits(13, 29);
                context->rsvd_bits_mask[1][1] = exb_bit_rsvd |
                        rsvd_bits(maxphyaddr, 51) |
                        rsvd_bits(13, 20);              /* large page */
@@ -2361,8 +2483,8 @@ int kvm_mmu_load(struct kvm_vcpu *vcpu)
        spin_unlock(&vcpu->kvm->mmu_lock);
        if (r)
                goto out;
+       /* set_cr3() should ensure TLB has been flushed */
        kvm_x86_ops->set_cr3(vcpu, vcpu->arch.mmu.root_hpa);
-       kvm_mmu_flush_tlb(vcpu);
 out:
        return r;
 }
@@ -2400,11 +2522,8 @@ static void mmu_pte_write_new_pte(struct kvm_vcpu *vcpu,
                                  const void *new)
 {
        if (sp->role.level != PT_PAGE_TABLE_LEVEL) {
-               if (!vcpu->arch.update_pte.largepage ||
-                   sp->role.glevels == PT32_ROOT_LEVEL) {
-                       ++vcpu->kvm->stat.mmu_pde_zapped;
-                       return;
-               }
+               ++vcpu->kvm->stat.mmu_pde_zapped;
+               return;
         }
 
        ++vcpu->kvm->stat.mmu_pte_updated;
@@ -2450,8 +2569,6 @@ static void mmu_guess_page_from_pte_write(struct kvm_vcpu *vcpu, gpa_t gpa,
        u64 gpte = 0;
        pfn_t pfn;
 
-       vcpu->arch.update_pte.largepage = 0;
-
        if (bytes != 4 && bytes != 8)
                return;
 
@@ -2479,10 +2596,6 @@ static void mmu_guess_page_from_pte_write(struct kvm_vcpu *vcpu, gpa_t gpa,
                return;
        gfn = (gpte & PT64_BASE_ADDR_MASK) >> PAGE_SHIFT;
 
-       if (is_large_pte(gpte) && is_largepage_backed(vcpu, gfn)) {
-               gfn &= ~(KVM_PAGES_PER_HPAGE-1);
-               vcpu->arch.update_pte.largepage = 1;
-       }
        vcpu->arch.update_pte.mmu_seq = vcpu->kvm->mmu_notifier_seq;
        smp_rmb();
        pfn = gfn_to_pfn(vcpu->kvm, gfn);
@@ -2625,7 +2738,10 @@ int kvm_mmu_unprotect_page_virt(struct kvm_vcpu *vcpu, gva_t gva)
        gpa_t gpa;
        int r;
 
-       gpa = vcpu->arch.mmu.gva_to_gpa(vcpu, gva);
+       if (tdp_enabled)
+               return 0;
+
+       gpa = kvm_mmu_gva_to_gpa_read(vcpu, gva, NULL);
 
        spin_lock(&vcpu->kvm->mmu_lock);
        r = kvm_mmu_unprotect_page(vcpu->kvm, gpa >> PAGE_SHIFT);
@@ -2636,7 +2752,8 @@ EXPORT_SYMBOL_GPL(kvm_mmu_unprotect_page_virt);
 
 void __kvm_mmu_free_some_pages(struct kvm_vcpu *vcpu)
 {
-       while (vcpu->kvm->arch.n_free_mmu_pages < KVM_REFILL_PAGES) {
+       while (vcpu->kvm->arch.n_free_mmu_pages < KVM_REFILL_PAGES &&
+              !list_empty(&vcpu->kvm->arch.active_mmu_pages)) {
                struct kvm_mmu_page *sp;
 
                sp = container_of(vcpu->kvm->arch.active_mmu_pages.prev,
@@ -2664,7 +2781,7 @@ int kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gva_t cr2, u32 error_code)
        if (r)
                goto out;
 
-       er = emulate_instruction(vcpu, vcpu->run, cr2, error_code, 0);
+       er = emulate_instruction(vcpu, cr2, error_code, 0);
 
        switch (er) {
        case EMULATE_DONE:
@@ -2675,6 +2792,7 @@ int kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gva_t cr2, u32 error_code)
        case EMULATE_FAIL:
                vcpu->run->exit_reason = KVM_EXIT_INTERNAL_ERROR;
                vcpu->run->internal.suberror = KVM_INTERNAL_ERROR_EMULATION;
+               vcpu->run->internal.ndata = 0;
                return 0;
        default:
                BUG();
@@ -2716,12 +2834,6 @@ static int alloc_mmu_pages(struct kvm_vcpu *vcpu)
 
        ASSERT(vcpu);
 
-       if (vcpu->kvm->arch.n_requested_mmu_pages)
-               vcpu->kvm->arch.n_free_mmu_pages =
-                                       vcpu->kvm->arch.n_requested_mmu_pages;
-       else
-               vcpu->kvm->arch.n_free_mmu_pages =
-                                       vcpu->kvm->arch.n_alloc_mmu_pages;
        /*
         * When emulating 32-bit mode, cr3 is only 32 bits even on x86_64.
         * Therefore we need to allocate shadow page tables in the first
@@ -2729,16 +2841,13 @@ static int alloc_mmu_pages(struct kvm_vcpu *vcpu)
         */
        page = alloc_page(GFP_KERNEL | __GFP_DMA32);
        if (!page)
-               goto error_1;
+               return -ENOMEM;
+
        vcpu->arch.mmu.pae_root = page_address(page);
        for (i = 0; i < 4; ++i)
                vcpu->arch.mmu.pae_root[i] = INVALID_PAGE;
 
        return 0;
-
-error_1:
-       free_mmu_pages(vcpu);
-       return -ENOMEM;
 }
 
 int kvm_mmu_create(struct kvm_vcpu *vcpu)
@@ -2818,10 +2927,9 @@ static int mmu_shrink(int nr_to_scan, gfp_t gfp_mask)
        spin_lock(&kvm_lock);
 
        list_for_each_entry(kvm, &vm_list, vm_list) {
-               int npages;
+               int npages, idx;
 
-               if (!down_read_trylock(&kvm->slots_lock))
-                       continue;
+               idx = srcu_read_lock(&kvm->srcu);
                spin_lock(&kvm->mmu_lock);
                npages = kvm->arch.n_alloc_mmu_pages -
                         kvm->arch.n_free_mmu_pages;
@@ -2834,7 +2942,7 @@ static int mmu_shrink(int nr_to_scan, gfp_t gfp_mask)
                nr_to_scan--;
 
                spin_unlock(&kvm->mmu_lock);
-               up_read(&kvm->slots_lock);
+               srcu_read_unlock(&kvm->srcu, idx);
        }
        if (kvm_freed)
                list_move_tail(&kvm_freed->vm_list, &vm_list);
@@ -2901,9 +3009,11 @@ unsigned int kvm_mmu_calculate_mmu_pages(struct kvm *kvm)
        int i;
        unsigned int nr_mmu_pages;
        unsigned int  nr_pages = 0;
+       struct kvm_memslots *slots;
 
-       for (i = 0; i < kvm->nmemslots; i++)
-               nr_pages += kvm->memslots[i].npages;
+       slots = rcu_dereference(kvm->memslots);
+       for (i = 0; i < slots->nmemslots; i++)
+               nr_pages += slots->memslots[i].npages;
 
        nr_mmu_pages = nr_pages * KVM_PERMILLE_MMU_PAGES / 1000;
        nr_mmu_pages = max(nr_mmu_pages,
@@ -3033,6 +3143,24 @@ out:
        return r;
 }
 
+int kvm_mmu_get_spte_hierarchy(struct kvm_vcpu *vcpu, u64 addr, u64 sptes[4])
+{
+       struct kvm_shadow_walk_iterator iterator;
+       int nr_sptes = 0;
+
+       spin_lock(&vcpu->kvm->mmu_lock);
+       for_each_shadow_entry(vcpu, addr, iterator) {
+               sptes[iterator.level-1] = *iterator.sptep;
+               nr_sptes++;
+               if (!is_shadow_present_pte(*iterator.sptep))
+                       break;
+       }
+       spin_unlock(&vcpu->kvm->mmu_lock);
+
+       return nr_sptes;
+}
+EXPORT_SYMBOL_GPL(kvm_mmu_get_spte_hierarchy);
+
 #ifdef AUDIT
 
 static const char *audit_msg;
@@ -3058,12 +3186,11 @@ static void __mmu_spte_walk(struct kvm *kvm, struct kvm_mmu_page *sp,
                u64 ent = sp->spt[i];
 
                if (is_shadow_present_pte(ent)) {
-                       if (sp->role.level > 1 && !is_large_pte(ent)) {
+                       if (!is_last_spte(ent, sp->role.level)) {
                                struct kvm_mmu_page *child;
                                child = page_header(ent & PT64_BASE_ADDR_MASK);
                                __mmu_spte_walk(kvm, child, fn);
-                       }
-                       if (sp->role.level == 1)
+                       } else
                                fn(kvm, sp, &sp->spt[i]);
                }
        }
@@ -3108,20 +3235,19 @@ static void audit_mappings_page(struct kvm_vcpu *vcpu, u64 page_pte,
                        continue;
 
                va = canonicalize(va);
-               if (level > 1) {
-                       if (ent == shadow_notrap_nonpresent_pte)
-                               printk(KERN_ERR "audit: (%s) nontrapping pte"
-                                      " in nonleaf level: levels %d gva %lx"
-                                      " level %d pte %llx\n", audit_msg,
-                                      vcpu->arch.mmu.root_level, va, level, ent);
-                       else
-                               audit_mappings_page(vcpu, ent, va, level - 1);
-               } else {
-                       gpa_t gpa = vcpu->arch.mmu.gva_to_gpa(vcpu, va);
+               if (is_shadow_present_pte(ent) && !is_last_spte(ent, level))
+                       audit_mappings_page(vcpu, ent, va, level - 1);
+               else {
+                       gpa_t gpa = kvm_mmu_gva_to_gpa_read(vcpu, va, NULL);
                        gfn_t gfn = gpa >> PAGE_SHIFT;
                        pfn_t pfn = gfn_to_pfn(vcpu->kvm, gfn);
                        hpa_t hpa = (hpa_t)pfn << PAGE_SHIFT;
 
+                       if (is_error_pfn(pfn)) {
+                               kvm_release_pfn_clean(pfn);
+                               continue;
+                       }
+
                        if (is_shadow_present_pte(ent)
                            && (ent & PT64_BASE_ADDR_MASK) != hpa)
                                printk(KERN_ERR "xx audit error: (%s) levels %d"
@@ -3157,10 +3283,12 @@ static void audit_mappings(struct kvm_vcpu *vcpu)
 static int count_rmaps(struct kvm_vcpu *vcpu)
 {
        int nmaps = 0;
-       int i, j, k;
+       int i, j, k, idx;
 
+       idx = srcu_read_lock(&kvm->srcu);
+       slots = rcu_dereference(kvm->memslots);
        for (i = 0; i < KVM_MEMORY_SLOTS; ++i) {
-               struct kvm_memory_slot *m = &vcpu->kvm->memslots[i];
+               struct kvm_memory_slot *m = &slots->memslots[i];
                struct kvm_rmap_desc *d;
 
                for (j = 0; j < m->npages; ++j) {
@@ -3183,6 +3311,7 @@ static int count_rmaps(struct kvm_vcpu *vcpu)
                        }
                }
        }
+       srcu_read_unlock(&kvm->srcu, idx);
        return nmaps;
 }
 
@@ -3208,7 +3337,8 @@ void inspect_spte_has_rmap(struct kvm *kvm, struct kvm_mmu_page *sp, u64 *sptep)
                        return;
                }
 
-               rmapp = gfn_to_rmap(kvm, rev_sp->gfns[sptep - rev_sp->spt], 0);
+               rmapp = gfn_to_rmap(kvm, rev_sp->gfns[sptep - rev_sp->spt],
+                                   is_large_pte(*sptep));
                if (!*rmapp) {
                        if (!printk_ratelimit())
                                return;
@@ -3260,20 +3390,28 @@ static void audit_write_protection(struct kvm_vcpu *vcpu)
        struct kvm_mmu_page *sp;
        struct kvm_memory_slot *slot;
        unsigned long *rmapp;
+       u64 *spte;
        gfn_t gfn;
 
        list_for_each_entry(sp, &vcpu->kvm->arch.active_mmu_pages, link) {
                if (sp->role.direct)
                        continue;
+               if (sp->unsync)
+                       continue;
 
                gfn = unalias_gfn(vcpu->kvm, sp->gfn);
                slot = gfn_to_memslot_unaliased(vcpu->kvm, sp->gfn);
                rmapp = &slot->rmap[gfn - slot->base_gfn];
-               if (*rmapp)
-                       printk(KERN_ERR "%s: (%s) shadow page has writable"
-                              " mappings: gfn %lx role %x\n",
+
+               spte = rmap_next(vcpu->kvm, rmapp, NULL);
+               while (spte) {
+                       if (*spte & PT_WRITABLE_MASK)
+                               printk(KERN_ERR "%s: (%s) shadow page has "
+                               "writable mappings: gfn %lx role %x\n",
                               __func__, audit_msg, sp->gfn,
                               sp->role.word);
+                       spte = rmap_next(vcpu->kvm, rmapp, spte);
+               }
        }
 }
 
@@ -3285,7 +3423,8 @@ static void kvm_mmu_audit(struct kvm_vcpu *vcpu, const char *msg)
        audit_msg = msg;
        audit_rmap(vcpu);
        audit_write_protection(vcpu);
-       audit_mappings(vcpu);
+       if (strcmp("pre pte write", audit_msg) != 0)
+               audit_mappings(vcpu);
        audit_writable_sptes_have_rmaps(vcpu);
        dbg = olddbg;
 }