From b1e37377a14a593474144b578d0ebce4ea5b9543 Mon Sep 17 00:00:00 2001
From: Sean Parkinson <sean@wolfssl.com>
Date: Thu, 27 Oct 2022 17:00:42 +1000
Subject: [PATCH] TLS performance fix: ForceZero minimization

Don't ForceZero the output buffer before free.
ForceZero it when encryption fails.

ShrinkInputBuffer needs to zeroize input buffer even if not currently
encrypting as it may be using the buffer on wolfSSL object reuse.

Fix SP to zeroize the whole buffer.

Fix DH to check cBuf when WOLFSSL_CHECK_MEM_ZERO defined.
---
 src/internal.c         | 29 +++++++++++++----------------
 src/tls13.c            |  9 +++++++++
 wolfcrypt/src/dh.c     |  2 ++
 wolfcrypt/src/sp_int.c |  2 +-
 4 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 2d83cdbbb9..4d4d28353b 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -9784,11 +9784,6 @@ static int wolfSSLReceive(WOLFSSL* ssl, byte* buf, word32 sz)
 void ShrinkOutputBuffer(WOLFSSL* ssl)
 {
     WOLFSSL_MSG("Shrinking output buffer");
-    if (IsEncryptionOn(ssl, 0)) {
-        ForceZero(ssl->buffers.outputBuffer.buffer -
-            ssl->buffers.outputBuffer.offset,
-            ssl->buffers.outputBuffer.bufferSize);
-    }
     XFREE(ssl->buffers.outputBuffer.buffer - ssl->buffers.outputBuffer.offset,
           ssl->heap, DYNAMIC_TYPE_OUT_BUFFER);
     ssl->buffers.outputBuffer.buffer = ssl->buffers.outputBuffer.staticBuffer;
@@ -9819,11 +9814,9 @@ void ShrinkInputBuffer(WOLFSSL* ssl, int forcedFree)
                usedLength);
     }
 
-    if (IsEncryptionOn(ssl, 1) || forcedFree) {
-        ForceZero(ssl->buffers.inputBuffer.buffer -
-            ssl->buffers.inputBuffer.offset,
-            ssl->buffers.inputBuffer.bufferSize);
-    }
+    ForceZero(ssl->buffers.inputBuffer.buffer -
+        ssl->buffers.inputBuffer.offset,
+        ssl->buffers.inputBuffer.bufferSize);
     XFREE(ssl->buffers.inputBuffer.buffer - ssl->buffers.inputBuffer.offset,
           ssl->heap, DYNAMIC_TYPE_IN_BUFFER);
     ssl->buffers.inputBuffer.buffer = ssl->buffers.inputBuffer.staticBuffer;
@@ -9968,11 +9961,6 @@ static WC_INLINE int GrowOutputBuffer(WOLFSSL* ssl, int size)
                ssl->buffers.outputBuffer.length);
 
     if (ssl->buffers.outputBuffer.dynamicFlag) {
-        if (IsEncryptionOn(ssl, 0)) {
-            ForceZero(ssl->buffers.outputBuffer.buffer -
-                ssl->buffers.outputBuffer.offset,
-                ssl->buffers.outputBuffer.bufferSize);
-        }
         XFREE(ssl->buffers.outputBuffer.buffer -
               ssl->buffers.outputBuffer.offset, ssl->heap,
               DYNAMIC_TYPE_OUT_BUFFER);
@@ -20819,8 +20807,17 @@ int BuildMessage(WOLFSSL* ssl, byte* output, int outSz, const byte* input,
     #endif
             }
 
-            if (ret != 0)
+            if (ret != 0) {
+            #ifdef WOLFSSL_ASYNC_CRYPT
+                if (ret != WC_PENDING_E)
+            #endif
+                {
+                    /* Zeroize plaintext. */
+                    ForceZero(output + args->headerSz,
+                        (word16)(args->size - args->digestSz));
+                }
                 goto exit_buildmsg;
+            }
             ssl->options.buildMsgState = BUILD_MSG_ENCRYPTED_VERIFY_MAC;
         }
         FALL_THROUGH;
diff --git a/src/tls13.c b/src/tls13.c
index 7cbe43cc5d..4a6bcfd90d 100644
--- a/src/tls13.c
+++ b/src/tls13.c
@@ -3008,6 +3008,15 @@ int BuildTls13Message(WOLFSSL* ssl, byte* output, int outSz, const byte* input,
                 output += args->headerSz;
                 ret = EncryptTls13(ssl, output, output, args->size, aad,
                                    (word16)args->headerSz, asyncOkay);
+                if (ret != 0) {
+                #ifdef WOLFSSL_ASYNC_CRYPT
+                    if (ret != WC_PENDING_E)
+                #endif
+                    {
+                        /* Zeroize plaintext. */
+                        ForceZero(output, args->size);
+                    }
+                }
 #ifdef WOLFSSL_DTLS13
                 if (ret == 0 && ssl->options.dtls) {
                     /* AAD points to the header. Reuse the variable  */
diff --git a/wolfcrypt/src/dh.c b/wolfcrypt/src/dh.c
index 900729e374..bf54ecc580 100644
--- a/wolfcrypt/src/dh.c
+++ b/wolfcrypt/src/dh.c
@@ -1161,6 +1161,8 @@ static int GeneratePrivateDh186(DhKey* key, WC_RNG* rng, byte* priv,
     ForceZero(cBuf, cSz);
 #if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
     XFREE(cBuf, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
+#elif defined(WOLFSSL_CHECK_MEM_ZERO)
+    wc_MemZero_Check(cBuf, cSz);
 #endif
 
     /* tmpQ: M = min(2^N,q) - 1 */
diff --git a/wolfcrypt/src/sp_int.c b/wolfcrypt/src/sp_int.c
index 0130cf3e28..8a1d125fab 100644
--- a/wolfcrypt/src/sp_int.c
+++ b/wolfcrypt/src/sp_int.c
@@ -4687,7 +4687,7 @@ void sp_forcezero(sp_int* a)
 {
     if (a != NULL) {
         /* Ensure all data zeroized - data not zeroed when used decreases. */
-        ForceZero(a->dp, a->used * sizeof(sp_int_digit));
+        ForceZero(a->dp, a->size * sizeof(sp_int_digit));
         _sp_zero(a);
     #ifdef HAVE_WOLF_BIGINT
         wc_bigint_zero(&a->raw);