ANDROID: KVM: arm64: iommu: Add context for guest teardown

As memory management for host and guests are different we rely on
__get_vcpu to know which methods to use, however in some contexts
as in free_domain/block_dev, the guest will not be loaded and we'd
need the memory reclaim to be done through the correct pool.
To fix this we add a static per cpu variable that can be set as a
fallback if vcpu is not loaded to check if we are in a teardown
context.

Bug: 357781595
Bug: 348382247
Bug: 236685427
Change-Id: I790d92a8dd6112bf40302bc0d6ec10582362db65
Signed-off-by: Mostafa Saleh <smostafa@google.com>

arch/arm64/kvm/hyp/nvhe/iommu/iommu.c

@@ -20,6 +20,10 @@
 struct kvm_iommu_ops *kvm_iommu_ops;
 void **kvm_hyp_iommu_domains;
 
+/* Hypervisor is non-preemptable, so cur_context can be per cpu. */
+DEFINE_PER_CPU(struct pkvm_hyp_vcpu *, __cur_context);
+#define cur_context (*this_cpu_ptr(&__cur_context))
+
 /*
  * Common pool that can be used by IOMMU driver to allocate pages.
  */
@@ -80,7 +84,8 @@ struct pkvm_hyp_vcpu *__get_vcpu(void)
 
 	if (vcpu)
 		return container_of(vcpu, struct pkvm_hyp_vcpu, vcpu);
-	return NULL;
+	/* Maybe guest is not loaded but we are in teardown context. */
+	return cur_context;
 }
 
 int iommu_pkvm_unuse_dma(u64 phys_addr, size_t size)