From afbe1b267032d2f4847e1d4523f57a6664c85e26 Mon Sep 17 00:00:00 2001
From: Dave Pate <dpate@yubico.com>
Date: Tue, 22 Jan 2019 07:53:22 -0800
Subject: [PATCH] lib: handle realloc failures safely

---
 lib/util.c | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/lib/util.c b/lib/util.c
index f14cc43..de4caf9 100644
--- a/lib/util.c
+++ b/lib/util.c
@@ -218,6 +218,7 @@ ykpiv_rc ykpiv_util_list_keys(ykpiv_state *state, uint8_t *key_count, ykpiv_key
   ykpiv_rc res = YKPIV_OK;
   ykpiv_key *pKey = NULL;
   uint8_t *pData = NULL;
+  uint8_t *pTemp = NULL;
   size_t cbData = 0;
   size_t offset = 0;
   uint8_t buf[CB_BUF_MAX];
@@ -282,10 +283,13 @@ ykpiv_rc ykpiv_util_list_keys(ykpiv_state *state, uint8_t *key_count, ykpiv_key
       cbRealloc = (sizeof(ykpiv_key) + cbBuf - 1) > (cbData - offset) ? MAX((sizeof(ykpiv_key) + cbBuf - 1) - (cbData - offset), CB_PAGE) : 0;
 
       if (0 != cbRealloc) {
-        if (NULL == (pData = _ykpiv_realloc(state, pData, cbData + cbRealloc))) {
+        if (!(pTemp = _ykpiv_realloc(state, pData, cbData + cbRealloc))) {
+          /* realloc failed, pData will be freed in cleanup */
           res = YKPIV_MEMORY_ERROR;
           goto Cleanup;
         }
+        pData = pTemp;
+        pTemp = NULL;
       }
 
       cbData += cbRealloc;
@@ -556,6 +560,7 @@ ykpiv_rc ykpiv_util_read_msroots(ykpiv_state *state, uint8_t **data, size_t *dat
   int object_id = 0;
   uint8_t tag = 0;
   uint8_t *pData = NULL;
+  uint8_t *pTemp = NULL;
   size_t cbData = 0;
   size_t cbRealloc = 0;
   size_t offset = 0;
@@ -606,10 +611,13 @@ ykpiv_rc ykpiv_util_read_msroots(ykpiv_state *state, uint8_t **data, size_t *dat
     cbRealloc = len > (cbData - offset) ? len - (cbData - offset) : 0;
 
     if (0 != cbRealloc) {
-      if (NULL == (pData = _ykpiv_realloc(state, pData, cbData + cbRealloc))) {
+      if (!(pTemp = _ykpiv_realloc(state, pData, cbData + cbRealloc))) {
+        /* realloc failed, pData will be freed in cleanup */
         res = YKPIV_MEMORY_ERROR;
         goto Cleanup;
       }
+      pData = pTemp;
+      pTemp = NULL;
     }
 
     cbData += cbRealloc;