Merge remote-tracking branch 'origin/pr/629' into baremetal
diff --git a/configs/baremetal.h b/configs/baremetal.h
index 9929d56..e4e3a2a 100644
--- a/configs/baremetal.h
+++ b/configs/baremetal.h
@@ -106,7 +106,7 @@
 #define MBEDTLS_SSL_CONF_RECV mbedtls_net_recv
 #define MBEDTLS_SSL_CONF_SEND mbedtls_net_send
 #define MBEDTLS_SSL_CONF_RECV_TIMEOUT mbedtls_net_recv_timeout
-#define MBEDTLS_SSL_CONF_RNG mbedtls_hmac_drbg_random
+#define MBEDTLS_SSL_CONF_RNG rng_wrap
 #define MBEDTLS_SSL_CONF_MIN_MINOR_VER MBEDTLS_SSL_MINOR_VERSION_3
 #define MBEDTLS_SSL_CONF_MAX_MINOR_VER MBEDTLS_SSL_MINOR_VERSION_3
 #define MBEDTLS_SSL_CONF_MIN_MAJOR_VER MBEDTLS_SSL_MAJOR_VERSION_3
diff --git a/configs/baremetal_test.h b/configs/baremetal_test.h
index 1b87474..82c0ed1 100644
--- a/configs/baremetal_test.h
+++ b/configs/baremetal_test.h
@@ -51,8 +51,6 @@
 
 /* ssl_client2 and ssl_server2 use CTR-DRBG so far. */
 #define MBEDTLS_CTR_DRBG_C
-#undef MBEDTLS_SSL_CONF_RNG
-#define MBEDTLS_SSL_CONF_RNG mbedtls_ctr_drbg_random
 
 /* The ticket implementation hardcodes AES-GCM */
 #define MBEDTLS_GCM_C
diff --git a/include/mbedtls/config.h b/include/mbedtls/config.h
index 243678c..930f7c7 100644
--- a/include/mbedtls/config.h
+++ b/include/mbedtls/config.h
@@ -3666,11 +3666,15 @@
 //#define MBEDTLS_SSL_CONF_RECV_TIMEOUT mbedtls_net_recv_timeout
 
 /* The PRNG to use by the SSL module. If defined, this must
- * evaluate to the name on externally defined function with signature
+ * evaluate to the name of an externally defined function with signature
  * int f_rng(void *, unsigned char *, size_t),
- * e.g. mbedtls_ctr_drbg_random or mbedtls_hmac_drbg_random.
+ * which ignores its first parameter (the stack will always
+ * pass NULL to this function).
+ * NB: The reason not to drop the context parameter entirely
+ * is for compatibility with the Crypto and X.509 libraries
+ * which use RNG callbacks including context parameters.
  */
-//#define MBEDTLS_SSL_CONF_RNG mbedtls_ctr_drbg_random
+//#define MBEDTLS_SSL_CONF_RNG rng_wrap
 
 /* TLS version */
 //#define MBEDTLS_SSL_CONF_MIN_MINOR_VER MBEDTLS_SSL_MINOR_VERSION_3
diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 1119438..b6a3b60 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -1002,8 +1002,8 @@
 #if !defined(MBEDTLS_SSL_CONF_RNG)
     /** Callback for getting (pseudo-)random numbers                        */
     int  (*f_rng)(void *, unsigned char *, size_t);
-#endif /* !MBEDTLS_SSL_CONF_RNG */
     void *p_rng;                    /*!< context for the RNG function       */
+#endif /* !MBEDTLS_SSL_CONF_RNG */
 
 #if defined(MBEDTLS_SSL_SRV_C) && !defined(MBEDTLS_SSL_NO_SESSION_CACHE)
     /** Callback to retrieve a session from the cache                       */
@@ -1603,15 +1603,6 @@
 void mbedtls_ssl_conf_rng( mbedtls_ssl_config *conf,
                   int (*f_rng)(void *, unsigned char *, size_t),
                   void *p_rng );
-#else /* !MBEDTLS_SSL_CONF_RNG */
-/**
- * \brief          Set the random number generator callback context.
- *
- * \param conf     SSL configuration
- * \param p_rng    RNG parameter
- */
-void mbedtls_ssl_conf_rng_ctx( mbedtls_ssl_config *conf,
-                               void *p_rng );
 #endif /* MBEDTLS_SSL_CONF_RNG */
 
 #if defined(MBEDTLS_DEBUG_C)
diff --git a/include/mbedtls/ssl_internal.h b/include/mbedtls/ssl_internal.h
index 192bc7a..ce7c9a5 100644
--- a/include/mbedtls/ssl_internal.h
+++ b/include/mbedtls/ssl_internal.h
@@ -404,12 +404,10 @@
     enum { /* this complements ssl->state with info on intra-state operations */
         ssl_ecrs_none = 0,              /*!< nothing going on (yet)         */
         ssl_ecrs_crt_verify,            /*!< Certificate: crt_verify()      */
-        ssl_ecrs_ske_start_processing,  /*!< ServerKeyExchange: pk_verify() */
         ssl_ecrs_cke_ecdh_calc_secret,  /*!< ClientKeyExchange: ECDH step 2 */
         ssl_ecrs_crt_vrfy_sign,         /*!< CertificateVerify: pk_sign()   */
     } ecrs_state;                       /*!< current (or last) operation    */
     mbedtls_x509_crt *ecrs_peer_cert;   /*!< The peer's CRT chain.          */
-    size_t ecrs_n;                      /*!< place for saving a length      */
 #endif
 #if defined(MBEDTLS_X509_CRT_PARSE_C) && \
     !defined(MBEDTLS_SSL_KEEP_PEER_CERTIFICATE)
@@ -988,6 +986,8 @@
 void mbedtls_ssl_optimize_checksum( mbedtls_ssl_context *ssl,
                             mbedtls_ssl_ciphersuite_handle_t ciphersuite_info );
 
+int mbedtls_ssl_build_pms( mbedtls_ssl_context *ssl );
+
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
 int mbedtls_ssl_psk_derive_premaster( mbedtls_ssl_context *ssl, mbedtls_key_exchange_type_t key_ex );
 #endif
@@ -1565,8 +1565,12 @@
 {
     return( conf->f_rng );
 }
-#else /* !MBEDTLS_SSL_CONF_RNG */
 
+static inline void* mbedtls_ssl_conf_get_prng( mbedtls_ssl_config const *conf )
+{
+    return( conf->p_rng );
+}
+#else /* !MBEDTLS_SSL_CONF_RNG */
 #define mbedtls_ssl_conf_rng_func MBEDTLS_SSL_CONF_RNG
 extern int mbedtls_ssl_conf_rng_func( void*, unsigned char*, size_t );
 
@@ -1576,6 +1580,12 @@
     ((void) conf);
     return ((mbedtls_frng_t*) mbedtls_ssl_conf_rng_func);
 }
+
+static inline void* mbedtls_ssl_conf_get_prng( mbedtls_ssl_config const *conf )
+{
+    ((void) conf);
+    return( NULL );
+}
 #endif /* MBEDTLS_SSL_CONF_RNG */
 
 static inline int mbedtls_ssl_conf_get_max_major_ver(
@@ -1821,4 +1831,6 @@
 #endif
 }
 
+#define MBEDTLS_SSL_CHK(f) do { if( ( ret = f ) < 0 ) goto cleanup; } while( 0 )
+
 #endif /* ssl_internal.h */
diff --git a/library/ssl_cli.c b/library/ssl_cli.c
index cdb35f3..050d8b5 100644
--- a/library/ssl_cli.c
+++ b/library/ssl_cli.c
@@ -375,7 +375,7 @@
         ret = mbedtls_ecjpake_write_round_one( &ssl->handshake->ecjpake_ctx,
                                         p + 2, end - p - 2, &kkpp_len,
                                         mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                        ssl->conf->p_rng );
+                                        mbedtls_ssl_conf_get_prng( ssl->conf ) );
         if( ret != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1 , "mbedtls_ecjpake_write_round_one", ret );
@@ -735,7 +735,7 @@
     MBEDTLS_SSL_DEBUG_MSG( 3, ( "client hello, current time: %lu", t ) );
 #else
     if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-          ( ssl->conf->p_rng, p, 4 ) ) != 0 )
+          ( mbedtls_ssl_conf_get_prng( ssl->conf ), p, 4 ) ) != 0 )
     {
         return( ret );
     }
@@ -744,7 +744,7 @@
 #endif /* MBEDTLS_HAVE_TIME */
 
     if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-          ( ssl->conf->p_rng, p, 28 ) ) != 0 )
