From 0ff4b0f7f98740197944e67303c244aba2578fa1 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Tue, 19 Jun 2018 21:31:50 +0200
Subject: [PATCH] psa_import_key: validate symmetric key size

When importing a symmetric key, validate that the key size is valid
for the given key type.

Non-supported key types may no longer be imported.
---
 library/psa_crypto.c                    | 61 +++++++++++++++++++++++--
 tests/suites/test_suite_psa_crypto.data | 16 +++++++
 2 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e41e51287..4d2f8d05b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -346,6 +346,57 @@ static psa_ecc_curve_t mbedtls_ecc_group_to_psa( mbedtls_ecp_group_id grpid )
     }
 }
 
+static psa_status_t prepare_raw_data_slot( psa_key_type_t type,
+                                           size_t bits,
+                                           struct raw_data *raw )
+{
+    /* Check that the bit size is acceptable for the key type */
+    switch( type )
+    {
+        case PSA_KEY_TYPE_RAW_DATA:
+#if defined(MBEDTLS_MD_C)
+        case PSA_KEY_TYPE_HMAC:
+#endif
+            break;
+#if defined(MBEDTLS_AES_C)
+        case PSA_KEY_TYPE_AES:
+            if( bits != 128 && bits != 192 && bits != 256 )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif
+#if defined(MBEDTLS_CAMELLIA_C)
+        case PSA_KEY_TYPE_CAMELLIA:
+            if( bits != 128 && bits != 192 && bits != 256 )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif
+#if defined(MBEDTLS_DES_C)
+        case PSA_KEY_TYPE_DES:
+            if( bits != 64 && bits != 128 && bits != 192 )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif
+#if defined(MBEDTLS_ARC4_C)
+        case PSA_KEY_TYPE_ARC4:
+            if( bits < 8 || bits > 2048 )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif
+        default:
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+
+    /* Allocate memory for the key */
+    raw->bytes = PSA_BITS_TO_BYTES( bits );
+    raw->data = mbedtls_calloc( 1, raw->bytes );
+    if( raw->data == NULL )
+    {
+        raw->bytes = 0;
+        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    }
+    return( PSA_SUCCESS );
+}
+
 psa_status_t psa_import_key( psa_key_slot_t key,
                              psa_key_type_t type,
                              const uint8_t *data,
@@ -361,14 +412,16 @@ psa_status_t psa_import_key( psa_key_slot_t key,
 
     if( PSA_KEY_TYPE_IS_RAW_BYTES( type ) )
     {
+        psa_status_t status;
         /* Ensure that a bytes-to-bit conversion won't overflow. */
         if( data_length > SIZE_MAX / 8 )
             return( PSA_ERROR_NOT_SUPPORTED );
-        slot->data.raw.data = mbedtls_calloc( 1, data_length );
-        if( slot->data.raw.data == NULL )
-            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+        status = prepare_raw_data_slot( type,
+                                        PSA_BYTES_TO_BITS( data_length ),
+                                        &slot->data.raw );
+        if( status != PSA_SUCCESS )
+            return( status );
         memcpy( slot->data.raw.data, data, data_length );
-        slot->data.raw.bytes = data_length;
     }
     else
 #if defined(MBEDTLS_PK_PARSE_C)
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index ca8ea2336..00add7d64 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -13,6 +13,22 @@ import_export:"2a":PSA_KEY_TYPE_RAW_DATA:PSA_ALG_CBC_BASE | PSA_ALG_BLOCK_CIPHER
 PSA import/export raw: 2 bytes, buffer too small
 import_export:"2a2b":PSA_KEY_TYPE_RAW_DATA:PSA_ALG_CBC_BASE | PSA_ALG_BLOCK_CIPHER_PAD_NONE:PSA_KEY_USAGE_EXPORT:16:-1:PSA_ERROR_BUFFER_TOO_SMALL:1
 
+PSA import/export AES-128
+depends_on:MBEDTLS_AES_C
+import_export:"0123456789abcdef0123456789abcdef":PSA_KEY_TYPE_AES:PSA_ALG_CTR:PSA_KEY_USAGE_EXPORT:128:0:PSA_SUCCESS:1
+
+PSA import/export AES-192
+depends_on:MBEDTLS_AES_C
+import_export:"0123456789abcdef0123456789abcdef0123456789abcdef":PSA_KEY_TYPE_AES:PSA_ALG_CTR:PSA_KEY_USAGE_EXPORT:192:0:PSA_SUCCESS:1
+
+PSA import/export AES-256
+depends_on:MBEDTLS_AES_C
+import_export:"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef":PSA_KEY_TYPE_AES:PSA_ALG_CTR:PSA_KEY_USAGE_EXPORT:256:0:PSA_SUCCESS:1
+
+PSA import AES: bad key size
+depends_on:MBEDTLS_AES_C
+import:"0123456789abcdef":PSA_KEY_TYPE_AES:PSA_ERROR_INVALID_ARGUMENT
+
 PSA import/export RSA public key: good, 1024-bit
 depends_on:MBEDTLS_PK_C:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C
 import_export:"30819f300d06092a864886f70d010101050003818d0030818902818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc30203010001":PSA_KEY_TYPE_RSA_PUBLIC_KEY:PSA_ALG_CBC_BASE | PSA_ALG_BLOCK_CIPHER_PAD_NONE:PSA_KEY_USAGE_EXPORT:1024:0:PSA_SUCCESS:1