3 files changed, 36 insertions, 8 deletions
diff --git a/core/crypto/aes-gcm-sw.c b/core/crypto/aes-gcm-sw.c
index 5c59bf16..b651d5e0 100644
--- a/core/crypto/aes-gcm-sw.c
+++ b/core/crypto/aes-gcm-sw.c
@@ -79,12 +79,17 @@ void __weak
 internal_aes_gcm_encrypt_block(const struct internal_aes_gcm_key *ek,
 			       const void *src, void *dst)
 {
-	crypto_aes_enc_block(ek->data, ek->rounds, src, dst);
+	size_t ek_len = sizeof(ek->data);
+
+	crypto_aes_enc_block(ek->data, ek_len, ek->rounds, src, dst);
 }
 
 TEE_Result __weak
 internal_aes_gcm_expand_enc_key(const void *key, size_t key_len,
 				struct internal_aes_gcm_key *ek)
 {
-	return crypto_aes_expand_enc_key(key, key_len, ek->data, &ek->rounds);
+	size_t ek_len = sizeof(ek->data);
+
+	return crypto_aes_expand_enc_key(key, key_len, ek->data, ek_len,
+					&ek->rounds);
 }
diff --git a/core/include/crypto/crypto.h b/core/include/crypto/crypto.h
index 54a5f74f..a57d9cc1 100644
--- a/core/include/crypto/crypto.h
+++ b/core/include/crypto/crypto.h
@@ -316,9 +316,27 @@ void crypto_rng_add_event(enum crypto_rng_src sid, unsigned int *pnum,
  */
 TEE_Result crypto_rng_read(void *buf, size_t len);
 
+/*
+ * crypto_aes_expand_enc_key() - Expand an AES key
+ * @key:	AES key buffer
+ * @key_len:	Size of the the @key buffer in bytes
+ * @enc_key:	Expanded AES encryption key buffer
+ * @enc_keylen: Size of the @enc_key buffer in bytes
+ * @rounds:	Number of rounds to be used during encryption
+ */
 TEE_Result crypto_aes_expand_enc_key(const void *key, size_t key_len,
-				     void *enc_key, unsigned int *rounds);
-void crypto_aes_enc_block(const void *enc_key, unsigned int rounds,
-			  const void *src, void *dst);
+				     void *enc_key, size_t enc_keylen,
+				     unsigned int *rounds);
+
+/*
+ * crypto_aes_enc_block() - Encrypt an AES block
+ * @enc_key:	Expanded AES encryption key
+ * @enc_keylen:	Size of @enc_key in bytes
+ * @rounds:	Number of rounds
+ * @src:	Source buffer of one AES block (16 bytes)
+ * @dst:	Destination buffer of one AES block (16 bytes)
+ */
+void crypto_aes_enc_block(const void *enc_key, size_t enc_keylen,
+			  unsigned int rounds, const void *src, void *dst);
 
 #endif /* __CRYPTO_CRYPTO_H */
diff --git a/core/lib/libtomcrypt/src/tee_ltc_provider.c b/core/lib/libtomcrypt/src/tee_ltc_provider.c
index cae91b7d..75373fac 100644
--- a/core/lib/libtomcrypt/src/tee_ltc_provider.c
+++ b/core/lib/libtomcrypt/src/tee_ltc_provider.c
@@ -1441,10 +1441,14 @@ TEE_Result hash_sha512_256_compute(uint8_t *digest, const uint8_t *data,
 #endif
 
 TEE_Result crypto_aes_expand_enc_key(const void *key, size_t key_len,
-				     void *enc_key, unsigned int *rounds)
+				     void *enc_key, size_t enc_keylen,
+				     unsigned int *rounds)
 {
 	symmetric_key skey;
 
+	if (enc_keylen < sizeof(skey.rijndael.eK))
+		return TEE_ERROR_BAD_PARAMETERS;
+
 	if (aes_setup(key, key_len, 0, &skey))
 		return TEE_ERROR_BAD_PARAMETERS;
 
@@ -1453,11 +1457,12 @@ TEE_Result crypto_aes_expand_enc_key(const void *key, size_t key_len,
 	return TEE_SUCCESS;
 }
 
-void crypto_aes_enc_block(const void *enc_key, unsigned int rounds,
-			  const void *src, void *dst)
+void crypto_aes_enc_block(const void *enc_key, size_t enc_keylen __maybe_unused,
+			  unsigned int rounds, const void *src, void *dst)
 {
 	symmetric_key skey;
 
+	assert(enc_keylen >= sizeof(skey.rijndael.eK));
 	memcpy(skey.rijndael.eK, enc_key, sizeof(skey.rijndael.eK));
 	skey.rijndael.Nr = rounds;
 	if (aes_ecb_encrypt(src, dst, &skey))