KVM: MMU: fast page fault

index 29ad6f9..151fcddc 100644
@@ -1850,9 +1850,9 @@ static void link_shadow_page(u64 *sptep, struct kvm_mmu_page *sp)
 {
 	u64 spte;
 
-	spte = __pa(sp->spt)
-		| PT_PRESENT_MASK | PT_ACCESSED_MASK
-		| PT_WRITABLE_MASK | PT_USER_MASK;
+	spte = __pa(sp->spt) | PT_PRESENT_MASK | PT_WRITABLE_MASK |
+	       shadow_user_mask | shadow_x_mask | shadow_accessed_mask;
+
 	mmu_spte_set(sptep, spte);
 }
 

@@ -2541,11 +2541,7 @@ static int __direct_map(struct kvm_vcpu *vcpu, gpa_t v, int write,
 				return -ENOMEM;
 			}
 
-			mmu_spte_set(iterator.sptep,
-				     __pa(sp->spt)
-				     | PT_PRESENT_MASK | PT_WRITABLE_MASK
-				     | shadow_user_mask | shadow_x_mask
-				     | shadow_accessed_mask);
+			link_shadow_page(iterator.sptep, sp);
 		}
 	}
 	return emulate;

index 151fcddc..aee13c6 100644
@@ -1368,9 +1368,6 @@ static unsigned kvm_page_table_hashfn(gfn_t gfn)
 static void mmu_page_add_parent_pte(struct kvm_vcpu *vcpu,
 				    struct kvm_mmu_page *sp, u64 *parent_pte)
 {
-	if (!parent_pte)
-		return;
-
 	pte_list_add(vcpu, parent_pte, &sp->parent_ptes);
 }
 

@@ -1401,7 +1398,6 @@ static struct kvm_mmu_page *kvm_mmu_alloc_page(struct kvm_vcpu *vcpu,
 	list_add(&sp->link, &vcpu->kvm->arch.active_mmu_pages);
 	bitmap_zero(sp->slot_bitmap, KVM_MEM_SLOTS_NUM);
 	sp->parent_ptes = 0;
-	mmu_page_add_parent_pte(vcpu, sp, parent_pte);
 	kvm_mod_used_mmu_pages(vcpu->kvm, +1);
 	return sp;
 }

