ANDROID: ufs, block: fix crypto power management and move into block layer

The call to pm_runtime_get_sync() in ufshcd_program_key() can deadlock
because it waits for the UFS controller to be resumed, but it can itself
be reached while resuming the UFS controller via:

- ufshcd_runtime_resume()
  - ufshcd_resume()
    - ufshcd_reset_and_restore()
      - ufshcd_host_reset_and_restore()
        - ufshcd_hba_enable()
          - ufshcd_hba_execute_hce()
            - ufshcd_hba_start()
              - ufshcd_crypto_enable()
                - keyslot_manager_reprogram_all_keys()
                  - ufshcd_crypto_keyslot_program()
                    - ufshcd_program_key()

But pm_runtime_get_sync() *is* needed when evicting a key.  Also, on
pre-4.20 kernels it's needed when programming a keyslot for a bio since
the block layer used to resume the device in a different place.

Thus, it's hard for drivers to know what to do in .keyslot_program() and
.keyslot_evict().  In old kernels it may even be impossible unless we
were to pass more information down from the keyslot_manager.

There's also another possible deadlock: keyslot programming and eviction
take ksm->lock for write and then resume the device, which may result in
ksm->lock being taken again via the above call stack.  To fix this, we
should resume the device before taking ksm->lock.

Fix these problems by moving to a better design where the block layer
(namely, the keyslot manager) handles runtime power management instead
of drivers.  This is analogous to the block layer's existing runtime
power management support (blk-pm), which handles resuming devices when
bios are submitted to them so that drivers don't need to handle it.

Test: Tested on coral with:
        echo 5 > /sys/bus/platform/devices/1d84000.ufshc/rpm_lvl
        sleep 30
        touch /data && sync  # hangs before this fix
  Also verified via kvm-xfstests that blk-crypto-fallback continues
  to work both with and without CONFIG_PM=y.

Bug: 137270441
Bug: 149368295
Change-Id: I6bc9fb81854afe7edf490d71796ee68a61f7cbc8
Signed-off-by: Eric Biggers <ebiggers@google.com>
diff --git a/block/blk-crypto-fallback.c b/block/blk-crypto-fallback.c
index 5039087..7886851 100644
--- a/block/blk-crypto-fallback.c
+++ b/block/blk-crypto-fallback.c
@@ -609,7 +609,7 @@
 		crypto_mode_supported[i] = 0xFFFFFFFF;
 	crypto_mode_supported[BLK_ENCRYPTION_MODE_INVALID] = 0;
 
-	blk_crypto_ksm = keyslot_manager_create(blk_crypto_num_keyslots,
+	blk_crypto_ksm = keyslot_manager_create(NULL, blk_crypto_num_keyslots,
 						&blk_crypto_ksm_ll_ops,
 						crypto_mode_supported, NULL);
 	if (!blk_crypto_ksm)
diff --git a/block/keyslot-manager.c b/block/keyslot-manager.c
index 1436426..6b563c8 100644
--- a/block/keyslot-manager.c
+++ b/block/keyslot-manager.c
@@ -29,6 +29,7 @@
 #include <linux/keyslot-manager.h>
 #include <linux/atomic.h>
 #include <linux/mutex.h>
+#include <linux/pm_runtime.h>
 #include <linux/wait.h>
 #include <linux/blkdev.h>
 
@@ -45,6 +46,11 @@
 	unsigned int crypto_mode_supported[BLK_ENCRYPTION_MODE_MAX];
 	void *ll_priv_data;
 
+#ifdef CONFIG_PM
+	/* Device for runtime power management (NULL if none) */
+	struct device *dev;
+#endif
+
 	/* Protects programming and evicting keys from the device */
 	struct rw_semaphore lock;
 
@@ -71,8 +77,60 @@
 	return ksm->num_slots == 0;
 }
 
