Crypto: Refactor the API dispatcher interface to reduce code size

This patch restructures the way the underlying APIs that
implements the PSA Crypto APIs are interfaced to the TF-M
Crypto service through a thin shim layer. The size of this
layer is reduced by nearly 45% on the default configuration.
Also, it removes the check for parameter number and size on
each function call as that is a redundant check as per the
overall threat model of the interaction between the crypto
service and the partition manager.

Signed-off-by: Antonio de Angelis <antonio.deangelis@arm.com>
Change-Id: I07165bad00346cd12cf63620532f55da0c5d1262
diff --git a/interface/include/tfm_crypto_defs.h b/interface/include/tfm_crypto_defs.h
index c2f4d7c..1d26c06 100644
--- a/interface/include/tfm_crypto_defs.h
+++ b/interface/include/tfm_crypto_defs.h
@@ -20,8 +20,7 @@
 /**
  * \brief This type is used to overcome a limitation in the number of maximum
  *        IOVECs that can be used especially in psa_aead_encrypt and
- *        psa_aead_decrypt. To be removed in case the AEAD APIs number of
- *        parameters passed gets restructured
+ *        psa_aead_decrypt.
  */
 #define TFM_CRYPTO_MAX_NONCE_LENGTH (16u)
 struct tfm_crypto_aead_pack_input {
@@ -34,99 +33,236 @@
  *
  */
 struct tfm_crypto_pack_iovec {
-    uint32_t srv_id;             /*!< Crypto service ID used to dispatch the
-                                  *   request
-                                  */
-    uint16_t step;               /*!< Key derivation step */
-    psa_key_id_t key_id;         /*!< Key id */
-    psa_algorithm_t alg;         /*!< Algorithm */
-    uint32_t op_handle;          /*!< Frontend context handle associated to a
-                                  *   multipart operation
-                                  */
-    size_t capacity;             /*!< Key derivation capacity */
-    size_t ad_length;            /*!< Additional Data length for multipart AEAD */
-    size_t plaintext_length;     /*!< Plaintext length for multipart AEAD */
+    uint32_t function_id;    /*!< Used to identify the function in the
+                              *   API dispatcher to the service backend
+                              */
+    uint16_t step;           /*!< Key derivation step */
+    psa_key_id_t key_id;     /*!< Key id */
+    psa_algorithm_t alg;     /*!< Algorithm */
+    uint32_t op_handle;      /*!< Frontend context handle associated to a
+                              *   multipart operation
+                              */
+    size_t capacity;         /*!< Key derivation capacity */
+    size_t ad_length;        /*!< Additional Data length for multipart AEAD */
+    size_t plaintext_length; /*!< Plaintext length for multipart AEAD */
 
-    struct tfm_crypto_aead_pack_input aead_in; /*!< FixMe: Temporarily used for
-                                                *   AEAD until the API is
-                                                *   restructured
-                                                */
+    struct tfm_crypto_aead_pack_input aead_in; /*!< Packs AEAD-related inputs */
 };
 
 /**
- * \brief Define a progressive numerical value for each SID which can be used
- *        when dispatching the requests to the service. Note: This has to
- *        match exactly with the list of APIs defined in tfm_crypto_api.h by
- *        the LIST_TFM_CRYPTO_UNIFORM_SIGNATURE_API X macro.
+ * \brief X macro describing each of the available PSA Crypto APIs in terms of
+ *        group ID (\ref tfm_crypto_group_id) and multipart function type
+ *        (\ref tfm_crypto_function_type)
  */
-enum {
-    TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID = (0u),
-    TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
-    TFM_CRYPTO_OPEN_KEY_SID,
-    TFM_CRYPTO_CLOSE_KEY_SID,
-    TFM_CRYPTO_IMPORT_KEY_SID,
-    TFM_CRYPTO_DESTROY_KEY_SID,
-    TFM_CRYPTO_EXPORT_KEY_SID,
-    TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
-    TFM_CRYPTO_PURGE_KEY_SID,
-    TFM_CRYPTO_COPY_KEY_SID,
-    TFM_CRYPTO_HASH_COMPUTE_SID,
-    TFM_CRYPTO_HASH_COMPARE_SID,
-    TFM_CRYPTO_HASH_SETUP_SID,
-    TFM_CRYPTO_HASH_UPDATE_SID,
-    TFM_CRYPTO_HASH_FINISH_SID,
-    TFM_CRYPTO_HASH_VERIFY_SID,
-    TFM_CRYPTO_HASH_ABORT_SID,
-    TFM_CRYPTO_HASH_CLONE_SID,
-    TFM_CRYPTO_MAC_COMPUTE_SID,
-    TFM_CRYPTO_MAC_VERIFY_SID,
-    TFM_CRYPTO_MAC_SIGN_SETUP_SID,
-    TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
-    TFM_CRYPTO_MAC_UPDATE_SID,
-    TFM_CRYPTO_MAC_SIGN_FINISH_SID,
-    TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
-    TFM_CRYPTO_MAC_ABORT_SID,
-    TFM_CRYPTO_CIPHER_ENCRYPT_SID,
-    TFM_CRYPTO_CIPHER_DECRYPT_SID,
-    TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
-    TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
-    TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
-    TFM_CRYPTO_CIPHER_SET_IV_SID,
-    TFM_CRYPTO_CIPHER_UPDATE_SID,
-    TFM_CRYPTO_CIPHER_FINISH_SID,
-    TFM_CRYPTO_CIPHER_ABORT_SID,
-    TFM_CRYPTO_AEAD_ENCRYPT_SID,
-    TFM_CRYPTO_AEAD_DECRYPT_SID,
-    TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
-    TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
-    TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
-    TFM_CRYPTO_AEAD_SET_NONCE_SID,
-    TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
-    TFM_CRYPTO_AEAD_UPDATE_AD_SID,
-    TFM_CRYPTO_AEAD_UPDATE_SID,
-    TFM_CRYPTO_AEAD_FINISH_SID,
-    TFM_CRYPTO_AEAD_VERIFY_SID,
-    TFM_CRYPTO_AEAD_ABORT_SID,
-    TFM_CRYPTO_SIGN_MESSAGE_SID,
-    TFM_CRYPTO_VERIFY_MESSAGE_SID,
-    TFM_CRYPTO_SIGN_HASH_SID,
-    TFM_CRYPTO_VERIFY_HASH_SID,
-    TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
-    TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
-    TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
-    TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
-    TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
-    TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
-    TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
-    TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
-    TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
-    TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
-    TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
-    TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
-    TFM_CRYPTO_GENERATE_RANDOM_SID,
-    TFM_CRYPTO_GENERATE_KEY_SID,
-    TFM_CRYPTO_SID_MAX,
+#define TFM_CRYPTO_SERVICE_API_DESCRIPTION          \
+    X(TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,            \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,          \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_OPEN_KEY_SID,                      \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_CLOSE_KEY_SID,                     \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_IMPORT_KEY_SID,                    \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_DESTROY_KEY_SID,                   \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_EXPORT_KEY_SID,                    \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,             \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_PURGE_KEY_SID,                     \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_COPY_KEY_SID,                      \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_GENERATE_KEY_SID,                  \
+        TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_HASH_COMPUTE_SID,                  \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_HASH_COMPARE_SID,                  \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_HASH_SETUP_SID,                    \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_HASH_UPDATE_SID,                   \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_HASH_CLONE_SID,                    \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_HASH_FINISH_SID,                   \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_HASH_VERIFY_SID,                   \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_HASH_ABORT_SID,                    \
+        TFM_CRYPTO_GROUP_ID_HASH,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_MAC_COMPUTE_SID,                   \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_MAC_VERIFY_SID,                    \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_MAC_SIGN_SETUP_SID,                \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_MAC_VERIFY_SETUP_SID,              \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_MAC_UPDATE_SID,                    \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_MAC_SIGN_FINISH_SID,               \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_MAC_VERIFY_FINISH_SID,             \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_MAC_ABORT_SID,                     \
+        TFM_CRYPTO_GROUP_ID_MAC,                    \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_CIPHER_ENCRYPT_SID,                \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_CIPHER_DECRYPT_SID,                \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,          \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,          \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_CIPHER_GENERATE_IV_SID,            \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_CIPHER_SET_IV_SID,                 \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_CIPHER_UPDATE_SID,                 \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_CIPHER_FINISH_SID,                 \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_CIPHER_ABORT_SID,                  \
+        TFM_CRYPTO_GROUP_ID_CIPHER,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_ENCRYPT_SID,                  \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_AEAD_DECRYPT_SID,                  \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,            \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,            \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,           \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_SET_NONCE_SID,                \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_SET_LENGTHS_SID,              \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_UPDATE_AD_SID,                \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_UPDATE_SID,                   \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_FINISH_SID,                   \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_VERIFY_SID,                   \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_AEAD_ABORT_SID,                    \
+        TFM_CRYPTO_GROUP_ID_AEAD,                   \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_ASYMMETRIC_SIGN_MESSAGE_SID,       \
+        TFM_CRYPTO_GROUP_ID_ASYM_SIGN,              \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_ASYMMETRIC_VERIFY_MESSAGE_SID,     \
+        TFM_CRYPTO_GROUP_ID_ASYM_SIGN,              \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_ASYMMETRIC_SIGN_HASH_SID,          \
+        TFM_CRYPTO_GROUP_ID_ASYM_SIGN,              \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_ASYMMETRIC_VERIFY_HASH_SID,        \
+        TFM_CRYPTO_GROUP_ID_ASYM_SIGN,              \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,            \
+        TFM_CRYPTO_GROUP_ID_ASYM_ENCRYPT,           \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,            \
+        TFM_CRYPTO_GROUP_ID_ASYM_ENCRYPT,           \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,             \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,          \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_SETUP)             \
+    X(TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,   \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,   \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,    \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,      \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,  \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,   \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,     \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,          \
+        TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,         \
+        TFM_CRYPTO_FUNCTION_TYPE_LOOKUP)            \
+    X(TFM_CRYPTO_GENERATE_RANDOM_SID,               \
+        TFM_CRYPTO_GROUP_ID_RANDOM,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+    X(TFM_CRYPTO_API_DISPATCHER_SID,                \
+        TFM_CRYPTO_GROUP_ID_RANDOM,                 \
+        TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART)     \
+/**
+ * \brief Numerical progressive value identifying a function API exposed through
+ *        the interfaces (S or NS). It's used to dispatch the requests from S/NS
+ *        to the corresponding API implementation in the Crypto service backend.
+ */
+#define X(function_id, group_id, function_type) function_id,
+enum tfm_crypto_function_id {
+    TFM_CRYPTO_SERVICE_API_DESCRIPTION
 };
+#undef X
 
 /**
  * \brief Define an invalid value for an SID
diff --git a/interface/src/tfm_crypto_func_api.c b/interface/src/tfm_crypto_func_api.c
index 66759fc..66615fd 100644
--- a/interface/src/tfm_crypto_func_api.c
+++ b/interface/src/tfm_crypto_func_api.c
@@ -11,15 +11,14 @@
 #include "psa/crypto.h"
 #include "tfm_ns_interface.h"
 
-#define API_DISPATCH(srv_name, srv_id)                               \
-        tfm_ns_interface_dispatch((veneer_fn)srv_name##_veneer,      \
-        (uint32_t)in_vec, IOVEC_LEN(in_vec),                         \
-        (uint32_t)out_vec, IOVEC_LEN(out_vec))
-
-#define API_DISPATCH_NO_OUTVEC(srv_name, srv_id)                     \
-        tfm_ns_interface_dispatch((veneer_fn)srv_name##_veneer,      \
-        (uint32_t)in_vec, IOVEC_LEN(in_vec),                         \
-        (uint32_t)NULL, 0)
+#define API_DISPATCH(in_vec, out_vec)          \
+    tfm_crypto_api_dispatcher_veneer(          \
+                in_vec, IOVEC_LEN(in_vec),     \
+                out_vec, IOVEC_LEN(out_vec))
+#define API_DISPATCH_NO_OUTVEC(in_vec)         \
+    tfm_crypto_api_dispatcher_veneer(          \
+                in_vec, IOVEC_LEN(in_vec),     \
+                NULL, 0)
 
 psa_status_t psa_crypto_init(void)
 {
@@ -33,32 +32,30 @@
                           psa_key_id_t *key)
 {
     const struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_OPEN_KEY_SID,
+        .function_id = TFM_CRYPTO_OPEN_KEY_SID,
+        .key_id = id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
-        {.base = &id, .len = sizeof(psa_key_id_t)},
     };
     psa_outvec out_vec[] = {
         {.base = key, .len = sizeof(psa_key_id_t)},
     };
 
-    return API_DISPATCH(tfm_crypto_open_key,
-                        TFM_CRYPTO_OPEN_KEY);
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_close_key(psa_key_id_t key)
 {
     const struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CLOSE_KEY_SID,
+        .function_id = TFM_CRYPTO_CLOSE_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    return API_DISPATCH_NO_OUTVEC(tfm_crypto_close_key,
-                                  TFM_CRYPTO_CLOSE_KEY);
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_import_key(const psa_key_attributes_t *attributes,
@@ -66,9 +63,8 @@
                             size_t data_length,
                             psa_key_id_t *key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_IMPORT_KEY_SID,
+        .function_id = TFM_CRYPTO_IMPORT_KEY_SID,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -79,35 +75,27 @@
         {.base = key, .len = sizeof(psa_key_id_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_import_key,
-                          TFM_CRYPTO_IMPORT_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_destroy_key(psa_key_id_t key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_DESTROY_KEY_SID,
+        .function_id = TFM_CRYPTO_DESTROY_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_destroy_key,
-                                    TFM_CRYPTO_DESTROY_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_get_key_attributes(psa_key_id_t key,
                                     psa_key_attributes_t *attributes)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
+        .function_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
@@ -117,15 +105,13 @@
         {.base = attributes, .len = sizeof(psa_key_attributes_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_get_key_attributes,
-                          TFM_CRYPTO_GET_KEY_ATTRIBUTES);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 void psa_reset_key_attributes(psa_key_attributes_t *attributes)
 {
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
+        .function_id = TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -134,8 +120,7 @@
         {.base = attributes, .len = sizeof(psa_key_attributes_t)},
     };
 
-    (void)API_DISPATCH(tfm_crypto_reset_key_attributes,
-                       TFM_CRYPTO_RESET_KEY_ATTRIBUTES);
+    (void)API_DISPATCH(in_vec, out_vec);
     return;
 }
 
@@ -146,7 +131,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_EXPORT_KEY_SID,
+        .function_id = TFM_CRYPTO_EXPORT_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
@@ -156,8 +141,7 @@
         {.base = data, .len = data_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_export_key,
-                          TFM_CRYPTO_EXPORT_KEY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *data_length = out_vec[0].len;
 
@@ -171,7 +155,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
+        .function_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
         .key_id = key,
     };
 
@@ -182,8 +166,7 @@
         {.base = data, .len = data_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_export_public_key,
-                          TFM_CRYPTO_EXPORT_PUBLIC_KEY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *data_length = out_vec[0].len;
 
@@ -192,27 +175,23 @@
 
 psa_status_t psa_purge_key(psa_key_id_t key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_PURGE_KEY_SID,
+        .function_id = TFM_CRYPTO_PURGE_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_purge_key,
-                                    TFM_CRYPTO_PURGE_KEY);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_copy_key(psa_key_id_t source_key,
                           const psa_key_attributes_t *attributes,
                           psa_key_id_t *target_key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_COPY_KEY_SID,
+        .function_id = TFM_CRYPTO_COPY_KEY_SID,
         .key_id = source_key,
     };
 
@@ -225,10 +204,7 @@
         {.base = target_key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_copy_key,
-                          TFM_CRYPTO_COPY_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_generate_iv(psa_cipher_operation_t *operation,
@@ -238,7 +214,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
+        .function_id = TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
         .op_handle = operation->handle,
     };
 
@@ -246,14 +222,12 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = iv, .len = iv_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_generate_iv,
-                          TFM_CRYPTO_CIPHER_GENERATE_IV);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *iv_length = out_vec[1].len;
+    *iv_length = out_vec[0].len;
 
     return status;
 }
@@ -262,9 +236,8 @@
                                const unsigned char *iv,
                                size_t iv_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_SET_IV_SID,
+        .function_id = TFM_CRYPTO_CIPHER_SET_IV_SID,
         .op_handle = operation->handle,
     };
 
@@ -272,23 +245,16 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = iv, .len = iv_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_cipher_set_iv,
-                          TFM_CRYPTO_CIPHER_SET_IV);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_cipher_encrypt_setup(psa_cipher_operation_t *operation,
                                       psa_key_id_t key,
                                       psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -301,19 +267,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_encrypt_setup,
-                          TFM_CRYPTO_CIPHER_ENCRYPT_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_decrypt_setup(psa_cipher_operation_t *operation,
                                       psa_key_id_t key,
                                       psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -326,10 +288,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_decrypt_setup,
-                          TFM_CRYPTO_CIPHER_DECRYPT_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_update(psa_cipher_operation_t *operation,
@@ -341,7 +300,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_UPDATE_SID,
+        .function_id = TFM_CRYPTO_CIPHER_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -350,23 +309,20 @@
         {.base = input, .len = input_length},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_update,
-                          TFM_CRYPTO_CIPHER_UPDATE);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *output_length = out_vec[1].len;
+    *output_length = out_vec[0].len;
 
     return status;
 }
 
 psa_status_t psa_cipher_abort(psa_cipher_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ABORT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -377,10 +333,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_abort,
-                          TFM_CRYPTO_CIPHER_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_finish(psa_cipher_operation_t *operation,
@@ -390,7 +343,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_FINISH_SID,
+        .function_id = TFM_CRYPTO_CIPHER_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -402,8 +355,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_finish,
-                          TFM_CRYPTO_CIPHER_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[1].len;
 
@@ -413,9 +365,8 @@
 psa_status_t psa_hash_setup(psa_hash_operation_t *operation,
                             psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_SETUP_SID,
+        .function_id = TFM_CRYPTO_HASH_SETUP_SID,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -427,19 +378,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_setup,
-                          TFM_CRYPTO_HASH_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_update(psa_hash_operation_t *operation,
                              const uint8_t *input,
                              size_t input_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_UPDATE_SID,
+        .function_id = TFM_CRYPTO_HASH_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -447,14 +394,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_hash_update,
-                          TFM_CRYPTO_HASH_UPDATE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_hash_finish(psa_hash_operation_t *operation,
@@ -464,7 +405,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_FINISH_SID,
+        .function_id = TFM_CRYPTO_HASH_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -476,8 +417,7 @@
         {.base = hash, .len = hash_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_finish,
-                          TFM_CRYPTO_HASH_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *hash_length = out_vec[1].len;
 
@@ -488,9 +428,8 @@
                              const uint8_t *hash,
                              size_t hash_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_VERIFY_SID,
+        .function_id = TFM_CRYPTO_HASH_VERIFY_SID,
         .op_handle = operation->handle,
     };
 
@@ -502,17 +441,13 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_verify,
-                          TFM_CRYPTO_HASH_VERIFY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_abort(psa_hash_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_ABORT_SID,
+        .function_id = TFM_CRYPTO_HASH_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -523,18 +458,14 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_abort,
-                          TFM_CRYPTO_HASH_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_clone(const psa_hash_operation_t *source_operation,
                             psa_hash_operation_t *target_operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_CLONE_SID,
+        .function_id = TFM_CRYPTO_HASH_CLONE_SID,
         .op_handle = source_operation->handle,
     };
 
@@ -549,10 +480,7 @@
         return PSA_ERROR_BAD_STATE;
     }
 
-    status = API_DISPATCH(tfm_crypto_hash_clone,
-                          TFM_CRYPTO_HASH_CLONE);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_compute(psa_algorithm_t alg,
@@ -564,7 +492,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_COMPUTE_SID,
+        .function_id = TFM_CRYPTO_HASH_COMPUTE_SID,
         .alg = alg,
     };
 
@@ -577,8 +505,7 @@
         {.base = hash, .len = hash_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_compute,
-                          TFM_CRYPTO_HASH_COMPUTE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *hash_length = out_vec[0].len;
 
@@ -591,9 +518,8 @@
                               const uint8_t *hash,
                               size_t hash_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_COMPARE_SID,
+        .function_id = TFM_CRYPTO_HASH_COMPARE_SID,
         .alg = alg,
     };
 
@@ -603,19 +529,15 @@
         {.base = hash, .len = hash_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_hash_compare,
-                          TFM_CRYPTO_HASH_COMPARE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_mac_sign_setup(psa_mac_operation_t *operation,
                                 psa_key_id_t key,
                                 psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
+        .function_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -628,19 +550,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_sign_setup,
-                          TFM_CRYPTO_MAC_SIGN_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_verify_setup(psa_mac_operation_t *operation,
                                   psa_key_id_t key,
                                   psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -653,19 +571,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_verify_setup,
-                          TFM_CRYPTO_MAC_VERIFY_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_update(psa_mac_operation_t *operation,
                             const uint8_t *input,
                             size_t input_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_UPDATE_SID,
+        .function_id = TFM_CRYPTO_MAC_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -673,14 +587,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_mac_update,
-                          TFM_CRYPTO_MAC_UPDATE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_mac_sign_finish(psa_mac_operation_t *operation,
@@ -690,7 +598,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_SIGN_FINISH_SID,
+        .function_id = TFM_CRYPTO_MAC_SIGN_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -702,8 +610,7 @@
         {.base = mac, .len = mac_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_sign_finish,
-                          TFM_CRYPTO_MAC_SIGN_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *mac_length = out_vec[1].len;
 
@@ -714,9 +621,8 @@
                                    const uint8_t *mac,
                                    size_t mac_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -728,17 +634,13 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_verify_finish,
-                          TFM_CRYPTO_MAC_VERIFY_FINISH);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_abort(psa_mac_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_ABORT_SID,
+        .function_id = TFM_CRYPTO_MAC_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -749,10 +651,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_abort,
-                          TFM_CRYPTO_MAC_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_aead_encrypt(psa_key_id_t key,
@@ -769,7 +668,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
+        .function_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
         .key_id = key,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = 0}
@@ -803,8 +702,7 @@
     in_vec[0].base = &iov;
     in_vec[0].len = sizeof(struct tfm_crypto_pack_iovec);
 
-    status = API_DISPATCH(tfm_crypto_aead_encrypt,
-                          TFM_CRYPTO_AEAD_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *ciphertext_length = out_vec[0].len;
 
@@ -825,7 +723,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
+        .function_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
         .key_id = key,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = 0}
@@ -859,8 +757,7 @@
     in_vec[0].base = &iov;
     in_vec[0].len = sizeof(struct tfm_crypto_pack_iovec);
 
-    status = API_DISPATCH(tfm_crypto_aead_decrypt,
-                          TFM_CRYPTO_AEAD_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *plaintext_length = out_vec[0].len;
 
@@ -871,9 +768,8 @@
                                     psa_key_id_t key,
                                     psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -886,18 +782,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_encrypt_setup,
-                          TFM_CRYPTO_AEAD_ENCRYPT_SETUP);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_aead_decrypt_setup(psa_aead_operation_t *operation,
                                     psa_key_id_t key,
                                     psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -910,9 +803,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_decrypt_setup,
-                          TFM_CRYPTO_AEAD_DECRYPT_SETUP);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_aead_generate_nonce(psa_aead_operation_t *operation,
@@ -922,7 +813,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
+        .function_id = TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
         .op_handle = operation->handle,
     };
 
@@ -930,14 +821,13 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = nonce, .len = nonce_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_generate_nonce,
-                          TFM_CRYPTO_AEAD_GENERATE_NONCE);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *nonce_length = out_vec[1].len;
+    *nonce_length = out_vec[0].len;
+
     return status;
 }
 
@@ -945,9 +835,8 @@
                                 const uint8_t *nonce,
                                 size_t nonce_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_SET_NONCE_SID,
+        .function_id = TFM_CRYPTO_AEAD_SET_NONCE_SID,
         .op_handle = operation->handle,
     };
 
@@ -955,22 +844,16 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = nonce, .len = nonce_length}
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_set_nonce,
-                          TFM_CRYPTO_AEAD_SET_NONCE);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_aead_set_lengths(psa_aead_operation_t *operation,
                                   size_t ad_length,
                                   size_t plaintext_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
+        .function_id = TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
         .ad_length = ad_length,
         .plaintext_length = plaintext_length,
         .op_handle = operation->handle,
@@ -979,22 +862,16 @@
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_set_lengths,
-                          TFM_CRYPTO_AEAD_SET_LENGTHS);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_aead_update_ad(psa_aead_operation_t *operation,
                                 const uint8_t *input,
                                 size_t input_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_UPDATE_AD_SID,
+        .function_id = TFM_CRYPTO_AEAD_UPDATE_AD_SID,
         .op_handle = operation->handle,
     };
 
@@ -1002,13 +879,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_update_ad,
-                          TFM_CRYPTO_AEAD_UPDATE_AD);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_aead_update(psa_aead_operation_t *operation,
@@ -1020,7 +892,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_UPDATE_SID,
+        .function_id = TFM_CRYPTO_AEAD_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -1029,14 +901,13 @@
         {.base = input, .len = input_length}
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_update,
-                          TFM_CRYPTO_AEAD_UPDATE);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *output_length = out_vec[1].len;
+    *output_length = out_vec[0].len;
+
     return status;
 }
 
@@ -1050,7 +921,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_FINISH_SID,
+        .function_id = TFM_CRYPTO_AEAD_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -1063,11 +934,11 @@
         {.base = ciphertext, .len = ciphertext_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_finish,
-                          TFM_CRYPTO_AEAD_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *ciphertext_length = out_vec[2].len;
     *tag_length = out_vec[1].len;
+
     return status;
 }
 
@@ -1080,7 +951,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_VERIFY_SID,
+        .function_id = TFM_CRYPTO_AEAD_VERIFY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1093,18 +964,17 @@
         {.base = plaintext, .len = plaintext_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_verify,
-                          TFM_CRYPTO_AEAD_VERIFY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *plaintext_length = out_vec[1].len;
+
     return status;
 }
 
 psa_status_t psa_aead_abort(psa_aead_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ABORT_SID,
+        .function_id = TFM_CRYPTO_AEAD_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -1115,9 +985,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_abort,
-                          TFM_CRYPTO_AEAD_ABORT);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_sign_message(psa_key_id_t key,
@@ -1130,7 +998,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_SIGN_MESSAGE_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_SIGN_MESSAGE_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1143,10 +1011,10 @@
         {.base = signature, .len = signature_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_sign_message,
-                          TFM_CRYPTO_SIGN_MESSAGE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *signature_length = out_vec[0].len;
+
     return status;
 }
 
@@ -1157,9 +1025,8 @@
                                 const uint8_t *signature,
                                 size_t signature_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_VERIFY_MESSAGE_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_VERIFY_MESSAGE_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1170,10 +1037,7 @@
         {.base = signature, .len = signature_length}
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_verify_message,
-                                    TFM_CRYPTO_VERIFY_MESSAGE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_sign_hash(psa_key_id_t key,
@@ -1186,7 +1050,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_SIGN_HASH_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_SIGN_HASH_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1199,8 +1063,7 @@
         {.base = signature, .len = signature_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_sign_hash,
-                          TFM_CRYPTO_SIGN_HASH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *signature_length = out_vec[0].len;
 
@@ -1214,9 +1077,8 @@
                              const uint8_t *signature,
                              size_t signature_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_VERIFY_HASH_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_VERIFY_HASH_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1227,10 +1089,7 @@
         {.base = signature, .len = signature_length}
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_verify_hash,
-                                    TFM_CRYPTO_VERIFY_HASH);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_asymmetric_encrypt(psa_key_id_t key,
@@ -1245,7 +1104,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1265,8 +1124,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_asymmetric_encrypt,
-                          TFM_CRYPTO_ASYMMETRIC_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
 
@@ -1285,7 +1143,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1305,8 +1163,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_asymmetric_decrypt,
-                          TFM_CRYPTO_ASYMMETRIC_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
 
@@ -1317,9 +1174,8 @@
                                 const psa_key_derivation_operation_t *operation,
                                 size_t *capacity)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1331,10 +1187,7 @@
         {.base = capacity, .len = sizeof(size_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_get_capacity,
-                          TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_output_bytes(
@@ -1342,9 +1195,8 @@
                                       uint8_t *output,
                                       size_t output_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
         .op_handle = operation->handle,
     };
 
@@ -1356,10 +1208,7 @@
         {.base = output, .len = output_length},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_output_bytes,
-                          TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_input_key(
@@ -1367,9 +1216,8 @@
                                       psa_key_derivation_step_t step,
                                       psa_key_id_t key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
         .key_id = key,
         .step = step,
         .op_handle = operation->handle,
@@ -1379,17 +1227,13 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_input_key,
-                                    TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_key_derivation_abort(psa_key_derivation_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -1401,10 +1245,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_abort,
-                          TFM_CRYPTO_KEY_DERIVATION_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_key_agreement(
@@ -1414,9 +1255,8 @@
                                       const uint8_t *peer_key,
                                       size_t peer_key_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
         .key_id = private_key,
         .step = step,
         .op_handle = operation->handle,
@@ -1427,18 +1267,14 @@
         {.base = peer_key, .len = peer_key_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_key_agreement,
-                                    TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_generate_random(uint8_t *output,
                                  size_t output_size)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GENERATE_RANDOM_SID,
+        .function_id = TFM_CRYPTO_GENERATE_RANDOM_SID,
     };
 
     psa_invec in_vec[] = {
@@ -1453,18 +1289,14 @@
         return PSA_SUCCESS;
     }
 
-    status = API_DISPATCH(tfm_crypto_generate_random,
-                          TFM_CRYPTO_GENERATE_RANDOM);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
                               psa_key_id_t *key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GENERATE_KEY_SID,
+        .function_id = TFM_CRYPTO_GENERATE_KEY_SID,
     };
 
     psa_invec in_vec[] = {
@@ -1476,10 +1308,7 @@
         {.base = key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_generate_key,
-                          TFM_CRYPTO_GENERATE_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_compute(psa_key_id_t key,
@@ -1492,7 +1321,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_COMPUTE_SID,
+        .function_id = TFM_CRYPTO_MAC_COMPUTE_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1505,10 +1334,10 @@
         {.base = mac, .len = mac_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_compute,
-                          TFM_CRYPTO_MAC_COMPUTE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *mac_length = out_vec[0].len;
+
     return status;
 }
 
@@ -1519,9 +1348,8 @@
                             const uint8_t *mac,
                             const size_t mac_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1532,10 +1360,7 @@
         {.base = mac, .len = mac_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_mac_verify,
-                                    TFM_CRYPTO_MAC_VERIFY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_cipher_encrypt(psa_key_id_t key,
@@ -1548,7 +1373,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ENCRYPT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ENCRYPT_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1561,10 +1386,10 @@
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_encrypt,
-                          TFM_CRYPTO_CIPHER_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
+
     return status;
 }
 
@@ -1578,7 +1403,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_DECRYPT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_DECRYPT_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1591,10 +1416,10 @@
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_decrypt,
-                          TFM_CRYPTO_CIPHER_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
+
     return status;
 }
 
@@ -1608,7 +1433,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
+        .function_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
         .alg = alg,
         .key_id = private_key
     };
@@ -1622,8 +1447,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_raw_key_agreement,
-                          TFM_CRYPTO_RAW_KEY_AGREEMENT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
 
@@ -1633,9 +1457,8 @@
 psa_status_t psa_key_derivation_setup(psa_key_derivation_operation_t *operation,
                                       psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -1647,18 +1470,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_setup,
-                          TFM_CRYPTO_KEY_DERIVATION_SETUP);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_set_capacity(
                                       psa_key_derivation_operation_t *operation,
                                       size_t capacity)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
         .capacity = capacity,
         .op_handle = operation->handle,
     };
@@ -1667,9 +1487,7 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_set_capacity,
-                                    TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_key_derivation_input_bytes(
@@ -1678,9 +1496,8 @@
                                       const uint8_t *data,
                                       size_t data_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
         .step = step,
         .op_handle = operation->handle,
     };
@@ -1690,9 +1507,7 @@
         {.base = data, .len = data_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_input_bytes,
-                                    TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_key_derivation_output_key(
@@ -1700,9 +1515,8 @@
                                       psa_key_derivation_operation_t *operation,
                                       psa_key_id_t *key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1715,7 +1529,5 @@
         {.base = key, .len = sizeof(psa_key_id_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_output_key,
-                          TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
diff --git a/interface/src/tfm_crypto_ipc_api.c b/interface/src/tfm_crypto_ipc_api.c
index da5358a..cd1f7f5 100644
--- a/interface/src/tfm_crypto_ipc_api.c
+++ b/interface/src/tfm_crypto_ipc_api.c
@@ -7,19 +7,18 @@
 
 #include "tfm_crypto_defs.h"
 #include "psa/crypto.h"
-#include "tfm_ns_interface.h"
-#include "psa_manifest/sid.h"
 #include "psa/client.h"
+#include "psa_manifest/sid.h"
+#include "tfm_ns_interface.h"
 
-#define API_DISPATCH(srv_name, srv_id)                          \
-    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,                   \
-        in_vec, IOVEC_LEN(in_vec),                              \
-        out_vec, IOVEC_LEN(out_vec))
-
-#define API_DISPATCH_NO_OUTVEC(srv_name, srv_id)                \
-    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,                   \
-        in_vec, IOVEC_LEN(in_vec),                              \
-        (psa_outvec *)NULL, 0)
+#define API_DISPATCH(in_vec, out_vec)          \
+    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,  \
+             in_vec, IOVEC_LEN(in_vec),        \
+             out_vec, IOVEC_LEN(out_vec))
+#define API_DISPATCH_NO_OUTVEC(in_vec)         \
+    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,  \
+             in_vec, IOVEC_LEN(in_vec),        \
+             (psa_outvec *)NULL, 0)
 
 psa_status_t psa_crypto_init(void)
 {
@@ -32,39 +31,31 @@
 psa_status_t psa_open_key(psa_key_id_t id,
                           psa_key_id_t *key)
 {
-    psa_status_t status;
     const struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_OPEN_KEY_SID,
+        .function_id = TFM_CRYPTO_OPEN_KEY_SID,
+        .key_id = id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
-        {.base = &id, .len = sizeof(psa_key_id_t)},
     };
     psa_outvec out_vec[] = {
         {.base = key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_open_key,
-                          TFM_CRYPTO_OPEN_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_close_key(psa_key_id_t key)
 {
-    psa_status_t status;
     const struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CLOSE_KEY_SID,
+        .function_id = TFM_CRYPTO_CLOSE_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_close_key,
-                                    TFM_CRYPTO_CLOSE_KEY);;
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_import_key(const psa_key_attributes_t *attributes,
@@ -72,9 +63,8 @@
                             size_t data_length,
                             psa_key_id_t *key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_IMPORT_KEY_SID,
+        .function_id = TFM_CRYPTO_IMPORT_KEY_SID,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -85,35 +75,27 @@
         {.base = key, .len = sizeof(psa_key_id_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_import_key,
-                          TFM_CRYPTO_IMPORT_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_destroy_key(psa_key_id_t key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_DESTROY_KEY_SID,
+        .function_id = TFM_CRYPTO_DESTROY_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_destroy_key,
-                                    TFM_CRYPTO_DESTROY_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_get_key_attributes(psa_key_id_t key,
                                     psa_key_attributes_t *attributes)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
+        .function_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
@@ -123,15 +105,13 @@
         {.base = attributes, .len = sizeof(psa_key_attributes_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_get_key_attributes,
-                          TFM_CRYPTO_GET_KEY_ATTRIBUTES);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 void psa_reset_key_attributes(psa_key_attributes_t *attributes)
 {
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
+        .function_id = TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -140,8 +120,7 @@
         {.base = attributes, .len = sizeof(psa_key_attributes_t)},
     };
 
-    (void)API_DISPATCH(tfm_crypto_reset_key_attributes,
-                       TFM_CRYPTO_RESET_KEY_ATTRIBUTES);
+    (void)API_DISPATCH(in_vec, out_vec);
     return;
 }
 
@@ -152,7 +131,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_EXPORT_KEY_SID,
+        .function_id = TFM_CRYPTO_EXPORT_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
@@ -162,8 +141,7 @@
         {.base = data, .len = data_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_export_key,
-                          TFM_CRYPTO_EXPORT_KEY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *data_length = out_vec[0].len;
 
@@ -177,7 +155,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
+        .function_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
         .key_id = key,
     };
 
@@ -188,8 +166,7 @@
         {.base = data, .len = data_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_export_public_key,
-                          TFM_CRYPTO_EXPORT_PUBLIC_KEY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *data_length = out_vec[0].len;
 
@@ -198,27 +175,23 @@
 
 psa_status_t psa_purge_key(psa_key_id_t key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_PURGE_KEY_SID,
+        .function_id = TFM_CRYPTO_PURGE_KEY_SID,
         .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_purge_key,
-                                    TFM_CRYPTO_PURGE_KEY);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_copy_key(psa_key_id_t source_key,
                           const psa_key_attributes_t *attributes,
                           psa_key_id_t *target_key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_COPY_KEY_SID,
+        .function_id = TFM_CRYPTO_COPY_KEY_SID,
         .key_id = source_key,
     };
 
@@ -231,10 +204,7 @@
         {.base = target_key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_copy_key,
-                          TFM_CRYPTO_COPY_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_generate_iv(psa_cipher_operation_t *operation,
@@ -244,7 +214,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
+        .function_id = TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
         .op_handle = operation->handle,
     };
 
@@ -252,14 +222,12 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = iv, .len = iv_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_generate_iv,
-                          TFM_CRYPTO_CIPHER_GENERATE_IV);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *iv_length = out_vec[1].len;
+    *iv_length = out_vec[0].len;
 
     return status;
 }
@@ -268,9 +236,8 @@
                                const unsigned char *iv,
                                size_t iv_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_SET_IV_SID,
+        .function_id = TFM_CRYPTO_CIPHER_SET_IV_SID,
         .op_handle = operation->handle,
     };
 
@@ -278,23 +245,16 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = iv, .len = iv_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_cipher_set_iv,
-                          TFM_CRYPTO_CIPHER_SET_IV);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_cipher_encrypt_setup(psa_cipher_operation_t *operation,
                                       psa_key_id_t key,
                                       psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -307,19 +267,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_encrypt_setup,
-                          TFM_CRYPTO_CIPHER_ENCRYPT_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_decrypt_setup(psa_cipher_operation_t *operation,
                                       psa_key_id_t key,
                                       psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -332,10 +288,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_decrypt_setup,
-                          TFM_CRYPTO_CIPHER_DECRYPT_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_update(psa_cipher_operation_t *operation,
@@ -347,7 +300,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_UPDATE_SID,
+        .function_id = TFM_CRYPTO_CIPHER_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -356,23 +309,20 @@
         {.base = input, .len = input_length},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_update,
-                          TFM_CRYPTO_CIPHER_UPDATE);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *output_length = out_vec[1].len;
+    *output_length = out_vec[0].len;
 
     return status;
 }
 
 psa_status_t psa_cipher_abort(psa_cipher_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ABORT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -383,10 +333,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_abort,
-                          TFM_CRYPTO_CIPHER_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_cipher_finish(psa_cipher_operation_t *operation,
@@ -396,7 +343,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_FINISH_SID,
+        .function_id = TFM_CRYPTO_CIPHER_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -408,8 +355,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_finish,
-                          TFM_CRYPTO_CIPHER_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[1].len;
 
@@ -419,9 +365,8 @@
 psa_status_t psa_hash_setup(psa_hash_operation_t *operation,
                             psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_SETUP_SID,
+        .function_id = TFM_CRYPTO_HASH_SETUP_SID,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -433,19 +378,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_setup,
-                          TFM_CRYPTO_HASH_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_update(psa_hash_operation_t *operation,
                              const uint8_t *input,
                              size_t input_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_UPDATE_SID,
+        .function_id = TFM_CRYPTO_HASH_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -453,14 +394,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_hash_update,
-                          TFM_CRYPTO_HASH_UPDATE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_hash_finish(psa_hash_operation_t *operation,
@@ -470,7 +405,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_FINISH_SID,
+        .function_id = TFM_CRYPTO_HASH_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -482,8 +417,7 @@
         {.base = hash, .len = hash_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_finish,
-                          TFM_CRYPTO_HASH_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *hash_length = out_vec[1].len;
 
@@ -494,9 +428,8 @@
                              const uint8_t *hash,
                              size_t hash_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_VERIFY_SID,
+        .function_id = TFM_CRYPTO_HASH_VERIFY_SID,
         .op_handle = operation->handle,
     };
 
@@ -508,17 +441,13 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_verify,
-                          TFM_CRYPTO_HASH_VERIFY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_abort(psa_hash_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_ABORT_SID,
+        .function_id = TFM_CRYPTO_HASH_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -529,18 +458,14 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_abort,
-                          TFM_CRYPTO_HASH_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_clone(const psa_hash_operation_t *source_operation,
                             psa_hash_operation_t *target_operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_CLONE_SID,
+        .function_id = TFM_CRYPTO_HASH_CLONE_SID,
         .op_handle = source_operation->handle,
     };
 
@@ -555,10 +480,7 @@
         return PSA_ERROR_BAD_STATE;
     }
 
-    status = API_DISPATCH(tfm_crypto_hash_clone,
-                          TFM_CRYPTO_HASH_CLONE);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_hash_compute(psa_algorithm_t alg,
@@ -570,7 +492,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_COMPUTE_SID,
+        .function_id = TFM_CRYPTO_HASH_COMPUTE_SID,
         .alg = alg,
     };
 
@@ -583,8 +505,7 @@
         {.base = hash, .len = hash_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_compute,
-                          TFM_CRYPTO_HASH_COMPUTE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *hash_length = out_vec[0].len;
 
@@ -597,9 +518,8 @@
                               const uint8_t *hash,
                               size_t hash_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_COMPARE_SID,
+        .function_id = TFM_CRYPTO_HASH_COMPARE_SID,
         .alg = alg,
     };
 
@@ -609,19 +529,15 @@
         {.base = hash, .len = hash_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_hash_compare,
-                          TFM_CRYPTO_HASH_COMPARE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_mac_sign_setup(psa_mac_operation_t *operation,
                                 psa_key_id_t key,
                                 psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
+        .function_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -634,19 +550,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_sign_setup,
-                          TFM_CRYPTO_MAC_SIGN_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_verify_setup(psa_mac_operation_t *operation,
                                   psa_key_id_t key,
                                   psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -659,19 +571,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_verify_setup,
-                          TFM_CRYPTO_MAC_VERIFY_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_update(psa_mac_operation_t *operation,
                             const uint8_t *input,
                             size_t input_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_UPDATE_SID,
+        .function_id = TFM_CRYPTO_MAC_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -679,14 +587,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_mac_update,
-                          TFM_CRYPTO_MAC_UPDATE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_mac_sign_finish(psa_mac_operation_t *operation,
@@ -696,7 +598,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_SIGN_FINISH_SID,
+        .function_id = TFM_CRYPTO_MAC_SIGN_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -708,8 +610,7 @@
         {.base = mac, .len = mac_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_sign_finish,
-                          TFM_CRYPTO_MAC_SIGN_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *mac_length = out_vec[1].len;
 
@@ -720,9 +621,8 @@
                                    const uint8_t *mac,
                                    size_t mac_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -734,17 +634,13 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_verify_finish,
-                          TFM_CRYPTO_MAC_VERIFY_FINISH);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_abort(psa_mac_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_ABORT_SID,
+        .function_id = TFM_CRYPTO_MAC_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -755,10 +651,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_abort,
-                          TFM_CRYPTO_MAC_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_aead_encrypt(psa_key_id_t key,
@@ -775,7 +668,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
+        .function_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
         .key_id = key,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = 0}
@@ -836,7 +729,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
+        .function_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
         .key_id = key,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = 0}
@@ -889,7 +782,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -902,8 +795,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_encrypt_setup,
-                          TFM_CRYPTO_AEAD_ENCRYPT_SETUP);
+    status = API_DISPATCH(in_vec, out_vec);
     return status;
 }
 
@@ -913,7 +805,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -926,8 +818,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_decrypt_setup,
-                          TFM_CRYPTO_AEAD_DECRYPT_SETUP);
+    status = API_DISPATCH(in_vec, out_vec);
     return status;
 }
 
@@ -938,7 +829,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
+        .function_id = TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
         .op_handle = operation->handle,
     };
 
@@ -946,14 +837,12 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = nonce, .len = nonce_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_generate_nonce,
-                          TFM_CRYPTO_AEAD_GENERATE_NONCE);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *nonce_length = out_vec[1].len;
+    *nonce_length = out_vec[0].len;
     return status;
 }
 
@@ -963,7 +852,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_SET_NONCE_SID,
+        .function_id = TFM_CRYPTO_AEAD_SET_NONCE_SID,
         .op_handle = operation->handle,
     };
 
@@ -971,12 +860,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = nonce, .len = nonce_length}
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_set_nonce,
-                          TFM_CRYPTO_AEAD_SET_NONCE);
+    status = API_DISPATCH_NO_OUTVEC(in_vec);
     return status;
 }
 
@@ -986,7 +871,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
+        .function_id = TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
         .ad_length = ad_length,
         .plaintext_length = plaintext_length,
         .op_handle = operation->handle,
@@ -995,12 +880,8 @@
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_set_lengths,
-                          TFM_CRYPTO_AEAD_SET_LENGTHS);
+    status = API_DISPATCH_NO_OUTVEC(in_vec);
     return status;
 }
 
@@ -1010,7 +891,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_UPDATE_AD_SID,
+        .function_id = TFM_CRYPTO_AEAD_UPDATE_AD_SID,
         .op_handle = operation->handle,
     };
 
@@ -1023,9 +904,6 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
     size_t in_len = IOVEC_LEN(in_vec);
 
@@ -1033,7 +911,7 @@
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, IOVEC_LEN(out_vec));
+                      NULL, 0);
     return status;
 }
 
@@ -1046,7 +924,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_UPDATE_SID,
+        .function_id = TFM_CRYPTO_AEAD_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -1060,7 +938,6 @@
         {.base = input, .len = input_length}
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = output, .len = output_size},
     };
 
@@ -1072,7 +949,7 @@
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
                       out_vec, IOVEC_LEN(out_vec));
 
-    *output_length = out_vec[1].len;
+    *output_length = out_vec[0].len;
     return status;
 }
 
@@ -1086,7 +963,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_FINISH_SID,
+        .function_id = TFM_CRYPTO_AEAD_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -1117,13 +994,14 @@
                       in_vec, IOVEC_LEN(in_vec),
                       out_vec, out_len);
 
-    *tag_length = out_vec[1].len;
-
     if (out_len == 3) {
         *ciphertext_length = out_vec[2].len;
     } else {
         *ciphertext_length = 0;
     }
+
+    *tag_length = out_vec[1].len;
+
     return status;
 }
 
@@ -1136,7 +1014,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_VERIFY_SID,
+        .function_id = TFM_CRYPTO_AEAD_VERIFY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1151,7 +1029,7 @@
     };
     psa_outvec out_vec[] = {
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
-        {.base = plaintext, .len = plaintext_size},
+        {.base = plaintext, .len = plaintext_size}
     };
 
     size_t out_len = IOVEC_LEN(out_vec);
@@ -1177,9 +1055,8 @@
 
 psa_status_t psa_aead_abort(psa_aead_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ABORT_SID,
+        .function_id = TFM_CRYPTO_AEAD_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -1190,9 +1067,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_abort,
-                          TFM_CRYPTO_AEAD_ABORT);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_sign_message(psa_key_id_t key,
@@ -1205,7 +1080,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_SIGN_MESSAGE_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_SIGN_MESSAGE_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1218,8 +1093,7 @@
         {.base = signature, .len = signature_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_sign_message,
-                          TFM_CRYPTO_SIGN_MESSAGE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *signature_length = out_vec[0].len;
     return status;
@@ -1232,9 +1106,8 @@
                                 const uint8_t *signature,
                                 size_t signature_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_VERIFY_MESSAGE_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_VERIFY_MESSAGE_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1245,10 +1118,7 @@
         {.base = signature, .len = signature_length}
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_verify_message,
-                                    TFM_CRYPTO_VERIFY_MESSAGE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_sign_hash(psa_key_id_t key,
@@ -1261,7 +1131,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_SIGN_HASH_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_SIGN_HASH_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1274,8 +1144,7 @@
         {.base = signature, .len = signature_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_sign_hash,
-                          TFM_CRYPTO_SIGN_HASH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *signature_length = out_vec[0].len;
 
@@ -1289,9 +1158,8 @@
                              const uint8_t *signature,
                              size_t signature_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_VERIFY_HASH_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_VERIFY_HASH_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1302,10 +1170,7 @@
         {.base = signature, .len = signature_length}
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_verify_hash,
-                                    TFM_CRYPTO_VERIFY_HASH);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_asymmetric_encrypt(psa_key_id_t key,
@@ -1320,7 +1185,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1365,7 +1230,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
         .key_id = key,
         .alg = alg
     };
@@ -1402,9 +1267,8 @@
                                 const psa_key_derivation_operation_t *operation,
                                 size_t *capacity)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1416,10 +1280,7 @@
         {.base = capacity, .len = sizeof(size_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_get_capacity,
-                          TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_output_bytes(
@@ -1427,9 +1288,8 @@
                                       uint8_t *output,
                                       size_t output_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
         .op_handle = operation->handle,
     };
 
@@ -1441,10 +1301,7 @@
         {.base = output, .len = output_length},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_output_bytes,
-                          TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_input_key(
@@ -1452,9 +1309,8 @@
                                       psa_key_derivation_step_t step,
                                       psa_key_id_t key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
         .key_id = key,
         .step = step,
         .op_handle = operation->handle,
@@ -1464,17 +1320,13 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_input_key,
-                                    TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_key_derivation_abort(psa_key_derivation_operation_t *operation)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -1486,10 +1338,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_abort,
-                          TFM_CRYPTO_KEY_DERIVATION_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_key_agreement(
@@ -1499,9 +1348,8 @@
                                       const uint8_t *peer_key,
                                       size_t peer_key_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
         .key_id = private_key,
         .step = step,
         .op_handle = operation->handle,
@@ -1512,18 +1360,14 @@
         {.base = peer_key, .len = peer_key_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_key_agreement,
-                                    TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_generate_random(uint8_t *output,
                                  size_t output_size)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GENERATE_RANDOM_SID,
+        .function_id = TFM_CRYPTO_GENERATE_RANDOM_SID,
     };
 
     psa_invec in_vec[] = {
@@ -1538,18 +1382,14 @@
         return PSA_SUCCESS;
     }
 
-    status = API_DISPATCH(tfm_crypto_generate_random,
-                          TFM_CRYPTO_GENERATE_RANDOM);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
                               psa_key_id_t *key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GENERATE_KEY_SID,
+        .function_id = TFM_CRYPTO_GENERATE_KEY_SID,
     };
 
     psa_invec in_vec[] = {
@@ -1561,10 +1401,7 @@
         {.base = key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_generate_key,
-                          TFM_CRYPTO_GENERATE_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_mac_compute(psa_key_id_t key,
@@ -1577,7 +1414,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_COMPUTE_SID,
+        .function_id = TFM_CRYPTO_MAC_COMPUTE_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1590,8 +1427,7 @@
         {.base = mac, .len = mac_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_compute,
-                          TFM_CRYPTO_MAC_COMPUTE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *mac_length = out_vec[0].len;
     return status;
@@ -1604,9 +1440,8 @@
                             const uint8_t *mac,
                             const size_t mac_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1617,10 +1452,7 @@
         {.base = mac, .len = mac_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_mac_verify,
-                                    TFM_CRYPTO_MAC_VERIFY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_cipher_encrypt(psa_key_id_t key,
@@ -1633,7 +1465,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ENCRYPT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ENCRYPT_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1646,8 +1478,7 @@
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_encrypt,
-                          TFM_CRYPTO_CIPHER_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
     return status;
@@ -1663,7 +1494,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_DECRYPT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_DECRYPT_SID,
         .key_id = key,
         .alg = alg,
     };
@@ -1676,8 +1507,7 @@
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_decrypt,
-                          TFM_CRYPTO_CIPHER_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
     return status;
@@ -1693,7 +1523,7 @@
 {
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
+        .function_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
         .alg = alg,
         .key_id = private_key
     };
@@ -1707,8 +1537,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_raw_key_agreement,
-                          TFM_CRYPTO_RAW_KEY_AGREEMENT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
 
@@ -1718,9 +1547,8 @@
 psa_status_t psa_key_derivation_setup(psa_key_derivation_operation_t *operation,
                                       psa_algorithm_t alg)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -1732,18 +1560,15 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_setup,
-                          TFM_CRYPTO_KEY_DERIVATION_SETUP);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
 
 psa_status_t psa_key_derivation_set_capacity(
                                       psa_key_derivation_operation_t *operation,
                                       size_t capacity)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
         .capacity = capacity,
         .op_handle = operation->handle,
     };
@@ -1752,9 +1577,7 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_set_capacity,
-                                    TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_key_derivation_input_bytes(
@@ -1763,9 +1586,8 @@
                                       const uint8_t *data,
                                       size_t data_length)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
         .step = step,
         .op_handle = operation->handle,
     };
@@ -1775,9 +1597,7 @@
         {.base = data, .len = data_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_input_bytes,
-                                    TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES);
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 }
 
 psa_status_t psa_key_derivation_output_key(
@@ -1785,9 +1605,8 @@
                                       psa_key_derivation_operation_t *operation,
                                       psa_key_id_t *key)
 {
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1800,7 +1619,5 @@
         {.base = key, .len = sizeof(psa_key_id_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_output_key,
-                          TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 }
diff --git a/lib/ext/mbedcrypto/0005-Fix-key_id-and-owner_id-accessor-macros.patch b/lib/ext/mbedcrypto/0005-Fix-key_id-and-owner_id-accessor-macros.patch
new file mode 100644
index 0000000..1d59f2b
--- /dev/null
+++ b/lib/ext/mbedcrypto/0005-Fix-key_id-and-owner_id-accessor-macros.patch
@@ -0,0 +1,31 @@
+From e76fe5287bdd23eb67ec45cb0c629a178733b0cc Mon Sep 17 00:00:00 2001
+From: Antonio de Angelis <Antonio.deAngelis@arm.com>
+Date: Thu, 5 May 2022 14:11:32 +0100
+Subject: [PATCH 1/1] Fix key_id and owner_id accessor macros
+
+The accessor macros for key_id and owner_id in the mbedtls_svc_key_id_t
+need to have the MBEDTLS_PRIVATE() specifier as these fields are private
+
+Signed-off-by: Antonio de Angelis <antonio.deangelis@arm.com>
+---
+ include/psa/crypto_values.h | 4 ++--
+ 1 file changed, 2 insertions(+), 2 deletions(-)
+
+diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
+index 5a903f86a..2d39625de 100644
+--- a/include/psa/crypto_values.h
++++ b/include/psa/crypto_values.h
+@@ -2242,8 +2242,8 @@ static inline int mbedtls_svc_key_id_is_null( mbedtls_svc_key_id_t key )
+ #else /* MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER */
+ 
+ #define MBEDTLS_SVC_KEY_ID_INIT ( (mbedtls_svc_key_id_t){ 0, 0 } )
+-#define MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ( ( id ).key_id )
+-#define MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( id ) ( ( id ).owner )
++#define MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ( ( id ).MBEDTLS_PRIVATE(key_id) )
++#define MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( id ) ( ( id ).MBEDTLS_PRIVATE(owner) )
+ 
+ /** Utility to initialize a key identifier at runtime.
+  *
+-- 
+2.25.1
+
diff --git a/secure_fw/partitions/crypto/crypto_aead.c b/secure_fw/partitions/crypto/crypto_aead.c
index 914aa14..2f9d2d5 100644
--- a/secure_fw/partitions/crypto/crypto_aead.c
+++ b/secure_fw/partitions/crypto/crypto_aead.c
@@ -12,561 +12,203 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
 /*!
- * \defgroup public_psa Public functions, PSA
- *
+ * \defgroup tfm_crypto_api_shim_layer Set of functions implementing a thin shim
+ *                                  layer between the TF-M Crypto service
+ *                                  frontend and the underlying library which
+ *                                  implements the PSA Crypto APIs
+ *                                  (i.e. mbed TLS)
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_aead_encrypt(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
+psa_status_t tfm_crypto_aead_interface(psa_invec in_vec[],
+                                       psa_outvec out_vec[],
+                                       mbedtls_svc_key_id_t *encoded_key)
 {
 #ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
+#endif
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    const struct tfm_crypto_aead_pack_input *aead_pack_input = &iov->aead_in;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *nonce = aead_pack_input->nonce;
-    size_t nonce_length = aead_pack_input->nonce_length;
-    const uint8_t *plaintext = in_vec[1].base;
-    size_t plaintext_length = in_vec[1].len;
-    uint8_t *ciphertext = out_vec[0].base;
-    size_t ciphertext_size = out_vec[0].len;
-    const uint8_t *additional_data = in_vec[2].base;
-    size_t additional_data_length = in_vec[2].len;
-    mbedtls_svc_key_id_t encoded_key;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_aead_operation_t *operation = NULL;
+    uint32_t handle = iov->op_handle;
+    uint32_t *handle_out = NULL;
+    enum tfm_crypto_function_type function_type =
+                        TFM_CRYPTO_GET_FUNCTION_TYPE(iov->function_id);
 
-    /* Initialise ciphertext_length to zero. */
-    out_vec[0].len = 0;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
+    if (function_type != TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART) {
+        handle_out = (out_vec && out_vec[0].base != NULL) ?
+                                                out_vec[0].base : &handle;
+        *handle_out = handle;
+        status = tfm_crypto_operation_handling(TFM_CRYPTO_AEAD_OPERATION,
+                                               function_type,
+                                               handle_out,
+                                               (void **)&operation);
+        if (status != PSA_SUCCESS) {
+            return (iov->function_id == TFM_CRYPTO_AEAD_ABORT_SID) ?
+                    PSA_SUCCESS : status;
+        }
     }
 
-    return psa_aead_encrypt(encoded_key, alg, nonce, nonce_length,
+    switch (iov->function_id) {
+    case TFM_CRYPTO_AEAD_ENCRYPT_SID:
+    {
+        const struct tfm_crypto_aead_pack_input *aead_pack_input =
+                                                                 &iov->aead_in;
+        const uint8_t *nonce = aead_pack_input->nonce;
+        size_t nonce_length = aead_pack_input->nonce_length;
+        const uint8_t *plaintext = in_vec[1].base;
+        size_t plaintext_length = in_vec[1].len;
+        uint8_t *ciphertext = out_vec[0].base;
+        size_t ciphertext_size = out_vec[0].len;
+        const uint8_t *additional_data = in_vec[2].base;
+        size_t additional_data_length = in_vec[2].len;
+        /* Initialise ciphertext_length to zero. */
+        out_vec[0].len = 0;
+
+        status = psa_aead_encrypt(*encoded_key, iov->alg, nonce, nonce_length,
                             additional_data, additional_data_length,
                             plaintext, plaintext_length,
                             ciphertext, ciphertext_size, &out_vec[0].len);
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_decrypt(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
     }
+    break;
+    case TFM_CRYPTO_AEAD_DECRYPT_SID:
+    {
+        const struct tfm_crypto_aead_pack_input *aead_pack_input =
+                                                                 &iov->aead_in;
+        const uint8_t *nonce = aead_pack_input->nonce;
+        size_t nonce_length = aead_pack_input->nonce_length;
+        const uint8_t *ciphertext = in_vec[1].base;
+        size_t ciphertext_length = in_vec[1].len;
+        uint8_t *plaintext = out_vec[0].base;
+        size_t plaintext_size = out_vec[0].len;
+        const uint8_t *additional_data = in_vec[2].base;
+        size_t additional_data_length = in_vec[2].len;
+        /* Initialise plaintext_length to zero. */
+        out_vec[0].len = 0;
 
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    const struct tfm_crypto_aead_pack_input *aead_pack_input = &iov->aead_in;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *nonce = aead_pack_input->nonce;
-    size_t nonce_length = aead_pack_input->nonce_length;
-    const uint8_t *ciphertext = in_vec[1].base;
-    size_t ciphertext_length = in_vec[1].len;
-    uint8_t *plaintext = out_vec[0].base;
-    size_t plaintext_size = out_vec[0].len;
-    const uint8_t *additional_data = in_vec[2].base;
-    size_t additional_data_length = in_vec[2].len;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Initialise plaintext_length to zero. */
-    out_vec[0].len = 0;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_aead_decrypt(encoded_key, alg, nonce, nonce_length,
+        status = psa_aead_decrypt(*encoded_key, iov->alg, nonce, nonce_length,
                             additional_data, additional_data_length,
                             ciphertext, ciphertext_length,
                             plaintext, plaintext_size, &out_vec[0].len);
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_encrypt_setup(psa_invec in_vec[],
-                                           size_t in_len,
-                                           psa_outvec out_vec[],
-                                           size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
     }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_AEAD_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
+    break;
+    case TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID:
+    {
+        status = psa_aead_encrypt_setup(operation, *encoded_key, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
     }
-
-    *handle_out = handle;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        goto exit;
+    break;
+    case TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID:
+    {
+        status = psa_aead_decrypt_setup(operation, *encoded_key, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
     }
+    break;
+    case TFM_CRYPTO_AEAD_ABORT_SID:
+    {
+        status = psa_aead_abort(operation);
 
-    status = psa_aead_encrypt_setup(operation, encoded_key, alg);
-    if (status != PSA_SUCCESS) {
-        goto exit;
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        } else {
+            status = tfm_crypto_operation_release(handle_out);
+        }
+    }
+    break;
+    case TFM_CRYPTO_AEAD_FINISH_SID:
+    {
+        uint8_t *ciphertext = out_vec[2].base;
+        size_t ciphertext_size = out_vec[2].len;
+        uint8_t *tag = out_vec[1].base;
+        size_t tag_size = out_vec[1].len;
+        /* Initialise tag and ciphertext lengths to zero */
+        out_vec[1].len = 0;
+        out_vec[2].len = 0;
+
+        status = psa_aead_finish(operation,
+                                 ciphertext, ciphertext_size, &out_vec[2].len,
+                                 tag, tag_size, &out_vec[1].len);
+        if (status == PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_AEAD_GENERATE_NONCE_SID:
+    {
+        uint8_t *nonce = out_vec[0].base;
+        size_t nonce_size = out_vec[0].len;
+        /* Initialise nonce length to zero */
+        out_vec[0].len = 0;
+
+        status = psa_aead_generate_nonce(operation,
+                                         nonce,
+                                         nonce_size,
+                                         &out_vec[0].len);
+    }
+    break;
+    case TFM_CRYPTO_AEAD_SET_NONCE_SID:
+    {
+        const uint8_t *nonce = in_vec[1].base;
+        size_t nonce_size = in_vec[1].len;
+
+        status = psa_aead_set_nonce(operation, nonce, nonce_size);
+    }
+    break;
+    case TFM_CRYPTO_AEAD_SET_LENGTHS_SID:
+    {
+        status = psa_aead_set_lengths(operation, iov->ad_length,
+                                      iov->plaintext_length);
+    }
+    break;
+    case TFM_CRYPTO_AEAD_UPDATE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        uint8_t *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+
+        status = psa_aead_update(operation, input, input_length,
+                                 output, output_size, &out_vec[0].len);
+    }
+    break;
+    case TFM_CRYPTO_AEAD_UPDATE_AD_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+
+        status = psa_aead_update_ad(operation, input, input_length);
+    }
+    break;
+    case TFM_CRYPTO_AEAD_VERIFY_SID:
+    {
+        const uint8_t *tag = in_vec[1].base;
+        size_t tag_length = in_vec[1].len;
+        uint8_t *plaintext = out_vec[1].base;
+        size_t plaintext_size = out_vec[1].len;
+
+        status = psa_aead_verify(operation,
+                                 plaintext, plaintext_size, &out_vec[1].len,
+                                 tag, tag_length);
+
+        if (status == PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    default:
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
 
     return status;
-
-exit:
+release_operation_and_return:
     /* Release the operation context, ignore if the operation fails. */
     (void)tfm_crypto_operation_release(handle_out);
     return status;
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_decrypt_setup(psa_invec in_vec[],
-                                           size_t in_len,
-                                           psa_outvec out_vec[],
-                                           size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_AEAD_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *handle_out = handle;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
-    status = psa_aead_decrypt_setup(operation, encoded_key, alg);
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
-    return status;
-
-exit:
-    /* Release the operation context, ignore if the operation fails. */
-    (void)tfm_crypto_operation_release(handle_out);
-    return status;
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_abort(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        /* Operation does not exist, so abort has no effect */
-        return PSA_SUCCESS;
-    }
-
-    status = psa_aead_abort(operation);
-
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-        return status;
-    }
-
-    return tfm_crypto_operation_release(handle_out);
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_finish(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 2, 3);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    uint8_t *ciphertext = out_vec[2].base;
-    size_t ciphertext_size = out_vec[2].len;
-    uint8_t *tag = out_vec[1].base;
-    size_t tag_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Initialise tag and ciphertext lengths to zero */
-    out_vec[1].len = 0;
-    out_vec[2].len = 0;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_aead_finish(operation,
-                             ciphertext, ciphertext_size, &out_vec[2].len,
-                             tag, tag_size, &out_vec[1].len);
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
-    return status;
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_generate_nonce(psa_invec in_vec[],
-                                            size_t in_len,
-                                            psa_outvec out_vec[],
-                                            size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 2, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    uint8_t *nonce = out_vec[1].base;
-    size_t nonce_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Initialise nonce length to zero */
-    out_vec[1].len = 0;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_aead_generate_nonce(operation,
-                                   nonce,
-                                   nonce_size,
-                                   &out_vec[1].len);
-
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_set_nonce(psa_invec in_vec[],
-                                       size_t in_len,
-                                       psa_outvec out_vec[],
-                                       size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *nonce = in_vec[1].base;
-    size_t nonce_size = in_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_aead_set_nonce(operation, nonce, nonce_size);
-
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_set_lengths(psa_invec in_vec[],
-                                         size_t in_len,
-                                         psa_outvec out_vec[],
-                                         size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    size_t ad_length = iov->ad_length;
-    size_t plaintext_length = iov->plaintext_length;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-   return psa_aead_set_lengths(operation, ad_length, plaintext_length);
-
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_update(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    uint8_t *output = out_vec[1].base;
-    size_t output_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_aead_update(operation, input, input_length,
-                           output, output_size, &out_vec[1].len);
-
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_update_ad(psa_invec in_vec[],
-                                       size_t in_len,
-                                       psa_outvec out_vec[],
-                                       size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_aead_update_ad(operation, input, input_length);
-
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_aead_verify(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_aead_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *tag = in_vec[1].base;
-    size_t tag_length = in_vec[1].len;
-    uint8_t *plaintext = out_vec[1].base;
-    size_t plaintext_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_AEAD_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_aead_verify(operation,
-                             plaintext, plaintext_size, &out_vec[1].len,
-                             tag, tag_length);
-
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
-    return status;
-#endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_alloc.c b/secure_fw/partitions/crypto/crypto_alloc.c
index 76fa8c2..7e3afe0 100644
--- a/secure_fw/partitions/crypto/crypto_alloc.c
+++ b/secure_fw/partitions/crypto/crypto_alloc.c
@@ -83,8 +83,9 @@
 }
 
 /*!
- * \defgroup public Public functions
- *
+ * \defgroup alloc Function that implement allocation and deallocation of
+ *                 contexts to be stored in the secure world for multipart
+ *                 operations
  */
 
 /*!@{*/
@@ -191,4 +192,35 @@
 
     return PSA_ERROR_BAD_STATE;
 }
+
+psa_status_t tfm_crypto_operation_handling(enum tfm_crypto_operation_type type,
+                                    enum tfm_crypto_function_type function_type,
+                                    uint32_t *handle,
+                                    void **ctx)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    /* Multipart context retrieving and handling if required */
+    switch (function_type) {
+    case TFM_CRYPTO_FUNCTION_TYPE_SETUP:
+        /* Allocate the operation context in the secure world */
+        status = tfm_crypto_operation_alloc(type,
+                                            handle,
+                                            ctx);
+        break;
+    case TFM_CRYPTO_FUNCTION_TYPE_LOOKUP:
+        /* Look up the corresponding operation context */
+        status = tfm_crypto_operation_lookup(type,
+                                             *handle,
+                                             ctx);
+        break;
+    /* All the other APIs don't deal with multipart */
+    case TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART:
+        status = PSA_ERROR_INVALID_ARGUMENT;
+    default:
+        break;
+    }
+
+    return status;
+}
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_asymmetric.c b/secure_fw/partitions/crypto/crypto_asymmetric.c
index 6640403..a8507b2 100644
--- a/secure_fw/partitions/crypto/crypto_asymmetric.c
+++ b/secure_fw/partitions/crypto/crypto_asymmetric.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
  *
@@ -12,243 +12,114 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
 /*!
- * \defgroup public_psa Public functions, PSA
+ * \addtogroup tfm_crypto_api_shim_layer
  *
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_sign_message(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
+psa_status_t tfm_crypto_asymmetric_interface(psa_invec in_vec[],
+                                             psa_outvec out_vec[],
+                                             mbedtls_svc_key_id_t *encoded_key)
 {
-#ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    uint8_t *signature = out_vec[0].base;
-    size_t signature_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
+#ifdef TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISALBED
+    if (TFM_CRYPTO_IS_GROUP_ID(iov->function_id,
+                               TFM_CRYPTO_GROUP_ID_ASYM_ENCRYPT)) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+#endif
+#ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISALBED
+    if (TFM_CRYPTO_IS_GROUP_ID(iov->function_id,
+                               TFM_CRYPTO_GROUP_ID_ASYM_SIGN)) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+#endif
+
+    switch (iov->function_id) {
+    case TFM_CRYPTO_ASYMMETRIC_SIGN_MESSAGE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        uint8_t *signature = out_vec[0].base;
+        size_t signature_size = out_vec[0].len;
+
+        status = psa_sign_message(*encoded_key, iov->alg, input, input_length,
+                                  signature, signature_size, &(out_vec[0].len));
+    }
+    break;
+    case TFM_CRYPTO_ASYMMETRIC_VERIFY_MESSAGE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        const uint8_t *signature = in_vec[2].base;
+        size_t signature_length = in_vec[2].len;
+
+        status = psa_verify_message(*encoded_key, iov->alg, input, input_length,
+                                    signature, signature_length);
+    }
+    break;
+    case TFM_CRYPTO_ASYMMETRIC_SIGN_HASH_SID:
+    {
+        const uint8_t *hash = in_vec[1].base;
+        size_t hash_length = in_vec[1].len;
+        uint8_t *signature = out_vec[0].base;
+        size_t signature_size = out_vec[0].len;
+
+        status = psa_sign_hash(*encoded_key, iov->alg, hash, hash_length,
+                               signature, signature_size, &(out_vec[0].len));
+    }
+    break;
+    case TFM_CRYPTO_ASYMMETRIC_VERIFY_HASH_SID:
+    {
+        const uint8_t *hash = in_vec[1].base;
+        size_t hash_length = in_vec[1].len;
+        const uint8_t *signature = in_vec[2].base;
+        size_t signature_length = in_vec[2].len;
+
+        status = psa_verify_hash(*encoded_key, iov->alg, hash, hash_length,
+                                 signature, signature_length);
+    }
+    break;
+    case TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        const uint8_t *salt = in_vec[2].base;
+        size_t salt_length = in_vec[2].len;
+        uint8_t *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+
+        status = psa_asymmetric_encrypt(*encoded_key, iov->alg,
+                                        input, input_length,
+                                        salt, salt_length,
+                                        output, output_size,
+                                        &(out_vec[0].len));
+    }
+    break;
+    case TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        const uint8_t *salt = in_vec[2].base;
+        size_t salt_length = in_vec[2].len;
+        uint8_t *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+
+        status = psa_asymmetric_decrypt(*encoded_key, iov->alg,
+                                        input, input_length,
+                                        salt, salt_length,
+                                        output, output_size,
+                                        &(out_vec[0].len));
+    }
+    break;
+    default:
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
 
-    return psa_sign_message(encoded_key, alg, input, input_length,
-                            signature, signature_size, &(out_vec[0].len));
-#endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_verify_message(psa_invec in_vec[],
-                                       size_t in_len,
-                                       psa_outvec out_vec[],
-                                       size_t out_len)
-{
-#ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    const uint8_t *signature = in_vec[2].base;
-    size_t signature_length = in_vec[2].len;
-    mbedtls_svc_key_id_t encoded_key;
-    psa_status_t status;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_verify_message(encoded_key, alg, input, input_length,
-                              signature, signature_length);
-#endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_sign_hash(psa_invec in_vec[],
-                                  size_t in_len,
-                                  psa_outvec out_vec[],
-                                  size_t out_len)
-{
-#ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *hash = in_vec[1].base;
-    size_t hash_length = in_vec[1].len;
-    uint8_t *signature = out_vec[0].base;
-    size_t signature_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-    psa_status_t status;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_sign_hash(encoded_key, alg, hash, hash_length,
-                         signature, signature_size, &(out_vec[0].len));
-#endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_verify_hash(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *hash = in_vec[1].base;
-    size_t hash_length = in_vec[1].len;
-    const uint8_t *signature = in_vec[2].base;
-    size_t signature_length = in_vec[2].len;
-    mbedtls_svc_key_id_t encoded_key;
-    psa_status_t status;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_verify_hash(encoded_key, alg, hash, hash_length,
-                           signature, signature_length);
-#endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_asymmetric_encrypt(psa_invec in_vec[],
-                                           size_t in_len,
-                                           psa_outvec out_vec[],
-                                           size_t out_len)
-{
-#ifdef TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    const uint8_t *salt = in_vec[2].base;
-    size_t salt_length = in_vec[2].len;
-    uint8_t *output = out_vec[0].base;
-    size_t output_size = out_vec[0].len;
-    psa_key_type_t type;
-    size_t key_bits;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_get_key_attributes(encoded_key, &key_attributes);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    key_bits = psa_get_key_bits(&key_attributes);
-    type = psa_get_key_type(&key_attributes);
-
-    psa_reset_key_attributes(&key_attributes);
-
-    /* Check that the output buffer is large enough */
-    if (output_size < PSA_ASYMMETRIC_ENCRYPT_OUTPUT_SIZE(type, key_bits, alg)) {
-        return PSA_ERROR_BUFFER_TOO_SMALL;
-    }
-
-    return psa_asymmetric_encrypt(encoded_key, alg, input, input_length,
-                                  salt, salt_length,
-                                  output, output_size, &(out_vec[0].len));
-#endif /* TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_asymmetric_decrypt(psa_invec in_vec[],
-                                           size_t in_len,
-                                           psa_outvec out_vec[],
-                                           size_t out_len)
-{
-#ifdef TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    const uint8_t *salt = in_vec[2].base;
-    size_t salt_length = in_vec[2].len;
-    uint8_t *output = out_vec[0].base;
-    size_t output_size = out_vec[0].len;
-    psa_status_t status;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_asymmetric_decrypt(encoded_key, alg, input, input_length,
-                                  salt, salt_length,
-                                  output, output_size, &(out_vec[0].len));
-#endif /* TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED */
+    return status;
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_cipher.c b/secure_fw/partitions/crypto/crypto_cipher.c
index a45d239..70ce182c 100644
--- a/secure_fw/partitions/crypto/crypto_cipher.c
+++ b/secure_fw/partitions/crypto/crypto_cipher.c
@@ -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
  *
@@ -12,400 +12,145 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
 /*!
- * \defgroup public_psa Public functions, PSA
+ * \addtogroup tfm_crypto_api_shim_layer
  *
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_cipher_generate_iv(psa_invec in_vec[],
-                                           size_t in_len,
-                                           psa_outvec out_vec[],
-                                           size_t out_len)
+psa_status_t tfm_crypto_cipher_interface(psa_invec in_vec[],
+                                         psa_outvec out_vec[],
+                                         mbedtls_svc_key_id_t *encoded_key)
 {
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
+#endif
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    unsigned char *iv = out_vec[1].base;
-    size_t iv_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_CIPHER_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *handle_out = handle;
-
-    return psa_cipher_generate_iv(operation, iv, iv_size, &out_vec[1].len);
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_set_iv(psa_invec in_vec[],
-                                      size_t in_len,
-                                      psa_outvec out_vec[],
-                                      size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
     uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const unsigned char *iv = in_vec[1].base;
-    size_t iv_length = in_vec[1].len;
+    uint32_t *handle_out = NULL;
+    enum tfm_crypto_function_type function_type =
+                        TFM_CRYPTO_GET_FUNCTION_TYPE(iov->function_id);
 
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_CIPHER_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
+    if (function_type != TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART) {
+        handle_out = (out_vec && out_vec[0].base != NULL) ?
+                                                out_vec[0].base : &handle;
+        *handle_out = handle;
+        status = tfm_crypto_operation_handling(TFM_CRYPTO_CIPHER_OPERATION,
+                                               function_type,
+                                               handle_out,
+                                               (void **)&operation);
+        if (status != PSA_SUCCESS) {
+            return (iov->function_id == TFM_CRYPTO_CIPHER_ABORT_SID) ?
+                    PSA_SUCCESS : status;
+        }
     }
 
-    return psa_cipher_set_iv(operation, iv, iv_length);
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
+    switch (iov->function_id) {
+    case TFM_CRYPTO_CIPHER_GENERATE_IV_SID:
+    {
+        unsigned char *iv = out_vec[0].base;
+        size_t iv_size = out_vec[0].len;
 
-psa_status_t tfm_crypto_cipher_encrypt_setup(psa_invec in_vec[],
-                                             size_t in_len,
-                                             psa_outvec out_vec[],
-                                             size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((out_vec[0].len != sizeof(uint32_t)) ||
-        (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
+        status = psa_cipher_generate_iv(operation,
+                                        iv, iv_size, &out_vec[0].len);
     }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    mbedtls_svc_key_id_t encoded_key;
+    break;
+    case TFM_CRYPTO_CIPHER_SET_IV_SID:
+    {
+        const unsigned char *iv = in_vec[1].base;
+        size_t iv_length = in_vec[1].len;
 
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_CIPHER_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
+        status = psa_cipher_set_iv(operation, iv, iv_length);
     }
-    *handle_out = handle;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        goto exit;
+    break;
+    case TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID:
+    {
+        status = psa_cipher_encrypt_setup(operation, *encoded_key, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
     }
+    break;
+    case TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID:
+    {
+        status = psa_cipher_decrypt_setup(operation, *encoded_key, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_CIPHER_UPDATE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        unsigned char *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+        /* Initialise the output_length to zero */
+        out_vec[0].len = 0;
 
-    status = psa_cipher_encrypt_setup(operation, encoded_key, alg);
-    if (status != PSA_SUCCESS) {
-        goto exit;
+        status = psa_cipher_update(operation, input, input_length,
+                                   output, output_size, &out_vec[0].len);
+    }
+    break;
+    case TFM_CRYPTO_CIPHER_FINISH_SID:
+    {
+        uint8_t *output = out_vec[1].base;
+        size_t output_size = out_vec[1].len;
+        /* Initialise the output_length to zero */
+        out_vec[1].len = 0;
+
+        status = psa_cipher_finish(operation,
+                                   output, output_size, &out_vec[1].len);
+        if (status == PSA_SUCCESS) {
+            /* In case of success automatically release the operation */
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_CIPHER_ABORT_SID:
+    {
+        status = psa_cipher_abort(operation);
+
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        } else {
+            status = tfm_crypto_operation_release(handle_out);
+        }
+    }
+    break;
+    case TFM_CRYPTO_CIPHER_ENCRYPT_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        uint8_t *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+
+        status = psa_cipher_encrypt(*encoded_key, iov->alg, input, input_length,
+                                    output, output_size, &out_vec[0].len);
+    }
+    break;
+    case TFM_CRYPTO_CIPHER_DECRYPT_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        uint8_t *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+
+        status = psa_cipher_decrypt(*encoded_key, iov->alg, input, input_length,
+                                    output, output_size, &out_vec[0].len);
+    }
+    break;
+    default:
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
 
     return status;
-
-exit:
+release_operation_and_return:
     /* Release the operation context, ignore if the operation fails. */
     (void)tfm_crypto_operation_release(handle_out);
     return status;
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_decrypt_setup(psa_invec in_vec[],
-                                             size_t in_len,
-                                             psa_outvec out_vec[],
-                                             size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((out_vec[0].len != sizeof(uint32_t)) ||
-        (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_CIPHER_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *handle_out = handle;
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
-    status = psa_cipher_decrypt_setup(operation, encoded_key, alg);
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
-    return status;
-
-exit:
-    /* Release the operation context, ignore if the operation fails. */
-    (void)tfm_crypto_operation_release(handle_out);
-    return status;
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_update(psa_invec in_vec[],
-                                      size_t in_len,
-                                      psa_outvec out_vec[],
-                                      size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    unsigned char *output = out_vec[1].base;
-    size_t output_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Initialise the output_length to zero */
-    out_vec[1].len = 0;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_CIPHER_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_cipher_update(operation, input, input_length,
-                             output, output_size, &out_vec[1].len);
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_finish(psa_invec in_vec[],
-                                      size_t in_len,
-                                      psa_outvec out_vec[],
-                                      size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    unsigned char *output = out_vec[1].base;
-    size_t output_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Initialise the output_length to zero */
-    out_vec[1].len = 0;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_CIPHER_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_cipher_finish(operation, output, output_size, &out_vec[1].len);
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
-    return status;
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_abort(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_cipher_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_CIPHER_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        /* Operation does not exist, so abort has no effect */
-        return PSA_SUCCESS;
-    }
-
-    status = psa_cipher_abort(operation);
-
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-        return status;
-    }
-
-    return tfm_crypto_operation_release(handle_out);
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_encrypt(psa_invec in_vec[],
-                                       size_t in_len,
-                                       psa_outvec out_vec[],
-                                       size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    uint8_t *output = out_vec[0].base;
-    size_t output_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_cipher_encrypt(encoded_key, alg, input, input_length, output,
-                              output_size, &out_vec[0].len);
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_cipher_decrypt(psa_invec in_vec[],
-                                       size_t in_len,
-                                       psa_outvec out_vec[],
-                                       size_t out_len)
-{
-#ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    uint8_t *output = out_vec[0].base;
-    size_t output_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_cipher_decrypt(encoded_key, alg, input, input_length, output,
-                              output_size, &out_vec[0].len);
-#endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_hash.c b/secure_fw/partitions/crypto/crypto_hash.c
index 6c2d27a..46a3244 100644
--- a/secure_fw/partitions/crypto/crypto_hash.c
+++ b/secure_fw/partitions/crypto/crypto_hash.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -12,326 +12,150 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
 /*!
- * \defgroup public_psa Public functions, PSA
+ * \addtogroup tfm_crypto_api_shim_layer
  *
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_hash_setup(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
+psa_status_t tfm_crypto_hash_interface(psa_invec in_vec[],
+                                       psa_outvec out_vec[])
 {
 #ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_hash_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((out_vec[0].len != sizeof(uint32_t)) ||
-        (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
+#endif
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_hash_operation_t *operation = NULL;
     uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_algorithm_t alg = iov->alg;
+    uint32_t *handle_out = NULL;
+    /* Used to save the target_operation during cloning */
+    uint32_t handle_clone = 0;
+    enum tfm_crypto_function_type function_type =
+                        TFM_CRYPTO_GET_FUNCTION_TYPE(iov->function_id);
 
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_HASH_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
+    if (function_type != TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART) {
+        handle_out = (out_vec && out_vec[0].base != NULL) ?
+                                                out_vec[0].base : &handle;
+        *handle_out = handle;
+        status = tfm_crypto_operation_handling(TFM_CRYPTO_HASH_OPERATION,
+                                               function_type,
+                                               handle_out,
+                                               (void **)&operation);
+        if (status != PSA_SUCCESS) {
+            return (iov->function_id == TFM_CRYPTO_HASH_ABORT_SID) ?
+                    PSA_SUCCESS : status;
+        }
     }
 
-    *handle_out = handle;
+    switch (iov->function_id) {
+    case TFM_CRYPTO_HASH_SETUP_SID:
+    {
+        status = psa_hash_setup(operation, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_HASH_UPDATE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
 
-    status = psa_hash_setup(operation, alg);
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
+        status = psa_hash_update(operation, input, input_length);
+    }
+    break;
+    case TFM_CRYPTO_HASH_FINISH_SID:
+    {
+        uint8_t *hash = out_vec[1].base;
+        size_t hash_size = out_vec[1].len;
+        /* Initialise the output_length to zero */
+        out_vec[1].len = 0;
+
+        status = psa_hash_finish(operation, hash, hash_size, &out_vec[1].len);
+        if (status == PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_HASH_VERIFY_SID:
+    {
+        const uint8_t *hash = in_vec[1].base;
+        size_t hash_length = in_vec[1].len;
+
+        status = psa_hash_verify(operation, hash, hash_length);
+        if (status == PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_HASH_ABORT_SID:
+    {
+        status = psa_hash_abort(operation);
+
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        } else {
+            status = tfm_crypto_operation_release(handle_out);
+        }
+    }
+    break;
+    case TFM_CRYPTO_HASH_CLONE_SID:
+    {
+        psa_hash_operation_t *target_operation = NULL;
+
+        /* handle_out is initialised here for the hash clone as a special
+         * case which don't rely on the standard flow for handling
+         * multipart operation contexts
+         */
+        *handle_out = handle_clone;
+
+        /* Allocate the target operation context in the secure world */
+        status = tfm_crypto_operation_alloc(TFM_CRYPTO_HASH_OPERATION,
+                                            handle_out,
+                                            (void **)&target_operation);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+        status = psa_hash_clone(operation, target_operation);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_HASH_COMPUTE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        uint8_t *hash = out_vec[0].base;
+        size_t hash_size = out_vec[0].len;
+        /* Initialize hash_length to zero */
+        out_vec[0].len = 0;
+
+        status = psa_hash_compute(iov->alg, input, input_length,
+                                  hash, hash_size, &out_vec[0].len);
+    }
+    break;
+    case TFM_CRYPTO_HASH_COMPARE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        const uint8_t *hash = in_vec[2].base;
+        size_t hash_length = in_vec[2].len;
+
+        status = psa_hash_compare(iov->alg, input, input_length,
+                                  hash, hash_length);
+    }
+    break;
+    default:
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
 
     return status;
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_update(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_hash_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_HASH_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_hash_update(operation, input, input_length);
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_finish(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_hash_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    uint8_t *hash = out_vec[1].base;
-    size_t hash_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Initialise hash_length to zero */
-    out_vec[1].len = 0;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_HASH_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_hash_finish(operation, hash, hash_size, &out_vec[1].len);
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
+release_operation_and_return:
+    /* Release the operation context, ignore if the release fails. */
+    (void)tfm_crypto_operation_release(handle_out);
     return status;
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_verify(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_hash_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *hash = in_vec[1].base;
-    size_t hash_length = in_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_HASH_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_hash_verify(operation, hash, hash_length);
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
-    return status;
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_abort(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_hash_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_HASH_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        /* Operation does not exist, so abort has no effect */
-        return PSA_SUCCESS;
-    }
-
-    status = psa_hash_abort(operation);
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-        return status;
-    }
-
-    return tfm_crypto_operation_release(handle_out);
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_clone(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_hash_operation_t *source_operation = NULL;
-    psa_hash_operation_t *target_operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t source_handle = iov->op_handle;
-    uint32_t *target_handle = out_vec[0].base;
-
-    /* Look up the corresponding source operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_HASH_OPERATION,
-                                         source_handle,
-                                         (void **)&source_operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    /* Allocate the target operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_HASH_OPERATION,
-                                        target_handle,
-                                        (void **)&target_operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_hash_clone(source_operation, target_operation);
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_compute(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-         return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    uint8_t *hash = out_vec[0].base;
-    size_t hash_size = out_vec[0].len;
-
-    /* Initialize hash_length to zero */
-    out_vec[0].len = 0;
-    return psa_hash_compute(alg, input, input_length, hash, hash_size,
-                            &out_vec[0].len);
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_hash_compare(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
-{
-#ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-         return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    const uint8_t *hash = in_vec[2].base;
-    size_t hash_length = in_vec[2].len;
-
-    return psa_hash_compare(alg, input, input_length, hash, hash_length);
-#endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_init.c b/secure_fw/partitions/crypto/crypto_init.c
index 476bea5..077fceb 100644
--- a/secure_fw/partitions/crypto/crypto_init.c
+++ b/secure_fw/partitions/crypto/crypto_init.c
@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-3-Clause
  *
  */
+#include <stdbool.h>
 
 #include "tfm_mbedcrypto_include.h"
 
@@ -13,7 +14,7 @@
 
 /*
  * \brief This Mbed TLS include is needed to initialise the memory allocator
- *        inside the Mbed TLS layer of Mbed Crypto
+ *        of the library used for internal allocations
  */
 #include "mbedtls/memory_buffer_alloc.h"
 
@@ -29,6 +30,40 @@
 #include "crypto_hw.h"
 #endif /* CRYPTO_HW_ACCELERATOR */
 
+#ifndef MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
+#error "MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER must be selected in Mbed TLS config file"
+#endif
+
+/**
+ * \brief Type describing the properties of each function identifier, i.e. the
+ *        group ID and the function type
+ */
+struct tfm_crypto_api_descriptor {
+    uint8_t group_id     : 6; /*!< Value from \ref tfm_crypto_group_id */
+    uint8_t function_type: 2; /*!< Value from \ref tfm_crypto_function_type */
+};
+
+/**
+ * \brief This table contains the description of each of the function IDs
+ *        defined by \ref tfm_crypto_function_id
+ */
+#define X(_function_id, _group_id, _function_type) \
+    [_function_id] = {.group_id = _group_id, .function_type = _function_type},
+static const struct tfm_crypto_api_descriptor tfm_crypto_api_descriptor[] = {
+    TFM_CRYPTO_SERVICE_API_DESCRIPTION
+};
+#undef X
+
+enum tfm_crypto_function_type get_function_type_from_descriptor(enum tfm_crypto_function_id func)
+{
+    return tfm_crypto_api_descriptor[func].function_type;
+}
+
+enum tfm_crypto_group_id get_group_id_from_descriptor(enum tfm_crypto_function_id func)
+{
+    return tfm_crypto_api_descriptor[func].group_id;
+}
+
 #ifdef TFM_PSA_API
 #include <string.h>
 #include "psa/framework_feature.h"
@@ -36,16 +71,6 @@
 #include "psa_manifest/tfm_crypto.h"
 
 /**
- * \brief Table containing all the Uniform Signature API exposed
- *        by the TF-M Crypto partition
- */
-static const tfm_crypto_us_t sfid_func_table[TFM_CRYPTO_SID_MAX] = {
-#define X(api_name) api_name,
-LIST_TFM_CRYPTO_UNIFORM_SIGNATURE_API
-#undef X
-};
-
-/**
  * \brief Aligns a value x up to an alignment a.
  */
 #define ALIGN(x, a) (((x) + ((a) - 1)) & ~((a) - 1))
@@ -235,10 +260,6 @@
         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].len = sizeof(struct tfm_crypto_pack_iovec);
@@ -250,8 +271,8 @@
 
     tfm_crypto_set_caller_id(msg->client_id);
 
-    /* Call the uniform signature API */
-    status = sfid_func_table[iov.srv_id](in_vec, in_len, out_vec, out_len);
+    /* Call the dispatcher to the functions that implement the PSA Crypto API */
+    status = tfm_crypto_api_dispatcher(in_vec, in_len, out_vec, out_len);
 
 #if PSA_FRAMEWORK_HAS_MM_IOVEC == 1
     for (i = 0; i < out_len; i++) {
@@ -301,33 +322,34 @@
 
 static psa_status_t tfm_crypto_engine_init(void)
 {
-
 #ifdef CRYPTO_NV_SEED
+    LOG_INFFMT("[INF][Crypto] ");
 #ifdef TFM_PSA_API
+    LOG_INFFMT("Provisioning entropy seed... ");
     if (tfm_plat_crypto_provision_entropy_seed() != TFM_CRYPTO_NV_SEED_SUCCESS) {
         return PSA_ERROR_GENERIC_ERROR;
     }
+    LOG_INFFMT("\033[0;32mcomplete.\033[0m\r\n");
 #else
-    LOG_INFFMT("\033[1;31m[Crypto] ");
     LOG_INFFMT("TF-M in library mode uses a dummy NV seed. ");
     LOG_INFFMT("This is not suitable for production! ");
-    LOG_INFFMT("This device is \033[1;1mNOT SECURE");
-    LOG_INFFMT("\033[0m\r\n");
+    LOG_INFFMT("This device is \033[1;31mNOT SECURE\033[0m\r\n");
 #endif /* TFM_PSA_API */
 #endif /* CRYPTO_NV_SEED */
 
-    /* Initialise the Mbed Crypto memory allocator to use static
-     * memory allocation from the provided buffer instead of using
-     * the heap
+    /* Initialise the Mbed Crypto memory allocator to use static memory
+     * allocation from the provided buffer instead of using the heap
      */
     mbedtls_memory_buffer_alloc_init(mbedtls_mem_buf,
                                      TFM_CRYPTO_ENGINE_BUF_SIZE);
 
     /* Initialise the crypto accelerator if one is enabled */
 #ifdef CRYPTO_HW_ACCELERATOR
+    LOG_INFFMT("[INF][Crypto] Initialising HW accelerator... ");
     if (crypto_hw_accelerator_init() != 0) {
         return PSA_ERROR_HARDWARE_FAILURE;
     }
+    LOG_INFFMT("\033[0;32mcomplete.\033[0m\r\n");
 #endif /* CRYPTO_HW_ACCELERATOR */
 
     /* Previous function does not return any value, so just call the
@@ -370,3 +392,79 @@
     return PSA_ERROR_GENERIC_ERROR;
 }
 #endif
+
+psa_status_t tfm_crypto_api_dispatcher(psa_invec in_vec[],
+                                       size_t in_len,
+                                       psa_outvec out_vec[],
+                                       size_t out_len)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    int32_t caller_id = 0;
+    mbedtls_svc_key_id_t encoded_key = MBEDTLS_SVC_KEY_ID_INIT;
+    bool is_key_required = false;
+
+    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
+        return PSA_ERROR_PROGRAMMER_ERROR;
+    }
+
+    is_key_required = !(TFM_CRYPTO_IS_GROUP_ID(
+                            iov->function_id, TFM_CRYPTO_GROUP_ID_HASH) ||
+                        (iov->function_id == TFM_CRYPTO_GENERATE_RANDOM_SID));
+
+    if (is_key_required) {
+        status = tfm_crypto_get_caller_id(&caller_id);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+        /* The caller_id being set in the owner field is the partition ID
+         * of the calling partition
+         */
+        encoded_key = mbedtls_svc_key_id_make(caller_id, iov->key_id);
+    }
+
+    /* Dispatch to each sub-module based on the Group ID */
+    if (TFM_CRYPTO_IS_GROUP_ID(
+            iov->function_id, TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT)) {
+        status = tfm_crypto_key_management_interface(in_vec,
+                                                     out_vec,
+                                                     &encoded_key);
+    } else if (TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_HASH)) {
+        status = tfm_crypto_hash_interface(in_vec, out_vec);
+    } else if (TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_MAC)) {
+        status = tfm_crypto_mac_interface(in_vec,
+                                          out_vec,
+                                          &encoded_key);
+    } else if (TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_CIPHER)) {
+        status = tfm_crypto_cipher_interface(in_vec,
+                                             out_vec,
+                                             &encoded_key);
+    } else if (TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_AEAD)) {
+        status = tfm_crypto_aead_interface(in_vec,
+                                           out_vec,
+                                           &encoded_key);
+    } else if (TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_ASYM_SIGN) ||
+               TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_ASYM_ENCRYPT)) {
+        status = tfm_crypto_asymmetric_interface(in_vec,
+                                                 out_vec,
+                                                 &encoded_key);
+    } else if (TFM_CRYPTO_IS_GROUP_ID(
+                   iov->function_id, TFM_CRYPTO_GROUP_ID_KEY_DERIVATION)) {
+        status = tfm_crypto_key_derivation_interface(in_vec,
+                                                     out_vec,
+                                                     &encoded_key);
+    } else if (iov->function_id == TFM_CRYPTO_GENERATE_RANDOM_SID) {
+        status = tfm_crypto_random_interface(in_vec, out_vec);
+    } else {
+        LOG_ERRFMT("[ERR][Crypto] Unsupported request!\r\n");
+        status = PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    return status;
+}
diff --git a/secure_fw/partitions/crypto/crypto_key.c b/secure_fw/partitions/crypto/crypto_key.c
index b731f00..ac36339 100644
--- a/secure_fw/partitions/crypto/crypto_key.c
+++ b/secure_fw/partitions/crypto/crypto_key.c
@@ -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
  *
@@ -12,14 +12,14 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
 #ifndef MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
 #error "MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER must be selected in Mbed TLS config file"
 #endif
 
 /*!
- * \defgroup public Public functions
+ * \defgroup key_helpers Helper functions to translate between client and server
+ *                       view of the key attributes structure.
  *
  */
 /*!@{*/
@@ -77,25 +77,4 @@
 
     return PSA_SUCCESS;
 }
-
-psa_status_t tfm_crypto_encode_id_and_owner(psa_key_id_t key_id,
-                                            mbedtls_svc_key_id_t *enc_key_ptr)
-{
-    int32_t partition_id = 0;
-    psa_status_t status = tfm_crypto_get_caller_id(&partition_id);
-
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    /* If Null Pointer, return PSA_ERROR_PROGRAMMER_ERROR */
-    if (enc_key_ptr == NULL) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    /* Use the client key id as the key_id and its partition id as the owner */
-    *enc_key_ptr = mbedtls_svc_key_id_make(partition_id, key_id);
-
-    return PSA_SUCCESS;
-}
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_key_derivation.c b/secure_fw/partitions/crypto/crypto_key_derivation.c
index 993477d..50e12b7 100644
--- a/secure_fw/partitions/crypto/crypto_key_derivation.c
+++ b/secure_fw/partitions/crypto/crypto_key_derivation.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
  *
@@ -8,6 +8,7 @@
 #include <stddef.h>
 #include <stdint.h>
 #include <string.h>
+#include "tfm_sp_log.h"
 
 #include "tfm_mbedcrypto_include.h"
 
@@ -18,7 +19,10 @@
 #include "tfm_crypto_defs.h"
 
 #include "tfm_plat_crypto_keys.h"
-#include "tfm_crypto_private.h"
+
+#ifndef MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
+#error "MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER must be selected in Mbed TLS config file"
+#endif
 
 #ifdef TFM_PARTITION_TEST_PS
 #include "psa_manifest/pid.h"
@@ -98,6 +102,11 @@
     }
 
     /* Derive key material from the HUK and output it to the operation buffer */
+    /* The function below accesses an abstracted functionality at platform level
+     * to get a key derived from the HUK. The specific algorithm is abstracted
+     * away to the platform implementation (i.e. it can be HMAC-based KDF or
+     * CMAC based KDF, the latter for example when CC-312 is present)
+     */
     err = tfm_plat_get_huk_derived_key(tls12_prf->MBEDTLS_PRIVATE(label),
                                        tls12_prf->MBEDTLS_PRIVATE(label_length),
                                        NULL, 0,
@@ -106,6 +115,7 @@
     if (err != TFM_PLAT_ERR_SUCCESS) {
         return PSA_ERROR_HARDWARE_FAILURE;
     }
+    LOG_DBGFMT("\033[0;32mcomplete.\033[0m\r\n");
 
     return psa_import_key(attributes, tls12_prf->MBEDTLS_PRIVATE(output_block),
                           bytes, key_id);
@@ -130,430 +140,169 @@
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 
 /*!
- * \defgroup public_psa Public functions, PSA
+ * \addtogroup tfm_crypto_api_shim_layer
  *
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_key_derivation_setup(psa_invec in_vec[],
-                                             size_t in_len,
-                                             psa_outvec out_vec[],
-                                             size_t out_len)
+psa_status_t tfm_crypto_key_derivation_interface(psa_invec in_vec[],
+                                            psa_outvec out_vec[],
+                                            mbedtls_svc_key_id_t *encoded_key)
 {
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_key_derivation_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((out_vec[0].len != sizeof(uint32_t)) ||
-        (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
+#endif
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_derivation_operation_t *operation = NULL;
     uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_algorithm_t alg = iov->alg;
+    uint32_t *handle_out = NULL;
+    enum tfm_crypto_function_type function_type =
+                        TFM_CRYPTO_GET_FUNCTION_TYPE(iov->function_id);
 
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                        &handle,
+    if (function_type != TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART) {
+        handle_out = (out_vec && out_vec[0].base != NULL) ?
+                                                out_vec[0].base : &handle;
+        *handle_out = handle;
+        status = tfm_crypto_operation_handling(
+                                        TFM_CRYPTO_KEY_DERIVATION_OPERATION,
+                                        function_type,
+                                        handle_out,
                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
+        if (status != PSA_SUCCESS) {
+            return (iov->function_id == TFM_CRYPTO_KEY_DERIVATION_ABORT_SID) ?
+                    PSA_SUCCESS : status;
+        }
     }
 
-    *handle_out = handle;
+    switch (iov->function_id) {
+    case TFM_CRYPTO_KEY_DERIVATION_SETUP_SID:
+    {
+        if (iov->alg == TFM_CRYPTO_ALG_HUK_DERIVATION) {
+            status = tfm_crypto_huk_derivation_setup(operation, iov->alg);
+        } else {
+            status = psa_key_derivation_setup(operation, iov->alg);
+        }
 
-    if (alg == TFM_CRYPTO_ALG_HUK_DERIVATION) {
-        status = tfm_crypto_huk_derivation_setup(operation, alg);
-    } else {
-        status = psa_key_derivation_setup(operation, alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
     }
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-        return status;
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID:
+    {
+        size_t *capacity = out_vec[0].base;
+
+        status = psa_key_derivation_get_capacity(operation, capacity);
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID:
+    {
+        status = psa_key_derivation_set_capacity(operation, iov->capacity);
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID:
+    {
+        const uint8_t *data = in_vec[1].base;
+        size_t data_length = in_vec[1].len;
+
+        if (operation->MBEDTLS_PRIVATE(alg) == TFM_CRYPTO_ALG_HUK_DERIVATION) {
+            status = tfm_crypto_huk_derivation_input_bytes(operation,
+                                                           iov->step, data,
+                                                           data_length);
+        } else {
+            status = psa_key_derivation_input_bytes(operation, iov->step, data,
+                                                    data_length);
+        }
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID:
+    {
+        uint8_t *output = out_vec[0].base;
+        size_t output_length = out_vec[0].len;
+
+        status = psa_key_derivation_output_bytes(operation,
+                                                 output, output_length);
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID:
+    {
+         status = psa_key_derivation_input_key(operation,
+                                               iov->step, *encoded_key);
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID:
+    {
+        const struct psa_client_key_attributes_s *client_key_attr =
+                                                            in_vec[1].base;
+        psa_key_id_t *key_handle = out_vec[0].base;
+        psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+        int32_t partition_id = MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(*encoded_key);
+
+        status = tfm_crypto_key_attributes_from_client(client_key_attr,
+                                                       partition_id,
+                                                       &key_attributes);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+
+        if (operation->MBEDTLS_PRIVATE(alg) == TFM_CRYPTO_ALG_HUK_DERIVATION) {
+            status = tfm_crypto_huk_derivation_output_key(&key_attributes,
+                                                    operation, encoded_key);
+        } else {
+            status = psa_key_derivation_output_key(&key_attributes, operation,
+                                                   encoded_key);
+        }
+
+        *key_handle = encoded_key->MBEDTLS_PRIVATE(key_id);
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_ABORT_SID:
+    {
+        if (operation->MBEDTLS_PRIVATE(alg) == TFM_CRYPTO_ALG_HUK_DERIVATION) {
+            status = tfm_crypto_huk_derivation_abort(operation);
+        } else {
+            status = psa_key_derivation_abort(operation);
+        }
+
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        } else {
+            status = tfm_crypto_operation_release(handle_out);
+        }
+    }
+    break;
+    case TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID:
+    {
+        const uint8_t *peer_key = in_vec[1].base;
+        size_t peer_key_length = in_vec[1].len;
+
+        status = psa_key_derivation_key_agreement(operation, iov->step,
+                                                  *encoded_key,
+                                                  peer_key,
+                                                  peer_key_length);
+    }
+    break;
+    case TFM_CRYPTO_RAW_KEY_AGREEMENT_SID:
+    {
+        uint8_t *output = out_vec[0].base;
+        size_t output_size = out_vec[0].len;
+        const uint8_t *peer_key = in_vec[1].base;
+        size_t peer_key_length = in_vec[1].len;
+
+        status = psa_raw_key_agreement(iov->alg, *encoded_key,
+                                       peer_key, peer_key_length,
+                                       output, output_size, &out_vec[0].len);
+    }
+    break;
+    default:
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
 
     return status;
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_get_capacity(psa_invec in_vec[],
-                                                    size_t in_len,
-                                                    psa_outvec out_vec[],
-                                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(size_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    size_t *capacity = out_vec[0].base;
-    psa_key_derivation_operation_t *operation = NULL;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        *capacity = 0;
-        return status;
-    }
-
-    return psa_key_derivation_get_capacity(operation, capacity);
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_set_capacity(psa_invec in_vec[],
-                                                    size_t in_len,
-                                                    psa_outvec out_vec[],
-                                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    size_t capacity = iov->capacity;
-    psa_key_derivation_operation_t *operation = NULL;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_key_derivation_set_capacity(operation, capacity);
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_input_bytes(psa_invec in_vec[],
-                                                   size_t in_len,
-                                                   psa_outvec out_vec[],
-                                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    psa_key_derivation_step_t step = iov->step;
-    const uint8_t *data = in_vec[1].base;
-    size_t data_length = in_vec[1].len;
-    psa_key_derivation_operation_t *operation = NULL;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    if (operation->MBEDTLS_PRIVATE(alg) == TFM_CRYPTO_ALG_HUK_DERIVATION) {
-        return tfm_crypto_huk_derivation_input_bytes(operation, step, data,
-                                                     data_length);
-    } else {
-        return psa_key_derivation_input_bytes(operation, step, data,
-                                              data_length);
-    }
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_output_bytes(psa_invec in_vec[],
-                                                    size_t in_len,
-                                                    psa_outvec out_vec[],
-                                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    uint8_t *output = out_vec[0].base;
-    size_t output_length = out_vec[0].len;
-    psa_key_derivation_operation_t *operation = NULL;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_key_derivation_output_bytes(operation, output, output_length);
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_input_key(psa_invec in_vec[],
-                                                 size_t in_len,
-                                                 psa_outvec out_vec[],
-                                                 size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    psa_key_id_t key_id = iov->key_id;
-    psa_key_derivation_step_t step = iov->step;
-    psa_key_derivation_operation_t *operation = NULL;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_key_derivation_input_key(operation, step, encoded_key);
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_output_key(psa_invec in_vec[],
-                                                  size_t in_len,
-                                                  psa_outvec out_vec[],
-                                                  size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (in_vec[1].len != sizeof(struct psa_client_key_attributes_s)) ||
-        (out_vec[0].len != sizeof(psa_key_id_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
-    psa_key_derivation_operation_t *operation = NULL;
-    psa_key_id_t *key_handle = out_vec[0].base;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    int32_t partition_id;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_key_attributes_from_client(client_key_attr,
-                                                   partition_id,
-                                                   &key_attributes);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    if (operation->MBEDTLS_PRIVATE(alg) == TFM_CRYPTO_ALG_HUK_DERIVATION) {
-        status = tfm_crypto_huk_derivation_output_key(&key_attributes,
-                                                      operation, &encoded_key);
-    } else {
-        status = psa_key_derivation_output_key(&key_attributes, operation,
-                                               &encoded_key);
-    }
-
-    *key_handle = encoded_key.MBEDTLS_PRIVATE(key_id);
-
+release_operation_and_return:
+    /* Release the operation context, ignore if the operation fails. */
+    (void)tfm_crypto_operation_release(handle_out);
     return status;
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_abort(psa_invec in_vec[],
-                                             size_t in_len,
-                                             psa_outvec out_vec[],
-                                             size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_derivation_operation_t *operation = NULL;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        /* Operation does not exist, so abort has no effect */
-        return PSA_SUCCESS;
-    }
-
-    *handle_out = handle;
-
-    if (operation->MBEDTLS_PRIVATE(alg) == TFM_CRYPTO_ALG_HUK_DERIVATION) {
-        status = tfm_crypto_huk_derivation_abort(operation);
-    } else {
-        status = psa_key_derivation_abort(operation);
-    }
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-        return status;
-    }
-
-    status = tfm_crypto_operation_release(handle_out);
-
-    return status;
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_key_derivation_key_agreement(psa_invec in_vec[],
-                                                     size_t in_len,
-                                                     psa_outvec out_vec[],
-                                                     size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 0, 0);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    uint32_t handle = iov->op_handle;
-    psa_key_id_t private_key = iov->key_id;
-    const uint8_t *peer_key = in_vec[1].base;
-    size_t peer_key_length = in_vec[1].len;
-    psa_key_derivation_operation_t *operation = NULL;
-    psa_key_derivation_step_t step = iov->step;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_KEY_DERIVATION_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_encode_id_and_owner(private_key, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_key_derivation_key_agreement(operation, step,
-                                            encoded_key,
-                                            peer_key,
-                                            peer_key_length);
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_raw_key_agreement(psa_invec in_vec[],
-                                          size_t in_len,
-                                          psa_outvec out_vec[],
-                                          size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 0, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint8_t *output = out_vec[0].base;
-    size_t output_size = out_vec[0].len;
-    psa_algorithm_t alg = iov->alg;
-    psa_key_id_t private_key = iov->key_id;
-    const uint8_t *peer_key = in_vec[1].base;
-    size_t peer_key_length = in_vec[1].len;
-    mbedtls_svc_key_id_t encoded_key;
-    psa_status_t status;
-
-    status = tfm_crypto_encode_id_and_owner(private_key, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_raw_key_agreement(alg, encoded_key, peer_key, peer_key_length,
-                                 output, output_size, &out_vec[0].len);
-#endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_key_management.c b/secure_fw/partitions/crypto/crypto_key_management.c
index f545c07..0054a1a 100644
--- a/secure_fw/partitions/crypto/crypto_key_management.c
+++ b/secure_fw/partitions/crypto/crypto_key_management.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
  *
@@ -11,438 +11,171 @@
 #include "tfm_mbedcrypto_include.h"
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
-psa_status_t tfm_crypto_import_key(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
+#ifndef MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
+#error "MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER must be selected in Mbed TLS config file"
+#endif
+
+/*!
+ * \addtogroup tfm_crypto_api_shim_layer
+ *
+ */
+
+/*!@{*/
+psa_status_t tfm_crypto_key_management_interface(psa_invec in_vec[],
+                                            psa_outvec out_vec[],
+                                            mbedtls_svc_key_id_t *encoded_key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 3, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (in_vec[1].len != sizeof(struct psa_client_key_attributes_s)) ||
-        (out_vec[0].len != sizeof(psa_key_id_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
-    const uint8_t *data = in_vec[2].base;
-    size_t data_length = in_vec[2].len;
-    psa_key_id_t *psa_key = out_vec[0].base;
-
-    psa_status_t status;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_key_attributes_from_client(client_key_attr,
-                                                   partition_id,
-                                                   &key_attributes);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_import_key(&key_attributes, data, data_length, &encoded_key);
-
-    /* Update the imported key id */
-    *psa_key = encoded_key.MBEDTLS_PRIVATE(key_id);
-
-    return status;
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_open_key(psa_invec in_vec[],
-                                 size_t in_len,
-                                 psa_outvec out_vec[],
-                                 size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (in_vec[1].len != sizeof(psa_key_id_t)) ||
-        (out_vec[0].len != sizeof(psa_key_id_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    psa_key_id_t client_key_id = *((psa_key_id_t *)in_vec[1].base);
-    psa_key_id_t *key = out_vec[0].base;
-    psa_status_t status;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    /* Use the client key id as the key_id and its partition id as the owner */
-    encoded_key = mbedtls_svc_key_id_make(partition_id, client_key_id);
-
-    status = psa_open_key(encoded_key, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *key = encoded_key.MBEDTLS_PRIVATE(key_id);
-
-    return status;
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_close_key(psa_invec in_vec[],
-                                  size_t in_len,
-                                  psa_outvec out_vec[],
-                                  size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    (void)out_vec;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 0);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
+#endif
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    int32_t partition_id = MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(*encoded_key);
 
-    psa_key_id_t key = iov->key_id;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-    psa_status_t status;
+    switch (iov->function_id) {
+    case TFM_CRYPTO_IMPORT_KEY_SID:
+    {
+        const struct psa_client_key_attributes_s *client_key_attr =
+                                                                 in_vec[1].base;
+        const uint8_t *data = in_vec[2].base;
+        size_t data_length = in_vec[2].len;
+        psa_key_id_t *key_id = out_vec[0].base;
+        psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
+        status = tfm_crypto_key_attributes_from_client(client_key_attr,
+                                                       partition_id,
+                                                       &key_attributes);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+
+        status = psa_import_key(&key_attributes,
+                                data, data_length, encoded_key);
+        /* Update the imported key id */
+        *key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID(*encoded_key);
     }
+    break;
+    case TFM_CRYPTO_OPEN_KEY_SID:
+    {
+        psa_key_id_t *key_id = out_vec[0].base;
 
-    encoded_key = mbedtls_svc_key_id_make(partition_id, key);
-
-    return psa_close_key(encoded_key);
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_destroy_key(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    (void)out_vec;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 0);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
+        status = psa_open_key(*encoded_key, encoded_key);
+        *key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID(*encoded_key);
     }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key = iov->key_id;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-    psa_status_t status;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
+    break;
+    case TFM_CRYPTO_CLOSE_KEY_SID:
+    {
+        status = psa_close_key(*encoded_key);
     }
-
-    encoded_key = mbedtls_svc_key_id_make(partition_id, key);
-
-    return psa_destroy_key(encoded_key);
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_get_key_attributes(psa_invec in_vec[],
-                                           size_t in_len,
-                                           psa_outvec out_vec[],
-                                           size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(struct psa_client_key_attributes_s))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
+    break;
+    case TFM_CRYPTO_DESTROY_KEY_SID:
+    {
+        status = psa_destroy_key(*encoded_key);
     }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    break;
+    case TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID:
+    {
+        psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+        struct psa_client_key_attributes_s *client_key_attr = out_vec[0].base;
 
-    psa_key_id_t key = iov->key_id;
-    struct psa_client_key_attributes_s *client_key_attr = out_vec[0].base;
-    psa_status_t status;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
+        status = psa_get_key_attributes(*encoded_key, &key_attributes);
+        if (status == PSA_SUCCESS) {
+            status = tfm_crypto_key_attributes_to_client(&key_attributes,
+                                                         client_key_attr);
+        }
     }
+    break;
+    case TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID:
+    {
+        psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+        struct psa_client_key_attributes_s *client_key_attr = out_vec[0].base;
 
-    encoded_key = mbedtls_svc_key_id_make(partition_id, key);
+        status = tfm_crypto_key_attributes_from_client(client_key_attr,
+                                                       partition_id,
+                                                       &key_attributes);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
 
-    status = psa_get_key_attributes(encoded_key, &key_attributes);
-    if (status == PSA_SUCCESS) {
+        psa_reset_key_attributes(&key_attributes);
+
         status = tfm_crypto_key_attributes_to_client(&key_attributes,
                                                      client_key_attr);
     }
+    break;
+    case TFM_CRYPTO_EXPORT_KEY_SID:
+    {
+        uint8_t *data = out_vec[0].base;
+        size_t data_size = out_vec[0].len;
+
+        status = psa_export_key(*encoded_key, data, data_size,
+                                &(out_vec[0].len));
+    }
+    break;
+    case TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID:
+    {
+        uint8_t *data = out_vec[0].base;
+        size_t data_size = out_vec[0].len;
+
+        status = psa_export_public_key(*encoded_key, data, data_size,
+                                       &(out_vec[0].len));
+    }
+    break;
+    case TFM_CRYPTO_PURGE_KEY_SID:
+    {
+        status = psa_purge_key(*encoded_key);
+    }
+    break;
+    case TFM_CRYPTO_COPY_KEY_SID:
+    {
+        psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+        const struct psa_client_key_attributes_s *client_key_attr =
+                                                                 in_vec[1].base;
+        psa_key_id_t *target_key_id = out_vec[0].base;
+        mbedtls_svc_key_id_t target_key = MBEDTLS_SVC_KEY_ID_INIT;
+
+        status = tfm_crypto_key_attributes_from_client(client_key_attr,
+                                                       partition_id,
+                                                       &key_attributes);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+
+        status = psa_copy_key(*encoded_key, &key_attributes, &target_key);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+
+        *target_key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID(target_key);
+    }
+    break;
+    case TFM_CRYPTO_GENERATE_KEY_SID:
+    {
+        psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+        const struct psa_client_key_attributes_s *client_key_attr =
+                                                                 in_vec[1].base;
+        psa_key_id_t *key_handle = out_vec[0].base;
+
+        status = tfm_crypto_key_attributes_from_client(client_key_attr,
+                                                       partition_id,
+                                                       &key_attributes);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+
+        status = psa_generate_key(&key_attributes, encoded_key);
+        if (status != PSA_SUCCESS) {
+            return status;
+        }
+
+        *key_handle = MBEDTLS_SVC_KEY_ID_GET_KEY_ID(*encoded_key);
+    }
+    break;
+    default:
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
 
     return status;
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
-
-psa_status_t tfm_crypto_reset_key_attributes(psa_invec in_vec[],
-                                             size_t in_len,
-                                             psa_outvec out_vec[],
-                                             size_t out_len)
-{
-#if (TFM_CRYPTO_KEY_MODULE_DISABLED != 0)
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(struct psa_client_key_attributes_s))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-
-    struct psa_client_key_attributes_s *client_key_attr = out_vec[0].base;
-    psa_status_t status;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    int32_t partition_id;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_key_attributes_from_client(client_key_attr,
-                                                   partition_id,
-                                                   &key_attributes);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    psa_reset_key_attributes(&key_attributes);
-
-    return tfm_crypto_key_attributes_to_client(&key_attributes,
-                                               client_key_attr);
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_export_key(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_id_t key = iov->key_id;
-    uint8_t *data = out_vec[0].base;
-    size_t data_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-    psa_status_t status;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    encoded_key = mbedtls_svc_key_id_make(partition_id, key);
-
-    return psa_export_key(encoded_key, data, data_size,
-                          &(out_vec[0].len));
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_export_public_key(psa_invec in_vec[],
-                                          size_t in_len,
-                                          psa_outvec out_vec[],
-                                          size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key = iov->key_id;
-    uint8_t *data = out_vec[0].base;
-    size_t data_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-    psa_status_t status;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    encoded_key = mbedtls_svc_key_id_make(partition_id, key);
-
-    return psa_export_public_key(encoded_key, data, data_size,
-                                 &(out_vec[0].len));
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_purge_key(psa_invec in_vec[],
-                                  size_t in_len,
-                                  psa_outvec out_vec[],
-                                  size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    (void)out_vec;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 0);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key = iov->key_id;
-    mbedtls_svc_key_id_t encoded_key;
-    int32_t partition_id = 0;
-    psa_status_t status;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    encoded_key = mbedtls_svc_key_id_make(partition_id, key);
-
-    return psa_purge_key(encoded_key);
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_copy_key(psa_invec in_vec[],
-                                 size_t in_len,
-                                 psa_outvec out_vec[],
-                                 size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(psa_key_id_t)) ||
-        (in_vec[1].len != sizeof(struct psa_client_key_attributes_s))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-
-    psa_key_id_t source_key_id = iov->key_id;
-    psa_key_id_t *target_key_id = out_vec[0].base;
-    const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
-    psa_status_t status;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    int32_t partition_id = 0;
-    mbedtls_svc_key_id_t target_key;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_key_attributes_from_client(client_key_attr,
-                                                   partition_id,
-                                                   &key_attributes);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    encoded_key = mbedtls_svc_key_id_make(partition_id, source_key_id);
-
-    status = psa_copy_key(encoded_key, &key_attributes, &target_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *target_key_id = target_key.MBEDTLS_PRIVATE(key_id);
-
-    return status;
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_generate_key(psa_invec in_vec[],
-                                     size_t in_len,
-                                     psa_outvec out_vec[],
-                                     size_t out_len)
-{
-#ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (in_vec[1].len != sizeof(struct psa_client_key_attributes_s)) ||
-        (out_vec[0].len != sizeof(psa_key_id_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    psa_key_id_t *key_handle = out_vec[0].base;
-    const struct psa_client_key_attributes_s *client_key_attr = in_vec[1].base;
-    psa_status_t status;
-    psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
-    int32_t partition_id = 0;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_get_caller_id(&partition_id);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = tfm_crypto_key_attributes_from_client(client_key_attr,
-                                                   partition_id,
-                                                   &key_attributes);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_generate_key(&key_attributes, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *key_handle = encoded_key.MBEDTLS_PRIVATE(key_id);
-
-    return status;
-#endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
-}
+/*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_mac.c b/secure_fw/partitions/crypto/crypto_mac.c
index 4c9ba2d..a86d7be 100644
--- a/secure_fw/partitions/crypto/crypto_mac.c
+++ b/secure_fw/partitions/crypto/crypto_mac.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
  *
@@ -12,365 +12,133 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
 
 /*!
- * \defgroup public_psa Public functions, PSA
+ * \addtogroup tfm_crypto_api_shim_layer
  *
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_mac_sign_setup(psa_invec in_vec[],
-                                       size_t in_len,
-                                       psa_outvec out_vec[],
-                                       size_t out_len)
+psa_status_t tfm_crypto_mac_interface(psa_invec in_vec[],
+                                      psa_outvec out_vec[],
+                                      mbedtls_svc_key_id_t *encoded_key)
 {
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_mac_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((out_vec[0].len != sizeof(uint32_t)) ||
-        (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
+#endif
     const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_mac_operation_t *operation = NULL;
     uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    mbedtls_svc_key_id_t encoded_key;
+    uint32_t *handle_out = NULL;
+    enum tfm_crypto_function_type function_type =
+                        TFM_CRYPTO_GET_FUNCTION_TYPE(iov->function_id);
 
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_MAC_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
+    if (function_type != TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART) {
+        handle_out = (out_vec && out_vec[0].base != NULL) ?
+                                                out_vec[0].base : &handle;
+        *handle_out = handle;
+        status = tfm_crypto_operation_handling(TFM_CRYPTO_MAC_OPERATION,
+                                               function_type,
+                                               handle_out,
+                                               (void **)&operation);
+        if (status != PSA_SUCCESS) {
+            return (iov->function_id == TFM_CRYPTO_MAC_ABORT_SID) ?
+                    PSA_SUCCESS : status;
+        }
     }
 
-    *handle_out = handle;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        goto exit;
+    switch (iov->function_id) {
+    case TFM_CRYPTO_MAC_SIGN_SETUP_SID:
+    {
+        status = psa_mac_sign_setup(operation, *encoded_key, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
     }
+    break;
+    case TFM_CRYPTO_MAC_VERIFY_SETUP_SID:
+    {
+        status = psa_mac_verify_setup(operation, *encoded_key, iov->alg);
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_MAC_UPDATE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
 
-    status = psa_mac_sign_setup(operation, encoded_key, alg);
-    if (status != PSA_SUCCESS) {
-        goto exit;
+        status = psa_mac_update(operation, input, input_length);
+    }
+    break;
+    case TFM_CRYPTO_MAC_SIGN_FINISH_SID:
+    {
+        uint8_t *mac = out_vec[1].base;
+        size_t mac_size = out_vec[1].len;
+        /* Initialise mac_length to zero */
+        out_vec[1].len = 0;
+
+        status = psa_mac_sign_finish(operation, mac, mac_size, &out_vec[1].len);
+        if (status == PSA_SUCCESS) {
+            /* In case of success automatically release the operation */
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_MAC_VERIFY_FINISH_SID:
+    {
+        const uint8_t *mac = in_vec[1].base;
+        size_t mac_length = in_vec[1].len;
+
+        status = psa_mac_verify_finish(operation, mac, mac_length);
+        if (status == PSA_SUCCESS) {
+            goto release_operation_and_return;
+        }
+    }
+    break;
+    case TFM_CRYPTO_MAC_ABORT_SID:
+    {
+        status = psa_mac_abort(operation);
+
+        if (status != PSA_SUCCESS) {
+            goto release_operation_and_return;
+        } else {
+            status = tfm_crypto_operation_release(handle_out);
+        }
+    }
+    break;
+    case TFM_CRYPTO_MAC_COMPUTE_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        uint8_t *mac = out_vec[0].base;
+        size_t mac_size = out_vec[0].len;
+
+        status = psa_mac_compute(*encoded_key, iov->alg, input, input_length,
+                                 mac, mac_size, &out_vec[0].len);
+    }
+    break;
+    case TFM_CRYPTO_MAC_VERIFY_SID:
+    {
+        const uint8_t *input = in_vec[1].base;
+        size_t input_length = in_vec[1].len;
+        const uint8_t *mac = in_vec[2].base;
+        size_t mac_length = in_vec[2].len;
+
+        status = psa_mac_verify(*encoded_key, iov->alg, input, input_length,
+                                mac, mac_length);
+    }
+    break;
+    default:
+        status = PSA_ERROR_NOT_SUPPORTED;
     }
 
     return status;
-
-exit:
+release_operation_and_return:
     /* Release the operation context, ignore if the operation fails. */
     (void)tfm_crypto_operation_release(handle_out);
     return status;
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_verify_setup(psa_invec in_vec[],
-                                         size_t in_len,
-                                         psa_outvec out_vec[],
-                                         size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_mac_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((out_vec[0].len != sizeof(uint32_t)) ||
-        (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    mbedtls_svc_key_id_t encoded_key;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Allocate the operation context in the secure world */
-    status = tfm_crypto_operation_alloc(TFM_CRYPTO_MAC_OPERATION,
-                                        &handle,
-                                        (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    *handle_out = handle;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
-    status = psa_mac_verify_setup(operation, encoded_key, alg);
-    if (status != PSA_SUCCESS) {
-        goto exit;
-    }
-
-    return status;
-
-exit:
-    /* Release the operation context, ignore if the operation fails. */
-    (void)tfm_crypto_operation_release(handle_out);
-    return status;
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_update(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_mac_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_MAC_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_mac_update(operation, input, input_length);
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_sign_finish(psa_invec in_vec[],
-                                        size_t in_len,
-                                        psa_outvec out_vec[],
-                                        size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_mac_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 2);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    uint8_t *mac = out_vec[1].base;
-    size_t mac_size = out_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Initialise mac_length to zero */
-    out_vec[1].len = 0;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_MAC_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_mac_sign_finish(operation, mac, mac_size, &out_vec[1].len);
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
-    return status;
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_verify_finish(psa_invec in_vec[],
-                                          size_t in_len,
-                                          psa_outvec out_vec[],
-                                          size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_mac_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-    const uint8_t *mac = in_vec[1].base;
-    size_t mac_length = in_vec[1].len;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_MAC_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    status = psa_mac_verify_finish(operation, mac, mac_length);
-    if (status == PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-    }
-
-    return status;
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_abort(psa_invec in_vec[],
-                                  size_t in_len,
-                                  psa_outvec out_vec[],
-                                  size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-    psa_mac_operation_t *operation = NULL;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 1, 1);
-
-    if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
-        (out_vec[0].len != sizeof(uint32_t))) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    uint32_t handle = iov->op_handle;
-    uint32_t *handle_out = out_vec[0].base;
-
-    /* Init the handle in the operation with the one passed from the iov */
-    *handle_out = iov->op_handle;
-
-    /* Look up the corresponding operation context */
-    status = tfm_crypto_operation_lookup(TFM_CRYPTO_MAC_OPERATION,
-                                         handle,
-                                         (void **)&operation);
-    if (status != PSA_SUCCESS) {
-        /* Operation does not exist, so abort has no effect */
-        return PSA_SUCCESS;
-    }
-
-    status = psa_mac_abort(operation);
-
-    if (status != PSA_SUCCESS) {
-        /* Release the operation context, ignore if the operation fails. */
-        (void)tfm_crypto_operation_release(handle_out);
-        return status;
-    }
-
-    return tfm_crypto_operation_release(handle_out);
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_compute(psa_invec in_vec[],
-                                    size_t in_len,
-                                    psa_outvec out_vec[],
-                                    size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 1, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    uint8_t *mac = out_vec[0].base;
-    size_t mac_size = out_vec[0].len;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_mac_compute(encoded_key, alg, input, input_length, mac, mac_size,
-                           &out_vec[0].len);
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
-}
-
-psa_status_t tfm_crypto_mac_verify(psa_invec in_vec[],
-                                   size_t in_len,
-                                   psa_outvec out_vec[],
-                                   size_t out_len)
-{
-#ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
-    return PSA_ERROR_NOT_SUPPORTED;
-#else
-    psa_status_t status = PSA_SUCCESS;
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 3, out_len, 0, 0);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
-    const struct tfm_crypto_pack_iovec *iov = in_vec[0].base;
-    psa_key_id_t key_id = iov->key_id;
-    psa_algorithm_t alg = iov->alg;
-    const uint8_t *input = in_vec[1].base;
-    size_t input_length = in_vec[1].len;
-    const uint8_t *mac = in_vec[2].base;
-    size_t mac_length = in_vec[2].len;
-    mbedtls_svc_key_id_t encoded_key;
-
-    status = tfm_crypto_encode_id_and_owner(key_id, &encoded_key);
-    if (status != PSA_SUCCESS) {
-        return status;
-    }
-
-    return psa_mac_verify(encoded_key, alg, input, input_length, mac,
-                          mac_length);
-#endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_rng.c b/secure_fw/partitions/crypto/crypto_rng.c
index 7f35f7f..1dccddf 100644
--- a/secure_fw/partitions/crypto/crypto_rng.c
+++ b/secure_fw/partitions/crypto/crypto_rng.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2019-2022, Arm Limited. All rights reserved.
  * Copyright (c) 2021, Nordic Semiconductor ASA.
  *
  * SPDX-License-Identifier: BSD-3-Clause
@@ -13,35 +13,22 @@
 
 #include "tfm_crypto_api.h"
 #include "tfm_crypto_defs.h"
-#include "tfm_crypto_private.h"
-
 
 /*!
- * \defgroup public_psa Public functions, PSA
+ * \addtogroup tfm_crypto_api_shim_layer
  *
  */
 
 /*!@{*/
-psa_status_t tfm_crypto_generate_random(psa_invec in_vec[],
-                                        size_t in_len,
-                                        psa_outvec out_vec[],
-                                        size_t out_len)
+psa_status_t tfm_crypto_random_interface(psa_invec in_vec[],
+                                         psa_outvec out_vec[])
 {
 #ifdef TFM_CRYPTO_RNG_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
-#else
-
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 1, out_len, 0, 1);
-
-    if (in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) {
-        return PSA_ERROR_PROGRAMMER_ERROR;
-    }
+#endif
     uint8_t *output = out_vec[0].base;
     size_t output_size = out_vec[0].len;
 
     return psa_generate_random(output, output_size);
-#endif /* TFM_CRYPTO_RNG_MODULE_DISABLED */
 }
-
-
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/tfm_crypto.yaml b/secure_fw/partitions/crypto/tfm_crypto.yaml
index 747a29e..8d922eb 100644
--- a/secure_fw/partitions/crypto/tfm_crypto.yaml
+++ b/secure_fw/partitions/crypto/tfm_crypto.yaml
@@ -15,460 +15,12 @@
   "stack_size": "0x1B00",
   "secure_functions": [
     {
-      "name": "TFM_CRYPTO_GET_KEY_ATTRIBUTES",
-      "signal": "TFM_CRYPTO_GET_KEY_ATTRIBUTES",
+      "name": "TFM_CRYPTO_API_DISPATCHER",
+      "signal": "TFM_CRYPTO_API_DISPATCHER",
       "non_secure_clients": true,
       "version": 1,
       "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_OPEN_KEY",
-      "signal": "TFM_CRYPTO_OPEN_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CLOSE_KEY",
-      "signal": "TFM_CRYPTO_CLOSE_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_RESET_KEY_ATTRIBUTES",
-      "signal": "TFM_CRYPTO_RESET_KEY_ATTRIBUTES",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_IMPORT_KEY",
-      "signal": "TFM_CRYPTO_IMPORT_KEY",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_DESTROY_KEY",
-      "signal": "TFM_CRYPTO_DESTROY_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_EXPORT_KEY",
-      "signal": "TFM_CRYPTO_EXPORT_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_EXPORT_PUBLIC_KEY",
-      "signal": "TFM_CRYPTO_EXPORT_PUBLIC_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_PURGE_KEY",
-      "signal": "TFM_CRYPTO_PURGE_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_COPY_KEY",
-      "signal": "TFM_CRYPTO_COPY_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_COMPUTE",
-      "signal": "TFM_CRYPTO_HASH_COMPUTE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_COMPARE",
-      "signal": "TFM_CRYPTO_HASH_COMPARE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_SETUP",
-      "signal": "TFM_CRYPTO_HASH_SETUP",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_UPDATE",
-      "signal": "TFM_CRYPTO_HASH_UPDATE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_FINISH",
-      "signal": "TFM_CRYPTO_HASH_FINISH",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_VERIFY",
-      "signal": "TFM_CRYPTO_HASH_VERIFY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_ABORT",
-      "signal": "TFM_CRYPTO_HASH_ABORT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_HASH_CLONE",
-      "signal": "TFM_CRYPTO_HASH_CLONE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_COMPUTE",
-      "signal": "TFM_CRYPTO_MAC_COMPUTE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_VERIFY",
-      "signal": "TFM_CRYPTO_MAC_VERIFY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_SIGN_SETUP",
-      "signal": "TFM_CRYPTO_MAC_SIGN_SETUP",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_VERIFY_SETUP",
-      "signal": "TFM_CRYPTO_MAC_VERIFY_SETUP",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_UPDATE",
-      "signal": "TFM_CRYPTO_MAC_UPDATE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_SIGN_FINISH",
-      "signal": "TFM_CRYPTO_MAC_SIGN_FINISH",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_VERIFY_FINISH",
-      "signal": "TFM_CRYPTO_MAC_VERIFY_FINISH",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_MAC_ABORT",
-      "signal": "TFM_CRYPTO_MAC_ABORT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_ENCRYPT",
-      "signal": "TFM_CRYPTO_CIPHER_ENCRYPT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_DECRYPT",
-      "signal": "TFM_CRYPTO_CIPHER_DECRYPT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_ENCRYPT_SETUP",
-      "signal": "TFM_CRYPTO_CIPHER_ENCRYPT_SETUP",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_DECRYPT_SETUP",
-      "signal": "TFM_CRYPTO_CIPHER_DECRYPT_SETUP",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_GENERATE_IV",
-      "signal": "TFM_CRYPTO_CIPHER_GENERATE_IV",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_SET_IV",
-      "signal": "TFM_CRYPTO_CIPHER_SET_IV",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_UPDATE",
-      "signal": "TFM_CRYPTO_CIPHER_UPDATE",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_FINISH",
-      "signal": "TFM_CRYPTO_CIPHER_FINISH",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_CIPHER_ABORT",
-      "signal": "TFM_CRYPTO_CIPHER_ABORT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_ENCRYPT",
-      "signal": "TFM_CRYPTO_AEAD_ENCRYPT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_DECRYPT",
-      "signal": "TFM_CRYPTO_AEAD_DECRYPT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_ENCRYPT_SETUP",
-      "signal": "TFM_CRYPTO_AEAD_ENCRYPT_SETUP",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_DECRYPT_SETUP",
-      "signal": "TFM_CRYPTO_AEAD_DECRYPT_SETUP",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_GENERATE_NONCE",
-      "signal": "TFM_CRYPTO_AEAD_GENERATE_NONCE",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_SET_NONCE",
-      "signal": "TFM_CRYPTO_AEAD_SET_NONCE",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_SET_LENGTHS",
-      "signal": "TFM_CRYPTO_AEAD_SET_LENGTHS",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_UPDATE_AD",
-      "signal": "TFM_CRYPTO_AEAD_UPDATE_AD",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_UPDATE",
-      "signal": "TFM_CRYPTO_AEAD_UPDATE",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_FINISH",
-      "signal": "TFM_CRYPTO_AEAD_FINISH",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_VERIFY",
-      "signal": "TFM_CRYPTO_AEAD_VERIFY",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_AEAD_ABORT",
-      "signal": "TFM_CRYPTO_AEAD_ABORT",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_SIGN_MESSAGE",
-      "signal": "TFM_CRYPTO_SIGN_MESSAGE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_VERIFY_MESSAGE",
-      "signal": "TFM_CRYPTO_VERIFY_MESSAGE",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_SIGN_HASH",
-      "signal": "TFM_CRYPTO_SIGN_HASH",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_VERIFY_HASH",
-      "signal": "TFM_CRYPTO_VERIFY_HASH",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_ASYMMETRIC_ENCRYPT",
-      "signal": "TFM_CRYPTO_ASYMMETRIC_ENCRYPT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_ASYMMETRIC_DECRYPT",
-      "signal": "TFM_CRYPTO_ASYMMETRIC_DECRYPT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_SETUP",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_SETUP",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES",
-      "non_secure_clients": true,
-      "minor_version": 1,
-      "minor_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_KEY_DERIVATION_ABORT",
-      "signal": "TFM_CRYPTO_KEY_DERIVATION_ABORT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_RAW_KEY_AGREEMENT",
-      "signal": "TFM_CRYPTO_RAW_KEY_AGREEMENT",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_GENERATE_RANDOM",
-      "signal": "TFM_CRYPTO_GENERATE_RANDOM",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_CRYPTO_GENERATE_KEY",
-      "signal": "TFM_CRYPTO_GENERATE_KEY",
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
+    }
   ],
   "services" : [
     {
diff --git a/secure_fw/partitions/crypto/tfm_crypto_api.h b/secure_fw/partitions/crypto/tfm_crypto_api.h
index 42c643b..ef26901 100644
--- a/secure_fw/partitions/crypto/tfm_crypto_api.h
+++ b/secure_fw/partitions/crypto/tfm_crypto_api.h
@@ -14,21 +14,8 @@
 
 #include <stdint.h>
 #include "tfm_crypto_defs.h"
-#ifdef TFM_PSA_API
-#include "psa/service.h"
-
-/**
- * \brief This define is a function pointer type to the Uniform Signature API
- *        prototype.
- */
-typedef psa_status_t (*tfm_crypto_us_t)(psa_invec[],size_t,psa_outvec[],size_t);
-#endif
-
 #include "psa/crypto_client_struct.h"
 
-#define UNIFORM_SIGNATURE_API(api_name) \
-    psa_status_t api_name(psa_invec[], size_t, psa_outvec[], size_t)
-
 /**
  * \brief List of possible operation types supported by the TFM based
  *        implementation. This type is needed by the operation allocation,
@@ -48,6 +35,68 @@
 };
 
 /**
+ * \brief Type associated to a function for the TF-M Crypto service
+ *        API. It describes if an API is multipart related or not, i.e. and
+ *        which multipart operation it requires (i.e. setup or lookup)
+ */
+enum tfm_crypto_function_type {
+    TFM_CRYPTO_FUNCTION_TYPE_NON_MULTIPART = 0x0,
+    TFM_CRYPTO_FUNCTION_TYPE_SETUP,
+    TFM_CRYPTO_FUNCTION_TYPE_LOOKUP
+};
+
+/**
+ * \brief Type associated to the group of a function encoding. There can be
+ *        nine groups (Random, Key management, Hash, MAC, Cipher, AEAD,
+ *        Asym sign, Asym encrypt, Key derivation).
+ */
+enum tfm_crypto_group_id {
+    TFM_CRYPTO_GROUP_ID_RANDOM = 0x0,
+    TFM_CRYPTO_GROUP_ID_KEY_MANAGEMENT,
+    TFM_CRYPTO_GROUP_ID_HASH,
+    TFM_CRYPTO_GROUP_ID_MAC,
+    TFM_CRYPTO_GROUP_ID_CIPHER,
+    TFM_CRYPTO_GROUP_ID_AEAD,
+    TFM_CRYPTO_GROUP_ID_ASYM_SIGN,
+    TFM_CRYPTO_GROUP_ID_ASYM_ENCRYPT,
+    TFM_CRYPTO_GROUP_ID_KEY_DERIVATION,
+};
+
+/**
+ * \brief Accessor to the API descriptor table that returns function_type of the API
+ *
+ * \param[in] func Function ID for the API to retrieve the associated group
+ *
+ * \return Return values as described in \ref enum tfm_crypto_function_type
+ */
+enum tfm_crypto_function_type
+    get_function_type_from_descriptor(enum tfm_crypto_function_id func);
+
+/**
+ * \brief Accessor to the API descriptor table that returns function_type of the API
+ *
+ * \param[in] id Pointer to hold the ID of the caller
+ *
+ * \return Return values as described in \ref enum tfm_crypto_group_id
+ */
+enum tfm_crypto_group_id
+    get_group_id_from_descriptor(enum tfm_crypto_function_id func);
+
+/**
+ * \brief Macro to determine the group_id corresponding to a function_id by
+ *        accessing the tfm_crypto_api_descriptor table
+ */
+#define TFM_CRYPTO_IS_GROUP_ID(_function_id, _group_id) \
+    (get_group_id_from_descriptor((_function_id)) == (_group_id))
+
+/**
+ * \brief Macro to get the function_type associated to a function_id by
+ *        accessing the tfm_crypto_api_descriptor table
+ */
+#define TFM_CRYPTO_GET_FUNCTION_TYPE(_function_id) \
+    (get_function_type_from_descriptor((_function_id)))
+
+/**
  * \brief Initialise the service
  *
  * \return Return values as described in \ref psa_status_t
@@ -130,86 +179,131 @@
                                          uint32_t handle,
                                          void **ctx);
 /**
- * \brief Encodes the input key id and owner to output key
+ * \brief This function handles the operations to allocate or retrieve a
+ *        multipart context. It receives a operation and function type as
+ *        inputs and then returns the looked or allocated context using
+ *        the provided handle. If a context is allocated in a setup call
+ *        the returned handle is then updated with the allocated value
  *
- * \param[in]  key_id       Id of the key to encode
- * \param[out] enc_key_ptr  Pointer to encoded key with id and owner
+ * \param[in]      type          Type of the operation context to look up
+ * \param[in]      function_type Type of the function
+ * \param[in, out] handle        Pointer to the handle
+ * \param[out]     ctx           Double pointer to the context
  *
  * \return Return values as described in \ref psa_status_t
  */
-psa_status_t tfm_crypto_encode_id_and_owner(psa_key_id_t key_id,
-                                            mbedtls_svc_key_id_t *enc_key_ptr);
-
-#define LIST_TFM_CRYPTO_UNIFORM_SIGNATURE_API \
-    X(tfm_crypto_get_key_attributes)          \
-    X(tfm_crypto_reset_key_attributes)        \
-    X(tfm_crypto_open_key)                    \
-    X(tfm_crypto_close_key)                   \
-    X(tfm_crypto_import_key)                  \
-    X(tfm_crypto_destroy_key)                 \
-    X(tfm_crypto_export_key)                  \
-    X(tfm_crypto_export_public_key)           \
-    X(tfm_crypto_purge_key)                   \
-    X(tfm_crypto_copy_key)                    \
-    X(tfm_crypto_hash_compute)                \
-    X(tfm_crypto_hash_compare)                \
-    X(tfm_crypto_hash_setup)                  \
-    X(tfm_crypto_hash_update)                 \
-    X(tfm_crypto_hash_finish)                 \
-    X(tfm_crypto_hash_verify)                 \
-    X(tfm_crypto_hash_abort)                  \
-    X(tfm_crypto_hash_clone)                  \
-    X(tfm_crypto_mac_compute)                 \
-    X(tfm_crypto_mac_verify)                  \
-    X(tfm_crypto_mac_sign_setup)              \
-    X(tfm_crypto_mac_verify_setup)            \
-    X(tfm_crypto_mac_update)                  \
-    X(tfm_crypto_mac_sign_finish)             \
-    X(tfm_crypto_mac_verify_finish)           \
-    X(tfm_crypto_mac_abort)                   \
-    X(tfm_crypto_cipher_encrypt)              \
-    X(tfm_crypto_cipher_decrypt)              \
-    X(tfm_crypto_cipher_encrypt_setup)        \
-    X(tfm_crypto_cipher_decrypt_setup)        \
-    X(tfm_crypto_cipher_generate_iv)          \
-    X(tfm_crypto_cipher_set_iv)               \
-    X(tfm_crypto_cipher_update)               \
-    X(tfm_crypto_cipher_finish)               \
-    X(tfm_crypto_cipher_abort)                \
-    X(tfm_crypto_aead_encrypt)                \
-    X(tfm_crypto_aead_decrypt)                \
-    X(tfm_crypto_aead_encrypt_setup)          \
-    X(tfm_crypto_aead_decrypt_setup)          \
-    X(tfm_crypto_aead_generate_nonce)         \
-    X(tfm_crypto_aead_set_nonce)              \
-    X(tfm_crypto_aead_set_lengths)            \
-    X(tfm_crypto_aead_update_ad)              \
-    X(tfm_crypto_aead_update)                 \
-    X(tfm_crypto_aead_finish)                 \
-    X(tfm_crypto_aead_verify)                 \
-    X(tfm_crypto_aead_abort)                  \
-    X(tfm_crypto_sign_message)                \
-    X(tfm_crypto_verify_message)              \
-    X(tfm_crypto_sign_hash)                   \
-    X(tfm_crypto_verify_hash)                 \
-    X(tfm_crypto_asymmetric_encrypt)          \
-    X(tfm_crypto_asymmetric_decrypt)          \
-    X(tfm_crypto_key_derivation_setup)        \
-    X(tfm_crypto_key_derivation_get_capacity) \
-    X(tfm_crypto_key_derivation_set_capacity) \
-    X(tfm_crypto_key_derivation_input_bytes)  \
-    X(tfm_crypto_key_derivation_input_key)    \
-    X(tfm_crypto_key_derivation_key_agreement)\
-    X(tfm_crypto_key_derivation_output_bytes) \
-    X(tfm_crypto_key_derivation_output_key)   \
-    X(tfm_crypto_key_derivation_abort)        \
-    X(tfm_crypto_raw_key_agreement)           \
-    X(tfm_crypto_generate_random)             \
-    X(tfm_crypto_generate_key)                \
-
-#define X(api_name) UNIFORM_SIGNATURE_API(api_name);
-LIST_TFM_CRYPTO_UNIFORM_SIGNATURE_API
-#undef X
+psa_status_t tfm_crypto_operation_handling(enum tfm_crypto_operation_type type,
+                                    enum tfm_crypto_function_type function_type,
+                                    uint32_t *handle,
+                                    void **ctx);
+/**
+ * \brief This function acts as interface from the framework dispatching
+ *        calls to the set of functions that implement the PSA Crypto APIs.
+ *        It is based on the Uniform Signatures prototype.
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[in]  in_len   Length of the valid entries in in_vec
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  out_len  Length of the valid entries in out_vec
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_api_dispatcher(psa_invec in_vec[],
+                                       size_t in_len,
+                                       psa_outvec out_vec[],
+                                       size_t out_len);
+/**
+ * \brief This function acts as interface for the Key management module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  encoded_key Key encoded with partition_id and key_id
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_key_management_interface(psa_invec in_vec[],
+                                            psa_outvec out_vec[],
+                                            mbedtls_svc_key_id_t *encoded_key);
+/**
+ * \brief This function acts as interface for the MAC module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  encoded_key Key encoded with partition_id and key_id
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_mac_interface(psa_invec in_vec[],
+                                      psa_outvec out_vec[],
+                                      mbedtls_svc_key_id_t *encoded_key);
+/**
+ * \brief This function acts as interface for the Cipher module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  encoded_key Key encoded with partition_id and key_id
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_cipher_interface(psa_invec in_vec[],
+                                         psa_outvec out_vec[],
+                                         mbedtls_svc_key_id_t *encoded_key);
+/**
+ * \brief This function acts as interface for the AEAD module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  encoded_key Key encoded with partition_id and key_id
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_aead_interface(psa_invec in_vec[],
+                                       psa_outvec out_vec[],
+                                       mbedtls_svc_key_id_t *encoded_key);
+/**
+ * \brief This function acts as interface for the Asymmetric module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  encoded_key Key encoded with partition_id and key_id
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_asymmetric_interface(psa_invec in_vec[],
+                                             psa_outvec out_vec[],
+                                             mbedtls_svc_key_id_t *encoded_key);
+/**
+ * \brief This function acts as interface for the Key derivation module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ * \param[in]  encoded_key Key encoded with partition_id and key_id
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_key_derivation_interface(psa_invec in_vec[],
+                                            psa_outvec out_vec[],
+                                            mbedtls_svc_key_id_t *encoded_key);
+/**
+ * \brief This function acts as interface for the Random module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_random_interface(psa_invec in_vec[],
+                                         psa_outvec out_vec[]);
+/**
+ * \brief This function acts as interface for the Hash module
+ *
+ * \param[in]  in_vec   Array of invec parameters
+ * \param[out] out_vec  Array of outvec parameters
+ *
+ * \return Return values as described in \ref psa_status_t
+ */
+psa_status_t tfm_crypto_hash_interface(psa_invec in_vec[],
+                                       psa_outvec out_vec[]);
 
 #ifdef __cplusplus
 }
diff --git a/secure_fw/partitions/crypto/tfm_crypto_private.h b/secure_fw/partitions/crypto/tfm_crypto_private.h
deleted file mode 100644
index 34a41c8..0000000
--- a/secure_fw/partitions/crypto/tfm_crypto_private.h
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Copyright (c) 2020-2021, Arm Limited. All rights reserved.
- *
- * SPDX-License-Identifier: BSD-3-Clause
- *
- */
-
-#ifndef __TFM_CRYPTO_PRIVATE_H__
-#define __TFM_CRYPTO_PRIVATE_H__
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-#ifdef TFM_PSA_API
-
-#include <stdbool.h>
-
-/* \brief Verifies that in_len is in the [in_min, in_max] range
- *        AND out_len is in the [out_min, out_max] range. In
- *        case any of the two in_len or out_len is out of range,
- *        returns false. Returns true in case of success.
- */
-static inline bool tfm_crypto_private_check_length(
-    size_t in_len, size_t in_min, size_t in_max,
-    size_t out_len, size_t out_min, size_t out_max)
-{
-   if ((in_len >= in_min) && (in_len <= in_max) &&
-        (out_len >= out_min) && (out_len <= out_max)) {
-        return true;
-    }
-    return false;
-}
-
-/*
- * Validate the IOVEC[] lengths for IPC model. The tfm_crypto_call_srv()
- * reduces the entries in IOVEC[] which are empty from `in_len` and `out_len`.
- * This means that Crypto service APIs need to ensure that the `in_len`
- * and `out_len` are within the expected range.
- *
- * Also tfm_crypto_call_srv() ensures that all entries in IOVEC[] are
- * initialised. Hence all entries in IOVEC[] can be accessed to
- * initialize internal variables even if they are outside `in_len`
- * and `out_len`.
- */
-#define CRYPTO_IN_OUT_LEN_VALIDATE(in_len, in_min, in_max, out_len, out_min, out_max)   \
-            if (!tfm_crypto_private_check_length(                             \
-                    in_len, in_min, in_max, out_len, out_min, out_max)) {    \
-                return PSA_ERROR_PROGRAMMER_ERROR;                           \
-            }
-#else
-/*
- * Validate the IOVEC[] lengths for Library model.  Unlike the IPC model, the
- * service APIs expects to receive the exact of `in_len` and `out_len`
- * as expected by the API.
- */
-#define CRYPTO_IN_OUT_LEN_VALIDATE(in_len, in_min, in_max, out_len, out_min, out_max)    \
-            if (((in_len) != (in_max)) || ((out_len) != (out_max))) {     \
-                return PSA_ERROR_PROGRAMMER_ERROR;                        \
-            }
-#endif
-
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif /* __TFM_CRYPTO_PRIVATE_H__ */
diff --git a/secure_fw/partitions/crypto/tfm_crypto_secure_api.c b/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
index 5cc7e90..89ec16a 100644
--- a/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
+++ b/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
@@ -5,7 +5,6 @@
  *
  */
 
-#include "array.h"
 #include "tfm_crypto_defs.h"
 #include "psa/crypto.h"
 #ifdef TFM_PSA_API
@@ -16,26 +15,23 @@
 #endif
 
 #ifdef TFM_PSA_API
-
-#define API_DISPATCH(srv_name, srv_id)                         \
-    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,                  \
-        in_vec, ARRAY_SIZE(in_vec),                            \
-        out_vec, ARRAY_SIZE(out_vec))
-
-#define API_DISPATCH_NO_OUTVEC(srv_name, srv_id)               \
-    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,                  \
-        in_vec, ARRAY_SIZE(in_vec),                            \
-        (psa_outvec *)NULL, 0)
+#define API_DISPATCH(in_vec, out_vec)          \
+    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,  \
+             in_vec, IOVEC_LEN(in_vec),        \
+             out_vec, IOVEC_LEN(out_vec))
+#define API_DISPATCH_NO_OUTVEC(in_vec)         \
+    psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,  \
+             in_vec, IOVEC_LEN(in_vec),        \
+             (psa_outvec *)NULL, 0)
 #else
-#define API_DISPATCH(srv_name, srv_id)                         \
-    srv_name##_veneer(                                         \
-        in_vec, ARRAY_SIZE(in_vec),                            \
-        out_vec, ARRAY_SIZE(out_vec))
-
-#define API_DISPATCH_NO_OUTVEC(srv_name, srv_id)               \
-    srv_name##_veneer(                                         \
-        in_vec, ARRAY_SIZE(in_vec),                            \
-        NULL, 0)
+#define API_DISPATCH(in_vec, out_vec)          \
+    tfm_crypto_api_dispatcher_veneer(  \
+                in_vec, IOVEC_LEN(in_vec),     \
+                out_vec, IOVEC_LEN(out_vec))
+#define API_DISPATCH_NO_OUTVEC(in_vec)         \
+    tfm_crypto_api_dispatcher_veneer(  \
+                in_vec, IOVEC_LEN(in_vec),     \
+                NULL, 0)
 #endif /* TFM_PSA_API */
 
 psa_status_t psa_crypto_init(void)
@@ -47,62 +43,53 @@
 }
 
 psa_status_t psa_open_key(psa_key_id_t id,
-                          psa_key_id_t *key_id)
+                          psa_key_id_t *key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
-    struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_OPEN_KEY_SID,
+    const struct tfm_crypto_pack_iovec iov = {
+        .function_id = TFM_CRYPTO_OPEN_KEY_SID,
+        .key_id = id,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
-        {.base = &id, .len = sizeof(psa_key_id_t)},
     };
     psa_outvec out_vec[] = {
-        {.base = key_id, .len = sizeof(psa_key_id_t)},
+        {.base = key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_open_key,
-                          TFM_CRYPTO_OPEN_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_close_key(psa_key_id_t key_id)
+psa_status_t psa_close_key(psa_key_id_t key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
-    struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CLOSE_KEY_SID,
-        .key_id = key_id,
+    const struct tfm_crypto_pack_iovec iov = {
+        .function_id = TFM_CRYPTO_CLOSE_KEY_SID,
+        .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_close_key,
-                                    TFM_CRYPTO_CLOSE_KEY);;
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
 psa_status_t psa_import_key(const psa_key_attributes_t *attributes,
                             const uint8_t *data,
                             size_t data_length,
-                            psa_key_id_t *key_id)
+                            psa_key_id_t *key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_IMPORT_KEY_SID,
+        .function_id = TFM_CRYPTO_IMPORT_KEY_SID,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -110,47 +97,39 @@
         {.base = data, .len = data_length}
     };
     psa_outvec out_vec[] = {
-        {.base = key_id, .len = sizeof(psa_key_id_t)}
+        {.base = key, .len = sizeof(psa_key_id_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_import_key,
-                          TFM_CRYPTO_IMPORT_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_destroy_key(psa_key_id_t key_id)
+psa_status_t psa_destroy_key(psa_key_id_t key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_DESTROY_KEY_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_DESTROY_KEY_SID,
+        .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_destroy_key,
-                                    TFM_CRYPTO_DESTROY_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_get_key_attributes(psa_key_id_t key_id,
+psa_status_t psa_get_key_attributes(psa_key_id_t key,
                                     psa_key_attributes_t *attributes)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_GET_KEY_ATTRIBUTES_SID,
+        .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -159,10 +138,7 @@
         {.base = attributes, .len = sizeof(psa_key_attributes_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_get_key_attributes,
-                          TFM_CRYPTO_GET_KEY_ATTRIBUTES);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
@@ -172,7 +148,7 @@
     return PSA_ERROR_NOT_SUPPORTED;
 #else
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
+        .function_id = TFM_CRYPTO_RESET_KEY_ATTRIBUTES_SID,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -181,14 +157,12 @@
         {.base = attributes, .len = sizeof(psa_key_attributes_t)},
     };
 
-    (void)API_DISPATCH(tfm_crypto_reset_key_attributes,
-                       TFM_CRYPTO_RESET_KEY_ATTRIBUTES);
-
+    (void)API_DISPATCH(in_vec, out_vec);
     return;
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_export_key(psa_key_id_t key_id,
+psa_status_t psa_export_key(psa_key_id_t key,
                             uint8_t *data,
                             size_t data_size,
                             size_t *data_length)
@@ -198,8 +172,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_EXPORT_KEY_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_EXPORT_KEY_SID,
+        .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
@@ -208,8 +182,7 @@
         {.base = data, .len = data_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_export_key,
-                          TFM_CRYPTO_EXPORT_KEY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *data_length = out_vec[0].len;
 
@@ -217,7 +190,7 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_export_public_key(psa_key_id_t key_id,
+psa_status_t psa_export_public_key(psa_key_id_t key,
                                    uint8_t *data,
                                    size_t data_size,
                                    size_t *data_length)
@@ -227,8 +200,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_EXPORT_PUBLIC_KEY_SID,
+        .key_id = key,
     };
 
     psa_invec in_vec[] = {
@@ -238,8 +211,7 @@
         {.base = data, .len = data_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_export_public_key,
-                          TFM_CRYPTO_EXPORT_PUBLIC_KEY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *data_length = out_vec[0].len;
 
@@ -247,38 +219,33 @@
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_purge_key(psa_key_id_t key_id)
+psa_status_t psa_purge_key(psa_key_id_t key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_PURGE_KEY_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_PURGE_KEY_SID,
+        .key_id = key,
     };
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_purge_key,
-                                    TFM_CRYPTO_PURGE_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_copy_key(psa_key_id_t source_key_id,
+psa_status_t psa_copy_key(psa_key_id_t source_key,
                           const psa_key_attributes_t *attributes,
-                          psa_key_id_t *target_key_id)
+                          psa_key_id_t *target_key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_COPY_KEY_SID,
-        .key_id = source_key_id,
+        .function_id = TFM_CRYPTO_COPY_KEY_SID,
+        .key_id = source_key,
     };
 
     psa_invec in_vec[] = {
@@ -287,13 +254,11 @@
     };
 
     psa_outvec out_vec[] = {
-        {.base = target_key_id, .len = sizeof(psa_key_id_t)},
+        {.base = target_key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_copy_key,
-                          TFM_CRYPTO_COPY_KEY);
+    return API_DISPATCH(in_vec, out_vec);
 
-    return status;
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
@@ -307,7 +272,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
+        .function_id = TFM_CRYPTO_CIPHER_GENERATE_IV_SID,
         .op_handle = operation->handle,
     };
 
@@ -315,14 +280,12 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = iv, .len = iv_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_generate_iv,
-                          TFM_CRYPTO_CIPHER_GENERATE_IV);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *iv_length = out_vec[1].len;
+    *iv_length = out_vec[0].len;
 
     return status;
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
@@ -335,9 +298,8 @@
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_SET_IV_SID,
+        .function_id = TFM_CRYPTO_CIPHER_SET_IV_SID,
         .op_handle = operation->handle,
     };
 
@@ -345,28 +307,21 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = iv, .len = iv_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_cipher_set_iv,
-                          TFM_CRYPTO_CIPHER_SET_IV);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
 }
 
 psa_status_t psa_cipher_encrypt_setup(psa_cipher_operation_t *operation,
-                                      psa_key_id_t key_id,
+                                      psa_key_id_t key,
                                       psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_CIPHER_ENCRYPT_SETUP_SID,
+        .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -378,24 +333,20 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_encrypt_setup,
-                          TFM_CRYPTO_CIPHER_ENCRYPT_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
 }
 
 psa_status_t psa_cipher_decrypt_setup(psa_cipher_operation_t *operation,
-                                      psa_key_id_t key_id,
+                                      psa_key_id_t key,
                                       psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_CIPHER_DECRYPT_SETUP_SID,
+        .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -407,11 +358,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-
-    status = API_DISPATCH(tfm_crypto_cipher_decrypt_setup,
-                          TFM_CRYPTO_CIPHER_DECRYPT_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
 }
 
@@ -427,7 +374,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_UPDATE_SID,
+        .function_id = TFM_CRYPTO_CIPHER_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -436,14 +383,12 @@
         {.base = input, .len = input_length},
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_update,
-                          TFM_CRYPTO_CIPHER_UPDATE);
+    status = API_DISPATCH(in_vec, out_vec);
 
-    *output_length = out_vec[1].len;
+    *output_length = out_vec[0].len;
 
     return status;
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
@@ -454,9 +399,8 @@
 #ifdef TFM_CRYPTO_CIPHER_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ABORT_SID,
+        .function_id = TFM_CRYPTO_CIPHER_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -467,10 +411,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_abort,
-                          TFM_CRYPTO_CIPHER_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
 }
 
@@ -484,7 +425,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_FINISH_SID,
+        .function_id = TFM_CRYPTO_CIPHER_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -496,8 +437,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_finish,
-                          TFM_CRYPTO_CIPHER_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[1].len;
 
@@ -511,9 +451,8 @@
 #ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_SETUP_SID,
+        .function_id = TFM_CRYPTO_HASH_SETUP_SID,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -525,10 +464,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_setup,
-                          TFM_CRYPTO_HASH_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 
@@ -539,9 +475,8 @@
 #ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_UPDATE_SID,
+        .function_id = TFM_CRYPTO_HASH_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -549,14 +484,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_hash_update,
-                          TFM_CRYPTO_HASH_UPDATE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 
@@ -570,7 +499,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_FINISH_SID,
+        .function_id = TFM_CRYPTO_HASH_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -582,8 +511,7 @@
         {.base = hash, .len = hash_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_finish,
-                          TFM_CRYPTO_HASH_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *hash_length = out_vec[1].len;
 
@@ -598,9 +526,8 @@
 #ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_VERIFY_SID,
+        .function_id = TFM_CRYPTO_HASH_VERIFY_SID,
         .op_handle = operation->handle,
     };
 
@@ -612,10 +539,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_verify,
-                          TFM_CRYPTO_HASH_VERIFY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 
@@ -624,9 +548,8 @@
 #ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_ABORT_SID,
+        .function_id = TFM_CRYPTO_HASH_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -637,10 +560,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_abort,
-                          TFM_CRYPTO_HASH_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 
@@ -650,9 +570,8 @@
 #ifdef TFM_CRYPTO_HASH_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_CLONE_SID,
+        .function_id = TFM_CRYPTO_HASH_CLONE_SID,
         .op_handle = source_operation->handle,
     };
 
@@ -667,10 +586,7 @@
         return PSA_ERROR_BAD_STATE;
     }
 
-    status = API_DISPATCH(tfm_crypto_hash_clone,
-                          TFM_CRYPTO_HASH_CLONE);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 
@@ -686,7 +602,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_COMPUTE_SID,
+        .function_id = TFM_CRYPTO_HASH_COMPUTE_SID,
         .alg = alg,
     };
 
@@ -699,8 +615,7 @@
         {.base = hash, .len = hash_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_hash_compute,
-                          TFM_CRYPTO_HASH_COMPUTE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *hash_length = out_vec[0].len;
 
@@ -717,9 +632,8 @@
 #if (TFM_CRYPTO_HASH_MODULE_DISABLED != 0)
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_HASH_COMPARE_SID,
+        .function_id = TFM_CRYPTO_HASH_COMPARE_SID,
         .alg = alg,
     };
 
@@ -729,24 +643,20 @@
         {.base = hash, .len = hash_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_hash_compare,
-                                    TFM_CRYPTO_HASH_COMPARE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_HASH_MODULE_DISABLED */
 }
 
 psa_status_t psa_mac_sign_setup(psa_mac_operation_t *operation,
-                                psa_key_id_t key_id,
+                                psa_key_id_t key,
                                 psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_MAC_SIGN_SETUP_SID,
+        .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -758,24 +668,20 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_sign_setup,
-                          TFM_CRYPTO_MAC_SIGN_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
 psa_status_t psa_mac_verify_setup(psa_mac_operation_t *operation,
-                                  psa_key_id_t key_id,
+                                  psa_key_id_t key,
                                   psa_algorithm_t alg)
 {
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_SETUP_SID,
+        .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -787,10 +693,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_verify_setup,
-                          TFM_CRYPTO_MAC_VERIFY_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
@@ -801,9 +704,8 @@
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_UPDATE_SID,
+        .function_id = TFM_CRYPTO_MAC_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -811,14 +713,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
 
-    status = API_DISPATCH(tfm_crypto_mac_update,
-                          TFM_CRYPTO_MAC_UPDATE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
@@ -832,7 +728,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_SIGN_FINISH_SID,
+        .function_id = TFM_CRYPTO_MAC_SIGN_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -844,8 +740,7 @@
         {.base = mac, .len = mac_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_sign_finish,
-                          TFM_CRYPTO_MAC_SIGN_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *mac_length = out_vec[1].len;
 
@@ -860,9 +755,8 @@
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -874,10 +768,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_verify_finish,
-                          TFM_CRYPTO_MAC_VERIFY_FINISH);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
@@ -886,9 +777,8 @@
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_ABORT_SID,
+        .function_id = TFM_CRYPTO_MAC_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -899,14 +789,11 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_abort,
-                          TFM_CRYPTO_MAC_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
-psa_status_t psa_aead_encrypt(psa_key_id_t key_id,
+psa_status_t psa_aead_encrypt(psa_key_id_t key,
                               psa_algorithm_t alg,
                               const uint8_t *nonce,
                               size_t nonce_length,
@@ -923,8 +810,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_AEAD_ENCRYPT_SID,
+        .key_id = key,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = 0}
     };
@@ -958,16 +845,15 @@
     in_vec[0].len = sizeof(struct tfm_crypto_pack_iovec);
 
 #ifdef TFM_PSA_API
-    size_t in_len = ARRAY_SIZE(in_vec);
+    size_t in_len = IOVEC_LEN(in_vec);
 
     if (additional_data == NULL) {
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, ARRAY_SIZE(out_vec));
+                      out_vec, IOVEC_LEN(out_vec));
 #else
-    status = API_DISPATCH(tfm_crypto_aead_encrypt,
-                          TFM_CRYPTO_AEAD_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 #endif
 
     *ciphertext_length = out_vec[0].len;
@@ -976,7 +862,7 @@
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 
-psa_status_t psa_aead_decrypt(psa_key_id_t key_id,
+psa_status_t psa_aead_decrypt(psa_key_id_t key,
                               psa_algorithm_t alg,
                               const uint8_t *nonce,
                               size_t nonce_length,
@@ -993,8 +879,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_AEAD_DECRYPT_SID,
+        .key_id = key,
         .alg = alg,
         .aead_in = {.nonce = {0}, .nonce_length = 0}
     };
@@ -1028,16 +914,15 @@
     in_vec[0].len = sizeof(struct tfm_crypto_pack_iovec);
 
 #ifdef TFM_PSA_API
-    size_t in_len = ARRAY_SIZE(in_vec);
+    size_t in_len = IOVEC_LEN(in_vec);
 
     if (additional_data == NULL) {
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, ARRAY_SIZE(out_vec));
+                      out_vec, IOVEC_LEN(out_vec));
 #else
-    status = API_DISPATCH(tfm_crypto_aead_decrypt,
-                          TFM_CRYPTO_AEAD_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 #endif
 
     *plaintext_length = out_vec[0].len;
@@ -1055,7 +940,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_AEAD_ENCRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -1068,8 +953,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_encrypt_setup,
-                          TFM_CRYPTO_AEAD_ENCRYPT_SETUP);
+    status = API_DISPATCH(in_vec, out_vec);
     return status;
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
@@ -1083,7 +967,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
+        .function_id = TFM_CRYPTO_AEAD_DECRYPT_SETUP_SID,
         .key_id = key,
         .alg = alg,
         .op_handle = operation->handle,
@@ -1096,8 +980,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_decrypt_setup,
-                          TFM_CRYPTO_AEAD_DECRYPT_SETUP);
+    status = API_DISPATCH(in_vec, out_vec);
     return status;
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
@@ -1112,7 +995,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
+        .function_id = TFM_CRYPTO_AEAD_GENERATE_NONCE_SID,
         .op_handle = operation->handle,
     };
 
@@ -1124,8 +1007,7 @@
         {.base = nonce, .len = nonce_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_generate_nonce,
-                          TFM_CRYPTO_AEAD_GENERATE_NONCE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *nonce_length = out_vec[1].len;
     return status;
@@ -1141,7 +1023,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_SET_NONCE_SID,
+        .function_id = TFM_CRYPTO_AEAD_SET_NONCE_SID,
         .op_handle = operation->handle,
     };
 
@@ -1149,12 +1031,8 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = nonce, .len = nonce_length}
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_set_nonce,
-                          TFM_CRYPTO_AEAD_SET_NONCE);
+    status = API_DISPATCH_NO_OUTVEC(in_vec);
     return status;
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
@@ -1168,7 +1046,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
+        .function_id = TFM_CRYPTO_AEAD_SET_LENGTHS_SID,
         .ad_length = ad_length,
         .plaintext_length = plaintext_length,
         .op_handle = operation->handle,
@@ -1177,12 +1055,8 @@
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
-    status = API_DISPATCH(tfm_crypto_aead_set_lengths,
-                          TFM_CRYPTO_AEAD_SET_LENGTHS);
+    status = API_DISPATCH_NO_OUTVEC(in_vec);
     return status;
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
@@ -1196,7 +1070,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_UPDATE_AD_SID,
+        .function_id = TFM_CRYPTO_AEAD_UPDATE_AD_SID,
         .op_handle = operation->handle,
     };
 
@@ -1209,21 +1083,17 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
     };
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)}
-    };
 
 #ifdef TFM_PSA_API
-    size_t in_len = ARRAY_SIZE(in_vec);
+    size_t in_len = IOVEC_LEN(in_vec);
 
     if (input == NULL) {
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, ARRAY_SIZE(out_vec));
+                      NULL, 0);
 #else
-    status = API_DISPATCH(tfm_crypto_aead_update_ad,
-                          TFM_CRYPTO_AEAD_UPDATE_AD);
+    status = API_DISPATCH_NO_OUTVEC(in_vec);
 #endif
 
     return status;
@@ -1242,7 +1112,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_UPDATE_SID,
+        .function_id = TFM_CRYPTO_AEAD_UPDATE_SID,
         .op_handle = operation->handle,
     };
 
@@ -1256,24 +1126,22 @@
         {.base = input, .len = input_length}
     };
     psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
         {.base = output, .len = output_size},
     };
 
 #ifdef TFM_PSA_API
-    size_t in_len = ARRAY_SIZE(in_vec);
+    size_t in_len = IOVEC_LEN(in_vec);
 
     if (input == NULL) {
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, ARRAY_SIZE(out_vec));
+                      out_vec, IOVEC_LEN(out_vec));
 #else
-    status = API_DISPATCH(tfm_crypto_aead_update,
-                          TFM_CRYPTO_AEAD_UPDATE);
+    status = API_DISPATCH(in_vec, out_vec);
 #endif
 
-    *output_length = out_vec[1].len;
+    *output_length = out_vec[0].len;
     return status;
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
@@ -1291,7 +1159,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_FINISH_SID,
+        .function_id = TFM_CRYPTO_AEAD_FINISH_SID,
         .op_handle = operation->handle,
     };
 
@@ -1310,7 +1178,7 @@
     };
 
 #ifdef TFM_PSA_API
-    size_t out_len = ARRAY_SIZE(out_vec);
+    size_t out_len = IOVEC_LEN(out_vec);
 
     if (ciphertext == NULL || ciphertext_size == 0) {
         out_len--;
@@ -1320,7 +1188,7 @@
     }
 
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,
-                      in_vec, ARRAY_SIZE(in_vec),
+                      in_vec, IOVEC_LEN(in_vec),
                       out_vec, out_len);
 
     if (out_len == 3) {
@@ -1328,10 +1196,8 @@
     } else {
         *ciphertext_length = 0;
     }
-
 #else
-    status = API_DISPATCH(tfm_crypto_aead_finish,
-                          TFM_CRYPTO_AEAD_FINISH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *ciphertext_length = out_vec[2].len;
 #endif
@@ -1354,7 +1220,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_VERIFY_SID,
+        .function_id = TFM_CRYPTO_AEAD_VERIFY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1373,7 +1239,7 @@
     };
 
 #ifdef TFM_PSA_API
-    size_t out_len = ARRAY_SIZE(out_vec);
+    size_t out_len = IOVEC_LEN(out_vec);
 
     if (plaintext == NULL || plaintext_size == 0) {
         out_len--;
@@ -1381,8 +1247,9 @@
     if ((out_len == 2) && (plaintext_length == NULL)) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }
+
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL,
-                      in_vec, ARRAY_SIZE(in_vec),
+                      in_vec, IOVEC_LEN(in_vec),
                       out_vec, out_len);
 
     if (out_len == 2) {
@@ -1392,8 +1259,7 @@
     }
 
 #else
-    status = API_DISPATCH(tfm_crypto_aead_verify,
-                          TFM_CRYPTO_AEAD_VERIFY);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *plaintext_length = out_vec[1].len;
 #endif
@@ -1407,9 +1273,8 @@
 #ifdef TFM_CRYPTO_AEAD_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_AEAD_ABORT_SID,
+        .function_id = TFM_CRYPTO_AEAD_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -1420,13 +1285,11 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_abort,
-                          TFM_CRYPTO_AEAD_ABORT);
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
 
-psa_status_t psa_sign_message(psa_key_id_t key_id,
+psa_status_t psa_sign_message(psa_key_id_t key,
                               psa_algorithm_t alg,
                               const uint8_t *input,
                               size_t input_length,
@@ -1439,8 +1302,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_SIGN_MESSAGE_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_SIGN_MESSAGE_SID,
+        .key_id = key,
         .alg = alg,
     };
 
@@ -1452,15 +1315,14 @@
         {.base = signature, .len = signature_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_sign_message,
-                          TFM_CRYPTO_SIGN_MESSAGE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *signature_length = out_vec[0].len;
     return status;
 #endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
 }
 
-psa_status_t psa_verify_message(psa_key_id_t key_id,
+psa_status_t psa_verify_message(psa_key_id_t key,
                                 psa_algorithm_t alg,
                                 const uint8_t *input,
                                 size_t input_length,
@@ -1470,10 +1332,9 @@
 #ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_VERIFY_MESSAGE_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_VERIFY_MESSAGE_SID,
+        .key_id = key,
         .alg = alg
     };
 
@@ -1483,14 +1344,11 @@
         {.base = signature, .len = signature_length}
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_verify_message,
-                                    TFM_CRYPTO_VERIFY_MESSAGE);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
 }
 
-psa_status_t psa_sign_hash(psa_key_id_t key_id,
+psa_status_t psa_sign_hash(psa_key_id_t key,
                            psa_algorithm_t alg,
                            const uint8_t *hash,
                            size_t hash_length,
@@ -1503,8 +1361,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_SIGN_HASH_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_SIGN_HASH_SID,
+        .key_id = key,
         .alg = alg,
     };
 
@@ -1516,8 +1374,7 @@
         {.base = signature, .len = signature_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_sign_hash,
-                          TFM_CRYPTO_SIGN_HASH);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *signature_length = out_vec[0].len;
 
@@ -1525,7 +1382,7 @@
 #endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
 }
 
-psa_status_t psa_verify_hash(psa_key_id_t key_id,
+psa_status_t psa_verify_hash(psa_key_id_t key,
                              psa_algorithm_t alg,
                              const uint8_t *hash,
                              size_t hash_length,
@@ -1535,10 +1392,9 @@
 #ifdef TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_VERIFY_HASH_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_VERIFY_HASH_SID,
+        .key_id = key,
         .alg = alg
     };
 
@@ -1548,14 +1404,11 @@
         {.base = signature, .len = signature_length}
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_verify_hash,
-                                    TFM_CRYPTO_VERIFY_HASH);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED */
 }
 
-psa_status_t psa_asymmetric_encrypt(psa_key_id_t key_id,
+psa_status_t psa_asymmetric_encrypt(psa_key_id_t key,
                                     psa_algorithm_t alg,
                                     const uint8_t *input,
                                     size_t input_length,
@@ -1570,8 +1423,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_ENCRYPT_SID,
+        .key_id = key,
         .alg = alg
     };
 
@@ -1591,16 +1444,15 @@
     };
 
 #ifdef TFM_PSA_API
-    size_t in_len = ARRAY_SIZE(in_vec);
+    size_t in_len = IOVEC_LEN(in_vec);
 
     if (salt == NULL) {
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, ARRAY_SIZE(out_vec));
+                      out_vec, IOVEC_LEN(out_vec));
 #else
-    status = API_DISPATCH(tfm_crypto_asymmetric_encrypt,
-                          TFM_CRYPTO_ASYMMETRIC_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 #endif
 
     *output_length = out_vec[0].len;
@@ -1609,7 +1461,7 @@
 #endif /* TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED */
 }
 
-psa_status_t psa_asymmetric_decrypt(psa_key_id_t key_id,
+psa_status_t psa_asymmetric_decrypt(psa_key_id_t key,
                                     psa_algorithm_t alg,
                                     const uint8_t *input,
                                     size_t input_length,
@@ -1624,8 +1476,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_ASYMMETRIC_DECRYPT_SID,
+        .key_id = key,
         .alg = alg
     };
 
@@ -1645,16 +1497,15 @@
     };
 
 #ifdef TFM_PSA_API
-    size_t in_len = ARRAY_SIZE(in_vec);
+    size_t in_len = IOVEC_LEN(in_vec);
 
     if (salt == NULL) {
         in_len--;
     }
     status = psa_call(TFM_CRYPTO_HANDLE, PSA_IPC_CALL, in_vec, in_len,
-                      out_vec, ARRAY_SIZE(out_vec));
+                      out_vec, IOVEC_LEN(out_vec));
 #else
-    status = API_DISPATCH(tfm_crypto_asymmetric_decrypt,
-                          TFM_CRYPTO_ASYMMETRIC_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 #endif
 
     *output_length = out_vec[0].len;
@@ -1670,9 +1521,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY_SID,
         .op_handle = operation->handle,
     };
 
@@ -1684,10 +1534,7 @@
         {.base = capacity, .len = sizeof(size_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_get_capacity,
-                          TFM_CRYPTO_KEY_DERIVATION_GET_CAPACITY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -1699,9 +1546,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES_SID,
         .op_handle = operation->handle,
     };
 
@@ -1713,25 +1559,21 @@
         {.base = output, .len = output_length},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_output_bytes,
-                          TFM_CRYPTO_KEY_DERIVATION_OUTPUT_BYTES);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
 psa_status_t psa_key_derivation_input_key(
                                       psa_key_derivation_operation_t *operation,
                                       psa_key_derivation_step_t step,
-                                      psa_key_id_t key_id)
+                                      psa_key_id_t key)
 {
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY_SID,
+        .key_id = key,
         .step = step,
         .op_handle = operation->handle,
     };
@@ -1740,10 +1582,7 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_input_key,
-                                    TFM_CRYPTO_KEY_DERIVATION_INPUT_KEY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -1752,9 +1591,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_ABORT_SID,
         .op_handle = operation->handle,
     };
 
@@ -1766,10 +1604,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_abort,
-                          TFM_CRYPTO_KEY_DERIVATION_ABORT);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -1783,9 +1618,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT_SID,
         .key_id = private_key,
         .step = step,
         .op_handle = operation->handle,
@@ -1796,14 +1630,7 @@
         {.base = peer_key, .len = peer_key_length},
     };
 
-    psa_outvec out_vec[] = {
-        {.base = &(operation->handle), .len = sizeof(uint32_t)},
-    };
-
-    status = API_DISPATCH(tfm_crypto_key_derivation_key_agreement,
-                          TFM_CRYPTO_KEY_DERIVATION_KEY_AGREEMENT);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -1813,9 +1640,8 @@
 #ifdef TFM_CRYPTO_RNG_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GENERATE_RANDOM_SID,
+        .function_id = TFM_CRYPTO_GENERATE_RANDOM_SID,
     };
 
     psa_invec in_vec[] = {
@@ -1830,22 +1656,18 @@
         return PSA_SUCCESS;
     }
 
-    status = API_DISPATCH(tfm_crypto_generate_random,
-                          TFM_CRYPTO_GENERATE_RANDOM);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_RNG_MODULE_DISABLED */
 }
 
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
-                              psa_key_id_t *key_id)
+                              psa_key_id_t *key)
 {
 #ifdef TFM_CRYPTO_KEY_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_GENERATE_KEY_SID,
+        .function_id = TFM_CRYPTO_GENERATE_KEY_SID,
     };
 
     psa_invec in_vec[] = {
@@ -1854,17 +1676,14 @@
     };
 
     psa_outvec out_vec[] = {
-        {.base = key_id, .len = sizeof(psa_key_id_t)},
+        {.base = key, .len = sizeof(psa_key_id_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_generate_key,
-                          TFM_CRYPTO_GENERATE_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_MODULE_DISABLED */
 }
 
-psa_status_t psa_mac_compute(psa_key_id_t key_id,
+psa_status_t psa_mac_compute(psa_key_id_t key,
                              psa_algorithm_t alg,
                              const uint8_t *input,
                              size_t input_length,
@@ -1877,8 +1696,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_COMPUTE_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_MAC_COMPUTE_SID,
+        .key_id = key,
         .alg = alg,
     };
 
@@ -1890,15 +1709,14 @@
         {.base = mac, .len = mac_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_mac_compute,
-                          TFM_CRYPTO_MAC_COMPUTE);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *mac_length = out_vec[0].len;
     return status;
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
-psa_status_t psa_mac_verify(psa_key_id_t key_id,
+psa_status_t psa_mac_verify(psa_key_id_t key,
                             psa_algorithm_t alg,
                             const uint8_t *input,
                             size_t input_length,
@@ -1908,10 +1726,9 @@
 #ifdef TFM_CRYPTO_MAC_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_MAC_VERIFY_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_MAC_VERIFY_SID,
+        .key_id = key,
         .alg = alg,
     };
 
@@ -1921,14 +1738,11 @@
         {.base = mac, .len = mac_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_mac_verify,
-                                    TFM_CRYPTO_MAC_VERIFY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_MAC_MODULE_DISABLED */
 }
 
-psa_status_t psa_cipher_encrypt(psa_key_id_t key_id,
+psa_status_t psa_cipher_encrypt(psa_key_id_t key,
                                 psa_algorithm_t alg,
                                 const uint8_t *input,
                                 size_t input_length,
@@ -1941,8 +1755,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_ENCRYPT_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_CIPHER_ENCRYPT_SID,
+        .key_id = key,
         .alg = alg,
     };
 
@@ -1951,18 +1765,17 @@
         {.base = input, .len = input_length},
     };
     psa_outvec out_vec[] = {
-        {.base = output, .len = output_size},
+        {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_encrypt,
-                          TFM_CRYPTO_CIPHER_ENCRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
     return status;
 #endif /* TFM_CRYPTO_CIPHER_MODULE_DISABLED */
 }
 
-psa_status_t psa_cipher_decrypt(psa_key_id_t key_id,
+psa_status_t psa_cipher_decrypt(psa_key_id_t key,
                                 psa_algorithm_t alg,
                                 const uint8_t *input,
                                 size_t input_length,
@@ -1975,8 +1788,8 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_CIPHER_DECRYPT_SID,
-        .key_id = key_id,
+        .function_id = TFM_CRYPTO_CIPHER_DECRYPT_SID,
+        .key_id = key,
         .alg = alg,
     };
 
@@ -1985,11 +1798,10 @@
         {.base = input, .len = input_length},
     };
     psa_outvec out_vec[] = {
-        {.base = output, .len = output_size},
+        {.base = output, .len = output_size}
     };
 
-    status = API_DISPATCH(tfm_crypto_cipher_decrypt,
-                          TFM_CRYPTO_CIPHER_DECRYPT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
     return status;
@@ -2009,7 +1821,7 @@
 #else
     psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
+        .function_id = TFM_CRYPTO_RAW_KEY_AGREEMENT_SID,
         .alg = alg,
         .key_id = private_key
     };
@@ -2023,8 +1835,7 @@
         {.base = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_raw_key_agreement,
-                          TFM_CRYPTO_RAW_KEY_AGREEMENT);
+    status = API_DISPATCH(in_vec, out_vec);
 
     *output_length = out_vec[0].len;
 
@@ -2038,9 +1849,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_SETUP_SID,
         .alg = alg,
         .op_handle = operation->handle,
     };
@@ -2052,10 +1862,7 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)},
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_setup,
-                          TFM_CRYPTO_KEY_DERIVATION_SETUP);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -2066,9 +1873,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY_SID,
         .capacity = capacity,
         .op_handle = operation->handle,
     };
@@ -2077,10 +1883,7 @@
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_set_capacity,
-                                    TFM_CRYPTO_KEY_DERIVATION_SET_CAPACITY);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
@@ -2093,9 +1896,8 @@
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES_SID,
         .step = step,
         .op_handle = operation->handle,
     };
@@ -2105,24 +1907,20 @@
         {.base = data, .len = data_length},
     };
 
-    status = API_DISPATCH_NO_OUTVEC(tfm_crypto_key_derivation_input_bytes,
-                                    TFM_CRYPTO_KEY_DERIVATION_INPUT_BYTES);
-
-    return status;
+    return API_DISPATCH_NO_OUTVEC(in_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }
 
 psa_status_t psa_key_derivation_output_key(
                                       const psa_key_attributes_t *attributes,
                                       psa_key_derivation_operation_t *operation,
-                                      psa_key_id_t *key_id)
+                                      psa_key_id_t *key)
 {
 #ifdef TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED
     return PSA_ERROR_NOT_SUPPORTED;
 #else
-    psa_status_t status;
     struct tfm_crypto_pack_iovec iov = {
-        .srv_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
+        .function_id = TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY_SID,
         .op_handle = operation->handle,
     };
 
@@ -2132,12 +1930,9 @@
     };
 
     psa_outvec out_vec[] = {
-        {.base = key_id, .len = sizeof(psa_key_id_t)}
+        {.base = key, .len = sizeof(psa_key_id_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_key_derivation_output_key,
-                          TFM_CRYPTO_KEY_DERIVATION_OUTPUT_KEY);
-
-    return status;
+    return API_DISPATCH(in_vec, out_vec);
 #endif /* TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED */
 }