ANDROID: KVM: arm64: devices: Takeover assignable devices struct

On pKVM setup fix the assignable devices address to hyp va and unmap
it from host as it can't be trusted after this HVC.

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

arch/arm64/include/asm/kvm_asm.h

@@ -74,6 +74,7 @@ enum __kvm_host_smccc_func {
 	__KVM_HOST_SMCCC_FUNC___pkvm_init_module,
 	__KVM_HOST_SMCCC_FUNC___pkvm_register_hcall,
 	__KVM_HOST_SMCCC_FUNC___pkvm_iommu_init,
+	__KVM_HOST_SMCCC_FUNC___pkvm_devices_init,
 	__KVM_HOST_SMCCC_FUNC___pkvm_prot_finalize,
 
 	/* Hypercalls available after pKVM finalisation */

arch/arm64/kvm/hyp/include/nvhe/pkvm.h

@@ -181,4 +181,6 @@ static inline int pkvm_init_power_domain(struct kvm_power_domain *pd,
 	}
 }
 
+int pkvm_init_devices(void);
+
 #endif /* __ARM64_KVM_NVHE_PKVM_H__ */

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

@@ -3,6 +3,27 @@
  * Copyright (C) 2023 Google LLC
  * Author: Mostafa Saleh <smostafa@google.com>
  */
+#include <nvhe/mem_protect.h>
+#include <nvhe/mm.h>
+#include <nvhe/pkvm.h>
+
+#include <kvm/device.h>
 
 struct pkvm_device *registered_devices;
 unsigned long registered_devices_nr;
+
+int pkvm_init_devices(void)
+{
+	size_t dev_sz;
+	int ret;
+
+	registered_devices = kern_hyp_va(registered_devices);
+	dev_sz = PAGE_ALIGN(size_mul(sizeof(struct pkvm_device),
+				     registered_devices_nr));
+
+	ret = __pkvm_host_donate_hyp(hyp_virt_to_phys(registered_devices) >> PAGE_SHIFT,
+				     dev_sz >> PAGE_SHIFT);
+	if (ret)
+		registered_devices_nr = 0;
+	return ret;
+}

arch/arm64/kvm/hyp/nvhe/hyp-main.c

@@ -1728,6 +1728,15 @@ static void handle___pkvm_ptdump(struct kvm_cpu_context *host_ctxt)
 		cpu_reg(host_ctxt, 0) = SMCCC_RET_NOT_SUPPORTED;
 }
 
+static void handle___pkvm_devices_init(struct kvm_cpu_context *host_ctxt)
+{
+	/*
+	 * Devices must be initialised after the IOMMUs driver is initialised.
+	 * We do this in a separate HVC to avoid complexity.
+	 */
+	cpu_reg(host_ctxt, 1) = pkvm_init_devices();
+}
+
 static void handle___pkvm_host_iommu_map_sg(struct kvm_cpu_context *host_ctxt)
 {
 	unsigned long ret;
@@ -1766,6 +1775,7 @@ static const hcall_t host_hcall[] = {
 	HANDLE_FUNC(__pkvm_init_module),
 	HANDLE_FUNC(__pkvm_register_hcall),
 	HANDLE_FUNC(__pkvm_iommu_init),
+	HANDLE_FUNC(__pkvm_devices_init),
 	HANDLE_FUNC(__pkvm_prot_finalize),
 
 	HANDLE_FUNC(__pkvm_host_share_hyp),

arch/arm64/kvm/pkvm.c

@@ -662,6 +662,10 @@ static int __init finalize_pkvm(void)
 		pkvm_firmware_rmem_clear();
 	}
 
+	ret = kvm_call_hyp_nvhe(__pkvm_devices_init);
+	if (ret)
+		pr_warn("Assignable devices failed to initialize in the hypervisor %d", ret);
+
 	/*
 	 * Exclude HYP sections from kmemleak so that they don't get peeked
 	 * at, which would end badly once inaccessible.