Replace mbedtls_psa_random_state( ... ) by MBEDTLS_PSA_RANDOM_STATE

Hide the obtention of the pointer to the RNG state behind a macro.

To make it possible to use this macro in a module other than
psa_crypto.c, which will happen in the future, make sure that the
definition of the macro does not reference internal variables of
psa_crypto.c. For this purpose, in the internal-DRBG case, export a
symbol containing the address of the DRBG state.

When the RNG state is a pointer a DRBG state, just keep this pointer
in a variable: there's no need to store a pointer to a larger structure.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index b3d7eb4..bb38475 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -124,6 +124,11 @@
 
 static psa_global_data_t global_data;
 
+#if !defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG)
+mbedtls_psa_drbg_context_t *const mbedtls_psa_random_state =
+    &global_data.rng.drbg;
+#endif
+
 #define GUARD_MODULE_INITIALIZED        \
     if( global_data.initialized == 0 )  \
         return( PSA_ERROR_BAD_STATE );
@@ -884,7 +889,7 @@
             /* Calculate the public key */
             status = mbedtls_to_psa_error(
                 mbedtls_ecp_mul( &ecp->grp, &ecp->Q, &ecp->d, &ecp->grp.G,
-                                 mbedtls_psa_get_random, mbedtls_psa_random_state( &global_data.rng ) ) );
+                                 mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE ) );
             if( status != PSA_SUCCESS )
                 return( status );
         }
@@ -3667,7 +3672,7 @@
                                  MBEDTLS_MD_NONE );
         ret = mbedtls_rsa_pkcs1_sign( rsa,
                                       mbedtls_psa_get_random,
-                                      mbedtls_psa_random_state( &global_data.rng ),
+                                      MBEDTLS_PSA_RANDOM_STATE,
                                       MBEDTLS_RSA_PRIVATE,
                                       md_alg,
                                       (unsigned int) hash_length,
@@ -3682,7 +3687,7 @@
         mbedtls_rsa_set_padding( rsa, MBEDTLS_RSA_PKCS_V21, md_alg );
         ret = mbedtls_rsa_rsassa_pss_sign( rsa,
                                            mbedtls_psa_get_random,
-                                           mbedtls_psa_random_state( &global_data.rng ),
+                                           MBEDTLS_PSA_RANDOM_STATE,
                                            MBEDTLS_RSA_PRIVATE,
                                            MBEDTLS_MD_NONE,
                                            (unsigned int) hash_length,
@@ -3725,7 +3730,7 @@
                                  MBEDTLS_MD_NONE );
         ret = mbedtls_rsa_pkcs1_verify( rsa,
                                         mbedtls_psa_get_random,
-                                        mbedtls_psa_random_state( &global_data.rng ),
+                                        MBEDTLS_PSA_RANDOM_STATE,
                                         MBEDTLS_RSA_PUBLIC,
                                         md_alg,
                                         (unsigned int) hash_length,
@@ -3740,7 +3745,7 @@
         mbedtls_rsa_set_padding( rsa, MBEDTLS_RSA_PKCS_V21, md_alg );
         ret = mbedtls_rsa_rsassa_pss_verify( rsa,
                                              mbedtls_psa_get_random,
-                                             mbedtls_psa_random_state( &global_data.rng ),
+                                             MBEDTLS_PSA_RANDOM_STATE,
                                              MBEDTLS_RSA_PUBLIC,
                                              MBEDTLS_MD_NONE,
                                              (unsigned int) hash_length,
@@ -3798,7 +3803,7 @@
                                                      &ecp->d, hash,
                                                      hash_length, md_alg,
                                                      mbedtls_psa_get_random,
-                                                     mbedtls_psa_random_state( &global_data.rng ) ) );
+                                                     MBEDTLS_PSA_RANDOM_STATE ) );
     }
     else
 #endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA) */
@@ -3807,7 +3812,7 @@
         MBEDTLS_MPI_CHK( mbedtls_ecdsa_sign( &ecp->grp, &r, &s, &ecp->d,
                                              hash, hash_length,
                                              mbedtls_psa_get_random,
-                                             mbedtls_psa_random_state( &global_data.rng ) ) );
+                                             MBEDTLS_PSA_RANDOM_STATE ) );
     }
 
     MBEDTLS_MPI_CHK( mbedtls_mpi_write_binary( &r,
@@ -3852,7 +3857,7 @@
     {
         MBEDTLS_MPI_CHK(
             mbedtls_ecp_mul( &ecp->grp, &ecp->Q, &ecp->d, &ecp->grp.G,
-                             mbedtls_psa_get_random, mbedtls_psa_random_state( &global_data.rng ) ) );
+                             mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE ) );
     }
 
     ret = mbedtls_ecdsa_verify( &ecp->grp, hash, hash_length,
@@ -4153,7 +4158,7 @@
             status = mbedtls_to_psa_error(
                     mbedtls_rsa_pkcs1_encrypt( rsa,
                                                mbedtls_psa_get_random,
-                                               mbedtls_psa_random_state( &global_data.rng ),
+                                               MBEDTLS_PSA_RANDOM_STATE,
                                                MBEDTLS_RSA_PUBLIC,
                                                input_length,
                                                input,
@@ -4168,7 +4173,7 @@
             status = mbedtls_to_psa_error(
                 mbedtls_rsa_rsaes_oaep_encrypt( rsa,
                                                 mbedtls_psa_get_random,
-                                                mbedtls_psa_random_state( &global_data.rng ),
+                                                MBEDTLS_PSA_RANDOM_STATE,
                                                 MBEDTLS_RSA_PUBLIC,
                                                 salt, salt_length,
                                                 input_length,
@@ -4260,7 +4265,7 @@
             status = mbedtls_to_psa_error(
                 mbedtls_rsa_pkcs1_decrypt( rsa,
                                            mbedtls_psa_get_random,
-                                           mbedtls_psa_random_state( &global_data.rng ),
+                                           MBEDTLS_PSA_RANDOM_STATE,
                                            MBEDTLS_RSA_PRIVATE,
                                            output_length,
                                            input,
@@ -4276,7 +4281,7 @@
             status = mbedtls_to_psa_error(
                 mbedtls_rsa_rsaes_oaep_decrypt( rsa,
                                                 mbedtls_psa_get_random,
-                                                mbedtls_psa_random_state( &global_data.rng ),
+                                                MBEDTLS_PSA_RANDOM_STATE,
                                                 MBEDTLS_RSA_PRIVATE,
                                                 salt, salt_length,
                                                 output_length,
@@ -4510,7 +4515,7 @@
         status = PSA_ERROR_BUFFER_TOO_SMALL;
         goto exit;
     }
-    ret = mbedtls_psa_get_random( mbedtls_psa_random_state( &global_data.rng ),
+    ret = mbedtls_psa_get_random( MBEDTLS_PSA_RANDOM_STATE,
                                   iv, operation->iv_size );
     if( ret != 0 )
     {
@@ -6127,7 +6132,7 @@
                                   shared_secret_length,
                                   shared_secret, shared_secret_size,
                                   mbedtls_psa_get_random,
-                                  mbedtls_psa_random_state( &global_data.rng ) ) );
+                                  MBEDTLS_PSA_RANDOM_STATE ) );
     if( status != PSA_SUCCESS )
         goto exit;
     if( PSA_BITS_TO_BYTES( bits ) != *shared_secret_length )
@@ -6336,7 +6341,7 @@
                                 MBEDTLS_ENTROPY_SOURCE_STRONG );
 #endif
 
-    mbedtls_psa_drbg_init( mbedtls_psa_random_state( rng ) );
+    mbedtls_psa_drbg_init( MBEDTLS_PSA_RANDOM_STATE );
 #endif /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
 }
 
@@ -6347,7 +6352,7 @@
 #if defined(MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG)
     memset( rng, 0, sizeof( *rng ) );
 #else /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
-    mbedtls_psa_drbg_free( mbedtls_psa_random_state( rng ) );
+    mbedtls_psa_drbg_free( MBEDTLS_PSA_RANDOM_STATE );
     rng->entropy_free( &rng->entropy );
 #endif /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
 }
@@ -6362,7 +6367,8 @@
     return( PSA_SUCCESS );
 #else /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
     const unsigned char drbg_seed[] = "PSA";
-    int ret = mbedtls_psa_drbg_seed( rng, drbg_seed, sizeof( drbg_seed ) - 1 );
+    int ret = mbedtls_psa_drbg_seed( &rng->entropy,
+                                     drbg_seed, sizeof( drbg_seed ) - 1 );
     return mbedtls_to_psa_error( ret );
 #endif /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
 }
@@ -6393,7 +6399,7 @@
     while( output_size > MBEDTLS_PSA_RANDOM_MAX_REQUEST )
     {
         ret = mbedtls_psa_get_random(
-            mbedtls_psa_random_state( &global_data.rng ),
+            MBEDTLS_PSA_RANDOM_STATE,
             output, MBEDTLS_PSA_RANDOM_MAX_REQUEST );
         if( ret != 0 )
             return( mbedtls_to_psa_error( ret ) );
@@ -6401,7 +6407,7 @@
         output_size -= MBEDTLS_PSA_RANDOM_MAX_REQUEST;
     }
 
-    ret = mbedtls_psa_get_random( mbedtls_psa_random_state( &global_data.rng ),
+    ret = mbedtls_psa_get_random( MBEDTLS_PSA_RANDOM_STATE,
                                   output, output_size );
     return( mbedtls_to_psa_error( ret ) );
 #endif /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
@@ -6510,7 +6516,7 @@
         mbedtls_rsa_init( &rsa, MBEDTLS_RSA_PKCS_V15, MBEDTLS_MD_NONE );
         ret = mbedtls_rsa_gen_key( &rsa,
                                    mbedtls_psa_get_random,
-                                   mbedtls_psa_random_state( &global_data.rng ),
+                                   MBEDTLS_PSA_RANDOM_STATE,
                                    (unsigned int) bits,
                                    exponent );
         if( ret != 0 )
@@ -6556,7 +6562,7 @@
         mbedtls_ecp_keypair_init( &ecp );
         ret = mbedtls_ecp_gen_key( grp_id, &ecp,
                                    mbedtls_psa_get_random,
-                                   mbedtls_psa_random_state( &global_data.rng ) );
+                                   MBEDTLS_PSA_RANDOM_STATE );
         if( ret != 0 )
         {
             mbedtls_ecp_keypair_free( &ecp );
diff --git a/library/psa_crypto_random.h b/library/psa_crypto_random.h
index 0f13c13..cc1222a 100644
--- a/library/psa_crypto_random.h
+++ b/library/psa_crypto_random.h
@@ -42,11 +42,7 @@
         return( MBEDTLS_ERR_ENTROPY_SOURCE_FAILED );
 }
 
-static inline void *mbedtls_psa_random_state( mbedtls_psa_random_context_t *rng )
-{
-    (void) rng;
-    return( NULL );
-}
+#define MBEDTLS_PSA_RANDOM_STATE NULL
 
 #else /* MBEDTLS_PSA_CRYPTO_EXTERNAL_RNG */
 
@@ -135,13 +131,11 @@
 /** Return random data.
  *
  * This function is suitable as the \p f_rng parameter to Mbed TLS functions
- * that require a random generator. Use mbedtls_psa_random_state() to
+ * that require a random generator. Use #MBEDTLS_PSA_RANDOM_STATE to
  * obtain the \p p_rng parameter.
  *
  * \param p_rng         The DRBG context. This must be
- *                      mbedtls_psa_random_state( \c rng )
- *                      where \c rng is a pointer to a
- *                      ::mbedtls_psa_random_context_t structure.
+ *                      #MBEDTLS_PSA_RANDOM_STATE.
  * \param output        The buffer to fill.
  * \param output_len    The length of the buffer in bytes.
  *                      It must be at most #MBEDTLS_PSA_RANDOM_MAX_REQUEST.
@@ -170,21 +164,26 @@
 #define MBEDTLS_PSA_RANDOM_MAX_REQUEST MBEDTLS_HMAC_DRBG_MAX_REQUEST
 #endif
 
-/** Retrieve the DRBG state from the PSA RNG state.
+/** A pointer to the PSA DRBG state.
  *
- * \param rng           Pointer to the PSA random generator state.
- *
- * \return              The DRBG state (\c p_rng argument ).
+ * This variable is only intended to be used through the macro
+ * #MBEDTLS_PSA_RANDOM_STATE.
  */
-static inline mbedtls_psa_drbg_context_t *mbedtls_psa_random_state(
-    mbedtls_psa_random_context_t *rng )
-{
-    return( &rng->drbg );
-}
+extern mbedtls_psa_drbg_context_t *const mbedtls_psa_random_state;
+
+/** A pointer to the PSA DRBG state.
+ *
+ * This macro expnds to an expression that is suitable as the \c p_rng
+ * parameter to pass to mbedtls_psa_get_random().
+ *
+ * This macro exists in all configurations where the psa_crypto module is
+ * enabled. Its expansion depends on the configuration.
+ */
+#define MBEDTLS_PSA_RANDOM_STATE mbedtls_psa_random_state
 
 /** Seed the PSA DRBG.
  *
- * \param rng           DRBG context to be seeded.
+ * \param entropy       An entropy context to read the seed from.
  * \param custom        The personalization string.
  *                      This can be \c NULL, in which case the personalization
  *                      string is empty regardless of the value of \p len.
@@ -194,21 +193,21 @@
  * \return              An Mbed TLS error code (\c MBEDTLS_ERR_xxx) on failure.
  */
 static inline int mbedtls_psa_drbg_seed(
-    mbedtls_psa_random_context_t *rng,
+    mbedtls_entropy_context *entropy,
     const unsigned char *custom, size_t len )
 {
 #if defined(MBEDTLS_CTR_DRBG_C)
-    return( mbedtls_ctr_drbg_seed( mbedtls_psa_random_state( rng ),
+    return( mbedtls_ctr_drbg_seed( MBEDTLS_PSA_RANDOM_STATE,
                                    mbedtls_entropy_func,
-                                   &rng->entropy,
+                                   entropy,
                                    custom, len ) );
 #elif defined(MBEDTLS_HMAC_DRBG_C)
     const mbedtls_md_info_t *md_info =
         mbedtls_md_info_from_type( MBEDTLS_PSA_HMAC_DRBG_MD_TYPE );
-    return( mbedtls_hmac_drbg_seed( mbedtls_psa_random_state( rng ),
+    return( mbedtls_hmac_drbg_seed( MBEDTLS_PSA_RANDOM_STATE,
                                     md_info,
                                     mbedtls_entropy_func,
-                                    &rng->entropy,
+                                    entropy,
                                     custom, len ) );
 #endif
 }