From 7d067dfec6ddb931399950c9c2dfb1510ee10e00 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Tue, 25 Aug 2015 15:55:09 -0400
Subject: [PATCH] defragment the handshake messages in TLS

---
 src/internal.c     | 61 +++++++++++++++++++++++++++++++++++++++++-----
 wolfssl/internal.h |  4 +++
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 73d837847..5916e738f 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -1829,7 +1829,9 @@ void FreeArrays(WOLFSSL* ssl, int keep)
         XMEMCPY(ssl->session.sessionID, ssl->arrays->sessionID, ID_LEN);
         ssl->session.sessionIDSz = ssl->arrays->sessionIDSz;
     }
-    XFREE(ssl->arrays, ssl->heap, DYNAMIC_TYPE_ARRAYS);
+    if (ssl->arrays)
+        XFREE(ssl->arrays->pendingMsg, ssl->heap, DYNAMIC_TYPE_ARRAYS);
+    XFREE(ssl->arrays, ssl->heap, DYNAMIC_TYPE_CERT);
     ssl->arrays = NULL;
 }
 
@@ -5106,16 +5108,63 @@ static int DoHandShakeMsgType(WOLFSSL* ssl, byte* input, word32* inOutIdx,
 static int DoHandShakeMsg(WOLFSSL* ssl, byte* input, word32* inOutIdx,
                           word32 totalSz)
 {
-    byte   type;
-    word32 size;
     int    ret = 0;
 
     WOLFSSL_ENTER("DoHandShakeMsg()");
 
-    if (GetHandShakeHeader(ssl, input, inOutIdx, &type, &size, totalSz) != 0)
-        return PARSE_ERROR;
+    /* If there is a pending fragmented handshake message, pending message size
+     * will be non-zero. */
+    if (ssl->arrays->pendingMsgSz == 0) {
+        byte   type;
+        word32 size;
 
-    ret = DoHandShakeMsgType(ssl, input, inOutIdx, type, size, totalSz);
+        if (GetHandShakeHeader(ssl,input, inOutIdx, &type, &size, totalSz) != 0)
+            return PARSE_ERROR;
+
+        /* size is the size of the certificate message payload */
+        if (totalSz - HANDSHAKE_HEADER_SZ < size) {
+            ssl->arrays->pendingMsgType = type;
+            ssl->arrays->pendingMsgSz = size + HANDSHAKE_HEADER_SZ;
+            ssl->arrays->pendingMsg = (byte*)XMALLOC(size + HANDSHAKE_HEADER_SZ,
+                                                     ssl->heap,
+                                                     DYNAMIC_TYPE_ARRAYS);
+            XMEMCPY(ssl->arrays->pendingMsg,
+                    input + *inOutIdx - HANDSHAKE_HEADER_SZ, totalSz);
+            ssl->arrays->pendingMsgOffset = totalSz;
+            *inOutIdx += totalSz - HANDSHAKE_HEADER_SZ;
+            return 0;
+        }
+
+        ret = DoHandShakeMsgType(ssl, input, inOutIdx, type, size, totalSz);
+    }
+    else {
+        if (totalSz + ssl->arrays->pendingMsgOffset
+                                                  > ssl->arrays->pendingMsgSz) {
+
+            return BUFFER_ERROR;
+        }
+        else {
+            XMEMCPY(ssl->arrays->pendingMsg + ssl->arrays->pendingMsgOffset,
+                    input + *inOutIdx, totalSz);
+            ssl->arrays->pendingMsgOffset += totalSz;
+            *inOutIdx += totalSz;
+        }
+
+        if (ssl->arrays->pendingMsgOffset == ssl->arrays->pendingMsgSz)
+        {
+            word32 idx = 0;
+            ret = DoHandShakeMsgType(ssl,
+                                     ssl->arrays->pendingMsg
+                                                          + HANDSHAKE_HEADER_SZ,
+                                     &idx, ssl->arrays->pendingMsgType,
+                                     ssl->arrays->pendingMsgSz
+                                                          - HANDSHAKE_HEADER_SZ,
+                                     ssl->arrays->pendingMsgSz);
+            XFREE(ssl->arrays->pendingMsg, ssl->heap, DYNAMIC_TYPE_ARRAYS);
+            ssl->arrays->pendingMsg = NULL;
+            ssl->arrays->pendingMsgSz = 0;
+        }
+    }
 
     WOLFSSL_LEAVE("DoHandShakeMsg()", ret);
     return ret;
diff --git a/wolfssl/internal.h b/wolfssl/internal.h
index 6fe87a9f1..93e3b3aa7 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -2067,6 +2067,8 @@ typedef struct Options {
 
 typedef struct Arrays {
     word32          preMasterSz;        /* differs for DH, actual size */
+    word32          pendingMsgSz;
+    word32          pendingMsgOffset;
 #ifndef NO_PSK
     word32          psk_keySz;          /* acutal size */
     char            client_identity[MAX_PSK_ID_LEN];
@@ -2083,6 +2085,8 @@ typedef struct Arrays {
     byte            cookie[MAX_COOKIE_LEN];
     byte            cookieSz;
 #endif
+    byte            pendingMsgType;
+    byte*           pendingMsg;
 } Arrays;
 
 #ifndef ASN_NAME_MAX