Merge pull request #2931 from piotr-now/memory-info

Avoid allocating 0-length buffers for PSK. Add memory debug information to ssl_client2.
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 69a46b9..e450207 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -9171,8 +9171,13 @@
     ssl_conf_remove_psk( conf );
 
     /* Check and set raw PSK */
-    if( psk == NULL || psk_len > MBEDTLS_PSK_MAX_LEN )
+    if( psk == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
+    if( psk_len == 0 )
+        return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
+    if( psk_len > MBEDTLS_PSK_MAX_LEN )
+        return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
+
     if( ( conf->psk = mbedtls_calloc( 1, psk_len ) ) == NULL )
         return( MBEDTLS_ERR_SSL_ALLOC_FAILED );
     conf->psk_len = psk_len;
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 558fa28..8f0d3b5 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -54,6 +54,10 @@
 }
 #else
 
+#if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
+#include "mbedtls/memory_buffer_alloc.h"
+#endif
+
 #include "mbedtls/net_sockets.h"
 #include "mbedtls/ssl.h"
 #include "mbedtls/entropy.h"
@@ -73,6 +77,10 @@
 #include <stdlib.h>
 #include <string.h>
 
+/* Size of memory to be allocated for the heap, when using the library's memory
+ * management and MBEDTLS_MEMORY_BUFFER_ALLOC_C is enabled. */
+#define MEMORY_HEAP_SIZE      120000
+
 #define MAX_REQUEST_SIZE      20000
 #define MAX_REQUEST_SIZE_STR "20000"
 
@@ -194,7 +202,8 @@
 
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
 #define USAGE_PSK_RAW                                               \
-    "    psk=%%s              default: \"\" (in hex, without 0x)\n" \
+    "    psk=%%s              default: \"\" (disabled)\n"     \
+    "                          The PSK values are in hex, without 0x.\n" \
     "    psk_identity=%%s     default: \"Client_identity\"\n"
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 #define USAGE_PSK_SLOT                          \
@@ -1129,6 +1138,11 @@
 #if defined(MBEDTLS_SSL_ALPN)
     const char *alpn_list[ALPN_LIST_SIZE];
 #endif
+
+#if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
+    unsigned char alloc_buf[MEMORY_HEAP_SIZE];
+#endif
+
 #if defined(MBEDTLS_ECP_C)
     mbedtls_ecp_group_id curve_list[CURVE_LIST_SIZE];
     const mbedtls_ecp_curve_info *curve_cur;
@@ -1178,6 +1192,10 @@
     eap_tls_keys eap_tls_keying;
 #endif
 
+#if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
+    mbedtls_memory_buffer_alloc_init( alloc_buf, sizeof(alloc_buf) );
+#endif
+
     /*
      * Make sure memory references are valid.
      */
@@ -2374,13 +2392,16 @@
     }
     else
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
-    if( ( ret = mbedtls_ssl_conf_psk( &conf, psk, psk_len,
-                             (const unsigned char *) opt.psk_identity,
-                             strlen( opt.psk_identity ) ) ) != 0 )
+    if( psk_len > 0 )
     {
-        mbedtls_printf( " failed\n  ! mbedtls_ssl_conf_psk returned %d\n\n",
-                        ret );
-        goto exit;
+        ret = mbedtls_ssl_conf_psk( &conf, psk, psk_len,
+                             (const unsigned char *) opt.psk_identity,
+                             strlen( opt.psk_identity ) );
+        if( ret != 0 )
+        {
+            mbedtls_printf( " failed\n  ! mbedtls_ssl_conf_psk returned %d\n\n", ret );
+            goto exit;
+        }
     }
 #endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED */
 
@@ -3287,6 +3308,13 @@
 #endif /* MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED &&
           MBEDTLS_USE_PSA_CRYPTO */
 
+#if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
+#if defined(MBEDTLS_MEMORY_DEBUG)
+    mbedtls_memory_buffer_alloc_status();
+#endif
+    mbedtls_memory_buffer_alloc_free();
+#endif
+
 #if defined(_WIN32)
     mbedtls_printf( "  + Press Enter to exit this program.\n" );
     fflush( stdout ); getchar();
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 27ffbb9..27f2312 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -259,7 +259,8 @@
 
 #if defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
 #define USAGE_PSK_RAW                                               \
-    "    psk=%%s              default: \"\" (in hex, without 0x)\n"     \
+    "    psk=%%s              default: \"\" (disabled)\n"     \
+    "                          The PSK values are in hex, without 0x.\n" \
     "    psk_list=%%s         default: \"\"\n"                          \
     "                          A list of (PSK identity, PSK value) pairs.\n" \
     "                          The PSK values are in hex, without 0x.\n" \
@@ -3364,12 +3365,16 @@
         }
         else
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
-        if( ( ret = mbedtls_ssl_conf_psk( &conf, psk, psk_len,
-                                     (const unsigned char *) opt.psk_identity,
-                                     strlen( opt.psk_identity ) ) ) != 0 )
+        if( psk_len > 0 )
         {
-            mbedtls_printf( "  failed\n  mbedtls_ssl_conf_psk returned -0x%04X\n\n", - ret );
-            goto exit;
+            ret = mbedtls_ssl_conf_psk( &conf, psk, psk_len,
+                                     (const unsigned char *) opt.psk_identity,
+                                     strlen( opt.psk_identity ) );
+            if( ret != 0 )
+            {
+                mbedtls_printf( "  failed\n  mbedtls_ssl_conf_psk returned -0x%04X\n\n", - ret );
+                goto exit;
+            }
         }
     }