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 <string.h> |
12 | #include <stdio.h> |
13 | #include <stdlib.h> |
14 | |
15 | #define ZSTD_STATIC_LINKING_ONLY |
16 | #include "util.h" |
17 | #include "zstd.h" |
18 | |
19 | #define CHECK(cond, ...) \ |
20 | do { \ |
21 | if (!(cond)) { \ |
22 | fprintf(stderr, "%s:%d CHECK(%s) failed: ", __FILE__, __LINE__, #cond); \ |
23 | fprintf(stderr, "" __VA_ARGS__); \ |
24 | fprintf(stderr, "\n"); \ |
25 | exit(1); \ |
26 | } \ |
27 | } while (0) |
28 | |
29 | static void usage(char const *program) { |
30 | fprintf(stderr, "USAGE: %s FILE.zst PREFIX\n", program); |
31 | fprintf(stderr, "FILE.zst: A zstd compressed file with multiple frames\n"); |
32 | fprintf(stderr, "PREFIX: The output prefix. Uncompressed files will be " |
33 | "created named ${PREFIX}0 ${PREFIX}1...\n\n"); |
34 | fprintf(stderr, "This program takes concatenated zstd frames and " |
35 | "decompresses them into individual files.\n"); |
36 | fprintf(stderr, "E.g. files created with a command like: zstd -r directory " |
37 | "-o file.zst\n"); |
38 | } |
39 | |
40 | typedef struct { |
41 | char *data; |
42 | size_t size; |
43 | size_t frames; |
44 | size_t maxFrameSize; |
45 | } ZstdFrames; |
46 | |
47 | static ZstdFrames readFile(char const *fileName) { |
48 | U64 const fileSize = UTIL_getFileSize(fileName); |
49 | CHECK(fileSize != UTIL_FILESIZE_UNKNOWN, "Unknown file size!"); |
50 | |
51 | char *const data = (char *)malloc(fileSize); |
52 | CHECK(data != NULL, "Allocation failed"); |
53 | |
54 | FILE *file = fopen(fileName, "rb"); |
55 | CHECK(file != NULL, "fopen failed"); |
56 | |
57 | size_t const readSize = fread(data, 1, fileSize, file); |
58 | CHECK(readSize == fileSize, "fread failed"); |
59 | |
60 | fclose(file); |
61 | ZstdFrames frames; |
62 | frames.data = (char *)data; |
63 | frames.size = fileSize; |
64 | frames.frames = 0; |
65 | |
66 | size_t index; |
67 | size_t maxFrameSize = 0; |
68 | for (index = 0; index < fileSize;) { |
69 | size_t const frameSize = |
70 | ZSTD_findFrameCompressedSize(data + index, fileSize - index); |
71 | CHECK(!ZSTD_isError(frameSize), "Bad zstd frame: %s", |
72 | ZSTD_getErrorName(frameSize)); |
73 | if (frameSize > maxFrameSize) |
74 | maxFrameSize = frameSize; |
75 | frames.frames += 1; |
76 | index += frameSize; |
77 | } |
78 | CHECK(index == fileSize, "Zstd file corrupt!"); |
79 | frames.maxFrameSize = maxFrameSize; |
80 | |
81 | return frames; |
82 | } |
83 | |
84 | static int computePadding(size_t numFrames) { |
85 | return snprintf(NULL, 0, "%u", (unsigned)numFrames); |
86 | } |
87 | |
88 | int main(int argc, char **argv) { |
89 | if (argc != 3) { |
90 | usage(argv[0]); |
91 | exit(1); |
92 | } |
93 | char const *const zstdFile = argv[1]; |
94 | char const *const prefix = argv[2]; |
95 | |
96 | ZstdFrames frames = readFile(zstdFile); |
97 | |
98 | if (frames.frames <= 1) { |
99 | fprintf( |
100 | stderr, |
101 | "%s only has %u zstd frame. Simply use `zstd -d` to decompress it.\n", |
102 | zstdFile, (unsigned)frames.frames); |
103 | exit(1); |
104 | } |
105 | |
106 | int const padding = computePadding(frames.frames - 1); |
107 | |
108 | size_t const outFileNameSize = strlen(prefix) + padding + 1; |
109 | char* outFileName = malloc(outFileNameSize); |
110 | CHECK(outFileName != NULL, "Allocation failure"); |
111 | |
112 | size_t const bufferSize = 128 * 1024; |
113 | void *buffer = malloc(bufferSize); |
114 | CHECK(buffer != NULL, "Allocation failure"); |
115 | |
116 | ZSTD_DCtx* dctx = ZSTD_createDCtx(); |
117 | CHECK(dctx != NULL, "Allocation failure"); |
118 | |
119 | fprintf(stderr, "Recovering %u files...\n", (unsigned)frames.frames); |
120 | |
121 | size_t index; |
122 | size_t frame = 0; |
123 | for (index = 0; index < frames.size; ++frame) { |
124 | size_t const frameSize = |
125 | ZSTD_findFrameCompressedSize(frames.data + index, frames.size - index); |
126 | |
127 | int const ret = snprintf(outFileName, outFileNameSize, "%s%0*u", prefix, padding, (unsigned)frame); |
128 | CHECK(ret >= 0 && (size_t)ret <= outFileNameSize, "snprintf failed!"); |
129 | |
130 | FILE* outFile = fopen(outFileName, "wb"); |
131 | CHECK(outFile != NULL, "fopen failed"); |
132 | |
133 | ZSTD_DCtx_reset(dctx, ZSTD_reset_session_only); |
134 | ZSTD_inBuffer in = {frames.data + index, frameSize, 0}; |
135 | while (in.pos < in.size) { |
136 | ZSTD_outBuffer out = {buffer, bufferSize, 0}; |
137 | CHECK(!ZSTD_isError(ZSTD_decompressStream(dctx, &out, &in)), "decompression failed"); |
138 | size_t const writeSize = fwrite(out.dst, 1, out.pos, outFile); |
139 | CHECK(writeSize == out.pos, "fwrite failed"); |
140 | } |
141 | fclose(outFile); |
142 | fprintf(stderr, "Recovered %s\n", outFileName); |
143 | index += frameSize; |
144 | } |
145 | fprintf(stderr, "Complete\n"); |
146 | |
147 | free(outFileName); |
148 | ZSTD_freeDCtx(dctx); |
149 | free(buffer); |
150 | free(frames.data); |
151 | return 0; |
152 | } |