From 007871cbf720374e5024f0eec74d8c5ce5453b04 Mon Sep 17 00:00:00 2001
From: Mostafa Saleh <smostafa@google.com>
Date: Sat, 22 Jun 2024 15:13:06 +0000
Subject: [PATCH] ANDROID: drivers/arm-smmu-v3-kvm: Add dev_block_dma

When a device is assigned dev_block_dma is called to put DMA in
quiescent state.

Bug: 357781595
Bug: 348382247
Change-Id: Ibcf054dab9bb3e193c45887b5cb82845a1072148
Signed-off-by: Mostafa Saleh <smostafa@google.com>
---
 .../iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c   |  5 ++--
 .../iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c  | 26 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c
index cc7169bd913c..cae6eafc2459 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-kvm.c
@@ -169,7 +169,8 @@ static struct iommu_domain *kvm_arm_smmu_domain_alloc(unsigned type)
 	 */
 	if (type != IOMMU_DOMAIN_DMA &&
 	    type != IOMMU_DOMAIN_UNMANAGED &&
-	    type != IOMMU_DOMAIN_IDENTITY)
+	    type != IOMMU_DOMAIN_IDENTITY &&
+		type != IOMMU_DOMAIN_BLOCKED)
 		return ERR_PTR(-EOPNOTSUPP);
 
 	kvm_smmu_domain = kzalloc(sizeof(*kvm_smmu_domain), GFP_KERNEL);
@@ -352,7 +353,7 @@ static int kvm_arm_smmu_set_dev_pasid(struct iommu_domain *domain,
 	host_smmu = smmu_to_host(smmu);
 
 	ret = kvm_arm_smmu_detach_dev_pasid(host_smmu, master, pasid);
-	if (ret)
+	if (ret || (domain->type == IOMMU_DOMAIN_BLOCKED))
 		return ret;
 
 	mutex_lock(&kvm_smmu_domain->init_mutex);
diff --git a/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c b/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c
index f01149970c72..16f89df3466a 100644
--- a/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c
@@ -1550,6 +1550,31 @@ static bool smmu_dabt_device(struct hyp_arm_smmu_v3_device *smmu,
 	return true;
 }
 
+static int smmu_dev_block_dma(struct kvm_hyp_iommu *iommu, u32 sid, bool is_host2guest)
+{
+	struct hyp_arm_smmu_v3_device *smmu = to_smmu(iommu);
+	static struct arm_smmu_ste *dst;
+	int ret = 0;
+
+	kvm_iommu_lock(iommu);
+	dst = smmu_get_ste_ptr(smmu, sid);
+
+	/*
+	 * VFIO will attach the device to a blocking domain, this will make the
+	 * kernel driver detach the device which should be have zeroed STE.
+	 * So, if this is not the current state of the device, something
+	 * went wrong.
+	 * For guests, we need to do more as guests might not exit cleanly
+	 * and the device might be translating, so we have to actually block
+	 * the device and clean the STE/CD.
+	 */
+	if (dst->data[0])
+		ret = -EINVAL;
+
+	kvm_iommu_unlock(iommu);
+	return ret;
+}
+
 static bool smmu_dabt_handler(struct kvm_cpu_context *host_ctxt, u64 esr, u64 addr)
 {
 	struct hyp_arm_smmu_v3_device *smmu;
@@ -1688,4 +1713,5 @@ struct kvm_iommu_ops smmu_ops = {
 	.suspend			= smmu_suspend,
 	.resume				= smmu_resume,
 	.host_stage2_idmap		= smmu_host_stage2_idmap,
+	.dev_block_dma			= smmu_dev_block_dma,
 };