Defragment incoming TLS handshake messages

Signed-off-by: Deomid rojer Ryabkov <rojer@rojer.me>
diff --git a/library/ssl_msg.c b/library/ssl_msg.c
index 7000e93..1c548ec 100644
--- a/library/ssl_msg.c
+++ b/library/ssl_msg.c
@@ -3225,7 +3225,11 @@
         return MBEDTLS_ERR_SSL_INVALID_RECORD;
     }
 
-    ssl->in_hslen = mbedtls_ssl_hs_hdr_len(ssl) + ssl_get_hs_total_len(ssl);
+    if (ssl->in_hslen == 0) {
+        ssl->in_hslen = mbedtls_ssl_hs_hdr_len(ssl) + ssl_get_hs_total_len(ssl);
+        ssl->in_hsfraglen = 0;
+        ssl->in_hshdr = ssl->in_hdr;
+    }
 
     MBEDTLS_SSL_DEBUG_MSG(3, ("handshake message: msglen ="
                               " %" MBEDTLS_PRINTF_SIZET ", type = %u, hslen = %"
@@ -3291,10 +3295,59 @@
         }
     } else
 #endif /* MBEDTLS_SSL_PROTO_DTLS */
-    /* With TLS we don't handle fragmentation (for now) */
-    if (ssl->in_msglen < ssl->in_hslen) {
-        MBEDTLS_SSL_DEBUG_MSG(1, ("TLS handshake fragmentation not supported"));
-        return MBEDTLS_ERR_SSL_FEATURE_UNAVAILABLE;
+    {
+        int ret;
+        const size_t hs_remain = ssl->in_hslen - ssl->in_hsfraglen;
+        const size_t msg_hslen = (hs_remain <= ssl->in_msglen ? hs_remain : ssl->in_msglen);
+
+        MBEDTLS_SSL_DEBUG_MSG(3,
+                              ("handshake fragment: %" MBEDTLS_PRINTF_SIZET " .. %"
+                               MBEDTLS_PRINTF_SIZET " of %"
+                               MBEDTLS_PRINTF_SIZET " msglen %" MBEDTLS_PRINTF_SIZET,
+                               ssl->in_hsfraglen, ssl->in_hsfraglen + msg_hslen,
+                               ssl->in_hslen, ssl->in_msglen));
+        (void) msg_hslen;
+        if (ssl->in_msglen < hs_remain) {
+            ssl->in_hsfraglen += ssl->in_msglen;
+            ssl->in_hdr = ssl->in_msg + ssl->in_msglen;
+            ssl->in_msglen = 0;
+            mbedtls_ssl_update_in_pointers(ssl);
+            return MBEDTLS_ERR_SSL_CONTINUE_PROCESSING;
+        }
+        if (ssl->in_hshdr != ssl->in_hdr) {
+            /*
+             * At ssl->in_hshdr we have a sequence of records that cover the next handshake
+             * record, each with its own record header that we need to remove.
+             * Note that the reassembled record size may not equal the size of the message,
+             * there maybe bytes from the next message following it.
+             */
+            size_t merged_rec_len = 0;
+            unsigned char *p = ssl->in_hshdr, *q = NULL;
+            do {
+                mbedtls_record rec;
+                ret = ssl_parse_record_header(ssl, p, mbedtls_ssl_in_hdr_len(ssl), &rec);
+                if (ret != 0) {
+                    return ret;
+                }
+                merged_rec_len += rec.data_len;
+                p = rec.buf + rec.buf_len;
+                if (q != NULL) {
+                    memmove(q, rec.buf + rec.data_offset, rec.data_len);
+                    q += rec.data_len;
+                } else {
+                    q = p;
+                }
+            } while (merged_rec_len < ssl->in_hslen);
+            ssl->in_hdr = ssl->in_hshdr;
+            mbedtls_ssl_update_in_pointers(ssl);
+            ssl->in_msglen = merged_rec_len;
+            /* Adjust message length. */
+            MBEDTLS_PUT_UINT16_BE(merged_rec_len, ssl->in_len, 0);
+            ssl->in_hsfraglen = 0;
+            ssl->in_hshdr = NULL;
+            MBEDTLS_SSL_DEBUG_BUF(4, "reassembled record",
+                                  ssl->in_hdr, mbedtls_ssl_in_hdr_len(ssl) + merged_rec_len);
+        }
     }
 
     return 0;
@@ -4639,6 +4692,16 @@
             return MBEDTLS_ERR_SSL_INTERNAL_ERROR;
         }
 
