From b33cc7688ee7857a49b53c45f09450246b7e4e7e Mon Sep 17 00:00:00 2001
From: Darryl Green <darryl.green@arm.com>
Date: Thu, 28 Nov 2019 14:29:44 +0000
Subject: [PATCH] Add I/O buffer length fields to mbedtls_ssl_context

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
Signed-off-by: Darryl Green <darryl.green@arm.com>
---
 include/mbedtls/config.h     |  7 +++++
 include/mbedtls/ssl.h        |  6 ++++
 library/ssl_msg.c            | 55 ++++++++++++++++++++++++++----------
 library/ssl_tls.c            | 43 ++++++++++++++++++++++------
 library/version_features.c   |  3 ++
 programs/test/query_config.c |  8 ++++++
 6 files changed, 99 insertions(+), 23 deletions(-)

diff --git a/include/mbedtls/config.h b/include/mbedtls/config.h
index 70a97405a..ac03884a7 100644
--- a/include/mbedtls/config.h
+++ b/include/mbedtls/config.h
@@ -1779,6 +1779,13 @@
  */
 #define MBEDTLS_SSL_TRUNCATED_HMAC
 
+/**
+ * \def MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH
+ *
+ * Enable modifying the maximum I/O buffer size.
+ */
+//#define MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH
+
 /**
  * \def MBEDTLS_SSL_TRUNCATED_HMAC_COMPAT
  *
diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 1666f8c86..b3c5288c0 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -1215,6 +1215,9 @@ struct mbedtls_ssl_context
     int in_msgtype;             /*!< record header: message type      */
     size_t in_msglen;           /*!< record header: message length    */
     size_t in_left;             /*!< amount of data read so far       */
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t in_buf_len;          /*!< length of input buffer           */
+#endif
 #if defined(MBEDTLS_SSL_PROTO_DTLS)
     uint16_t in_epoch;          /*!< DTLS epoch for incoming records  */
     size_t next_record_offset;  /*!< offset of the next record in datagram
@@ -1254,6 +1257,9 @@ struct mbedtls_ssl_context
     int out_msgtype;            /*!< record header: message type      */
     size_t out_msglen;          /*!< record header: message length    */
     size_t out_left;            /*!< amount of data not yet written   */
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t out_buf_len;         /*!< length of output buffer          */
+#endif
 
     unsigned char cur_out_ctr[8]; /*!<  Outgoing record sequence  number. */
 
diff --git a/library/ssl_msg.c b/library/ssl_msg.c
index 9c2d61509..df0d0fcad 100644
--- a/library/ssl_msg.c
+++ b/library/ssl_msg.c
@@ -179,11 +179,16 @@ static int ssl_next_record_is_in_datagram( mbedtls_ssl_context *ssl );
 static size_t ssl_get_maximum_datagram_size( mbedtls_ssl_context const *ssl )
 {
     size_t mtu = mbedtls_ssl_get_current_mtu( ssl );
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t out_buf_len = ssl->out_buf_len;
+#else
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN;
+#endif
 
-    if( mtu != 0 && mtu < MBEDTLS_SSL_OUT_BUFFER_LEN )
+    if( mtu != 0 && mtu < out_buf_len )
         return( mtu );
 
-    return( MBEDTLS_SSL_OUT_BUFFER_LEN );
+    return( out_buf_len );
 }
 
 static int ssl_get_remaining_space_in_datagram( mbedtls_ssl_context const *ssl )
@@ -1574,6 +1579,11 @@ static int ssl_compress_buf( mbedtls_ssl_context *ssl )
     ptrdiff_t bytes_written = ssl->out_msg - ssl->out_buf;
     size_t len_pre = ssl->out_msglen;
     unsigned char *msg_pre = ssl->compress_buf;
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t out_buf_len = ssl->out_buf_len;
+#else
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN;
+#endif
 
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> compress buf" ) );
 
@@ -1591,7 +1601,7 @@ static int ssl_compress_buf( mbedtls_ssl_context *ssl )
     ssl->transform_out->ctx_deflate.next_in = msg_pre;
     ssl->transform_out->ctx_deflate.avail_in = len_pre;
     ssl->transform_out->ctx_deflate.next_out = msg_post;
-    ssl->transform_out->ctx_deflate.avail_out = MBEDTLS_SSL_OUT_BUFFER_LEN - bytes_written;
+    ssl->transform_out->ctx_deflate.avail_out = out_buf_len - bytes_written;
 
     ret = deflate( &ssl->transform_out->ctx_deflate, Z_SYNC_FLUSH );
     if( ret != Z_OK )
@@ -1600,7 +1610,7 @@ static int ssl_compress_buf( mbedtls_ssl_context *ssl )
         return( MBEDTLS_ERR_SSL_COMPRESSION_FAILED );
     }
 
-    ssl->out_msglen = MBEDTLS_SSL_OUT_BUFFER_LEN -
+    ssl->out_msglen = out_buf_len -
                       ssl->transform_out->ctx_deflate.avail_out - bytes_written;
 
     MBEDTLS_SSL_DEBUG_MSG( 3, ( "after compression: msglen = %d, ",
@@ -1621,6 +1631,11 @@ static int ssl_decompress_buf( mbedtls_ssl_context *ssl )
     ptrdiff_t header_bytes = ssl->in_msg - ssl->in_buf;
     size_t len_pre = ssl->in_msglen;
     unsigned char *msg_pre = ssl->compress_buf;
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t in_buf_len = ssl->in_buf_len;
+#else
+    size_t in_buf_len = MBEDTLS_SSL_IN_BUFFER_LEN;
+#endif
 
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> decompress buf" ) );
 
@@ -1638,8 +1653,7 @@ static int ssl_decompress_buf( mbedtls_ssl_context *ssl )
     ssl->transform_in->ctx_inflate.next_in = msg_pre;
     ssl->transform_in->ctx_inflate.avail_in = len_pre;
     ssl->transform_in->ctx_inflate.next_out = msg_post;
-    ssl->transform_in->ctx_inflate.avail_out = MBEDTLS_SSL_IN_BUFFER_LEN -
-                                               header_bytes;
+    ssl->transform_in->ctx_inflate.avail_out = in_buf_len - header_bytes;
 
     ret = inflate( &ssl->transform_in->ctx_inflate, Z_SYNC_FLUSH );
     if( ret != Z_OK )
@@ -1648,7 +1662,7 @@ static int ssl_decompress_buf( mbedtls_ssl_context *ssl )
         return( MBEDTLS_ERR_SSL_COMPRESSION_FAILED );
     }
 
-    ssl->in_msglen = MBEDTLS_SSL_IN_BUFFER_LEN -
+    ssl->in_msglen = in_buf_len -
                      ssl->transform_in->ctx_inflate.avail_out - header_bytes;
 
     MBEDTLS_SSL_DEBUG_MSG( 3, ( "after decompression: msglen = %d, ",
@@ -1682,6 +1696,11 @@ int mbedtls_ssl_fetch_input( mbedtls_ssl_context *ssl, size_t nb_want )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t len;
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t in_buf_len = ssl->in_buf_len;
+#else
+    size_t in_buf_len = MBEDTLS_SSL_IN_BUFFER_LEN;
+#endif
 
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> fetch input" ) );
 
@@ -1692,7 +1711,7 @@ int mbedtls_ssl_fetch_input( mbedtls_ssl_context *ssl, size_t nb_want )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
     }
 
-    if( nb_want > MBEDTLS_SSL_IN_BUFFER_LEN - (size_t)( ssl->in_hdr - ssl->in_buf ) )
+    if( nb_want > in_buf_len - (size_t)( ssl->in_hdr - ssl->in_buf ) )
     {
         MBEDTLS_SSL_DEBUG_MSG( 1, ( "requesting more data than fits" ) );
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
@@ -1778,7 +1797,7 @@ int mbedtls_ssl_fetch_input( mbedtls_ssl_context *ssl, size_t nb_want )
         }
         else
         {
-            len = MBEDTLS_SSL_IN_BUFFER_LEN - ( ssl->in_hdr - ssl->in_buf );
+            len = in_buf_len - ( ssl->in_hdr - ssl->in_buf );
 
             if( ssl->state != MBEDTLS_SSL_HANDSHAKE_OVER )
                 timeout = ssl->handshake->retransmit_timeout;
@@ -2523,7 +2542,11 @@ int mbedtls_ssl_write_record( mbedtls_ssl_context *ssl, uint8_t force_flush )
     {
         unsigned i;
         size_t protected_record_size;
-
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+        size_t out_buf_len = ssl->out_buf_len;
+#else
+        size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN;
+#endif
         /* Skip writing the record content type to after the encryption,
          * as it may change when using the CID extension. */
 
@@ -2539,8 +2562,7 @@ int mbedtls_ssl_write_record( mbedtls_ssl_context *ssl, uint8_t force_flush )
             mbedtls_record rec;
 
             rec.buf         = ssl->out_iv;
-            rec.buf_len     = MBEDTLS_SSL_OUT_BUFFER_LEN -
-                ( ssl->out_iv - ssl->out_buf );
+            rec.buf_len     = out_buf_len - ( ssl->out_iv - ssl->out_buf );
             rec.data_len    = ssl->out_msglen;
             rec.data_offset = ssl->out_msg - rec.buf;
 
@@ -4216,7 +4238,11 @@ static int ssl_load_buffered_record( mbedtls_ssl_context *ssl )
     unsigned char * rec;
     size_t rec_len;
     unsigned rec_epoch;
-
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t in_buf_len = ssl->in_buf_len;
+#else
+    size_t in_buf_len = MBEDTLS_SSL_IN_BUFFER_LEN;
+#endif
     if( ssl->conf->transport != MBEDTLS_SSL_TRANSPORT_DATAGRAM )
         return( 0 );
 
@@ -4246,8 +4272,7 @@ static int ssl_load_buffered_record( mbedtls_ssl_context *ssl )
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "Found buffered record from current epoch - load" ) );
 
     /* Double-check that the record is not too large */
-    if( rec_len > MBEDTLS_SSL_IN_BUFFER_LEN -
-        (size_t)( ssl->in_hdr - ssl->in_buf ) )
+    if( rec_len > in_buf_len - (size_t)( ssl->in_hdr - ssl->in_buf ) )
     {
         MBEDTLS_SSL_DEBUG_MSG( 1, ( "should never happen" ) );
         return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 63bc5c850..60ffa612d 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -3729,6 +3729,8 @@ int mbedtls_ssl_setup( mbedtls_ssl_context *ssl,
                        const mbedtls_ssl_config *conf )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    size_t in_buf_len = MBEDTLS_SSL_IN_BUFFER_LEN;
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN;
 
     ssl->conf = conf;
 
@@ -3739,18 +3741,24 @@ int mbedtls_ssl_setup( mbedtls_ssl_context *ssl,
     /* Set to NULL in case of an error condition */
     ssl->out_buf = NULL;
 
-    ssl->in_buf = mbedtls_calloc( 1, MBEDTLS_SSL_IN_BUFFER_LEN );
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    ssl->in_buf_len = in_buf_len;
+#endif
+    ssl->in_buf = mbedtls_calloc( 1, in_buf_len );
     if( ssl->in_buf == NULL )
     {
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "alloc(%d bytes) failed", MBEDTLS_SSL_IN_BUFFER_LEN) );
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "alloc(%d bytes) failed", in_buf_len ) );
         ret = MBEDTLS_ERR_SSL_ALLOC_FAILED;
         goto error;
     }
 
