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 } 153