@@ -1765,12 +1761,8 @@ static struct kvm_mmu_page *kvm_mmu_get_page(struct kvm_vcpu *vcpu,
 		if (sp->unsync && kvm_sync_page_transient(vcpu, sp))
 			break;
 
-		mmu_page_add_parent_pte(vcpu, sp, parent_pte);
-		if (sp->unsync_children) {
+		if (sp->unsync_children)
 			kvm_make_request(KVM_REQ_MMU_SYNC, vcpu);
-			kvm_mmu_mark_parents_unsync(sp);
-		} else if (sp->unsync)
-			kvm_mmu_mark_parents_unsync(sp);
 
 		__clear_sp_write_flooding_count(sp);
 		trace_kvm_mmu_get_page(sp, false);

@@ -1797,6 +1789,35 @@ static struct kvm_mmu_page *kvm_mmu_get_page(struct kvm_vcpu *vcpu,
 	return sp;
 }
 
+static void link_shadow_page(u64 *sptep, struct kvm_mmu_page *sp)
+{
+	u64 spte;
+
+	spte = __pa(sp->spt) | PT_PRESENT_MASK | PT_WRITABLE_MASK |
+	       shadow_user_mask | shadow_x_mask | shadow_accessed_mask;
+
+	mmu_spte_set(sptep, spte);
+}
+
+static struct kvm_mmu_page *
+mmu_spte_establish(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn, gva_t gaddr,
+		   unsigned level, int direct, unsigned access)
+{
+	struct kvm_mmu_page *sp;
+
+	WARN_ON(is_shadow_present_pte(*spte));
+
+	sp = kvm_mmu_get_page(vcpu, gfn, gaddr, level, direct, access, spte);
+
+	link_shadow_page(spte, sp);
+	mmu_page_add_parent_pte(vcpu, sp, spte);
+
+	if (sp->unsync_children || sp->unsync)
+		kvm_mmu_mark_parents_unsync(sp);
+
+	return sp;
+}
+
 static void shadow_walk_init(struct kvm_shadow_walk_iterator *iterator,
 			     struct kvm_vcpu *vcpu, u64 addr)
 {

@@ -1846,16 +1867,6 @@ static void shadow_walk_next(struct kvm_shadow_walk_iterator *iterator)
 	return __shadow_walk_next(iterator, *iterator->sptep);
 }
 
-static void link_shadow_page(u64 *sptep, struct kvm_mmu_page *sp)
-{
-	u64 spte;
-
-	spte = __pa(sp->spt) | PT_PRESENT_MASK | PT_WRITABLE_MASK |
-	       shadow_user_mask | shadow_x_mask | shadow_accessed_mask;
-
-	mmu_spte_set(sptep, spte);
-}
-
 static void drop_large_spte(struct kvm_vcpu *vcpu, u64 *sptep)
 {
 	if (is_large_pte(*sptep)) {

@@ -1921,11 +1932,6 @@ static void kvm_mmu_page_unlink_children(struct kvm *kvm,
 		mmu_page_zap_pte(kvm, sp, sp->spt + i);
 }
 
-static void kvm_mmu_put_page(struct kvm_mmu_page *sp, u64 *parent_pte)
-{
-	mmu_page_remove_parent_pte(sp, parent_pte);
-}
-
 static void kvm_mmu_unlink_parents(struct kvm *kvm, struct kvm_mmu_page *sp)
 {
 	u64 *sptep;

@@ -2511,7 +2517,6 @@ static int __direct_map(struct kvm_vcpu *vcpu, gpa_t v, int write,
 			bool prefault)
 {
 	struct kvm_shadow_walk_iterator iterator;
-	struct kvm_mmu_page *sp;
 	int emulate = 0;
 	gfn_t pseudo_gfn;
 

@@ -2532,16 +2537,9 @@ static int __direct_map(struct kvm_vcpu *vcpu, gpa_t v, int write,
 
 			base_addr &= PT64_LVL_ADDR_MASK(iterator.level);
 			pseudo_gfn = base_addr >> PAGE_SHIFT;
-			sp = kvm_mmu_get_page(vcpu, pseudo_gfn, iterator.addr,
-					      iterator.level - 1,
-					      1, ACC_ALL, iterator.sptep);
-			if (!sp) {
-				pgprintk("nonpaging_map: ENOMEM\n");
-				kvm_release_pfn_clean(pfn);
-				return -ENOMEM;
-			}
-
-			link_shadow_page(iterator.sptep, sp);
+			mmu_spte_establish(vcpu, iterator.sptep, pseudo_gfn,
+					   iterator.addr, iterator.level - 1,
+					   1, ACC_ALL);
 		}
 	}
 	return emulate;

index df5a703..e3cf43b 100644
@@ -503,8 +503,9 @@ static u64 *FNAME(fetch)(struct kvm_vcpu *vcpu, gva_t addr,
 		sp = NULL;
 		if (!is_shadow_present_pte(*it.sptep)) {
 			table_gfn = gw->table_gfn[it.level - 2];
-			sp = kvm_mmu_get_page(vcpu, table_gfn, addr, it.level-1,
-					      false, access, it.sptep);
+			sp = mmu_spte_establish(vcpu, it.sptep, table_gfn,
+						addr, it.level-1, false,
+						access);
 		}
 
 		/*

@@ -513,9 +514,6 @@ static u64 *FNAME(fetch)(struct kvm_vcpu *vcpu, gva_t addr,
 		 */
 		if (FNAME(gpte_changed)(vcpu, gw, it.level - 1))
 			goto out_gpte_changed;
-
-		if (sp)
-			link_shadow_page(it.sptep, sp);
 	}
 
 	for (;

@@ -533,9 +531,8 @@ static u64 *FNAME(fetch)(struct kvm_vcpu *vcpu, gva_t addr,
 
 		direct_gfn = gw->gfn & ~(KVM_PAGES_PER_HPAGE(it.level) - 1);
 
-		sp = kvm_mmu_get_page(vcpu, direct_gfn, addr, it.level-1,
-				      true, direct_access, it.sptep);
-		link_shadow_page(it.sptep, sp);
+		mmu_spte_establish(vcpu, it.sptep, direct_gfn, addr, it.level-1,
+				      true, direct_access);
 	}
 
 	clear_sp_write_flooding_count(it.sptep);

@@ -548,7 +545,8 @@ static u64 *FNAME(fetch)(struct kvm_vcpu *vcpu, gva_t addr,
 
 out_gpte_changed:
 	if (sp)
-		kvm_mmu_put_page(sp, it.sptep);
+		drop_parent_pte(sp, it.sptep);
+
 	kvm_release_pfn_clean(pfn);
 	return NULL;
 }

index aee13c6..91518b6 100644
@@ -993,17 +993,25 @@ struct rmap_iterator {
  */
 static u64 *rmap_get_first(unsigned long rmap, struct rmap_iterator *iter)
 {
+	u64 *sptep;
+
 	if (!rmap)
 		return NULL;
 
 	if (!(rmap & 1)) {
 		iter->desc = NULL;
-		return (u64 *)rmap;
+		sptep = (u64 *)rmap;
+
+		goto exit;
 	}
 
 	iter->desc = (struct pte_list_desc *)(rmap & ~1ul);
 	iter->pos = 0;
-	return iter->desc->sptes[iter->pos];
+	sptep = iter->desc->sptes[iter->pos];
+
+exit:
+	WARN_ON(sptep && !is_shadow_present_pte(*sptep));
+	return sptep;
 }
 
 /*

@@ -1013,14 +1021,14 @@ static u64 *rmap_get_first(unsigned long rmap, struct rmap_iterator *iter)
  */
 static u64 *rmap_get_next(struct rmap_iterator *iter)
 {
+	u64 *sptep = NULL;
+
 	if (iter->desc) {
 		if (iter->pos < PTE_LIST_EXT - 1) {
-			u64 *sptep;
-
 			++iter->pos;
 			sptep = iter->desc->sptes[iter->pos];
 			if (sptep)
-				return sptep;
+				goto exit;
 		}
 
 		iter->desc = iter->desc->more;

@@ -1028,11 +1036,14 @@ static u64 *rmap_get_next(struct rmap_iterator *iter)
 		if (iter->desc) {
 			iter->pos = 0;
 			/* desc->sptes[0] cannot be NULL */
-			return iter->desc->sptes[iter->pos];
+			sptep = iter->desc->sptes[iter->pos];
+			goto exit;
 		}
 	}
 
-	return NULL;
+exit:
+	WARN_ON(sptep && !is_shadow_present_pte(*sptep));
+	return sptep;
 }
 
 static void drop_spte(struct kvm *kvm, u64 *sptep)

@@ -1048,7 +1059,6 @@ static int __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level
 	int write_protected = 0;
 
 	for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
-		BUG_ON(!(*sptep & PT_PRESENT_MASK));
 		rmap_printk("rmap_write_protect: spte %p %llx\n", sptep, *sptep);
 
 		if (!is_writable_pte(*sptep)) {

@@ -1123,7 +1133,6 @@ static int kvm_unmap_rmapp(struct kvm *kvm, unsigned long *rmapp,
 	int need_tlb_flush = 0;
 
 	while ((sptep = rmap_get_first(*rmapp, &iter))) {
-		BUG_ON(!(*sptep & PT_PRESENT_MASK));
 		rmap_printk("kvm_rmap_unmap_hva: spte %p %llx\n", sptep, *sptep);
 
 		drop_spte(kvm, sptep);

@@ -1147,7 +1156,6 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 	new_pfn = pte_pfn(*ptep);
 
 	for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
-		BUG_ON(!is_shadow_present_pte(*sptep));
 		rmap_printk("kvm_set_pte_rmapp: spte %p %llx\n", sptep, *sptep);
 
 		need_flush = 1;

@@ -1242,14 +1250,11 @@ static int kvm_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
 		return kvm_unmap_rmapp(kvm, rmapp, data);
 
 	for (sptep = rmap_get_first(*rmapp, &iter); sptep;
-	     sptep = rmap_get_next(&iter)) {
-		BUG_ON(!(*sptep & PT_PRESENT_MASK));
-
+	     sptep = rmap_get_next(&iter))
 		if (*sptep & PT_ACCESSED_MASK) {
 			young = 1;
 			clear_bit(PT_ACCESSED_SHIFT, (unsigned long *)sptep);
 		}
-	}
 
 	return young;
 }

@@ -1270,14 +1275,11 @@ static int kvm_test_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
 		goto out;
 
 	for (sptep = rmap_get_first(*rmapp, &iter); sptep;
-	     sptep = rmap_get_next(&iter)) {
-		BUG_ON(!(*sptep & PT_PRESENT_MASK));
-
+	     sptep = rmap_get_next(&iter))
 		if (*sptep & PT_ACCESSED_MASK) {
 			young = 1;
 			break;
 		}
-	}
 out:
 	return young;
 }

index 91518b6..7589e56 100644
@@ -1052,11 +1052,12 @@ static void drop_spte(struct kvm *kvm, u64 *sptep)
 		rmap_remove(kvm, sptep);
 }
 
-static int __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
+static bool
+__rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
 {
 	u64 *sptep;
 	struct rmap_iterator iter;
-	int write_protected = 0;
+	bool write_protected = false;
 
 	for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
 		rmap_printk("rmap_write_protect: spte %p %llx\n", sptep, *sptep);

@@ -1076,7 +1077,7 @@ static int __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level
 			sptep = rmap_get_first(*rmapp, &iter);
 		}
 
-		write_protected = 1;
+		write_protected = true;
 	}
 
 	return write_protected;

@@ -1107,12 +1108,12 @@ void kvm_mmu_write_protect_pt_masked(struct kvm *kvm,
 	}
 }
 
-static int rmap_write_protect(struct kvm *kvm, u64 gfn)
+static bool rmap_write_protect(struct kvm *kvm, u64 gfn)
 {
 	struct kvm_memory_slot *slot;
 	unsigned long *rmapp;
 	int i;
-	int write_protected = 0;
+	bool write_protected = false;
 
 	slot = gfn_to_memslot(kvm, gfn);
 

@@ -1689,7 +1690,7 @@ static void mmu_sync_children(struct kvm_vcpu *vcpu,
 
 	kvm_mmu_pages_init(parent, &parents, &pages);
 	while (mmu_unsync_walk(parent, &pages)) {
-		int protected = 0;
+		bool protected = false;
 
 		for_each_sp(pages, sp, parents, i)
 			protected |= rmap_write_protect(vcpu->kvm, sp->gfn);

index 7589e56..a1c3628 100644
@@ -1052,6 +1052,34 @@ static void drop_spte(struct kvm *kvm, u64 *sptep)
 		rmap_remove(kvm, sptep);
 }
 
+/* Return true if the spte is dropped. */
+static bool spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
+			       bool *flush)
+{
+	u64 spte = *sptep;
+
+	if (!is_writable_pte(spte))
+		return false;
+
+	*flush |= true;
+
+	if (large) {
+		pgprintk("rmap_write_protect(large): spte %p %llx\n",
+			 spte, *spte);
+		BUG_ON(!is_large_pte(spte));
+
+		drop_spte(kvm, sptep);
+		--kvm->stat.lpages;
+		return true;
+	}
+
+	rmap_printk("rmap_write_protect: spte %p %llx\n", spte, *spte);
+	spte = spte & ~PT_WRITABLE_MASK;
+	mmu_spte_update(sptep, spte);
+
+	return false;
+}
+
 static bool
 __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
 {

@@ -1060,24 +1088,13 @@ __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
 	bool write_protected = false;
 
 	for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
-		rmap_printk("rmap_write_protect: spte %p %llx\n", sptep, *sptep);
-
-		if (!is_writable_pte(*sptep)) {
-			sptep = rmap_get_next(&iter);
-			continue;
-		}
-
-		if (level == PT_PAGE_TABLE_LEVEL) {
-			mmu_spte_update(sptep, *sptep & ~PT_WRITABLE_MASK);
-			sptep = rmap_get_next(&iter);
-		} else {
-			BUG_ON(!is_large_pte(*sptep));
-			drop_spte(kvm, sptep);
-			--kvm->stat.lpages;
+		if (spte_write_protect(kvm, sptep, level > PT_PAGE_TABLE_LEVEL,
+			  &write_protected)) {
 			sptep = rmap_get_first(*rmapp, &iter);
+			continue;
 		}
 
-		write_protected = true;
+		sptep = rmap_get_next(&iter);
 	}
 
 	return write_protected;

@@ -3898,6 +3915,7 @@ int kvm_mmu_setup(struct kvm_vcpu *vcpu)
 void kvm_mmu_slot_remove_write_access(struct kvm *kvm, int slot)
 {
 	struct kvm_mmu_page *sp;
+	bool flush = false;
 
 	list_for_each_entry(sp, &kvm->arch.active_mmu_pages, link) {
 		int i;

@@ -3912,16 +3930,8 @@ void kvm_mmu_slot_remove_write_access(struct kvm *kvm, int slot)
 			      !is_last_spte(pt[i], sp->role.level))
 				continue;
 
-			if (is_large_pte(pt[i])) {
-				drop_spte(kvm, &pt[i]);
-				--kvm->stat.lpages;
-				continue;
-			}
-
-			/* avoid RMW */
-			if (is_writable_pte(pt[i]))
-				mmu_spte_update(&pt[i],
-						pt[i] & ~PT_WRITABLE_MASK);
+			spte_write_protect(kvm, &pt[i],
+					   is_large_pte(pt[i]), &flush);
 		}
 	}
 	kvm_flush_remote_tlbs(kvm);

index 52f6856..2c98057 100644
@@ -4735,6 +4735,7 @@ static int handle_ept_violation(struct kvm_vcpu *vcpu)
 {
 	unsigned long exit_qualification;
 	gpa_t gpa;
+	u32 error_code;
 	int gla_validity;
 
 	exit_qualification = vmcs_readl(EXIT_QUALIFICATION);

@@ -4759,7 +4760,13 @@ static int handle_ept_violation(struct kvm_vcpu *vcpu)
 
 	gpa = vmcs_read64(GUEST_PHYSICAL_ADDRESS);
 	trace_kvm_page_fault(gpa, exit_qualification);
-	return kvm_mmu_page_fault(vcpu, gpa, exit_qualification & 0x3, NULL, 0);
+
+	/* It is a write fault? */
+	error_code = exit_qualification & (1U << 1);
+	/* ept page table is present? */
+	error_code |= (exit_qualification >> 3) & 0x1;
+
+	return kvm_mmu_page_fault(vcpu, gpa, error_code, NULL, 0);
 }
 
 static u64 ept_rsvd_mask(u64 spte, int level)

index a1c3628..4e91e94 100644
@@ -900,26 +900,110 @@ static void pte_list_remove(u64 *spte, unsigned long *pte_list)
 	}
 }
 
-typedef void (*pte_list_walk_fn) (u64 *spte);
-static void pte_list_walk(unsigned long *pte_list, pte_list_walk_fn fn)
-{
+/*
+ * Used by the following functions to iterate through the sptes linked by a
+ * pte_list. All fields are private and not assumed to be used outside.
+ */
+struct spte_iterator {
+	/* private fields */
+	unsigned long *pte_list;
+	/* holds the sptep if not NULL */
 	struct pte_list_desc *desc;
-	int i;
+	/* index of the sptep, -1 means the walking does not start. */
+	int pos;
+};
 
-	if (!*pte_list)
-		return;
+static void pte_list_walk_init(struct spte_iterator *iter,
+			       unsigned long *pte_list)
+{
+	iter->pte_list = pte_list;
+	iter->pos = -1;
+}
+
+static void pte_list_walk_check_restart(struct spte_iterator *iter, u64 *spte)
+{
+	/*
+	 * If the spte is zapped, we should set 'iter->pos = -1' to restart
+	 * the walk.
+	 */
+	if (!is_shadow_present_pte(*spte))
+		iter->pos = -1;
+}
+
+static u64 *pte_list_first(struct spte_iterator *iter)
+{
+	unsigned long pte_list = *iter->pte_list;
+	u64 *sptep;
+
+	if (!pte_list)
+		return NULL;
+
+	if (!(pte_list & 1)) {
+		iter->desc = NULL;
+		iter->pos = 0;
+		sptep = (u64 *)pte_list;
+
+		goto exit;
+	}
+
+	iter->desc = (struct pte_list_desc *)(pte_list & ~1ul);
+	iter->pos = 0;
+	sptep = iter->desc->sptes[iter->pos];
+
+exit:
+	WARN_ON(sptep && !is_shadow_present_pte(*sptep));
+	return sptep;
+}
 
-	if (!(*pte_list & 1))
-		return fn((u64 *)*pte_list);
+static u64 *pte_list_next(struct spte_iterator *iter)
+{
+	u64 *sptep = NULL;
 
-	desc = (struct pte_list_desc *)(*pte_list & ~1ul);
-	while (desc) {
-		for (i = 0; i < PTE_LIST_EXT && desc->sptes[i]; ++i)
-			fn(desc->sptes[i]);
-		desc = desc->more;
+	if (iter->pos < 0)
+		return pte_list_first(iter);
+
+	if (iter->desc) {
+		if (iter->pos < PTE_LIST_EXT - 1) {
+			++iter->pos;
+			sptep = iter->desc->sptes[iter->pos];
+			if (sptep)
+				goto exit;
+		}
+
+		iter->desc = iter->desc->more;
+
+		if (iter->desc) {
+			iter->pos = 0;
+			/* desc->sptes[0] cannot be NULL */
+			sptep = iter->desc->sptes[iter->pos];
+			goto exit;
+		}
 	}
+
+exit:
+	WARN_ON(sptep && !is_shadow_present_pte(*sptep));
+	return sptep;
 }
 
+
+#define for_each_pte_list_spte(pte_list, iter, spte)	\
+	for (pte_list_walk_init(iter, pte_list);	\
+	      (spte = pte_list_next(iter)) != NULL;	\
+		pte_list_walk_check_restart(iter, spte))
+
+typedef void (*pte_list_walk_fn) (u64 *spte);
+static void pte_list_walk(unsigned long *pte_list, pte_list_walk_fn fn)
+{
+	struct spte_iterator iter;
+	u64 *spte;
+
+	for_each_pte_list_spte(pte_list, &iter, spte)
+		fn(spte);
+}
+
+#define for_each_rmap_spte(rmap, iter, spte)	\
+		for_each_pte_list_spte(rmap, iter, spte)
+
 static unsigned long *__gfn_to_rmap(gfn_t gfn, int level,
 				    struct kvm_memory_slot *slot)
 {

@@ -974,92 +1058,19 @@ static void rmap_remove(struct kvm *kvm, u64 *spte)
 	pte_list_remove(spte, rmapp);
 }
 
-/*
- * Used by the following functions to iterate through the sptes linked by a
- * rmap.  All fields are private and not assumed to be used outside.
- */
-struct rmap_iterator {
-	/* private fields */
-	struct pte_list_desc *desc;	/* holds the sptep if not NULL */
-	int pos;			/* index of the sptep */
-};
-
-/*
- * Iteration must be started by this function.  This should also be used after
- * removing/dropping sptes from the rmap link because in such cases the
- * information in the itererator may not be valid.
- *
- * Returns sptep if found, NULL otherwise.
- */
-static u64 *rmap_get_first(unsigned long rmap, struct rmap_iterator *iter)
-{
-	u64 *sptep;
-
-	if (!rmap)
-		return NULL;
-
-	if (!(rmap & 1)) {
-		iter->desc = NULL;
-		sptep = (u64 *)rmap;
-
-		goto exit;
-	}
-
-	iter->desc = (struct pte_list_desc *)(rmap & ~1ul);
-	iter->pos = 0;
-	sptep = iter->desc->sptes[iter->pos];
-
-exit:
-	WARN_ON(sptep && !is_shadow_present_pte(*sptep));
-	return sptep;
-}
-
-/*
- * Must be used with a valid iterator: e.g. after rmap_get_first().
- *
- * Returns sptep if found, NULL otherwise.
- */
-static u64 *rmap_get_next(struct rmap_iterator *iter)
-{
-	u64 *sptep = NULL;
-
-	if (iter->desc) {
-		if (iter->pos < PTE_LIST_EXT - 1) {
-			++iter->pos;
-			sptep = iter->desc->sptes[iter->pos];
-			if (sptep)
-				goto exit;
-		}
-
-		iter->desc = iter->desc->more;
-
-		if (iter->desc) {
-			iter->pos = 0;
-			/* desc->sptes[0] cannot be NULL */
-			sptep = iter->desc->sptes[iter->pos];
-			goto exit;
-		}
-	}
-
-exit:
-	WARN_ON(sptep && !is_shadow_present_pte(*sptep));
-	return sptep;
-}
-
 static void drop_spte(struct kvm *kvm, u64 *sptep)
 {
 	if (mmu_spte_clear_track_bits(sptep))
 		rmap_remove(kvm, sptep);
 }
 
-/* Return true if the spte is dropped. */
-static bool spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
+static void spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
 			       bool *flush)
 {
 	u64 spte = *sptep;
 
 	if (!is_writable_pte(spte))
-		return false;
+		return;
 
 	*flush |= true;
 

@@ -1070,32 +1081,24 @@ static bool spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
 
 		drop_spte(kvm, sptep);
 		--kvm->stat.lpages;
-		return true;
+		return;
 	}
 
 	rmap_printk("rmap_write_protect: spte %p %llx\n", spte, *spte);
 	spte = spte & ~PT_WRITABLE_MASK;
 	mmu_spte_update(sptep, spte);
-
-	return false;
 }
 
 static bool
 __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
 {
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 	bool write_protected = false;
 
-	for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
-		if (spte_write_protect(kvm, sptep, level > PT_PAGE_TABLE_LEVEL,
-			  &write_protected)) {
-			sptep = rmap_get_first(*rmapp, &iter);
-			continue;
-		}
-
-		sptep = rmap_get_next(&iter);
-	}
+	for_each_rmap_spte(rmapp, &iter, sptep)
+		spte_write_protect(kvm, sptep, level > PT_PAGE_TABLE_LEVEL,
+			  &write_protected);
 
 	return write_protected;
 }