-    ssl->out_buf = mbedtls_calloc( 1, MBEDTLS_SSL_OUT_BUFFER_LEN );
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    ssl->out_buf_len = out_buf_len;
+#endif
+    ssl->out_buf = mbedtls_calloc( 1, out_buf_len );
     if( ssl->out_buf == NULL )
     {
-        MBEDTLS_SSL_DEBUG_MSG( 1, ( "alloc(%d bytes) failed", MBEDTLS_SSL_OUT_BUFFER_LEN) );
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "alloc(%d bytes) failed", out_buf_len ) );
         ret = MBEDTLS_ERR_SSL_ALLOC_FAILED;
         goto error;
     }
@@ -3768,6 +3776,10 @@ error:
 
     ssl->conf = NULL;
 
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    ssl->in_buf_len = 0;
+    ssl->out_buf_len = 0;
+#endif
     ssl->in_buf = NULL;
     ssl->out_buf = NULL;
 
@@ -3796,6 +3808,13 @@ error:
 int mbedtls_ssl_session_reset_int( mbedtls_ssl_context *ssl, int partial )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t in_buf_len = ssl->in_buf_len;
+    size_t out_buf_len = ssl->out_buf_len;
+#else
+    size_t in_buf_len = MBEDTLS_SSL_IN_BUFFER_LEN;
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN;
+#endif
 
 #if !defined(MBEDTLS_SSL_DTLS_CLIENT_PORT_REUSE) ||     \
     !defined(MBEDTLS_SSL_SRV_C)
