crypto: tegra: Reserve keyslots to allocate dynamically

The HW supports only storing 15 keys at a time. This limits the number
of tfms that can work without failutes. Reserve keyslots to solve this
and use the reserved ones during the encryption/decryption operation.
This allow users to have the capability of hardware protected keys
and faster operations if there are limited number of tfms while not
halting the operation if there are more tfms.

Bug 4883011

Signed-off-by: Akhil R <akhilrajeev@nvidia.com>
Change-Id: I7de17eb1acf3b5a9f55a42e9df2aa8b64e20cb6d
Reviewed-on: https://git-master.nvidia.com/r/c/linux-nv-oot/+/3347661
GVS: buildbot_gerritrpt <buildbot_gerritrpt@nvidia.com>
Reviewed-by: Laxman Dewangan <ldewangan@nvidia.com>

drivers/crypto/tegra/tegra-se-aes.c

@@ -33,6 +33,9 @@ struct tegra_aes_ctx {
 	u32 ivsize;
 	u32 key1_id;
 	u32 key2_id;
+	u32 keylen;
+	u8 key1[AES_MAX_KEY_SIZE];
+	u8 key2[AES_MAX_KEY_SIZE];
 };
 
 struct tegra_aes_reqctx {
@@ -58,6 +61,7 @@ struct tegra_aead_ctx {
 	u32 verify_alg;
 	u32 keylen;
 	u32 key_id;
+	u8 key[AES_MAX_KEY_SIZE];
 };
 
 struct tegra_aead_reqctx {
@@ -84,6 +88,8 @@ struct tegra_cmac_ctx {
 	u32 alg;
 	u32 final_alg;
 	u32 key_id;
+	u32 keylen;
+	u8 key[AES_MAX_KEY_SIZE];
 	struct crypto_shash *fallback_tfm;
 };
 
@@ -285,20 +291,29 @@ static int tegra_aes_do_one_req(struct crypto_engine *engine, void *areq)
 	int ret;
 
 	/* Keys in ctx might be stored in KDS. Copy it to request ctx */
-	rctx->key1_id = tegra_key_get_idx(ctx->se, ctx->key1_id);
+	if (ctx->key1_id)
-	if (!rctx->key1_id) {
+		rctx->key1_id = tegra_key_get_idx(ctx->se, ctx->key1_id);
-		ret = -ENOMEM;
+
-		goto out;
+	/* Use reserved keyslots if keyslots are unavailable */
+	if (!ctx->key1_id || !rctx->key1_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key1,
+					ctx->keylen, ctx->alg, &rctx->key1_id);
+		if (ret)
+			goto out;
 	}
 
 	rctx->key2_id = 0;
 
 	/* If there are 2 keys stored (for XTS), retrieve them both */
-	if (ctx->key2_id) {
+	if (ctx->alg == SE_ALG_XTS) {
-		rctx->key2_id = tegra_key_get_idx(ctx->se, ctx->key2_id);
+		if (ctx->key2_id)
-		if (!rctx->key2_id) {
+			rctx->key2_id = tegra_key_get_idx(ctx->se, ctx->key2_id);
-			ret = -ENOMEM;
+
-			goto key1_free;
+		if (!ctx->key2_id || !rctx->key2_id) {
+			ret = tegra_key_submit_reserved_xts(ctx->se, ctx->key2,
+						ctx->keylen, ctx->alg, &rctx->key2_id);
+			if (ret)
+				goto key1_free;
 		}
 	}
 
@@ -340,10 +355,14 @@ static int tegra_aes_do_one_req(struct crypto_engine *engine, void *areq)
 			  rctx->datbuf.buf, rctx->datbuf.addr);
 
 key2_free:
-	if (rctx->key2_id != ctx->key2_id)
+	if (tegra_key_is_reserved(rctx->key2_id))
+		tegra_key_invalidate_reserved(ctx->se, rctx->key2_id, ctx->alg);
+	else if (rctx->key2_id != ctx->key2_id)
 		tegra_key_invalidate(ctx->se, rctx->key2_id, ctx->alg);
 key1_free:
-	if (rctx->key1_id != ctx->key1_id)
+	if (tegra_key_is_reserved(rctx->key1_id))
+		tegra_key_invalidate_reserved(ctx->se, rctx->key1_id, ctx->alg);
+	else if (rctx->key1_id != ctx->key1_id)
 		tegra_key_invalidate(ctx->se, rctx->key1_id, ctx->alg);
 out:
 	crypto_finalize_skcipher_request(se->engine, req, ret);
@@ -371,6 +390,7 @@ static int tegra_aes_cra_init(struct crypto_skcipher *tfm)
 	ctx->se = se_alg->se_dev;
 	ctx->key1_id = 0;
 	ctx->key2_id = 0;
+	ctx->keylen = 0;
 
 	algname = crypto_tfm_alg_name(&tfm->base);
 	ret = se_algname_to_algid(algname);
@@ -405,13 +425,20 @@ static int tegra_aes_setkey(struct crypto_skcipher *tfm,
 			    const u8 *key, u32 keylen)
 {
 	struct tegra_aes_ctx *ctx = crypto_skcipher_ctx(tfm);
+	int ret;
 
 	if (aes_check_keylen(keylen)) {
 		dev_dbg(ctx->se->dev, "invalid key length (%d)\n", keylen);
 		return -EINVAL;
 	}
 
-	return tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key1_id);
+	ret = tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key1_id);
+	if (ret) {
+		ctx->keylen = keylen;
+		memcpy(ctx->key1, key, keylen);
+	}
+
+	return 0;
 }
 
 static int tegra_xts_setkey(struct crypto_skcipher *tfm,
@@ -429,11 +456,17 @@ static int tegra_xts_setkey(struct crypto_skcipher *tfm,
 
 	ret = tegra_key_submit(ctx->se, key, len,
 			       ctx->alg, &ctx->key1_id);
-	if (ret)
+	if (ret) {
-		return ret;
+		ctx->keylen = len;
+		memcpy(ctx->key1, key, len);
+	}
 
-	return tegra_key_submit(ctx->se, key + len, len,
+	ret = tegra_key_submit(ctx->se, key + len, len,
 			       ctx->alg, &ctx->key2_id);
+	if (ret) {
+		ctx->keylen = len;
+		memcpy(ctx->key2, key + len, len);
+	}
 
 	return 0;
 }
@@ -1416,9 +1449,16 @@ static int tegra_ccm_do_one_req(struct crypto_engine *engine, void *areq)
 		rctx->cryptlen = req->cryptlen - ctx->authsize;
 
 	/* Keys in ctx might be stored in KDS. Copy it to local keyslot */
-	rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
+	if (ctx->key_id)
-	if (!rctx->key_id)
+		rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
-		goto out;
+
+	/* Use reserved keyslots if keyslots are unavailable */
+	if (!ctx->key_id || !rctx->key_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key,
+					ctx->keylen, ctx->alg, &rctx->key_id);
+		if (ret)
+			goto out;
+	}
 
 	rctx->inbuf.size = rctx->assoclen + rctx->authsize + rctx->cryptlen + 100;
 	/* Allocate buffers required */
@@ -1472,8 +1512,9 @@ inbuf_free:
 	dma_free_coherent(ctx->se->dev, rctx->inbuf.size,
 			  rctx->inbuf.buf, rctx->inbuf.addr);
 key_free:
-	/* Free the keyslot if it is cloned for this request */
+	if (tegra_key_is_reserved(rctx->key_id))
-	if (rctx->key_id != ctx->key_id)
+		tegra_key_invalidate_reserved(ctx->se, rctx->key_id, ctx->alg);
+	else if (rctx->key_id != ctx->key_id)
 		tegra_key_invalidate(ctx->se, rctx->key_id, ctx->alg);
 out:
 	crypto_finalize_aead_request(ctx->se->engine, req, ret);
@@ -1502,9 +1543,16 @@ static int tegra_gcm_do_one_req(struct crypto_engine *engine, void *areq)
 
 
 	/* Keys in ctx might be stored in KDS. Copy it to local keyslot */
-	rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
+	if (ctx->key_id)
-	if (!rctx->key_id)
+		rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
-		goto key_err;
+
+	/* Use reserved keyslots if keyslots are unavailable */
+	if (!ctx->key_id || !rctx->key_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key,
+					ctx->keylen, ctx->alg, &rctx->key_id);
+		if (ret)
+			goto key_err;
+	}
 
 	/* Allocate buffers required */
 	rctx->inbuf.size = rctx->assoclen + rctx->authsize + rctx->cryptlen;
@@ -1552,7 +1600,9 @@ static int tegra_gcm_do_one_req(struct crypto_engine *engine, void *areq)
 		ret = tegra_gcm_do_verify(ctx, rctx);
 
 out:
-	if (rctx->key_id != ctx->key_id)
+	if (tegra_key_is_reserved(rctx->key_id))
+		tegra_key_invalidate_reserved(ctx->se, rctx->key_id, ctx->alg);
+	else if (rctx->key_id != ctx->key_id)
 		tegra_key_invalidate(ctx->se, rctx->key_id, ctx->alg);
 key_err:
 	dma_free_coherent(ctx->se->dev, rctx->outbuf.size,
@@ -1588,6 +1638,7 @@ static int tegra_ccm_cra_init(struct crypto_aead *tfm)
 
 	ctx->se = se_alg->se_dev;
 	ctx->key_id = 0;
+	ctx->keylen = 0;
 
 	ret = se_algname_to_algid(algname);
 	if (ret < 0) {
@@ -1622,6 +1673,7 @@ static int tegra_gcm_cra_init(struct crypto_aead *tfm)
 
 	ctx->se = se_alg->se_dev;
 	ctx->key_id = 0;
+	ctx->keylen = 0;
 
 	ctx->alg = SE_ALG_GCM;
 	ctx->final_alg = SE_ALG_GCM_FINAL;
@@ -1705,13 +1757,20 @@ static int tegra_aead_setkey(struct crypto_aead *tfm,
 			     const u8 *key, u32 keylen)
 {
 	struct tegra_aead_ctx *ctx = crypto_aead_ctx(tfm);
+	int ret;
 
 	if (aes_check_keylen(keylen)) {
 		dev_dbg(ctx->se->dev, "invalid key length (%d)\n", keylen);
 		return -EINVAL;
 	}
 
-	return tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	ret = tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	if (ret) {
+		ctx->keylen = keylen;
+		memcpy(ctx->key, key, keylen);
+	}
+
+	return 0;
 }
 
 static unsigned int tegra_cmac_prep_cmd(struct tegra_se *se, struct tegra_cmac_reqctx *rctx)
@@ -1986,6 +2045,18 @@ static int tegra_cmac_do_one_req(struct crypto_engine *engine, void *areq)
 		rctx->task &= ~SHA_INIT;
 	}
 
+	/* Keys in ctx might be stored in KDS. Copy it to local keyslot */
+	if (ctx->key_id)
+		rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
+
+	/* Use reserved keyslots if keyslots are unavailable */
+	if (!ctx->key_id || !rctx->key_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key,
+					ctx->keylen, ctx->alg, &rctx->key_id);
+		if (ret)
+			goto out;
+	}
+
 	if (rctx->task & SHA_UPDATE) {
 		ret = tegra_cmac_do_update(req);
 		if (ret)
@@ -2003,6 +2074,11 @@ static int tegra_cmac_do_one_req(struct crypto_engine *engine, void *areq)
 	}
 
 out:
+	if (tegra_key_is_reserved(rctx->key_id))
+		tegra_key_invalidate_reserved(ctx->se, rctx->key_id, ctx->alg);
+	else if (rctx->key_id != ctx->key_id)
+		tegra_key_invalidate(ctx->se, rctx->key_id, ctx->alg);
+
 	crypto_finalize_hash_request(se->engine, req, ret);
 
 	return 0;
@@ -2075,6 +2151,7 @@ static int tegra_cmac_setkey(struct crypto_ahash *tfm, const u8 *key,
 			     unsigned int keylen)
 {
 	struct tegra_cmac_ctx *ctx = crypto_ahash_ctx(tfm);
+	int ret;
 
 	if (aes_check_keylen(keylen)) {
 		dev_dbg(ctx->se->dev, "invalid key length (%d)\n", keylen);
@@ -2084,7 +2161,13 @@ static int tegra_cmac_setkey(struct crypto_ahash *tfm, const u8 *key,
 	if (ctx->fallback_tfm)
 		crypto_shash_setkey(ctx->fallback_tfm, key, keylen);
 
-	return tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	ret = tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	if (ret) {
+		ctx->keylen = keylen;
+		memcpy(ctx->key, key, keylen);
+	}
+
+	return 0;
 }
 
 static int tegra_cmac_init(struct ahash_request *req)

drivers/crypto/tegra/tegra-se-hash.c

@@ -655,13 +655,18 @@ static int tegra_hmac_setkey(struct crypto_ahash *tfm, const u8 *key,
 			     unsigned int keylen)
 {
 	struct tegra_sha_ctx *ctx = crypto_ahash_ctx(tfm);
+	int ret;
 
 	if (aes_check_keylen(keylen))
 		return tegra_hmac_fallback_setkey(ctx, key, keylen);
 
+	ret = tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	if (ret)
+		return tegra_hmac_fallback_setkey(ctx, key, keylen);
+
 	ctx->fallback = false;
 
-	return tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	return 0;
 }
 
 static int tegra_sha_init(struct ahash_request *req)

drivers/crypto/tegra/tegra-se-key.c

@@ -314,3 +314,20 @@ int tegra_key_submit(struct tegra_se *se, const u8 *key, u32 keylen, u32 alg, u3
 
 	return 0;
 }
+
+void tegra_key_invalidate_reserved(struct tegra_se *se, u32 keyid, u32 alg)
+{
+	u8 zkey[AES_MAX_KEY_SIZE] = {0};
+
+	if (!keyid)
+		return;
+
+	/* Overwrite the key with 0s */
+	tegra_key_insert(se, zkey, AES_MAX_KEY_SIZE, keyid, alg);
+}
+
+inline int tegra_key_submit_reserved(struct tegra_se *se, const u8 *key,
+				     u32 keylen, u32 alg, u32 *keyid)
+{
+	return tegra_key_insert(se, key, keylen, *keyid, alg);
+}

drivers/crypto/tegra/tegra-se-sm4.c

@@ -34,6 +34,8 @@ struct tegra_sm4_ctx {
 	u32 ivsize;
 	u32 key1_id;
 	u32 key2_id;
+	u8 key1[AES_MAX_KEY_SIZE];
+	u8 key2[AES_MAX_KEY_SIZE];
 };
 
 struct tegra_sm4_reqctx {
@@ -58,6 +60,7 @@ struct tegra_sm4_gcm_ctx {
 	u32 verify_alg;
 	u32 keylen;
 	u32 key_id;
+	u8 key[AES_MAX_KEY_SIZE];
 };
 
 struct tegra_sm4_gcm_reqctx {
@@ -84,6 +87,8 @@ struct tegra_sm4_cmac_ctx {
 	u32 alg;
 	u32 final_alg;
 	u32 key_id;
+	u32 keylen;
+	u8 key[AES_MAX_KEY_SIZE];
 	struct crypto_shash *fallback_tfm;
 };
 
@@ -199,20 +204,30 @@ static int tegra_sm4_do_one_req(struct crypto_engine *engine, void *areq)
 	iv = (u32 *)req->iv;
 
 	/* Keys in ctx might be stored in KDS. Copy it to request ctx */
-	rctx->key1_id = tegra_key_get_idx(ctx->se, ctx->key1_id);
+	if (ctx->key1_id)
-	if (!rctx->key1_id) {
+		rctx->key1_id = tegra_key_get_idx(ctx->se, ctx->key1_id);
-		ret = -ENOMEM;
+
-		goto out;
+	/* Use reserved keyslots if keyslots are unavailable */
+	if (!ctx->key1_id || !rctx->key1_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key1,
+					ctx->keylen, ctx->alg, &rctx->key1_id);
+		if (ret)
+			goto out;
 	}
 
 	rctx->key2_id = 0;
 
 	/* If there are 2 keys stored (for XTS), retrieve them both */
-	if (ctx->key2_id) {
+	if (ctx->alg == SE_ALG_SM4_XTS) {
-		rctx->key2_id = tegra_key_get_idx(ctx->se, ctx->key2_id);
+		if (ctx->key2_id)
-		if (!rctx->key2_id) {
+			rctx->key1_id = tegra_key_get_idx(ctx->se, ctx->key1_id);
-			ret = -ENOMEM;
+
-			goto key1_free;
+		/* Use reserved keyslots if keyslots are unavailable */
+		if (!ctx->key2_id || !rctx->key2_id) {
+			ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key2,
+						ctx->keylen, ctx->alg, &rctx->key2_id);
+			if (ret)
+				goto key1_free;
 		}
 	}
 
@@ -255,11 +270,14 @@ static int tegra_sm4_do_one_req(struct crypto_engine *engine, void *areq)
 			  rctx->datbuf.buf, rctx->datbuf.addr);
 
 key2_free:
-	/* Free the keyslots if it is cloned for this request */
+	if (tegra_key_is_reserved(rctx->key2_id))
-	if (rctx->key2_id != ctx->key2_id)
+		tegra_key_invalidate_reserved(ctx->se, rctx->key2_id, ctx->alg);
+	else if (rctx->key2_id != ctx->key2_id)
 		tegra_key_invalidate(ctx->se, rctx->key2_id, ctx->alg);
 key1_free:
-	if (rctx->key1_id != ctx->key1_id)
+	if (tegra_key_is_reserved(rctx->key1_id))
+		tegra_key_invalidate_reserved(ctx->se, rctx->key1_id, ctx->alg);
+	else if (rctx->key1_id != ctx->key1_id)
 		tegra_key_invalidate(ctx->se, rctx->key1_id, ctx->alg);
 out:
 	crypto_finalize_skcipher_request(se->engine, req, ret);
@@ -286,6 +304,7 @@ static int tegra_sm4_cra_init(struct crypto_skcipher *tfm)
 	ctx->se = se_alg->se_dev;
 	ctx->key1_id = 0;
 	ctx->key2_id = 0;
+	ctx->keylen = 0;
 	ctx->ivsize = crypto_skcipher_ivsize(tfm);
 
 	algname = crypto_tfm_alg_name(&tfm->base);
@@ -319,13 +338,20 @@ static int tegra_sm4_setkey(struct crypto_skcipher *tfm,
 			       const u8 *key, u32 keylen)
 {
 	struct tegra_sm4_ctx *ctx = crypto_skcipher_ctx(tfm);
+	int ret;
 
 	if (aes_check_keylen(keylen)) {
 		dev_err(ctx->se->dev, "key length validation failed\n");
 		return -EINVAL;
 	}
 
-	return tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key1_id);
+	ret = tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key1_id);
+	if (ret) {
+		ctx->keylen = keylen;
+		memcpy(ctx->key1, key, keylen);
+	}
+
+	return 0;
 }
 
 static int tegra_sm4_xts_setkey(struct crypto_skcipher *tfm,
@@ -343,13 +369,17 @@ static int tegra_sm4_xts_setkey(struct crypto_skcipher *tfm,
 
 	ret = tegra_key_submit(ctx->se, key, len,
 			ctx->alg, &ctx->key1_id);
-	if (ret)
+	if (ret) {
-		return ret;
+		ctx->keylen = len;
+		memcpy(ctx->key1, key, len);
+	}
 
 	ret = tegra_key_submit(ctx->se, key + len, len,
 			ctx->alg, &ctx->key2_id);
-	if (ret)
+	if (ret) {
-		return ret;
+		ctx->keylen = len;
+		memcpy(ctx->key2, key + len, len);
+	}
 
 	return 0;
 }
@@ -950,7 +980,7 @@ static int tegra_sm4_gcm_do_one_req(struct crypto_engine *engine, void *areq)
 	struct tegra_sm4_gcm_ctx *ctx = crypto_aead_ctx(tfm);
 	struct tegra_sm4_gcm_reqctx *rctx = aead_request_ctx(req);
 	struct tegra_se *se = ctx->se;
-	int ret, keyid;
+	int ret;
 
 	rctx->src_sg = req->src;
 	rctx->dst_sg = req->dst;
@@ -963,11 +993,15 @@ static int tegra_sm4_gcm_do_one_req(struct crypto_engine *engine, void *areq)
 		rctx->cryptlen = req->cryptlen - ctx->authsize;
 
 	/* Keys in ctx might be stored in KDS. Copy it to local keyslot */
-	keyid = tegra_key_get_idx(ctx->se, ctx->key_id);
+	if (ctx->key_id)
-	if (!keyid)
+		rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
-		goto out;
 
-	rctx->key_id = keyid;
+	if (!ctx->key_id || !rctx->key_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key,
+					ctx->keylen, ctx->alg, &rctx->key_id);
+		if (ret)
+			goto out;
+	}
 
 	rctx->inbuf.size = rctx->assoclen + rctx->authsize + rctx->cryptlen;
 	rctx->inbuf.buf = dma_alloc_coherent(ctx->se->dev, rctx->inbuf.size,
@@ -1016,8 +1050,9 @@ inbuf_free:
 	dma_free_coherent(ctx->se->dev, rctx->inbuf.size,
 			  rctx->inbuf.buf, rctx->inbuf.addr);
 key_free:
-	/* Free the keyslot if it is cloned for this request */
+	if (tegra_key_is_reserved(rctx->key_id))
-	if (rctx->key_id != ctx->key_id)
+		tegra_key_invalidate_reserved(ctx->se, rctx->key_id, ctx->alg);
+	else if (rctx->key_id != ctx->key_id)
 		tegra_key_invalidate(ctx->se, rctx->key_id, ctx->alg);
 
 out:
@@ -1042,6 +1077,7 @@ static int tegra_sm4_gcm_cra_init(struct crypto_aead *tfm)
 
 	ctx->se = se_alg->se_dev;
 	ctx->key_id = 0;
+	ctx->keylen = 0;
 	ctx->alg = SE_ALG_SM4_GCM;
 	ctx->final_alg = SE_ALG_SM4_GCM_FINAL;
 	ctx->verify_alg = SE_ALG_SM4_GCM_VERIFY;
@@ -1202,13 +1238,6 @@ static int tegra_sm4_cmac_do_init(struct ahash_request *req)
 	rctx->blk_size = crypto_ahash_blocksize(tfm);
 	rctx->digest.size = crypto_ahash_digestsize(tfm);
 
-	/* Retrieve the key slot for CMAC */
-	if (ctx->key_id) {
-		rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
-		if (!rctx->key_id)
-			return -ENOMEM;
-	}
-
 	rctx->digest.buf = dma_alloc_coherent(se->dev, rctx->digest.size,
 				&rctx->digest.addr, GFP_KERNEL);
 	if (!rctx->digest.buf)
@@ -1232,9 +1261,6 @@ resbuf_fail:
 	dma_free_coherent(se->dev, rctx->blk_size, rctx->digest.buf,
 				rctx->digest.addr);
 digbuf_fail:
-	if (rctx->key_id != ctx->key_id)
-		tegra_key_invalidate(ctx->se, rctx->key_id, ctx->alg);
-
 	return -ENOMEM;
 }
 
@@ -1384,6 +1410,17 @@ static int tegra_sm4_cmac_do_one_req(struct crypto_engine *engine, void *areq)
 		rctx->task &= ~SHA_INIT;
 	}
 
+	/* Retrieve the key slot for CMAC */
+	if (ctx->key_id)
+		rctx->key_id = tegra_key_get_idx(ctx->se, ctx->key_id);
+
+	if (!ctx->key_id || !rctx->key_id) {
+		ret = tegra_key_submit_reserved_aes(ctx->se, ctx->key,
+					ctx->keylen, ctx->alg, &rctx->key_id);
+		if (ret)
+			goto out;
+	}
+
 	if (rctx->task & SHA_UPDATE) {
 		ret = tegra_sm4_cmac_do_update(req);
 		if (ret)
@@ -1401,6 +1438,11 @@ static int tegra_sm4_cmac_do_one_req(struct crypto_engine *engine, void *areq)
 	}
 
 out:
+	if (tegra_key_is_reserved(rctx->key_id))
+		tegra_key_invalidate_reserved(ctx->se, rctx->key_id, ctx->alg);
+	else if (rctx->key_id != ctx->key_id)
+		tegra_key_invalidate(ctx->se, rctx->key_id, ctx->alg);
+
 	crypto_finalize_hash_request(se->engine, req, ret);
 
 	return ret;
@@ -1447,13 +1489,20 @@ static int tegra_sm4_cmac_setkey(struct crypto_ahash *tfm, const u8 *key,
 			     unsigned int keylen)
 {
 	struct tegra_sm4_cmac_ctx *ctx = crypto_ahash_ctx(tfm);
+	int ret;
 
 	if (aes_check_keylen(keylen)) {
 		dev_err(ctx->se->dev, "key length validation failed\n");
 		return -EINVAL;
 	}
 
-	return tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	ret = tegra_key_submit(ctx->se, key, keylen, ctx->alg, &ctx->key_id);
+	if (ret) {
+		ctx->keylen = keylen;
+		memcpy(ctx->key, key, keylen);
+	}
+
+	return 0;
 }
 
 static int tegra_sm4_cmac_init(struct ahash_request *req)

drivers/crypto/tegra/tegra-se.h

@@ -537,6 +537,9 @@
 #define SHA_UPDATE	BIT(2)
 #define SHA_FINAL	BIT(3)
 
+#define TEGRA_AES_RESERVED_KSLT			14
+#define TEGRA_XTS_RESERVED_KSLT			15
+
 #ifdef NV_CONFTEST_REMOVE_STRUCT_CRYPTO_ENGINE_CTX
 #define CRYPTO_REGISTER(alg, x) \
 		crypto_engine_register_##alg(x)
@@ -772,6 +775,9 @@ void tegra_deinit_sm4(struct tegra_se *se);
 int tegra_key_submit(struct tegra_se *se, const u8 *key,
 		     u32 keylen, u32 alg, u32 *keyid);
 void tegra_key_invalidate(struct tegra_se *se, u32 keyid, u32 alg);
+int tegra_key_submit_reserved(struct tegra_se *se, const u8 *key,
+				u32 keylen, u32 alg, u32 *keyid);
+void tegra_key_invalidate_reserved(struct tegra_se *se, u32 keyid, u32 alg);
 unsigned int tegra_key_get_idx(struct tegra_se *se, u32 keyid);
 int tegra_se_host1x_submit(struct tegra_se *se, struct tegra_se_cmdbuf *cmdbuf, u32 size);
 
@@ -779,6 +785,26 @@ u32 tegra_kds_get_id(void);
 void tegra_kds_free_id(u32 keyid);
 bool tegra_key_in_kds(u32 keyid);
 
+static inline int tegra_key_submit_reserved_aes(struct tegra_se *se, const u8 *key,
+						u32 keylen, u32 alg, u32 *keyid)
+{
+	*keyid = TEGRA_AES_RESERVED_KSLT;
+	return tegra_key_submit_reserved(se, key, keylen, alg, keyid);
+}
+
+static inline int tegra_key_submit_reserved_xts(struct tegra_se *se, const u8 *key,
+						u32 keylen, u32 alg, u32 *keyid)
+{
+	*keyid = TEGRA_XTS_RESERVED_KSLT;
+	return tegra_key_submit_reserved(se, key, keylen, alg, keyid);
+}
+
+static inline bool tegra_key_is_reserved(u32 keyid)
+{
+	return ((keyid == TEGRA_AES_RESERVED_KSLT) ||
+		(keyid == TEGRA_XTS_RESERVED_KSLT));
+}
+
 /* HOST1x OPCODES */
 static inline u32 host1x_opcode_setpayload(unsigned int payload)
 {