Home | History | Annotate | Download | only in tools 
      1 /* Copyright 2014 Google Inc. All Rights Reserved.
      2 
      3    Distributed under MIT license.
      4    See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
      5 */
      6 
      7 /* Example main() function for Brotli library. */
      8 
      9 #include <fcntl.h>
     10 #include <stdio.h>
     11 #include <stdlib.h>
     12 #include <string.h>
     13 #include <sys/stat.h>
     14 #include <sys/types.h>
     15 #include <time.h>
     16 
     17 #include <brotli/decode.h>
     18 #include <brotli/encode.h>
     19 
     20 #if !defined(_WIN32)
     21 #include <unistd.h>
     22 #include <utime.h>
     23 #else
     24 #include <io.h>
     25 #include <share.h>
     26 #include <sys/utime.h>
     27 
     28 #define MAKE_BINARY(FILENO) (_setmode((FILENO), _O_BINARY), (FILENO))
     29 
     30 #if !defined(__MINGW32__)
     31 #define STDIN_FILENO MAKE_BINARY(_fileno(stdin))
     32 #define STDOUT_FILENO MAKE_BINARY(_fileno(stdout))
     33 #define S_IRUSR S_IREAD
     34 #define S_IWUSR S_IWRITE
     35 #endif
     36 
     37 #define fdopen _fdopen
     38 #define unlink _unlink
     39 #define utimbuf _utimbuf
     40 #define utime _utime
     41 
     42 #define fopen ms_fopen
     43 #define open ms_open
     44 
     45 #define chmod(F, P) (0)
     46 #define chown(F, O, G) (0)
     47 
     48 #if defined(_MSC_VER) && (_MSC_VER >= 1400)
     49 #define fseek _fseeki64
     50 #define ftell _ftelli64
     51 #endif
     52 
     53 static FILE* ms_fopen(const char *filename, const char *mode) {
     54   FILE* result = 0;
     55   fopen_s(&result, filename, mode);
     56   return result;
     57 }
     58 
     59 static int ms_open(const char *filename, int oflag, int pmode) {
     60   int result = -1;
     61   _sopen_s(&result, filename, oflag | O_BINARY, _SH_DENYNO, pmode);
     62   return result;
     63 }
     64 #endif  /* WIN32 */
     65 
     66 static int ParseQuality(const char* s, int* quality) {
     67   if (s[0] >= '0' && s[0] <= '9') {
     68     *quality = s[0] - '0';
     69     if (s[1] >= '0' && s[1] <= '9') {
     70       *quality = *quality * 10 + s[1] - '0';
     71       return (s[2] == 0) ? 1 : 0;
     72     }
     73     return (s[1] == 0) ? 1 : 0;
     74   }
     75   return 0;
     76 }
     77 
     78 static void ParseArgv(int argc, char **argv,
     79                       char **input_path,
     80                       char **output_path,
     81                       char **dictionary_path,
     82                       int *force,
     83                       int *quality,
     84                       int *decompress,
     85                       int *repeat,
     86                       int *verbose,
     87                       int *lgwin,
     88                       int *copy_stat) {
     89   int k;
     90   *force = 0;
     91   *input_path = 0;
     92   *output_path = 0;
     93   *repeat = 1;
     94   *verbose = 0;
     95   *lgwin = 22;
     96   *copy_stat = 1;
     97   {
     98     size_t argv0_len = strlen(argv[0]);
     99     *decompress =
    100         argv0_len >= 5 && strcmp(&argv[0][argv0_len - 5], "unbro") == 0;
    101   }
    102   for (k = 1; k < argc; ++k) {
    103     if (!strcmp("--force", argv[k]) ||
    104         !strcmp("-f", argv[k])) {
    105       if (*force != 0) {
    106         goto error;
    107       }
    108       *force = 1;
    109       continue;
    110     } else if (!strcmp("--decompress", argv[k]) ||
    111                !strcmp("--uncompress", argv[k]) ||
    112                !strcmp("-d", argv[k])) {
    113       *decompress = 1;
    114       continue;
    115     } else if (!strcmp("--verbose", argv[k]) ||
    116                !strcmp("-v", argv[k])) {
    117       if (*verbose != 0) {
    118         goto error;
    119       }
    120       *verbose = 1;
    121       continue;
    122     } else if (!strcmp("--no-copy-stat", argv[k])) {
    123       if (*copy_stat == 0) {
    124         goto error;
    125       }
    126       *copy_stat = 0;
    127       continue;
    128     }
    129     if (k < argc - 1) {
    130       if (!strcmp("--input", argv[k]) ||
    131           !strcmp("--in", argv[k]) ||
    132           !strcmp("-i", argv[k])) {
    133         if (*input_path != 0) {
    134           goto error;
    135         }
    136         *input_path = argv[k + 1];
    137         ++k;
    138         continue;
    139       } else if (!strcmp("--output", argv[k]) ||
    140                  !strcmp("--out", argv[k]) ||
    141                  !strcmp("-o", argv[k])) {
    142         if (*output_path != 0) {
    143           goto error;
    144         }
    145         *output_path = argv[k + 1];
    146         ++k;
    147         continue;
    148       } else if (!strcmp("--custom-dictionary", argv[k])) {
    149         if (*dictionary_path != 0) {
    150           goto error;
    151         }
    152         *dictionary_path = argv[k + 1];
    153         ++k;
    154         continue;
    155       } else if (!strcmp("--quality", argv[k]) ||
    156                  !strcmp("-q", argv[k])) {
    157         if (!ParseQuality(argv[k + 1], quality)) {
    158           goto error;
    159         }
    160         ++k;
    161         continue;
    162       } else if (!strcmp("--repeat", argv[k]) ||
    163                  !strcmp("-r", argv[k])) {
    164         if (!ParseQuality(argv[k + 1], repeat)) {
    165           goto error;
    166         }
    167         ++k;
    168         continue;
    169       } else if (!strcmp("--window", argv[k]) ||
    170                   !strcmp("-w", argv[k])) {
    171         if (!ParseQuality(argv[k + 1], lgwin)) {
    172           goto error;
    173         }
    174         if (*lgwin < 10 || *lgwin >= 25) {
    175           goto error;
    176         }
    177         ++k;
    178         continue;
    179       }
    180     }
    181     goto error;
    182   }
    183   return;
    184 error:
    185   fprintf(stderr,
    186           "Usage: %s [--force] [--quality n] [--decompress]"
    187           " [--input filename] [--output filename] [--repeat iters]"
    188           " [--verbose] [--window n] [--custom-dictionary filename]"
    189           " [--no-copy-stat]\n",
    190           argv[0]);
    191   exit(1);
    192 }
    193 
    194 static FILE* OpenInputFile(const char* input_path) {
    195   FILE* f;
    196   if (input_path == 0) {
    197     return fdopen(STDIN_FILENO, "rb");
    198   }
    199   f = fopen(input_path, "rb");
    200   if (f == 0) {
    201     perror("fopen");
    202     exit(1);
    203   }
    204   return f;
    205 }
    206 
    207 static FILE *OpenOutputFile(const char *output_path, const int force) {
    208   int fd;
    209   if (output_path == 0) {
    210     return fdopen(STDOUT_FILENO, "wb");
    211   }
    212   fd = open(output_path, O_CREAT | (force ? 0 : O_EXCL) | O_WRONLY | O_TRUNC,
    213             S_IRUSR | S_IWUSR);
    214   if (fd < 0) {
    215     if (!force) {
    216       struct stat statbuf;
    217       if (stat(output_path, &statbuf) == 0) {
    218         fprintf(stderr, "output file exists\n");
    219         exit(1);
    220       }
    221     }
    222     perror("open");
    223     exit(1);
    224   }
    225   return fdopen(fd, "wb");
    226 }
    227 
    228 static int64_t FileSize(const char *path) {
    229   FILE *f = fopen(path, "rb");
    230   int64_t retval;
    231   if (f == NULL) {
    232     return -1;
    233   }
    234   if (fseek(f, 0L, SEEK_END) != 0) {
    235     fclose(f);
    236     return -1;
    237   }
    238   retval = ftell(f);
    239   if (fclose(f) != 0) {
    240     return -1;
    241   }
    242   return retval;
    243 }
    244 
    245 /* Copy file times and permissions.
    246    TODO: this is a "best effort" implementation; honest cross-platform
    247    fully featured implementation is way too hacky; add more hacks by request. */
    248 static void CopyStat(const char* input_path, const char* output_path) {
    249   struct stat statbuf;
    250   struct utimbuf times;
    251   int res;
    252   if (input_path == 0 || output_path == 0) {
    253     return;
    254   }
    255   if (stat(input_path, &statbuf) != 0) {
    256     return;
    257   }
    258   times.actime = statbuf.st_atime;
    259   times.modtime = statbuf.st_mtime;
    260   utime(output_path, &times);
    261   res = chmod(output_path, statbuf.st_mode & (S_IRWXU | S_IRWXG | S_IRWXO));
    262   if (res != 0)
    263     perror("chmod failed");
    264   res = chown(output_path, (uid_t)-1, statbuf.st_gid);
    265   if (res != 0)
    266     perror("chown failed");
    267   res = chown(output_path, statbuf.st_uid, (gid_t)-1);
    268   if (res != 0)
    269     perror("chown failed");
    270 }
    271 
    272 /* Result ownersip is passed to caller.
    273    |*dictionary_size| is set to resulting buffer size. */
    274 static uint8_t* ReadDictionary(const char* path, size_t* dictionary_size) {
    275   static const int kMaxDictionarySize = (1 << 24) - 16;
    276   FILE *f = fopen(path, "rb");
    277   int64_t file_size_64;
    278   uint8_t* buffer;
    279   size_t bytes_read;
    280 
    281   if (f == NULL) {
    282     perror("fopen");
    283     exit(1);
    284   }
    285 
    286   file_size_64 = FileSize(path);
    287   if (file_size_64 == -1) {
    288     fprintf(stderr, "could not get size of dictionary file");
    289     exit(1);
    290   }
    291 
    292   if (file_size_64 > kMaxDictionarySize) {
    293     fprintf(stderr, "dictionary is larger than maximum allowed: %d\n",
    294             kMaxDictionarySize);
    295     exit(1);
    296   }
    297   *dictionary_size = (size_t)file_size_64;
    298 
    299   buffer = (uint8_t*)malloc(*dictionary_size);
    300   if (!buffer) {
    301     fprintf(stderr, "could not read dictionary: out of memory\n");
    302     exit(1);
    303   }
    304   bytes_read = fread(buffer, sizeof(uint8_t), *dictionary_size, f);
    305   if (bytes_read != *dictionary_size) {
    306     fprintf(stderr, "could not read dictionary\n");
    307     exit(1);
    308   }
    309   fclose(f);
    310   return buffer;
    311 }
    312 
    313 static const size_t kFileBufferSize = 65536;
    314 
    315 static int Decompress(FILE* fin, FILE* fout, const char* dictionary_path) {
    316   /* Dictionary should be kept during first rounds of decompression. */
    317   uint8_t* dictionary = NULL;
    318   uint8_t* input;
    319   uint8_t* output;
    320   size_t available_in;
    321   const uint8_t* next_in;
    322   size_t available_out = kFileBufferSize;
    323   uint8_t* next_out;
    324   BrotliDecoderResult result = BROTLI_DECODER_RESULT_ERROR;
    325   BrotliDecoderState* s = BrotliDecoderCreateInstance(NULL, NULL, NULL);
    326   if (!s) {
    327     fprintf(stderr, "out of memory\n");
    328     return 0;
    329   }
    330   if (dictionary_path != NULL) {
    331     size_t dictionary_size = 0;
    332     dictionary = ReadDictionary(dictionary_path, &dictionary_size);
    333     BrotliDecoderSetCustomDictionary(s, dictionary_size, dictionary);
    334   }
    335   input = (uint8_t*)malloc(kFileBufferSize);
    336   output = (uint8_t*)malloc(kFileBufferSize);
    337   if (!input || !output) {
    338     fprintf(stderr, "out of memory\n");
    339     goto end;
    340   }
    341   next_out = output;
    342   result = BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT;
    343   while (1) {
    344     if (result == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT) {
    345       if (feof(fin)) {
    346         break;
    347       }
    348       available_in = fread(input, 1, kFileBufferSize, fin);
    349       next_in = input;
    350       if (ferror(fin)) {
    351         break;
    352       }
    353     } else if (result == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT) {
    354       fwrite(output, 1, kFileBufferSize, fout);
    355       if (ferror(fout)) {
    356         break;
    357       }
    358       available_out = kFileBufferSize;
    359       next_out = output;
    360     } else {
    361       break; /* Error or success. */
    362     }
    363     result = BrotliDecoderDecompressStream(
    364         s, &available_in, &next_in, &available_out, &next_out, 0);
    365   }
    366   if (next_out != output) {
    367     fwrite(output, 1, (size_t)(next_out - output), fout);
    368   }
    369 
    370   if ((result == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT) || ferror(fout)) {
    371     fprintf(stderr, "failed to write output\n");
    372   } else if (result != BROTLI_DECODER_RESULT_SUCCESS) {
    373     /* Error or needs more input. */
    374     fprintf(stderr, "corrupt input\n");
    375   }
    376 
    377 end:
    378   free(dictionary);
    379   free(input);
    380   free(output);
    381   BrotliDecoderDestroyInstance(s);
    382   return (result == BROTLI_DECODER_RESULT_SUCCESS) ? 1 : 0;
    383 }
    384 
    385 static int Compress(int quality, int lgwin, FILE* fin, FILE* fout,
    386     const char *dictionary_path) {
    387   BrotliEncoderState* s = BrotliEncoderCreateInstance(0, 0, 0);
    388   uint8_t* buffer = (uint8_t*)malloc(kFileBufferSize << 1);
    389   uint8_t* input = buffer;
    390   uint8_t* output = buffer + kFileBufferSize;
    391   size_t available_in = 0;
    392   const uint8_t* next_in = NULL;
    393   size_t available_out = kFileBufferSize;
    394   uint8_t* next_out = output;
    395   int is_eof = 0;
    396   int is_ok = 1;
    397 
    398   if (!s || !buffer) {
    399     is_ok = 0;
    400     goto finish;
    401   }
    402 
    403   BrotliEncoderSetParameter(s, BROTLI_PARAM_QUALITY, (uint32_t)quality);
    404   BrotliEncoderSetParameter(s, BROTLI_PARAM_LGWIN, (uint32_t)lgwin);
    405   if (dictionary_path != NULL) {
    406     size_t dictionary_size = 0;
    407     uint8_t* dictionary = ReadDictionary(dictionary_path, &dictionary_size);
    408     BrotliEncoderSetCustomDictionary(s, dictionary_size, dictionary);
    409     free(dictionary);
    410   }
    411 
    412   while (1) {
    413     if (available_in == 0 && !is_eof) {
    414       available_in = fread(input, 1, kFileBufferSize, fin);
    415       next_in = input;
    416       if (ferror(fin)) break;
    417       is_eof = feof(fin);
    418     }
    419 
    420     if (!BrotliEncoderCompressStream(s,
    421         is_eof ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS,
    422         &available_in, &next_in, &available_out, &next_out, NULL)) {
    423       is_ok = 0;
    424       break;
    425     }
    426 
    427     if (available_out != kFileBufferSize) {
    428       size_t out_size = kFileBufferSize - available_out;
    429       fwrite(output, 1, out_size, fout);
    430       if (ferror(fout)) break;
    431       available_out = kFileBufferSize;
    432       next_out = output;
    433     }
    434 
    435     if (BrotliEncoderIsFinished(s)) break;
    436   }
    437 
    438 finish:
    439   free(buffer);
    440   BrotliEncoderDestroyInstance(s);
    441 
    442   if (!is_ok) {
    443     /* Should detect OOM? */
    444     fprintf(stderr, "failed to compress data\n");
    445     return 0;
    446   } else if (ferror(fout)) {
    447     fprintf(stderr, "failed to write output\n");
    448     return 0;
    449   } else if (ferror(fin)) {
    450     fprintf(stderr, "failed to read input\n");
    451     return 0;
    452   }
    453   return 1;
    454 }
    455 
    456 int main(int argc, char** argv) {
    457   char *input_path = 0;
    458   char *output_path = 0;
    459   char *dictionary_path = 0;
    460   int force = 0;
    461   int quality = 11;
    462   int decompress = 0;
    463   int repeat = 1;
    464   int verbose = 0;
    465   int lgwin = 0;
    466   int copy_stat = 1;
    467   clock_t clock_start;
    468   int i;
    469   ParseArgv(argc, argv, &input_path, &output_path, &dictionary_path, &force,
    470             &quality, &decompress, &repeat, &verbose, &lgwin, &copy_stat);
    471   clock_start = clock();
    472   for (i = 0; i < repeat; ++i) {
    473     FILE* fin = OpenInputFile(input_path);
    474     FILE* fout = OpenOutputFile(output_path, force || (repeat > 1));
    475     int is_ok = 0;
    476     if (decompress) {
    477       is_ok = Decompress(fin, fout, dictionary_path);
    478     } else {
    479       is_ok = Compress(quality, lgwin, fin, fout, dictionary_path);
    480     }
    481     if (!is_ok) {
    482       unlink(output_path);
    483       exit(1);
    484     }
    485     if (fclose(fout) != 0) {
    486       perror("fclose");
    487       exit(1);
    488     }
    489     /* TOCTOU violation, but otherwise it is impossible to set file times. */
    490     if (copy_stat && (i + 1 == repeat)) {
    491       CopyStat(input_path, output_path);
    492     }
    493     if (fclose(fin) != 0) {
    494       perror("fclose");
    495       exit(1);
    496     }
    497   }
    498   if (verbose) {
    499     clock_t clock_end = clock();
    500     double duration = (double)(clock_end - clock_start) / CLOCKS_PER_SEC;
    501     int64_t uncompressed_size;
    502     double uncompressed_bytes_in_MB;
    503     if (duration < 1e-9) {
    504       duration = 1e-9;
    505     }
    506     uncompressed_size = FileSize(decompress ? output_path : input_path);
    507     if (uncompressed_size == -1) {
    508       fprintf(stderr, "failed to determine uncompressed file size\n");
    509       exit(1);
    510     }
    511     uncompressed_bytes_in_MB =
    512         (double)(repeat * uncompressed_size) / (1024.0 * 1024.0);
    513     if (decompress) {
    514       printf("Brotli decompression speed: ");
    515     } else {
    516       printf("Brotli compression speed: ");
    517     }
    518     printf("%g MB/s\n", uncompressed_bytes_in_MB / duration);
    519   }
    520   return 0;
    521 }
    522