From 0feaaab922f6636bd1dc3ed0324d35c43b87cf31 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Wed, 13 Aug 2014 19:31:59 -0700
Subject: [PATCH] Moved the handshake data into its own structure so it can be
 freed when the handshake completes.

---
 src/internal.c     | 12 ++++++------
 src/ssh.c          | 29 +++++++++++++++++++++++------
 wolfssh/internal.h | 25 +++++++++++++++----------
 wolfssh/ssh.h      |  9 +++++----
 4 files changed, 49 insertions(+), 26 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index f1e2ba3..739e861 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -537,7 +537,7 @@ static int DoKexInit(WOLFSSH* ssh, uint8_t* buf, uint32_t len, uint32_t* idx)
     begin += 4 + skipSz;
 
     /* First KEX Packet Follows */
-    ssh->kexPacketFollows = buf[begin];
+    ssh->handshake->kexPacketFollows = buf[begin];
     begin += 1;
 
     /* Skip the "for future use" length. */
@@ -656,15 +656,15 @@ int ProcessClientVersion(WOLFSSH* ssh)
         return WS_VERSION_E;
     }
 
-    ssh->peerId = (char*)WMALLOC(ssh->inputBuffer.length-1, ssh->ctx->heap, WOLFSSH_ID_TYPE);
-    if (ssh->peerId == NULL) {
+    ssh->handshake->peerId = (char*)WMALLOC(ssh->inputBuffer.length-1, ssh->ctx->heap, WOLFSSH_ID_TYPE);
+    if (ssh->handshake->peerId == NULL) {
         return WS_MEMORY_E;
     }
 
-    WMEMCPY(ssh->peerId, ssh->inputBuffer.buffer, ssh->inputBuffer.length-2);
-    ssh->peerId[ssh->inputBuffer.length - 1] = 0;
+    WMEMCPY(ssh->handshake->peerId, ssh->inputBuffer.buffer, ssh->inputBuffer.length-2);
+    ssh->handshake->peerId[ssh->inputBuffer.length - 1] = 0;
     ssh->inputBuffer.idx += ssh->inputBuffer.length;
-    WLOG(WS_LOG_DEBUG, "%s", ssh->peerId);
+    WLOG(WS_LOG_DEBUG, "%s", ssh->handshake->peerId);
 
     return WS_SUCCESS;
 }
diff --git a/src/ssh.c b/src/ssh.c
index ccffded..ca1b8da 100644
--- a/src/ssh.c
+++ b/src/ssh.c
@@ -114,12 +114,21 @@ void wolfSSH_CTX_free(WOLFSSH_CTX* ctx)
 
 static WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
 {
+    HandshakeInfo* handshake;
+
     WLOG(WS_LOG_DEBUG, "Enter SshInit()");
 
     if (ssh == NULL)
         return ssh;
 
+    handshake = (HandshakeInfo*)WMALLOC(sizeof(HandshakeInfo), ctx->heap, WOLFSSH_HANDSHAKE_TYPE);
+    if (handshake == NULL) {
+        wolfSSH_free(ssh);
+        return NULL;
+    }
+
     WMEMSET(ssh, 0, sizeof(WOLFSSH));  /* default init to zeros */
+    WMEMSET(handshake, 0, sizeof(HandshakeInfo));
 
     ssh->ctx         = ctx;
     ssh->rfd         = -1;         /* set to invalid */
@@ -131,12 +140,16 @@ static WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
     ssh->publicKeyId   = ID_NONE;
     ssh->encryptionId  = ID_NONE;
     ssh->integrityId   = ID_NONE;
-    ssh->pendingKeyExchangeId = ID_NONE;
-    ssh->pendingPublicKeyId   = ID_NONE;
-    ssh->pendingEncryptionId  = ID_NONE;
-    ssh->pendingIntegrityId   = ID_NONE;
+    ssh->handshake = handshake;
+    handshake->keyExchangeId = ID_NONE;
+    handshake->publicKeyId   = ID_NONE;
+    handshake->encryptionId  = ID_NONE;
+    handshake->integrityId   = ID_NONE;
+
     if (BufferInit(&ssh->inputBuffer, 0, ctx->heap) != WS_SUCCESS ||
-        BufferInit(&ssh->outputBuffer, 0, ctx->heap) != WS_SUCCESS) {
+        BufferInit(&ssh->outputBuffer, 0, ctx->heap) != WS_SUCCESS ||
+        InitSha(&ssh->handshake->hash) != 0) {
+
         wolfSSH_free(ssh);
         ssh = NULL;
     }
@@ -174,9 +187,13 @@ static void SshResourceFree(WOLFSSH* ssh, void* heap)
     (void)heap;
 
     WLOG(WS_LOG_DEBUG, "Enter sshResourceFree()");
-    WFREE(ssh->peerId, heap, WOLFSSH_ID_TYPE);
     ShrinkBuffer(&ssh->inputBuffer, 1);
     ShrinkBuffer(&ssh->outputBuffer, 1);
+    if (ssh->handshake) {
+        XFREE(ssh->handshake->peerId, heap, WOLFSSH_ID_TYPE);
+        XMEMSET(ssh->handshake, 0, sizeof(HandshakeInfo));
+        XFREE(ssh->handshake, heap, WOLFSSH_HANDSHAKE_TYPE);
+    }
 }
 
 
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index 5e3c25d..22cc67c 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -121,6 +121,19 @@ struct WOLFSSH_CTX {
 };
 
 
+typedef struct HandshakeInfo {
+    char*          peerId;
+    uint8_t        keyExchangeId;
+    uint8_t        publicKeyId;
+    uint8_t        encryptionId;
+    uint8_t        integrityId;
+    uint8_t        kexPacketFollows;
+
+    Sha            hash;
+    uint8_t        session_id[SHA_DIGEST_SIZE];
+} HandshakeInfo;
+
+
 /* our wolfSSH session */
 struct WOLFSSH {
     WOLFSSH_CTX*   ctx;            /* owner context */
@@ -146,21 +159,13 @@ struct WOLFSSH {
     uint8_t        publicKeyId;
     uint8_t        encryptionId;
     uint8_t        integrityId;
-    uint8_t        kexPacketFollows;
-
-    char*          peerId;
-
-    uint8_t        pendingKeyExchangeId;
-    uint8_t        pendingPublicKeyId;
-    uint8_t        pendingEncryptionId;
-    uint8_t        pendingIntegrityId;
 
     Buffer         inputBuffer;
     Buffer         outputBuffer;
 
-    Sha            handshakeHash;
-    uint8_t        session_id[SHA_DIGEST_SIZE];
     uint8_t        H[SHA_DIGEST_SIZE];
+
+    HandshakeInfo* handshake;
 };
 
 
diff --git a/wolfssh/ssh.h b/wolfssh/ssh.h
index de18e3e..564955d 100644
--- a/wolfssh/ssh.h
+++ b/wolfssh/ssh.h
@@ -87,10 +87,11 @@ enum WS_EndpointTypes {
 
 /* dynamic memory types */
 enum WS_DynamicTypes {
-    WOLFSSH_CTX_TYPE    = 1,
-    WOLFSSH_TYPE        = 2,
-    WOLFSSH_TYPE_BUFFER = 3,
-    WOLFSSH_ID_TYPE     = 4
+    WOLFSSH_CTX_TYPE       = 1,
+    WOLFSSH_TYPE           = 2,
+    WOLFSSH_TYPE_BUFFER    = 3,
+    WOLFSSH_ID_TYPE        = 4,
+    WOLFSSH_HANDSHAKE_TYPE = 5
 };