@@ -3851,14 +3870,14 @@ int mbedtls_ssl_session_reset_int( mbedtls_ssl_context *ssl, int partial )
     ssl->session_in = NULL;
     ssl->session_out = NULL;
 
-    memset( ssl->out_buf, 0, MBEDTLS_SSL_OUT_BUFFER_LEN );
+    memset( ssl->out_buf, 0, out_buf_len );
 
 #if defined(MBEDTLS_SSL_DTLS_CLIENT_PORT_REUSE) && defined(MBEDTLS_SSL_SRV_C)
     if( partial == 0 )
 #endif /* MBEDTLS_SSL_DTLS_CLIENT_PORT_REUSE && MBEDTLS_SSL_SRV_C */
     {
         ssl->in_left = 0;
-        memset( ssl->in_buf, 0, MBEDTLS_SSL_IN_BUFFER_LEN );
+        memset( ssl->in_buf, 0, in_buf_len );
     }
 
 #if defined(MBEDTLS_SSL_HW_RECORD_ACCEL)
@@ -6463,6 +6482,14 @@ int mbedtls_ssl_context_load( mbedtls_ssl_context *context,
  */
 void mbedtls_ssl_free( mbedtls_ssl_context *ssl )
 {
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    size_t in_buf_len = ssl->in_buf_len;
+    size_t out_buf_len = ssl->out_buf_len;
+#else
+    size_t in_buf_len = MBEDTLS_SSL_IN_BUFFER_LEN;
+    size_t out_buf_len = MBEDTLS_SSL_OUT_BUFFER_LEN;
+#endif
+
     if( ssl == NULL )
         return;
 
@@ -6470,13 +6497,13 @@ void mbedtls_ssl_free( mbedtls_ssl_context *ssl )
 
     if( ssl->out_buf != NULL )
     {
-        mbedtls_platform_zeroize( ssl->out_buf, MBEDTLS_SSL_OUT_BUFFER_LEN );
+        mbedtls_platform_zeroize( ssl->out_buf, out_buf_len );
         mbedtls_free( ssl->out_buf );
     }
 
     if( ssl->in_buf != NULL )
     {
-        mbedtls_platform_zeroize( ssl->in_buf, MBEDTLS_SSL_IN_BUFFER_LEN );
+        mbedtls_platform_zeroize( ssl->in_buf, in_buf_len );
         mbedtls_free( ssl->in_buf );
     }
 
diff --git a/library/version_features.c b/library/version_features.c
index cc47dacc9..7b2914f5f 100644
--- a/library/version_features.c
+++ b/library/version_features.c
@@ -537,6 +537,9 @@ static const char * const features[] = {
 #if defined(MBEDTLS_SSL_TRUNCATED_HMAC)
     "MBEDTLS_SSL_TRUNCATED_HMAC",
 #endif /* MBEDTLS_SSL_TRUNCATED_HMAC */
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    "MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH",
+#endif /* MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH */
 #if defined(MBEDTLS_SSL_TRUNCATED_HMAC_COMPAT)
     "MBEDTLS_SSL_TRUNCATED_HMAC_COMPAT",
 #endif /* MBEDTLS_SSL_TRUNCATED_HMAC_COMPAT */
diff --git a/programs/test/query_config.c b/programs/test/query_config.c
index 71693c787..df532fde0 100644
--- a/programs/test/query_config.c
+++ b/programs/test/query_config.c
@@ -1474,6 +1474,14 @@ int query_config( const char *config )
     }
 #endif /* MBEDTLS_SSL_TRUNCATED_HMAC */
 
+#if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
+    if( strcmp( "MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH", config ) == 0 )
+    {
+        MACRO_EXPANSION_TO_STR( MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH );
+        return( 0 );
+    }
+#endif /* MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH */
+
 #if defined(MBEDTLS_SSL_TRUNCATED_HMAC_COMPAT)
     if( strcmp( "MBEDTLS_SSL_TRUNCATED_HMAC_COMPAT", config ) == 0 )
     {