+#ifdef CONFIG_PM
+static inline void keyslot_manager_set_dev(struct keyslot_manager *ksm,
+					   struct device *dev)
+{
+	ksm->dev = dev;
+}
+
+/* If there's an underlying device and it's suspended, resume it. */
+static inline void keyslot_manager_pm_get(struct keyslot_manager *ksm)
+{
+	if (ksm->dev)
+		pm_runtime_get_sync(ksm->dev);
+}
+
+static inline void keyslot_manager_pm_put(struct keyslot_manager *ksm)
+{
+	if (ksm->dev)
+		pm_runtime_put_sync(ksm->dev);
+}
+#else /* CONFIG_PM */
+static inline void keyslot_manager_set_dev(struct keyslot_manager *ksm,
+					   struct device *dev)
+{
+}
+
+static inline void keyslot_manager_pm_get(struct keyslot_manager *ksm)
+{
+}
+
+static inline void keyslot_manager_pm_put(struct keyslot_manager *ksm)
+{
+}
+#endif /* !CONFIG_PM */
+
+static inline void keyslot_manager_hw_enter(struct keyslot_manager *ksm)
+{
+	/*
+	 * Calling into the driver requires ksm->lock held and the device
+	 * resumed.  But we must resume the device first, since that can acquire
+	 * and release ksm->lock via keyslot_manager_reprogram_all_keys().
+	 */
+	keyslot_manager_pm_get(ksm);
+	down_write(&ksm->lock);
+}
+
+static inline void keyslot_manager_hw_exit(struct keyslot_manager *ksm)
+{
+	up_write(&ksm->lock);
+	keyslot_manager_pm_put(ksm);
+}
+
 /**
  * keyslot_manager_create() - Create a keyslot manager
+ * @dev: Device for runtime power management (NULL if none)
  * @num_slots: The number of key slots to manage.
  * @ksm_ll_ops: The struct keyslot_mgmt_ll_ops for the device that this keyslot
  *		manager will use to perform operations like programming and
@@ -92,7 +150,9 @@
  * Context: May sleep
  * Return: Pointer to constructed keyslot manager or NULL on error.
  */
-struct keyslot_manager *keyslot_manager_create(unsigned int num_slots,
+struct keyslot_manager *keyslot_manager_create(
+	struct device *dev,
+	unsigned int num_slots,
 	const struct keyslot_mgmt_ll_ops *ksm_ll_ops,
 	const unsigned int crypto_mode_supported[BLK_ENCRYPTION_MODE_MAX],
 	void *ll_priv_data)
@@ -118,6 +178,7 @@
 	memcpy(ksm->crypto_mode_supported, crypto_mode_supported,
 	       sizeof(ksm->crypto_mode_supported));
 	ksm->ll_priv_data = ll_priv_data;
+	keyslot_manager_set_dev(ksm, dev);
 
 	init_rwsem(&ksm->lock);
 
@@ -226,10 +287,10 @@
 		return slot;
 
 	for (;;) {
-		down_write(&ksm->lock);
+		keyslot_manager_hw_enter(ksm);
 		slot = find_and_grab_keyslot(ksm, key);
 		if (slot != -ENOKEY) {
-			up_write(&ksm->lock);
+			keyslot_manager_hw_exit(ksm);
 			return slot;
 		}
 
@@ -240,7 +301,7 @@
 		if (!list_empty(&ksm->idle_slots))
 			break;
 
-		up_write(&ksm->lock);
+		keyslot_manager_hw_exit(ksm);
 		wait_event(ksm->idle_slots_wait_queue,
 			   !list_empty(&ksm->idle_slots));
 	}
@@ -252,7 +313,7 @@
 	err = ksm->ksm_ll_ops.keyslot_program(ksm, key, slot);
 	if (err) {
 		wake_up(&ksm->idle_slots_wait_queue);
-		up_write(&ksm->lock);
+		keyslot_manager_hw_exit(ksm);
 		return err;
 	}
 
@@ -266,7 +327,7 @@
 
 	remove_slot_from_lru_list(ksm, slot);
 
-	up_write(&ksm->lock);
+	keyslot_manager_hw_exit(ksm);
 	return slot;
 }
 
@@ -368,15 +429,16 @@
 
 	if (keyslot_manager_is_passthrough(ksm)) {
 		if (ksm->ksm_ll_ops.keyslot_evict) {
-			down_write(&ksm->lock);
+			keyslot_manager_hw_enter(ksm);
 			err = ksm->ksm_ll_ops.keyslot_evict(ksm, key, -1);
-			up_write(&ksm->lock);
+			keyslot_manager_hw_exit(ksm);
 			return err;
 		}
 		return 0;
 	}
 
-	down_write(&ksm->lock);
+	keyslot_manager_hw_enter(ksm);
+
 	slot = find_keyslot(ksm, key);
 	if (slot < 0) {
 		err = slot;
@@ -396,7 +458,7 @@
 	memzero_explicit(&slotp->key, sizeof(slotp->key));
 	err = 0;
 out_unlock:
-	up_write(&ksm->lock);
+	keyslot_manager_hw_exit(ksm);
 	return err;
 }
 
@@ -416,6 +478,7 @@
 	if (WARN_ON(keyslot_manager_is_passthrough(ksm)))
 		return;
 
