Implement the VFIO_DEVICE_FEATURE_P2P_REGISTER. A bitmask is added to
track the registered regions.

Post-registration, the BAR is added to the device's P2P pool and is
available to be used with standard page-based APIs.

Signed-off-by: Pranjal Shrivastava <praan@google.com>
---
 drivers/vfio/pci/vfio_pci_core.c | 39 ++++++++++++++++++++++++++++++++
 include/linux/vfio_pci_core.h    |  1 +
 2 files changed, 40 insertions(+)

diff --git a/drivers/vfio/pci/vfio_pci_core.c b/drivers/vfio/pci/vfio_pci_core.c
index a28f1e99362c..1e922e3aaeb3 100644
--- a/drivers/vfio/pci/vfio_pci_core.c
+++ b/drivers/vfio/pci/vfio_pci_core.c
@@ -1554,6 +1554,43 @@ static int vfio_pci_core_feature_token(struct vfio_pci_core_device *vdev,
 	return 0;
 }
 
+static int vfio_pci_core_feature_p2p_register(struct vfio_pci_core_device *vdev,
+					      u32 flags, u32 __user *arg,
+					      size_t argsz)
+{
+	u32 bar_index;
+	int ret;
+
+	ret = vfio_check_feature(flags, argsz, VFIO_DEVICE_FEATURE_SET,
+				 sizeof(bar_index));
+	if (ret != 1)
+		return ret;
+
+	if (copy_from_user(&bar_index, arg, sizeof(bar_index)))
+		return -EFAULT;
+
+	if (bar_index >= PCI_STD_NUM_BARS)
+		return -EINVAL;
+
+	if (!(pci_resource_flags(vdev->pdev, bar_index) & IORESOURCE_MEM))
+		return -EINVAL;
+
+	/* Already registered */
+	if (vdev->p2p_registered_bars & (1 << bar_index))
+		return 0;
+
+	ret = pci_p2pdma_add_resource(vdev->pdev, bar_index, 0, 0);
+	if (ret && ret != -EEXIST)
+		return ret;
+
+	vdev->p2p_registered_bars |= (1 << bar_index);
+
+	pci_info(vdev->pdev, "BAR %d registered for ZONE_DEVICE P2P\n",
+		 bar_index);
+
+	return 0;
+}
+
 int vfio_pci_core_ioctl_feature(struct vfio_device *device, u32 flags,
 				void __user *arg, size_t argsz)
 {
@@ -1572,6 +1609,8 @@ int vfio_pci_core_ioctl_feature(struct vfio_device *device, u32 flags,
 		return vfio_pci_core_feature_token(vdev, flags, arg, argsz);
 	case VFIO_DEVICE_FEATURE_DMA_BUF:
 		return vfio_pci_core_feature_dma_buf(vdev, flags, arg, argsz);
+	case VFIO_DEVICE_FEATURE_P2P_REGISTER:
+		return vfio_pci_core_feature_p2p_register(vdev, flags, arg, argsz);
 	default:
 		return -ENOTTY;
 	}
diff --git a/include/linux/vfio_pci_core.h b/include/linux/vfio_pci_core.h
index 5fc6ce4dd786..cae7f069e5b6 100644
--- a/include/linux/vfio_pci_core.h
+++ b/include/linux/vfio_pci_core.h
@@ -112,6 +112,7 @@ struct vfio_pci_core_device {
 	struct vfio_pci_region	*region;
 	u8			msi_qmax;
 	u8			msix_bar;
+	u8			p2p_registered_bars;
 	u16			msix_size;
 	u32			msix_offset;
 	u32			rbar[7];
-- 
2.54.0.1099.g489fc7bff1-goog