@@ -1147,10 +1150,10 @@ static int kvm_unmap_rmapp(struct kvm *kvm, unsigned long *rmapp,
 			   unsigned long data)
 {
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 	int need_tlb_flush = 0;
 
-	while ((sptep = rmap_get_first(*rmapp, &iter))) {
+	for_each_rmap_spte(rmapp, &iter, sptep) {
 		rmap_printk("kvm_rmap_unmap_hva: spte %p %llx\n", sptep, *sptep);
 
 		drop_spte(kvm, sptep);

@@ -1164,7 +1167,7 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 			     unsigned long data)
 {
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 	int need_flush = 0;
 	u64 new_spte;
 	pte_t *ptep = (pte_t *)data;

@@ -1173,15 +1176,14 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 	WARN_ON(pte_huge(*ptep));
 	new_pfn = pte_pfn(*ptep);
 
-	for (sptep = rmap_get_first(*rmapp, &iter); sptep;) {
+	for_each_rmap_spte(rmapp, &iter, sptep) {
 		rmap_printk("kvm_set_pte_rmapp: spte %p %llx\n", sptep, *sptep);
 
 		need_flush = 1;
 
-		if (pte_write(*ptep)) {
+		if (pte_write(*ptep))
 			drop_spte(kvm, sptep);
-			sptep = rmap_get_first(*rmapp, &iter);
-		} else {
+		else {
 			new_spte = *sptep & ~PT64_BASE_ADDR_MASK;
 			new_spte |= (u64)new_pfn << PAGE_SHIFT;
 

@@ -1191,7 +1193,6 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 
 			mmu_spte_clear_track_bits(sptep);
 			mmu_spte_set(sptep, new_spte);
-			sptep = rmap_get_next(&iter);
 		}
 	}
 

@@ -1254,7 +1255,7 @@ static int kvm_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
 			 unsigned long data)
 {
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 	int young = 0;
 
 	/*

@@ -1267,8 +1268,7 @@ static int kvm_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
 	if (!shadow_accessed_mask)
 		return kvm_unmap_rmapp(kvm, rmapp, data);
 
-	for (sptep = rmap_get_first(*rmapp, &iter); sptep;
-	     sptep = rmap_get_next(&iter))
+	for_each_rmap_spte(rmapp, &iter, sptep)
 		if (*sptep & PT_ACCESSED_MASK) {
 			young = 1;
 			clear_bit(PT_ACCESSED_SHIFT, (unsigned long *)sptep);

@@ -1281,7 +1281,7 @@ static int kvm_test_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
 			      unsigned long data)
 {
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 	int young = 0;
 
 	/*

@@ -1292,8 +1292,7 @@ static int kvm_test_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
 	if (!shadow_accessed_mask)
 		goto out;
 
-	for (sptep = rmap_get_first(*rmapp, &iter); sptep;
-	     sptep = rmap_get_next(&iter))
+	for_each_rmap_spte(rmapp, &iter, sptep)
 		if (*sptep & PT_ACCESSED_MASK) {
 			young = 1;
 			break;

@@ -1955,9 +1954,9 @@ static void kvm_mmu_page_unlink_children(struct kvm *kvm,
 static void kvm_mmu_unlink_parents(struct kvm *kvm, struct kvm_mmu_page *sp)
 {
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 
-	while ((sptep = rmap_get_first(sp->parent_ptes, &iter)))
+	for_each_pte_list_spte(&sp->parent_ptes, &iter, sptep)
 		drop_parent_pte(sp, sptep);
 }
 

index 7d7d0b9..a42bd85 100644
@@ -193,7 +193,7 @@ static void audit_write_protection(struct kvm *kvm, struct kvm_mmu_page *sp)
 	struct kvm_memory_slot *slot;
 	unsigned long *rmapp;
 	u64 *sptep;
-	struct rmap_iterator iter;
+	struct spte_iterator iter;
 
 	if (sp->role.direct || sp->unsync || sp->role.invalid)
 		return;

@@ -201,13 +201,11 @@ static void audit_write_protection(struct kvm *kvm, struct kvm_mmu_page *sp)
 	slot = gfn_to_memslot(kvm, sp->gfn);
 	rmapp = &slot->rmap[sp->gfn - slot->base_gfn];
 
-	for (sptep = rmap_get_first(*rmapp, &iter); sptep;
-	     sptep = rmap_get_next(&iter)) {
+	for_each_rmap_spte(rmapp, &iter, sptep)
 		if (is_writable_pte(*sptep))
 			audit_printk(kvm, "shadow page has writable "
 				     "mappings: gfn %llx role %x\n",
 				     sp->gfn, sp->role.word);
-	}
 }
 
 static void audit_sp(struct kvm *kvm, struct kvm_mmu_page *sp)

index 4e91e94..53e92de 100644
@@ -796,13 +796,24 @@ static int mapping_level(struct kvm_vcpu *vcpu, gfn_t large_gfn)
 	return level - 1;
 }
 
+#define PTE_LIST_DESC		(0x1ull)
+#define PTE_LIST_FLAG_MASK	(0x3ull)
+
+static void
+pte_list_decode(const unsigned long *pte_list, unsigned long *map,
+		unsigned long *flags)
+{
+	*map = *pte_list & (~PTE_LIST_FLAG_MASK);
+	*flags = *pte_list & PTE_LIST_FLAG_MASK;
+}
+
 /*
  * Pte mapping structures:
  *
- * If pte_list bit zero is zero, then pte_list point to the spte.
+ * If PTE_LIST_DESC bit is zero, then pte_list point to the spte.
  *
- * If pte_list bit zero is one, (then pte_list & ~1) points to a struct
- * pte_list_desc containing more mappings.
+ * If PTE_LIST_DESC bit is one, (then pte_list & ~PTE_LIST_FLAG_MASK) points
+ * to a struct pte_list_desc containing more mappings.
  *
  * Returns the number of pte entries before the spte was added or zero if
  * the spte was not added.

@@ -812,39 +823,52 @@ static int pte_list_add(struct kvm_vcpu *vcpu, u64 *spte,
 			unsigned long *pte_list)
 {
 	struct pte_list_desc *desc;
+	unsigned long map, flags;
 	int i, count = 0;
 
-	if (!*pte_list) {
+	pte_list_decode(pte_list, &map, &flags);
+
+	if (!map) {
 		rmap_printk("pte_list_add: %p %llx 0->1\n", spte, *spte);
-		*pte_list = (unsigned long)spte;
-	} else if (!(*pte_list & 1)) {
+		WARN_ON(flags & PTE_LIST_DESC);
+		*pte_list = (unsigned long)spte | flags;
+
+		 return 0;
+	}
+
+	if (!(flags & PTE_LIST_DESC)) {
 		rmap_printk("pte_list_add: %p %llx 1->many\n", spte, *spte);
 		desc = mmu_alloc_pte_list_desc(vcpu);
-		desc->sptes[0] = (u64 *)*pte_list;
+		desc->sptes[0] = (u64 *)map;
 		desc->sptes[1] = spte;
-		*pte_list = (unsigned long)desc | 1;
-		++count;
-	} else {
-		rmap_printk("pte_list_add: %p %llx many->many\n", spte, *spte);
-		desc = (struct pte_list_desc *)(*pte_list & ~1ul);
-		while (desc->sptes[PTE_LIST_EXT-1] && desc->more) {
-			desc = desc->more;
-			count += PTE_LIST_EXT;
-		}
-		if (desc->sptes[PTE_LIST_EXT-1]) {
-			desc->more = mmu_alloc_pte_list_desc(vcpu);
-			desc = desc->more;
-		}
-		for (i = 0; desc->sptes[i]; ++i)
-			++count;
-		desc->sptes[i] = spte;
+		*pte_list = (unsigned long)desc | flags | PTE_LIST_DESC;
+
+		return 1;
+	}
+
+	rmap_printk("pte_list_add: %p %llx many->many\n", spte, *spte);
+	desc = (struct pte_list_desc *)map;
+	while (desc->sptes[PTE_LIST_EXT - 1] && desc->more) {
+		desc = desc->more;
+		count += PTE_LIST_EXT;
 	}
+
+	if (desc->sptes[PTE_LIST_EXT - 1]) {
+		desc->more = mmu_alloc_pte_list_desc(vcpu);
+		desc = desc->more;
+	}
+
+	for (i = 0; desc->sptes[i]; ++i)
+		++count;
+	desc->sptes[i] = spte;
+
 	return count;
 }
 
 static void
 pte_list_desc_remove_entry(unsigned long *pte_list, struct pte_list_desc *desc,
-			   int i, struct pte_list_desc *prev_desc)
+			   int i, struct pte_list_desc *prev_desc,
+			   unsigned long flags)
 {
 	int j;
 

@@ -855,12 +879,13 @@ pte_list_desc_remove_entry(unsigned long *pte_list, struct pte_list_desc *desc,
 	if (j != 0)
 		return;
 	if (!prev_desc && !desc->more)
-		*pte_list = (unsigned long)desc->sptes[0];
+		*pte_list = (unsigned long)desc->sptes[0] |
+					(flags & (~PTE_LIST_DESC)) ;
 	else
 		if (prev_desc)
 			prev_desc->more = desc->more;
 		else
-			*pte_list = (unsigned long)desc->more | 1;
+			*pte_list = (unsigned long)desc->more | flags;
 	mmu_free_pte_list_desc(desc);
 }
 

@@ -868,36 +893,42 @@ static void pte_list_remove(u64 *spte, unsigned long *pte_list)
 {
 	struct pte_list_desc *desc;
 	struct pte_list_desc *prev_desc;
+	unsigned long map, flags;
 	int i;
 
-	if (!*pte_list) {
+	pte_list_decode(pte_list, &map, &flags);
+
+	if (!map) {
 		printk(KERN_ERR "pte_list_remove: %p 0->BUG\n", spte);
 		BUG();
-	} else if (!(*pte_list & 1)) {
+		return;
+	}
+
+	if (!(flags & PTE_LIST_DESC)) {
 		rmap_printk("pte_list_remove:  %p 1->0\n", spte);
-		if ((u64 *)*pte_list != spte) {
+		if ((u64 *)map != spte) {
 			printk(KERN_ERR "pte_list_remove:  %p 1->BUG\n", spte);
 			BUG();
 		}
-		*pte_list = 0;
-	} else {
-		rmap_printk("pte_list_remove:  %p many->many\n", spte);
-		desc = (struct pte_list_desc *)(*pte_list & ~1ul);
-		prev_desc = NULL;
-		while (desc) {
-			for (i = 0; i < PTE_LIST_EXT && desc->sptes[i]; ++i)
-				if (desc->sptes[i] == spte) {
-					pte_list_desc_remove_entry(pte_list,
-							       desc, i,
-							       prev_desc);
-					return;
-				}
-			prev_desc = desc;
-			desc = desc->more;
-		}
-		pr_err("pte_list_remove: %p many->many\n", spte);
-		BUG();
+		*pte_list = flags;
+		return;
+	}
+
+	rmap_printk("pte_list_remove:  %p many->many\n", spte);
+	desc = (struct pte_list_desc *)map;
+	prev_desc = NULL;
+	while (desc) {
+		for (i = 0; i < PTE_LIST_EXT && desc->sptes[i]; ++i)
+			if (desc->sptes[i] == spte) {
+				pte_list_desc_remove_entry(pte_list,
+						desc, i, prev_desc, flags);
+				return;
+			}
+		prev_desc = desc;
+		desc = desc->more;
 	}
+	pr_err("pte_list_remove: %p many->many\n", spte);
+	BUG();
 }
 
 /*

@@ -932,21 +963,22 @@ static void pte_list_walk_check_restart(struct spte_iterator *iter, u64 *spte)
 
 static u64 *pte_list_first(struct spte_iterator *iter)
 {
-	unsigned long pte_list = *iter->pte_list;
+	unsigned long map, flags;
 	u64 *sptep;
 
-	if (!pte_list)
+	pte_list_decode(iter->pte_list, &map, &flags);
+
+	if (!map)
 		return NULL;
 
-	if (!(pte_list & 1)) {
+	if (!(flags & PTE_LIST_DESC)) {
 		iter->desc = NULL;
 		iter->pos = 0;
-		sptep = (u64 *)pte_list;
-
+		sptep = (u64 *)map;
 		goto exit;
 	}
 
-	iter->desc = (struct pte_list_desc *)(pte_list & ~1ul);
+	iter->desc = (struct pte_list_desc *)map;
 	iter->pos = 0;
 	sptep = iter->desc->sptes[iter->pos];
 

index 53e92de..0c6e92d 100644
@@ -2293,6 +2293,9 @@ static int mmu_need_write_protect(struct kvm_vcpu *vcpu, gfn_t gfn,
 	struct hlist_node *node;
 	bool need_unsync = false;
 
+	if (!vcpu->kvm->arch.indirect_shadow_pages)
+		return 0;
+
 	for_each_gfn_indirect_valid_sp(vcpu->kvm, s, gfn, node) {
 		if (!can_unsync)
 			return 1;

index 0c6e92d..8b71908 100644
@@ -796,7 +796,9 @@ static int mapping_level(struct kvm_vcpu *vcpu, gfn_t large_gfn)
 	return level - 1;
 }
 
-#define PTE_LIST_DESC		(0x1ull)
+#define PTE_LIST_DESC_BIT	0
+#define PTE_LIST_WP_BIT	1
+#define PTE_LIST_DESC		(1 << PTE_LIST_DESC_BIT)
 #define PTE_LIST_FLAG_MASK	(0x3ull)
 
 static void

@@ -1067,6 +1069,12 @@ static bool rmap_can_add(struct kvm_vcpu *vcpu)
 	return mmu_memory_cache_free_objects(cache);
 }
 
+static void host_page_write_protect(u64 *spte, unsigned long *rmapp)
+{
+	if (!(*spte & SPTE_HOST_WRITEABLE))
+		__test_and_set_bit(PTE_LIST_WP_BIT, rmapp);
+}
+
 static int rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
 {
 	struct kvm_mmu_page *sp;

@@ -1074,7 +1082,10 @@ static int rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
 
 	sp = page_header(__pa(spte));
 	kvm_mmu_page_set_gfn(sp, spte - sp->spt, gfn);
+
 	rmapp = gfn_to_rmap(vcpu->kvm, gfn, sp->role.level);
+	host_page_write_protect(spte, rmapp);
+
 	return pte_list_add(vcpu, spte, rmapp);
 }
 

@@ -1164,16 +1175,23 @@ static bool rmap_write_protect(struct kvm *kvm, u64 gfn)
 {
 	struct kvm_memory_slot *slot;
 	unsigned long *rmapp;
-	int i;
+	int i = PT_PAGE_TABLE_LEVEL;
 	bool write_protected = false;
 
 	slot = gfn_to_memslot(kvm, gfn);
+	rmapp = __gfn_to_rmap(gfn, i, slot);
+
+	if (__test_and_set_bit(PTE_LIST_WP_BIT, rmapp))
+		return false;
+
+	do {
+		write_protected |= __rmap_write_protect(kvm, rmapp, i++);
+
+		if (i >= PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES)
+			break;
 
-	for (i = PT_PAGE_TABLE_LEVEL;
-	     i < PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES; ++i) {
 		rmapp = __gfn_to_rmap(gfn, i, slot);
-		write_protected |= __rmap_write_protect(kvm, rmapp, i);
-	}
+	} while (true);
 
 	return write_protected;
 }

@@ -1225,6 +1243,7 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 
 			mmu_spte_clear_track_bits(sptep);
 			mmu_spte_set(sptep, new_spte);
+			host_page_write_protect(sptep, rmapp);
 		}
 	}
 

@@ -2291,9 +2310,15 @@ static int mmu_need_write_protect(struct kvm_vcpu *vcpu, gfn_t gfn,
 {
 	struct kvm_mmu_page *s;
 	struct hlist_node *node;
+	unsigned long *rmap;
 	bool need_unsync = false;
 
+	rmap = gfn_to_rmap(vcpu->kvm, gfn, PT_PAGE_TABLE_LEVEL);
+
 	if (!vcpu->kvm->arch.indirect_shadow_pages)
+		goto write_free;
+
+	if (!test_bit(PTE_LIST_WP_BIT, rmap))
 		return 0;
 
 	for_each_gfn_indirect_valid_sp(vcpu->kvm, s, gfn, node) {

@@ -2309,6 +2334,9 @@ static int mmu_need_write_protect(struct kvm_vcpu *vcpu, gfn_t gfn,
 	}
 	if (need_unsync)
 		kvm_unsync_pages(vcpu, gfn);
+
+write_free:
+	__clear_bit(PTE_LIST_WP_BIT, rmap);
 	return 0;
 }
 

index 8b71908..1a06776 100644
@@ -145,7 +145,8 @@ module_param(dbg, bool, 0644);
 #define CREATE_TRACE_POINTS
 #include "mmutrace.h"
 
-#define SPTE_HOST_WRITEABLE (1ULL << PT_FIRST_AVAIL_BITS_SHIFT)
+#define SPTE_HOST_WRITEABLE	(1ULL << PT_FIRST_AVAIL_BITS_SHIFT)
+#define SPTE_ALLOW_WRITE	(1ULL << (PT_FIRST_AVAIL_BITS_SHIFT + 1))
 
 #define SHADOW_PT_INDEX(addr, level) PT64_INDEX(addr, level)
 

@@ -1237,9 +1238,8 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 			new_spte = *sptep & ~PT64_BASE_ADDR_MASK;
 			new_spte |= (u64)new_pfn << PAGE_SHIFT;
 
-			new_spte &= ~PT_WRITABLE_MASK;
-			new_spte &= ~SPTE_HOST_WRITEABLE;
-			new_spte &= ~shadow_accessed_mask;
+			new_spte &= ~(PT_WRITABLE_MASK | SPTE_HOST_WRITEABLE |
+				      shadow_accessed_mask | SPTE_ALLOW_WRITE);
 
 			mmu_spte_clear_track_bits(sptep);
 			mmu_spte_set(sptep, new_spte);

@@ -2386,7 +2386,7 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 			goto done;
 		}
 
-		spte |= PT_WRITABLE_MASK;
+		spte |= PT_WRITABLE_MASK | SPTE_ALLOW_WRITE;
 
 		if (!vcpu->arch.mmu.direct_map
 		    && !(pte_access & ACC_WRITE_MASK)) {

@@ -2415,8 +2415,7 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 				 __func__, gfn);
 			ret = 1;
 			pte_access &= ~ACC_WRITE_MASK;
-			if (is_writable_pte(spte))
-				spte &= ~PT_WRITABLE_MASK;
+			spte &= ~PT_WRITABLE_MASK;
 		}
 	}
 

index 1a06776..578a1e2 100644
@@ -147,6 +147,7 @@ module_param(dbg, bool, 0644);
 
 #define SPTE_HOST_WRITEABLE	(1ULL << PT_FIRST_AVAIL_BITS_SHIFT)
 #define SPTE_ALLOW_WRITE	(1ULL << (PT_FIRST_AVAIL_BITS_SHIFT + 1))
+#define SPTE_WRITE_PROTECT	(1ULL << (PT_FIRST_AVAIL_BITS_SHIFT + 2))
 
 #define SHADOW_PT_INDEX(addr, level) PT64_INDEX(addr, level)
 

@@ -1108,33 +1109,49 @@ static void drop_spte(struct kvm *kvm, u64 *sptep)
 		rmap_remove(kvm, sptep);
 }
 
+static bool spte_wp_by_dirty_log(u64 spte)
+{
+	WARN_ON(is_writable_pte(spte));
+
+	return (spte & SPTE_ALLOW_WRITE) && !(spte & SPTE_WRITE_PROTECT);
+}
+
 static void spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
-			       bool *flush)
+			       bool *flush, bool page_table_protect)
 {
 	u64 spte = *sptep;
 
-	if (!is_writable_pte(spte))
-		return;
+	if (is_writable_pte(spte)) {
+		*flush |= true;
 
-	*flush |= true;
+		if (large) {
+			pgprintk("rmap_write_protect(large): spte %p %llx\n",
+				 spte, *spte);
+			BUG_ON(!is_large_pte(spte));
 
-	if (large) {
-		pgprintk("rmap_write_protect(large): spte %p %llx\n",
-			 spte, *spte);
-		BUG_ON(!is_large_pte(spte));
+			drop_spte(kvm, sptep);
+			--kvm->stat.lpages;
+			return;
+		}
 
-		drop_spte(kvm, sptep);
-		--kvm->stat.lpages;
-		return;
+		goto reset_spte;
 	}
 
+	if (page_table_protect && spte_wp_by_dirty_log(spte))
+		goto reset_spte;
+
+	return;
+
+reset_spte:
 	rmap_printk("rmap_write_protect: spte %p %llx\n", spte, *spte);
 	spte = spte & ~PT_WRITABLE_MASK;
+	if (page_table_protect)
+		spte |= SPTE_WRITE_PROTECT;
 	mmu_spte_update(sptep, spte);
 }
 
-static bool
-__rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
+static bool __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp,
+				 int level, bool page_table_protect)
 {
 	u64 *sptep;
 	struct spte_iterator iter;

@@ -1142,7 +1159,7 @@ __rmap_write_protect(struct kvm *kvm, unsigned long *rmapp, int level)
 
 	for_each_rmap_spte(rmapp, &iter, sptep)
 		spte_write_protect(kvm, sptep, level > PT_PAGE_TABLE_LEVEL,
-			  &write_protected);
+			  &write_protected, page_table_protect);
 
 	return write_protected;
 }

@@ -1165,7 +1182,7 @@ void kvm_mmu_write_protect_pt_masked(struct kvm *kvm,
 
 	while (mask) {
 		rmapp = &slot->rmap[gfn_offset + __ffs(mask)];
-		__rmap_write_protect(kvm, rmapp, PT_PAGE_TABLE_LEVEL);
+		__rmap_write_protect(kvm, rmapp, PT_PAGE_TABLE_LEVEL, false);
 
 		/* clear the first set bit */
 		mask &= mask - 1;

@@ -1186,7 +1203,7 @@ static bool rmap_write_protect(struct kvm *kvm, u64 gfn)
 		return false;
 
 	do {
-		write_protected |= __rmap_write_protect(kvm, rmapp, i++);
+		write_protected |= __rmap_write_protect(kvm, rmapp, i++, true);
 
 		if (i >= PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES)
 			break;

@@ -1239,7 +1256,8 @@ static int kvm_set_pte_rmapp(struct kvm *kvm, unsigned long *rmapp,
 			new_spte |= (u64)new_pfn << PAGE_SHIFT;
 
 			new_spte &= ~(PT_WRITABLE_MASK | SPTE_HOST_WRITEABLE |
-				      shadow_accessed_mask | SPTE_ALLOW_WRITE);
+				      shadow_accessed_mask | SPTE_ALLOW_WRITE |
+				      SPTE_WRITE_PROTECT);
 
 			mmu_spte_clear_track_bits(sptep);
 			mmu_spte_set(sptep, new_spte);

@@ -2416,6 +2434,7 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 			ret = 1;
 			pte_access &= ~ACC_WRITE_MASK;
 			spte &= ~PT_WRITABLE_MASK;
+			spte |= SPTE_WRITE_PROTECT;
 		}
 	}
 

@@ -3992,7 +4011,7 @@ void kvm_mmu_slot_remove_write_access(struct kvm *kvm, int slot)
 				continue;
 
 			spte_write_protect(kvm, &pt[i],
-					   is_large_pte(pt[i]), &flush);
+					   is_large_pte(pt[i]), &flush, false);
 		}
 	}
 	kvm_flush_remote_tlbs(kvm);

index 578a1e2..efa5d59 100644
@@ -2323,15 +2323,45 @@ static void kvm_unsync_pages(struct kvm_vcpu *vcpu,  gfn_t gfn)
 	}
 }
 
