From accfd80c375d1c313241bbfe18b22ec7f62dcd19 Mon Sep 17 00:00:00 2001
From: Yann Collet <yann.collet.73@gmail.com>
Date: Mon, 15 Feb 2016 19:33:16 +0100
Subject: [PATCH] Added : option `-t` to test compressed files integrity

---
 lib/zstd_internal.h |  1 -
 lib/zstd_static.h   |  6 ++++++
 programs/fileio.c   | 31 ++++++++++++++++++++-----------
 programs/fileio.h   |  1 -
 programs/zstdcli.c  | 12 ++++++++----
 5 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/lib/zstd_internal.h b/lib/zstd_internal.h
index 318fd55f..26fc8578 100644
--- a/lib/zstd_internal.h
+++ b/lib/zstd_internal.h
@@ -50,7 +50,6 @@
 /*-*************************************
 *  Common constants
 ***************************************/
-#define ZSTD_MAGICNUMBER 0xFD2FB525   /* v0.5 */
 #define ZSTD_DICT_MAGIC  0xEC30A435
 
 #define KB *(1 <<10)
diff --git a/lib/zstd_static.h b/lib/zstd_static.h
index 258327b1..61216535 100644
--- a/lib/zstd_static.h
+++ b/lib/zstd_static.h
@@ -48,6 +48,12 @@ extern "C" {
 #include "mem.h"
 
 
+/*-*************************************
+*  Constants
+***************************************/
+#define ZSTD_MAGICNUMBER 0xFD2FB525   /* v0.5 */
+
+
 /*-*************************************
 *  Types
 ***************************************/
diff --git a/programs/fileio.c b/programs/fileio.c
index 003e1bd5..e2ad9e84 100644
--- a/programs/fileio.c
+++ b/programs/fileio.c
@@ -587,7 +587,10 @@ static int FIO_decompressSrcFile(dRess_t ress, const char* srcFileName)
             continue;
         }
 #endif   /* ZSTD_LEGACY_SUPPORT */
-
+        if (MEM_readLE32(ress.srcBuffer) !=  ZSTD_MAGICNUMBER) {
+            DISPLAYLEVEL(1, "zstd: %s: not in zstd format \n", srcFileName);
+            return 1;
+        }
         filesize += FIO_decompressFrame(ress, dstFile, srcFile, toRead);
     }
 
