Merge pull request #57 from Patater/check-generator-validity

psa: Check generator validity before read
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 325abde..84c0e88 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3621,6 +3621,12 @@
 psa_status_t psa_get_generator_capacity(const psa_crypto_generator_t *generator,
                                         size_t *capacity)
 {
+    if( generator->alg == 0 )
+    {
+        /* This is a blank generator. */
+        return PSA_ERROR_BAD_STATE;
+    }
+
     *capacity = generator->capacity;
     return( PSA_SUCCESS );
 }
@@ -3850,6 +3856,12 @@
 {
     psa_status_t status;
 
+    if( generator->alg == 0 )
+    {
+        /* This is a blank generator. */
+        return PSA_ERROR_BAD_STATE;
+    }
+
     if( output_length > generator->capacity )
     {
         generator->capacity = 0;
@@ -3858,11 +3870,10 @@
         status = PSA_ERROR_INSUFFICIENT_DATA;
         goto exit;
     }
-    if( output_length == 0 &&
-        generator->capacity == 0 && generator->alg == 0 )
+    if( output_length == 0 && generator->capacity == 0 )
     {
-        /* Edge case: this is a blank or finished generator, and 0
-         * bytes were requested. The right error in this case could
+        /* Edge case: this is a finished generator, and 0 bytes
+         * were requested. The right error in this case could
          * be either INSUFFICIENT_CAPACITY or BAD_STATE. Return
          * INSUFFICIENT_CAPACITY, which is right for a finished
          * generator, for consistency with the case when
@@ -3911,7 +3922,13 @@
 exit:
     if( status != PSA_SUCCESS )
     {
+        /* Preserve the algorithm upon errors, but clear all sensitive state.
+         * This allows us to differentiate between exhausted generators and
+         * blank generators, so we can return PSA_ERROR_BAD_STATE on blank
+         * generators. */
+        psa_algorithm_t alg = generator->alg;
         psa_generator_abort( generator );
+        generator->alg = alg;
         memset( output, '!', output_length );
     }
     return( status );
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 92b6fb0..cccc870 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -3519,13 +3519,13 @@
 
     memset( &zero, 0, sizeof( zero ) );
 
-    /* A default generator should have no capacity. */
-    PSA_ASSERT( psa_get_generator_capacity( &func, &capacity ) );
-    TEST_EQUAL( capacity, 0 );
-    PSA_ASSERT( psa_get_generator_capacity( &init, &capacity ) );
-    TEST_EQUAL( capacity, 0 );
-    PSA_ASSERT( psa_get_generator_capacity( &zero, &capacity ) );
-    TEST_EQUAL( capacity, 0 );
+    /* A default generator should not be able to report its capacity. */
+    TEST_EQUAL( psa_get_generator_capacity( &func, &capacity ),
+                PSA_ERROR_BAD_STATE );
+    TEST_EQUAL( psa_get_generator_capacity( &init, &capacity ),
+                PSA_ERROR_BAD_STATE );
+    TEST_EQUAL( psa_get_generator_capacity( &zero, &capacity ),
+                PSA_ERROR_BAD_STATE );
 
     /* A default generator should be abortable without error. */
     PSA_ASSERT( psa_generator_abort(&func) );
@@ -3632,18 +3632,18 @@
     psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
 
     TEST_ASSERT( psa_generator_read( &generator, output_buffer, buffer_size )
-                 == PSA_ERROR_INSUFFICIENT_DATA ); // should be PSA_ERROR_BAD_STATE:#183
+                 == PSA_ERROR_BAD_STATE );
 
     TEST_ASSERT( psa_get_generator_capacity( &generator, &capacity )
-                 == PSA_SUCCESS ); // should be PSA_ERROR_BAD_STATE:#183
+                 == PSA_ERROR_BAD_STATE );
 
     PSA_ASSERT( psa_generator_abort( &generator ) );
 
     TEST_ASSERT( psa_generator_read( &generator, output_buffer, buffer_size )
-                 == PSA_ERROR_INSUFFICIENT_DATA ); // should be PSA_ERROR_BAD_STATE:#183
+                 == PSA_ERROR_BAD_STATE );
 
     TEST_ASSERT( psa_get_generator_capacity( &generator, &capacity )
-                 == PSA_SUCCESS );// should be PSA_ERROR_BAD_STATE:#183
+                 == PSA_ERROR_BAD_STATE );
 
 exit:
     psa_generator_abort( &generator );