+	/* This is for device initialization, so don't resume the device */
 	down_write(&ksm->lock);
 	for (slot = 0; slot < ksm->num_slots; slot++) {
 		const struct keyslot *slotp = &ksm->slots[slot];
@@ -455,6 +518,7 @@
 
 /**
  * keyslot_manager_create_passthrough() - Create a passthrough keyslot manager
+ * @dev: Device for runtime power management (NULL if none)
  * @ksm_ll_ops: The struct keyslot_mgmt_ll_ops
  * @crypto_mode_supported: Bitmasks for supported encryption modes
  * @ll_priv_data: Private data passed as is to the functions in ksm_ll_ops.
@@ -471,6 +535,7 @@
  * Return: Pointer to constructed keyslot manager or NULL on error.
  */
 struct keyslot_manager *keyslot_manager_create_passthrough(
+	struct device *dev,
 	const struct keyslot_mgmt_ll_ops *ksm_ll_ops,
 	const unsigned int crypto_mode_supported[BLK_ENCRYPTION_MODE_MAX],
 	void *ll_priv_data)
@@ -485,6 +550,7 @@
 	memcpy(ksm->crypto_mode_supported, crypto_mode_supported,
 	       sizeof(ksm->crypto_mode_supported));
 	ksm->ll_priv_data = ll_priv_data;
+	keyslot_manager_set_dev(ksm, dev);
 
 	init_rwsem(&ksm->lock);
 
@@ -544,15 +610,15 @@
 {
 	int err;
 
-	down_write(&ksm->lock);
 	if (ksm->ksm_ll_ops.derive_raw_secret) {
+		keyslot_manager_hw_enter(ksm);
 		err = ksm->ksm_ll_ops.derive_raw_secret(ksm, wrapped_key,
 							wrapped_key_size,
 							secret, secret_size);
+		keyslot_manager_hw_exit(ksm);
 	} else {
 		err = -EOPNOTSUPP;
 	}
-	up_write(&ksm->lock);
 
 	return err;
 }
diff --git a/drivers/md/dm.c b/drivers/md/dm.c
index ad059494..7dc5881 100644
--- a/drivers/md/dm.c
+++ b/drivers/md/dm.c
@@ -2315,7 +2315,8 @@
 	 */
 	memset(mode_masks, 0xFF, sizeof(mode_masks));
 
-	md->queue->ksm = keyslot_manager_create_passthrough(&dm_ksm_ll_ops,
+	md->queue->ksm = keyslot_manager_create_passthrough(NULL,
+							    &dm_ksm_ll_ops,
 							    mode_masks, md);
 	if (!md->queue->ksm)
 		return -ENOMEM;
diff --git a/drivers/scsi/ufs/ufshcd-crypto.c b/drivers/scsi/ufs/ufshcd-crypto.c
index 276b49a..e3de448 100644
--- a/drivers/scsi/ufs/ufshcd-crypto.c
+++ b/drivers/scsi/ufs/ufshcd-crypto.c
@@ -125,7 +125,6 @@
 	u32 slot_offset = hba->crypto_cfg_register + slot * sizeof(*cfg);
 	int err;
 
-	pm_runtime_get_sync(hba->dev);
 	ufshcd_hold(hba, false);
 
 	if (hba->vops->program_key) {
@@ -155,7 +154,6 @@
 	err = 0;
 out:
 	ufshcd_release(hba);
-	pm_runtime_put_sync(hba->dev);
 	return err;
 }
 
@@ -337,8 +335,8 @@
 
 	ufshcd_clear_all_keyslots(hba);
 
-	hba->ksm = keyslot_manager_create(ufshcd_num_keyslots(hba), ksm_ops,
-					  crypto_modes_supported, hba);
+	hba->ksm = keyslot_manager_create(hba->dev, ufshcd_num_keyslots(hba),
+					  ksm_ops, crypto_modes_supported, hba);
 
 	if (!hba->ksm) {
 		err = -ENOMEM;
diff --git a/include/linux/keyslot-manager.h b/include/linux/keyslot-manager.h
index 6d32a03..2f4aac2 100644
--- a/include/linux/keyslot-manager.h
+++ b/include/linux/keyslot-manager.h
@@ -41,7 +41,9 @@
 				 u8 *secret, unsigned int secret_size);
 };
 
-struct keyslot_manager *keyslot_manager_create(unsigned int num_slots,
+struct keyslot_manager *keyslot_manager_create(
+	struct device *dev,
+	unsigned int num_slots,
 	const struct keyslot_mgmt_ll_ops *ksm_ops,
 	const unsigned int crypto_mode_supported[BLK_ENCRYPTION_MODE_MAX],
 	void *ll_priv_data);
@@ -67,6 +69,7 @@
 void keyslot_manager_destroy(struct keyslot_manager *ksm);
 
 struct keyslot_manager *keyslot_manager_create_passthrough(
+	struct device *dev,
 	const struct keyslot_mgmt_ll_ops *ksm_ops,
 	const unsigned int crypto_mode_supported[BLK_ENCRYPTION_MODE_MAX],
 	void *ll_priv_data);