From 4dc604cab8f5d5e89454d2afeffe32f2001c41aa Mon Sep 17 00:00:00 2001
From: Sen Huang <senhuang96@fb.com>
Date: Wed, 2 Oct 2019 18:34:42 -0400
Subject: [PATCH] Addressed comments on malloc

---
 programs/fileio.c  |  6 ------
 programs/util.c    | 12 ++++++++++--
 programs/zstdcli.c |  5 +++++
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/programs/fileio.c b/programs/fileio.c
index c0300f6a..6006af01 100644
--- a/programs/fileio.c
+++ b/programs/fileio.c
@@ -1453,9 +1453,6 @@ int FIO_compressMultipleFilenames(FIO_prefs_t* const prefs, const char** inFileN
     }   }
 
     FIO_freeCResources(ress);
-    if (dstFileNamesTable)
-        UTIL_freeDestinationFilenameTable(dstFileNamesTable, nbFiles);
-        
     return error;
 }
 
@@ -2276,9 +2273,6 @@ FIO_decompressMultipleFilenames(FIO_prefs_t* const prefs,
     }
 
     FIO_freeDResources(ress);
-    if (outDirName)
-        UTIL_freeDestinationFilenameTable(dstFileNamesTable, nbFiles);
-
     return error;
 }
 
diff --git a/programs/util.c b/programs/util.c
index 83fa70ba..c12a3188 100644
--- a/programs/util.c
+++ b/programs/util.c
@@ -140,14 +140,22 @@ void UTIL_createDestinationDirTable(char** dstFilenameTable, const char** filena
         filenameBegin = strrchr(filenameTable[u], c[0]);
         if (filenameBegin == NULL) {
             filename = (char*) malloc((strlen(filenameTable[u])+1) * sizeof(char));
+            if (!filename) {
+                UTIL_DISPLAYLEVEL(1, "Unable to allocate space for filename str\n");
+                continue;
+            }
             strcpy(filename, filenameTable[u]);
         } else {
-            filename = (char*) malloc((strlen(filenameBegin+1)) * sizeof(char));
+            filename = (char*) malloc((strlen(filenameBegin+1)+1) * sizeof(char));
+            if (!filename) {
+                UTIL_DISPLAYLEVEL(1, "Unable to allocate space for filename str\n");
+                continue;
+            }
             strcpy(filename, filenameBegin+1);
         }
 
         finalPathLen += strlen(filename);
-        dstFilenameTable[u] = (char*) malloc((finalPathLen+2) * sizeof(char));
+        dstFilenameTable[u] = (char*) malloc((finalPathLen+3) * sizeof(char));
         if (!dstFilenameTable[u]) {
             UTIL_DISPLAYLEVEL(1, "Unable to allocate space for file destination str\n");
             free(filename);
diff --git a/programs/zstdcli.c b/programs/zstdcli.c
index d9e1be66..91e67a2d 100644
--- a/programs/zstdcli.c
+++ b/programs/zstdcli.c
@@ -1192,6 +1192,9 @@ int main(int argCount, const char* argv[])
           operationResult = FIO_compressFilename(prefs, outFileName, filenameTable[0], dictFileName, cLevel, compressionParams);
         else
           operationResult = FIO_compressMultipleFilenames(prefs, filenameTable, outDirName, dstFilenameTable, filenameIdx, outFileName, suffix, dictFileName, cLevel, compressionParams);
+        
+        if (dstFilenameTable)
+            UTIL_freeDestinationFilenameTable(dstFilenameTable, filenameIdx);
 #else
         (void)suffix; (void)adapt; (void)rsyncable; (void)ultra; (void)cLevel; (void)ldmFlag; (void)literalCompressionMode; (void)targetCBlockSize; (void)streamSrcSize; (void)srcSizeHint; /* not used when ZSTD_NOCOMPRESS set */
         DISPLAY("Compression not supported \n");
@@ -1224,6 +1227,8 @@ int main(int argCount, const char* argv[])
             operationResult = FIO_decompressFilename(prefs, outFileName, filenameTable[0], dictFileName);
         else
             operationResult = FIO_decompressMultipleFilenames(prefs, filenameTable, filenameIdx, outDirName, dstFilenameTable, outFileName, dictFileName);
+        if (dstFilenameTable)
+            UTIL_freeDestinationFilenameTable(dstFilenameTable, filenameIdx);
 #else
         DISPLAY("Decompression not supported \n");
 #endif