block_cipher: add PSA dispatch if possible

"if possible" means:
- PSA has been initialized
- requested key type is available in PSA

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/include/mbedtls/block_cipher.h b/include/mbedtls/block_cipher.h
index 154ae26..d76d860 100644
--- a/include/mbedtls/block_cipher.h
+++ b/include/mbedtls/block_cipher.h
@@ -24,6 +24,10 @@
 #include "mbedtls/camellia.h"
 #endif
 
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+#include "psa/crypto.h"
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -35,8 +39,24 @@
     MBEDTLS_BLOCK_CIPHER_ID_ARIA,      /**< The Aria cipher. */
 } mbedtls_block_cipher_id_t;
 
+/**
+ * Used internally to indicate whether a context uses legacy or PSA.
+ *
+ * Internal use only.
+ */
+typedef enum {
+    MBEDTLS_BLOCK_CIPHER_ENGINE_LEGACY = 0,
+    MBEDTLS_BLOCK_CIPHER_ENGINE_PSA,
+} mbedtls_block_cipher_engine_t;
+
 typedef struct {
     mbedtls_block_cipher_id_t MBEDTLS_PRIVATE(id);
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+    mbedtls_block_cipher_engine_t engine;
+    psa_cipher_operation_t psa_operation;
+    psa_key_type_t psa_key_type;
+    mbedtls_svc_key_id_t psa_key_id;
+#endif
     union {
         unsigned dummy; /* Make the union non-empty even with no supported algorithms. */
 #if defined(MBEDTLS_AES_C)
diff --git a/library/block_cipher.c b/library/block_cipher.c
index 1118d3a..0fd78ab 100644
--- a/library/block_cipher.c
+++ b/library/block_cipher.c
@@ -11,12 +11,53 @@
 
 #include "common.h"
 
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+#include "psa_crypto_core.h"
+#include "psa/crypto.h"
+#include "psa_util_internal.h"
+#endif
+
 #include "block_cipher_internal.h"
 
 #if defined(MBEDTLS_BLOCK_CIPHER_C)
 
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+static psa_key_type_t psa_key_type_from_cipher_id(mbedtls_cipher_id_t cipher_id)
+{
+    switch (cipher_id) {
+#if defined(MBEDTLS_BLOCK_CIPHER_AES_VIA_PSA)
+        case MBEDTLS_CIPHER_ID_AES:
+            return PSA_KEY_TYPE_AES;
+#endif
+#if defined(MBEDTLS_BLOCK_CIPHER_ARIA_VIA_PSA)
+        case MBEDTLS_CIPHER_ID_ARIA:
+            return PSA_KEY_TYPE_ARIA;
+#endif
+#if defined(MBEDTLS_BLOCK_CIPHER_CAMELLIA_VIA_PSA)
+        case MBEDTLS_CIPHER_ID_CAMELLIA:
+            return PSA_KEY_TYPE_CAMELLIA;
+#endif
+        default:
+            return PSA_KEY_TYPE_NONE;
+    }
+}
+
+int mbedtls_cipher_error_from_psa(psa_status_t status)
+{
+    return PSA_TO_MBEDTLS_ERR_LIST(status, psa_to_cipher_errors,
+                                   psa_generic_status_to_mbedtls);
+}
+#endif /* MBEDTLS_BLOCK_CIPHER_SOME_PSA */
+
 void mbedtls_block_cipher_free(mbedtls_block_cipher_context_t *ctx)
 {
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+    if (ctx->engine == MBEDTLS_BLOCK_CIPHER_ENGINE_PSA) {
+        psa_cipher_abort(&ctx->psa_operation);
+        psa_destroy_key(ctx->psa_key_id);
+        return;
+    }
+#endif
     switch (ctx->id) {
 #if defined(MBEDTLS_AES_C)
         case MBEDTLS_BLOCK_CIPHER_ID_AES:
@@ -42,6 +83,17 @@
 int mbedtls_block_cipher_setup(mbedtls_block_cipher_context_t *ctx,
                                mbedtls_cipher_id_t cipher_id)
 {
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+    if (psa_can_do_cipher(cipher_id)) {
+        ctx->psa_key_type = psa_key_type_from_cipher_id(cipher_id);
+        if (ctx->psa_key_type != PSA_KEY_TYPE_NONE) {
+            ctx->engine = MBEDTLS_BLOCK_CIPHER_ENGINE_PSA;
+            return 0;
+        }
+    }
+    ctx->engine = MBEDTLS_BLOCK_CIPHER_ENGINE_LEGACY;
+#endif
+
     switch (cipher_id) {
 #if defined(MBEDTLS_AES_C)
         case MBEDTLS_CIPHER_ID_AES:
@@ -70,6 +122,32 @@
                                 const unsigned char *key,
                                 unsigned key_bitlen)
 {
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+    if (ctx->engine == MBEDTLS_BLOCK_CIPHER_ENGINE_PSA) {
+        psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
+        psa_status_t status;
+
+        psa_set_key_type(&key_attr, ctx->psa_key_type);
+        psa_set_key_bits(&key_attr, key_bitlen);
+        psa_set_key_algorithm(&key_attr, PSA_ALG_ECB_NO_PADDING);
+        psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_ENCRYPT);
+
+        status = psa_import_key(&key_attr, key, key_bitlen/8, &ctx->psa_key_id);
+        if (status != PSA_SUCCESS) {
+            return mbedtls_cipher_error_from_psa(status);
+        }
+        psa_reset_key_attributes(&key_attr);
+
+        status = psa_cipher_encrypt_setup(&ctx->psa_operation, ctx->psa_key_id,
+                                          PSA_ALG_ECB_NO_PADDING);
+        if (status != PSA_SUCCESS) {
+            return mbedtls_cipher_error_from_psa(status);
+        }
+
+        return 0;
+    }
+#endif /* MBEDTLS_BLOCK_CIPHER_SOME_PSA */
+
     switch (ctx->id) {
 #if defined(MBEDTLS_AES_C)
         case MBEDTLS_BLOCK_CIPHER_ID_AES:
@@ -92,6 +170,20 @@
                                  const unsigned char input[16],
                                  unsigned char output[16])
 {
+#if defined(MBEDTLS_BLOCK_CIPHER_SOME_PSA)
+    if (ctx->engine == MBEDTLS_BLOCK_CIPHER_ENGINE_PSA) {
+        psa_status_t status;
+        size_t olen;
+
+        status = psa_cipher_encrypt(ctx->psa_key_id, PSA_ALG_ECB_NO_PADDING,
+                                    input, 16, output, 16, &olen);
+        if (status != PSA_SUCCESS) {
+            return mbedtls_cipher_error_from_psa(status);
+        }
+        return 0;
+    }
+#endif /* MBEDTLS_BLOCK_CIPHER_SOME_PSA */
+
     switch (ctx->id) {
 #if defined(MBEDTLS_AES_C)
         case MBEDTLS_BLOCK_CIPHER_ID_AES: