ANDROID: KVM: arm64: devices: Teardown assigned devices

Teardown devices when VM is toredown, at the moment, we only
set the device back to host, in the next patches more logic
is added as device reset and IOMMU blocking.

Unlike memory, we don't rely on a hypercall for each page,
mainly because:
- We don't need to track the pages in list/tree
- MMIO space is usually very small, so it shouldn't impact
  performance.

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

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

@@ -163,6 +163,7 @@ int pkvm_init_scmi_pd(struct kvm_power_domain *pd,
 		      const struct kvm_power_domain_ops *ops);
 
 bool pkvm_device_request_mmio(struct pkvm_hyp_vcpu *hyp_vcpu, u64 *exit_code);
+void pkvm_devices_teardown(struct pkvm_hyp_vm *vm);
 
 /*
  * Register a power domain. When the hypervisor catches power requests from the

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

@@ -302,3 +302,32 @@ out_inval:
 	smccc_set_retval(vcpu, SMCCC_RET_INVALID_PARAMETER, 0, 0, 0);
 	return true;
 }
+
+static void pkvm_devices_reclaim_device(struct pkvm_device *dev)
+{
+	int i;
+
+	for (i = 0 ; i < dev->nr_resources ; ++i) {
+		struct pkvm_dev_resource *res = &dev->resources[i];
+
+		hyp_spin_lock(&host_mmu.lock);
+		WARN_ON(host_stage2_set_owner_locked(res->base, res->size, PKVM_ID_HOST));
+		hyp_spin_unlock(&host_mmu.lock);
+	}
+}
+
+void pkvm_devices_teardown(struct pkvm_hyp_vm *vm)
+{
+	int i;
+
+	hyp_spin_lock(&device_spinlock);
+	for (i = 0 ; i < registered_devices_nr ; ++i) {
+		struct pkvm_device *dev = &registered_devices[i];
+
+		if (dev->ctxt != vm)
+			continue;
+		dev->ctxt = NULL;
+		pkvm_devices_reclaim_device(dev);
+	}
+	hyp_spin_unlock(&device_spinlock);
+}

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

@@ -931,6 +931,8 @@ int __pkvm_finalize_teardown_vm(pkvm_handle_t handle)
 	remove_vm_table_entry(handle);
 	hyp_write_unlock(&vm_table_lock);
 
+	pkvm_devices_teardown(hyp_vm);
+
 	/*
 	 * At this point, the VM has been detached from the VM table and
 	 * has a refcount of 0 so we're free to tear it down without