Add support for SHA-3 in PSA

Signed-off-by: Dave Rodgman <dave.rodgman@arm.com>
diff --git a/include/mbedtls/config_psa.h b/include/mbedtls/config_psa.h
index b7e8947..b1d3862 100644
--- a/include/mbedtls/config_psa.h
+++ b/include/mbedtls/config_psa.h
@@ -113,6 +113,10 @@
 #define PSA_WANT_ALG_SHA_512 1
 #endif
 
+#if defined(MBEDTLS_SHA3_C)
+#define PSA_WANT_ALG_SHA_3 1
+#endif
+
 
 /****************************************************************/
 /* Require built-in implementations based on PSA requirements */
@@ -270,6 +274,11 @@
 #define MBEDTLS_SHA512_C
 #endif
 
+#if defined(PSA_WANT_ALG_SHA_3) && !defined(MBEDTLS_PSA_ACCEL_ALG_SHA_3)
+#define MBEDTLS_PSA_BUILTIN_ALG_SHA_3 1
+#define MBEDTLS_SHA3_C
+#endif
+
 #if defined(PSA_WANT_ALG_PBKDF2_HMAC)
 #if !defined(MBEDTLS_PSA_ACCEL_ALG_PBKDF2_HMAC)
 #define MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC 1
@@ -865,6 +874,11 @@
 #define PSA_WANT_ALG_SHA_512 1
 #endif
 
+#if defined(MBEDTLS_SHA3_C)
+#define MBEDTLS_PSA_BUILTIN_ALG_SHA_3 1
+#define PSA_WANT_ALG_SHA_3 1
+#endif
+
 #if defined(MBEDTLS_AES_C)
 #define PSA_WANT_KEY_TYPE_AES 1
 #define MBEDTLS_PSA_BUILTIN_KEY_TYPE_AES 1
diff --git a/include/psa/crypto_builtin_primitives.h b/include/psa/crypto_builtin_primitives.h
index f3e438d..ea7da9a 100644
--- a/include/psa/crypto_builtin_primitives.h
+++ b/include/psa/crypto_builtin_primitives.h
@@ -45,6 +45,7 @@
 #include "mbedtls/sha1.h"
 #include "mbedtls/sha256.h"
 #include "mbedtls/sha512.h"
+#include "mbedtls/sha3.h"
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_MD5) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_RIPEMD160) || \
@@ -52,7 +53,8 @@
     defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_224) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_256) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_384) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_512)
+    defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_512) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
 #define MBEDTLS_PSA_BUILTIN_HASH
 #endif
 
@@ -77,6 +79,9 @@
         defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_384)
         mbedtls_sha512_context sha512;
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
+        mbedtls_sha3_context sha3;
+#endif
     } MBEDTLS_PRIVATE(ctx);
 } mbedtls_psa_hash_operation_t;
 
diff --git a/include/psa/crypto_config.h b/include/psa/crypto_config.h
index 9da28de..e79f217 100644
--- a/include/psa/crypto_config.h
+++ b/include/psa/crypto_config.h
@@ -86,6 +86,7 @@
 #define PSA_WANT_ALG_SHA_256                    1
 #define PSA_WANT_ALG_SHA_384                    1
 #define PSA_WANT_ALG_SHA_512                    1
+#define PSA_WANT_ALG_SHA_3                      1
 #define PSA_WANT_ALG_STREAM_CIPHER              1
 #define PSA_WANT_ALG_TLS12_PRF                  1
 #define PSA_WANT_ALG_TLS12_PSK_TO_MS            1
diff --git a/library/psa_crypto_hash.c b/library/psa_crypto_hash.c
index ddf7094..44df552 100644
--- a/library/psa_crypto_hash.c
+++ b/library/psa_crypto_hash.c
@@ -74,6 +74,14 @@
             mbedtls_sha512_free(&operation->ctx.sha512);
             break;
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
+        case PSA_ALG_SHA3_224:
+        case PSA_ALG_SHA3_256:
+        case PSA_ALG_SHA3_384:
+        case PSA_ALG_SHA3_512:
+            mbedtls_sha3_free(&operation->ctx.sha3);
+            break;
+#endif
         default:
             return PSA_ERROR_BAD_STATE;
     }
@@ -135,6 +143,24 @@
             ret = mbedtls_sha512_starts(&operation->ctx.sha512, 0);
             break;
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
+        case PSA_ALG_SHA3_224:
+            mbedtls_sha3_init(&operation->ctx.sha3);
+            ret = mbedtls_sha3_starts(&operation->ctx.sha3, MBEDTLS_SHA3_224);
+            break;
+        case PSA_ALG_SHA3_256:
+            mbedtls_sha3_init(&operation->ctx.sha3);
+            ret = mbedtls_sha3_starts(&operation->ctx.sha3, MBEDTLS_SHA3_256);
+            break;
+        case PSA_ALG_SHA3_384:
+            mbedtls_sha3_init(&operation->ctx.sha3);
+            ret = mbedtls_sha3_starts(&operation->ctx.sha3, MBEDTLS_SHA3_384);
+            break;
+        case PSA_ALG_SHA3_512:
+            mbedtls_sha3_init(&operation->ctx.sha3);
+            ret = mbedtls_sha3_starts(&operation->ctx.sha3, MBEDTLS_SHA3_512);
+            break;
+#endif
         default:
             return PSA_ALG_IS_HASH(alg) ?
                    PSA_ERROR_NOT_SUPPORTED :
@@ -197,6 +223,15 @@
                                  &source_operation->ctx.sha512);
             break;
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
+        case PSA_ALG_SHA3_224:
+        case PSA_ALG_SHA3_256:
+        case PSA_ALG_SHA3_384:
+        case PSA_ALG_SHA3_512:
+            mbedtls_sha3_clone(&target_operation->ctx.sha3,
+                               &source_operation->ctx.sha3);
+            break;
+#endif
         default:
             (void) source_operation;
             (void) target_operation;
@@ -257,6 +292,15 @@
                                         input, input_length);
             break;
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
+        case PSA_ALG_SHA3_224:
+        case PSA_ALG_SHA3_256:
+        case PSA_ALG_SHA3_384:
+        case PSA_ALG_SHA3_512:
+            ret = mbedtls_sha3_update(&operation->ctx.sha3,
+                                        input, input_length);
+            break;
+#endif
         default:
             (void) input;
             (void) input_length;
@@ -327,6 +371,14 @@
             ret = mbedtls_sha512_finish(&operation->ctx.sha512, hash);
             break;
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_3)
+        case PSA_ALG_SHA3_224:
+        case PSA_ALG_SHA3_256:
+        case PSA_ALG_SHA3_384:
+        case PSA_ALG_SHA3_512:
+            ret = mbedtls_sha3_finish(&operation->ctx.sha3, hash, hash_size);
+            break;
+#endif
         default:
             (void) hash;
             return PSA_ERROR_BAD_STATE;