Implement RSA-PSK key exchange
diff --git a/library/ssl_cli.c b/library/ssl_cli.c
index 77a18ed..4da7149 100644
--- a/library/ssl_cli.c
+++ b/library/ssl_cli.c
@@ -1183,6 +1183,72 @@
           POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED
           POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
 
+#if defined(POLARSSL_KEY_EXCHANGE_RSA_ENABLED) ||                           \
+    defined(POLARSSL_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( ssl_context *ssl,
+                                    size_t offset, size_t *olen,
+                                    size_t pms_offset )
+{
+    int ret;
+    size_t len_bytes = ssl->minor_ver == SSL_MINOR_VERSION_0 ? 0 : 2;
+    unsigned char *p = ssl->handshake->premaster + pms_offset;
+
+    /*
+     * Generate (part of) the pre-master as
+     *  struct {
+     *      ProtocolVersion client_version;
+     *      opaque random[46];
+     *  } PreMasterSecret;
+     */
+    p[0] = (unsigned char) ssl->max_major_ver;
+    p[1] = (unsigned char) ssl->max_minor_ver;
+
+    if( ( ret = ssl->f_rng( ssl->p_rng, p + 2, 46 ) ) != 0 )
+    {
+        SSL_DEBUG_RET( 1, "f_rng", ret );
+        return( ret );
+    }
+
+    ssl->handshake->pmslen = 48;
+
+    /*
+     * Now write it out, encrypted
+     */
+    if( ! pk_can_do( &ssl->session_negotiate->peer_cert->pk,
+                POLARSSL_PK_RSA ) )
+    {
+        SSL_DEBUG_MSG( 1, ( "certificate key type mismatch" ) );
+        return( POLARSSL_ERR_SSL_PK_TYPE_MISMATCH );
+    }
+
+    if( ( ret = pk_encrypt( &ssl->session_negotiate->peer_cert->pk,
+                            p, ssl->handshake->pmslen,
+                            ssl->out_msg + offset + len_bytes, olen,
+                            SSL_MAX_CONTENT_LEN - offset - len_bytes,
+                            ssl->f_rng, ssl->p_rng ) ) != 0 )
+    {
+        SSL_DEBUG_RET( 1, "rsa_pkcs1_encrypt", ret );
+        return( ret );
+    }
+
+#if defined(POLARSSL_SSL_PROTO_TLS1) || defined(POLARSSL_SSL_PROTO_TLS1_1) || \
+    defined(POLARSSL_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      );
+        *olen += 2;
+    }
+#endif
+
+    return( 0 );
+}
+#endif /* POLARSSL_KEY_EXCHANGE_RSA_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED */
+
 #if defined(POLARSSL_SSL_PROTO_TLS1_2)
 #if defined(POLARSSL_KEY_EXCHANGE_DHE_RSA_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_ECDHE_RSA_ENABLED) ||                     \
@@ -1258,6 +1324,7 @@
         ciphersuite_info->key_exchange != POLARSSL_KEY_EXCHANGE_ECDHE_RSA &&
         ciphersuite_info->key_exchange != POLARSSL_KEY_EXCHANGE_ECDHE_ECDSA &&
         ciphersuite_info->key_exchange != POLARSSL_KEY_EXCHANGE_PSK &&
+        ciphersuite_info->key_exchange != POLARSSL_KEY_EXCHANGE_RSA_PSK &&
         ciphersuite_info->key_exchange != POLARSSL_KEY_EXCHANGE_DHE_PSK &&
         ciphersuite_info->key_exchange != POLARSSL_KEY_EXCHANGE_ECDHE_PSK )
     {
@@ -1324,8 +1391,10 @@
     else
 #endif /* POLARSSL_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
           POLARSSL_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED */
-#if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED)
-    if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_PSK )
+#if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED) ||                           \
+    defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED)
+    if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_PSK ||
+        ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_RSA_PSK )
     {
         unsigned char *p = ssl->in_msg + 4;
         unsigned char *end = ssl->in_msg + ssl->in_hslen;
@@ -1337,7 +1406,8 @@
         }
     }
     else
-#endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED */
+#endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED */
 #if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
     if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_DHE_PSK )
     {
@@ -1803,9 +1873,11 @@
 #endif /* POLARSSL_KEY_EXCHANGE_ECDHE_RSA_ENABLED ||
           POLARSSL_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED */
 #if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED) ||                           \
+    defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED)
     if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_PSK ||
+        ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_RSA_PSK ||
         ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_DHE_PSK ||
         ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_ECDHE_PSK )
     {
@@ -1830,6 +1902,14 @@
         }
         else
 #endif
