From 4fc6e7ca0cfe439061869d1e1cd83da8cf533a93 Mon Sep 17 00:00:00 2001
From: makejian <makejian@xiaomi.com>
Date: Fri, 20 Oct 2023 17:07:33 +0800
Subject: [PATCH] mbedtls-alt/aes-alt: Illegal parameter detection for aes
 related functions

(1)aes-xts sets the key length to only 256 and 512 bits
(2)when the key length of aes-xts is 512 bits, MAX_KEY_SIZE needs to be expanded to 64 bytes.
(3)check invalid input length and mode

Signed-off-by: makejian <makejian@xiaomi.com>
---
 crypto/mbedtls/include/aes_alt.h |  2 +-
 crypto/mbedtls/source/aes_alt.c  | 83 +++++++++++++++++++++++++++++++-
 2 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/crypto/mbedtls/include/aes_alt.h b/crypto/mbedtls/include/aes_alt.h
index c920b2bd3..590773829 100644
--- a/crypto/mbedtls/include/aes_alt.h
+++ b/crypto/mbedtls/include/aes_alt.h
@@ -30,7 +30,7 @@
  * Pre-processor Definitions
  ****************************************************************************/
 
-#define MAX_KEY_SIZE      36
+#define MAX_KEY_SIZE      64
 
 typedef struct mbedtls_aes_context
 {
diff --git a/crypto/mbedtls/source/aes_alt.c b/crypto/mbedtls/source/aes_alt.c
index 34ec26604..57d6cc80e 100644
--- a/crypto/mbedtls/source/aes_alt.c
+++ b/crypto/mbedtls/source/aes_alt.c
@@ -51,6 +51,18 @@ int mbedtls_aes_setkey_enc(FAR mbedtls_aes_context *ctx,
                            FAR const unsigned char *key,
                            unsigned int keybits)
 {
+  switch (keybits)
+    {
+      case 128:
+        break;
+      case 192:
+        break;
+      case 256:
+        break;
+      default:
+        return MBEDTLS_ERR_AES_INVALID_KEY_LENGTH;
+    }
+
   memcpy(ctx->key, key, keybits / 8);
   ctx->dev.session.key = (caddr_t)ctx->key;
   ctx->dev.session.keylen = keybits / 8;
@@ -74,6 +86,11 @@ int mbedtls_aes_crypt_ecb(FAR mbedtls_aes_context *ctx,
   int ret;
   unsigned char iv[16];
 
+  if (mode != MBEDTLS_AES_ENCRYPT && mode != MBEDTLS_AES_DECRYPT)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_CBC;
   ret = cryptodev_get_session(&ctx->dev);
   if (ret != 0)
@@ -107,6 +124,16 @@ int mbedtls_aes_crypt_cbc(mbedtls_aes_context *ctx,
 {
   int ret;
 
+  if (mode != MBEDTLS_AES_ENCRYPT && mode != MBEDTLS_AES_DECRYPT)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
+  if (length % 16)
+    {
+      return MBEDTLS_ERR_AES_INVALID_INPUT_LENGTH;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_CBC;
   ret = cryptodev_get_session(&ctx->dev);
   if (ret != 0)
@@ -141,6 +168,11 @@ int mbedtls_aes_crypt_ctr(FAR mbedtls_aes_context *ctx,
 {
   int ret;
 
+  if (*nc_off > 0x0f)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_CTR;
   memcpy(ctx->key + ctx->dev.session.keylen,
          nonce_counter, NONCE_LENGTH);
@@ -183,14 +215,22 @@ int mbedtls_aes_xts_setkey_enc(FAR mbedtls_aes_xts_context *ctx,
                                FAR const unsigned char *key,
                                unsigned int keybits)
 {
-  return mbedtls_aes_setkey_enc(ctx, key, keybits);
+  if (keybits != 256 && keybits != 512)
+    {
+      return MBEDTLS_ERR_AES_INVALID_KEY_LENGTH;
+    }
+
+  memcpy(ctx->key, key, keybits / 8);
+  ctx->dev.session.key = (caddr_t)ctx->key;
+  ctx->dev.session.keylen = keybits / 8;
+  return 0;
 }
 
 int mbedtls_aes_xts_setkey_dec(FAR mbedtls_aes_xts_context *ctx,
                                FAR const unsigned char *key,
                                unsigned int keybits)
 {
-  return mbedtls_aes_setkey_dec(ctx, key, keybits);
+  return mbedtls_aes_xts_setkey_enc(ctx, key, keybits);
 }
 
 int mbedtls_aes_crypt_xts(FAR mbedtls_aes_xts_context *ctx,
@@ -203,6 +243,25 @@ int mbedtls_aes_crypt_xts(FAR mbedtls_aes_xts_context *ctx,
   int ret;
   unsigned char iv[16];
 
+  if (mode != MBEDTLS_AES_ENCRYPT && mode != MBEDTLS_AES_DECRYPT)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
+  /* Data units must be at least 16 bytes long. */
+
+  if (length < 16)
+    {
+      return MBEDTLS_ERR_AES_INVALID_INPUT_LENGTH;
+    }
+
+  /* NIST SP 800-38E disallows data units larger than 2**20 blocks. */
+
+  if (length > (1 << 20) * 16)
+    {
+      return MBEDTLS_ERR_AES_INVALID_INPUT_LENGTH;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_XTS;
   ret = cryptodev_get_session(&ctx->dev);
   if (ret != 0)
@@ -238,6 +297,16 @@ int mbedtls_aes_crypt_cfb128(FAR mbedtls_aes_context *ctx,
 {
   int ret;
 
+  if (mode != MBEDTLS_AES_ENCRYPT && mode != MBEDTLS_AES_DECRYPT)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
+  if (*iv_off > 15)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_CFB_128;
   ret = cryptodev_get_session(&ctx->dev);
   if (ret != 0)
@@ -273,6 +342,11 @@ int mbedtls_aes_crypt_cfb8(FAR mbedtls_aes_context *ctx,
 {
   int ret;
 
+  if (mode != MBEDTLS_AES_ENCRYPT && mode != MBEDTLS_AES_DECRYPT)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_CFB_8;
   ret = cryptodev_get_session(&ctx->dev);
   if (ret != 0)
@@ -306,6 +380,11 @@ int mbedtls_aes_crypt_ofb(FAR mbedtls_aes_context *ctx,
 {
   int ret;
 
+  if (*iv_off > 15)
+    {
+      return MBEDTLS_ERR_AES_BAD_INPUT_DATA;
+    }
+
   ctx->dev.session.cipher = CRYPTO_AES_OFB;
   ret = cryptodev_get_session(&ctx->dev);
   if (ret != 0)