From c6b6907066b016a2e2babb9ea6940d6a8c202575 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Tue, 20 Nov 2018 21:42:52 +0100
Subject: [PATCH] Make library init and deinit more robust to errors

Allow mbedtls_psa_crypto_free to be called twice, or without a prior
call to psa_crypto_init. Keep track of the initialization state more
precisely in psa_crypto_init so that mbedtls_psa_crypto_free knows
what to do.
---
 library/psa_crypto.c                          | 51 ++++++++++++++-----
 tests/suites/test_suite_psa_crypto_init.data  |  6 +++
 .../test_suite_psa_crypto_init.function       | 13 +++++
 3 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 291dcdb0d..4c0ac1213 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -146,12 +146,21 @@ static int key_type_is_raw_bytes( psa_key_type_t type )
     return( PSA_KEY_TYPE_IS_UNSTRUCTURED( type ) );
 }
 
+enum rng_state
+{
+    RNG_NOT_INITIALIZED = 0,
+    RNG_INITIALIZED,
+    RNG_SEEDED,
+};
+
 typedef struct
 {
-    int initialized;
     mbedtls_entropy_context entropy;
     mbedtls_ctr_drbg_context ctr_drbg;
     key_slot_t key_slots[PSA_KEY_SLOT_COUNT];
+    unsigned initialized : 1;
+    enum rng_state rng_state : 2;
+    unsigned key_slots_initialized : 1;
 } psa_global_data_t;
 
 static psa_global_data_t global_data;
@@ -4433,18 +4442,26 @@ void mbedtls_psa_crypto_free( void )
     psa_key_slot_t key;
     key_slot_t *slot;
     psa_status_t status;
-
-    for( key = 1; key <= PSA_KEY_SLOT_COUNT; key++ )
+    if( global_data.key_slots_initialized )
     {
-        status = psa_get_key_slot( key, &slot );
-        if( status != PSA_SUCCESS )
-            continue;
-        psa_remove_key_data_from_memory( slot );
-        /* Zeroize the slot to wipe metadata such as policies. */
-        mbedtls_zeroize( slot, sizeof( *slot ) );
+        for( key = 1; key <= PSA_KEY_SLOT_COUNT; key++ )
+        {
+            status = psa_get_key_slot( key, &slot );
+            if( status != PSA_SUCCESS )
+                continue;
+            psa_remove_key_data_from_memory( slot );
+            /* Zeroize the slot to wipe metadata such as policies. */
+            mbedtls_zeroize( slot, sizeof( *slot ) );
+        }
     }
-    mbedtls_ctr_drbg_free( &global_data.ctr_drbg );
-    mbedtls_entropy_free( &global_data.entropy );
+    if( global_data.rng_state != RNG_NOT_INITIALIZED )
+    {
+        mbedtls_ctr_drbg_free( &global_data.ctr_drbg );
+        mbedtls_entropy_free( &global_data.entropy );
+    }
+    /* Wipe all remaining data, including configuration.
+     * In particular, this sets all state indicator to the value
+     * indicating "uninitialized". */
     mbedtls_zeroize( &global_data, sizeof( global_data ) );
 }
 
@@ -4453,20 +4470,30 @@ psa_status_t psa_crypto_init( void )
     int ret;
     const unsigned char drbg_seed[] = "PSA";
 
+    /* Double initialization is explicitly allowed. */
     if( global_data.initialized != 0 )
         return( PSA_SUCCESS );
 
     mbedtls_zeroize( &global_data, sizeof( global_data ) );
+
+    /* Initialize the random generator. */
     mbedtls_entropy_init( &global_data.entropy );
     mbedtls_ctr_drbg_init( &global_data.ctr_drbg );
-
+    global_data.rng_state = RNG_INITIALIZED;
     ret = mbedtls_ctr_drbg_seed( &global_data.ctr_drbg,
                                  mbedtls_entropy_func,
                                  &global_data.entropy,
                                  drbg_seed, sizeof( drbg_seed ) - 1 );
     if( ret != 0 )
         goto exit;
+    global_data.rng_state = RNG_SEEDED;
 
+    /* Initialize the key slots. Zero-initialization has made all key
+     * slots empty, so there is nothing to do. In a future version we will
+     * load data from storage. */
+    global_data.key_slots_initialized = 1;
+
+    /* All done. */
     global_data.initialized = 1;
 
 exit:
diff --git a/tests/suites/test_suite_psa_crypto_init.data b/tests/suites/test_suite_psa_crypto_init.data
index 61f067d06..e44111814 100644
--- a/tests/suites/test_suite_psa_crypto_init.data
+++ b/tests/suites/test_suite_psa_crypto_init.data
@@ -1,6 +1,12 @@
 PSA init/deinit
 init_deinit:2
 
+PSA deinit without init
+deinit_without_init:0
+
+PSA deinit twice
+deinit_without_init:1
+
 No random without init
 validate_module_init_generate_random:0
 
diff --git a/tests/suites/test_suite_psa_crypto_init.function b/tests/suites/test_suite_psa_crypto_init.function
index 7fccc13d2..7cb10c0a1 100644
--- a/tests/suites/test_suite_psa_crypto_init.function
+++ b/tests/suites/test_suite_psa_crypto_init.function
@@ -29,6 +29,19 @@ void init_deinit( int count )
 }
 /* END_CASE */
 
+/* BEGIN_CASE */
+void deinit_without_init( int count )
+{
+    int i;
+    for( i = 0; i < count; i++ )
+    {
+        TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
+        mbedtls_psa_crypto_free( );
+    }
+    mbedtls_psa_crypto_free( );
+}
+/* END_CASE */
+
 /* BEGIN_CASE */
 void validate_module_init_generate_random( int count )
 {