vfio-pci/zdev: fix possible segmentation fault issue

stable inclusion
from stable-5.10.20
commit 9eb145398e6fd31f1224a70a69db5ee0882210ec
bugzilla: 50608

--------------------------------

[ Upstream commit 7e31d6dc ]

In case allocation fails, we must behave correctly and exit with error.

Fixes: e6b817d4 ("vfio-pci/zdev: Add zPCI capabilities to VFIO_DEVICE_GET_INFO")
Signed-off-by: NMax Gurtovoy <mgurtovoy@nvidia.com>
Reviewed-by: NCornelia Huck <cohuck@redhat.com>
Reviewed-by: NMatthew Rosato <mjrosato@linux.ibm.com>
Signed-off-by: NAlex Williamson <alex.williamson@redhat.com>
Signed-off-by: NSasha Levin <sashal@kernel.org>
Signed-off-by: NChen Jun <chenjun102@huawei.com>
Acked-by: NXie XiuQi <xiexiuqi@huawei.com>
Signed-off-by: NZheng Zengkai <zhengzengkai@huawei.com>

drivers/vfio/pci/vfio_pci_zdev.c

@@ -74,6 +74,8 @@ static int zpci_util_cap(struct zpci_dev *zdev, struct vfio_pci_device *vdev,
 	int ret;
 
 	cap = kmalloc(cap_size, GFP_KERNEL);
+	if (!cap)
+		return -ENOMEM;
 
 	cap->header.id = VFIO_DEVICE_INFO_CAP_ZPCI_UTIL;
 	cap->header.version = 1;
@@ -98,6 +100,8 @@ static int zpci_pfip_cap(struct zpci_dev *zdev, struct vfio_pci_device *vdev,
 	int ret;
 
 	cap = kmalloc(cap_size, GFP_KERNEL);
+	if (!cap)
+		return -ENOMEM;
 
 	cap->header.id = VFIO_DEVICE_INFO_CAP_ZPCI_PFIP;
 	cap->header.version = 1;