xref: /freebsd/sys/contrib/zstd/doc/educational_decoder/harness.c (revision 4b50c451720d8b427757a6da1dd2bb4c52cd9e35)
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 static size_t read_file(const char *path, u8 **ptr)
25 {
26     FILE* const f = fopen(path, "rb");
27     if (!f) {
28         fprintf(stderr, "failed to open file %s \n", path);
29         exit(1);
30     }
31 
32     fseek(f, 0L, SEEK_END);
33     size_t const size = (size_t)ftell(f);
34     rewind(f);
35 
36     *ptr = malloc(size);
37     if (!ptr) {
38         fprintf(stderr, "failed to allocate memory to hold %s \n", path);
39         exit(1);
40     }
41 
42     size_t const read = fread(*ptr, 1, size, f);
43     if (read != size) {  /* must read everything in one pass */
44         fprintf(stderr, "error while reading file %s \n", path);
45         exit(1);
46     }
47 
48     fclose(f);
49 
50     return read;
51 }
52 
53 static void write_file(const char *path, const u8 *ptr, size_t size)
54 {
55     FILE* const f = fopen(path, "wb");
56     if (!f) {
57         fprintf(stderr, "failed to open file %s \n", path);
58         exit(1);
59     }
60 
61     size_t written = 0;
62     while (written < size) {
63         written += fwrite(ptr+written, 1, size, f);
64         if (ferror(f)) {
65             fprintf(stderr, "error while writing file %s\n", path);
66             exit(1);
67     }   }
68 
69     fclose(f);
70 }
71 
72 int main(int argc, char **argv)
73 {
74     if (argc < 3) {
75         fprintf(stderr, "usage: %s <file.zst> <out_path> [dictionary] \n",
76                 argv[0]);
77 
78         return 1;
79     }
80 
81     u8* input;
82     size_t const input_size = read_file(argv[1], &input);
83 
84     u8* dict = NULL;
85     size_t dict_size = 0;
86     if (argc >= 4) {
87         dict_size = read_file(argv[3], &dict);
88     }
89 
90     size_t out_capacity = ZSTD_get_decompressed_size(input, input_size);
91     if (out_capacity == (size_t)-1) {
92         out_capacity = MAX_COMPRESSION_RATIO * input_size;
93         fprintf(stderr, "WARNING: Compressed data does not contain "
94                         "decompressed size, going to assume the compression "
95                         "ratio is at most %d (decompressed size of at most "
96                         "%u) \n",
97                 MAX_COMPRESSION_RATIO, (unsigned)out_capacity);
98     }
99     if (out_capacity > MAX_OUTPUT_SIZE) {
100         fprintf(stderr,
101                 "Required output size too large for this implementation \n");
102         return 1;
103     }
104 
105     u8* const output = malloc(out_capacity);
106     if (!output) {
107         fprintf(stderr, "failed to allocate memory \n");
108         return 1;
109     }
110 
111     dictionary_t* const parsed_dict = create_dictionary();
112     if (dict) {
113         parse_dictionary(parsed_dict, dict, dict_size);
114     }
115     size_t const decompressed_size =
116         ZSTD_decompress_with_dict(output, out_capacity,
117                                   input, input_size,
118                                   parsed_dict);
119 
120     free_dictionary(parsed_dict);
121 
122     write_file(argv[2], output, decompressed_size);
123 
124     free(input);
125     free(output);
126     free(dict);
127     return 0;
128 }
129