Protect setting of premaster_generated flag

The flag is used for tracking if the premaster has
been succesfully generated. Note that when resuming
a session, the flag should not be used when trying to
notice if all the key generation/derivation has been done.
diff --git a/library/ssl_cli.c b/library/ssl_cli.c
index 6b95cfa..08c65d2 100644
--- a/library/ssl_cli.c
+++ b/library/ssl_cli.c
@@ -2342,8 +2342,9 @@
                                          unsigned char* out,
                                          unsigned add_length_tag )
 {
-    int ret;
+    volatile int ret;
 
+    ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_UNSET;
     /*
      * Generate (part of) the pre-master secret as
      *  struct {
@@ -2364,14 +2365,21 @@
                                mbedtls_ssl_conf_get_max_minor_ver( ssl->conf ),
                                ssl->conf->transport, out );
 
-    if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-          ( mbedtls_ssl_conf_get_prng( ssl->conf ), out + 2, 46 ) ) != 0 )
+    ret = mbedtls_ssl_conf_get_frng( ssl->conf )
+          ( mbedtls_ssl_conf_get_prng( ssl->conf ), out + 2, 46 );
+
+    if( ret == 0 )
     {
-        MBEDTLS_SSL_DEBUG_RET( 1, "f_rng", ret );
-        return( ret );
+        mbedtls_platform_enforce_volatile_reads();
+        if( ret == 0 )
+        {
+            ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+            return( 0 );
+        }
     }
 
-    return( 0 );
+    MBEDTLS_SSL_DEBUG_RET( 1, "f_rng", ret );
+    return( ret );
 }
 
 /*
@@ -2383,11 +2391,12 @@
                                         unsigned char *out, size_t buflen,
                                         size_t *olen )
 {
-    int ret;
+    volatile int ret;
     size_t len_bytes = mbedtls_ssl_get_minor_ver( ssl ) ==
         MBEDTLS_SSL_MINOR_VERSION_0 ? 0 : 2;
     mbedtls_pk_context *peer_pk = NULL;
 
+    ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_UNSET;
     if( buflen < len_bytes )
     {
         MBEDTLS_SSL_DEBUG_MSG( 1, ( "buffer too small for encrypted pms" ) );
@@ -2427,16 +2436,27 @@
         goto cleanup;
     }
 
-    if( ( ret = mbedtls_pk_encrypt( peer_pk,
+    ret = mbedtls_pk_encrypt( peer_pk,
                             ppms, 48, out + len_bytes,
                             olen, buflen - len_bytes,
                             mbedtls_ssl_conf_get_frng( ssl->conf ),
-                            mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
+                            mbedtls_ssl_conf_get_prng( ssl->conf ) );
+
+    if( ret == 0 )
+    {
+        mbedtls_platform_enforce_volatile_reads();
+        if( ret == 0 )
+        {
+            ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+        }
+    }
+    else
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_rsa_pkcs1_encrypt", ret );
         goto cleanup;
     }
 
+
 #if defined(MBEDTLS_SSL_PROTO_TLS1) || defined(MBEDTLS_SSL_PROTO_TLS1_1) || \
     defined(MBEDTLS_SSL_PROTO_TLS1_2)
     if( len_bytes == 2 )
diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index e349ed8..3053818 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -3978,7 +3978,9 @@
     unsigned char mask;
     size_t i, peer_pmslen;
     unsigned int diff;
+    volatile unsigned int pmscounter = 0;
 
+    ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_UNSET;
     /* In case of a failure in decryption, the decryption may write less than
      * 2 bytes of output, but we always read the first two bytes. It doesn't
      * matter in the end because diff will be nonzero in that case due to
@@ -4056,7 +4058,19 @@
     /* Set pms to either the true or the fake PMS, without
      * data-dependent branches. */
     for( i = 0; i < ssl->handshake->pmslen; i++ )
+    {
         pms[i] = ( mask & fake_pms[i] ) | ( (~mask) & peer_pms[i] );
+        pmscounter++;
+    }
+
+    if( pmscounter == ssl->handshake->pmslen )
+    {
+        mbedtls_platform_enforce_volatile_reads();
+        if( pmscounter == ssl->handshake->pmslen )
+        {
+            ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+        }
+    }
 
     return( 0 );
 }
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index d7ad696..66772f4 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1881,7 +1881,7 @@
 
 int mbedtls_ssl_derive_keys( mbedtls_ssl_context *ssl )
 {
-    int ret;
+    volatile int ret;
 
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> derive keys" ) );
 
@@ -1958,11 +1958,11 @@
 
 int mbedtls_ssl_build_pms( mbedtls_ssl_context *ssl )
 {
-    int ret;
+    volatile int ret = MBEDTLS_ERR_SSL_INTERNAL_ERROR;
 
     mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
         mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
-
+    ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_UNSET;
 #if defined(MBEDTLS_USE_TINYCRYPT)
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
         == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA                       ||
@@ -1980,6 +1980,7 @@
             return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
         if( ret != UECC_SUCCESS )
             return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+        ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
 
         ssl->handshake->pmslen = NUM_ECC_BYTES;
     }
@@ -1989,12 +1990,26 @@
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
         == MBEDTLS_KEY_EXCHANGE_DHE_RSA )
     {
-        if( ( ret = mbedtls_dhm_calc_secret( &ssl->handshake->dhm_ctx,
+        ret = mbedtls_dhm_calc_secret( &ssl->handshake->dhm_ctx,
                                       ssl->handshake->premaster,
                                       MBEDTLS_PREMASTER_SIZE,
                                      &ssl->handshake->pmslen,
                                       mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                      mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
+                                      mbedtls_ssl_conf_get_prng( ssl->conf ) );
+        if( ret == 0 )
+        {
+            mbedtls_platform_enforce_volatile_reads();
+            if( ret == 0 )
+            {
+                ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+            }
+            else
+            {
+                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );
+                return( ret );
+            }
+        }
+        else
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );
             return( ret );
@@ -2018,12 +2033,26 @@
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
         == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
-        if( ( ret = mbedtls_ecdh_calc_secret( &ssl->handshake->ecdh_ctx,
+        ret = mbedtls_ecdh_calc_secret( &ssl->handshake->ecdh_ctx,
                                       &ssl->handshake->pmslen,
                                        ssl->handshake->premaster,
                                        MBEDTLS_MPI_MAX_SIZE,
                                        mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                       mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
+                                       mbedtls_ssl_conf_get_prng( ssl->conf ) );
+        if( ret == 0 )
+        {
+            mbedtls_platform_enforce_volatile_reads();
+            if( ret == 0 )
+            {
+                ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+            }
+            else
+            {
+                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
+                return( ret );
+            }
+        }
+        else
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
             return( ret );
@@ -2039,8 +2068,22 @@
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
     if( mbedtls_ssl_ciphersuite_uses_psk( ciphersuite_info ) )
     {
-        if( ( ret = mbedtls_ssl_psk_derive_premaster( ssl,
-            mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ) ) != 0 )
+        ret = mbedtls_ssl_psk_derive_premaster( ssl,
+            mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) );
+        if( ret == 0 )
+        {
+            mbedtls_platform_enforce_volatile_reads();
+            if( ret == 0 )
+            {
+                ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+            }
+            else
+            {
+                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
+                return( ret );
+            }
+        }
+        else
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
             return( ret );
@@ -2056,7 +2099,20 @@
                 ssl->handshake->premaster, 32, &ssl->handshake->pmslen,
                 mbedtls_ssl_conf_get_frng( ssl->conf ),
                 mbedtls_ssl_conf_get_prng( ssl->conf ) );
-        if( ret != 0 )
+        if( ret == 0 )
+        {
+            mbedtls_platform_enforce_volatile_reads();
+            if( ret == 0 )
+            {
+                ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+            }
+            else
+            {
+                MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
+                return( ret );
+            }
+        }
+        else
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
             return( ret );
@@ -2091,6 +2147,7 @@
     unsigned char *end = p + sizeof( ssl->handshake->premaster );
     const unsigned char *psk = ssl->conf->psk;
     size_t psk_len = ssl->conf->psk_len;
+    ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_UNSET;
 
     /* If the psk callback was called, use its result */
     if( ssl->handshake->psk != NULL )
@@ -2214,6 +2271,8 @@
 
     ssl->handshake->pmslen = p - ssl->handshake->premaster;
 
+    ssl->handshake->premaster_generated = MBEDTLS_SSL_FI_FLAG_SET;
+
     return( 0 );
 }
 #endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED */