From fbadac84cbb78684624b6528702ed73cb6a7d7b9 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Tue, 2 Jul 2024 13:23:39 -0700
Subject: [PATCH] Reduce SCP Allocations

1. Create the SCP file and recv msg buffer once, and just keep it.
2. Whitespace.
---
 src/internal.c |  4 +++-
 src/wolfscp.c  | 36 +++++++++++++-----------------------
 2 files changed, 16 insertions(+), 24 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 3eab81dba..07c28cfd1 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -1043,6 +1043,8 @@ WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
     ssh->scpRequestState = SCP_PARSE_COMMAND;
     ssh->scpConfirmMsg   = NULL;
     ssh->scpConfirmMsgSz = 0;
+    ssh->scpRecvMsg      = NULL;
+    ssh->scpRecvMsgSz    = 0;
     ssh->scpRecvCtx      = NULL;
     #if !defined(WOLFSSH_SCP_USER_CALLBACKS) && !defined(NO_FILESYSTEM)
     ssh->scpSendCtx      = &(ssh->scpSendCbCtx);
@@ -7143,7 +7145,7 @@ static int DoUserAuthRequestPublicKey(WOLFSSH* ssh, WS_UserAuthData* authData,
                 case WOLFSSH_USERAUTH_PARTIAL_SUCCESS:
                     WLOG(WS_LOG_DEBUG, "DUARPK: user auth partial success");
                     break;
-                    
+
                 case WOLFSSH_USERAUTH_WOULD_BLOCK:
                     WLOG(WS_LOG_DEBUG, "DUARPK: userauth callback would block");
                     break;
diff --git a/src/wolfscp.c b/src/wolfscp.c
index 480918767..1eeeb1a97 100644
--- a/src/wolfscp.c
+++ b/src/wolfscp.c
@@ -98,7 +98,8 @@ int DoScpSink(WOLFSSH* ssh)
 
                 ssh->scpConfirm = ssh->ctx->scpRecvCb(ssh,
                         WOLFSSH_SCP_NEW_REQUEST, ssh->scpBasePath,
-                        NULL, 0, 0, 0, 0, NULL, 0, 0, wolfSSH_GetScpRecvCtx(ssh));
+                        NULL, 0, 0, 0, 0, NULL, 0, 0,
+                        wolfSSH_GetScpRecvCtx(ssh));
                 continue;
 
             case SCP_RECEIVE_MESSAGE:
@@ -188,12 +189,9 @@ int DoScpSink(WOLFSSH* ssh)
                         ssh->scpATime, ssh->scpFileSz, ssh->scpFileBuffer,
                         ssh->scpFileBufferSz, ssh->scpFileOffset,
                         wolfSSH_GetScpRecvCtx(ssh));
-
                 ssh->scpFileOffset += ssh->scpFileBufferSz;
 
-                /* shrink and reset recv buffer */
-                WFREE(ssh->scpFileBuffer, ssh->ctx->heap, DYNTYPE_BUFFER);
-                ssh->scpFileBuffer = NULL;
+                /* reset recv buffer */
                 ssh->scpFileBufferSz = 0;
 
                 if (ssh->scpConfirm != WS_SCP_CONTINUE) {
@@ -1445,8 +1443,6 @@ int ReceiveScpMessage(WOLFSSH* ssh)
             break;
     }
 
-    WFREE(ssh->scpRecvMsg, ssh->ctx->heap, DYNTYPE_STRING);
-    ssh->scpRecvMsg = NULL;
     ssh->scpRecvMsgSz = 0;
 
     return ret;
@@ -1455,34 +1451,28 @@ int ReceiveScpMessage(WOLFSSH* ssh)
 int ReceiveScpFile(WOLFSSH* ssh)
 {
     int partSz, ret = WS_SUCCESS;
-    byte* part;
 
     if (ssh == NULL)
         return WS_BAD_ARGUMENT;
 
+    /* We don't want to over-read the buffer. The file data is
+     * terminated by the sender with a nul which is checked later. */
     partSz = min(ssh->scpFileSz - ssh->scpFileOffset, DEFAULT_SCP_BUFFER_SZ);
 
     /* don't even bother reading if read size is 0 */
     if (partSz == 0) return ret;
 
-    part = (byte*)WMALLOC(partSz, ssh->ctx->heap, DYNTYPE_BUFFER);
-    if (part == NULL)
-        ret = WS_MEMORY_E;
+    if (ssh->scpFileBuffer == NULL) {
+        ssh->scpFileBuffer = (byte*)WMALLOC(DEFAULT_SCP_BUFFER_SZ,
+                ssh->ctx->heap, DYNTYPE_BUFFER);
+        if (ssh->scpFileBuffer == NULL)
+            ret = WS_MEMORY_E;
+    }
 
     if (ret == WS_SUCCESS) {
-        WMEMSET(part, 0, partSz);
-
-        ret = wolfSSH_stream_read(ssh, part, partSz);
+        ret = wolfSSH_stream_read(ssh, ssh->scpFileBuffer, partSz);
         if (ret > 0) {
-            if (ssh->scpFileBuffer != NULL) {
-                WFREE(ssh->scpFileBuffer, ssh->ctx->heap, DYNTYPE_BUFFER);
-                ssh->scpFileBuffer = NULL;
-                ssh->scpFileBufferSz = 0;
-            }
-            ssh->scpFileBuffer = part;
             ssh->scpFileBufferSz = ret;
-        } else {
-            WFREE(part, ssh->ctx->heap, DYNTYPE_BUFFER);
         }
     }
 
@@ -2122,7 +2112,7 @@ static int GetFileStats(void *fs, ScpSendCtx* ctx, const char* fileName,
                         word64* mTime, word64* aTime, int* fileMode)
 {
     int ret = WS_SUCCESS;
-    
+
     WOLFSSH_UNUSED(fs);
 
     if (ctx == NULL || fileName == NULL || mTime == NULL ||