@@ -648,13 +651,24 @@ int FIO_decompressMultipleFilenames(const char** srcNamesTable, unsigned nbFiles
 	if (dstFileName==NULL) EXM_THROW(70, "not enough memory for dstFileName");
     ress = FIO_createDResources(dictFileName);
 
-    if (suffix) {
+    if (!strcmp(suffix, stdoutmark) || !strcmp(suffix, nulmark)) {
+        ress.dstFile = FIO_openDstFile(suffix);
+        if (ress.dstFile == 0) EXM_THROW(71, "cannot open %s", suffix);
+        for (u=0; u<nbFiles; u++)
+            missingFiles += FIO_decompressSrcFile(ress, srcNamesTable[u]);
+        if (fclose(ress.dstFile)) EXM_THROW(39, "Write error : cannot properly close %s", stdoutmark);
+    } else {
         for (u=0; u<nbFiles; u++) {   /* create dstFileName */
             const char* srcFileName = srcNamesTable[u];
             size_t sfnSize = strlen(srcFileName);
             const char* suffixPtr = srcFileName + sfnSize - suffixSize;
-            if (dfnSize <= sfnSize-suffixSize+1) { free(dstFileName); dfnSize = sfnSize + 20; dstFileName = (char*)malloc(dfnSize); if (dstFileName==NULL) EXM_THROW(71, "not enough memory for dstFileName"); }
-            if (sfnSize <= suffixSize  ||  strcmp(suffixPtr, suffix) != 0) {
+            if (dfnSize+suffixSize <= sfnSize+1) {
+                free(dstFileName);
+                dfnSize = sfnSize + 20;
+                dstFileName = (char*)malloc(dfnSize);
+                if (dstFileName==NULL) EXM_THROW(71, "not enough memory for dstFileName");
+            }
+            if (sfnSize <= suffixSize || strcmp(suffixPtr, suffix) != 0) {
                 DISPLAYLEVEL(1, "zstd: %s: unknown suffix (%4s expected) -- ignored \n", srcFileName, suffix);
                 skippedFiles++;
                 continue;
@@ -663,15 +677,10 @@ int FIO_decompressMultipleFilenames(const char** srcNamesTable, unsigned nbFiles
             dstFileName[sfnSize-suffixSize] = '\0';
 
             missingFiles += FIO_decompressFile_extRess(ress, dstFileName, srcFileName);
-        }
-    } else {
-        ress.dstFile = stdout;
-        for (u=0; u<nbFiles; u++)
-            missingFiles += FIO_decompressSrcFile(ress, srcNamesTable[u]);
-        if (fclose(ress.dstFile)) EXM_THROW(39, "Write error : cannot properly close %s", stdoutmark);
-    }
+    }   }
 
     FIO_freeDResources(ress);
     free(dstFileName);
     return missingFiles + skippedFiles;
 }
+
diff --git a/programs/fileio.h b/programs/fileio.h
index caae8768..ee3cf227 100644
--- a/programs/fileio.h
+++ b/programs/fileio.h
@@ -64,7 +64,6 @@ int FIO_decompressFilename (const char* outfilename, const char* infilename, con
 *  Multiple File functions
 ***************************************/
 /** FIO_compressMultipleFilenames() :
-    if `suffix == NULL`, output is stdout.
     @return : nb of missing files */
 int FIO_compressMultipleFilenames(const char** srcNamesTable, unsigned nbFiles,
                                   const char* suffix,
diff --git a/programs/zstdcli.c b/programs/zstdcli.c
index cd285f51..7a6dae99 100644
--- a/programs/zstdcli.c
+++ b/programs/zstdcli.c
@@ -115,7 +115,7 @@ static int usage(const char* programName)
     DISPLAY( " -#     : # compression level (1-%u, default:1) \n", ZSTD_maxCLevel());
     DISPLAY( " -d     : decompression \n");
     DISPLAY( " -D file: use `file` as Dictionary \n");
-    DISPLAY( " -o file: result stored into `file` (only possible if 1 input file) \n");
+    DISPLAY( " -o file: result stored into `file` (only if 1 input file) \n");
     DISPLAY( " -f     : overwrite output without prompting \n");
     DISPLAY( " -h/-H  : display help/long help and exit\n");
     return 0;
@@ -128,6 +128,7 @@ static int usage_advanced(const char* programName)
     DISPLAY( "\n");
     DISPLAY( "Advanced arguments :\n");
     DISPLAY( " -V     : display Version number and exit\n");
+    DISPLAY( " -t     : test compressed file integrity \n");
     DISPLAY( " -v     : verbose mode\n");
     DISPLAY( " -q     : suppress warnings; specify twice to suppress errors too\n");
     DISPLAY( " -c     : force write to standard output, even if it is the console\n");
@@ -266,6 +267,9 @@ int main(int argCount, const char** argv)
                     /* keep source file (default anyway, so useless; for gzip/xz compatibility) */
                 case 'k': argument++; break;
 
+                    /* test compressed file */
+                case 't': decode=1; outFileName=nulmark; FIO_overwriteMode(); argument++; break;
+
                     /* dictionary name */
                 case 'o': nextArgumentIsOutFileName=1; argument++; break;
 
@@ -380,7 +384,7 @@ int main(int argCount, const char** argv)
     if (outFileName && !strcmp(outFileName, stdoutmark) && IS_CONSOLE(stdout) && !forceStdout) return badusage(programName);
 
     /* user-selected output filename, only possible with a single file */
-    if (outFileName && strcmp(outFileName,stdoutmark) && (filenameIdx>1)) {
+    if (outFileName && strcmp(outFileName,stdoutmark) && strcmp(outFileName,nulmark) && (filenameIdx>1)) {
         DISPLAY("Too many files (%u) on the command line. \n", filenameIdx);
         return filenameIdx;
     }
@@ -395,12 +399,12 @@ int main(int argCount, const char** argv)
       if (filenameIdx==1 && outFileName)
         operationResult = FIO_decompressFilename(outFileName, filenameTable[0], dictFileName);
       else
-        operationResult = FIO_decompressMultipleFilenames(filenameTable, filenameIdx, forceStdout ? NULL : ZSTD_EXTENSION, dictFileName);
+        operationResult = FIO_decompressMultipleFilenames(filenameTable, filenameIdx, outFileName ? outFileName : ZSTD_EXTENSION, dictFileName);
     } else {  /* compression */
         if (filenameIdx==1 && outFileName)
           operationResult = FIO_compressFilename(outFileName, filenameTable[0], dictFileName, cLevel);
         else
-          operationResult = FIO_compressMultipleFilenames(filenameTable, filenameIdx, forceStdout ? NULL : ZSTD_EXTENSION, dictFileName, cLevel);
+          operationResult = FIO_compressMultipleFilenames(filenameTable, filenameIdx, outFileName ? outFileName : ZSTD_EXTENSION, dictFileName, cLevel);
     }
 
 _end: