ceph/src/zstd/doc/educational_decoder/harness.c

   1 /*
   2  * Copyright (c) 2017-present, 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  */
   9
  10 #include <stdio.h>
  11 #include <stdlib.h>
  12
  13 #include "zstd_decompress.h"
  14
  15 typedef unsigned char u8;
  16
  17 // If the data doesn't have decompressed size with it, fallback on assuming the
  18 // compression ratio is at most 16
  19 #define MAX_COMPRESSION_RATIO (16)
  20
  21 // Protect against allocating too much memory for output
  22 #define MAX_OUTPUT_SIZE ((size_t)1024 * 1024 * 1024)
  23
  24 u8 *input;
  25 u8 *output;
  26 u8 *dict;
  27
  28 size_t read_file(const char *path, u8 **ptr) {
  29     FILE *f = fopen(path, "rb");
  30     if (!f) {
  31         fprintf(stderr, "failed to open file %s\n", path);
  32         exit(1);
  33     }
  34
  35     fseek(f, 0L, SEEK_END);
  36     size_t size = ftell(f);
  37     rewind(f);
  38
  39     *ptr = malloc(size);
  40     if (!ptr) {
  41         fprintf(stderr, "failed to allocate memory to hold %s\n", path);
  42         exit(1);
  43     }
  44
  45     size_t pos = 0;
  46     while (!feof(f)) {
  47         size_t read = fread(&(*ptr)[pos], 1, size, f);
  48         if (ferror(f)) {
  49             fprintf(stderr, "error while reading file %s\n", path);
  50             exit(1);
  51         }
  52         pos += read;
  53     }
  54
  55     fclose(f);
  56
  57     return pos;
  58 }
  59
  60 void write_file(const char *path, const u8 *ptr, size_t size) {
  61     FILE *f = fopen(path, "wb");
  62
  63     size_t written = 0;
  64     while (written < size) {
  65         written += fwrite(&ptr[written], 1, size, f);
  66         if (ferror(f)) {
  67             fprintf(stderr, "error while writing file %s\n", path);
  68             exit(1);
  69         }
  70     }
  71
  72     fclose(f);
  73 }
  74
  75 int main(int argc, char **argv) {
  76     if (argc < 3) {
  77         fprintf(stderr, "usage: %s <file.zst> <out_path> [dictionary]\n",
  78                 argv[0]);
  79
  80         return 1;
  81     }
  82
  83     size_t input_size = read_file(argv[1], &input);
  84     size_t dict_size = 0;
  85     if (argc >= 4) {
  86         dict_size = read_file(argv[3], &dict);
  87     }
  88
  89     size_t decompressed_size = ZSTD_get_decompressed_size(input, input_size);
  90     if (decompressed_size == (size_t)-1) {
  91         decompressed_size = MAX_COMPRESSION_RATIO * input_size;
  92         fprintf(stderr, "WARNING: Compressed data does not contain "
  93                         "decompressed size, going to assume the compression "
  94                         "ratio is at most %d (decompressed size of at most "
  95                         "%zu)\n",
  96                 MAX_COMPRESSION_RATIO, decompressed_size);
  97     }
  98     if (decompressed_size > MAX_OUTPUT_SIZE) {
  99         fprintf(stderr,
 100                 "Required output size too large for this implementation\n");
 101         return 1;
 102     }
 103     output = malloc(decompressed_size);
 104     if (!output) {
 105         fprintf(stderr, "failed to allocate memory\n");
 106         return 1;
 107     }
 108
 109     dictionary_t* const parsed_dict = create_dictionary();
 110     if (dict) {
 111         parse_dictionary(parsed_dict, dict, dict_size);
 112     }
 113     size_t decompressed =
 114         ZSTD_decompress_with_dict(output, decompressed_size,
 115                                   input, input_size, parsed_dict);
 116
 117     free_dictionary(parsed_dict);
 118
 119     write_file(argv[2], output, decompressed);
 120
 121     free(input);
 122     free(output);
 123     free(dict);
 124     input = output = dict = NULL;
 125 }