+#if defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED)
+        if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_RSA_PSK )
+        {
+            if( ( ret = ssl_write_encrypted_pms( ssl, i, &n, 2 ) ) != 0 )
+                return( ret );
+        }
+        else
+#endif
 #if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
         if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_DHE_PSK )
         {
@@ -1885,56 +1965,21 @@
     }
     else
 #endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED ||
           POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED ||
           POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
 #if defined(POLARSSL_KEY_EXCHANGE_RSA_ENABLED)
     if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_RSA )
     {
-        /*
-         * RSA key exchange -- send rsa_public(pkcs1 v1.5(premaster))
-         */
-        ssl->handshake->premaster[0] = (unsigned char) ssl->max_major_ver;
-        ssl->handshake->premaster[1] = (unsigned char) ssl->max_minor_ver;
-        ssl->handshake->pmslen = 48;
-
-        ret = ssl->f_rng( ssl->p_rng, ssl->handshake->premaster + 2,
-                          ssl->handshake->pmslen - 2 );
-        if( ret != 0 )
+        i = 4;
+        if( ( ret = ssl_write_encrypted_pms( ssl, i, &n, 0 ) ) != 0 )
             return( ret );
-
-        if( ! pk_can_do( &ssl->session_negotiate->peer_cert->pk,
-                         POLARSSL_PK_RSA ) )
-        {
-            SSL_DEBUG_MSG( 1, ( "certificate key type mismatch" ) );
-            return( POLARSSL_ERR_SSL_PK_TYPE_MISMATCH );
-        }
-
-        i = ssl->minor_ver == SSL_MINOR_VERSION_0 ? 4 : 6;
-
-        ret = pk_encrypt( &ssl->session_negotiate->peer_cert->pk,
-                ssl->handshake->premaster, ssl->handshake->pmslen,
-                ssl->out_msg + i, &n, SSL_BUFFER_LEN,
-                ssl->f_rng, ssl->p_rng );
-        if( ret != 0 )
-        {
-            SSL_DEBUG_RET( 1, "rsa_pkcs1_encrypt", ret );
-            return( ret );
-        }
-
-#if defined(POLARSSL_SSL_PROTO_TLS1) || defined(POLARSSL_SSL_PROTO_TLS1_1) || \
-    defined(POLARSSL_SSL_PROTO_TLS1_2)
-        if( ssl->minor_ver != SSL_MINOR_VERSION_0 )
-        {
-            ssl->out_msg[4] = (unsigned char)( n >> 8 );
-            ssl->out_msg[5] = (unsigned char)( n      );
-        }
-#endif
-
     }
     else
 #endif /* POLARSSL_KEY_EXCHANGE_RSA_ENABLED */
     {
         ((void) ciphersuite_info);
+        SSL_DEBUG_MSG( 1, ( "should never happen" ) );
         return( POLARSSL_ERR_SSL_FEATURE_UNAVAILABLE );
     }
 
diff --git a/library/ssl_srv.c b/library/ssl_srv.c
index 336add2..6db92d6 100644
--- a/library/ssl_srv.c
+++ b/library/ssl_srv.c
@@ -2288,10 +2288,14 @@
           POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED */
 
 #if defined(POLARSSL_KEY_EXCHANGE_RSA_ENABLED)
