From d6a5bfb53dd96c3f691ebdfbedeed989ad6b079e Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Mon, 5 Oct 2015 15:31:39 -0700
Subject: [PATCH] Revert "revert defragment of handshake messages in TLS"

This reverts commit 6d21d328fb2e7e2c2b3e93bc521b4af111830c64.
---
 src/internal.c     | 63 ++++++++++++++++++++++++++++++++++++++++++----
 wolfssl/internal.h |  3 +++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 9321a840e..09c6d23f7 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -1956,6 +1956,10 @@ void FreeArrays(WOLFSSL* ssl, int keep)
         XMEMCPY(ssl->session.sessionID, ssl->arrays->sessionID, ID_LEN);
         ssl->session.sessionIDSz = ssl->arrays->sessionIDSz;
     }
+    if (ssl->arrays) {
+        XFREE(ssl->arrays->pendingMsg, ssl->heap, DYNAMIC_TYPE_ARRAYS);
+        ssl->arrays->pendingMsg = NULL;
+    }
     XFREE(ssl->arrays, ssl->heap, DYNAMIC_TYPE_CERT);
     ssl->arrays = NULL;
 }
@@ -5259,16 +5263,65 @@ 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);
+            if (ssl->arrays->pendingMsg == NULL)
+                return MEMORY_E;
+            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 99831e599..420c0f677 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -2098,7 +2098,10 @@ typedef struct Options {
 } Options;
 
 typedef struct Arrays {
+    byte*           pendingMsg;         /* defrag buffer */
     word32          preMasterSz;        /* differs for DH, actual size */
+    word32          pendingMsgSz;       /* defrag buffer size */
+    word32          pendingMsgOffset;   /* current offset into defrag buffer */
 #ifndef NO_PSK
     word32          psk_keySz;          /* acutal size */
     char            client_identity[MAX_PSK_ID_LEN];