xref: /freebsd/sys/contrib/zstd/programs/util.c (revision 5ff13fbc199bdf5f0572845351c68ee5ca828e71)
1 /*
2  * Copyright (c) Przemyslaw Skibinski, Yann Collet, Facebook, Inc.
3  * All rights reserved.
4  *
5  * This source code is licensed under both the BSD-style license (found in the
6  * LICENSE file in the root directory of this source tree) and the GPLv2 (found
7  * in the COPYING file in the root directory of this source tree).
8  * You may select, at your option, one of the above-listed licenses.
9  */
10 
11 #if defined (__cplusplus)
12 extern "C" {
13 #endif
14 
15 
16 /*-****************************************
17 *  Dependencies
18 ******************************************/
19 #include "util.h"       /* note : ensure that platform.h is included first ! */
20 #include <stdlib.h>     /* malloc, realloc, free */
21 #include <stdio.h>      /* fprintf */
22 #include <time.h>       /* clock_t, clock, CLOCKS_PER_SEC, nanosleep */
23 #include <errno.h>
24 #include <assert.h>
25 
26 #if defined(_WIN32)
27 #  include <sys/utime.h>  /* utime */
28 #  include <io.h>         /* _chmod */
29 #else
30 #  include <unistd.h>     /* chown, stat */
31 #  if PLATFORM_POSIX_VERSION < 200809L || !defined(st_mtime)
32 #    include <utime.h>    /* utime */
33 #  else
34 #    include <fcntl.h>    /* AT_FDCWD */
35 #    include <sys/stat.h> /* utimensat */
36 #  endif
37 #endif
38 
39 #if defined(_MSC_VER) || defined(__MINGW32__) || defined (__MSVCRT__)
40 #include <direct.h>     /* needed for _mkdir in windows */
41 #endif
42 
43 #if defined(__linux__) || (PLATFORM_POSIX_VERSION >= 200112L)  /* opendir, readdir require POSIX.1-2001 */
44 #  include <dirent.h>       /* opendir, readdir */
45 #  include <string.h>       /* strerror, memcpy */
46 #endif /* #ifdef _WIN32 */
47 
48 /*-****************************************
49 *  Internal Macros
50 ******************************************/
51 
52 /* CONTROL is almost like an assert(), but is never disabled.
53  * It's designed for failures that may happen rarely,
54  * but we don't want to maintain a specific error code path for them,
55  * such as a malloc() returning NULL for example.
56  * Since it's always active, this macro can trigger side effects.
57  */
58 #define CONTROL(c)  {         \
59     if (!(c)) {               \
60         UTIL_DISPLAYLEVEL(1, "Error : %s, %i : %s",  \
61                           __FILE__, __LINE__, #c);   \
62         exit(1);              \
63 }   }
64 
65 /* console log */
66 #define UTIL_DISPLAY(...)         fprintf(stderr, __VA_ARGS__)
67 #define UTIL_DISPLAYLEVEL(l, ...) { if (g_utilDisplayLevel>=l) { UTIL_DISPLAY(__VA_ARGS__); } }
68 
69 /* A modified version of realloc().
70  * If UTIL_realloc() fails the original block is freed.
71  */
UTIL_realloc(void * ptr,size_t size)72 UTIL_STATIC void* UTIL_realloc(void *ptr, size_t size)
73 {
74     void *newptr = realloc(ptr, size);
75     if (newptr) return newptr;
76     free(ptr);
77     return NULL;
78 }
79 
80 #if defined(_MSC_VER)
81     #define chmod _chmod
82 #endif
83 
84 
85 /*-****************************************
86 *  Console log
87 ******************************************/
88 int g_utilDisplayLevel;
89 
UTIL_requireUserConfirmation(const char * prompt,const char * abortMsg,const char * acceptableLetters,int hasStdinInput)90 int UTIL_requireUserConfirmation(const char* prompt, const char* abortMsg,
91                                  const char* acceptableLetters, int hasStdinInput) {
92     int ch, result;
93 
94     if (hasStdinInput) {
95         UTIL_DISPLAY("stdin is an input - not proceeding.\n");
96         return 1;
97     }
98 
99     UTIL_DISPLAY("%s", prompt);
100     ch = getchar();
101     result = 0;
102     if (strchr(acceptableLetters, ch) == NULL) {
103         UTIL_DISPLAY("%s", abortMsg);
104         result = 1;
105     }
106     /* flush the rest */
107     while ((ch!=EOF) && (ch!='\n'))
108         ch = getchar();
109     return result;
110 }
111 
112 
113 /*-*************************************
114 *  Constants
115 ***************************************/
116 #define LIST_SIZE_INCREASE   (8*1024)
117 #define MAX_FILE_OF_FILE_NAMES_SIZE (1<<20)*50
118 
119 
120 /*-*************************************
121 *  Functions
122 ***************************************/
123 
UTIL_stat(const char * filename,stat_t * statbuf)124 int UTIL_stat(const char* filename, stat_t* statbuf)
125 {
126 #if defined(_MSC_VER)
127     return !_stat64(filename, statbuf);
128 #elif defined(__MINGW32__) && defined (__MSVCRT__)
129     return !_stati64(filename, statbuf);
130 #else
131     return !stat(filename, statbuf);
132 #endif
133 }
134 
UTIL_isRegularFile(const char * infilename)135 int UTIL_isRegularFile(const char* infilename)
136 {
137     stat_t statbuf;
138     return UTIL_stat(infilename, &statbuf) && UTIL_isRegularFileStat(&statbuf);
139 }
140 
UTIL_isRegularFileStat(const stat_t * statbuf)141 int UTIL_isRegularFileStat(const stat_t* statbuf)
142 {
143 #if defined(_MSC_VER)
144     return (statbuf->st_mode & S_IFREG) != 0;
145 #else
146     return S_ISREG(statbuf->st_mode) != 0;
147 #endif
148 }
149 
150 /* like chmod, but avoid changing permission of /dev/null */
UTIL_chmod(char const * filename,const stat_t * statbuf,mode_t permissions)151 int UTIL_chmod(char const* filename, const stat_t* statbuf, mode_t permissions)
152 {
153     stat_t localStatBuf;
154     if (statbuf == NULL) {
155         if (!UTIL_stat(filename, &localStatBuf)) return 0;
156         statbuf = &localStatBuf;
157     }
158     if (!UTIL_isRegularFileStat(statbuf)) return 0; /* pretend success, but don't change anything */
159     return chmod(filename, permissions);
160 }
161 
162 /* set access and modification times */
UTIL_utime(const char * filename,const stat_t * statbuf)163 int UTIL_utime(const char* filename, const stat_t *statbuf)
164 {
165     int ret;
166     /* We check that st_mtime is a macro here in order to give us confidence
167      * that struct stat has a struct timespec st_mtim member. We need this
168      * check because there are some platforms that claim to be POSIX 2008
169      * compliant but which do not have st_mtim... */
170 #if (PLATFORM_POSIX_VERSION >= 200809L) && defined(st_mtime)
171     /* (atime, mtime) */
172     struct timespec timebuf[2] = { {0, UTIME_NOW} };
173     timebuf[1] = statbuf->st_mtim;
174     ret = utimensat(AT_FDCWD, filename, timebuf, 0);
175 #else
176     struct utimbuf timebuf;
177     timebuf.actime = time(NULL);
178     timebuf.modtime = statbuf->st_mtime;
179     ret = utime(filename, &timebuf);
180 #endif
181     errno = 0;
182     return ret;
183 }
184 
UTIL_setFileStat(const char * filename,const stat_t * statbuf)185 int UTIL_setFileStat(const char *filename, const stat_t *statbuf)
186 {
187     int res = 0;
188 
189     stat_t curStatBuf;
190     if (!UTIL_stat(filename, &curStatBuf) || !UTIL_isRegularFileStat(&curStatBuf))
191         return -1;
192 
193     /* set access and modification times */
194     res += UTIL_utime(filename, statbuf);
195 
196 #if !defined(_WIN32)
197     res += chown(filename, statbuf->st_uid, statbuf->st_gid);  /* Copy ownership */
198 #endif
199 
200     res += UTIL_chmod(filename, &curStatBuf, statbuf->st_mode & 07777);  /* Copy file permissions */
201 
202     errno = 0;
203     return -res; /* number of errors is returned */
204 }
205 
UTIL_isDirectory(const char * infilename)206 int UTIL_isDirectory(const char* infilename)
207 {
208     stat_t statbuf;
209     return UTIL_stat(infilename, &statbuf) && UTIL_isDirectoryStat(&statbuf);
210 }
211 
UTIL_isDirectoryStat(const stat_t * statbuf)212 int UTIL_isDirectoryStat(const stat_t* statbuf)
213 {
214 #if defined(_MSC_VER)
215     return (statbuf->st_mode & _S_IFDIR) != 0;
216 #else
217     return S_ISDIR(statbuf->st_mode) != 0;
218 #endif
219 }
220 
UTIL_compareStr(const void * p1,const void * p2)221 int UTIL_compareStr(const void *p1, const void *p2) {
222     return strcmp(* (char * const *) p1, * (char * const *) p2);
223 }
224 
UTIL_isSameFile(const char * fName1,const char * fName2)225 int UTIL_isSameFile(const char* fName1, const char* fName2)
226 {
227     assert(fName1 != NULL); assert(fName2 != NULL);
228 #if defined(_MSC_VER) || defined(_WIN32)
229     /* note : Visual does not support file identification by inode.
230      *        inode does not work on Windows, even with a posix layer, like msys2.
231      *        The following work-around is limited to detecting exact name repetition only,
232      *        aka `filename` is considered different from `subdir/../filename` */
233     return !strcmp(fName1, fName2);
234 #else
235     {   stat_t file1Stat;
236         stat_t file2Stat;
237         return UTIL_stat(fName1, &file1Stat)
238             && UTIL_stat(fName2, &file2Stat)
239             && (file1Stat.st_dev == file2Stat.st_dev)
240             && (file1Stat.st_ino == file2Stat.st_ino);
241     }
242 #endif
243 }
244 
245 /* UTIL_isFIFO : distinguish named pipes */
UTIL_isFIFO(const char * infilename)246 int UTIL_isFIFO(const char* infilename)
247 {
248 /* macro guards, as defined in : https://linux.die.net/man/2/lstat */
249 #if PLATFORM_POSIX_VERSION >= 200112L
250     stat_t statbuf;
251     if (UTIL_stat(infilename, &statbuf) && UTIL_isFIFOStat(&statbuf)) return 1;
252 #endif
253     (void)infilename;
254     return 0;
255 }
256 
257 /* UTIL_isFIFO : distinguish named pipes */
UTIL_isFIFOStat(const stat_t * statbuf)258 int UTIL_isFIFOStat(const stat_t* statbuf)
259 {
260 /* macro guards, as defined in : https://linux.die.net/man/2/lstat */
261 #if PLATFORM_POSIX_VERSION >= 200112L
262     if (S_ISFIFO(statbuf->st_mode)) return 1;
263 #endif
264     (void)statbuf;
265     return 0;
266 }
267 
268 /* UTIL_isBlockDevStat : distinguish named pipes */
UTIL_isBlockDevStat(const stat_t * statbuf)269 int UTIL_isBlockDevStat(const stat_t* statbuf)
270 {
271 /* macro guards, as defined in : https://linux.die.net/man/2/lstat */
272 #if PLATFORM_POSIX_VERSION >= 200112L
273     if (S_ISBLK(statbuf->st_mode)) return 1;
274 #endif
275     (void)statbuf;
276     return 0;
277 }
278 
UTIL_isLink(const char * infilename)279 int UTIL_isLink(const char* infilename)
280 {
281 /* macro guards, as defined in : https://linux.die.net/man/2/lstat */
282 #if PLATFORM_POSIX_VERSION >= 200112L
283     stat_t statbuf;
284     int const r = lstat(infilename, &statbuf);
285     if (!r && S_ISLNK(statbuf.st_mode)) return 1;
286 #endif
287     (void)infilename;
288     return 0;
289 }
290 
UTIL_getFileSize(const char * infilename)291 U64 UTIL_getFileSize(const char* infilename)
292 {
293     stat_t statbuf;
294     if (!UTIL_stat(infilename, &statbuf)) return UTIL_FILESIZE_UNKNOWN;
295     return UTIL_getFileSizeStat(&statbuf);
296 }
297 
UTIL_getFileSizeStat(const stat_t * statbuf)298 U64 UTIL_getFileSizeStat(const stat_t* statbuf)
299 {
300     if (!UTIL_isRegularFileStat(statbuf)) return UTIL_FILESIZE_UNKNOWN;
301 #if defined(_MSC_VER)
302     if (!(statbuf->st_mode & S_IFREG)) return UTIL_FILESIZE_UNKNOWN;
303 #elif defined(__MINGW32__) && defined (__MSVCRT__)
304     if (!(statbuf->st_mode & S_IFREG)) return UTIL_FILESIZE_UNKNOWN;
305 #else
306     if (!S_ISREG(statbuf->st_mode)) return UTIL_FILESIZE_UNKNOWN;
307 #endif
308     return (U64)statbuf->st_size;
309 }
310 
UTIL_makeHumanReadableSize(U64 size)311 UTIL_HumanReadableSize_t UTIL_makeHumanReadableSize(U64 size)
312 {
313     UTIL_HumanReadableSize_t hrs;
314 
315     if (g_utilDisplayLevel > 3) {
316         /* In verbose mode, do not scale sizes down, except in the case of
317          * values that exceed the integral precision of a double. */
318         if (size >= (1ull << 53)) {
319             hrs.value = (double)size / (1ull << 20);
320             hrs.suffix = " MiB";
321             /* At worst, a double representation of a maximal size will be
322              * accurate to better than tens of kilobytes. */
323             hrs.precision = 2;
324         } else {
325             hrs.value = (double)size;
326             hrs.suffix = " B";
327             hrs.precision = 0;
328         }
329     } else {
330         /* In regular mode, scale sizes down and use suffixes. */
331         if (size >= (1ull << 60)) {
332             hrs.value = (double)size / (1ull << 60);
333             hrs.suffix = " EiB";
334         } else if (size >= (1ull << 50)) {
335             hrs.value = (double)size / (1ull << 50);
336             hrs.suffix = " PiB";
337         } else if (size >= (1ull << 40)) {
338             hrs.value = (double)size / (1ull << 40);
339             hrs.suffix = " TiB";
340         } else if (size >= (1ull << 30)) {
341             hrs.value = (double)size / (1ull << 30);
342             hrs.suffix = " GiB";
343         } else if (size >= (1ull << 20)) {
344             hrs.value = (double)size / (1ull << 20);
345             hrs.suffix = " MiB";
346         } else if (size >= (1ull << 10)) {
347             hrs.value = (double)size / (1ull << 10);
348             hrs.suffix = " KiB";
349         } else {
350             hrs.value = (double)size;
351             hrs.suffix = " B";
352         }
353 
354         if (hrs.value >= 100 || (U64)hrs.value == size) {
355             hrs.precision = 0;
356         } else if (hrs.value >= 10) {
357             hrs.precision = 1;
358         } else if (hrs.value > 1) {
359             hrs.precision = 2;
360         } else {
361             hrs.precision = 3;
362         }
363     }
364 
365     return hrs;
366 }
367 
UTIL_getTotalFileSize(const char * const * fileNamesTable,unsigned nbFiles)368 U64 UTIL_getTotalFileSize(const char* const * fileNamesTable, unsigned nbFiles)
369 {
370     U64 total = 0;
371     unsigned n;
372     for (n=0; n<nbFiles; n++) {
373         U64 const size = UTIL_getFileSize(fileNamesTable[n]);
374         if (size == UTIL_FILESIZE_UNKNOWN) return UTIL_FILESIZE_UNKNOWN;
375         total += size;
376     }
377     return total;
378 }
379 
380 
381 /* condition : @file must be valid, and not have reached its end.
382  * @return : length of line written into @buf, ended with `\0` instead of '\n',
383  *           or 0, if there is no new line */
readLineFromFile(char * buf,size_t len,FILE * file)384 static size_t readLineFromFile(char* buf, size_t len, FILE* file)
385 {
386     assert(!feof(file));
387     if ( fgets(buf, (int) len, file) == NULL ) return 0;
388     {   size_t linelen = strlen(buf);
389         if (strlen(buf)==0) return 0;
390         if (buf[linelen-1] == '\n') linelen--;
391         buf[linelen] = '\0';
392         return linelen+1;
393     }
394 }
395 
396 /* Conditions :
397  *   size of @inputFileName file must be < @dstCapacity
398  *   @dst must be initialized
399  * @return : nb of lines
400  *       or -1 if there's an error
401  */
402 static int
readLinesFromFile(void * dst,size_t dstCapacity,const char * inputFileName)403 readLinesFromFile(void* dst, size_t dstCapacity,
404             const char* inputFileName)
405 {
406     int nbFiles = 0;
407     size_t pos = 0;
408     char* const buf = (char*)dst;
409     FILE* const inputFile = fopen(inputFileName, "r");
410 
411     assert(dst != NULL);
412 
413     if(!inputFile) {
414         if (g_utilDisplayLevel >= 1) perror("zstd:util:readLinesFromFile");
415         return -1;
416     }
417 
418     while ( !feof(inputFile) ) {
419         size_t const lineLength = readLineFromFile(buf+pos, dstCapacity-pos, inputFile);
420         if (lineLength == 0) break;
421         assert(pos + lineLength < dstCapacity);
422         pos += lineLength;
423         ++nbFiles;
424     }
425 
426     CONTROL( fclose(inputFile) == 0 );
427 
428     return nbFiles;
429 }
430 
431 /*Note: buf is not freed in case function successfully created table because filesTable->fileNames[0] = buf*/
432 FileNamesTable*
UTIL_createFileNamesTable_fromFileName(const char * inputFileName)433 UTIL_createFileNamesTable_fromFileName(const char* inputFileName)
434 {
435     size_t nbFiles = 0;
436     char* buf;
437     size_t bufSize;
438     size_t pos = 0;
439     stat_t statbuf;
440 
441     if (!UTIL_stat(inputFileName, &statbuf) || !UTIL_isRegularFileStat(&statbuf))
442         return NULL;
443 
444     {   U64 const inputFileSize = UTIL_getFileSizeStat(&statbuf);
445         if(inputFileSize > MAX_FILE_OF_FILE_NAMES_SIZE)
446             return NULL;
447         bufSize = (size_t)(inputFileSize + 1); /* (+1) to add '\0' at the end of last filename */
448     }
449 
450     buf = (char*) malloc(bufSize);
451     CONTROL( buf != NULL );
452 
453     {   int const ret_nbFiles = readLinesFromFile(buf, bufSize, inputFileName);
454 
455         if (ret_nbFiles <= 0) {
456           free(buf);
457           return NULL;
458         }
459         nbFiles = (size_t)ret_nbFiles;
460     }
461 
462     {   const char** filenamesTable = (const char**) malloc(nbFiles * sizeof(*filenamesTable));
463         CONTROL(filenamesTable != NULL);
464 
465         {   size_t fnb;
466             for (fnb = 0, pos = 0; fnb < nbFiles; fnb++) {
467                 filenamesTable[fnb] = buf+pos;
468                 pos += strlen(buf+pos)+1;  /* +1 for the finishing `\0` */
469         }   }
470         assert(pos <= bufSize);
471 
472         return UTIL_assembleFileNamesTable(filenamesTable, nbFiles, buf);
473     }
474 }
475 
476 static FileNamesTable*
UTIL_assembleFileNamesTable2(const char ** filenames,size_t tableSize,size_t tableCapacity,char * buf)477 UTIL_assembleFileNamesTable2(const char** filenames, size_t tableSize, size_t tableCapacity, char* buf)
478 {
479     FileNamesTable* const table = (FileNamesTable*) malloc(sizeof(*table));
480     CONTROL(table != NULL);
481     table->fileNames = filenames;
482     table->buf = buf;
483     table->tableSize = tableSize;
484     table->tableCapacity = tableCapacity;
485     return table;
486 }
487 
488 FileNamesTable*
UTIL_assembleFileNamesTable(const char ** filenames,size_t tableSize,char * buf)489 UTIL_assembleFileNamesTable(const char** filenames, size_t tableSize, char* buf)
490 {
491     return UTIL_assembleFileNamesTable2(filenames, tableSize, tableSize, buf);
492 }
493 
UTIL_freeFileNamesTable(FileNamesTable * table)494 void UTIL_freeFileNamesTable(FileNamesTable* table)
495 {
496     if (table==NULL) return;
497     free((void*)table->fileNames);
498     free(table->buf);
499     free(table);
500 }
501 
UTIL_allocateFileNamesTable(size_t tableSize)502 FileNamesTable* UTIL_allocateFileNamesTable(size_t tableSize)
503 {
504     const char** const fnTable = (const char**)malloc(tableSize * sizeof(*fnTable));
505     FileNamesTable* fnt;
506     if (fnTable==NULL) return NULL;
507     fnt = UTIL_assembleFileNamesTable(fnTable, tableSize, NULL);
508     fnt->tableSize = 0;   /* the table is empty */
509     return fnt;
510 }
511 
UTIL_refFilename(FileNamesTable * fnt,const char * filename)512 void UTIL_refFilename(FileNamesTable* fnt, const char* filename)
513 {
514     assert(fnt->tableSize < fnt->tableCapacity);
515     fnt->fileNames[fnt->tableSize] = filename;
516     fnt->tableSize++;
517 }
518 
getTotalTableSize(FileNamesTable * table)519 static size_t getTotalTableSize(FileNamesTable* table)
520 {
521     size_t fnb = 0, totalSize = 0;
522     for(fnb = 0 ; fnb < table->tableSize && table->fileNames[fnb] ; ++fnb) {
523         totalSize += strlen(table->fileNames[fnb]) + 1; /* +1 to add '\0' at the end of each fileName */
524     }
525     return totalSize;
526 }
527 
528 FileNamesTable*
UTIL_mergeFileNamesTable(FileNamesTable * table1,FileNamesTable * table2)529 UTIL_mergeFileNamesTable(FileNamesTable* table1, FileNamesTable* table2)
530 {
531     unsigned newTableIdx = 0;
532     size_t pos = 0;
533     size_t newTotalTableSize;
534     char* buf;
535 
536     FileNamesTable* const newTable = UTIL_assembleFileNamesTable(NULL, 0, NULL);
537     CONTROL( newTable != NULL );
538 
539     newTotalTableSize = getTotalTableSize(table1) + getTotalTableSize(table2);
540 
541     buf = (char*) calloc(newTotalTableSize, sizeof(*buf));
542     CONTROL ( buf != NULL );
543 
544     newTable->buf = buf;
545     newTable->tableSize = table1->tableSize + table2->tableSize;
546     newTable->fileNames = (const char **) calloc(newTable->tableSize, sizeof(*(newTable->fileNames)));
547     CONTROL ( newTable->fileNames != NULL );
548 
549     {   unsigned idx1;
550         for( idx1=0 ; (idx1 < table1->tableSize) && table1->fileNames[idx1] && (pos < newTotalTableSize); ++idx1, ++newTableIdx) {
551             size_t const curLen = strlen(table1->fileNames[idx1]);
552             memcpy(buf+pos, table1->fileNames[idx1], curLen);
553             assert(newTableIdx <= newTable->tableSize);
554             newTable->fileNames[newTableIdx] = buf+pos;
555             pos += curLen+1;
556     }   }
557 
558     {   unsigned idx2;
559         for( idx2=0 ; (idx2 < table2->tableSize) && table2->fileNames[idx2] && (pos < newTotalTableSize) ; ++idx2, ++newTableIdx) {
560             size_t const curLen = strlen(table2->fileNames[idx2]);
561             memcpy(buf+pos, table2->fileNames[idx2], curLen);
562             assert(newTableIdx <= newTable->tableSize);
563             newTable->fileNames[newTableIdx] = buf+pos;
564             pos += curLen+1;
565     }   }
566     assert(pos <= newTotalTableSize);
567     newTable->tableSize = newTableIdx;
568 
569     UTIL_freeFileNamesTable(table1);
570     UTIL_freeFileNamesTable(table2);
571 
572     return newTable;
573 }
574 
575 #ifdef _WIN32
UTIL_prepareFileList(const char * dirName,char ** bufStart,size_t * pos,char ** bufEnd,int followLinks)576 static int UTIL_prepareFileList(const char* dirName,
577                                 char** bufStart, size_t* pos,
578                                 char** bufEnd, int followLinks)
579 {
580     char* path;
581     size_t dirLength, pathLength;
582     int nbFiles = 0;
583     WIN32_FIND_DATAA cFile;
584     HANDLE hFile;
585 
586     dirLength = strlen(dirName);
587     path = (char*) malloc(dirLength + 3);
588     if (!path) return 0;
589 
590     memcpy(path, dirName, dirLength);
591     path[dirLength] = '\\';
592     path[dirLength+1] = '*';
593     path[dirLength+2] = 0;
594 
595     hFile=FindFirstFileA(path, &cFile);
596     if (hFile == INVALID_HANDLE_VALUE) {
597         UTIL_DISPLAYLEVEL(1, "Cannot open directory '%s'\n", dirName);
598         return 0;
599     }
600     free(path);
601 
602     do {
603         size_t const fnameLength = strlen(cFile.cFileName);
604         path = (char*) malloc(dirLength + fnameLength + 2);
605         if (!path) { FindClose(hFile); return 0; }
606         memcpy(path, dirName, dirLength);
607         path[dirLength] = '\\';
608         memcpy(path+dirLength+1, cFile.cFileName, fnameLength);
609         pathLength = dirLength+1+fnameLength;
610         path[pathLength] = 0;
611         if (cFile.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
612             if ( strcmp (cFile.cFileName, "..") == 0
613               || strcmp (cFile.cFileName, ".") == 0 )
614                 continue;
615             /* Recursively call "UTIL_prepareFileList" with the new path. */
616             nbFiles += UTIL_prepareFileList(path, bufStart, pos, bufEnd, followLinks);
617             if (*bufStart == NULL) { free(path); FindClose(hFile); return 0; }
618         } else if ( (cFile.dwFileAttributes & FILE_ATTRIBUTE_NORMAL)
619                  || (cFile.dwFileAttributes & FILE_ATTRIBUTE_ARCHIVE)
620                  || (cFile.dwFileAttributes & FILE_ATTRIBUTE_COMPRESSED) ) {
621             if (*bufStart + *pos + pathLength >= *bufEnd) {
622                 ptrdiff_t const newListSize = (*bufEnd - *bufStart) + LIST_SIZE_INCREASE;
623                 *bufStart = (char*)UTIL_realloc(*bufStart, newListSize);
624                 if (*bufStart == NULL) { free(path); FindClose(hFile); return 0; }
625                 *bufEnd = *bufStart + newListSize;
626             }
627             if (*bufStart + *pos + pathLength < *bufEnd) {
628                 memcpy(*bufStart + *pos, path, pathLength+1 /* include final \0 */);
629                 *pos += pathLength + 1;
630                 nbFiles++;
631         }   }
632         free(path);
633     } while (FindNextFileA(hFile, &cFile));
634 
635     FindClose(hFile);
636     return nbFiles;
637 }
638 
639 #elif defined(__linux__) || (PLATFORM_POSIX_VERSION >= 200112L)  /* opendir, readdir require POSIX.1-2001 */
640 
UTIL_prepareFileList(const char * dirName,char ** bufStart,size_t * pos,char ** bufEnd,int followLinks)641 static int UTIL_prepareFileList(const char *dirName,
642                                 char** bufStart, size_t* pos,
643                                 char** bufEnd, int followLinks)
644 {
645     DIR* dir;
646     struct dirent * entry;
647     size_t dirLength;
648     int nbFiles = 0;
649 
650     if (!(dir = opendir(dirName))) {
651         UTIL_DISPLAYLEVEL(1, "Cannot open directory '%s': %s\n", dirName, strerror(errno));
652         return 0;
653     }
654 
655     dirLength = strlen(dirName);
656     errno = 0;
657     while ((entry = readdir(dir)) != NULL) {
658         char* path;
659         size_t fnameLength, pathLength;
660         if (strcmp (entry->d_name, "..") == 0 ||
661             strcmp (entry->d_name, ".") == 0) continue;
662         fnameLength = strlen(entry->d_name);
663         path = (char*) malloc(dirLength + fnameLength + 2);
664         if (!path) { closedir(dir); return 0; }
665         memcpy(path, dirName, dirLength);
666 
667         path[dirLength] = '/';
668         memcpy(path+dirLength+1, entry->d_name, fnameLength);
669         pathLength = dirLength+1+fnameLength;
670         path[pathLength] = 0;
671 
672         if (!followLinks && UTIL_isLink(path)) {
673             UTIL_DISPLAYLEVEL(2, "Warning : %s is a symbolic link, ignoring\n", path);
674             free(path);
675             continue;
676         }
677 
678         if (UTIL_isDirectory(path)) {
679             nbFiles += UTIL_prepareFileList(path, bufStart, pos, bufEnd, followLinks);  /* Recursively call "UTIL_prepareFileList" with the new path. */
680             if (*bufStart == NULL) { free(path); closedir(dir); return 0; }
681         } else {
682             if (*bufStart + *pos + pathLength >= *bufEnd) {
683                 ptrdiff_t newListSize = (*bufEnd - *bufStart) + LIST_SIZE_INCREASE;
684                 assert(newListSize >= 0);
685                 *bufStart = (char*)UTIL_realloc(*bufStart, (size_t)newListSize);
686                 *bufEnd = *bufStart + newListSize;
687                 if (*bufStart == NULL) { free(path); closedir(dir); return 0; }
688             }
689             if (*bufStart + *pos + pathLength < *bufEnd) {
690                 memcpy(*bufStart + *pos, path, pathLength + 1);  /* with final \0 */
691                 *pos += pathLength + 1;
692                 nbFiles++;
693         }   }
694         free(path);
695         errno = 0; /* clear errno after UTIL_isDirectory, UTIL_prepareFileList */
696     }
697 
698     if (errno != 0) {
699         UTIL_DISPLAYLEVEL(1, "readdir(%s) error: %s \n", dirName, strerror(errno));
700         free(*bufStart);
701         *bufStart = NULL;
702     }
703     closedir(dir);
704     return nbFiles;
705 }
706 
707 #else
708 
UTIL_prepareFileList(const char * dirName,char ** bufStart,size_t * pos,char ** bufEnd,int followLinks)709 static int UTIL_prepareFileList(const char *dirName,
710                                 char** bufStart, size_t* pos,
711                                 char** bufEnd, int followLinks)
712 {
713     (void)bufStart; (void)bufEnd; (void)pos; (void)followLinks;
714     UTIL_DISPLAYLEVEL(1, "Directory %s ignored (compiled without _WIN32 or _POSIX_C_SOURCE) \n", dirName);
715     return 0;
716 }
717 
718 #endif /* #ifdef _WIN32 */
719 
UTIL_isCompressedFile(const char * inputName,const char * extensionList[])720 int UTIL_isCompressedFile(const char *inputName, const char *extensionList[])
721 {
722   const char* ext = UTIL_getFileExtension(inputName);
723   while(*extensionList!=NULL)
724   {
725     const int isCompressedExtension = strcmp(ext,*extensionList);
726     if(isCompressedExtension==0)
727       return 1;
728     ++extensionList;
729   }
730    return 0;
731 }
732 
733 /*Utility function to get file extension from file */
UTIL_getFileExtension(const char * infilename)734 const char* UTIL_getFileExtension(const char* infilename)
735 {
736    const char* extension = strrchr(infilename, '.');
737    if(!extension || extension==infilename) return "";
738    return extension;
739 }
740 
pathnameHas2Dots(const char * pathname)741 static int pathnameHas2Dots(const char *pathname)
742 {
743     /* We need to figure out whether any ".." present in the path is a whole
744      * path token, which is the case if it is bordered on both sides by either
745      * the beginning/end of the path or by a directory separator.
746      */
747     const char *needle = pathname;
748     while (1) {
749         needle = strstr(needle, "..");
750 
751         if (needle == NULL) {
752             return 0;
753         }
754 
755         if ((needle == pathname || needle[-1] == PATH_SEP)
756          && (needle[2] == '\0' || needle[2] == PATH_SEP)) {
757             return 1;
758         }
759 
760         /* increment so we search for the next match */
761         needle++;
762     };
763     return 0;
764 }
765 
isFileNameValidForMirroredOutput(const char * filename)766 static int isFileNameValidForMirroredOutput(const char *filename)
767 {
768     return !pathnameHas2Dots(filename);
769 }
770 
771 
772 #define DIR_DEFAULT_MODE 0755
getDirMode(const char * dirName)773 static mode_t getDirMode(const char *dirName)
774 {
775     stat_t st;
776     if (!UTIL_stat(dirName, &st)) {
777         UTIL_DISPLAY("zstd: failed to get DIR stats %s: %s\n", dirName, strerror(errno));
778         return DIR_DEFAULT_MODE;
779     }
780     if (!UTIL_isDirectoryStat(&st)) {
781         UTIL_DISPLAY("zstd: expected directory: %s\n", dirName);
782         return DIR_DEFAULT_MODE;
783     }
784     return st.st_mode;
785 }
786 
makeDir(const char * dir,mode_t mode)787 static int makeDir(const char *dir, mode_t mode)
788 {
789 #if defined(_MSC_VER) || defined(__MINGW32__) || defined (__MSVCRT__)
790     int ret = _mkdir(dir);
791     (void) mode;
792 #else
793     int ret = mkdir(dir, mode);
794 #endif
795     if (ret != 0) {
796         if (errno == EEXIST)
797             return 0;
798         UTIL_DISPLAY("zstd: failed to create DIR %s: %s\n", dir, strerror(errno));
799     }
800     return ret;
801 }
802 
803 /* this function requires a mutable input string */
convertPathnameToDirName(char * pathname)804 static void convertPathnameToDirName(char *pathname)
805 {
806     size_t len = 0;
807     char* pos = NULL;
808     /* get dir name from pathname similar to 'dirname()' */
809     assert(pathname != NULL);
810 
811     /* remove trailing '/' chars */
812     len = strlen(pathname);
813     assert(len > 0);
814     while (pathname[len] == PATH_SEP) {
815         pathname[len] = '\0';
816         len--;
817     }
818     if (len == 0) return;
819 
820     /* if input is a single file, return '.' instead. i.e.
821      * "xyz/abc/file.txt" => "xyz/abc"
822        "./file.txt"       => "."
823        "file.txt"         => "."
824      */
825     pos = strrchr(pathname, PATH_SEP);
826     if (pos == NULL) {
827         pathname[0] = '.';
828         pathname[1] = '\0';
829     } else {
830         *pos = '\0';
831     }
832 }
833 
834 /* pathname must be valid */
trimLeadingRootChar(const char * pathname)835 static const char* trimLeadingRootChar(const char *pathname)
836 {
837     assert(pathname != NULL);
838     if (pathname[0] == PATH_SEP)
839         return pathname + 1;
840     return pathname;
841 }
842 
843 /* pathname must be valid */
trimLeadingCurrentDirConst(const char * pathname)844 static const char* trimLeadingCurrentDirConst(const char *pathname)
845 {
846     assert(pathname != NULL);
847     if ((pathname[0] == '.') && (pathname[1] == PATH_SEP))
848         return pathname + 2;
849     return pathname;
850 }
851 
852 static char*
trimLeadingCurrentDir(char * pathname)853 trimLeadingCurrentDir(char *pathname)
854 {
855     /* 'union charunion' can do const-cast without compiler warning */
856     union charunion {
857         char *chr;
858         const char* cchr;
859     } ptr;
860     ptr.cchr = trimLeadingCurrentDirConst(pathname);
861     return ptr.chr;
862 }
863 
864 /* remove leading './' or '/' chars here */
trimPath(const char * pathname)865 static const char * trimPath(const char *pathname)
866 {
867     return trimLeadingRootChar(
868             trimLeadingCurrentDirConst(pathname));
869 }
870 
mallocAndJoin2Dir(const char * dir1,const char * dir2)871 static char* mallocAndJoin2Dir(const char *dir1, const char *dir2)
872 {
873     const size_t dir1Size = strlen(dir1);
874     const size_t dir2Size = strlen(dir2);
875     char *outDirBuffer, *buffer, trailingChar;
876 
877     assert(dir1 != NULL && dir2 != NULL);
878     outDirBuffer = (char *) malloc(dir1Size + dir2Size + 2);
879     CONTROL(outDirBuffer != NULL);
880 
881     memcpy(outDirBuffer, dir1, dir1Size);
882     outDirBuffer[dir1Size] = '\0';
883 
884     if (dir2[0] == '.')
885         return outDirBuffer;
886 
887     buffer = outDirBuffer + dir1Size;
888     trailingChar = *(buffer - 1);
889     if (trailingChar != PATH_SEP) {
890         *buffer = PATH_SEP;
891         buffer++;
892     }
893     memcpy(buffer, dir2, dir2Size);
894     buffer[dir2Size] = '\0';
895 
896     return outDirBuffer;
897 }
898 
899 /* this function will return NULL if input srcFileName is not valid name for mirrored output path */
UTIL_createMirroredDestDirName(const char * srcFileName,const char * outDirRootName)900 char* UTIL_createMirroredDestDirName(const char* srcFileName, const char* outDirRootName)
901 {
902     char* pathname = NULL;
903     if (!isFileNameValidForMirroredOutput(srcFileName))
904         return NULL;
905 
906     pathname = mallocAndJoin2Dir(outDirRootName, trimPath(srcFileName));
907 
908     convertPathnameToDirName(pathname);
909     return pathname;
910 }
911 
912 static int
mirrorSrcDir(char * srcDirName,const char * outDirName)913 mirrorSrcDir(char* srcDirName, const char* outDirName)
914 {
915     mode_t srcMode;
916     int status = 0;
917     char* newDir = mallocAndJoin2Dir(outDirName, trimPath(srcDirName));
918     if (!newDir)
919         return -ENOMEM;
920 
921     srcMode = getDirMode(srcDirName);
922     status = makeDir(newDir, srcMode);
923     free(newDir);
924     return status;
925 }
926 
927 static int
mirrorSrcDirRecursive(char * srcDirName,const char * outDirName)928 mirrorSrcDirRecursive(char* srcDirName, const char* outDirName)
929 {
930     int status = 0;
931     char* pp = trimLeadingCurrentDir(srcDirName);
932     char* sp = NULL;
933 
934     while ((sp = strchr(pp, PATH_SEP)) != NULL) {
935         if (sp != pp) {
936             *sp = '\0';
937             status = mirrorSrcDir(srcDirName, outDirName);
938             if (status != 0)
939                 return status;
940             *sp = PATH_SEP;
941         }
942         pp = sp + 1;
943     }
944     status = mirrorSrcDir(srcDirName, outDirName);
945     return status;
946 }
947 
948 static void
makeMirroredDestDirsWithSameSrcDirMode(char ** srcDirNames,unsigned nbFile,const char * outDirName)949 makeMirroredDestDirsWithSameSrcDirMode(char** srcDirNames, unsigned nbFile, const char* outDirName)
950 {
951     unsigned int i = 0;
952     for (i = 0; i < nbFile; i++)
953         mirrorSrcDirRecursive(srcDirNames[i], outDirName);
954 }
955 
956 static int
firstIsParentOrSameDirOfSecond(const char * firstDir,const char * secondDir)957 firstIsParentOrSameDirOfSecond(const char* firstDir, const char* secondDir)
958 {
959     size_t firstDirLen  = strlen(firstDir),
960            secondDirLen = strlen(secondDir);
961     return firstDirLen <= secondDirLen &&
962            (secondDir[firstDirLen] == PATH_SEP || secondDir[firstDirLen] == '\0') &&
963            0 == strncmp(firstDir, secondDir, firstDirLen);
964 }
965 
compareDir(const void * pathname1,const void * pathname2)966 static int compareDir(const void* pathname1, const void* pathname2) {
967     /* sort it after remove the leading '/'  or './'*/
968     const char* s1 = trimPath(*(char * const *) pathname1);
969     const char* s2 = trimPath(*(char * const *) pathname2);
970     return strcmp(s1, s2);
971 }
972 
973 static void
makeUniqueMirroredDestDirs(char ** srcDirNames,unsigned nbFile,const char * outDirName)974 makeUniqueMirroredDestDirs(char** srcDirNames, unsigned nbFile, const char* outDirName)
975 {
976     unsigned int i = 0, uniqueDirNr = 0;
977     char** uniqueDirNames = NULL;
978 
979     if (nbFile == 0)
980         return;
981 
982     uniqueDirNames = (char** ) malloc(nbFile * sizeof (char *));
983     CONTROL(uniqueDirNames != NULL);
984 
985     /* if dirs is "a/b/c" and "a/b/c/d", we only need call:
986      * we just need "a/b/c/d" */
987     qsort((void *)srcDirNames, nbFile, sizeof(char*), compareDir);
988 
989     uniqueDirNr = 1;
990     uniqueDirNames[uniqueDirNr - 1] = srcDirNames[0];
991     for (i = 1; i < nbFile; i++) {
992         char* prevDirName = srcDirNames[i - 1];
993         char* currDirName = srcDirNames[i];
994 
995         /* note: we always compare trimmed path, i.e.:
996          * src dir of "./foo" and "/foo" will be both saved into:
997          * "outDirName/foo/" */
998         if (!firstIsParentOrSameDirOfSecond(trimPath(prevDirName),
999                                             trimPath(currDirName)))
1000             uniqueDirNr++;
1001 
1002         /* we need maintain original src dir name instead of trimmed
1003          * dir, so we can retrieve the original src dir's mode_t */
1004         uniqueDirNames[uniqueDirNr - 1] = currDirName;
1005     }
1006 
1007     makeMirroredDestDirsWithSameSrcDirMode(uniqueDirNames, uniqueDirNr, outDirName);
1008 
1009     free(uniqueDirNames);
1010 }
1011 
1012 static void
makeMirroredDestDirs(char ** srcFileNames,unsigned nbFile,const char * outDirName)1013 makeMirroredDestDirs(char** srcFileNames, unsigned nbFile, const char* outDirName)
1014 {
1015     unsigned int i = 0;
1016     for (i = 0; i < nbFile; ++i)
1017         convertPathnameToDirName(srcFileNames[i]);
1018     makeUniqueMirroredDestDirs(srcFileNames, nbFile, outDirName);
1019 }
1020 
UTIL_mirrorSourceFilesDirectories(const char ** inFileNames,unsigned int nbFile,const char * outDirName)1021 void UTIL_mirrorSourceFilesDirectories(const char** inFileNames, unsigned int nbFile, const char* outDirName)
1022 {
1023     unsigned int i = 0, validFilenamesNr = 0;
1024     char** srcFileNames = (char **) malloc(nbFile * sizeof (char *));
1025     CONTROL(srcFileNames != NULL);
1026 
1027     /* check input filenames is valid */
1028     for (i = 0; i < nbFile; ++i) {
1029         if (isFileNameValidForMirroredOutput(inFileNames[i])) {
1030             char* fname = STRDUP(inFileNames[i]);
1031             CONTROL(fname != NULL);
1032             srcFileNames[validFilenamesNr++] = fname;
1033         }
1034     }
1035 
1036     if (validFilenamesNr > 0) {
1037         makeDir(outDirName, DIR_DEFAULT_MODE);
1038         makeMirroredDestDirs(srcFileNames, validFilenamesNr, outDirName);
1039     }
1040 
1041     for (i = 0; i < validFilenamesNr; i++)
1042         free(srcFileNames[i]);
1043     free(srcFileNames);
1044 }
1045 
1046 FileNamesTable*
UTIL_createExpandedFNT(const char * const * inputNames,size_t nbIfns,int followLinks)1047 UTIL_createExpandedFNT(const char* const* inputNames, size_t nbIfns, int followLinks)
1048 {
1049     unsigned nbFiles;
1050     char* buf = (char*)malloc(LIST_SIZE_INCREASE);
1051     char* bufend = buf + LIST_SIZE_INCREASE;
1052 
1053     if (!buf) return NULL;
1054 
1055     {   size_t ifnNb, pos;
1056         for (ifnNb=0, pos=0, nbFiles=0; ifnNb<nbIfns; ifnNb++) {
1057             if (!UTIL_isDirectory(inputNames[ifnNb])) {
1058                 size_t const len = strlen(inputNames[ifnNb]);
1059                 if (buf + pos + len >= bufend) {
1060                     ptrdiff_t newListSize = (bufend - buf) + LIST_SIZE_INCREASE;
1061                     assert(newListSize >= 0);
1062                     buf = (char*)UTIL_realloc(buf, (size_t)newListSize);
1063                     if (!buf) return NULL;
1064                     bufend = buf + newListSize;
1065                 }
1066                 if (buf + pos + len < bufend) {
1067                     memcpy(buf+pos, inputNames[ifnNb], len+1);  /* including final \0 */
1068                     pos += len + 1;
1069                     nbFiles++;
1070                 }
1071             } else {
1072                 nbFiles += (unsigned)UTIL_prepareFileList(inputNames[ifnNb], &buf, &pos, &bufend, followLinks);
1073                 if (buf == NULL) return NULL;
1074     }   }   }
1075 
1076     /* note : even if nbFiles==0, function returns a valid, though empty, FileNamesTable* object */
1077 
1078     {   size_t ifnNb, pos;
1079         size_t const fntCapacity = nbFiles + 1;  /* minimum 1, allows adding one reference, typically stdin */
1080         const char** const fileNamesTable = (const char**)malloc(fntCapacity * sizeof(*fileNamesTable));
1081         if (!fileNamesTable) { free(buf); return NULL; }
1082 
1083         for (ifnNb = 0, pos = 0; ifnNb < nbFiles; ifnNb++) {
1084             fileNamesTable[ifnNb] = buf + pos;
1085             if (buf + pos > bufend) { free(buf); free((void*)fileNamesTable); return NULL; }
1086             pos += strlen(fileNamesTable[ifnNb]) + 1;
1087         }
1088         return UTIL_assembleFileNamesTable2(fileNamesTable, nbFiles, fntCapacity, buf);
1089     }
1090 }
1091 
1092 
UTIL_expandFNT(FileNamesTable ** fnt,int followLinks)1093 void UTIL_expandFNT(FileNamesTable** fnt, int followLinks)
1094 {
1095     FileNamesTable* const newFNT = UTIL_createExpandedFNT((*fnt)->fileNames, (*fnt)->tableSize, followLinks);
1096     CONTROL(newFNT != NULL);
1097     UTIL_freeFileNamesTable(*fnt);
1098     *fnt = newFNT;
1099 }
1100 
UTIL_createFNT_fromROTable(const char ** filenames,size_t nbFilenames)1101 FileNamesTable* UTIL_createFNT_fromROTable(const char** filenames, size_t nbFilenames)
1102 {
1103     size_t const sizeof_FNTable = nbFilenames * sizeof(*filenames);
1104     const char** const newFNTable = (const char**)malloc(sizeof_FNTable);
1105     if (newFNTable==NULL) return NULL;
1106     memcpy((void*)newFNTable, filenames, sizeof_FNTable);  /* void* : mitigate a Visual compiler bug or limitation */
1107     return UTIL_assembleFileNamesTable(newFNTable, nbFilenames, NULL);
1108 }
1109 
1110 
1111 /*-****************************************
1112 *  count the number of cores
1113 ******************************************/
1114 
1115 #if defined(_WIN32) || defined(WIN32)
1116 
1117 #include <windows.h>
1118 
1119 typedef BOOL(WINAPI* LPFN_GLPI)(PSYSTEM_LOGICAL_PROCESSOR_INFORMATION, PDWORD);
1120 
CountSetBits(ULONG_PTR bitMask)1121 DWORD CountSetBits(ULONG_PTR bitMask)
1122 {
1123     DWORD LSHIFT = sizeof(ULONG_PTR)*8 - 1;
1124     DWORD bitSetCount = 0;
1125     ULONG_PTR bitTest = (ULONG_PTR)1 << LSHIFT;
1126     DWORD i;
1127 
1128     for (i = 0; i <= LSHIFT; ++i)
1129     {
1130         bitSetCount += ((bitMask & bitTest)?1:0);
1131         bitTest/=2;
1132     }
1133 
1134     return bitSetCount;
1135 }
1136 
UTIL_countCores(int logical)1137 int UTIL_countCores(int logical)
1138 {
1139     static int numCores = 0;
1140     if (numCores != 0) return numCores;
1141 
1142     {   LPFN_GLPI glpi;
1143         BOOL done = FALSE;
1144         PSYSTEM_LOGICAL_PROCESSOR_INFORMATION buffer = NULL;
1145         PSYSTEM_LOGICAL_PROCESSOR_INFORMATION ptr = NULL;
1146         DWORD returnLength = 0;
1147         size_t byteOffset = 0;
1148 
1149 #if defined(_MSC_VER)
1150 /* Visual Studio does not like the following cast */
1151 #   pragma warning( disable : 4054 )  /* conversion from function ptr to data ptr */
1152 #   pragma warning( disable : 4055 )  /* conversion from data ptr to function ptr */
1153 #endif
1154         glpi = (LPFN_GLPI)(void*)GetProcAddress(GetModuleHandle(TEXT("kernel32")),
1155                                                "GetLogicalProcessorInformation");
1156 
1157         if (glpi == NULL) {
1158             goto failed;
1159         }
1160 
1161         while(!done) {
1162             DWORD rc = glpi(buffer, &returnLength);
1163             if (FALSE == rc) {
1164                 if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
1165                     if (buffer)
1166                         free(buffer);
1167                     buffer = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION)malloc(returnLength);
1168 
1169                     if (buffer == NULL) {
1170                         perror("zstd");
1171                         exit(1);
1172                     }
1173                 } else {
1174                     /* some other error */
1175                     goto failed;
1176                 }
1177             } else {
1178                 done = TRUE;
1179         }   }
1180 
1181         ptr = buffer;
1182 
1183         while (byteOffset + sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION) <= returnLength) {
1184 
1185             if (ptr->Relationship == RelationProcessorCore) {
1186                 if (logical)
1187                     numCores += CountSetBits(ptr->ProcessorMask);
1188                 else
1189                     numCores++;
1190             }
1191 
1192             ptr++;
1193             byteOffset += sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION);
1194         }
1195 
1196         free(buffer);
1197 
1198         return numCores;
1199     }
1200 
1201 failed:
1202     /* try to fall back on GetSystemInfo */
1203     {   SYSTEM_INFO sysinfo;
1204         GetSystemInfo(&sysinfo);
1205         numCores = sysinfo.dwNumberOfProcessors;
1206         if (numCores == 0) numCores = 1; /* just in case */
1207     }
1208     return numCores;
1209 }
1210 
1211 #elif defined(__APPLE__)
1212 
1213 #include <sys/sysctl.h>
1214 
1215 /* Use apple-provided syscall
1216  * see: man 3 sysctl */
UTIL_countCores(int logical)1217 int UTIL_countCores(int logical)
1218 {
1219     static S32 numCores = 0; /* apple specifies int32_t */
1220     if (numCores != 0) return numCores;
1221 
1222     {   size_t size = sizeof(S32);
1223         int const ret = sysctlbyname(logical ? "hw.logicalcpu" : "hw.physicalcpu", &numCores, &size, NULL, 0);
1224         if (ret != 0) {
1225             if (errno == ENOENT) {
1226                 /* entry not present, fall back on 1 */
1227                 numCores = 1;
1228             } else {
1229                 perror("zstd: can't get number of cpus");
1230                 exit(1);
1231             }
1232         }
1233 
1234         return numCores;
1235     }
1236 }
1237 
1238 #elif defined(__linux__)
1239 
1240 /* parse /proc/cpuinfo
1241  * siblings / cpu cores should give hyperthreading ratio
1242  * otherwise fall back on sysconf */
UTIL_countCores(int logical)1243 int UTIL_countCores(int logical)
1244 {
1245     static int numCores = 0;
1246 
1247     if (numCores != 0) return numCores;
1248 
1249     numCores = (int)sysconf(_SC_NPROCESSORS_ONLN);
1250     if (numCores == -1) {
1251         /* value not queryable, fall back on 1 */
1252         return numCores = 1;
1253     }
1254 
1255     /* try to determine if there's hyperthreading */
1256     {   FILE* const cpuinfo = fopen("/proc/cpuinfo", "r");
1257 #define BUF_SIZE 80
1258         char buff[BUF_SIZE];
1259 
1260         int siblings = 0;
1261         int cpu_cores = 0;
1262         int ratio = 1;
1263 
1264         if (cpuinfo == NULL) {
1265             /* fall back on the sysconf value */
1266             return numCores;
1267         }
1268 
1269         /* assume the cpu cores/siblings values will be constant across all
1270          * present processors */
1271         while (!feof(cpuinfo)) {
1272             if (fgets(buff, BUF_SIZE, cpuinfo) != NULL) {
1273                 if (strncmp(buff, "siblings", 8) == 0) {
1274                     const char* const sep = strchr(buff, ':');
1275                     if (sep == NULL || *sep == '\0') {
1276                         /* formatting was broken? */
1277                         goto failed;
1278                     }
1279 
1280                     siblings = atoi(sep + 1);
1281                 }
1282                 if (strncmp(buff, "cpu cores", 9) == 0) {
1283                     const char* const sep = strchr(buff, ':');
1284                     if (sep == NULL || *sep == '\0') {
1285                         /* formatting was broken? */
1286                         goto failed;
1287                     }
1288 
1289                     cpu_cores = atoi(sep + 1);
1290                 }
1291             } else if (ferror(cpuinfo)) {
1292                 /* fall back on the sysconf value */
1293                 goto failed;
1294         }   }
1295         if (siblings && cpu_cores && siblings > cpu_cores) {
1296             ratio = siblings / cpu_cores;
1297         }
1298 
1299         if (ratio && numCores > ratio && !logical) {
1300             numCores = numCores / ratio;
1301         }
1302 
1303 failed:
1304         fclose(cpuinfo);
1305         return numCores;
1306     }
1307 }
1308 
1309 #elif defined(__FreeBSD__)
1310 
1311 #include <sys/param.h>
1312 #include <sys/sysctl.h>
1313 
1314 /* Use physical core sysctl when available
1315  * see: man 4 smp, man 3 sysctl */
UTIL_countCores(int logical)1316 int UTIL_countCores(int logical)
1317 {
1318     static int numCores = 0; /* freebsd sysctl is native int sized */
1319 #if __FreeBSD_version >= 1300008
1320     static int perCore = 1;
1321 #endif
1322     if (numCores != 0) return numCores;
1323 
1324 #if __FreeBSD_version >= 1300008
1325     {   size_t size = sizeof(numCores);
1326         int ret = sysctlbyname("kern.smp.cores", &numCores, &size, NULL, 0);
1327         if (ret == 0) {
1328             if (logical) {
1329                 ret = sysctlbyname("kern.smp.threads_per_core", &perCore, &size, NULL, 0);
1330                 /* default to physical cores if logical cannot be read */
1331                 if (ret == 0)
1332                     numCores *= perCore;
1333             }
1334 
1335             return numCores;
1336         }
1337         if (errno != ENOENT) {
1338             perror("zstd: can't get number of cpus");
1339             exit(1);
1340         }
1341         /* sysctl not present, fall through to older sysconf method */
1342     }
1343 #else
1344     /* suppress unused parameter warning */
1345     (void) logical;
1346 #endif
1347 
1348     numCores = (int)sysconf(_SC_NPROCESSORS_ONLN);
1349     if (numCores == -1) {
1350         /* value not queryable, fall back on 1 */
1351         numCores = 1;
1352     }
1353     return numCores;
1354 }
1355 
1356 #elif defined(__NetBSD__) || defined(__OpenBSD__) || defined(__DragonFly__) || defined(__CYGWIN__)
1357 
1358 /* Use POSIX sysconf
1359  * see: man 3 sysconf */
UTIL_countCores(int logical)1360 int UTIL_countCores(int logical)
1361 {
1362     static int numCores = 0;
1363 
1364     /* suppress unused parameter warning */
1365     (void)logical;
1366 
1367     if (numCores != 0) return numCores;
1368 
1369     numCores = (int)sysconf(_SC_NPROCESSORS_ONLN);
1370     if (numCores == -1) {
1371         /* value not queryable, fall back on 1 */
1372         return numCores = 1;
1373     }
1374     return numCores;
1375 }
1376 
1377 #else
1378 
UTIL_countCores(int logical)1379 int UTIL_countCores(int logical)
1380 {
1381     /* assume 1 */
1382     return 1;
1383 }
1384 
1385 #endif
1386 
UTIL_countPhysicalCores(void)1387 int UTIL_countPhysicalCores(void)
1388 {
1389     return UTIL_countCores(0);
1390 }
1391 
UTIL_countLogicalCores(void)1392 int UTIL_countLogicalCores(void)
1393 {
1394     return UTIL_countCores(1);
1395 }
1396 
1397 #if defined (__cplusplus)
1398 }
1399 #endif
1400