From dbbe886efb19f28557ac898ed903fe328c0eff63 Mon Sep 17 00:00:00 2001
From: Jacob Barthelmeh <jacob@wolfssl.com>
Date: Thu, 27 Dec 2018 10:00:43 -0700
Subject: [PATCH] add wolfSSH_SFTP_ClearState function for free'ing states

---
 src/wolfsftp.c | 49 +++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/wolfsftp.c b/src/wolfsftp.c
index 700c89e..c4f3561 100644
--- a/src/wolfsftp.c
+++ b/src/wolfsftp.c
@@ -38,6 +38,14 @@
 #endif
 
 
+/* enum for bit field with an ID of each of the state structures */
+enum WS_SFTP_STATE_ID {
+    STATE_ID_ALL   = 0, /* default to select all */
+    STATE_ID_LSTAT = 0x01,
+    STATE_ID_OPEN  = 0x02,
+    STATE_ID_GET   = 0x04,
+};
+
 enum WS_SFTP_LSTAT_STATE_ID {
     STATE_LSTAT_INIT
 };
@@ -95,6 +103,38 @@ static int SFTP_GetAttributes_Handle(WOLFSSH* ssh, byte* handle, int handleSz,
         WS_SFTP_FILEATRB* atr);
 static WS_SFTPNAME* wolfSSH_SFTPNAME_new(void* heap);
 
+
+/* Used to clear and free all states. Should be when returning errors or success
+ * Must be called when free'ing the SFTP. For now static since only used in
+ * wolfsftp.c
+ *
+ * Note: Most cases an error will free all states and a success will free
+ *       specific state ID.
+ */
+static void wolfSSH_SFTP_ClearState(WOLFSSH* ssh, enum WS_SFTP_STATE_ID state)
+{
+    if (ssh) {
+
+        if (state == 0)
+            state = ~state; /* set all bits hot */
+
+        if (state & STATE_ID_GET) {
+            XFREE(ssh->getState, ssh->ctx->heap, DYNTYPE_SFTP_STATE);
+            ssh->getState   = NULL;
+        }
+        if (state & STATE_ID_LSTAT) {
+            XFREE(ssh->lstatState, ssh->ctx->heap, DYNTYPE_SFTP_STATE);
+            ssh->lstatState = NULL;
+        }
+
+        if (state & STATE_ID_OPEN) {
+            XFREE(ssh->openState, ssh->ctx->heap, DYNTYPE_SFTP_STATE);
+            ssh->openState  = NULL;
+        }
+    }
+}
+
+
 /* Gets packet header information
  * request Id, type, and size of type specific data
  * return value is length of type specific data still on the wire to be read
@@ -107,12 +147,14 @@ static int SFTP_GetHeader(WOLFSSH* ssh, word32* reqId, byte* type)
 
     if (type == NULL || reqId == NULL || ssh == NULL) {
         WLOG(WS_LOG_SFTP, "NULL argument error");
+        wolfSSH_SFTP_ClearState(ssh, STATE_ID_ALL);
         return WS_BAD_ARGUMENT;
     }
 
     ret = wolfSSH_stream_read(ssh, buf, sizeof(buf));
-    if (ret < WOLFSSH_SFTP_HEADER) {
+    if (ret < WOLFSSH_SFTP_HEADER && ret != WS_WANT_READ) {
         WLOG(WS_LOG_SFTP, "Unable to read SFTP header");
+        wolfSSH_SFTP_ClearState(ssh, STATE_ID_ALL);
         return WS_FATAL_ERROR;
     }
 
@@ -164,13 +206,15 @@ static int SFTP_ServerRecvInit(WOLFSSH* ssh) {
 
     ato32(buf, &sz);
     if (sz < MSG_ID_SZ + UINT32_SZ) {
-       return WS_BUFFER_E;
+        wolfSSH_SFTP_ClearState(ssh, STATE_ID_ALL);
+        return WS_BUFFER_E;
     }
 
     /* compare versions supported */
     id = buf[LENGTH_SZ];
     if (id != WOLFSSH_FTP_INIT) {
         WLOG(WS_LOG_SFTP, "Unexpected SFTP type received");
+        wolfSSH_SFTP_ClearState(ssh, STATE_ID_ALL);
         return WS_BUFFER_E;
     }
 
@@ -4660,6 +4704,7 @@ int wolfSSH_SFTP_free(WOLFSSH* ssh)
         }
     }
 #endif
+    wolfSSH_SFTP_ClearState(ssh, STATE_ID_ALL);
     return WS_SUCCESS;
 }