Partitions: Convert PSA Partitions to SFN Model

This patch converts the following Secure Partitions to SFN Model:
- Crypto
- ITS
- PS
- Initial Attestation

Change-Id: Iecd1ed578681c78c5231064a7278b4fc967dca23
Signed-off-by: Kevin Peng <kevin.peng@arm.com>
diff --git a/secure_fw/partitions/crypto/crypto_init.c b/secure_fw/partitions/crypto/crypto_init.c
index 5c9965c..0fb8a69 100644
--- a/secure_fw/partitions/crypto/crypto_init.c
+++ b/secure_fw/partitions/crypto/crypto_init.c
@@ -111,14 +111,13 @@
     scratch.alloc_index = 0;
 }
 
-static psa_status_t tfm_crypto_call_srv(psa_msg_t *msg,
-                                        struct tfm_crypto_pack_iovec *iov,
-                                        const uint32_t srv_id)
+static psa_status_t tfm_crypto_call_srv(const psa_msg_t *msg)
 {
     psa_status_t status = PSA_SUCCESS;
     size_t in_len = PSA_MAX_IOVEC, out_len = PSA_MAX_IOVEC, i;
     psa_invec in_vec[PSA_MAX_IOVEC] = { {NULL, 0} };
     psa_outvec out_vec[PSA_MAX_IOVEC] = { {NULL, 0} };
+    struct tfm_crypto_pack_iovec iov = {0};
     void *alloc_buf_ptr = NULL;
 
     /* Check the number of in_vec filled */
@@ -130,8 +129,17 @@
     if (in_len < 1) {
         return PSA_ERROR_GENERIC_ERROR;
     }
+
+    if (psa_read(msg->handle, 0, &iov, sizeof(iov)) != sizeof(iov)) {
+        return PSA_ERROR_GENERIC_ERROR;
+    }
+
+    if (iov.srv_id >= TFM_CRYPTO_SID_MAX) {
+        return PSA_ERROR_GENERIC_ERROR;
+    }
+
     /* Initialise the first iovec with the IOV read when parsing */
-    in_vec[0].base = iov;
+    in_vec[0].base = &iov;
     in_vec[0].len = sizeof(struct tfm_crypto_pack_iovec);
 
     /* Alloc/read from the second element as the first is read when parsing */
@@ -143,10 +151,10 @@
             return status;
         }
         /* Read from the IPC framework inputs into the scratch */
-        (void) psa_read(msg->handle, i, alloc_buf_ptr, msg->in_size[i]);
+        in_vec[i].len =
+                       psa_read(msg->handle, i, alloc_buf_ptr, msg->in_size[i]);
         /* Populate the fields of the input to the secure function */
         in_vec[i].base = alloc_buf_ptr;
-        in_vec[i].len = msg->in_size[i];
     }
 
     /* Check the number of out_vec filled */
@@ -170,7 +178,7 @@
     (void)tfm_crypto_set_scratch_owner(msg->client_id);
 
     /* Call the uniform signature API */
-    status = sfid_func_table[srv_id](in_vec, in_len, out_vec, out_len);
+    status = sfid_func_table[iov.srv_id](in_vec, in_len, out_vec, out_len);
 
     /* Write into the IPC framework outputs from the scratch */
     for (i = 0; i < out_len; i++) {
@@ -183,70 +191,6 @@
     return status;
 }
 
