Crypto: Support optional inputs in aead_update and aead_update_ad

The PSA APIs psa_aead_update_ad() and psa_aead_update() support
the inputs as optional values (i.e. NULL/zero-length buffers),
hence make sure this case is handled correctly in IPC mode.

Signed-off-by: Antonio de Angelis <antonio.deangelis@arm.com>
Change-Id: I1ec39359557921769ca475e6c69fab26cb6ca901
diff --git a/interface/src/tfm_crypto_ipc_api.c b/interface/src/tfm_crypto_ipc_api.c
index 29fbf95..cd43112 100644
--- a/interface/src/tfm_crypto_ipc_api.c
+++ b/interface/src/tfm_crypto_ipc_api.c
@@ -1012,6 +1012,11 @@
         .op_handle = operation->handle,
     };
 
+    /* Sanitize the optional input */
+    if ((input == NULL) && (input_length != 0)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
@@ -1020,8 +1025,12 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_update_ad,
-                          TFM_CRYPTO_AEAD_UPDATE_AD);
+    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, IOVEC_LEN(out_vec));
     return status;
 }
 
@@ -1038,6 +1047,11 @@
         .op_handle = operation->handle,
     };
 
+    /* Sanitize the optional input */
+    if ((input == NULL) && (input_length != 0)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
@@ -1047,8 +1061,12 @@
         {.base  = output, .len = output_size},
     };
 
-    status = API_DISPATCH(tfm_crypto_aead_update,
-                          TFM_CRYPTO_AEAD_UPDATE);
+    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, IOVEC_LEN(out_vec));
 
     *output_length = out_vec[1].len;
     return status;
diff --git a/secure_fw/partitions/crypto/crypto_aead.c b/secure_fw/partitions/crypto/crypto_aead.c
index 85ec46b..ded596c 100644
--- a/secure_fw/partitions/crypto/crypto_aead.c
+++ b/secure_fw/partitions/crypto/crypto_aead.c
@@ -468,7 +468,7 @@
     psa_status_t status = PSA_SUCCESS;
     psa_aead_operation_t *operation = NULL;
 
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 2, 2);
+    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 1, 2, out_len, 2, 2);
 
     if ((in_vec[0].len != sizeof(struct tfm_crypto_pack_iovec)) ||
         (out_vec[0].len != sizeof(uint32_t))) {
@@ -517,7 +517,7 @@
     psa_status_t status = PSA_SUCCESS;
     psa_aead_operation_t *operation = NULL;
 
-    CRYPTO_IN_OUT_LEN_VALIDATE(in_len, 2, 2, out_len, 1, 1);
+    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))) {
diff --git a/secure_fw/partitions/crypto/tfm_crypto_secure_api.c b/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
index 8253a0c..a0cfed7 100644
--- a/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
+++ b/secure_fw/partitions/crypto/tfm_crypto_secure_api.c
@@ -1198,6 +1198,11 @@
         .op_handle = operation->handle,
     };
 
+    /* Sanitize the optional input */
+    if ((input == NULL) && (input_length != 0)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
@@ -1206,8 +1211,18 @@
         {.base = &(operation->handle), .len = sizeof(uint32_t)}
     };
 
+#ifdef TFM_PSA_API
+    size_t in_len = ARRAY_SIZE(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));
+#else
     status = API_DISPATCH(tfm_crypto_aead_update_ad,
                           TFM_CRYPTO_AEAD_UPDATE_AD);
+#endif
+
     return status;
 #endif /* TFM_CRYPTO_AEAD_MODULE_DISABLED */
 }
@@ -1228,6 +1243,11 @@
         .op_handle = operation->handle,
     };
 
+    /* Sanitize the optional input */
+    if ((input == NULL) && (input_length != 0)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     psa_invec in_vec[] = {
         {.base = &iov, .len = sizeof(struct tfm_crypto_pack_iovec)},
         {.base = input, .len = input_length}
@@ -1237,8 +1257,17 @@
         {.base  = output, .len = output_size},
     };
 
+#ifdef TFM_PSA_API
+    size_t in_len = ARRAY_SIZE(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));
+#else
     status = API_DISPATCH(tfm_crypto_aead_update,
                           TFM_CRYPTO_AEAD_UPDATE);
+#endif
 
     *output_length = out_vec[1].len;
     return status;