From 19244fc0c9cd1b5ce131791274b65bde4b4e7b72 Mon Sep 17 00:00:00 2001
From: Jacob Barthelmeh <jacob@wolfssl.com>
Date: Fri, 16 Jun 2017 11:52:52 -0600
Subject: [PATCH] fix memory management

---
 src/bio.c   | 2 +-
 src/ssl.c   | 8 ++++++--
 tests/api.c | 3 +++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/bio.c b/src/bio.c
index 574e6c8a3..da19fd5ec 100644
--- a/src/bio.c
+++ b/src/bio.c
@@ -673,7 +673,7 @@ long wolfSSL_BIO_get_mem_ptr(WOLFSSL_BIO *bio, WOLFSSL_BUF_MEM **ptr)
         return WOLFSSL_FAILURE;
     }
 
-    if (bio->type == WOLFSSL_BIO_FILE) {
+    if (bio->type == WOLFSSL_BIO_FILE || bio->type == WOLFSSL_BIO_SOCKET) {
         WOLFSSL_MSG("NO memory buffer for FILE type");
         return SSL_FAILURE;
     }
diff --git a/src/ssl.c b/src/ssl.c
index cef4228d1..f8ac0da1d 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -12198,7 +12198,8 @@ int wolfSSL_set_compression(WOLFSSL* ssl)
             XMEMSET(bio, 0, sizeof(WOLFSSL_BIO));
             bio->type   = method->type;
             bio->close  = BIO_CLOSE; /* default to close things */
-            if (method->type != WOLFSSL_BIO_FILE) {
+            if (method->type != WOLFSSL_BIO_FILE &&
+                    method->type != WOLFSSL_BIO_SOCKET) {
                 bio->mem_buf =(WOLFSSL_BUF_MEM*)XMALLOC(sizeof(WOLFSSL_BUF_MEM),
                                                        0, DYNAMIC_TYPE_OPENSSL);
                 if (bio->mem_buf == NULL) {
@@ -12239,7 +12240,7 @@ int wolfSSL_set_compression(WOLFSSL* ssl)
         bio->memLen = bio->wrSz = len;
         bio->mem    = (byte*)XMALLOC(len, 0, DYNAMIC_TYPE_OPENSSL);
         if (bio->mem == NULL) {
-            XFREE(bio, 0, DYNAMIC_TYPE_OPENSSL);
+            wolfSSL_BIO_free(bio);
             return NULL;
         }
         if (bio->mem_buf != NULL) {
@@ -21136,6 +21137,9 @@ int wolfSSL_RAND_egd(const char* nm)
 #endif
 
     if (nm == NULL) {
+    #ifdef WOLFSSL_SMALL_STACK
+        XFREE(buf, NULL, DYNAMIC_TYPE_TMP_BUFFER);
+    #endif
         return SSL_FATAL_ERROR;
     }
 
diff --git a/tests/api.c b/tests/api.c
index 5f5e948d9..93052d072 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -14933,6 +14933,7 @@ static void test_wolfSSL_PEM_read_bio(void)
     int  bytes;
     X509* x509;
     BIO*  bio = NULL;
+    BUF_MEM* buf;
 
     printf(testingFmt, "wolfSSL_PEM_read_bio()");
 
@@ -14944,8 +14945,10 @@ static void test_wolfSSL_PEM_read_bio(void)
     AssertNotNull(bio = BIO_new_mem_buf((void*)buff, bytes));
     AssertNotNull(x509 = PEM_read_bio_X509_AUX(bio, NULL, NULL, NULL));
     AssertIntEQ((int)BIO_set_fd(bio, 0, BIO_NOCLOSE), 1);
+    AssertIntEQ(SSL_SUCCESS, BIO_get_mem_ptr(bio, &buf));
 
     BIO_free(bio);
+    BUF_MEM_free(buf);
     X509_free(x509);
 
     printf(resultFmt, passed);