+          ( mbedtls_ssl_conf_get_prng( ssl->conf ), p, 28 ) ) != 0 )
     {
         return( ret );
     }
@@ -911,7 +911,7 @@
         ssl->session_negotiate->ticket_len != 0 )
     {
         ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-            ( ssl->conf->p_rng, ssl->session_negotiate->id, 32 );
+            ( mbedtls_ssl_conf_get_prng( ssl->conf ), ssl->session_negotiate->id, 32 );
 
         if( ret != 0 )
             return( ret );
@@ -2278,10 +2278,6 @@
                                   (const unsigned char **) p, end ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, ( "mbedtls_ecdh_read_params" ), ret );
-#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
-        if( ret == MBEDTLS_ERR_ECP_IN_PROGRESS )
-            ret = MBEDTLS_ERR_SSL_CRYPTO_IN_PROGRESS;
-#endif
         return( ret );
     }
 
@@ -2341,43 +2337,65 @@
 
 #if defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED) ||                           \
     defined(MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED)
-/*
- * Generate a pre-master secret and encrypt it with the server's RSA key
- */
-static int ssl_write_encrypted_pms( mbedtls_ssl_context *ssl,
-                                    size_t offset, size_t *olen,
-                                    size_t pms_offset )
+
+/* This function assumes that `out` points to a writable buffer
+ * of size 48 Bytes if `add_length_tag == 0` and size 50 Bytes
+ * if `add_length_tag != 0`. */
+static int ssl_rsa_generate_partial_pms( mbedtls_ssl_context *ssl,
+                                         unsigned char* out,
+                                         unsigned add_length_tag )
 {
     int ret;
-    size_t len_bytes = mbedtls_ssl_get_minor_ver( ssl ) == MBEDTLS_SSL_MINOR_VERSION_0 ? 0 : 2;
-    unsigned char *p = ssl->handshake->premaster + pms_offset;
-    mbedtls_pk_context *peer_pk = NULL;
-
-    if( offset + len_bytes > MBEDTLS_SSL_OUT_CONTENT_LEN )
-    {
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "buffer too small for encrypted pms" ) );
-        return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
-    }
 
     /*
-     * Generate (part of) the pre-master as
+     * Generate (part of) the pre-master secret as
      *  struct {
+     *      [ uint16 length(48) ]
      *      ProtocolVersion client_version;
      *      opaque random[46];
      *  } PreMasterSecret;
      */
+
+    if( add_length_tag )
+    {
+        out[0] = 0;
+        out[1] = 48;
+        out += 2;
+    }
+
     mbedtls_ssl_write_version( mbedtls_ssl_conf_get_max_major_ver( ssl->conf ),
                                mbedtls_ssl_conf_get_max_minor_ver( ssl->conf ),
-                               ssl->conf->transport, p );
+                               ssl->conf->transport, out );
 
     if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-          ( ssl->conf->p_rng, p + 2, 46 ) ) != 0 )
+          ( mbedtls_ssl_conf_get_prng( ssl->conf ), out + 2, 46 ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "f_rng", ret );
         return( ret );
     }
 
-    ssl->handshake->pmslen = 48;
+    return( 0 );
+}
+
+/*
+ * Encrypt the randomly chosen part of the Premaster Secret with the
+ * server's RSA key and write it to the provided buffer.
+ */
+static int ssl_rsa_encrypt_partial_pms( mbedtls_ssl_context *ssl,
+                                        unsigned char const *ppms,
+                                        unsigned char *out, size_t buflen,
+                                        size_t *olen )
+{
+    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;
+
+    if( buflen < len_bytes )
+    {
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "buffer too small for encrypted pms" ) );
+        return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
+    }
 
 #if !defined(MBEDTLS_SSL_KEEP_PEER_CERTIFICATE)
     /* Because the peer CRT pubkey is embedded into the handshake
@@ -2407,7 +2425,7 @@
     }
 
     /*
-     * Now write it out, encrypted
+     * Encrypt the partial premaster secret and write it out.
      */
     if( ! mbedtls_pk_can_do( peer_pk, MBEDTLS_PK_RSA ) )
     {
@@ -2417,11 +2435,10 @@
     }
 
     if( ( ret = mbedtls_pk_encrypt( peer_pk,
-                            p, ssl->handshake->pmslen,
-                            ssl->out_msg + offset + len_bytes, olen,
-                            MBEDTLS_SSL_OUT_CONTENT_LEN - offset - len_bytes,
+                            ppms, 48, out + len_bytes,
+                            olen, buflen - len_bytes,
                             mbedtls_ssl_conf_get_frng( ssl->conf ),
-                            ssl->conf->p_rng ) ) != 0 )
+                            mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_rsa_pkcs1_encrypt", ret );
         goto cleanup;
@@ -2431,8 +2448,8 @@
     defined(MBEDTLS_SSL_PROTO_TLS1_2)
     if( len_bytes == 2 )
     {
-        ssl->out_msg[offset+0] = (unsigned char)( *olen >> 8 );
-        ssl->out_msg[offset+1] = (unsigned char)( *olen      );
+        out[0] = (unsigned char)( *olen >> 8 );
+        out[1] = (unsigned char)( *olen      );
         *olen += 2;
     }
 #endif
@@ -2581,25 +2598,94 @@
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||
           MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED */
 
-static int ssl_parse_server_key_exchange( mbedtls_ssl_context *ssl )
+/*
+ *
+ * STATE HANDLING: ServerKeyExchange
+ *
+ */
+
+/*
+ * Overview
+ */
+
+/* Main entry point; orchestrates the other functions. */
+static int ssl_process_in_server_key_exchange( mbedtls_ssl_context *ssl );
+
+/* Coordination:
+ * Check if a ServerKeyExchange message is expected, and skip if not.
+ * Returns a negative code on failure, or
+ * - SSL_SRV_KEY_EXCHANGE_SKIP
+ * - SSL_SRV_KEY_EXCHANGE_EXPECTED
+ * indicating if a ServerKeyExchange message is expected or not.
+ */
+#define SSL_SRV_KEY_EXCHANGE_SKIP      0
+#define SSL_SRV_KEY_EXCHANGE_EXPECTED  1
+static int ssl_in_server_key_exchange_coordinate( mbedtls_ssl_context *ssl );
+/* Preparation
+ * If applicable, prepare DH parameters from Server certificate. */
+static int ssl_in_server_key_exchange_prepare( mbedtls_ssl_context *ssl );
+/* Parse SrvKeyExchange message and store contents
+ * (PSK or DH parameters) in handshake structure. */
+static int ssl_in_server_key_exchange_parse( mbedtls_ssl_context *ssl,
+                                          unsigned char *buf,
+                                          size_t buflen );
+/* Update the handshake state */
+static int ssl_in_server_key_exchange_postprocess( mbedtls_ssl_context *ssl );
+
+/*
+ * Implementation
+ */
+
+static int ssl_in_server_key_exchange_prepare( mbedtls_ssl_context *ssl )
+{
+    mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
+        mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
+    ((void) ciphersuite_info);
+
+    /* If applicable, extract static DH parameters from Server CRT. */
+
+#if defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) || \
+    defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
+        MBEDTLS_KEY_EXCHANGE_ECDH_RSA ||
+        mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
+        MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
+    {
+        int ret;
+
+        if( ( ret = ssl_get_ecdh_params_from_cert( ssl ) ) != 0 )
+        {
+            MBEDTLS_SSL_DEBUG_RET( 1, "ssl_get_ecdh_params_from_cert", ret );
+            mbedtls_ssl_pend_fatal_alert( ssl,
+                                     MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE );
+            return( ret );
+        }
+    }
+#endif /* MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED ||
+          MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED */
+
+    return( 0 );
+}
+
+static int ssl_in_server_key_exchange_coordinate( mbedtls_ssl_context *ssl )
 {
     int ret;
     mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
         mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
-    unsigned char *p = NULL, *end = NULL;
 
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> parse server key exchange" ) );
+    /* The ServerKeyExchange message is not used for
+     * - RSA or
+     * - static ECDH
+     * ciphersuites.
+     * It MAY be used in PSK or RSA-PSK.
+     */
 
 #if defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED)
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
         MBEDTLS_KEY_EXCHANGE_RSA )
     {
-        MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= skip parse server key exchange" ) );
-        ssl->state++;
-        return( 0 );
+        return( SSL_SRV_KEY_EXCHANGE_SKIP );
     }
-    ((void) p);
-    ((void) end);
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) || \
@@ -2609,79 +2695,56 @@
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
         MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
