Factor PSK pms computation to ssl_tls.c
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index d3b4ce3..e4ae682 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -825,6 +825,96 @@
 #endif /* POLARSSL_SHA512_C */
 #endif /* POLARSSL_SSL_PROTO_TLS1_2 */
 
+#if defined(POLARSSL_KEY_EXCHANGE_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 )
+{
+    int ret;
+    unsigned char *p = ssl->handshake->premaster;
+    unsigned char *end = p + sizeof( ssl->handshake->premaster );
+
+    /*
+     * PMS = struct {
+     *     opaque other_secret<0..2^16-1>;
+     *     opaque psk<0..2^16-1>;
+     * };
+     * with "other_secret" depending on the particular key exchange
+     */
+#if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED)
+    if( key_ex == POLARSSL_KEY_EXCHANGE_PSK )
+    {
+        if( end - p < 2 + (int) ssl->psk_len )
+            return( POLARSSL_ERR_SSL_BAD_INPUT_DATA );
+
+        *(p++) = (unsigned char)( ssl->psk_len >> 8 );
+        *(p++) = (unsigned char)( ssl->psk_len      );
+        p += ssl->psk_len;
+    }
+    else
+#endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED */
+#if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
+    if( key_ex == POLARSSL_KEY_EXCHANGE_DHE_PSK )
+    {
+        size_t len = ssl->handshake->dhm_ctx.len;
+
+        if( end - p < 2 + (int) len )
+            return( POLARSSL_ERR_SSL_BAD_INPUT_DATA );
+
+        *(p++) = (unsigned char)( len >> 8 );
+        *(p++) = (unsigned char)( len );
+        if( ( ret = dhm_calc_secret( &ssl->handshake->dhm_ctx,
+                                      p, &len, ssl->f_rng, ssl->p_rng ) ) != 0 )
+        {
+            SSL_DEBUG_RET( 1, "dhm_calc_secret", ret );
+            return( ret );
+        }
+        p += len;
+
+        SSL_DEBUG_MPI( 3, "DHM: K ", &ssl->handshake->dhm_ctx.K  );
+    }
+    else
+#endif /* POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED */
+#if defined(POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED)
+    if( key_ex == POLARSSL_KEY_EXCHANGE_ECDHE_PSK )
+    {
+        size_t zlen;
+
+        if( ( ret = ecdh_calc_secret( &ssl->handshake->ecdh_ctx, &zlen,
+                                       p + 2, end - (p + 2),
+                                       ssl->f_rng, ssl->p_rng ) ) != 0 )
+        {
+            SSL_DEBUG_RET( 1, "ecdh_calc_secret", ret );
+            return( ret );
+        }
+
+        *(p++) = (unsigned char)( zlen >> 8 );
+        *(p++) = (unsigned char)( zlen      );
+        p += zlen;
+
+        SSL_DEBUG_MPI( 3, "ECDH: z", &ssl->handshake->ecdh_ctx.z );
+    }
+    else
+#endif /* POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
+    {
+        SSL_DEBUG_MSG( 1, ( "should never happen" ) );
+        return( POLARSSL_ERR_SSL_FEATURE_UNAVAILABLE );
+    }
+
+    /* opaque psk<0..2^16-1>; */
+    *(p++) = (unsigned char)( ssl->psk_len >> 8 );
+    *(p++) = (unsigned char)( ssl->psk_len      );
+    memcpy( p, ssl->psk, ssl->psk_len );
+    p += ssl->psk_len;
+
+    ssl->handshake->pmslen = p - ssl->handshake->premaster;
+
+    return( 0 );
+}
+#endif /* POLARSSL_KEY_EXCHANGE_PSK_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED ||
+          POLARSSL_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
+
 #if defined(POLARSSL_SSL_PROTO_SSL3)
 /*
  * SSLv3.0 MAC functions