Home | History | Annotate | Line # | Download | only in recovery 
      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