-static int ssl_parse_encrypted_pms_secret( ssl_context *ssl )
+static int ssl_parse_encrypted_pms( ssl_context *ssl,
+                                    const unsigned char *p,
+                                    const unsigned char *end,
+                                    size_t pms_offset )
 {
-    int ret = POLARSSL_ERR_SSL_FEATURE_UNAVAILABLE;
-    size_t i, n = 0;
+    int ret;
+    size_t len = pk_get_len( ssl_own_key( ssl ) );
+    unsigned char *pms = ssl->handshake->premaster + pms_offset;
 
     if( ! pk_can_do( ssl_own_key( ssl ), POLARSSL_PK_RSA ) )
     {
@@ -2302,17 +2306,12 @@
     /*
      * Decrypt the premaster using own private RSA key
      */
-    i = 4;
-    n = pk_get_len( ssl_own_key( ssl ) );
-    ssl->handshake->pmslen = 48;
-
 #if defined(POLARSSL_SSL_PROTO_TLS1) || defined(POLARSSL_SSL_PROTO_TLS1_1) || \
     defined(POLARSSL_SSL_PROTO_TLS1_2)
     if( ssl->minor_ver != SSL_MINOR_VERSION_0 )
     {
-        i += 2;
-        if( ssl->in_msg[4] != ( ( n >> 8 ) & 0xFF ) ||
-            ssl->in_msg[5] != ( ( n      ) & 0xFF ) )
+        if( *p++ != ( ( len >> 8 ) & 0xFF ) ||
+            *p++ != ( ( len      ) & 0xFF ) )
         {
             SSL_DEBUG_MSG( 1, ( "bad client key exchange message" ) );
             return( POLARSSL_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
@@ -2320,21 +2319,20 @@
     }
 #endif
 
-    if( ssl->in_hslen != i + n )
+    if( p + len != end )
     {
         SSL_DEBUG_MSG( 1, ( "bad client key exchange message" ) );
         return( POLARSSL_ERR_SSL_BAD_HS_CLIENT_KEY_EXCHANGE );
     }
 
-    ret = pk_decrypt( ssl_own_key( ssl ),
-                      ssl->in_msg + i, n,
-                      ssl->handshake->premaster, &ssl->handshake->pmslen,
+    ret = pk_decrypt( ssl_own_key( ssl ), p, len,
+                      pms, &ssl->handshake->pmslen,
                       sizeof(ssl->handshake->premaster),
                       ssl->f_rng, ssl->p_rng );
 
     if( ret != 0 || ssl->handshake->pmslen != 48 ||
-        ssl->handshake->premaster[0] != ssl->handshake->max_major_ver ||
-        ssl->handshake->premaster[1] != ssl->handshake->max_minor_ver )
+        pms[0] != ssl->handshake->max_major_ver ||
+        pms[1] != ssl->handshake->max_minor_ver )
     {
         SSL_DEBUG_MSG( 1, ( "bad client key exchange message" ) );
 
@@ -2346,8 +2344,7 @@
          */
         ssl->handshake->pmslen = 48;
 
-        ret = ssl->f_rng( ssl->p_rng, ssl->handshake->premaster,
-                          ssl->handshake->pmslen );
+        ret = ssl->f_rng( ssl->p_rng, pms, ssl->handshake->pmslen );
         if( ret != 0 )
             return( ret );
     }
@@ -2541,6 +2538,33 @@
     }
     else
 #endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED */
+#if defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED)
+    if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_RSA_PSK )
+    {
+        unsigned char *p = ssl->in_msg + 4;
+        unsigned char *end = ssl->in_msg + ssl->in_msglen;
+
+        if( ( ret = ssl_parse_client_psk_identity( ssl, &p, end ) ) != 0 )
+        {
+            SSL_DEBUG_RET( 1, ( "ssl_parse_client_psk_identity" ), ret );
+            return( ret );
+        }
+
+        if( ( ret = ssl_parse_encrypted_pms( ssl, p, end, 2 ) ) != 0 )
+        {
+            SSL_DEBUG_RET( 1, ( "ssl_parse_encrypted_pms" ), ret );
+            return( ret );
+        }
+
+        if( ( ret = ssl_psk_derive_premaster( ssl,
+                        ciphersuite_info->key_exchange ) ) != 0 )
+        {
+            SSL_DEBUG_RET( 1, "ssl_psk_derive_premaster", ret );
+            return( ret );
+        }
+    }
+    else
+#endif /* POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED */
 #if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
     if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_DHE_PSK )
     {
@@ -2600,7 +2624,10 @@
 #if defined(POLARSSL_KEY_EXCHANGE_RSA_ENABLED)
     if( ciphersuite_info->key_exchange == POLARSSL_KEY_EXCHANGE_RSA )
     {
-        if( ( ret = ssl_parse_encrypted_pms_secret( ssl ) ) != 0 )
+        if( ( ret = ssl_parse_encrypted_pms( ssl,
+                                             ssl->in_msg + 4,
+                                             ssl->in_msg + ssl->in_msglen,
+                                             0 ) ) != 0 )
         {
             SSL_DEBUG_RET( 1, ( "ssl_parse_parse_ecrypted_pms_secret" ), ret );
             return( ret );
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 2be2016..e254f3a 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -826,6 +826,7 @@
 #endif /* POLARSSL_SSL_PROTO_TLS1_2 */
 
 #if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED) ||                           \
+    defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED)
 int ssl_psk_derive_premaster( ssl_context *ssl, key_exchange_type_t key_ex )
@@ -852,6 +853,19 @@
     }
     else
 #endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED */
+#if defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED)
+    if( key_ex == POLARSSL_KEY_EXCHANGE_RSA_PSK )
+    {
+        /*
+         * other_secret already set by the ClientKeyExchange message,
+         * and is 48 bytes long
+         */
+        *p++ = 0;
+        *p++ = 48;
+        p += 48;
+    }
+    else
+#endif /* POLARSSL_KEY_EXCHANGE_RSA_PKS_ENABLED */
 #if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
     if( key_ex == POLARSSL_KEY_EXCHANGE_DHE_PSK )
     {
@@ -913,6 +927,7 @@
     return( 0 );
 }
 #endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED ||
           POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED ||
           POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
 
@@ -3661,6 +3676,7 @@
 #endif /* POLARSSL_X509_CRT_PARSE_C */
 
 #if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED) ||                           \
+    defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED) ||                       \
     defined(POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED)
 int ssl_set_psk( ssl_context *ssl, const unsigned char *psk, size_t psk_len,
@@ -3699,6 +3715,7 @@
     ssl->p_psk = p_psk;
 }
 #endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED ||
           POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED ||
           POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED */