[PATCH v3 05/24] x86/virt/tdx: Enhance tdh_phymem_page_reclaim() to support huge pages

[Yan Zhao <yan.y.zhao@intel.com>]

Enhance the SEAMCALL wrapper tdh_phymem_page_reclaim() to support huge
pages by introducing new parameters: "folio", "start_idx", and "npages".
These parameters specify the physical memory to be reclaimed, starting from
the page at "start_idx" within a folio and spanning "npages" contiguous
PFNs. The specified memory must be entirely contained within a
single folio. Return TDX_SW_ERROR if the size of the reclaimed memory does
not match the specified size.

On the KVM side, introduce tdx_reclaim_folio() to invoke
tdh_phymem_page_reclaim() for reclaiming huge guest private pages. The
"reset" parameter in tdx_reclaim_folio() specifies whether
tdx_quirk_reset_folio() should be subsequently invoked within
tdx_reclaim_folio().  To facilitate reclaiming of 4KB pages, keep function
tdx_reclaim_page() and make it a helper for reclaiming normal TDX control
pages, and introduce a new helper tdx_reclaim_page_noreset() for reclaiming
the TDR page.

Opportunistically, rename rcx, rdx, r8 to tdx_pt, tdx_owner, tdx_size in
tdx_reclaim_folio() to improve readability.

Signed-off-by: Xiaoyao Li <xiaoyao.li@intel.com>
Signed-off-by: Isaku Yamahata <isaku.yamahata@intel.com>
Signed-off-by: Yan Zhao <yan.y.zhao@intel.com>
---
v3:
- Rebased to Sean's cleanup series. Dropped invoking tdx_reclaim_folio()
  in tdx_sept_remove_private_spte() due to no reclaiming is required in
  that path.
  However, keep introducing tdx_reclaim_folio() as it will be needed when
  the patches of removing guest private memory after releasing HKID are
  merged.
- tdx_reclaim_page_noclear() --> tdx_reclaim_page_noreset() and invoke
  tdx_quirk_reset_folio() instead in tdx_reclaim_folio() due to rebase.
- Check mismatch between the request size and the reclaimed size, and
  return TDX_SW_ERROR only after a successful TDH_PHYMEM_PAGE_RECLAIM.
  (Binbin)

RFC v2:
- Introduce new params "folio", "start_idx" and "npages" to wrapper
  tdh_phymem_page_reclaim().
- Move the checking of return size from KVM to x86/virt and return error.
- Rename tdx_reclaim_page() to tdx_reclaim_folio().
- Add two helper functions tdx_reclaim_page() tdx_reclaim_page_noclear()
  to faciliate the reclaiming of 4KB pages.

RFC v1:
- Rebased and split patch.
---
 arch/x86/include/asm/tdx.h  |  3 ++-
 arch/x86/kvm/vmx/tdx.c      | 27 +++++++++++++++++----------
 arch/x86/virt/vmx/tdx/tdx.c | 12 ++++++++++--
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/arch/x86/include/asm/tdx.h b/arch/x86/include/asm/tdx.h
index 669dd6d99821..abe484045132 100644
--- a/arch/x86/include/asm/tdx.h
+++ b/arch/x86/include/asm/tdx.h
@@ -261,7 +261,8 @@ u64 tdh_mng_init(struct tdx_td *td, u64 td_params, u64 *extended_err);
 u64 tdh_vp_init(struct tdx_vp *vp, u64 initial_rcx, u32 x2apicid);
 u64 tdh_vp_rd(struct tdx_vp *vp, u64 field, u64 *data);
 u64 tdh_vp_wr(struct tdx_vp *vp, u64 field, u64 data, u64 mask);
-u64 tdh_phymem_page_reclaim(struct page *page, u64 *tdx_pt, u64 *tdx_owner, u64 *tdx_size);
+u64 tdh_phymem_page_reclaim(struct folio *folio, unsigned long start_idx, unsigned long npages,
+			    u64 *tdx_pt, u64 *tdx_owner, u64 *tdx_size);
 u64 tdh_mem_track(struct tdx_td *tdr);
 u64 tdh_mem_page_remove(struct tdx_td *td, u64 gpa, u64 level, u64 *ext_err1, u64 *ext_err2);
 u64 tdh_phymem_cache_wb(bool resume);
diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index 5b499593edff..405afd2a56b7 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -318,33 +318,40 @@ static inline void tdx_disassociate_vp(struct kvm_vcpu *vcpu)
 })
 
 /* TDH.PHYMEM.PAGE.RECLAIM is allowed only when destroying the TD. */
-static int __tdx_reclaim_page(struct page *page)
+static int tdx_reclaim_folio(struct folio *folio, unsigned long start_idx,
+			     unsigned long npages, bool reset)
 {
-	u64 err, rcx, rdx, r8;
+	u64 err, tdx_pt, tdx_owner, tdx_size;
 
-	err = tdh_phymem_page_reclaim(page, &rcx, &rdx, &r8);
+	err = tdh_phymem_page_reclaim(folio, start_idx, npages, &tdx_pt,
+				      &tdx_owner, &tdx_size);
 
 	/*
 	 * No need to check for TDX_OPERAND_BUSY; all TD pages are freed
 	 * before the HKID is released and control pages have also been
 	 * released at this point, so there is no possibility of contention.
 	 */
-	if (TDX_BUG_ON_3(err, TDH_PHYMEM_PAGE_RECLAIM, rcx, rdx, r8, NULL))
+	if (TDX_BUG_ON_3(err, TDH_PHYMEM_PAGE_RECLAIM, tdx_pt, tdx_owner, tdx_size, NULL))
 		return -EIO;
 
+	if (reset)
+		tdx_quirk_reset_folio(folio, start_idx, npages);
 	return 0;
 }
 
 static int tdx_reclaim_page(struct page *page)
 {
-	int r;
+	struct folio *folio = page_folio(page);
 
-	r = __tdx_reclaim_page(page);
-	if (!r)
-		tdx_quirk_reset_page(page);
-	return r;
+	return tdx_reclaim_folio(folio, folio_page_idx(folio, page), 1, true);
 }
 
+static int tdx_reclaim_page_noreset(struct page *page)
+{
+	struct folio *folio = page_folio(page);
+
+	return tdx_reclaim_folio(folio, folio_page_idx(folio, page), 1, false);
+}
 
 /*
  * Reclaim the TD control page(s) which are crypto-protected by TDX guest's
@@ -583,7 +590,7 @@ static void tdx_reclaim_td_control_pages(struct kvm *kvm)
 	if (!kvm_tdx->td.tdr_page)
 		return;
 
-	if (__tdx_reclaim_page(kvm_tdx->td.tdr_page))
+	if (tdx_reclaim_page_noreset(kvm_tdx->td.tdr_page))
 		return;
 
 	/*
diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 20708f56b1a0..c12665389b67 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -1993,19 +1993,27 @@ EXPORT_SYMBOL_GPL(tdh_vp_init);
  * So despite the names, they must be interpted specially as described by the spec. Return
  * them only for error reporting purposes.
  */
-u64 tdh_phymem_page_reclaim(struct page *page, u64 *tdx_pt, u64 *tdx_owner, u64 *tdx_size)
+u64 tdh_phymem_page_reclaim(struct folio *folio, unsigned long start_idx,
+			    unsigned long npages, u64 *tdx_pt, u64 *tdx_owner,
+			    u64 *tdx_size)
 {
 	struct tdx_module_args args = {
-		.rcx = page_to_phys(page),
+		.rcx = page_to_phys(folio_page(folio, start_idx)),
 	};
 	u64 ret;
 
+	if (start_idx + npages > folio_nr_pages(folio))
+		return TDX_OPERAND_INVALID;
+
 	ret = seamcall_ret(TDH_PHYMEM_PAGE_RECLAIM, &args);
 
 	*tdx_pt = args.rcx;
 	*tdx_owner = args.rdx;
 	*tdx_size = args.r8;
 
+	if (!ret && npages != (1 << (*tdx_size) * PTE_SHIFT))
+		return TDX_SW_ERROR;
+
 	return ret;
 }
 EXPORT_SYMBOL_GPL(tdh_phymem_page_reclaim);
-- 
2.43.2