+/*
+ * If the gfn become write-free, we make all sptes which point to this
+ * gfn to be writable.
+ * Note: we should call mark_page_dirty for the gfn later.
+ */
+static void rmap_break_page_table_wp(struct kvm_memory_slot *slot, gfn_t gfn)
+{
+	struct spte_iterator iter;
+	u64 *sptep;
+	int i;
+
+	for (i = PT_PAGE_TABLE_LEVEL;
+	      i < PT_PAGE_TABLE_LEVEL + KVM_NR_PAGE_SIZES; i++) {
+		unsigned long *rmap = __gfn_to_rmap(gfn, i, slot);
+
+		for_each_rmap_spte(rmap, &iter, sptep) {
+			u64 spte = *sptep;
+
+			if (!is_writable_pte(spte) &&
+			      (spte & SPTE_ALLOW_WRITE)) {
+				spte &= ~SPTE_WRITE_PROTECT;
+				spte |= PT_WRITABLE_MASK;
+				mmu_spte_update(sptep, spte);
+			}
+		}
+	}
+}
+
 static int mmu_need_write_protect(struct kvm_vcpu *vcpu, gfn_t gfn,
 				  bool can_unsync)
 {
+	struct kvm_memory_slot *slot;
 	struct kvm_mmu_page *s;
 	struct hlist_node *node;
 	unsigned long *rmap;
 	bool need_unsync = false;
 
-	rmap = gfn_to_rmap(vcpu->kvm, gfn, PT_PAGE_TABLE_LEVEL);
+	slot = gfn_to_memslot(vcpu->kvm, gfn);
+	rmap =	__gfn_to_rmap(gfn, PT_PAGE_TABLE_LEVEL, slot);
 
 	if (!vcpu->kvm->arch.indirect_shadow_pages)
 		goto write_free;

@@ -2353,6 +2383,8 @@ static int mmu_need_write_protect(struct kvm_vcpu *vcpu, gfn_t gfn,
 	if (need_unsync)
 		kvm_unsync_pages(vcpu, gfn);
 
+	rmap_break_page_table_wp(slot, gfn);
+
 write_free:
 	__clear_bit(PTE_LIST_WP_BIT, rmap);
 	return 0;

index efa5d59..fc91667 100644
@@ -446,6 +446,13 @@ static bool __check_direct_spte_mmio_pf(u64 spte)
 }
 #endif
 
+static bool spte_wp_by_dirty_log(u64 spte)
+{
+	WARN_ON(is_writable_pte(spte));
+
+	return (spte & SPTE_ALLOW_WRITE) && !(spte & SPTE_WRITE_PROTECT);
+}
+
 static bool spte_has_volatile_bits(u64 spte)
 {
 	if (!shadow_accessed_mask)

@@ -454,9 +461,18 @@ static bool spte_has_volatile_bits(u64 spte)
 	if (!is_shadow_present_pte(spte))
 		return false;
 
-	if ((spte & shadow_accessed_mask) &&
-	      (!is_writable_pte(spte) || (spte & shadow_dirty_mask)))
-		return false;
+	if (spte & shadow_accessed_mask) {
+		if (is_writable_pte(spte))
+			return !(spte & shadow_dirty_mask);
+
+		/*
+		 * If the spte is write-protected by dirty-log, it may
+		 * be marked writable on fast page fault path, then CPU
+		 * can modify the Dirty bit.
+		 */
+		if (!spte_wp_by_dirty_log(spte))
+			return false;
+	}
 
 	return true;
 }

@@ -1109,26 +1125,18 @@ static void drop_spte(struct kvm *kvm, u64 *sptep)
 		rmap_remove(kvm, sptep);
 }
 
