Fix definition of MBEDTLS_MD_MAX_SIZE and MBEDTLS_MD_MAX_BLOCK_SIZE

Signed-off-by: Dave Rodgman <dave.rodgman@arm.com>
diff --git a/include/mbedtls/md.h b/include/mbedtls/md.h
index 67a5bd6..c040a48 100644
--- a/include/mbedtls/md.h
+++ b/include/mbedtls/md.h
@@ -169,13 +169,13 @@
  * and legacy, then assume the buffer's size is PSA_HASH_MAX_SIZE in another
  * part of the code based on PSA.
  */
-#if defined(MBEDTLS_MD_CAN_SHA512) || defined(MBEDTLS_SHA3_C)
+#if defined(MBEDTLS_MD_CAN_SHA512) || defined(MBEDTLS_MD_CAN_SHA3_512)
 #define MBEDTLS_MD_MAX_SIZE         64  /* longest known is SHA512 */
-#elif defined(MBEDTLS_MD_CAN_SHA384)
+#elif defined(MBEDTLS_MD_CAN_SHA384) || defined(MBEDTLS_MD_CAN_SHA3_384)
 #define MBEDTLS_MD_MAX_SIZE         48  /* longest known is SHA384 */
-#elif defined(MBEDTLS_MD_CAN_SHA256)
+#elif defined(MBEDTLS_MD_CAN_SHA256) || defined(MBEDTLS_MD_CAN_SHA3_256)
 #define MBEDTLS_MD_MAX_SIZE         32  /* longest known is SHA256 */
-#elif defined(MBEDTLS_MD_CAN_SHA224)
+#elif defined(MBEDTLS_MD_CAN_SHA224) || defined(MBEDTLS_MD_CAN_SHA3_224)
 #define MBEDTLS_MD_MAX_SIZE         28  /* longest known is SHA224 */
 #else
 #define MBEDTLS_MD_MAX_SIZE         20  /* longest known is SHA1 or RIPE MD-160
@@ -184,8 +184,14 @@
 
 #if defined(MBEDTLS_MD_CAN_SHA3_224)
 #define MBEDTLS_MD_MAX_BLOCK_SIZE         144 /* the longest known is SHA3-224 */
+#elif defined(MBEDTLS_MD_CAN_SHA3_256)
+#define MBEDTLS_MD_MAX_BLOCK_SIZE         136
 #elif defined(MBEDTLS_MD_CAN_SHA512)
 #define MBEDTLS_MD_MAX_BLOCK_SIZE         128
+#elif defined(MBEDTLS_MD_CAN_SHA3_384)
+#define MBEDTLS_MD_MAX_BLOCK_SIZE         104
+#elif defined(MBEDTLS_MD_CAN_SHA3_512)
+#define MBEDTLS_MD_MAX_BLOCK_SIZE         72
 #else
 #define MBEDTLS_MD_MAX_BLOCK_SIZE         64
 #endif