From dad36fa855e17a039e1f8f1a5671093a8255808b Mon Sep 17 00:00:00 2001
From: mohammad1603 <mohammad.abomokh@arm.com>
Date: Wed, 9 May 2018 02:24:42 -0700
Subject: [PATCH] add Key and Algorithm validation

---
 include/psa/crypto.h |  1 +
 library/psa_crypto.c | 38 +++++++++++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index cd86080c1..deeab4a64 100755
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -143,6 +143,7 @@ typedef uint32_t psa_key_type_t;
 #define PSA_KEY_TYPE_CATEGORY_MASK              ((psa_key_type_t)0x7e000000)
 #define PSA_KEY_TYPE_RAW_DATA                   ((psa_key_type_t)0x02000000)
 #define PSA_KEY_TYPE_CATEGORY_SYMMETRIC         ((psa_key_type_t)0x04000000)
+#define PSA_KEY_TYPE_CATEGORY_CIPHER            ((psa_key_type_t)0x04000000)
 #define PSA_KEY_TYPE_CATEGORY_ASYMMETRIC        ((psa_key_type_t)0x06000000)
 #define PSA_KEY_TYPE_PAIR_FLAG                  ((psa_key_type_t)0x01000000)
 
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 33e265766..7d70d534a 100755
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1488,7 +1488,8 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
     size_t key_bits;
     const mbedtls_cipher_info_t *cipher_info = NULL;
     unsigned char tag[16];
-
+    mbedtls_cipher_id_t cipher_id;
+    
     if( ciphertext_size < ( plaintext_length + sizeof( tag ) ) )
             return( PSA_ERROR_INVALID_ARGUMENT );
 
@@ -1497,6 +1498,15 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
         return( status );
     slot = &global_data.key_slots[key];
 
+    if ( key_type == PSA_KEY_TYPE_AES )
+    {
+        cipher_id = MBEDTLS_CIPHER_ID_AES;
+    }
+    else
+    {
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+
     //TODO: check key policy
 
     cipher_info = mbedtls_cipher_info_from_psa( alg, key_type, key_bits );
@@ -1507,13 +1517,11 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
            && PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type ) == 16 )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    operation->block_size = cipher_info->block_size;
-
     if( alg == PSA_ALG_GCM )
     {
         mbedtls_gcm_context gcm;
         mbedtls_gcm_init( &gcm );
-        ret = mbedtls_gcm_setkey( &gcm, cipher_info->base->cipher, 
+        ret = mbedtls_gcm_setkey( &gcm, cipher_id, 
                                 ( const unsigned char * )slot->data.raw.data, key_bits );
         if( ret != 0 )
         {
@@ -1541,7 +1549,7 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
             return( PSA_ERROR_INVALID_ARGUMENT );
 
         mbedtls_ccm_init( &ccm );
-        ret = mbedtls_ccm_setkey( &ccm, cipher_info->base->cipher, 
+        ret = mbedtls_ccm_setkey( &ccm, cipher_id, 
                                   slot->data.raw.data, key_bits );
         if( ret != 0 )
         {
@@ -1551,7 +1559,7 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
         ret = mbedtls_ccm_encrypt_and_tag( &ccm, plaintext_length, 
                        nonce , nonce_length, additional_data,
                        additional_data_length,
-                       plaintext, ciphertext, sizeof( tag ), tag );
+                       plaintext, ciphertext, tag, sizeof( tag ) );
         if( ret != 0 )
         {
             mbedtls_ccm_free( &ccm );
@@ -1585,6 +1593,7 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
     size_t key_bits;
     const mbedtls_cipher_info_t *cipher_info = NULL;
     unsigned char tag[16];
+    mbedtls_cipher_id_t cipher_id;
 
     if( plaintext_size < ciphertext_length )
             return( PSA_ERROR_INVALID_ARGUMENT );
@@ -1594,6 +1603,15 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
         return( status );
     slot = &global_data.key_slots[key];
 
+    if ( key_type == PSA_KEY_TYPE_AES )
+    {
+        cipher_id = MBEDTLS_CIPHER_ID_AES;
+    }
+    else
+    {
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+
     //TODO: check key policy
 
     cipher_info = mbedtls_cipher_info_from_psa( alg, key_type, key_bits );
@@ -1604,14 +1622,12 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
            && PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type ) == 16 )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    operation->block_size = cipher_info->block_size;
-
     if( alg == PSA_ALG_GCM )
     {
         mbedtls_gcm_context gcm;
 
         mbedtls_gcm_init( &gcm );
-        ret = mbedtls_gcm_setkey( &gcm, cipher_info->base->cipher, 
+        ret = mbedtls_gcm_setkey( &gcm, cipher_id, 
                                   slot->data.raw.data, key_bits );
         if( ret != 0 )
         {
@@ -1639,7 +1655,7 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
             return( PSA_ERROR_INVALID_ARGUMENT );
 
         mbedtls_ccm_init( &ccm );
-        ret = mbedtls_ccm_setkey( &ccm, cipher_info->base->cipher, 
+        ret = mbedtls_ccm_setkey( &ccm, cipher_id, 
                                   slot->data.raw.data, key_bits );
         if( ret != 0 )
         {
@@ -1649,7 +1665,7 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
         ret = mbedtls_ccm_auth_decrypt( &ccm, ciphertext_length, 
                        nonce , nonce_length, additional_data,
                        additional_data_length, ciphertext ,
-                       plaintext, sizeof( tag ), tag );
+                       plaintext, tag, sizeof( tag ) );
         if( ret != 0 )
         {
             mbedtls_ccm_free( &ccm );