Currently the code misuse GD_SUPPRESS_PART_SCAN flag
as it tries to use it as a switch for the auto partition scan.
To fully disable creation of partitions GENHD_FL_NO_PART must
be used, switch code to use it instead GD_SUPPRESS_PART_SCAN.

Using the new flags rules for partitions become:
 - Unprivileged daemons - never scan partitions, they have
GENHD_FL_NO_PART always set - they are devices without partitions.

 - Trusted daemons - by default have partitions disabled.
Partitions can be enabled via UBLK_F_PARTITIONS flag and
automatic initial scan can be requested via UBLK_F_AUTO_PART_SCAN flag.

Rework the code to work as described above:
- remove checks in ublk_partition_scan_work and rely on
the caller to schedule the work only if it has to be done.
- set GENHD_FL_NO_PART on unprivileged devices
- set GENHD_FL_NO_PART depending on UBLK_F_PARTITIONS on trusted
devices.

Fixes: 8443e2087e70 ("ublk: add UBLK_F_NO_AUTO_PART_SCAN feature flag")
Signed-off-by: Alexander Atanasov <alex@zazolabs.com>
---
 drivers/block/ublk_drv.c      | 34 +++++++++++++++++++---------------
 include/uapi/linux/ublk_cmd.h |  7 +++++--
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 3c918db4905c..d0c646eb9435 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -81,7 +81,8 @@
 		| (IS_ENABLED(CONFIG_BLK_DEV_INTEGRITY) ? UBLK_F_INTEGRITY : 0) \
 		| UBLK_F_SAFE_STOP_DEV \
 		| UBLK_F_BATCH_IO \
-		| UBLK_F_NO_AUTO_PART_SCAN)
+		| UBLK_F_PARTITIONS \
+		| UBLK_F_AUTO_PART_SCAN)
 
 #define UBLK_F_ALL_RECOVERY_FLAGS (UBLK_F_USER_RECOVERY \
 		| UBLK_F_USER_RECOVERY_REISSUE \
@@ -2361,14 +2362,9 @@ static void ublk_partition_scan_work(struct work_struct *work)
 	if (!disk)
 		return;
 
-	if (WARN_ON_ONCE(!test_and_clear_bit(GD_SUPPRESS_PART_SCAN,
-					     &disk->state)))
-		goto out;
-
 	mutex_lock(&disk->open_mutex);
 	bdev_disk_changed(disk, false);
 	mutex_unlock(&disk->open_mutex);
-out:
 	ublk_put_disk(disk);
 }
 
@@ -4409,9 +4405,13 @@ static int ublk_ctrl_start_dev(struct ublk_device *ub,
 	 * wait while holding ub->mutex, which can deadlock with other
 	 * operations that need the mutex. Defer partition scan to async
 	 * work.
-	 * For unprivileged daemons, keep GD_SUPPRESS_PART_SCAN set
-	 * permanently.
+	 * For unprivileged daemons, set GENHD_FL_NO_PART to
+	 * disable partitions.
 	 */
+	if (ub->unprivileged_daemons)
+		disk->flags |= GENHD_FL_NO_PART;
+
+	/* Disable partition scans */
 	set_bit(GD_SUPPRESS_PART_SCAN, &disk->state);
 
 	ublk_get_device(ub);
@@ -4429,13 +4429,17 @@ static int ublk_ctrl_start_dev(struct ublk_device *ub,
 
 	set_bit(UB_STATE_USED, &ub->state);
 
-	/* Skip partition scan if disabled by user */
-	if (ub->dev_info.flags & UBLK_F_NO_AUTO_PART_SCAN) {
-		clear_bit(GD_SUPPRESS_PART_SCAN, &disk->state);
-	} else {
-		/* Schedule async partition scan for trusted daemons */
-		if (!ub->unprivileged_daemons)
-			schedule_work(&ub->partition_scan_work);
+	if (!ub->unprivileged_daemons) {
+		if (!(ub->dev_info.flags & UBLK_F_PARTITIONS)) {
+			/* disable partitions on trusted device */
+			disk->flags |= GENHD_FL_NO_PART;
+		} else {
+			/* Enable partition scan on device */
+			clear_bit(GD_SUPPRESS_PART_SCAN, &disk->state);
+			/* Schedule auto partition scan if requested */
+			if (ub->dev_info.flags & UBLK_F_AUTO_PART_SCAN)
+				schedule_work(&ub->partition_scan_work);
+		}
 	}
 
 out_put_cdev:
diff --git a/include/uapi/linux/ublk_cmd.h b/include/uapi/linux/ublk_cmd.h
index a88876756805..467ecaa20ac1 100644
--- a/include/uapi/linux/ublk_cmd.h
+++ b/include/uapi/linux/ublk_cmd.h
@@ -367,8 +367,11 @@
  */
 #define UBLK_F_SAFE_STOP_DEV	(1ULL << 17)
 
-/* Disable automatic partition scanning when device is started */
-#define UBLK_F_NO_AUTO_PART_SCAN (1ULL << 18)
+/* Enable partitions on the device */
+#define UBLK_F_PARTITIONS (1ULL << 18)
+
+/* Perform automatic partition scanning when device is started */
+#define UBLK_F_AUTO_PART_SCAN (1ULL << 19)
 
 /* device state */
 #define UBLK_S_DEV_DEAD	0
-- 
2.43.0