Re: [PATCH v9 03/11] x86/virt/tdx: Add tdx_alloc/free_control_page() helpers

> +/*
> + * Calculate the arg needed for operating on the DPAMT backing for
> + * a given 4KB page.
> + */
> +static u64 pamt_2mb_arg(kvm_pfn_t pfn)
> +{
> +	/* Arg value will specify a 2MB region of physical address space. */
> +	unsigned long hpa_2mb = ALIGN_DOWN(pfn << PAGE_SHIFT, PMD_SIZE);
> +
> +	return hpa_2mb | TDX_PS_2M;
> +}

/* Helper for building dynamic PAMT seamcall() arguments. */
static u64 pamt_2mb_arg(kvm_pfn_t pfn)
{
	/* Find the 2MB-wide DPAMT region for 'pfn':  */
	unsigned long hpa_2mb = ALIGN_DOWN(pfn << PAGE_SHIFT, PMD_SIZE);

	/* Talk about why the flag is needed... */
	return hpa_2mb | TDX_PS_2M;
}


> +/* Add PAMT backing for the 2MB region surrounding the given pfn. */
> +static u64 tdh_phymem_pamt_add(kvm_pfn_t pfn, struct page **pamt_pages)
> +{
> +	struct tdx_module_args args = {
> +		.rcx = pamt_2mb_arg(pfn),
> +		.rdx = page_to_phys(pamt_pages[0]),
> +		.r8 = page_to_phys(pamt_pages[1]),

Super nit:          ^ should be vertically aligned

> +	};
> +
> +	return seamcall(TDH_PHYMEM_PAMT_ADD, &args);
> +}
> +
> +/* Remove PAMT backing for the 2MB region surrounding the given pfn. */
> +static u64 tdh_phymem_pamt_remove(kvm_pfn_t pfn, struct page **pamt_pages)
> +{
> +	struct tdx_module_args args = {
> +		.rcx = pamt_2mb_arg(pfn),
> +	};
> +	u64 ret;
> +
> +	ret = seamcall_ret(TDH_PHYMEM_PAMT_REMOVE, &args);
> +	if (ret)
> +		return ret;
> +
> +	/* Copy PAMT pages out of the struct per the TDX ABI */
> +	pamt_pages[0] = phys_to_page(args.rdx);
> +	pamt_pages[1] = phys_to_page(args.r8);
> +
> +	return 0;
> +}
> +
> +/* Allocate PAMT memory for the given page */
> +static int tdx_pamt_get(kvm_pfn_t pfn)
> +{
> +	struct page *pamt_pages[TDX_DPAMT_ENTRY_PAGE_CNT];
> +	u64 tdx_status;
> +	int ret;
> +
> +	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo))
> +		return 0;
> +
> +	ret = alloc_pamt_array(pamt_pages);
> +	if (ret)
> +		return ret;
> +
> +	tdx_status = tdh_phymem_pamt_add(pfn, pamt_pages);
> +	if (tdx_status != TDX_SUCCESS) {
> +		ret = -EIO;
> +		goto out_free;
> +	}
> +
> +	return 0;
> +
> +out_free:
> +	free_pamt_array(pamt_pages);
> +
> +	return ret;
> +}
> +
> +/* Free PAMT memory for the given page */
> +static void tdx_pamt_put(kvm_pfn_t pfn)
> +{
> +	struct page *pamt_pages[TDX_DPAMT_ENTRY_PAGE_CNT] = {};
> +	u64 tdx_status;
> +
> +	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo))
> +		return;
> +
> +	tdx_status = tdh_phymem_pamt_remove(pfn, pamt_pages);
> +
> +	/*
> +	 * Don't free pamt_pages as it could hold garbage when
> +	 * tdh_phymem_pamt_remove() fails.  Don't panic/BUG_ON(), as
> +	 * there is no risk of data corruption, but do yell loudly as
> +	 * failure indicates a kernel bug, memory is being leaked, and
> +	 * the dangling PAMT entry may cause future operations to fail.
> +	 */
> +	if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS))
> +		return;0

Whose fault is the bug here? Probably the TDX module?

> +	free_pamt_array(pamt_pages);
> +}
> +
> +/*
> + * Return a page that can be gifted to the TDX-Module for use as a "control"
> + * page, i.e. pages that are used for control structures for a given TDX
> + * guest, and thus obtain TDX protections, including PAMT tracking.
> + */
> +struct page *tdx_alloc_control_page(void)
> +{
> +	struct page *page;
> +
> +	page = alloc_page(GFP_KERNEL_ACCOUNT);
> +	if (!page)
> +		return NULL;
> +
> +	if (tdx_pamt_get(page_to_pfn(page))) {
> +		__free_page(page);
> +		return NULL;
> +	}

Here's where I start to get lost.

What are the rules around tdx_pamt_get()? tdx_alloc_control_page() is
obviously for one 4k page. But tdx_pamt_get() does an allocation for 2MB
regions. So it has to be able to fail to actually add PAMT in some
cases, gracefully.

Right?

How does that happen?

I know the changelog tried to hand wave this away a bit. But the
comments in here need to say something about tdx_pamt_get() being
unusable on its own. The normal convention would be to __ it, too.

> +	return page;
> +}
> +EXPORT_SYMBOL_FOR_KVM(tdx_alloc_control_page);
> +
> +/*
> + * Free a page that was gifted to the TDX-Module for use as a control

Nit "TDX Module" for consistency.

> + * page. After this, the page is no longer protected by TDX.
> + */
> +void tdx_free_control_page(struct page *page)
> +{
> +	if (!page)
> +		return;
> +
> +	tdx_pamt_put(page_to_pfn(page));
> +	__free_page(page);
> +}
> +EXPORT_SYMBOL_FOR_KVM(tdx_free_control_page);
> +
>  void tdx_sys_disable(void)
>  {
>  	struct tdx_module_args args = {};
> diff --git a/arch/x86/virt/vmx/tdx/tdx.h b/arch/x86/virt/vmx/tdx/tdx.h
> index bdfd0e1e337ac..a886c54decaad 100644
> --- a/arch/x86/virt/vmx/tdx/tdx.h
> +++ b/arch/x86/virt/vmx/tdx/tdx.h
> @@ -48,6 +48,8 @@
>  #define TDH_SYS_CONFIG			45
>  #define TDH_SYS_SHUTDOWN		52
>  #define TDH_SYS_UPDATE			53
> +#define TDH_PHYMEM_PAMT_ADD		58
> +#define TDH_PHYMEM_PAMT_REMOVE		59
>  #define TDH_SYS_DISABLE			69
>  
>  /*