From 41148bbeb73900a1fe8c5d7c64bbec5207d84c06 Mon Sep 17 00:00:00 2001
From: Mostafa Saleh <smostafa@google.com>
Date: Mon, 3 Feb 2025 10:15:44 +0000
Subject: [PATCH] ANDROID: KVM: arm64: iommu: Make allocation guest aware

In preparation to use IOMMU functions for guests, use guest
pools when allocating from guest context.
We have to rely on the current vcpu running from kvm_host_data
as modifying all allocation is quite intrusive and would make
the IOMMU drivers guest aware which we are trying to avoid.

Bug: 357781595
Bug: 348382247
Bug: 236685427
Change-Id: Iaa7b58971273b1adc3f2aa607c1bbc1f65516303
Signed-off-by: Mostafa Saleh <smostafa@google.com>
---
 arch/arm64/kvm/hyp/nvhe/iommu/iommu.c | 37 +++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c b/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c
index 656428235c9b..e389f81dbf4d 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/iommu.c
@@ -73,12 +73,23 @@ struct hyp_mgt_allocator_ops kvm_iommu_allocator_ops = {
 	.reclaimable = kvm_iommu_reclaimable,
 };
 
+/* Return current vcpu or NULL for host. */
+struct pkvm_hyp_vcpu *__get_vcpu(void)
+{
+	struct kvm_vcpu *vcpu = this_cpu_ptr(&kvm_host_data)->host_ctxt.__hyp_running_vcpu;
+
+	if (vcpu)
+		return container_of(vcpu, struct pkvm_hyp_vcpu, vcpu);
+	return NULL;
+}
+
 static void *__kvm_iommu_donate_pages(struct hyp_pool *pool, u8 order, int flags)
 {
 	void *p;
 	struct kvm_hyp_req *req = this_cpu_ptr(&host_hyp_reqs);
 	int ret;
 	size_t size = (1 << order) * PAGE_SIZE;
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
 
 	p = hyp_alloc_pages(pool, order);
 	if (p) {
@@ -100,6 +111,12 @@ static void *__kvm_iommu_donate_pages(struct hyp_pool *pool, u8 order, int flags
 		return p;
 	}
 
+	if (hyp_vcpu) {
+		req = pkvm_hyp_req_reserve(hyp_vcpu, KVM_HYP_REQ_TYPE_MEM);
+		if (WARN_ON(!req))
+			return NULL;
+	}
+
 	req->type = KVM_HYP_REQ_TYPE_MEM;
 	req->mem.dest = REQ_MEM_DEST_HYP_IOMMU;
 	req->mem.sz_alloc = size;
@@ -120,12 +137,28 @@ static void __kvm_iommu_reclaim_pages(struct hyp_pool *pool, void *p, u8 order)
 
 void *kvm_iommu_donate_pages(u8 order, int flags)
 {
-	return __kvm_iommu_donate_pages(&iommu_host_pool, order, flags);
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
+	struct hyp_pool *pool;
+
+	if (hyp_vcpu)
+		pool = &pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu)->iommu_pool;
+	else
+		pool = &iommu_host_pool;
+
+	return __kvm_iommu_donate_pages(pool, order, flags);
 }
 
 void kvm_iommu_reclaim_pages(void *p, u8 order)
 {
-	__kvm_iommu_reclaim_pages(&iommu_host_pool, p, order);
+	struct pkvm_hyp_vcpu *hyp_vcpu = __get_vcpu();
+	struct hyp_pool *pool;
+
+	if (hyp_vcpu)
+		pool = &pkvm_hyp_vcpu_to_hyp_vm(hyp_vcpu)->iommu_pool;
+	else
+		pool = &iommu_host_pool;
+
+	__kvm_iommu_reclaim_pages(pool, p, order);
 }
 
 void *kvm_iommu_donate_pages_atomic(u8 order)