boot: Add AES256 support for image encryption

Support only works when using mbedtls as the cryptographic library.

Signed-off-by: Salome Thirot <salome.thirot@arm.com>
diff --git a/boot/bootutil/include/bootutil/caps.h b/boot/bootutil/include/bootutil/caps.h
index 8ead3be..9bf060c 100644
--- a/boot/bootutil/include/bootutil/caps.h
+++ b/boot/bootutil/include/bootutil/caps.h
@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2017 Linaro Limited
+ * Copyright (c) 2021 Arm Limited
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,6 +48,7 @@
 #define BOOTUTIL_CAP_DOWNGRADE_PREVENTION   (1<<12)
 #define BOOTUTIL_CAP_ENC_X25519             (1<<13)
 #define BOOTUTIL_CAP_BOOTSTRAP              (1<<14)
+#define BOOTUTIL_CAP_AES256                 (1<<15)
 
 /*
  * Query the number of images this bootloader is configured for.  This
diff --git a/boot/bootutil/include/bootutil/crypto/aes_ctr.h b/boot/bootutil/include/bootutil/crypto/aes_ctr.h
index 8422f45..e69b037 100644
--- a/boot/bootutil/include/bootutil/crypto/aes_ctr.h
+++ b/boot/bootutil/include/bootutil/crypto/aes_ctr.h
@@ -21,11 +21,15 @@
 
 #if defined(MCUBOOT_USE_MBED_TLS)
     #include <mbedtls/aes.h>
-    #define BOOTUTIL_CRYPTO_AES_CTR_KEY_SIZE (16)
+    #include "bootutil/enc_key_public.h"
+    #define BOOTUTIL_CRYPTO_AES_CTR_KEY_SIZE BOOT_ENC_KEY_SIZE
     #define BOOTUTIL_CRYPTO_AES_CTR_BLOCK_SIZE (16)
 #endif /* MCUBOOT_USE_MBED_TLS */
 
 #if defined(MCUBOOT_USE_TINYCRYPT)
+    #if defined(MCUBOOT_AES_256)
+        #error "Cannot use AES-256 for encryption with Tinycrypt library."
+    #endif
     #include <string.h>
     #include <tinycrypt/aes.h>
     #include <tinycrypt/ctr_mode.h>
diff --git a/boot/bootutil/include/bootutil/crypto/aes_kw.h b/boot/bootutil/include/bootutil/crypto/aes_kw.h
index 925f46a..cf3194f 100644
--- a/boot/bootutil/include/bootutil/crypto/aes_kw.h
+++ b/boot/bootutil/include/bootutil/crypto/aes_kw.h
@@ -23,6 +23,9 @@
 #endif /* MCUBOOT_USE_MBED_TLS */
 
 #if defined(MCUBOOT_USE_TINYCRYPT)
+    #if defined(MCUBOOT_AES_256)
+        #error "Cannot use AES-256 for encryption with Tinycrypt library."
+    #endif
     #include <tinycrypt/aes.h>
     #include <tinycrypt/constants.h>
 #endif /* MCUBOOT_USE_TINYCRYPT */
diff --git a/boot/bootutil/include/bootutil/enc_key_public.h b/boot/bootutil/include/bootutil/enc_key_public.h
index 634f842..420d0d6 100644
--- a/boot/bootutil/include/bootutil/enc_key_public.h
+++ b/boot/bootutil/include/bootutil/enc_key_public.h
@@ -2,7 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  *
  * Copyright (c) 2018-2019 JUUL Labs
- * Copyright (c) 2019 Arm Limited
+ * Copyright (c) 2019-2021 Arm Limited
  * Copyright (c) 2021 Nordic Semiconductor ASA
  *
  * Original license:
@@ -32,12 +32,16 @@
 extern "C" {
 #endif
 
+#ifdef MCUBOOT_AES_256
+#define BOOT_ENC_KEY_SIZE       32
+#else
 #define BOOT_ENC_KEY_SIZE       16
+#endif
 
 #define TLV_ENC_RSA_SZ    256
-#define TLV_ENC_KW_SZ     24
-#define TLV_ENC_EC256_SZ  (65 + 32 + 16)
-#define TLV_ENC_X25519_SZ (32 + 32 + 16)
+#define TLV_ENC_KW_SZ     BOOT_ENC_KEY_SIZE + 8
+#define TLV_ENC_EC256_SZ  (65 + 32 + BOOT_ENC_KEY_SIZE)
+#define TLV_ENC_X25519_SZ (32 + 32 + BOOT_ENC_KEY_SIZE)
 
 #if defined(MCUBOOT_ENCRYPT_RSA)
 #define BOOT_ENC_TLV_SIZE TLV_ENC_RSA_SZ
@@ -53,4 +57,4 @@
 }
 #endif
 
-#endif /* BOOTUTIL_ENC_KEY_PUBLIC_H */
\ No newline at end of file
+#endif /* BOOTUTIL_ENC_KEY_PUBLIC_H */
diff --git a/boot/bootutil/include/bootutil/image.h b/boot/bootutil/include/bootutil/image.h
index 38304f0..fc94915 100644
--- a/boot/bootutil/include/bootutil/image.h
+++ b/boot/bootutil/include/bootutil/image.h
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2016-2019 Linaro LTD
  * Copyright (c) 2016-2019 JUUL Labs
- * Copyright (c) 2019-2020 Arm Limited
+ * Copyright (c) 2019-2021 Arm Limited
  *
  * Original license:
  *
@@ -50,7 +50,8 @@
  * Image header flags.
  */
 #define IMAGE_F_PIC                      0x00000001 /* Not supported. */
-#define IMAGE_F_ENCRYPTED                0x00000004 /* Encrypted. */
+#define IMAGE_F_ENCRYPTED_AES128         0x00000004 /* Encrypted using AES128. */
+#define IMAGE_F_ENCRYPTED_AES256         0x00000008 /* Encrypted using AES256. */
 #define IMAGE_F_NON_BOOTABLE             0x00000010 /* Split image app. */
 /*
  * Indicates that this image should be loaded into RAM instead of run
@@ -89,7 +90,7 @@
 #define IMAGE_TLV_RSA3072_PSS       0x23   /* RSA3072 of hash output */
 #define IMAGE_TLV_ED25519           0x24   /* ed25519 of hash output */
 #define IMAGE_TLV_ENC_RSA2048       0x30   /* Key encrypted with RSA-OAEP-2048 */
-#define IMAGE_TLV_ENC_KW128         0x31   /* Key encrypted with AES-KW-128 */
+#define IMAGE_TLV_ENC_KW            0x31   /* Key encrypted with AES-KW 128 or 256*/
 #define IMAGE_TLV_ENC_EC256         0x32   /* Key encrypted with ECIES-EC256 */
 #define IMAGE_TLV_ENC_X25519        0x33   /* Key encrypted with ECIES-X25519 */
 #define IMAGE_TLV_DEPENDENCY        0x40   /* Image depends on other image */
@@ -148,7 +149,8 @@
     uint16_t it_len;    /* Data length (not including TLV header). */
 };
 
-#define IS_ENCRYPTED(hdr) ((hdr)->ih_flags & IMAGE_F_ENCRYPTED)
+#define IS_ENCRYPTED(hdr) (((hdr)->ih_flags && IMAGE_F_ENCRYPTED_AES128) \
+                        || ((hdr)->ih_flags && IMAGE_F_ENCRYPTED_AES256))
 #define MUST_DECRYPT(fap, idx, hdr) \
     ((fap)->fa_id == FLASH_AREA_IMAGE_SECONDARY(idx) && IS_ENCRYPTED(hdr))
 
diff --git a/boot/bootutil/src/caps.c b/boot/bootutil/src/caps.c
index a21bc09..87bac1a 100644
--- a/boot/bootutil/src/caps.c
+++ b/boot/bootutil/src/caps.c
@@ -2,6 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  *
  * Copyright (c) 2017 Linaro Limited
+ * Copyright (c) 2021 Arm Limited
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -68,6 +69,9 @@
 #if defined(MCUBOOT_BOOTSTRAP)
     res |= BOOTUTIL_CAP_BOOTSTRAP;
 #endif
+#if defined(MCUBOOT_AES_256)
+    res |= BOOTUTIL_CAP_AES256;
+#endif
 
     return res;
 }
diff --git a/boot/bootutil/src/encrypted.c b/boot/bootutil/src/encrypted.c
index d14855b..278f51d 100644
--- a/boot/bootutil/src/encrypted.c
+++ b/boot/bootutil/src/encrypted.c
@@ -2,7 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  *
  * Copyright (c) 2018-2019 JUUL Labs
- * Copyright (c) 2019 Arm Limited
+ * Copyright (c) 2019-2021 Arm Limited
  */
 
 #include "mcuboot_config/mcuboot_config.h"
@@ -434,20 +434,20 @@
 #if defined(MCUBOOT_ENCRYPT_RSA)
 #    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_RSA2048
 #elif defined(MCUBOOT_ENCRYPT_KW)
-#    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_KW128
+#    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_KW
 #elif defined(MCUBOOT_ENCRYPT_EC256)
 #    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_EC256
 #    define EC_PUBK_INDEX       (0)
 #    define EC_TAG_INDEX        (65)
 #    define EC_CIPHERKEY_INDEX  (65 + 32)
-_Static_assert(EC_CIPHERKEY_INDEX + 16 == EXPECTED_ENC_LEN,
+_Static_assert(EC_CIPHERKEY_INDEX + BOOT_ENC_KEY_SIZE == EXPECTED_ENC_LEN,
         "Please fix ECIES-P256 component indexes");
 #elif defined(MCUBOOT_ENCRYPT_X25519)
 #    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_X25519
 #    define EC_PUBK_INDEX       (0)
 #    define EC_TAG_INDEX        (32)
 #    define EC_CIPHERKEY_INDEX  (32 + 32)
-_Static_assert(EC_CIPHERKEY_INDEX + 16 == EXPECTED_ENC_LEN,
+_Static_assert(EC_CIPHERKEY_INDEX + BOOT_ENC_KEY_SIZE == EXPECTED_ENC_LEN,
         "Please fix ECIES-X25519 component indexes");
 #endif
 
@@ -455,7 +455,7 @@
  * Decrypt an encryption key TLV.
  *
  * @param buf An encryption TLV read from flash (build time fixed length)
- * @param enckey An AES-128 key sized buffer to store to plain key.
+ * @param enckey An AES-128 or AES-256 key sized buffer to store to plain key.
  */
 int
 boot_enc_decrypt(const uint8_t *buf, uint8_t *enckey)
@@ -507,7 +507,7 @@
 
 #if defined(MCUBOOT_ENCRYPT_KW)
 
-    assert(*bootutil_enc_key.len == 16);
+    assert(*bootutil_enc_key.len == BOOT_ENC_KEY_SIZE);
     rc = key_unwrap(buf, enckey);
 
 #endif /* defined(MCUBOOT_ENCRYPT_KW) */
@@ -586,13 +586,13 @@
 
     bootutil_hmac_sha256_init(&hmac);
 
-    rc = bootutil_hmac_sha256_set_key(&hmac, &derived_key[16], 32);
+    rc = bootutil_hmac_sha256_set_key(&hmac, &derived_key[BOOT_ENC_KEY_SIZE], 32);
     if (rc != 0) {
         (void)bootutil_hmac_sha256_drop(&hmac);
         return -1;
     }
 
-    rc = bootutil_hmac_sha256_update(&hmac, &buf[EC_CIPHERKEY_INDEX], 16);
+    rc = bootutil_hmac_sha256_update(&hmac, &buf[EC_CIPHERKEY_INDEX], BOOT_ENC_KEY_SIZE);
     if (rc != 0) {
         (void)bootutil_hmac_sha256_drop(&hmac);
         return -1;