-static psa_status_t tfm_crypto_parse_msg(psa_msg_t *msg,
-                                         struct tfm_crypto_pack_iovec *iov,
-                                         uint32_t *srv_id_p)
-{
-    size_t read_size;
-
-    /* Read the in_vec[0] which holds the IOVEC always */
-    read_size = psa_read(msg->handle,
-                         0,
-                         iov,
-                         sizeof(struct tfm_crypto_pack_iovec));
-
-    if (read_size != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_GENERIC_ERROR;
-    }
-
-    if (iov->srv_id >= TFM_CRYPTO_SID_MAX) {
-        *srv_id_p = TFM_CRYPTO_SID_INVALID;
-        return PSA_ERROR_GENERIC_ERROR;
-    }
-
-    *srv_id_p = iov->srv_id;
-
-    return PSA_SUCCESS;
-}
-
-static void tfm_crypto_ipc_handler(void)
-{
-    psa_signal_t signals;
-    psa_msg_t msg;
-    psa_status_t status = PSA_SUCCESS;
-    uint32_t srv_id = TFM_CRYPTO_SID_INVALID;
-    struct tfm_crypto_pack_iovec iov = {0};
-
-    while (1) {
-        signals = psa_wait(PSA_WAIT_ANY, PSA_BLOCK);
-        if (signals & TFM_CRYPTO_SIGNAL) {
-            /* Extract the message */
-            if (psa_get(TFM_CRYPTO_SIGNAL, &msg) != PSA_SUCCESS) {
-                psa_panic();
-            }
-
-            /* Process the message type */
-            switch (msg.type) {
-            case PSA_IPC_CALL:
-                /* Parse the message */
-                status = tfm_crypto_parse_msg(&msg, &iov, &srv_id);
-                /* Call the dispatcher based on the SID passed as type */
-                if (status == PSA_SUCCESS) {
-                    status = tfm_crypto_call_srv(&msg, &iov, srv_id);
-                }
-                psa_reply(msg.handle, status);
-                break;
-            default:
-                psa_panic();
-            }
-        } else {
-            psa_panic();
-        }
-    }
-
-    /* NOTREACHED */
-    return;
-}
 #endif /* TFM_PSA_API */
 
 /**
@@ -333,15 +277,20 @@
     }
 
     /* Initialise the engine layer */
-    status = tfm_crypto_engine_init();
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
+    return tfm_crypto_engine_init();
+}
 
 #ifdef TFM_PSA_API
