[PATCH v2 5/7] crypto: ti - Fix use-after-free of dev_data on DTHEv2 driver removal

[T Pratham <t-pratham@ti.com>]

Each *_init_tfm() caches a pointer to the per-instance struct dthe_data
in its transform context (ctx->dev_data), but never takes a reference on
it. If there are tfms in progress when dthe_remove() is called, the devm
allocated dev_data gets freed. Then ctx->dev_data will point to a memory
that has been freed.

Add a refcnt to struct dthe_data, incrementing it atomically in
*_init_tfm() and decrementing atomically in *_exit_tfm().

dthe_remove() now polls this count, with a bounded timeout, so tfms
allocated before removal have a chance to be freed first. If the timeout
expires, it warns and proceeds anyway rather than blocking removal
indefinitely.

Fixes: 52f641bc63a46 ("crypto: ti - Add driver for DTHE V2 AES Engine (ECB, CBC)")
Signed-off-by: T Pratham <t-pratham@ti.com>
---
 drivers/crypto/ti/dthev2-aes.c    |  4 ++++
 drivers/crypto/ti/dthev2-common.c | 22 +++++++++++++++++++++-
 drivers/crypto/ti/dthev2-common.h |  5 +++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/drivers/crypto/ti/dthev2-aes.c b/drivers/crypto/ti/dthev2-aes.c
index 8998e02b0e86e..8ef54dacf63f2 100644
--- a/drivers/crypto/ti/dthev2-aes.c
+++ b/drivers/crypto/ti/dthev2-aes.c
@@ -124,6 +124,7 @@ static int dthe_cipher_init_tfm(struct crypto_skcipher *tfm)
 	if (IS_ERR(ctx->skcipher_fb)) {
 		dev_err(dev_data->dev, "fallback driver %s couldn't be loaded\n",
 			alg_name);
+		dthe_put_dev(ctx);
 		return PTR_ERR(ctx->skcipher_fb);
 	}
 
@@ -135,6 +136,7 @@ static void dthe_cipher_exit_tfm(struct crypto_skcipher *tfm)
 	struct dthe_tfm_ctx *ctx = crypto_skcipher_ctx(tfm);
 
 	crypto_free_sync_skcipher(ctx->skcipher_fb);
+	dthe_put_dev(ctx);
 }
 
 static int dthe_aes_setkey(struct crypto_skcipher *tfm, const u8 *key, unsigned int keylen)
@@ -577,6 +579,7 @@ static int dthe_aead_init_tfm(struct crypto_aead *tfm)
 	if (IS_ERR(ctx->aead_fb)) {
 		dev_err(dev_data->dev, "fallback driver %s couldn't be loaded\n",
 			alg_name);
+		dthe_put_dev(ctx);
 		return PTR_ERR(ctx->aead_fb);
 	}
 
@@ -588,6 +591,7 @@ static void dthe_aead_exit_tfm(struct crypto_aead *tfm)
 	struct dthe_tfm_ctx *ctx = crypto_aead_ctx(tfm);
 
 	crypto_free_sync_aead(ctx->aead_fb);
+	dthe_put_dev(ctx);
 }
 
 /**
diff --git a/drivers/crypto/ti/dthev2-common.c b/drivers/crypto/ti/dthev2-common.c
index 8628187a32e18..c5bb7fe7f9876 100644
--- a/drivers/crypto/ti/dthev2-common.c
+++ b/drivers/crypto/ti/dthev2-common.c
@@ -27,6 +27,10 @@
 
 #define DRIVER_NAME	"dthev2"
 
+/* Interval and timeout for polling dthe_data::refcnt on removal */
+#define DTHE_REFCNT_POLL_INTERVAL_US	20000
+#define DTHE_REFCNT_POLL_TIMEOUT_US	1000000
+
 static struct dthe_list dthe_dev_list = {
 	.dev_list = LIST_HEAD_INIT(dthe_dev_list.dev_list),
 	.lock = __SPIN_LOCK_UNLOCKED(dthe_dev_list.lock),
@@ -41,13 +45,21 @@ struct dthe_data *dthe_get_dev(struct dthe_tfm_ctx *ctx)
 
 	spin_lock(&dthe_dev_list.lock);
 	dev_data = list_first_entry_or_null(&dthe_dev_list.dev_list, struct dthe_data, list);
-	if (dev_data)
+	if (dev_data) {
 		list_move_tail(&dev_data->list, &dthe_dev_list.dev_list);
+		atomic_inc(&dev_data->refcnt);
+	}
 	spin_unlock(&dthe_dev_list.lock);
 
 	return dev_data;
 }
 
+void dthe_put_dev(struct dthe_tfm_ctx *ctx)
+{
+	atomic_dec(&ctx->dev_data->refcnt);
+	ctx->dev_data = NULL;
+}
+
 struct scatterlist *dthe_copy_sg(struct scatterlist *dst,
 				 struct scatterlist *src,
 				 unsigned int buflen)
@@ -200,6 +212,7 @@ static int dthe_probe(struct platform_device *pdev)
 static void dthe_remove(struct platform_device *pdev)
 {
 	struct dthe_data *dev_data = platform_get_drvdata(pdev);
+	int refcnt, ret;
 
 	spin_lock(&dthe_dev_list.lock);
 	list_del(&dev_data->list);
@@ -207,6 +220,13 @@ static void dthe_remove(struct platform_device *pdev)
 
 	dthe_unregister_algs();
 
+	ret = readx_poll_timeout(atomic_read, &dev_data->refcnt, refcnt, !refcnt,
+				 DTHE_REFCNT_POLL_INTERVAL_US, DTHE_REFCNT_POLL_TIMEOUT_US);
+	if (ret)
+		dev_warn(dev_data->dev,
+			 "removing with %d transform context(s) still active\n",
+			 refcnt);
+
 	crypto_engine_exit(dev_data->engine);
 
 	dma_release_channel(dev_data->dma_aes_rx);
diff --git a/drivers/crypto/ti/dthev2-common.h b/drivers/crypto/ti/dthev2-common.h
index 75d9a097650da..c4adeb34b14f3 100644
--- a/drivers/crypto/ti/dthev2-common.h
+++ b/drivers/crypto/ti/dthev2-common.h
@@ -18,6 +18,7 @@
 #include <crypto/internal/hash.h>
 #include <crypto/internal/skcipher.h>
 
+#include <linux/atomic.h>
 #include <linux/delay.h>
 #include <linux/dmaengine.h>
 #include <linux/dmapool.h>
@@ -53,6 +54,7 @@ enum dthe_aes_mode {
  * @dma_aes_rx: AES Rx DMA Channel
  * @dma_aes_tx: AES Tx DMA Channel
  * @dma_sha_tx: SHA Tx DMA Channel
+ * @refcnt: Count of transform contexts currently holding a reference to this instance
  */
 struct dthe_data {
 	struct device *dev;
@@ -64,6 +66,8 @@ struct dthe_data {
 	struct dma_chan *dma_aes_tx;
 
 	struct dma_chan *dma_sha_tx;
+
+	atomic_t refcnt;
 };
 
 /**
@@ -113,6 +117,7 @@ struct dthe_aes_req_ctx {
 /* Struct definitions end */
 
 struct dthe_data *dthe_get_dev(struct dthe_tfm_ctx *ctx);
+void dthe_put_dev(struct dthe_tfm_ctx *ctx);
 
 /**
  * dthe_copy_sg - Copy sg entries from src to dst
-- 
2.34.1