+        if (ssl->in_hsfraglen != 0) {
+            /* Not all handshake fragments have arrived, do not consume. */
+            MBEDTLS_SSL_DEBUG_MSG(3,
+                                  ("waiting for more fragments (%" MBEDTLS_PRINTF_SIZET " of %"
+                                   MBEDTLS_PRINTF_SIZET ", %" MBEDTLS_PRINTF_SIZET " left)",
+                                   ssl->in_hsfraglen, ssl->in_hslen,
+                                   ssl->in_hslen - ssl->in_hsfraglen));
+            return 0;
+        }
+
         /*
          * Get next Handshake message in the current record
          */
@@ -4664,6 +4727,7 @@
             ssl->in_msglen -= ssl->in_hslen;
             memmove(ssl->in_msg, ssl->in_msg + ssl->in_hslen,
                     ssl->in_msglen);
+            MBEDTLS_PUT_UINT16_BE(ssl->in_msglen, ssl->in_len, 0);
 
             MBEDTLS_SSL_DEBUG_BUF(4, "remaining content in record",
                                   ssl->in_msg, ssl->in_msglen);
@@ -5338,7 +5402,7 @@
     } else
 #endif
     {
-        ssl->in_ctr = ssl->in_hdr - MBEDTLS_SSL_SEQUENCE_NUMBER_LEN;
+        ssl->in_ctr = ssl->in_buf;
         ssl->in_len = ssl->in_hdr + 3;
 #if defined(MBEDTLS_SSL_DTLS_CONNECTION_ID)
         ssl->in_cid = ssl->in_len;
@@ -5354,24 +5418,35 @@
  * Setup an SSL context
  */
 
-void mbedtls_ssl_reset_in_out_pointers(mbedtls_ssl_context *ssl)
+void mbedtls_ssl_reset_in_pointers(mbedtls_ssl_context *ssl)
+{
+#if defined(MBEDTLS_SSL_PROTO_DTLS)
+    if (ssl->conf->transport == MBEDTLS_SSL_TRANSPORT_DATAGRAM) {
+        ssl->in_hdr = ssl->in_buf;
+    } else
+#endif
+    {
+        ssl->in_hdr = ssl->in_buf + MBEDTLS_SSL_SEQUENCE_NUMBER_LEN;
+    }
+
+    /* Derive other internal pointers. */
+    mbedtls_ssl_update_in_pointers(ssl);
+}
+
+void mbedtls_ssl_reset_out_pointers(mbedtls_ssl_context *ssl)
 {
     /* Set the incoming and outgoing record pointers. */
 #if defined(MBEDTLS_SSL_PROTO_DTLS)
     if (ssl->conf->transport == MBEDTLS_SSL_TRANSPORT_DATAGRAM) {
         ssl->out_hdr = ssl->out_buf;
-        ssl->in_hdr  = ssl->in_buf;
     } else
 #endif /* MBEDTLS_SSL_PROTO_DTLS */
     {
         ssl->out_ctr = ssl->out_buf;
-        ssl->out_hdr = ssl->out_buf + 8;
-        ssl->in_hdr  = ssl->in_buf  + 8;
+        ssl->out_hdr = ssl->out_buf + MBEDTLS_SSL_SEQUENCE_NUMBER_LEN;
     }
-
     /* Derive other internal pointers. */
     mbedtls_ssl_update_out_pointers(ssl, NULL /* no transform enabled */);
-    mbedtls_ssl_update_in_pointers(ssl);
 }
 
 /*