[RFC PATCH 09/13] KVM: TDX: Pin pages via get_page() right before ADD/AUG'ed to TDs

[isaku.yamahata@xxxxxxxxx]

From: Xiaoyao Li <xiaoyao.li@xxxxxxxxx>

When kvm_faultin_pfn(), it doesn't have the info regarding which page level
will the gfn be mapped at. Hence it doesn't know to pin a 4K page or a
2M page.

Move the guest private pages pinning logic right before
TDH_MEM_PAGE_ADD/AUG() since at that time it knows the page level info.

Signed-off-by: Xiaoyao Li <xiaoyao.li@xxxxxxxxx>
---
 arch/x86/kvm/vmx/tdx.c | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index 81d88b1e63ac..2fdf3aa70c57 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -1440,12 +1440,22 @@ static void tdx_measure_page(struct kvm_tdx *kvm_tdx, hpa_t gpa, int size)
 	}
 }
 
-static void tdx_unpin_pfn(struct kvm *kvm, kvm_pfn_t pfn)
+static void tdx_unpin(struct kvm *kvm, gfn_t gfn, kvm_pfn_t pfn,
+		      enum pg_level level)
 {
-	struct page *page = pfn_to_page(pfn);
+	struct kvm_memory_slot *slot = gfn_to_memslot(kvm, gfn);
+	int i;
+
+	for (i = 0; i < KVM_PAGES_PER_HPAGE(level); i++) {
+		struct page *page = pfn_to_page(pfn + i);
 
-	put_page(page);
-	WARN_ON(!page_count(page) && to_kvm_tdx(kvm)->hkid > 0);
+		put_page(page);
+		WARN_ON(!page_count(page) && to_kvm_tdx(kvm)->hkid > 0);
+	}
+	if (kvm_slot_can_be_private(slot)) {
+		/* Private slot case */
+		return;
+	}
 }
 
 static void __tdx_sept_set_private_spte(struct kvm *kvm, gfn_t gfn,
@@ -1473,7 +1483,7 @@ static void __tdx_sept_set_private_spte(struct kvm *kvm, gfn_t gfn,
 		err = tdh_mem_page_aug(kvm_tdx->tdr.pa, gpa, tdx_level, hpa, &out);
 		if (KVM_BUG_ON(err, kvm)) {
 			pr_tdx_error(TDH_MEM_PAGE_AUG, err, &out);
-			tdx_unpin_pfn(kvm, pfn);
+			tdx_unpin(kvm, gfn, pfn, level);
 		}
 		return;
 	}
@@ -1492,7 +1502,7 @@ static void __tdx_sept_set_private_spte(struct kvm *kvm, gfn_t gfn,
 	 * always uses vcpu 0's page table and protected by vcpu->mutex).
 	 */
 	if (KVM_BUG_ON(kvm_tdx->source_pa == INVALID_PAGE, kvm)) {
-		tdx_unpin_pfn(kvm, pfn);
+		tdx_unpin(kvm, gfn, pfn, level);
 		return;
 	}
 
@@ -1501,7 +1511,7 @@ static void __tdx_sept_set_private_spte(struct kvm *kvm, gfn_t gfn,
 	err = tdh_mem_page_add(kvm_tdx->tdr.pa, gpa, tdx_level, hpa, source_pa, &out);
 	if (KVM_BUG_ON(err, kvm)) {
 		pr_tdx_error(TDH_MEM_PAGE_ADD, err, &out);
-		tdx_unpin_pfn(kvm, pfn);
+		tdx_unpin(kvm, gfn, pfn, level);
 	} else if ((kvm_tdx->source_pa & KVM_TDX_MEASURE_MEMORY_REGION))
 		tdx_measure_page(kvm_tdx, gpa, KVM_HPAGE_SIZE(level));
 
@@ -1547,7 +1557,7 @@ static void tdx_sept_drop_private_spte(
 			if (WARN_ON_ONCE(err))
 				pr_tdx_error(TDH_PHYMEM_PAGE_WBINVD, err, NULL);
 			else
-				tdx_unpin(kvm, gfn + i, pfn + i);
+				tdx_unpin(kvm, gfn + i, pfn + i, PG_LEVEL_4K);
 			hpa += PAGE_SIZE;
 		}
 	} else {
@@ -1560,7 +1570,7 @@ static void tdx_sept_drop_private_spte(
 				       false, 0);
 		spin_unlock(&kvm_tdx->seamcall_lock);
 		if (!err)
-			tdx_unpin(kvm, gfn, pfn);
+			tdx_unpin(kvm, gfn, pfn, level);
 	}
 }
 
-- 
2.25.1