-    /* Should not return in normal operations */
-    tfm_crypto_ipc_handler();
-#endif
+psa_status_t tfm_crypto_sfn(const psa_msg_t *msg)
+{
+    /* Process the message type */
+    switch (msg->type) {
+    case PSA_IPC_CALL:
+        return tfm_crypto_call_srv(msg);
+    default:
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
 
-    return status;
+    return PSA_ERROR_GENERIC_ERROR;
 }
+#endif
diff --git a/secure_fw/partitions/crypto/tfm_crypto.yaml b/secure_fw/partitions/crypto/tfm_crypto.yaml
index 980849a..54cd06a 100644
--- a/secure_fw/partitions/crypto/tfm_crypto.yaml
+++ b/secure_fw/partitions/crypto/tfm_crypto.yaml
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2018-2021, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -10,8 +10,8 @@
   "name": "TFM_SP_CRYPTO",
   "type": "PSA-ROT",
   "priority": "NORMAL",
-  "model": "IPC",
-  "entry_point": "tfm_crypto_init",
+  "model": "SFN",
+  "entry_init": "tfm_crypto_init",
   "stack_size": "0x2000",
   "secure_functions": [
     {
diff --git a/secure_fw/partitions/initial_attestation/tfm_attest_req_mngr.c b/secure_fw/partitions/initial_attestation/tfm_attest_req_mngr.c
index d33db15..842c0f4 100644
--- a/secure_fw/partitions/initial_attestation/tfm_attest_req_mngr.c
+++ b/secure_fw/partitions/initial_attestation/tfm_attest_req_mngr.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -86,7 +86,7 @@
 
     bytes_read = psa_read(msg->handle, 0,
                           &challenge_size, msg->in_size[0]);
-    if (bytes_read != msg->in_size[0]) {
+    if (bytes_read != sizeof(challenge_size)) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
@@ -99,56 +99,22 @@
     return status;
 }
 
-/*
- * Fixme: Temporarily implement abort as infinite loop,
- * will replace it later.
- */
-static void tfm_abort(void)
+psa_status_t tfm_attestation_service_sfn(const psa_msg_t *msg)
 {
-    while (1)
-        ;
-}
-
-static void attest_signal_handle(psa_signal_t signal)
-{
-    psa_msg_t msg;
-    psa_status_t status;
-
-    status = psa_get(signal, &msg);
-    switch (msg.type) {
+    switch (msg->type) {
     case TFM_ATTEST_GET_TOKEN:
-        status = psa_attest_get_token(&msg);
-        psa_reply(msg.handle, status);
-        break;
+        return psa_attest_get_token(msg);
     case TFM_ATTEST_GET_TOKEN_SIZE:
-        status = psa_attest_get_token_size(&msg);
-        psa_reply(msg.handle, status);
-        break;
+        return psa_attest_get_token_size(msg);
     default:
-        tfm_abort();
+        return PSA_ERROR_NOT_SUPPORTED;
     }
+
+    return PSA_ERROR_GENERIC_ERROR;
 }
-#endif
+#endif /* TFM_PSA_API */
 
 psa_status_t attest_partition_init(void)
 {
-    psa_status_t err = attest_init();
-#ifdef TFM_PSA_API
-    psa_signal_t signals;
-
-    if (err != PSA_SUCCESS) {
-        tfm_abort();
-    }
-
-    while (1) {
-        signals = psa_wait(PSA_WAIT_ANY, PSA_BLOCK);
-        if (signals & TFM_ATTESTATION_SERVICE_SIGNAL) {
-            attest_signal_handle(TFM_ATTESTATION_SERVICE_SIGNAL);
-        } else {
-            tfm_abort();
-        }
-    }
-#else
-    return err;
-#endif
+    return attest_init();
 }
diff --git a/secure_fw/partitions/initial_attestation/tfm_initial_attestation.yaml b/secure_fw/partitions/initial_attestation/tfm_initial_attestation.yaml
index 9bb716d..b3f16b5 100644
--- a/secure_fw/partitions/initial_attestation/tfm_initial_attestation.yaml
+++ b/secure_fw/partitions/initial_attestation/tfm_initial_attestation.yaml
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2018-2021, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -10,8 +10,8 @@
   "name": "TFM_SP_INITIAL_ATTESTATION",
   "type": "PSA-ROT",
   "priority": "NORMAL",
-  "model": "IPC",
-  "entry_point": "attest_partition_init",
+  "model": "SFN",
+  "entry_init": "attest_partition_init",
   "stack_size": "0x0A80",
   "secure_functions": [
     {
diff --git a/secure_fw/partitions/internal_trusted_storage/tfm_internal_trusted_storage.yaml b/secure_fw/partitions/internal_trusted_storage/tfm_internal_trusted_storage.yaml
index 5907370..70aa6ca 100644
--- a/secure_fw/partitions/internal_trusted_storage/tfm_internal_trusted_storage.yaml
+++ b/secure_fw/partitions/internal_trusted_storage/tfm_internal_trusted_storage.yaml
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2019-2021, Arm Limited. All rights reserved.
+# Copyright (c) 2019-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -10,8 +10,8 @@
   "name": "TFM_SP_ITS",
   "type": "PSA-ROT",
   "priority": "NORMAL",
-  "model": "IPC",
-  "entry_point": "tfm_its_req_mngr_init",
+  "model": "SFN",
+  "entry_init": "tfm_its_entry",
   "stack_size": "0x680",
   "secure_functions": [
     {
diff --git a/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.c b/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.c
index db6d512..88300e0 100644
--- a/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.c
+++ b/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -180,38 +180,37 @@
 }
 
 #else /* !defined(TFM_PSA_API) */
-typedef psa_status_t (*its_func_t)(void);
-static psa_msg_t msg;
+static const psa_msg_t *p_msg;
 
-static psa_status_t tfm_its_set_ipc(void)
+static psa_status_t tfm_its_set_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
     size_t data_length;
     psa_storage_create_flags_t create_flags;
     size_t num;
 
-    if (msg.in_size[0] != sizeof(uid) ||
-        msg.in_size[2] != sizeof(create_flags)) {
+    if (msg->in_size[0] != sizeof(uid) ||
+        msg->in_size[2] != sizeof(create_flags)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    data_length = msg.in_size[1];
+    data_length = msg->in_size[1];
 
-    num = psa_read(msg.handle, 0, &uid, sizeof(uid));
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
     if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 2, &create_flags, sizeof(create_flags));
+    num = psa_read(msg->handle, 2, &create_flags, sizeof(create_flags));
     if (num != sizeof(create_flags)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    return tfm_its_set(msg.client_id, uid, data_length, create_flags);
+    return tfm_its_set(msg->client_id, uid, data_length, create_flags);
 }
 
-static psa_status_t tfm_its_get_ipc(void)
+static psa_status_t tfm_its_get_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
     size_t data_offset;
@@ -219,121 +218,78 @@
     size_t data_length;
     size_t num;
 
-    if (msg.in_size[0] != sizeof(uid) ||
-        msg.in_size[1] != sizeof(data_offset)) {
+    if (msg->in_size[0] != sizeof(uid) ||
+        msg->in_size[1] != sizeof(data_offset)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    data_size = msg.out_size[0];
+    data_size = msg->out_size[0];
 
-    num = psa_read(msg.handle, 0, &uid, sizeof(uid));
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
     if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 1, &data_offset, sizeof(data_offset));
+    num = psa_read(msg->handle, 1, &data_offset, sizeof(data_offset));
     if (num != sizeof(data_offset)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    return tfm_its_get(msg.client_id, uid, data_offset, data_size,
+    return tfm_its_get(msg->client_id, uid, data_offset, data_size,
                        &data_length);
 }
 
-static psa_status_t tfm_its_get_info_ipc(void)
+static psa_status_t tfm_its_get_info_req(const psa_msg_t *msg)
 {
     psa_status_t status;
     psa_storage_uid_t uid;
     struct psa_storage_info_t info;
     size_t num;
 
-    if (msg.in_size[0] != sizeof(uid) ||
-        msg.out_size[0] != sizeof(info)) {
+    if (msg->in_size[0] != sizeof(uid) ||
+        msg->out_size[0] != sizeof(info)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 0, &uid, sizeof(uid));
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
     if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    status = tfm_its_get_info(msg.client_id, uid, &info);
+    status = tfm_its_get_info(msg->client_id, uid, &info);
     if (status == PSA_SUCCESS) {
-        psa_write(msg.handle, 0, &info, sizeof(info));
+        psa_write(msg->handle, 0, &info, sizeof(info));
     }
 
     return status;
 }
 
-static psa_status_t tfm_its_remove_ipc(void)
+static psa_status_t tfm_its_remove_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
     size_t num;
 
-    if (msg.in_size[0] != sizeof(uid)) {
+    if (msg->in_size[0] != sizeof(uid)) {
         /* The input argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 0, &uid, sizeof(uid));
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
     if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    return tfm_its_remove(msg.client_id, uid);
+    return tfm_its_remove(msg->client_id, uid);
 }
 
-static void its_signal_handle(psa_signal_t signal)
-{
-    psa_status_t status;
-
-    status = psa_get(signal, &msg);
-    if (status != PSA_SUCCESS) {
-        return;
-    }
-
-    switch (msg.type) {
-    case TFM_ITS_SET:
-        status = tfm_its_set_ipc();
-        psa_reply(msg.handle, status);
-        break;
-    case TFM_ITS_GET:
-        status = tfm_its_get_ipc();
-        psa_reply(msg.handle, status);
-        break;
-    case TFM_ITS_GET_INFO:
-        status = tfm_its_get_info_ipc();
-        psa_reply(msg.handle, status);
-        break;
-    case TFM_ITS_REMOVE:
-        status = tfm_its_remove_ipc();
-        psa_reply(msg.handle, status);
-        break;
-    default:
-        psa_panic();
-    }
-}
 #endif /* !defined(TFM_PSA_API) */
 
-psa_status_t tfm_its_req_mngr_init(void)
+psa_status_t tfm_its_entry(void)
 {
 #ifdef TFM_PSA_API
-    psa_signal_t signals;
-
-    if (tfm_its_init() != PSA_SUCCESS) {
-        psa_panic();
-    }
-
-    while (1) {
-        signals = psa_wait(PSA_WAIT_ANY, PSA_BLOCK);
-        if (signals & TFM_INTERNAL_TRUSTED_STORAGE_SERVICE_SIGNAL) {
-            its_signal_handle(TFM_INTERNAL_TRUSTED_STORAGE_SERVICE_SIGNAL);
-        } else {
-            psa_panic();
-        }
-    }
+    return tfm_its_init();
 #else
     if (tfm_its_init() != PSA_SUCCESS) {
         return PSA_ERROR_GENERIC_ERROR;
@@ -343,10 +299,32 @@
 #endif
 }
 
+#ifdef TFM_PSA_API
+psa_status_t tfm_internal_trusted_storage_service_sfn(const psa_msg_t *msg)
+{
+    p_msg = msg;
+
+    switch (msg->type) {
+    case TFM_ITS_SET:
+        return tfm_its_set_req(msg);
+    case TFM_ITS_GET:
+        return tfm_its_get_req(msg);
+    case TFM_ITS_GET_INFO:
+        return tfm_its_get_info_req(msg);
+    case TFM_ITS_REMOVE:
+        return tfm_its_remove_req(msg);
+    default:
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    return PSA_ERROR_GENERIC_ERROR;
+}
+#endif /* TFM_PSA_API */
+
 size_t its_req_mngr_read(uint8_t *buf, size_t num_bytes)
 {
 #ifdef TFM_PSA_API
-    return psa_read(msg.handle, 1, buf, num_bytes);
+    return psa_read(p_msg->handle, 1, buf, num_bytes);
 #else
     (void)tfm_memcpy(buf, p_data, num_bytes);
     p_data += num_bytes;
@@ -357,7 +335,7 @@
 void its_req_mngr_write(const uint8_t *buf, size_t num_bytes)
 {
 #ifdef TFM_PSA_API
-    psa_write(msg.handle, 0, buf, num_bytes);
+    psa_write(p_msg->handle, 0, buf, num_bytes);
 #else
     (void)tfm_memcpy(p_data, buf, num_bytes);
     p_data += num_bytes;
diff --git a/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.h b/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.h
index 3389638..050eab7 100644
--- a/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.h
+++ b/secure_fw/partitions/internal_trusted_storage/tfm_its_req_mngr.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -16,6 +16,7 @@
 extern "C" {
 #endif
 
+#ifndef TFM_PSA_API
 /**
  * \brief Handles the set request.
  *
@@ -79,6 +80,7 @@
  */
 psa_status_t tfm_its_remove_req(psa_invec *in_vec, size_t in_len,
                                 psa_outvec *out_vec, size_t out_len);
+#endif /* NOT TFM_PSA_API */
 
 /**
  * \brief Reads asset data from the caller.
diff --git a/secure_fw/partitions/protected_storage/tfm_protected_storage.yaml b/secure_fw/partitions/protected_storage/tfm_protected_storage.yaml
index 8d71688..6e1cfef 100644
--- a/secure_fw/partitions/protected_storage/tfm_protected_storage.yaml
+++ b/secure_fw/partitions/protected_storage/tfm_protected_storage.yaml
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# Copyright (c) 2018-2021, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2022, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -10,8 +10,8 @@
   "name": "TFM_SP_PS",
   "type": "APPLICATION-ROT",
   "priority": "NORMAL",
-  "model": "IPC",
-  "entry_point": "tfm_ps_req_mngr_init",
+  "model": "SFN",
+  "entry_init": "tfm_ps_entry",
   "stack_size": "0x800",
   "secure_functions": [
     {
diff --git a/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.c b/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.c
index 1bd25da..845f5ae 100644
--- a/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.c
+++ b/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -70,7 +70,7 @@
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    if (in_vec[0].len != sizeof(psa_storage_uid_t)) {
+    if (in_vec[0].len != sizeof(uid)) {
         /* The input argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
@@ -80,7 +80,7 @@
     p_data = (void *)in_vec[1].base;
     data_length = in_vec[1].len;
 
-    if (in_vec[2].len != sizeof(psa_storage_create_flags_t)) {
+    if (in_vec[2].len != sizeof(create_flags)) {
         /* The input argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
@@ -115,7 +115,7 @@
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    if (in_vec[0].len != sizeof(psa_storage_uid_t)) {
+    if (in_vec[0].len != sizeof(uid)) {
         /* The input argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
@@ -160,14 +160,14 @@
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    if (in_vec[0].len != sizeof(psa_storage_uid_t)) {
+    if (in_vec[0].len != sizeof(uid)) {
         /* The input argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
     uid = *((psa_storage_uid_t *)in_vec[0].base);
 
-    if (out_vec[0].len != sizeof(struct psa_storage_info_t)) {
+    if (out_vec[0].len != sizeof(*p_info)) {
         /* The output argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
@@ -195,7 +195,7 @@
     }
 
     if ((in_len != 1) ||
-        (in_vec[0].len != sizeof(psa_storage_uid_t)) ||
+        (in_vec[0].len != sizeof(uid)) ||
         (out_len != 0)) {
         /* The number of arguments/output argument size are incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
@@ -241,65 +241,64 @@
 }
 
 #else /* !defined(TFM_PSA_API) */
-typedef psa_status_t (*ps_func_t)(void);
-static psa_msg_t msg;
+static const psa_msg_t *p_msg;
 
-static psa_status_t tfm_ps_set_ipc(void)
+static psa_status_t tfm_ps_set_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
     int32_t client_id;
     psa_storage_create_flags_t create_flags;
     size_t num = 0;
 
-    client_id = msg.client_id;
+    client_id = msg->client_id;
 
-    if (msg.in_size[0] != sizeof(psa_storage_uid_t) ||
-        msg.in_size[2] != sizeof(psa_storage_create_flags_t)) {
+    if (msg->in_size[0] != sizeof(uid) ||
+        msg->in_size[2] != sizeof(create_flags)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 0, &uid, msg.in_size[0]);
-    if (num != msg.in_size[0]) {
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
+    if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 2, &create_flags, msg.in_size[2]);
-    if (num != msg.in_size[2]) {
+    num = psa_read(msg->handle, 2, &create_flags, sizeof(create_flags));
+    if (num != sizeof(create_flags)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    return tfm_ps_set(client_id, uid, msg.in_size[1], create_flags);
+    return tfm_ps_set(client_id, uid, msg->in_size[1], create_flags);
 }
 
-static psa_status_t tfm_ps_get_ipc(void)
+static psa_status_t tfm_ps_get_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
     uint32_t data_offset;
     size_t num = 0;
     size_t p_data_length;
 
-    if (msg.in_size[0] != sizeof(psa_storage_uid_t) ||
-        msg.in_size[1] != sizeof(uint32_t)) {
+    if (msg->in_size[0] != sizeof(uid) ||
+        msg->in_size[1] != sizeof(data_offset)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 0, &uid,  msg.in_size[0]);
-    if (num != msg.in_size[0]) {
+    num = psa_read(msg->handle, 0, &uid,  sizeof(uid));
+    if (num != sizeof(psa_storage_uid_t)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 1, &data_offset, msg.in_size[1]);
-    if (num !=  msg.in_size[1]) {
+    num = psa_read(msg->handle, 1, &data_offset, sizeof(data_offset));
+    if (num != sizeof(data_offset)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    return tfm_ps_get(msg.client_id, uid, data_offset,  msg.out_size[0],
+    return tfm_ps_get(msg->client_id, uid, data_offset,  msg->out_size[0],
                       &p_data_length);
 }
 
-static psa_status_t tfm_ps_get_info_ipc(void)
+static psa_status_t tfm_ps_get_info_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
 
@@ -307,108 +306,86 @@
     size_t num = 0;
     psa_status_t status;
 
-    if (msg.in_size[0] != sizeof(psa_storage_uid_t) ||
-        msg.out_size[0] != sizeof(struct psa_storage_info_t)) {
+    if (msg->in_size[0] != sizeof(uid) ||
+        msg->out_size[0] != sizeof(info)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 0, &uid, msg.in_size[0]);
-    if (num != msg.in_size[0]) {
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
+    if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    status = tfm_ps_get_info(msg.client_id, uid, &info);
+    status = tfm_ps_get_info(msg->client_id, uid, &info);
 
     if (status == PSA_SUCCESS) {
-        psa_write(msg.handle, 0, &info, msg.out_size[0]);
+        psa_write(msg->handle, 0, &info, sizeof(info));
     }
     return status;
 }
 
-static psa_status_t tfm_ps_remove_ipc(void)
+static psa_status_t tfm_ps_remove_req(const psa_msg_t *msg)
 {
     psa_storage_uid_t uid;
     size_t num = 0;
 
-    if (msg.in_size[0] != sizeof(psa_storage_uid_t)) {
+    if (msg->in_size[0] != sizeof(uid)) {
         /* The size of one of the arguments is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    num = psa_read(msg.handle, 0, &uid, msg.in_size[0]);
-    if (num != msg.in_size[0]) {
+    num = psa_read(msg->handle, 0, &uid, sizeof(uid));
+    if (num != sizeof(uid)) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
-    return tfm_ps_remove(msg.client_id, uid);
+    return tfm_ps_remove(msg->client_id, uid);
 }
 
-static psa_status_t tfm_ps_get_support_ipc(void)
+static psa_status_t tfm_ps_get_support_req(const psa_msg_t *msg)
 {
     size_t out_size;
     uint32_t support_flags;
 
-    out_size = msg.out_size[0];
+    out_size = msg->out_size[0];
     if (out_size != sizeof(support_flags)) {
         /* The output argument size is incorrect */
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
 
     support_flags = tfm_ps_get_support();
-    psa_write(msg.handle, 0, &support_flags, out_size);
+    psa_write(msg->handle, 0, &support_flags, out_size);
     return PSA_SUCCESS;
 }
 
-static void ps_signal_handle(psa_signal_t signal)
+psa_status_t tfm_protected_storage_service_sfn(const psa_msg_t *msg)
 {
-    psa_status_t status;
+    p_msg = msg;
 
-    status = psa_get(signal, &msg);
-    switch (msg.type) {
+    switch (msg->type) {
     case TFM_PS_SET:
-        status = tfm_ps_set_ipc();
-        psa_reply(msg.handle, status);
-        break;
+        return tfm_ps_set_req(msg);
     case TFM_PS_GET:
-        status = tfm_ps_get_ipc();
-        psa_reply(msg.handle, status);
-        break;
+        return tfm_ps_get_req(msg);
     case TFM_PS_GET_INFO:
-        status = tfm_ps_get_info_ipc();
-        psa_reply(msg.handle, status);
-        break;
+        return tfm_ps_get_info_req(msg);
     case TFM_PS_REMOVE:
-        status = tfm_ps_remove_ipc();
-        psa_reply(msg.handle, status);
-        break;
+        return tfm_ps_remove_req(msg);
     case TFM_PS_GET_SUPPORT:
-        status = tfm_ps_get_support_ipc();
-        psa_reply(msg.handle, status);
-        break;
+        return tfm_ps_get_support_req(msg);
     default:
-        psa_panic();
+        return PSA_ERROR_PROGRAMMER_ERROR;
     }
+
+    return PSA_ERROR_GENERIC_ERROR;
 }
 #endif /* !defined(TFM_PSA_API) */
 
-psa_status_t tfm_ps_req_mngr_init(void)
+psa_status_t tfm_ps_entry(void)
 {
 #ifdef TFM_PSA_API
-    psa_signal_t signals;
-
-    if (tfm_ps_init() != PSA_SUCCESS) {
-        psa_panic();
-    }
-
-    while (1) {
-        signals = psa_wait(PSA_WAIT_ANY, PSA_BLOCK);
-        if (signals & TFM_PROTECTED_STORAGE_SERVICE_SIGNAL) {
-            ps_signal_handle(TFM_PROTECTED_STORAGE_SERVICE_SIGNAL);
-        } else {
-            psa_panic();
-        }
-    }
+    return tfm_ps_init();
 #else
     /* In library mode, initialisation is delayed until the first secure
      * function call, as calls to the Crypto service are required for
@@ -423,7 +400,7 @@
 #ifdef TFM_PSA_API
     size_t num = 0;
 
-    num = psa_read(msg.handle, 1, out_data, size);
+    num = psa_read(p_msg->handle, 1, out_data, size);
     if (num != size) {
         return PSA_ERROR_PROGRAMMER_ERROR;
     }
@@ -436,7 +413,7 @@
 void ps_req_mngr_write_asset_data(const uint8_t *in_data, uint32_t size)
 {
 #ifdef TFM_PSA_API
-    psa_write(msg.handle, 0, in_data, size);
+    psa_write(p_msg->handle, 0, in_data, size);
 #else /* TFM_PSA_API */
     (void)tfm_memcpy(p_data, in_data, size);
 #endif
diff --git a/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.h b/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.h
index 1dc3a4f..1c24d76 100644
--- a/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.h
+++ b/secure_fw/partitions/protected_storage/tfm_ps_req_mngr.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -17,6 +17,7 @@
 extern "C" {
 #endif
 
+#ifndef TFM_PSA_API
 /**
  * \brief Handles the set request.
  *
@@ -97,6 +98,7 @@
  */
 psa_status_t tfm_ps_get_support_req(psa_invec *in_vec, size_t in_len,
                                     psa_outvec *out_vec, size_t out_len);
+#endif /* NOT TFM_PSA_API */
 
 /**
  * \brief Takes an input buffer containing asset data and writes