-        if( ( ret = ssl_get_ecdh_params_from_cert( ssl ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "ssl_get_ecdh_params_from_cert", ret );
-            mbedtls_ssl_pend_fatal_alert( ssl,
-                                    MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE );
-            return( ret );
-        }
-
-        MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= skip parse server key exchange" ) );
-        ssl->state++;
-        return( 0 );
+        return( SSL_SRV_KEY_EXCHANGE_SKIP );
     }
-    ((void) p);
-    ((void) end);
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED ||
           MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED */
 
-#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
-    if( ssl->handshake->ecrs_enabled &&
-        ssl->handshake->ecrs_state == ssl_ecrs_ske_start_processing )
-    {
-        goto start_processing;
-    }
-#endif
-
-    if( ( ret = mbedtls_ssl_read_record( ssl, 1 ) ) != 0 )
-    {
-        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_read_record", ret );
-        return( ret );
-    }
-
-    if( ssl->in_msgtype != MBEDTLS_SSL_MSG_HANDSHAKE )
-    {
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad server key exchange message" ) );
-        mbedtls_ssl_pend_fatal_alert( ssl,
-                               MBEDTLS_SSL_ALERT_MSG_UNEXPECTED_MESSAGE );
-        return( MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE );
-    }
-
     /*
      * ServerKeyExchange may be skipped with PSK and RSA-PSK when the server
-     * doesn't use a psk_identity_hint
+     * doesn't use a psk_identity_hint. Peek at next message to decide whether
+     * the ServerKeyExchange is being skipped or not.
      */
-    if( ssl->in_msg[0] != MBEDTLS_SSL_HS_SERVER_KEY_EXCHANGE )
+
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_PSK ||
+        mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA_PSK )
     {
-        if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-            == MBEDTLS_KEY_EXCHANGE_PSK ||
-            mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-            == MBEDTLS_KEY_EXCHANGE_RSA_PSK )
+        if( ( ret = mbedtls_ssl_read_record( ssl, 1 ) ) != 0 )
+        {
+            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_read_record", ret );
+            return( ret );
+        }
+        ssl->keep_current_message = 1;
+
+        if( ssl->in_msgtype == MBEDTLS_SSL_MSG_HANDSHAKE &&
+            ssl->in_msg[0]  != MBEDTLS_SSL_HS_SERVER_KEY_EXCHANGE )
         {
             /* Current message is probably either
              * CertificateRequest or ServerHelloDone */
-            ssl->keep_current_message = 1;
-            goto exit;
+            return( SSL_SRV_KEY_EXCHANGE_SKIP );
         }
-
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "server key exchange message must "
-                                    "not be skipped" ) );
-        mbedtls_ssl_pend_fatal_alert( ssl,
-                               MBEDTLS_SSL_ALERT_MSG_UNEXPECTED_MESSAGE );
-
-        return( MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE );
     }
 
-#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
-    if( ssl->handshake->ecrs_enabled )
-        ssl->handshake->ecrs_state = ssl_ecrs_ske_start_processing;
+    return( SSL_SRV_KEY_EXCHANGE_EXPECTED );
+}
 
-start_processing:
-#endif
-    p   = ssl->in_msg + mbedtls_ssl_hs_hdr_len( ssl );
-    end = ssl->in_msg + ssl->in_hslen;
-    MBEDTLS_SSL_DEBUG_BUF( 3,   "server key exchange", p, end - p );
+static int ssl_in_server_key_exchange_parse( mbedtls_ssl_context *ssl,
+                                          unsigned char *buf,
+                                          size_t buflen )
+{
+    int ret;
+    unsigned char *p;
+    unsigned char *end;
+
+    mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
+        mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
+
+    p   = buf + mbedtls_ssl_hs_hdr_len( ssl );
+    end = buf + buflen;
+
+    MBEDTLS_SSL_DEBUG_BUF( 3, "server key exchange", p, end - p );
 
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
@@ -2772,6 +2835,10 @@
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
     {
+        ((void) ret);
+        ((void) p);
+        ((void) end);
+        ((void) ciphersuite_info);
         MBEDTLS_SSL_DEBUG_MSG( 1, ( "should never happen" ) );
         return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
     }
@@ -2961,12 +3028,73 @@
     }
 #endif /* MBEDTLS_KEY_EXCHANGE__WITH_SERVER_SIGNATURE__ENABLED */
 
-exit:
-    ssl->state++;
+    return( 0 );
+}
+
+static int ssl_in_server_key_exchange_postprocess( mbedtls_ssl_context *ssl )
+{
+    ssl->state = MBEDTLS_SSL_CERTIFICATE_REQUEST;
+    return( 0 );
+}
+
+static int ssl_process_in_server_key_exchange( mbedtls_ssl_context *ssl )
+{
+    int ret;
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> parse server key exchange" ) );
+
+    /* Preparation:
+     * Potentially extract DH parameters from Server's certificate.
+     *
+     * Consider: Why don't we do this as post-processing after
+     *           the server certificate has been read?
+     */
+    MBEDTLS_SSL_CHK( ssl_in_server_key_exchange_prepare( ssl ) );
+
+    /* Coordination:
+     * Check if we expect a ServerKeyExchange */
+    MBEDTLS_SSL_CHK( ssl_in_server_key_exchange_coordinate( ssl ) );
+
+    if( ret == SSL_SRV_KEY_EXCHANGE_EXPECTED )
+    {
+        /* Reading step */
+        if( ( ret = mbedtls_ssl_read_record( ssl, 1 ) ) != 0 )
+        {
+            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_read_record", ret );
+            return( ret );
+        }
+
+        if( ssl->in_msgtype != MBEDTLS_SSL_MSG_HANDSHAKE ||
+            ssl->in_msg[0]  != MBEDTLS_SSL_HS_SERVER_KEY_EXCHANGE )
+        {
+            MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad server key exchange message" ) );
+            mbedtls_ssl_pend_fatal_alert( ssl,
+                                     MBEDTLS_SSL_ALERT_MSG_UNEXPECTED_MESSAGE );
+            ret = MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE;
+            goto cleanup;
+        }
+        else
+        {
+            MBEDTLS_SSL_CHK( ssl_in_server_key_exchange_parse( ssl, ssl->in_msg,
+                                                         ssl->in_hslen ) );
+        }
+    }
+    else if( ret == SSL_SRV_KEY_EXCHANGE_SKIP )
+    {
+        MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= skip parse server key exchange" ) );
+    }
+
+    /* Update state */
+    MBEDTLS_SSL_CHK( ssl_in_server_key_exchange_postprocess( ssl ) );
+
+cleanup:
+
+#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
+    if( ret == MBEDTLS_ERR_SSL_CRYPTO_IN_PROGRESS )
+        ssl->keep_current_message = 1;
+#endif
 
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= parse server key exchange" ) );
-
-    return( 0 );
+    return( ret );
 }
 
 #if ! defined(MBEDTLS_KEY_EXCHANGE__CERT_REQ_ALLOWED__ENABLED)
@@ -3191,53 +3319,209 @@
     return( 0 );
 }
 
-static int ssl_write_client_key_exchange( mbedtls_ssl_context *ssl )
+/*
+ *
+ * STATE HANDLING: Client Key Exchange
+ *
+ */
+
+/*
+ * Overview
+ */
+
+/* Main entry point; orchestrates the other functions */
+static int ssl_process_out_client_key_exchange( mbedtls_ssl_context *ssl );
+
+/* Preparation
+ * - For ECDH: Generate client params and derive premaster secret
+ * - For RSA-suites: Encrypt PMS
+ * - For ECJPAKE: Do Round 2
+ */
+static int ssl_out_client_key_exchange_prepare( mbedtls_ssl_context *ssl );
+static int ssl_out_client_key_exchange_write( mbedtls_ssl_context *ssl,
+                                          unsigned char *buf,
+                                          size_t buflen,
+                                          size_t *olen );
+static int ssl_out_client_key_exchange_postprocess( mbedtls_ssl_context *ssl );
+
+/*
+ * Implementation
+ */
+
+static int ssl_process_out_client_key_exchange( mbedtls_ssl_context *ssl )
 {
-    int ret;
-    size_t i, n;
+    int ret = 0;
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> process client key exchange" ) );
+
+#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
+    if( ssl->handshake->ecrs_state == ssl_ecrs_cke_ecdh_calc_secret )
+        goto cli_key_exchange_postprocess;
+
+    if( ssl->handshake->ecrs_enabled )
+        mbedtls_ecdh_enable_restart( &ssl->handshake->ecdh_ctx );
+#endif
+
+    MBEDTLS_SSL_CHK( ssl_out_client_key_exchange_prepare( ssl ) );
+
+     /* Prepare CertificateVerify message in output buffer. */
+    MBEDTLS_SSL_CHK( ssl_out_client_key_exchange_write( ssl, ssl->out_msg,
+                                                 MBEDTLS_SSL_MAX_CONTENT_LEN,
+                                                 &ssl->out_msglen ) );
+
+    ssl->out_msgtype = MBEDTLS_SSL_MSG_HANDSHAKE;
+    ssl->out_msg[0]  = MBEDTLS_SSL_HS_CLIENT_KEY_EXCHANGE;
+
+    /* Calculate secrets and update state */
+#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
+    if( ssl->handshake->ecrs_enabled )
+        ssl->handshake->ecrs_state = ssl_ecrs_cke_ecdh_calc_secret;
+
+cli_key_exchange_postprocess:
+#endif
+
+    ret = ssl_out_client_key_exchange_postprocess( ssl );
+#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
+    if( ret == MBEDTLS_ERR_ECP_IN_PROGRESS )
+        ret = MBEDTLS_ERR_SSL_CRYPTO_IN_PROGRESS;
+#endif /* MBEDTLS_SSL__ECP_RESTARTABLE */
+    MBEDTLS_SSL_CHK( ret );
+
+    /* Dispatch message */
+
+    MBEDTLS_SSL_CHK( mbedtls_ssl_write_handshake_msg( ssl ) );
+
+    /* NOTE: For the new messaging layer, the postprocessing step
+     *       might come after the dispatching step if the latter
+     *       doesn't send the message immediately.
+     *       At the moment, we must do the postprocessing
+     *       prior to the dispatching because if the latter
+     *       returns WANT_WRITE, we want the handshake state
+     *       to be updated in order to not enter
+     *       this function again on retry. */
+
+cleanup:
+
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= process client key exchange" ) );
+    return( ret );
+}
+
+
+
+static int ssl_out_client_key_exchange_prepare( mbedtls_ssl_context *ssl )
+{
+    int ret = 0;
     mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
         mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
 
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> write client key exchange" ) );
+    ((void) ret);
+    ((void) ciphersuite_info);
+
+    /* TODO: The current API for DH and ECDH does not allow
+     * to separate public key generation from public key export.
+     *
+     * Ideally, we would like to pick the private (EC)DH keys
+     * in this preparation step, exporting the corresponding
+     * public key in the writing step only.
+     *
+     * The necessary extension of the (EC)DH API is being
+     * considered, but until then we perform the public
+     * generation + export in the writing step.
+     *
+     */
+
+#if defined(MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA_PSK )
+    {
+        /* For RSA-PSK, the premaster secret is composed of
+         * - Length tag with value 48, encoded as a uint16
+         * - 2 bytes indicating the TLS version
+         * - 46 randomly chosen bytes
+         * - the chosen PSK.
+         * The following call takes care of all but the PSK. */
+        ret = ssl_rsa_generate_partial_pms( ssl, ssl->handshake->premaster,
+                                            1 /* Add length tag */ );
+        if( ret != 0 )
+            return( ret );
+    }
+#endif /* MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED */
+
+#if defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA )
+    {
+        /* For RSA, the premaster secret is composed of
+         * - 2 bytes indicating the TLS version
+         * - 46 randomly chosen bytes
+         * which the following call generates. */
+        ret = ssl_rsa_generate_partial_pms( ssl, ssl->handshake->premaster,
+                                            0 /* Omit length tag */ );
+        if( ret != 0 )
+            return( ret );
+    }
+#endif /* MBEDTLS_KEY_EXCHANGE_RSA_ENABLED */
+
+    return( 0 );
+}
+
+/* Warning: Despite accepting a length argument, this function is currently
+ * still lacking some bounds checks and assumes that `buf` has length
+ * `MBEDTLS_SSL_OUT_CONTENT_LEN`. Eventually, it should be rewritten to work
+ * with any buffer + length pair, returning MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL
+ * on insufficient writing space. */
+static int ssl_out_client_key_exchange_write( mbedtls_ssl_context *ssl,
+                                          unsigned char *buf,
+                                          size_t buflen,
+                                          size_t *olen )
+{
+    int ret;
+    unsigned char *p, *end;
+    size_t n;
+    mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
+        mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
+
+    /* NOTE: This function will generate different messages
+     * when it's called multiple times, because it currently
+     * includes private/public key generation in case of
+     * (EC)DH ciphersuites.
+     *
+     * It is therefore not suitable to be registered as a callback
+     * for retransmission, if such get introduced at some point.
+     *
+     * Also see the documentation of ssl_out_client_key_exchange_prepare().
+     */
+
+    p   = buf + 4;
+    end = buf + buflen;
 
 #if defined(MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) == MBEDTLS_KEY_EXCHANGE_DHE_RSA )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_DHE_RSA )
     {
         /*
          * DHM key exchange -- send G^X mod P
          */
         n = ssl->handshake->dhm_ctx.len;
+        if( (size_t)( end - p ) < n + 2 )
+            return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
 
-        ssl->out_msg[4] = (unsigned char)( n >> 8 );
-        ssl->out_msg[5] = (unsigned char)( n      );
-        i = 6;
+        p[0] = (unsigned char)( n >> 8 );
+        p[1] = (unsigned char)( n      );
+        p += 2;
 
         ret = mbedtls_dhm_make_public( &ssl->handshake->dhm_ctx,
-                                (int) mbedtls_mpi_size( &ssl->handshake->dhm_ctx.P ),
-                               &ssl->out_msg[i], n,
-                                mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                ssl->conf->p_rng );
+                           (int) mbedtls_mpi_size( &ssl->handshake->dhm_ctx.P ),
+                           p, n, mbedtls_ssl_conf_get_frng( ssl->conf ),
+                           mbedtls_ssl_conf_get_prng( ssl->conf ) );
         if( ret != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_make_public", ret );
             return( ret );
         }
-
         MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: X ", &ssl->handshake->dhm_ctx.X  );
         MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: GX", &ssl->handshake->dhm_ctx.GX );
 
-        if( ( 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 ),
-                                      ssl->conf->p_rng ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );
-            return( ret );
-        }
-
-        MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: K ", &ssl->handshake->dhm_ctx.K  );
+        p += n;
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED */
@@ -3246,34 +3530,26 @@
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||                      \
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED)
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-        == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA                       ||
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-        == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA                     ||
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA ||
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-        == MBEDTLS_KEY_EXCHANGE_ECDH_RSA                        ||
+        == MBEDTLS_KEY_EXCHANGE_ECDH_RSA ||
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
         == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
         /*
-         * ECDH key exchange -- send client public value
+         * ECDH key exchange -- generate and send client public value
          */
-        i = 4;
 
-#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
-        if( ssl->handshake->ecrs_enabled )
-        {
-            if( ssl->handshake->ecrs_state == ssl_ecrs_cke_ecdh_calc_secret )
-                goto ecdh_calc_secret;
-
-            mbedtls_ecdh_enable_restart( &ssl->handshake->ecdh_ctx );
-        }
-#endif
-
+        /* NOTE: If we ever switch the ECDH stack/API to allow
+         * independent key generation and export, we should have
+         * generated our key share in the preparation step, and
+         * we'd only do the export here. */
         ret = mbedtls_ecdh_make_public( &ssl->handshake->ecdh_ctx,
-                                &n,
-                                &ssl->out_msg[i], 1000,
+                                &n, p, end - p,
                                 mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                ssl->conf->p_rng );
+                                mbedtls_ssl_conf_get_prng( ssl->conf ) );
         if( ret != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_make_public", ret );
@@ -3287,34 +3563,7 @@
         MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
                                 MBEDTLS_DEBUG_ECDH_Q );
 
-#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
-        if( ssl->handshake->ecrs_enabled )
-        {
-            ssl->handshake->ecrs_n = n;
-            ssl->handshake->ecrs_state = ssl_ecrs_cke_ecdh_calc_secret;
-        }
-
-ecdh_calc_secret:
-        if( ssl->handshake->ecrs_enabled )
-            n = ssl->handshake->ecrs_n;
-#endif
-        if( ( 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 ),
-                                       ssl->conf->p_rng ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
-#if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
-            if( ret == MBEDTLS_ERR_ECP_IN_PROGRESS )
-                ret = MBEDTLS_ERR_SSL_CRYPTO_IN_PROGRESS;
-#endif
-            return( ret );
-        }
-
-        MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
-                                MBEDTLS_DEBUG_ECDH_Z );
+        p += n;
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
@@ -3333,27 +3582,27 @@
             return( MBEDTLS_ERR_SSL_PRIVATE_KEY_REQUIRED );
         }
 
