Make PSK_LEN configurable and adjust PMS size
diff --git a/include/polarssl/config.h b/include/polarssl/config.h
index e83518a..d43365f 100644
--- a/include/polarssl/config.h
+++ b/include/polarssl/config.h
@@ -2154,6 +2154,7 @@
 /* SSL options */
 //#define SSL_MAX_CONTENT_LEN             16384 /**< Size of the input / output buffer */
 //#define SSL_DEFAULT_TICKET_LIFETIME     86400 /**< Lifetime of session tickets (if enabled) */
+//#define POLARSSL_PSK_MAX_LEN               32 /**< Max size of TLS pre-shared keys, in bytes (default 256 bits) */
 
 /**
  * Complete list of ciphersuites to use, in order of preference.
diff --git a/include/polarssl/ssl.h b/include/polarssl/ssl.h
index 9db0fc6..f1de499 100644
--- a/include/polarssl/ssl.h
+++ b/include/polarssl/ssl.h
@@ -34,6 +34,7 @@
 #endif
 #include "net.h"
 #include "bignum.h"
+#include "ecp.h"
 
 #include "ssl_ciphersuites.h"
 
@@ -409,12 +410,43 @@
 /*
  * Size defines
  */
-#if !defined(POLARSSL_MPI_MAX_SIZE)
-#define POLARSSL_PREMASTER_SIZE             512
-#else
-#define POLARSSL_PREMASTER_SIZE             POLARSSL_MPI_MAX_SIZE
+#if !defined(POLARSSL_PSK_MAX_LEN)
+#define POLARSSL_PSK_MAX_LEN            32 /* 256 bits */
 #endif
 
+/* Dummy type used only for its size */
+union _ssl_premaster_secret
+{
+#if defined(POLARSSL_KEY_EXCHANGE_RSA_ENABLED)
+    unsigned char _pms_rsa[48];                         /* RFC 5246 8.1.1 */
+#endif
+#if defined(POLARSSL_KEY_EXCHANGE_DHE_RSA_ENABLED)
+    unsigned char _pms_dhm[POLARSSL_MPI_MAX_SIZE];      /* RFC 5246 8.1.2 */
+#endif
+#if defined(POLARSSL_KEY_EXCHANGE_ECDHE_RSA_ENABLED)    || \
+    defined(POLARSSL_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED)  || \
+    defined(POLARSSL_KEY_EXCHANGE_ECDH_RSA_ENABLED)     || \
+    defined(POLARSSL_KEY_EXCHANGE_ECDH_ECDSA_ENABLED)
+    unsigned char _pms_ecdh[POLARSSL_ECP_MAX_BYTES];    /* RFC 4492 5.10 */
+#endif
+#if defined(POLARSSL_KEY_EXCHANGE_PSK_ENABLED)
+    unsigned char _pms_psk[4 + 2 * POLARSSL_PSK_MAX_LEN];       /* RFC 4279 2 */
+#endif
+#if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
+    unsigned char _pms_dhe_psk[4 + POLARSSL_MPI_MAX_SIZE
+                                 + POLARSSL_PSK_MAX_LEN];       /* RFC 4279 3 */
+#endif
+#if defined(POLARSSL_KEY_EXCHANGE_RSA_PSK_ENABLED)
+    unsigned char _pms_rsa_psk[52 + POLARSSL_PSK_MAX_LEN];      /* RFC 4279 4 */
+#endif
+#if defined(POLARSSL_KEY_EXCHANGE_DHE_PSK_ENABLED)
+    unsigned char _pms_ecdhe_psk[4 + POLARSSL_ECP_MAX_BYTES
+                                   + POLARSSL_PSK_MAX_LEN];     /* RFC 5489 2 */
+#endif
+};
+
+#define POLARSSL_PREMASTER_SIZE     sizeof( union _ssl_premaster_secret )
+
 #ifdef __cplusplus
 extern "C" {
 #endif
diff --git a/include/polarssl/ssl_ciphersuites.h b/include/polarssl/ssl_ciphersuites.h
index 5ecd5fe..c4f1ffe 100644
--- a/include/polarssl/ssl_ciphersuites.h
+++ b/include/polarssl/ssl_ciphersuites.h
@@ -233,6 +233,7 @@
 #define TLS_ECDHE_ECDSA_WITH_AES_128_CCM_8      0xC0AE  /**< TLS 1.2 */
 #define TLS_ECDHE_ECDSA_WITH_AES_256_CCM_8      0xC0AF  /**< TLS 1.2 */
 
+/* Reminder: update _ssl_premaster_secret when adding a new key exchange */
 typedef enum {
     POLARSSL_KEY_EXCHANGE_NONE = 0,
     POLARSSL_KEY_EXCHANGE_RSA,
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index a1428dc..373a1f5 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -3740,12 +3740,7 @@
     if( psk == NULL || psk_identity == NULL )
         return( POLARSSL_ERR_SSL_BAD_INPUT_DATA );
 
-    /*
-     * The length will be check later anyway, but in case it is obviously
-     * too large, better abort now. The PMS is as follows:
-     * other_len (2 bytes) + other + psk_len (2 bytes) + psk
-     */
-    if( psk_len + 4 > POLARSSL_PREMASTER_SIZE )
+    if( psk_len > POLARSSL_PSK_MAX_LEN )
         return( POLARSSL_ERR_SSL_BAD_INPUT_DATA );
 
     if( ssl->psk != NULL )
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 3af54f9..4682ee5 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -325,7 +325,7 @@
     int ret = 0, len, server_fd, i, written, frags;
     unsigned char buf[SSL_MAX_CONTENT_LEN + 1];
 #if defined(POLARSSL_KEY_EXCHANGE__SOME__PSK_ENABLED)
-    unsigned char psk[256];
+    unsigned char psk[POLARSSL_PSK_MAX_LEN];
     size_t psk_len = 0;
 #endif
 #if defined(POLARSSL_SSL_ALPN)
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index d5f01bc..935a7a5 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -126,8 +126,6 @@
     "<h2>PolarSSL Test Server</h2>\r\n" \
     "<p>Successful connection using: %s</p>\r\n" // LONG_RESPONSE
 
-#define MAX_PSK_LEN     256
-
 /*
  * Size of the basic I/O buffer. Able to hold our default response.
  *
@@ -460,7 +458,7 @@
     size_t j;
 
     *olen = strlen( input );
-    if( *olen % 2 != 0 || *olen / 2 > MAX_PSK_LEN )
+    if( *olen % 2 != 0 || *olen / 2 > POLARSSL_PSK_MAX_LEN )
         return( -1 );
     *olen /= 2;
 
@@ -484,7 +482,7 @@
 {
     const char *name;
     size_t key_len;
-    unsigned char key[MAX_PSK_LEN];
+    unsigned char key[POLARSSL_PSK_MAX_LEN];
     psk_entry *next;
 };
 
@@ -573,7 +571,7 @@
     int version_suites[4][2];
     unsigned char buf[IO_BUF_LEN];
 #if defined(POLARSSL_KEY_EXCHANGE__SOME__PSK_ENABLED)
-    unsigned char psk[MAX_PSK_LEN];
+    unsigned char psk[POLARSSL_PSK_MAX_LEN];
     size_t psk_len = 0;
     psk_entry *psk_info = NULL;
 #endif