deps/libchdr/deps/zstd-1.5.5/contrib/recovery/recover_directory.c

   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 }