-        i = 4;
         n = ssl->conf->psk_identity_len;
-
-        if( i + 2 + n > MBEDTLS_SSL_OUT_CONTENT_LEN )
+        if( buflen < n + 2 )
         {
             MBEDTLS_SSL_DEBUG_MSG( 1, ( "psk identity too long or "
                                         "SSL buffer too short" ) );
             return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
         }
 
-        ssl->out_msg[i++] = (unsigned char)( n >> 8 );
-        ssl->out_msg[i++] = (unsigned char)( n      );
+        p[0] = (unsigned char)( n >> 8 );
+        p[1] = (unsigned char)( n      );
+        p += 2;
 
-        memcpy( ssl->out_msg + i, ssl->conf->psk_identity, ssl->conf->psk_identity_len );
-        i += ssl->conf->psk_identity_len;
+        memcpy( p, ssl->conf->psk_identity, n );
+        p += ssl->conf->psk_identity_len;
 
 #if defined(MBEDTLS_KEY_EXCHANGE_PSK_ENABLED)
         if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
             == MBEDTLS_KEY_EXCHANGE_PSK )
         {
-            n = 0;
+            ((void) ret);
+            ((void) end);
         }
         else
 #endif
@@ -3361,8 +3610,11 @@
         if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
             == MBEDTLS_KEY_EXCHANGE_RSA_PSK )
         {
-            if( ( ret = ssl_write_encrypted_pms( ssl, i, &n, 2 ) ) != 0 )
+            if( ( ret = ssl_rsa_encrypt_partial_pms( ssl,
+                                                ssl->handshake->premaster + 2,
+                                                p, end - p, &n ) ) != 0 )
                 return( ret );
+            p += n;
         }
         else
 #endif
@@ -3373,28 +3625,32 @@
             /*
              * ClientDiffieHellmanPublic public (DHM send G^X mod P)
              */
+            ((void) end);
+
             n = ssl->handshake->dhm_ctx.len;
 
-            if( i + 2 + n > MBEDTLS_SSL_OUT_CONTENT_LEN )
+            if( buflen < n + 2 )
             {
                 MBEDTLS_SSL_DEBUG_MSG( 1, ( "psk identity or DHM size too long"
                                             " or SSL buffer too short" ) );
                 return( MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL );
             }
 
-            ssl->out_msg[i++] = (unsigned char)( n >> 8 );
-            ssl->out_msg[i++] = (unsigned char)( n      );
+            p[0] = (unsigned char)( n >> 8 );
+            p[1] = (unsigned char)( n      );
+            p += 2;
 
             ret = mbedtls_dhm_make_public( &ssl->handshake->dhm_ctx,
-                    (int) mbedtls_mpi_size( &ssl->handshake->dhm_ctx.P ),
-                    &ssl->out_msg[i], n,
-                    mbedtls_ssl_conf_get_frng( ssl->conf ),
-                    ssl->conf->p_rng );
+                           (int) mbedtls_mpi_size( &ssl->handshake->dhm_ctx.P ),
+                           p, n, mbedtls_ssl_conf_get_frng( ssl->conf ),
+                           mbedtls_ssl_conf_get_prng( ssl->conf ) );
             if( ret != 0 )
             {
                 MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_make_public", ret );
                 return( ret );
             }
+
+            p += n;
         }
         else
 #endif /* MBEDTLS_KEY_EXCHANGE_DHE_PSK_ENABLED */
@@ -3406,17 +3662,17 @@
              * ClientECDiffieHellmanPublic public;
              */
             ret = mbedtls_ecdh_make_public( &ssl->handshake->ecdh_ctx, &n,
-                    &ssl->out_msg[i], MBEDTLS_SSL_OUT_CONTENT_LEN - i,
-                    mbedtls_ssl_conf_get_frng( ssl->conf ),
-                    ssl->conf->p_rng );
+                                p, buflen,
+                                mbedtls_ssl_conf_get_frng( ssl->conf ),
+                                mbedtls_ssl_conf_get_prng( ssl->conf ) );
             if( ret != 0 )
             {
                 MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_make_public", ret );
                 return( ret );
             }
+            MBEDTLS_SSL_DEBUG_ECP( 3, "ECDH: Q", &ssl->handshake->ecdh_ctx.Q );
 
-            MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
-                                    MBEDTLS_DEBUG_ECDH_Q );
+            p += n;
         }
         else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
@@ -3424,51 +3680,34 @@
             MBEDTLS_SSL_DEBUG_MSG( 1, ( "should never happen" ) );
             return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
         }
-
-        if( ( ret = mbedtls_ssl_psk_derive_premaster( ssl,
-               mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
-            return( ret );
-        }
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_RSA )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA )
     {
-        i = 4;
-        if( ( ret = ssl_write_encrypted_pms( ssl, i, &n, 0 ) ) != 0 )
+        if( ( ret = ssl_rsa_encrypt_partial_pms( ssl, ssl->handshake->premaster,
+                                                 p, end - p, &n ) ) != 0 )
             return( ret );
+        p += n;
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_RSA_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_ECJPAKE )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
     {
-        i = 4;
-
         ret = mbedtls_ecjpake_write_round_two( &ssl->handshake->ecjpake_ctx,
-                ssl->out_msg + i, MBEDTLS_SSL_OUT_CONTENT_LEN - i, &n,
-                mbedtls_ssl_conf_get_frng( ssl->conf ),
-                ssl->conf->p_rng );
+                          p, end - p, &n,
+                          mbedtls_ssl_conf_get_frng( ssl->conf ),
+                          mbedtls_ssl_conf_get_prng( ssl->conf ) );
         if( ret != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_write_round_two", ret );
             return( ret );
         }
-
-        ret = mbedtls_ecjpake_derive_secret( &ssl->handshake->ecjpake_ctx,
-                ssl->handshake->premaster, 32, &ssl->handshake->pmslen,
-                mbedtls_ssl_conf_get_frng( ssl->conf ),
-                ssl->conf->p_rng );
-        if( ret != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
-            return( ret );
-        }
+        p += n;
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_RSA_ENABLED */
@@ -3478,20 +3717,21 @@
         return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
     }
 
-    ssl->out_msglen  = i + n;
-    ssl->out_msgtype = MBEDTLS_SSL_MSG_HANDSHAKE;
-    ssl->out_msg[0]  = MBEDTLS_SSL_HS_CLIENT_KEY_EXCHANGE;
+    *olen = p - buf;
+    return( 0 );
+}
 
-    ssl->state++;
+static int ssl_out_client_key_exchange_postprocess( mbedtls_ssl_context *ssl )
+{
+    int ret;
 
-    if( ( ret = mbedtls_ssl_write_handshake_msg( ssl ) ) != 0 )
+    if( ( ret = mbedtls_ssl_build_pms( ssl ) ) != 0 )
     {
-        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_write_handshake_msg", ret );
+        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_build_pms", ret );
         return( ret );
     }
 
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= write client key exchange" ) );
-
+    ssl->state = MBEDTLS_SSL_CERTIFICATE_VERIFY;
     return( 0 );
 }
 
@@ -3664,7 +3904,7 @@
                          md_alg, hash_start, hashlen,
                          ssl->out_msg + 6 + offset, &n,
                          mbedtls_ssl_conf_get_frng( ssl->conf ),
-                         ssl->conf->p_rng, rs_ctx ) ) != 0 )
+                         mbedtls_ssl_conf_get_prng( ssl->conf ), rs_ctx ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_pk_sign", ret );
 #if defined(MBEDTLS_SSL__ECP_RESTARTABLE)
@@ -3873,7 +4113,7 @@
            break;
 
        case MBEDTLS_SSL_SERVER_KEY_EXCHANGE:
-           ret = ssl_parse_server_key_exchange( ssl );
+           ret = ssl_process_in_server_key_exchange( ssl );
            break;
 
        case MBEDTLS_SSL_CERTIFICATE_REQUEST:
@@ -3896,7 +4136,7 @@
            break;
 
        case MBEDTLS_SSL_CLIENT_KEY_EXCHANGE:
-           ret = ssl_write_client_key_exchange( ssl );
+           ret = ssl_process_out_client_key_exchange( ssl );
            break;
 
        case MBEDTLS_SSL_CERTIFICATE_VERIFY:
diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index 5a6441f..b2b6144 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -2603,7 +2603,7 @@
     ret = mbedtls_ecjpake_write_round_one( &ssl->handshake->ecjpake_ctx,
                                         p + 2, end - p - 2, &kkpp_len,
                                         mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                        ssl->conf->p_rng );
+                                        mbedtls_ssl_conf_get_prng( ssl->conf ) );
     if( ret != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1 , "mbedtls_ecjpake_write_round_one", ret );
@@ -2783,7 +2783,7 @@
     MBEDTLS_SSL_DEBUG_MSG( 3, ( "server hello, current time: %lu", t ) );
 #else
     if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-          ( ssl->conf->p_rng, p, 4 ) ) != 0 )
+          ( mbedtls_ssl_conf_get_prng( ssl->conf ), p, 4 ) ) != 0 )
     {
         return( ret );
     }
@@ -2792,7 +2792,7 @@
 #endif /* MBEDTLS_HAVE_TIME */
 
     if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-          ( ssl->conf->p_rng, p, 28 ) ) != 0 )
+          ( mbedtls_ssl_conf_get_prng( ssl->conf ), p, 28 ) ) != 0 )
     {
         return( ret );
     }
@@ -2859,7 +2859,7 @@
         {
             ssl->session_negotiate->id_len = n = 32;
             if( ( ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-                  ( ssl->conf->p_rng, ssl->session_negotiate->id, n ) ) != 0 )
+                  ( mbedtls_ssl_conf_get_prng( ssl->conf ), ssl->session_negotiate->id, n ) ) != 0 )
             {
                 return( ret );
             }
@@ -3273,7 +3273,7 @@
             ssl->out_msg + ssl->out_msglen,
             MBEDTLS_SSL_OUT_CONTENT_LEN - ssl->out_msglen, &len,
             mbedtls_ssl_conf_get_frng( ssl->conf ),
-            ssl->conf->p_rng );
+            mbedtls_ssl_conf_get_prng( ssl->conf ) );
         if( ret != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_write_round_two", ret );
@@ -3337,7 +3337,7 @@
                   (int) mbedtls_mpi_size( &ssl->handshake->dhm_ctx.P ),
                   ssl->out_msg + ssl->out_msglen, &len,
                   mbedtls_ssl_conf_get_frng( ssl->conf ),
-                  ssl->conf->p_rng ) ) != 0 )
+                  mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_make_params", ret );
             return( ret );
@@ -3394,7 +3394,7 @@
                   ssl->out_msg + ssl->out_msglen,
                   MBEDTLS_SSL_OUT_CONTENT_LEN - ssl->out_msglen,
                   mbedtls_ssl_conf_get_frng( ssl->conf ),
-                  ssl->conf->p_rng ) ) != 0 )
+                  mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_make_params", ret );
             return( ret );
@@ -3580,7 +3580,7 @@
                                      ssl->out_msg + ssl->out_msglen + 2,
                                      signature_len,
                                      mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                     ssl->conf->p_rng ) ) != 0 )
+                                     mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_pk_sign", ret );
             return( ret );
@@ -3877,7 +3877,7 @@
     ret = mbedtls_pk_decrypt( private_key, p, len,
                               peer_pms, peer_pmslen, peer_pmssize,
                               mbedtls_ssl_conf_get_frng( ssl->conf ),
-                              ssl->conf->p_rng );
+                              mbedtls_ssl_conf_get_prng( ssl->conf ) );
     return( ret );
 }
 
@@ -3947,7 +3947,7 @@
      * regardless of whether it will ultimately influence the output or not.
      */
     ret = mbedtls_ssl_conf_get_frng( ssl->conf )
-        ( ssl->conf->p_rng, fake_pms, sizeof( fake_pms ) );
+        ( mbedtls_ssl_conf_get_prng( ssl->conf ), fake_pms, sizeof( fake_pms ) );
     if( ret != 0 )
     {
         /* It's ok to abort on an RNG failure, since this does not reveal
@@ -4041,49 +4041,87 @@
 }
 #endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED */
 
-static int ssl_parse_client_key_exchange( mbedtls_ssl_context *ssl )
+/*
+ *
+ * STATE HANDLING: Client Key Exchange
+ *
+ */
+
+/*
+ * Overview
+ */
+
+/* Main entry point; orchestrates the other functions. */
+static int ssl_process_in_client_key_exchange( mbedtls_ssl_context *ssl );
+
+static int ssl_in_client_key_exchange_parse( mbedtls_ssl_context *ssl,
+                                             unsigned char *buf,
+                                             size_t buflen );
+/* Update the handshake state */
+static int ssl_in_client_key_exchange_postprocess( mbedtls_ssl_context *ssl );
+
+/*
+ * Implementation
+ */
+
+static int ssl_process_in_client_key_exchange( mbedtls_ssl_context *ssl )
+{
+    int ret;
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> process client key exchange" ) );
+
+    /* The ClientKeyExchange message is never skipped. */
+
+    /* Reading step */
+    if( ( ret = mbedtls_ssl_read_record( ssl,
+                                         1 /* update checksum */ ) ) != 0 )
+    {
+        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_read_record", ret );
+        return( ret );
+    }
+
+    if( ssl->in_msgtype != MBEDTLS_SSL_MSG_HANDSHAKE ||
+        ssl->in_msg[0]  != MBEDTLS_SSL_HS_CLIENT_KEY_EXCHANGE )
+    {
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad client key exchange message" ) );
+        mbedtls_ssl_pend_fatal_alert( ssl,
+                                MBEDTLS_SSL_ALERT_MSG_UNEXPECTED_MESSAGE );
+        ret = MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE;
+        goto cleanup;
+    }
+
+    MBEDTLS_SSL_CHK( ssl_in_client_key_exchange_parse( ssl, ssl->in_msg,
+                                                       ssl->in_hslen ) );
+
+    /* Update state */
+    MBEDTLS_SSL_CHK( ssl_in_client_key_exchange_postprocess( ssl ) );
+
+cleanup:
+
+#if defined(MBEDTLS_SSL_ASYNC_PRIVATE)
+    if ( ret == MBEDTLS_ERR_SSL_ASYNC_IN_PROGRESS )
+        ssl->keep_current_message = 1;
+#endif /* MBEDTLS_SSL_ASYNC_PRIVATE */
+
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= process client key exchange" ) );
+    return( ret );
+}
+
+/* Warning: Despite accepting a length argument, this function is currently
+ * still lacking some bounds checks and assumes that `buf` has length
+ * `MBEDTLS_SSL_IN_CONTENT_LEN`. Eventually, it should be rewritten to work
+ * with any buffer + length pair, returning MBEDTLS_ERR_SSL_BUFFER_TOO_SMALL
+ * on insufficient parsing space. */
+static int ssl_in_client_key_exchange_parse( mbedtls_ssl_context *ssl,
+                                          unsigned char *buf,
+                                          size_t buflen )
 {
     int ret;
     mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
         mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
     unsigned char *p, *end;
 
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> parse client key exchange" ) );
-
-#if defined(MBEDTLS_SSL_ASYNC_PRIVATE) && \
-    ( defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED) || \
-      defined(MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED) )
-    if( ( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) == MBEDTLS_KEY_EXCHANGE_RSA_PSK ||
-          mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) == MBEDTLS_KEY_EXCHANGE_RSA ) &&
-        ( ssl->handshake->async_in_progress != 0 ) )
-    {
-        /* We've already read a record and there is an asynchronous
-         * operation in progress to decrypt it. So skip reading the
-         * record. */
-        MBEDTLS_SSL_DEBUG_MSG( 3, ( "will resume decryption of previously-read record" ) );
-    }
-    else
-#endif
-    if( ( ret = mbedtls_ssl_read_record( ssl, 1 ) ) != 0 )
-    {
-        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_read_record", ret );
-        return( ret );
-    }
-
-    p = ssl->in_msg + mbedtls_ssl_hs_hdr_len( ssl );
-    end = ssl->in_msg + ssl->in_hslen;
-
-    if( ssl->in_msgtype != MBEDTLS_SSL_MSG_HANDSHAKE )
-    {
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad client key exchange message" ) );
-        return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
-    }
-
-    if( ssl->in_msg[0] != MBEDTLS_SSL_HS_CLIENT_KEY_EXCHANGE )
-    {
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad client key exchange message" ) );
-        return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
-    }
+    p = buf + mbedtls_ssl_hs_hdr_len( ssl );
+    end = buf + buflen;
 
 #if defined(MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED)
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
@@ -4100,19 +4138,6 @@
             MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad client key exchange" ) );
             return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
         }
-
-        if( ( 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 ),
-                                      ssl->conf->p_rng ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );
-            return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE_CS );
-        }
-
-        MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: K ", &ssl->handshake->dhm_ctx.K  );
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED */
@@ -4121,37 +4146,24 @@
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||                      \
     defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED)
     if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-        == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA                       ||
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA ||
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-        == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA                     ||
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA ||
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
-        == MBEDTLS_KEY_EXCHANGE_ECDH_RSA                        ||
+        == MBEDTLS_KEY_EXCHANGE_ECDH_RSA ||
         mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
         == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
     {
         if( ( ret = mbedtls_ecdh_read_public( &ssl->handshake->ecdh_ctx,
                                       p, end - p) ) != 0 )
         {
+            ((void) ret);
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_read_public", ret );
             return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE_RP );
         }
 
         MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
                                 MBEDTLS_DEBUG_ECDH_QP );
-
-        if( ( 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 ),
-                                       ssl->conf->p_rng ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
-            return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE_CS );
-        }
-
-        MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
-                                MBEDTLS_DEBUG_ECDH_Z );
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
@@ -4159,8 +4171,8 @@
           MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED ||
           MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_PSK_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_PSK )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_PSK )
     {
         if( ( ret = ssl_parse_client_psk_identity( ssl, &p, end ) ) != 0 )
         {
@@ -4173,19 +4185,12 @@
             MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad client key exchange" ) );
             return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
         }
-
-        if( ( ret = mbedtls_ssl_psk_derive_premaster( ssl,
-               mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
-            return( ret );
-        }
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_PSK_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_RSA_PSK )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA_PSK )
     {
 #if defined(MBEDTLS_SSL_ASYNC_PRIVATE)
         if ( ssl->handshake->async_in_progress != 0 )
@@ -4211,19 +4216,12 @@
             MBEDTLS_SSL_DEBUG_RET( 1, ( "ssl_parse_encrypted_pms" ), ret );
             return( ret );
         }
-
-        if( ( ret = mbedtls_ssl_psk_derive_premaster( ssl,
-               mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
-            return( ret );
-        }
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_RSA_PSK_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_DHE_PSK_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_DHE_PSK )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_DHE_PSK )
     {
         if( ( ret = ssl_parse_client_psk_identity( ssl, &p, end ) ) != 0 )
         {
@@ -4241,19 +4239,12 @@
             MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad client key exchange" ) );
             return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
         }
-
-        if( ( ret = mbedtls_ssl_psk_derive_premaster( ssl,
-               mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
-            return( ret );
-        }
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_DHE_PSK_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_ECDHE_PSK )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_PSK )
     {
         if( ( ret = ssl_parse_client_psk_identity( ssl, &p, end ) ) != 0 )
         {
@@ -4268,21 +4259,13 @@
             return( MBEDTLS_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE_RP );
         }
 
-        MBEDTLS_SSL_DEBUG_ECDH( 3, &ssl->handshake->ecdh_ctx,
-                                MBEDTLS_DEBUG_ECDH_QP );
-
-        if( ( ret = mbedtls_ssl_psk_derive_premaster( ssl,
-               mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ) ) != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
-            return( ret );
-        }
+        MBEDTLS_SSL_DEBUG_ECP( 3, "ECDH: Qp ", &ssl->handshake->ecdh_ctx.Qp );
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_RSA )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA )
     {
         if( ( ret = ssl_parse_encrypted_pms( ssl, p, end, 0 ) ) != 0 )
         {
@@ -4293,8 +4276,8 @@
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_RSA_ENABLED */
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
-    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
-        MBEDTLS_KEY_EXCHANGE_ECJPAKE )
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
     {
         ret = mbedtls_ecjpake_read_round_two( &ssl->handshake->ecjpake_ctx,
                                               p, end - p );
@@ -4303,16 +4286,6 @@
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_read_round_two", ret );
             return( MBEDTLS_ERR_SSL_BAD_HS_SERVER_KEY_EXCHANGE );
         }
-
-        ret = mbedtls_ecjpake_derive_secret( &ssl->handshake->ecjpake_ctx,
-                ssl->handshake->premaster, 32, &ssl->handshake->pmslen,
-                mbedtls_ssl_conf_get_frng( ssl->conf ),
-                ssl->conf->p_rng );
-        if( ret != 0 )
-        {
-            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
-            return( ret );
-        }
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
@@ -4321,16 +4294,27 @@
         return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
     }
 
+    return( 0 );
+}
+
+/* Update the handshake state */
+static int ssl_in_client_key_exchange_postprocess( mbedtls_ssl_context *ssl )
+{
+    int ret;
+
+    if( ( ret = mbedtls_ssl_build_pms( ssl ) ) != 0 )
+    {
+        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_build_pms", ret );
+        return( ret );
+    }
+
     if( ( ret = mbedtls_ssl_derive_keys( ssl ) ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_derive_keys", ret );
         return( ret );
     }
 
-    ssl->state++;
-
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= parse client key exchange" ) );
-
+    ssl->state = MBEDTLS_SSL_CERTIFICATE_VERIFY;
     return( 0 );
 }
 
@@ -4706,7 +4690,7 @@
             break;
 
         case MBEDTLS_SSL_CLIENT_KEY_EXCHANGE:
-            ret = ssl_parse_client_key_exchange( ssl );
+            ret = ssl_process_in_client_key_exchange( ssl );
             break;
 
         case MBEDTLS_SSL_CERTIFICATE_VERIFY:
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index d5fdbba..509fc8c 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1653,6 +1653,111 @@
 #endif /* MBEDTLS_SHA512_C */
 #endif /* MBEDTLS_SSL_PROTO_TLS1_2 */
 
+int mbedtls_ssl_build_pms( mbedtls_ssl_context *ssl )
+{
+    int ret;
+
+    mbedtls_ssl_ciphersuite_handle_t ciphersuite_info =
+        mbedtls_ssl_handshake_get_ciphersuite( ssl->handshake );
+
+#if defined(MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_DHE_RSA )
+    {
+        if( ( 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_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );
+            return( ret );
+        }
+
+        MBEDTLS_SSL_DEBUG_MPI( 3, "DHM: K ", &ssl->handshake->dhm_ctx.K  );
+    }
+    else
+#endif /* MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED */
+#if defined(MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED) ||                     \
+    defined(MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED) ||                   \
+    defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) ||                      \
+    defined(MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_RSA                       ||
+        mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA                     ||
+        mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECDH_RSA                        ||
+        mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA )
+    {
+        if( ( 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_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
+            return( ret );
+        }
+
+        MBEDTLS_SSL_DEBUG_MPI( 3, "ECDH: z", &ssl->handshake->ecdh_ctx.z );
+    }
+    else
+#endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
+          MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED ||
+          MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED ||
+          MBEDTLS_KEY_EXCHANGE_ECDH_ECDSA_ENABLED */
+#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 )
+        {
+            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_psk_derive_premaster", ret );
+            return( ret );
+        }
+    }
+    else
+#endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED */
+#if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info ) ==
+        MBEDTLS_KEY_EXCHANGE_ECJPAKE )
+    {
+        ret = mbedtls_ecjpake_derive_secret( &ssl->handshake->ecjpake_ctx,
+                ssl->handshake->premaster, 32, &ssl->handshake->pmslen,
+                mbedtls_ssl_conf_get_frng( ssl->conf ),
+                mbedtls_ssl_conf_get_prng( ssl->conf ) );
+        if( ret != 0 )
+        {
+            MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
+            return( ret );
+        }
+    }
+    else
+#endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
+#if defined(MBEDTLS_KEY_EXCHANGE_RSA_ENABLED)
+    if( mbedtls_ssl_suite_get_key_exchange( ciphersuite_info )
+        == MBEDTLS_KEY_EXCHANGE_RSA )
+    {
+        ((void) ret);
+        /* The premaster secret has already been set by
+         * ssl_rsa_generate_partial_pms(). Only the
+         * PMS length needs to be set. */
+        ssl->handshake->pmslen = 48;
+    }
+    else
+#endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
+    {
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "should never happen" ) );
+        return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
+    }
+
+    return( 0 );
+}
+
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
 int mbedtls_ssl_psk_derive_premaster( mbedtls_ssl_context *ssl, mbedtls_key_exchange_type_t key_ex )
 {
@@ -1718,7 +1823,7 @@
         if( ( ret = mbedtls_dhm_calc_secret( &ssl->handshake->dhm_ctx,
                                       p + 2, end - ( p + 2 ), &len,
                                       mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                      ssl->conf->p_rng ) ) != 0 )
+                                      mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_dhm_calc_secret", ret );
             return( ret );
@@ -1740,7 +1845,7 @@
         if( ( ret = mbedtls_ecdh_calc_secret( &ssl->handshake->ecdh_ctx, &zlen,
                                        p + 2, end - ( p + 2 ),
                                        mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                       ssl->conf->p_rng ) ) != 0 )
+                                       mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
         {
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecdh_calc_secret", ret );
             return( ret );
@@ -4084,7 +4189,7 @@
 
             if( ( ret = mbedtls_ssl_encrypt_buf( ssl, ssl->transform_out, &rec,
                                          mbedtls_ssl_conf_get_frng( ssl->conf ),
-                                         ssl->conf->p_rng ) ) != 0 )
+                                         mbedtls_ssl_conf_get_prng( ssl->conf ) ) ) != 0 )
             {
                 MBEDTLS_SSL_DEBUG_RET( 1, "ssl_encrypt_buf", ret );
                 return( ret );
@@ -8405,12 +8510,6 @@
     conf->f_rng      = f_rng;
     conf->p_rng      = p_rng;
 }
-#else
-void mbedtls_ssl_conf_rng_ctx( mbedtls_ssl_config *conf,
-                               void *p_rng )
-{
-    conf->p_rng      = p_rng;
-}
 #endif
 
 #if defined(MBEDTLS_DEBUG_C)
diff --git a/programs/ssl/dtls_client.c b/programs/ssl/dtls_client.c
index 5bd5cbe..39a9f76 100644
--- a/programs/ssl/dtls_client.c
+++ b/programs/ssl/dtls_client.c
@@ -95,6 +95,20 @@
 }
 #endif /* MBEDTLS_DEBUG_C */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( int argc, char *argv[] )
 {
     int ret, len;
@@ -192,7 +206,13 @@
      * Production code should set a proper ca chain and use REQUIRED. */
     mbedtls_ssl_conf_authmode( &conf, MBEDTLS_SSL_VERIFY_OPTIONAL );
     mbedtls_ssl_conf_ca_chain( &conf, &cacert, NULL );
+
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/programs/ssl/dtls_server.c b/programs/ssl/dtls_server.c
index 84b905d..1dddf8e 100644
--- a/programs/ssl/dtls_server.c
+++ b/programs/ssl/dtls_server.c
@@ -104,6 +104,20 @@
 }
 #endif /* MBEDTLS_DEBUG_C */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( void )
 {
     int ret, len;
@@ -224,7 +238,12 @@
         goto exit;
     }
 
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/programs/ssl/mini_client.c b/programs/ssl/mini_client.c
index 4231c0c..87c48ff 100644
--- a/programs/ssl/mini_client.c
+++ b/programs/ssl/mini_client.c
@@ -166,6 +166,19 @@
     ssl_write_failed,
 };
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
 
 int main( void )
 {
@@ -212,7 +225,7 @@
 #if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
 #else
-    mbedtls_ssl_conf_rng_ctx( &conf, &ctr_drbg );
+    rng_ctx_global = &ctr_drbg;
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
diff --git a/programs/ssl/ssl_client1.c b/programs/ssl/ssl_client1.c
index 5f453dc..8629980 100644
--- a/programs/ssl/ssl_client1.c
+++ b/programs/ssl/ssl_client1.c
@@ -85,6 +85,20 @@
 }
 #endif /* MBEDTLS_DEBUG_C */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( void )
 {
     int ret = 1, len;
@@ -179,7 +193,13 @@
      * but makes interop easier in this simplified example */
     mbedtls_ssl_conf_authmode( &conf, MBEDTLS_SSL_VERIFY_OPTIONAL );
     mbedtls_ssl_conf_ca_chain( &conf, &cacert, NULL );
+
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 716263b..a7fd259 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -889,6 +889,20 @@
 }
 #endif /* MBEDTLS_SSL_DTLS_CONNECTION_ID */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( int argc, char *argv[] )
 {
     int ret = 0, len, tail_len, i, written, frags, retry_left;
@@ -1942,7 +1956,7 @@
 #if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
 #else
-    mbedtls_ssl_conf_rng_ctx( &conf, &ctr_drbg );
+    rng_ctx_global = &ctr_drbg;
 #endif
 
 #if defined(MBEDTLS_DEBUG_C)
diff --git a/programs/ssl/ssl_fork_server.c b/programs/ssl/ssl_fork_server.c
index 098761c..7033b86 100644
--- a/programs/ssl/ssl_fork_server.c
+++ b/programs/ssl/ssl_fork_server.c
@@ -102,6 +102,20 @@
 }
 #endif /* MBEDTLS_DEBUG_C */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( void )
 {
     int ret = 1, len, cnt = 0, pid;
@@ -196,7 +210,12 @@
         goto exit;
     }
 
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/programs/ssl/ssl_mail_client.c b/programs/ssl/ssl_mail_client.c
index b2e4f7f..205c93e 100644
--- a/programs/ssl/ssl_mail_client.c
+++ b/programs/ssl/ssl_mail_client.c
@@ -361,6 +361,20 @@
     while( 1 );
 }
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( int argc, char *argv[] )
 {
     int ret = 1, len;
@@ -620,7 +634,12 @@
      * but makes interop easier in this simplified example */
     mbedtls_ssl_conf_authmode( &conf, MBEDTLS_SSL_VERIFY_OPTIONAL );
 
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/programs/ssl/ssl_pthread_server.c b/programs/ssl/ssl_pthread_server.c
index fd6ca26..80bb049 100644
--- a/programs/ssl/ssl_pthread_server.c
+++ b/programs/ssl/ssl_pthread_server.c
@@ -325,6 +325,20 @@
     return( 0 );
 }
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( void )
 {
     int ret;
@@ -439,7 +453,12 @@
         goto exit;
     }
 
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_mutexed_debug, stdout );
 #endif
diff --git a/programs/ssl/ssl_server.c b/programs/ssl/ssl_server.c
index bf502a5..e13af91 100644
--- a/programs/ssl/ssl_server.c
+++ b/programs/ssl/ssl_server.c
@@ -97,6 +97,20 @@
 }
 #endif /* MBEDTLS_DEBUG_C */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( void )
 {
     int ret, len;
@@ -212,7 +226,12 @@
         goto exit;
     }
 
+#if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
     mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 3fa2b15..e0e4337 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -1505,6 +1505,20 @@
 }
 #endif /* MBEDTLS_SSL_DTLS_CONNECTION_ID */
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( int argc, char *argv[] )
 {
     int ret = 0, len, written, frags, exchanges_left;
@@ -2754,7 +2768,7 @@
 #if !defined(MBEDTLS_SSL_CONF_RNG)
     mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
 #else
-    mbedtls_ssl_conf_rng_ctx( &conf, &ctr_drbg );
+    rng_ctx_global = &ctr_drbg;
 #endif
 
 #if defined(MBEDTLS_DEBUG_C)
diff --git a/programs/x509/cert_app.c b/programs/x509/cert_app.c
index cdd77f2..9b40fa7 100644
--- a/programs/x509/cert_app.c
+++ b/programs/x509/cert_app.c
@@ -149,6 +149,20 @@
     return( 0 );
 }
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len );
+
+mbedtls_ctr_drbg_context *rng_ctx_global = NULL;
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    /* We expect the NULL parameter here. */
+    if( ctx != NULL )
+        return( -1 );
+
+    return( mbedtls_ctr_drbg_random( rng_ctx_global, dst, len ) );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 int main( int argc, char *argv[] )
 {
     int ret = 1;
@@ -424,7 +438,12 @@
         else
             mbedtls_ssl_conf_authmode( &conf, MBEDTLS_SSL_VERIFY_NONE );
 
-        mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#if !defined(MBEDTLS_SSL_CONF_RNG)
+    mbedtls_ssl_conf_rng( &conf, mbedtls_ctr_drbg_random, &ctr_drbg );
+#else
+    rng_ctx_global = &ctr_drbg;
+#endif
+
 #if defined(MBEDTLS_DEBUG_C)
         mbedtls_ssl_conf_dbg( &conf, my_debug, stdout );
 #endif
diff --git a/tests/suites/helpers.function b/tests/suites/helpers.function
index ccd4d42..f5b61f3 100644
--- a/tests/suites/helpers.function
+++ b/tests/suites/helpers.function
@@ -356,6 +356,18 @@
 }
 #endif
 
+#if defined(MBEDTLS_SSL_CONF_RNG)
+int rng_wrap( void *ctx, unsigned char *dst, size_t len )
+{
+    ((void) ctx);
+    ((void) dst);
+    ((void) len);
+    /* We don't expect test suites to use SSL functionality
+     * that would trigger the hardcoded RNG. */
+    return( -1 );
+}
+#endif /* MBEDTLS_SSL_CONF_RNG */
+
 #if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
 static int redirect_output( FILE** out_stream, const char* path )
 {