diff --git a/drivers/crypto/inside-secure/safexcel_cipher.c b/drivers/crypto/inside-secure/safexcel_cipher.c
index 407ebcd8d71f..1aade2fd6f1e 100644
--- a/drivers/crypto/inside-secure/safexcel_cipher.c
+++ b/drivers/crypto/inside-secure/safexcel_cipher.c
@@ -354,10 +354,10 @@ static int safexcel_aead_setkey(struct crypto_aead *ctfm, const u8 *key,
 {
 	struct crypto_tfm *tfm = crypto_aead_tfm(ctfm);
 	struct safexcel_cipher_ctx *ctx = crypto_tfm_ctx(tfm);
-	struct safexcel_ahash_export_state istate, ostate;
+	struct safexcel_ahash_export_state *state;
 	struct safexcel_crypto_priv *priv = ctx->priv;
+	struct crypto_aes_ctx *aes;
 	struct crypto_authenc_keys keys;
-	struct crypto_aes_ctx aes;
 	int err = -EINVAL;
 
 	if (unlikely(crypto_authenc_extractkeys(&keys, key, len)))
@@ -387,7 +387,14 @@ static int safexcel_aead_setkey(struct crypto_aead *ctfm, const u8 *key,
 			goto badkey_expflags;
 		break;
 	case SAFEXCEL_AES:
-		err = aes_expandkey(&aes, keys.enckey, keys.enckeylen);
+		aes = kzalloc(sizeof(*aes), GFP_KERNEL);
+		if (!aes) {
+			err = -ENOMEM;
+			goto badkey;
+		}
+
+		err = aes_expandkey(aes, keys.enckey, keys.enckeylen);
+		kfree(aes);
 		if (unlikely(err))
 			goto badkey;
 		break;
@@ -404,61 +411,71 @@ static int safexcel_aead_setkey(struct crypto_aead *ctfm, const u8 *key,
 	    memcmp(ctx->key, keys.enckey, keys.enckeylen))
 		ctx->base.needs_inv = true;
 
+	state = kzalloc(sizeof(struct safexcel_ahash_export_state) * 2, GFP_KERNEL);
+	if (!state) {
+		err = -ENOMEM;
+		goto badkey;
+	}
+
 	/* Auth key */
 	switch (ctx->hash_alg) {
 	case CONTEXT_CONTROL_CRYPTO_ALG_SHA1:
 		if (safexcel_hmac_setkey("safexcel-sha1", keys.authkey,
-					 keys.authkeylen, &istate, &ostate))
-			goto badkey;
+					 keys.authkeylen, &state[0], &state[1]))
+			goto badkey_free;
 		break;
 	case CONTEXT_CONTROL_CRYPTO_ALG_SHA224:
 		if (safexcel_hmac_setkey("safexcel-sha224", keys.authkey,
-					 keys.authkeylen, &istate, &ostate))
-			goto badkey;
+					 keys.authkeylen, &state[0], &state[1]))
+			goto badkey_free;
 		break;
 	case CONTEXT_CONTROL_CRYPTO_ALG_SHA256:
 		if (safexcel_hmac_setkey("safexcel-sha256", keys.authkey,
-					 keys.authkeylen, &istate, &ostate))
-			goto badkey;
+					 keys.authkeylen, &state[0], &state[1]))
+			goto badkey_free;
 		break;
 	case CONTEXT_CONTROL_CRYPTO_ALG_SHA384:
 		if (safexcel_hmac_setkey("safexcel-sha384", keys.authkey,
-					 keys.authkeylen, &istate, &ostate))
-			goto badkey;
+					 keys.authkeylen, &state[0], &state[1]))
+			goto badkey_free;
 		break;
 	case CONTEXT_CONTROL_CRYPTO_ALG_SHA512:
 		if (safexcel_hmac_setkey("safexcel-sha512", keys.authkey,
-					 keys.authkeylen, &istate, &ostate))
-			goto badkey;
+					 keys.authkeylen, &state[0], &state[1]))
+			goto badkey_free;
 		break;
 	case CONTEXT_CONTROL_CRYPTO_ALG_SM3:
 		if (safexcel_hmac_setkey("safexcel-sm3", keys.authkey,
-					 keys.authkeylen, &istate, &ostate))
+					 keys.authkeylen, &state[0], &state[1]))
 			goto badkey;
 		break;
 	default:
 		dev_err(priv->dev, "aead: unsupported hash algorithm\n");
-		goto badkey;
+		goto badkey_free;
 	}
 
 	crypto_aead_set_flags(ctfm, crypto_aead_get_flags(ctfm) &
 				    CRYPTO_TFM_RES_MASK);
 
 	if (priv->flags & EIP197_TRC_CACHE && ctx->base.ctxr_dma &&
-	    (memcmp(ctx->ipad, istate.state, ctx->state_sz) ||
-	     memcmp(ctx->opad, ostate.state, ctx->state_sz)))
+	    (memcmp(ctx->ipad, &state[0].state, ctx->state_sz) ||
+	     memcmp(ctx->opad, &state[1].state, ctx->state_sz)))
 		ctx->base.needs_inv = true;
 
 	/* Now copy the keys into the context */
 	memcpy(ctx->key, keys.enckey, keys.enckeylen);
 	ctx->key_len = keys.enckeylen;
 
-	memcpy(ctx->ipad, &istate.state, ctx->state_sz);
-	memcpy(ctx->opad, &ostate.state, ctx->state_sz);
+	memcpy(ctx->ipad, &state[0].state, ctx->state_sz);
+	memcpy(ctx->opad, &state[1].state, ctx->state_sz);
 
 	memzero_explicit(&keys, sizeof(keys));
+	kfree(state);
+
 	return 0;
 
+badkey_free:
+	kfree(state);
 badkey:
 	crypto_aead_set_flags(ctfm, CRYPTO_TFM_RES_BAD_KEY_LEN);
 badkey_expflags:
