FWU: Optimize the RAM cost in FWU partition

- Make PSA_FWU_MAX_BLOCK_SIZE and internal buf for block write
configurable.
- Decrease the stack size of FWU partition to a reasonable
value.

Signed-off-by: Sherry Zhang <sherry.zhang2@arm.com>
Change-Id: I8aa115d7824bf887b4fe0814c4bdb1dfd3407ce7
diff --git a/cmake/install.cmake b/cmake/install.cmake
index f7fa5bf..cd5e77c 100644
--- a/cmake/install.cmake
+++ b/cmake/install.cmake
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2020-2021, Arm Limited. All rights reserved.
+# Copyright (c) 2020-2022, Arm Limited. All rights reserved.
 # Copyright (c) 2020, Cypress Semiconductor Corporation. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
@@ -106,6 +106,7 @@
 
 if(TFM_PARTITION_FIRMWARE_UPDATE)
     install(FILES       ${INTERFACE_INC_DIR}/psa/update.h
+                        ${CMAKE_BINARY_DIR}/generated/interface/include/psa/fwu_config.h
             DESTINATION ${INSTALL_INTERFACE_INC_DIR}/psa)
     install(FILES       ${INTERFACE_INC_DIR}/tfm_fwu_defs.h
             DESTINATION ${INSTALL_INTERFACE_INC_DIR})
diff --git a/config/config_default.cmake b/config/config_default.cmake
index a094957..aa3512a 100755
--- a/config/config_default.cmake
+++ b/config/config_default.cmake
@@ -140,6 +140,8 @@
 set(FORWARD_PROT_MSG                    OFF         CACHE BOOL      "Whether to forward all PSA RoT messages to a Secure Enclave")
 set(TFM_PARTITION_FIRMWARE_UPDATE       OFF         CACHE BOOL      "Enable firmware update partition")
 set(TFM_FWU_BOOTLOADER_LIB              "mcuboot"   CACHE STRING    "Bootloader configure file for Firmware Update partition")
+set(PSA_FWU_MAX_BLOCK_SIZE              1024        CACHE STRING    "The maximum permitted size for block in psa_fwu_write, in bytes.")
+set(TFM_FWU_BUF_SIZE                    ""          CACHE STRING    "Size of the FWU internal data transfer buffer (defaults to PSA_FWU_MAX_BLOCK_SIZE if not set)")
 
 ################################## Dependencies ################################
 
diff --git a/interface/include/psa/fwu_config.h.in b/interface/include/psa/fwu_config.h.in
new file mode 100644
index 0000000..e7b9b51
--- /dev/null
+++ b/interface/include/psa/fwu_config.h.in
@@ -0,0 +1,13 @@
+/*
+ * Copyright (c) 2022, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#ifndef __FWU_CONFIG_H__
+#define __FWU_CONFIG_H__
+
+#cmakedefine PSA_FWU_MAX_BLOCK_SIZE @PSA_FWU_MAX_BLOCK_SIZE@
+
+#endif /* __FWU_CONFIG_H__ */
diff --git a/interface/include/psa/update.h b/interface/include/psa/update.h
index d7de5ac..5af0f1c 100644
--- a/interface/include/psa/update.h
+++ b/interface/include/psa/update.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -12,13 +12,13 @@
 #include <stdint.h>
 
 #include "psa/error.h"
+#include "psa/fwu_config.h"
 #include "tfm_fwu_defs.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-#define PSA_FWU_MAX_BLOCK_SIZE               1024
 #define TFM_FWU_INVALID_IMAGE_ID             0
 
 /* The maximum size of an image digest in bytes. This is dependent
diff --git a/secure_fw/partitions/firmware_update/CMakeLists.txt b/secure_fw/partitions/firmware_update/CMakeLists.txt
index bfce009..09d2ba7 100644
--- a/secure_fw/partitions/firmware_update/CMakeLists.txt
+++ b/secure_fw/partitions/firmware_update/CMakeLists.txt
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2021, Arm Limited. All rights reserved.
+# Copyright (c) 2021-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -57,6 +57,7 @@
     PRIVATE
         $<$<BOOL:${TFM_PSA_API}>:TFM_PSA_API>
         $<$<BOOL:${DEFAULT_MCUBOOT_FLASH_MAP}>:DEFAULT_MCUBOOT_FLASH_MAP>
+        $<$<BOOL:${TFM_FWU_BUF_SIZE}>:TFM_FWU_BUF_SIZE=${TFM_FWU_BUF_SIZE}>
 )
 
 ############################ Secure API ########################################
@@ -87,3 +88,7 @@
     INTERFACE
     TFM_PARTITION_FIRMWARE_UPDATE
 )
+# Generate firmware update feature
+configure_file(${CMAKE_SOURCE_DIR}/interface/include/psa/fwu_config.h.in
+               ${CMAKE_BINARY_DIR}/generated/interface/include/psa/fwu_config.h
+               @ONLY)
diff --git a/secure_fw/partitions/firmware_update/tfm_firmware_update.yaml b/secure_fw/partitions/firmware_update/tfm_firmware_update.yaml
index e2db153..bea6fcf 100644
--- a/secure_fw/partitions/firmware_update/tfm_firmware_update.yaml
+++ b/secure_fw/partitions/firmware_update/tfm_firmware_update.yaml
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2021, Arm Limited. All rights reserved.
+# Copyright (c) 2021-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -11,7 +11,7 @@
   "type": "PSA-ROT",
   "priority": "NORMAL",
   "entry_point": "tfm_fwu_init",
-  "stack_size": "0x2000",
+  "stack_size": "0x480",
   "secure_functions": [
     {
       "name": "TFM_FWU_WRITE",
diff --git a/secure_fw/partitions/firmware_update/tfm_fwu_req_mngr.c b/secure_fw/partitions/firmware_update/tfm_fwu_req_mngr.c
index a379481..7ea10c1 100644
--- a/secure_fw/partitions/firmware_update/tfm_fwu_req_mngr.c
+++ b/secure_fw/partitions/firmware_update/tfm_fwu_req_mngr.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -30,6 +30,12 @@
  */
 static tfm_fwu_ctx_t fwu_ctx[TFM_FWU_MAX_IMAGES];
 
+#ifdef TFM_PSA_API
+#ifndef TFM_FWU_BUF_SIZE
+#define TFM_FWU_BUF_SIZE PSA_FWU_MAX_BLOCK_SIZE
+#endif
+static uint8_t data_block[TFM_FWU_BUF_SIZE];
+#endif
 /**
  * \brief Check if the image is in FWU process, return the index if it is.
  */
@@ -277,13 +283,12 @@
 {
     psa_image_id_t image_id;
     size_t block_offset;
-    size_t data_length, num;
+    size_t data_length, write_size, num;
     psa_status_t status = PSA_SUCCESS;
-    uint8_t data_block[PSA_FWU_MAX_BLOCK_SIZE];
     uint8_t image_index;
 
     /* Check input parameters. */
-    if (msg.in_size[2] > sizeof(data_block)) {
+    if (msg.in_size[2] > PSA_FWU_MAX_BLOCK_SIZE) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
@@ -301,14 +306,6 @@
     if (num != sizeof(block_offset)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
-
-    tfm_memset(data_block, 0, sizeof(data_block));
-    data_length = msg.in_size[2];
-    num = psa_read(msg.handle, 2, data_block, data_length);
-    if (num != data_length) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
     if (get_image_index(image_id, &image_index)) {
         /* The image is in FWU process. */
         if ((fwu_ctx[image_index].image_state != PSA_IMAGE_CANDIDATE) &&
@@ -332,10 +329,27 @@
         }
     }
 
-    return tfm_internal_fwu_write(image_id,
-                                  block_offset,
-                                  data_block,
-                                  data_length);
+    tfm_memset(data_block, 0, sizeof(data_block));
+    data_length = msg.in_size[2];
+    while (data_length > 0) {
+        write_size = sizeof(data_block) <= data_length ?
+                     sizeof(data_block) : data_length;
+        num = psa_read(msg.handle, 2, data_block, write_size);
+        if (num != write_size) {
+            return PSA_ERROR_PROGRAMMER_ERROR;
+        }
+
+        status = tfm_internal_fwu_write(image_id,
+                                        block_offset,
+                                        data_block,
+                                        write_size);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+        data_length -= write_size;
+        block_offset += write_size;
+    }
+    return PSA_SUCCESS;
 }
 
 static psa_status_t tfm_fwu_install_ipc(void)