-static bool spte_wp_by_dirty_log(u64 spte)
-{
-	WARN_ON(is_writable_pte(spte));
-
-	return (spte & SPTE_ALLOW_WRITE) && !(spte & SPTE_WRITE_PROTECT);
-}
-
 static void spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
 			       bool *flush, bool page_table_protect)
 {
 	u64 spte = *sptep;
 
 	if (is_writable_pte(spte)) {
-		*flush |= true;
-
 		if (large) {
 			pgprintk("rmap_write_protect(large): spte %p %llx\n",
 				 spte, *spte);
 			BUG_ON(!is_large_pte(spte));
 
+			*flush |= true;
 			drop_spte(kvm, sptep);
 			--kvm->stat.lpages;
 			return;

@@ -1137,6 +1145,9 @@ static void spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
 		goto reset_spte;
 	}
 
+	/* We need flush tlbs in this case: the fast page fault path
+	 * can mark the spte writable after we read the sptep.
+	 */
 	if (page_table_protect && spte_wp_by_dirty_log(spte))
 		goto reset_spte;
 

@@ -1144,6 +1155,8 @@ static void spte_write_protect(struct kvm *kvm, u64 *sptep, bool large,
 
 reset_spte:
 	rmap_printk("rmap_write_protect: spte %p %llx\n", spte, *spte);
+
+	*flush |= true;
 	spte = spte & ~PT_WRITABLE_MASK;
 	if (page_table_protect)
 		spte |= SPTE_WRITE_PROTECT;

@@ -2767,18 +2780,172 @@ exit:
 	return ret;
 }
 
+static bool page_fault_can_be_fast(struct kvm_vcpu *vcpu, gfn_t gfn,
+				   u32 error_code)
+{
+	unsigned long *rmap;
+
+	/*
+	 * #PF can be fast only if the shadow page table is present and it
+	 * is caused by write-protect, that means we just need change the
+	 * W bit of the spte which can be done out of mmu-lock.
+	 */
+	if (!(error_code & PFERR_PRESENT_MASK) ||
+	      !(error_code & PFERR_WRITE_MASK))
+		return false;
+
+	rmap = gfn_to_rmap(vcpu->kvm, gfn, PT_PAGE_TABLE_LEVEL);
+
+	/* Quickly check the page can be writable. */
+	if (test_bit(PTE_LIST_WP_BIT, ACCESS_ONCE(rmap)))
+		return false;
+
+	return true;
+}
+
+static bool
+fast_pf_fix_indirect_spte(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp,
+			  u64 *sptep, u64 spte, gfn_t gfn)
+{
+	pfn_t pfn;
+	bool ret = false;
+
+	/*
+	 * For the indirect spte, it is hard to get a stable gfn since
+	 * we just use a cmpxchg to avoid all the races which is not
+	 * enough to avoid the ABA problem: the host can arbitrarily
+	 * change spte and the mapping from gfn to pfh.
+	 *
+	 * What we do is call gfn_to_pfn_atomic to bind the gfn and the
+	 * pfn because after the call:
+	 * - we have held the refcount of pfn that means the pfn can not
+	 *   be freed and be reused for another gfn.
+	 * - the pfn is writable that means it can not be shared by different
+	 *   gfn.
+	 */
+	pfn = gfn_to_pfn_atomic(vcpu->kvm, gfn);
+
+	/* The host page is swapped out or merged. */
+	if (mmu_invalid_pfn(pfn))
+		goto exit;
+
+	ret = true;
+
+	if (pfn != spte_to_pfn(spte))
+		goto exit;
+
+	if (cmpxchg64(sptep, spte, spte | PT_WRITABLE_MASK) == spte)
+		mark_page_dirty(vcpu->kvm, gfn);
+
+exit:
+	kvm_release_pfn_clean(pfn);
+	return ret;
+}
+
+static bool
+fast_pf_fix_direct_spte(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp,
+			u64 *sptep, u64 spte)
+{
+	gfn_t gfn;
+
+	WARN_ON(!sp->role.direct);
+
+	/*
+	 * The gfn of direct spte is stable since it is calculated
+	 * by sp->gfn.
+	 */
+	gfn = kvm_mmu_page_get_gfn(sp, sptep - sp->spt);
+
+	if (cmpxchg64(sptep, spte, spte | PT_WRITABLE_MASK) == spte)
+		mark_page_dirty(vcpu->kvm, gfn);
+
+	return true;
+}
+
+/*
+ * Return value:
+ * - true: let the vcpu to access on the same address again.
+ * - false: let the real page fault path to fix it.
+ */
+static bool fast_page_fault(struct kvm_vcpu *vcpu, gva_t gva, gfn_t gfn,
+			    int level, u32 error_code)
+{
+	struct kvm_shadow_walk_iterator iterator;
+	struct kvm_mmu_page *sp;
+	bool ret = false;
+	u64 spte = 0ull;
+
+	if (!page_fault_can_be_fast(vcpu, gfn, error_code))
+		return false;
+
+	walk_shadow_page_lockless_begin(vcpu);
+	for_each_shadow_entry_lockless(vcpu, gva, iterator, spte)
+		if (!is_shadow_present_pte(spte) || iterator.level < level)
+			break;
+
+	/*
+	 * If the mapping has been changed, let the vcpu fault on the
+	 * same address again.
+	 */
+	if (!is_rmap_spte(spte)) {
+		ret = true;
+		goto exit;
+	}
+
+	if (!is_last_spte(spte, level))
+		goto exit;
+
+	/*
+	 * Check if it is a spurious fault caused by TLB lazily flushed.
+	 *
+	 * Need not check the access of upper level table entries since
+	 * they are always ACC_ALL.
+	 */
+	 if (is_writable_pte(spte)) {
+		ret = true;
+		goto exit;
+	}
+
+	/*
+	 * If the page fault is caused by write but host do not allow
+	 * to write the page, we need cow the host page.
+	 */
+	if (!(spte & SPTE_ALLOW_WRITE))
+		goto exit;
+
+	/*
+	 * Currently, to simplify the code, only the spte write-protected
+	 * by dirty-log can be fast fixed.
+	 */
+	if (!spte_wp_by_dirty_log(spte))
+		goto exit;
+
+	sp = page_header(__pa(iterator.sptep));
+
+	if (sp->role.direct)
+		ret = fast_pf_fix_direct_spte(vcpu, sp, iterator.sptep, spte);
+	else
+		ret = fast_pf_fix_indirect_spte(vcpu, sp, iterator.sptep,
+						spte, gfn);
+
+exit:
+	walk_shadow_page_lockless_end(vcpu);
+
+	return ret;
+}
+
 static bool try_async_pf(struct kvm_vcpu *vcpu, bool prefault, gfn_t gfn,
 			 gva_t gva, pfn_t *pfn, bool write, bool *writable);
 
-static int nonpaging_map(struct kvm_vcpu *vcpu, gva_t v, int write, gfn_t gfn,
-			 bool prefault)
+static int nonpaging_map(struct kvm_vcpu *vcpu, gva_t v, u32 error_code,
+			 gfn_t gfn, bool prefault)
 {
 	int r;
 	int level;
 	int force_pt_level;
 	pfn_t pfn;
 	unsigned long mmu_seq;
-	bool map_writable;
+	bool map_writable, write = error_code & PFERR_WRITE_MASK;
 
 	force_pt_level = mapping_level_dirty_bitmap(vcpu, gfn);
 	if (likely(!force_pt_level)) {

@@ -2795,6 +2962,9 @@ static int nonpaging_map(struct kvm_vcpu *vcpu, gva_t v, int write, gfn_t gfn,
 	} else
 		level = PT_PAGE_TABLE_LEVEL;
 
+	if (fast_page_fault(vcpu, v, gfn, level, error_code))
+		return 0;
+
 	mmu_seq = vcpu->kvm->mmu_notifier_seq;
 	smp_rmb();
 

@@ -3183,7 +3353,7 @@ static int nonpaging_page_fault(struct kvm_vcpu *vcpu, gva_t gva,
 	gfn = gva >> PAGE_SHIFT;
 
 	return nonpaging_map(vcpu, gva & PAGE_MASK,
-			     error_code & PFERR_WRITE_MASK, gfn, prefault);
+			     error_code, gfn, prefault);
 }
 
 static int kvm_arch_setup_async_pf(struct kvm_vcpu *vcpu, gva_t gva, gfn_t gfn)

@@ -3263,6 +3433,9 @@ static int tdp_page_fault(struct kvm_vcpu *vcpu, gva_t gpa, u32 error_code,
 	} else
 		level = PT_PAGE_TABLE_LEVEL;
 
+	if (fast_page_fault(vcpu, gpa, gfn, level, error_code))
+		return 0;
+
 	mmu_seq = vcpu->kvm->mmu_notifier_seq;
 	smp_rmb();
 

index e3cf43b..c85dbbe 100644
@@ -615,6 +615,9 @@ static int FNAME(page_fault)(struct kvm_vcpu *vcpu, gva_t addr, u32 error_code,
 		walker.gfn = walker.gfn & ~(KVM_PAGES_PER_HPAGE(level) - 1);
 	}
 
+	if (fast_page_fault(vcpu, addr, walker.gfn, level, error_code))
+		return 0;
+
 	mmu_seq = vcpu->kvm->mmu_notifier_seq;
 	smp_rmb();
 

index fc91667..589b507 100644
@@ -2929,6 +2929,8 @@ static bool fast_page_fault(struct kvm_vcpu *vcpu, gva_t gva, gfn_t gfn,
 						spte, gfn);
 
 exit:
+	trace_fast_page_fault(vcpu, gva, gfn, error_code, iterator.sptep,
+			      spte, ret);
 	walk_shadow_page_lockless_end(vcpu);
 
 	return ret;