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