From ef43ce6e25054718cee7e0837f127b39b9cb34db Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Tue, 9 Oct 2018 08:24:12 -0400
Subject: [PATCH] Dtls: change the way unlimited mtu is set for client hello
 messages

---
 library/ssl_tls.c | 22 ++++++----------------
 1 file changed, 6 insertions(+), 16 deletions(-)

diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 38826f93f..353c3471b 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -2948,7 +2948,6 @@ int mbedtls_ssl_resend( mbedtls_ssl_context *ssl )
 int mbedtls_ssl_flight_transmit( mbedtls_ssl_context *ssl )
 {
     int ret;
-    uint16_t mtu_temp = 0;
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> mbedtls_ssl_flight_transmit" ) );
 
     if( ssl->handshake->retransmit_state != MBEDTLS_SSL_RETRANS_SENDING )
@@ -2983,15 +2982,6 @@ int mbedtls_ssl_flight_transmit( mbedtls_ssl_context *ssl )
             ssl_swap_epochs( ssl );
         }
 
-        /* Disable handshake mtu for client hello message to avoid fragmentation.
-         * Setting it back after calling mbedtls_ssl_write_record */
-        if( ssl->out_msg[0] == MBEDTLS_SSL_HS_CLIENT_HELLO )
-        {
-            mtu_temp = ssl->handshake->mtu;
-            ssl->handshake->mtu = 0;
-            MBEDTLS_SSL_DEBUG_MSG( 2, ( "disabling fragmentation of ClientHello message" ) );
-        }
-
         ret = ssl_get_remaining_payload_in_datagram( ssl );
         if( ret < 0 )
             return( ret );
@@ -3090,12 +3080,6 @@ int mbedtls_ssl_flight_transmit( mbedtls_ssl_context *ssl )
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_write_record", ret );
             return( ret );
         }
-
-        if( mtu_temp != 0 )
-        {
-            ssl->handshake->mtu = mtu_temp;
-            mtu_temp = 0;
-        }
     }
 
     if( ( ret = mbedtls_ssl_flush_output( ssl ) ) != 0 )
@@ -7924,6 +7908,12 @@ size_t mbedtls_ssl_get_max_frag_len( const mbedtls_ssl_context *ssl )
 #if defined(MBEDTLS_SSL_PROTO_DTLS)
 static size_t ssl_get_current_mtu( const mbedtls_ssl_context *ssl )
 {
+    /* Return unlimited mtu for client hello messages to avoid fragmentation. */
+    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_CLIENT &&
+        ( ssl->state == MBEDTLS_SSL_CLIENT_HELLO ||
+          ssl->state == MBEDTLS_SSL_SERVER_HELLO ) )
+        return ( 0 );
+
     if( ssl->handshake == NULL || ssl->handshake->mtu == 0 )
         return( ssl->mtu );