Add set_rid_user and unset_rid_user support

index 080278c8154d..f85ce6f01408 100644
@@ -248,6 +248,7 @@ struct iommu_iotlb_gather {
  *                      driver init to device driver init (default no)
  * @dev_enable/disable_feat: per device entries to enable/disable
  *                               iommu specific features.
+ * @set/unset_rid_user: per device entries to set/unset user space Request ID
  * @page_response: handle page request response
  * @def_domain_type: device default domain type, return value:
  *		- IOMMU_DOMAIN_IDENTITY: must use an identity domain

@@ -291,6 +292,9 @@ struct iommu_ops {
 	int (*dev_enable_feat)(struct device *dev, enum iommu_dev_features f);
 	int (*dev_disable_feat)(struct device *dev, enum iommu_dev_features f);
 
+	int (*set_rid_user)(struct device *dev, u32 rid, u32 rid_base);
+	int (*unset_rid_user)(struct device *dev);
+
 	int (*page_response)(struct device *dev,
 			     struct iommu_fault_event *evt,
 			     struct iommu_page_response *msg);

index 207b0e3fd9c8..db756e5cd77b 100644
@@ -735,6 +735,45 @@ void iommufd_device_detach(struct iommufd_device *idev)
 }
 EXPORT_SYMBOL_NS_GPL(iommufd_device_detach, IOMMUFD);
 
+int iommufd_device_set_rid(struct iommufd_ucmd *ucmd)
+{
+	struct iommufd_device_set_rid *cmd = ucmd->cmd;
+	struct iommufd_device *idev;
+	const struct iommu_ops *ops;
+	int rc = 0;
+
+	if (!cmd->rid)
+		return -EINVAL;
+
+	idev = iommufd_get_device(ucmd, cmd->dev_id);
+	if (IS_ERR(idev))
+		return PTR_ERR(idev);
+
+	ops = dev_iommu_ops(idev->dev);
+	if (ops && ops->set_rid_user)
+		rc = ops->set_rid_user(idev->dev, cmd->rid, cmd->rid_base);
+	iommufd_put_object(&idev->obj);
+	return rc;
+}
+
+int iommufd_device_unset_rid(struct iommufd_ucmd *ucmd)
+{
+	struct iommufd_device_unset_rid *cmd = ucmd->cmd;
+	struct iommufd_device *idev;
+	const struct iommu_ops *ops;
+	int rc = 0;
+
+	idev = iommufd_get_device(ucmd, cmd->dev_id);
+	if (IS_ERR(idev))
+		return PTR_ERR(idev);
+
+	ops = dev_iommu_ops(idev->dev);
+	if (ops && ops->unset_rid_user)
+		rc = ops->unset_rid_user(idev->dev);
+	iommufd_put_object(&idev->obj);
+	return rc;
+}
+
 void iommufd_access_destroy_object(struct iommufd_object *obj)
 {
 	struct iommufd_access *access =

index e52fd9d8c160..d65617727c0c 100644
@@ -320,6 +320,8 @@ iommufd_get_device(struct iommufd_ucmd *ucmd, u32 id)
 
 void iommufd_device_destroy(struct iommufd_object *obj);
 int iommufd_device_get_hw_info(struct iommufd_ucmd *ucmd);
+int iommufd_device_set_rid(struct iommufd_ucmd *ucmd);
+int iommufd_device_unset_rid(struct iommufd_ucmd *ucmd);
 
 extern const u64 iommufd_hwpt_type_bitmaps[];
 

index 7ec3ceac01b3..113fac4d4de0 100644
@@ -317,6 +317,10 @@ static const struct iommufd_ioctl_op iommufd_ioctl_ops[] = {
 		 val64),
 	IOCTL_OP(IOMMU_VFIO_IOAS, iommufd_vfio_ioas, struct iommu_vfio_ioas,
 		 __reserved),
+	IOCTL_OP(IOMMU_DEVICE_SET_RID, iommufd_device_set_rid,
+		 struct iommufd_device_set_rid, rid_base),
+	IOCTL_OP(IOMMU_DEVICE_UNSET_RID, iommufd_device_unset_rid,
+		 struct iommufd_device_unset_rid, dev_id),
 #ifdef CONFIG_IOMMUFD_TEST
 	IOCTL_OP(IOMMU_TEST_CMD, iommufd_test, struct iommu_test_cmd, last),
 #endif

index 59a44ceff27d..5188a4318e22 100644
@@ -48,6 +48,8 @@ enum {
 	IOMMUFD_CMD_HWPT_ALLOC,
 	IOMMUFD_CMD_DEVICE_GET_HW_INFO,
 	IOMMUFD_CMD_HWPT_INVALIDATE,
+	IOMMUFD_CMD_DEVICE_SET_RID,
+	IOMMUFD_CMD_DEVICE_UNSET_RID,
 };
 
 /**

@@ -482,4 +484,33 @@ struct iommu_hwpt_invalidate {
 	__aligned_u64 data_uptr;
 };
 #define IOMMU_HWPT_INVALIDATE _IO(IOMMUFD_TYPE, IOMMUFD_CMD_HWPT_INVALIDATE)
+
+/**
+ * struct iommufd_device_set_rid - ioctl(IOMMU_DEVICE_SET_RID)
+ * @size: sizeof(struct iommufd_device_set_rid)
+ * @dev_id: The device to set a Request ID
+ * @rid: PCI Request ID
+ *       Bits [15:8] are the Bus number.
+ *       Bits [7:3] are the Device number.
+ *       Bits [2:0] are the Function number.
+ * @rid_base: Base ID to map PCI Request ID
+ */
+struct iommufd_device_set_rid {
+	__u32 size;
+	__u32 dev_id;
+	__u32 rid;
+	__u32 rid_base;
+};
+#define IOMMU_DEVICE_SET_RID _IO(IOMMUFD_TYPE, IOMMUFD_CMD_DEVICE_SET_RID)
+
+/**
+ * struct iommufd_device_unset_rid - ioctl(IOMMU_DEVICE_UNSET_RID)
+ * @size: sizeof(struct iommufd_device_unset_rid)
+ * @dev_id: The device to unset a Request ID
+ */
+struct iommufd_device_unset_rid {
+	__u32 size;
+	__u32 dev_id;
+};
+#define IOMMU_DEVICE_UNSET_RID _IO(IOMMUFD_TYPE, IOMMUFD_CMD_DEVICE_UNSET_RID)
 #endif

index f2425b0f0cd6..17bb90e78547 100644
@@ -2843,6 +2843,48 @@ static void arm_smmu_remove_dev_pasid(struct device *dev, ioasid_t pasid)
 	arm_smmu_sva_remove_dev_pasid(domain, dev, pasid);
 }
 
+static int arm_smmu_set_rid_user(struct device *dev, u32 rid, u32 rid_base)
+{
+	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
+	struct arm_smmu_stream *stream = &master->streams[0];
+	struct arm_smmu_device *smmu = master->smmu;
+	u32 sid_user = rid_base | rid;
+	int ret = 0;
+
+	if (!sid_user)
+		return -EINVAL;
+
+	ret = xa_alloc(&smmu->streams_user, &sid_user, stream,
+			XA_LIMIT(sid_user, sid_user), GFP_KERNEL_ACCOUNT);
+	if (ret)
+		return ret;
+	stream->id_user = sid_user;
+	return 0;
+}
+
+static int arm_smmu_unset_rid_user(struct device *dev)
+{
+	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
+	struct arm_smmu_stream *stream = &master->streams[0];
+	struct arm_smmu_device *smmu = master->smmu;
+	u32 sid_user = stream->id_user;
+
+	if (!sid_user)
+		return -ENODEV;
+
+	xa_lock(&smmu->streams_user);
+	stream = __xa_erase(&smmu->streams_user, sid_user);
+	if (stream != master->streams) {
+		WARN_ON(__xa_alloc(&smmu->streams_user, &sid_user, stream,
+				   XA_LIMIT(sid_user, sid_user),
+				   GFP_KERNEL_ACCOUNT));
+		xa_unlock(&smmu->streams_user);
+		return -EINVAL;
+	}
+	xa_unlock(&smmu->streams_user);
+	return 0;
+}
+
 static struct iommu_ops arm_smmu_ops = {
 	.capable		= arm_smmu_capable,
 	.domain_alloc		= arm_smmu_domain_alloc,

@@ -2854,6 +2896,8 @@ static struct iommu_ops arm_smmu_ops = {
 	.remove_dev_pasid	= arm_smmu_remove_dev_pasid,
 	.dev_enable_feat	= arm_smmu_dev_enable_feature,
 	.dev_disable_feat	= arm_smmu_dev_disable_feature,
+	.set_rid_user		= arm_smmu_set_rid_user,
+	.unset_rid_user		= arm_smmu_unset_rid_user,
 	.page_response		= arm_smmu_page_response,
 	.def_domain_type	= arm_smmu_def_domain_type,
 	.pgsize_bitmap		= -1UL, /* Restricted during device attach */

@@ -3082,6 +3126,7 @@ static int arm_smmu_init_structures(struct arm_smmu_device *smmu)
 
 	mutex_init(&smmu->streams_mutex);
 	smmu->streams = RB_ROOT;
+	xa_init_flags(&smmu->streams_user, XA_FLAGS_ALLOC1 | XA_FLAGS_ACCOUNT);
 
 	ret = arm_smmu_init_queues(smmu);
 	if (ret)

index 8d772ea8a583..8ccea2e84bbb 100644
@@ -674,10 +674,12 @@ struct arm_smmu_device {
 
 	struct rb_root			streams;
 	struct mutex			streams_mutex;
+	struct xarray			streams_user;
 };
 
 struct arm_smmu_stream {
 	u32				id;
+	u32				id_user;
 	struct arm_smmu_master		*master;
 	struct rb_node			node;
 };