[PATCH V2 1/5] ublk_drv: avoid to leak ublk device in case that add_disk fails

[Ming Lei <ming.lei@redhat.com>]

->free_disk is only called after disk is added successfully, so
not hold ublk device reference count until add_disk is done.

Fixes: 6d9e6dfdf3b2 ("ublk: defer disk allocation")
Signed-off-by: Ming Lei <ming.lei@redhat.com>
---
 drivers/block/ublk_drv.c | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 255b2de46a24..b30d6c3355e8 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -215,8 +215,11 @@ static void ublk_free_disk(struct gendisk *disk)
 {
 	struct ublk_device *ub = disk->private_data;
 
-	clear_bit(UB_STATE_USED, &ub->state);
-	put_device(&ub->cdev_dev);
+	/* only called for added/used disk */
+	if (test_bit(UB_STATE_USED, &ub->state)) {
+		clear_bit(UB_STATE_USED, &ub->state);
+		put_device(&ub->cdev_dev);
+	}
 }
 
 static const struct block_device_operations ub_fops = {
@@ -1181,12 +1184,12 @@ static int ublk_ctrl_start_dev(struct io_uring_cmd *cmd)
 
 	ub->dev_info.ublksrv_pid = ublksrv_pid;
 	ub->ub_disk = disk;
-	get_device(&ub->cdev_dev);
 	ret = add_disk(disk);
 	if (ret) {
 		put_disk(disk);
 		goto out_unlock;
 	}
+	get_device(&ub->cdev_dev);
 	set_bit(UB_STATE_USED, &ub->state);
 	ub->dev_info.state = UBLK_S_DEV_LIVE;
 out_unlock:
-- 
2.31.1