From 7220cae93c9dc5a820c18e1a9a2329f97c6256ec Mon Sep 17 00:00:00 2001
From: Paul Elliott <paul.elliott@arm.com>
Date: Tue, 22 Jun 2021 17:25:57 +0100
Subject: [PATCH] Ensure generate nonce unavailable in decrypt

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
---
 include/psa/crypto_struct.h                 |  3 ++-
 library/psa_crypto.c                        |  9 ++++++++-
 tests/suites/test_suite_psa_crypto.function | 13 +++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 0f74c5481d..e05c846ff8 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -172,11 +172,12 @@ struct psa_aead_operation_s
     unsigned int lengths_set : 1;
     unsigned int ad_started : 1;
     unsigned int body_started : 1;
+    unsigned int is_encrypt : 1;
 
     psa_driver_aead_context_t ctx;
 };
 
-#define PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, {0}}
+#define PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, {0}}
 static inline struct psa_aead_operation_s psa_aead_operation_init( void )
 {
     const struct psa_aead_operation_s v = PSA_AEAD_OPERATION_INIT;
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 8dc6aad534..aec22c79cf 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3469,7 +3469,10 @@ exit:
         status = unlock_status;
 
     if( status == PSA_SUCCESS )
+    {
         operation->alg = psa_aead_get_base_algorithm( alg );
+        operation->is_encrypt = 1;
+    }
     else
         psa_aead_abort( operation );
 
@@ -3531,7 +3534,10 @@ exit:
         status = unlock_status;
 
     if( status == PSA_SUCCESS )
+    {
         operation->alg = psa_aead_get_base_algorithm( alg );
+        operation->is_encrypt = 0;
+    }
     else
         psa_aead_abort( operation );
 
@@ -3556,7 +3562,7 @@ psa_status_t psa_aead_generate_nonce( psa_aead_operation_t *operation,
     }
 
     if( operation->nonce_set || operation->ad_started ||
-        operation->body_started )
+        operation->body_started  || operation->is_encrypt == 0 )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3881,6 +3887,7 @@ psa_status_t psa_aead_abort( psa_aead_operation_t *operation )
     operation->lengths_set = 0;
     operation->ad_started = 0;
     operation->body_started = 0;
+    operation->is_encrypt = 0;
 
     return( status );
 }
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 2a2f2e61c4..38545bccc6 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -3913,6 +3913,19 @@ void aead_multipart_state_test( int key_type_arg, data_t *key_data,
 
     psa_aead_abort( &operation );
 
+    /* Test for generating nonce in decrypt setup. */
+
+    operation = psa_aead_operation_init( );
+
+    PSA_ASSERT( psa_aead_decrypt_setup( &operation, key, alg ) );
+
+    TEST_EQUAL( psa_aead_generate_nonce( &operation, nonce_buffer,
+                                         PSA_AEAD_NONCE_MAX_SIZE,
+                                         &nonce_length ),
+                PSA_ERROR_BAD_STATE );
+
+    psa_aead_abort( &operation );
+
     /* Test for setting lengths twice. */
 
     operation = psa_aead_operation_init( );