From acf3156faca7ecf1bfc520a94eb06b5663e795f5 Mon Sep 17 00:00:00 2001
From: Juliusz Sosinowicz <juliusz@wolfssl.com>
Date: Fri, 18 Sep 2020 18:40:30 +0200
Subject: [PATCH] Dynamically allocate memory in wolfSSL_i2d_PKCS7_bio

---
 src/ssl.c             | 18 ++++++++++++++++--
 wolfcrypt/src/pkcs7.c | 21 ++++++++++++++++++++-
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/ssl.c b/src/ssl.c
index ecf8976e5..d8eed8e49 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -50697,7 +50697,7 @@ PKCS7* wolfSSL_d2i_PKCS7_bio(WOLFSSL_BIO* bio, PKCS7** p7)
 
 int wolfSSL_i2d_PKCS7_bio(WOLFSSL_BIO *bio, PKCS7 *p7)
 {
-    byte output[4096];
+    byte* output;
     int len;
     WOLFSSL_ENTER("wolfSSL_i2d_PKCS7_bio");
 
@@ -50706,16 +50706,30 @@ int wolfSSL_i2d_PKCS7_bio(WOLFSSL_BIO *bio, PKCS7 *p7)
         return WOLFSSL_FAILURE;
     }
 
-    if ((len = wc_PKCS7_EncodeSignedData(p7, output, sizeof(output))) < 0) {
+    if ((len = wc_PKCS7_EncodeSignedData(p7, NULL, 0)) < 0) {
         WOLFSSL_MSG("wc_PKCS7_EncodeSignedData error");
         return WOLFSSL_FAILURE;
     }
 
+    output = (byte*)XMALLOC(len, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    if (!output) {
+        WOLFSSL_MSG("malloc error");
+        return WOLFSSL_FAILURE;
+    }
+
+    if ((len = wc_PKCS7_EncodeSignedData(p7, output, len)) < 0) {
+        WOLFSSL_MSG("wc_PKCS7_EncodeSignedData error");
+        XFREE(output, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+        return WOLFSSL_FAILURE;
+    }
+
     if (wolfSSL_BIO_write(bio, output, len) <= 0) {
         WOLFSSL_MSG("wolfSSL_BIO_write error");
+        XFREE(output, NULL, DYNAMIC_TYPE_TMP_BUFFER);
         return WOLFSSL_FAILURE;
     }
 
+    XFREE(output, NULL, DYNAMIC_TYPE_TMP_BUFFER);
     return WOLFSSL_SUCCESS;
 }
 
diff --git a/wolfcrypt/src/pkcs7.c b/wolfcrypt/src/pkcs7.c
index 78a368ef0..b11660029 100644
--- a/wolfcrypt/src/pkcs7.c
+++ b/wolfcrypt/src/pkcs7.c
@@ -2235,6 +2235,7 @@ static int wc_PKCS7_SignedDataBuildSignature(PKCS7* pkcs7,
 
 
 /* build PKCS#7 signedData content type */
+/* To get the output size then set output = 0 and *outputSz = 0 */
 static int PKCS7_EncodeSigned(PKCS7* pkcs7, ESD* esd,
     const byte* hashBuf, word32 hashSz, byte* output, word32* outputSz,
     byte* output2, word32* output2Sz)
@@ -2270,7 +2271,7 @@ static int PKCS7_EncodeSigned(PKCS7* pkcs7, ESD* esd,
     byte signingTime[MAX_TIME_STRING_SZ];
 
     if (pkcs7 == NULL || pkcs7->hashOID == 0 ||
-        output == NULL || outputSz == NULL || *outputSz == 0 || hashSz == 0 ||
+        outputSz == NULL || hashSz == 0 ||
         hashBuf == NULL) {
         return BAD_FUNC_ARG;
     }
@@ -2511,6 +2512,11 @@ static int PKCS7_EncodeSigned(PKCS7* pkcs7, ESD* esd,
         #ifdef WOLFSSL_SMALL_STACK
             XFREE(esd, pkcs7->heap, DYNAMIC_TYPE_TMP_BUFFER);
         #endif
+            if (*outputSz == 0 && *output2Sz == 0) {
+                *outputSz = totalSz;
+                *output2Sz = total2Sz;
+                return 0;
+            }
             return BUFFER_E;
         }
 
@@ -2528,6 +2534,19 @@ static int PKCS7_EncodeSigned(PKCS7* pkcs7, ESD* esd,
             XFREE(flatSignedAttribs, pkcs7->heap, DYNAMIC_TYPE_PKCS7);
     #ifdef WOLFSSL_SMALL_STACK
         XFREE(esd, pkcs7->heap, DYNAMIC_TYPE_TMP_BUFFER);
+    #endif
+        if (*outputSz == 0) {
+            *outputSz = totalSz;
+            return totalSz;
+        }
+        return BUFFER_E;
+    }
+
+    if (output == NULL) {
+        if (pkcs7->signedAttribsSz != 0)
+            XFREE(flatSignedAttribs, pkcs7->heap, DYNAMIC_TYPE_PKCS7);
+    #ifdef WOLFSSL_SMALL_STACK
+        XFREE(esd, pkcs7->heap, DYNAMIC_TYPE_TMP_BUFFER);
     #endif
         return BUFFER_E;
     }