cmd/zstd/unzstd.c

   1 // unzstd -- .warc.zst decompressor
   2 // Copyright (C) 2021-2024 Sergey Matveev <stargrave@stargrave.org>
   3 //
   4 // This program is free software: you can redistribute it and/or modify
   5 // it under the terms of the GNU General Public License as published by
   6 // the Free Software Foundation, version 3 of the License.
   7 //
   8 // This program is distributed in the hope that it will be useful,
   9 // but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 // GNU General Public License for more details.
  12 //
  13 // You should have received a copy of the GNU General Public License
  14 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
  15
  16 // https://iipc.github.io/warc-specifications/specifications/warc-zstd/
  17
  18 #include <errno.h>
  19 #include <inttypes.h>
  20 #include <stdbool.h>
  21 #include <stdio.h>
  22 #include <stdlib.h>
  23 #include <string.h>
  24
  25 #include <zstd.h>
  26
  27 #ifdef __FreeBSD__
  28 #include "capsicum.c.in"
  29 #include <capsicum_helpers.h>
  30 #include <err.h>
  31 #include <sysexits.h>
  32 #endif // __FreeBSD__
  33
  34 static uint32_t
  35 le32dec(const char buf[4])
  36 {
  37     return (uint32_t)(buf[3]) << 24 | (uint32_t)(buf[2]) << 16 |
  38            (uint32_t)(buf[1]) << 8 | (uint32_t)(buf[0]);
  39 }
  40
  41 int
  42 main(int argc, char **argv)
  43 {
  44     FILE *fdOff = fdopen(3, "wb");
  45 #ifdef __FreeBSD__
  46     if ((fdOff != NULL) && (caph_limit_stream(3, CAPH_WRITE)) != 0) {
  47         errx(EX_OSERR, "can not caph_limit_stream(3)");
  48     }
  49     capsicum_start();
  50 #endif // __FreeBSD__
  51
  52     ZSTD_DCtx *ctx = ZSTD_createDCtx();
  53     if (ctx == NULL) {
  54         fputs("can not initialize ZSTD_DCtx\n", stderr);
  55         return 1;
  56     }
  57     int rc = EXIT_FAILURE;
  58     char *bufIn = NULL;
  59     char *bufOut = NULL;
  60     const size_t bufInSize = ZSTD_DStreamInSize();
  61     bufIn = malloc(bufInSize);
  62     if (bufIn == NULL) {
  63         fputs("no memory\n", stderr);
  64         goto Exit;
  65     }
  66     const size_t bufOutSize = ZSTD_DStreamOutSize();
  67     bufOut = malloc(bufOutSize);
  68     if (bufOut == NULL) {
  69         fputs("no memory\n", stderr);
  70         goto Exit;
  71     }
  72
  73     unsigned long long bufSize = 0;
  74
  75     ZSTD_inBuffer bIn = {bufIn, 0, 0};
  76     ZSTD_outBuffer bOut = {bufOut, 0, 0};
  77
  78     bool isEmpty = true;
  79     bool lastBlock = false;
  80     size_t n = 0;
  81     size_t written = 0;
  82     size_t offset = 0;
  83     size_t offsetPrev = 0;
  84     size_t zCode = 0;
  85 ReadAgain:
  86     for (;;) {
  87         n = fread(bufIn, 1, bufInSize, stdin);
  88         if (n != bufInSize) {
  89             if (feof(stdin)) {
  90                 lastBlock = true;
  91             } else {
  92                 perror("can not fread(FILE)");
  93                 goto Exit;
  94             }
  95         }
  96         if (n >= 8 && le32dec(bufIn) == 0x184D2A5D) {
  97             // dictionary frame
  98             size_t dictSize = (size_t)le32dec(bufIn + 4);
  99             char *dict = malloc(dictSize);
 100             if (dict == NULL) {
 101                 fprintf(stderr, "insufficient memory for dictionary: %zu\n", dictSize);
 102                 goto Exit;
 103             }
 104             const size_t alreadyRead = n - 8;
 105             memcpy(dict, bufIn + 8, alreadyRead);
 106             errno = 0;
 107             n = fread(dict + alreadyRead, 1, dictSize - alreadyRead, stdin);
 108             if (n != dictSize - alreadyRead) {
 109                 perror("can not read dictionary data");
 110                 free(dict);
 111                 goto Exit;
 112             }
 113             offset = dictSize + 8;
 114             offsetPrev = offset;
 115             if (fdOff != NULL) {
 116                 fprintf(fdOff, "%zu\t0\n", offset);
 117             }
 118             uint32_t hdr = le32dec(dict);
 119             switch (hdr) {
 120             case ZSTD_MAGIC_DICTIONARY:
 121                 zCode = ZSTD_DCtx_loadDictionary(ctx, dict, dictSize);
 122                 free(dict);
 123                 if ((zCode != 0) && (ZSTD_isError(zCode))) {
 124                     fprintf(
 125                         stderr,
 126                         "can not load dictionary: %s\n",
 127                         ZSTD_getErrorName(zCode));
 128                     goto Exit;
 129                 }
 130                 goto ReadAgain;
 131             case ZSTD_MAGICNUMBER:
 132                 bufSize = ZSTD_getFrameContentSize(dict, dictSize);
 133                 switch (bufSize) {
 134                 case ZSTD_CONTENTSIZE_UNKNOWN:
 135                 case ZSTD_CONTENTSIZE_ERROR:
 136                     fprintf(stderr, "can not determine dictionary's size\n");
 137                     free(dict);
 138                     goto Exit;
 139                 }
 140                 char *buf = malloc(bufSize);
 141                 if (buf == NULL) {
 142                     fprintf(
 143                         stderr, "insufficient memory for dictionary: %llu\n", bufSize);
 144                     free(dict);
 145                     goto Exit;
 146                 }
 147                 zCode = ZSTD_decompress(buf, bufSize, dict, dictSize);
 148                 free(dict);
 149                 if (ZSTD_isError(zCode)) {
 150                     fprintf(
 151                         stderr,
 152                         "can not decompress dictionary: %s\n",
 153                         ZSTD_getErrorName(zCode));
 154                     free(buf);
 155                     goto Exit;
 156                 }
 157                 zCode = ZSTD_DCtx_loadDictionary(ctx, buf, zCode);
 158                 free(buf);
 159                 if ((zCode != 0) && (ZSTD_isError(zCode))) {
 160                     fprintf(
 161                         stderr,
 162                         "can not load dictionary: %s\n",
 163                         ZSTD_getErrorName(zCode));
 164                     goto Exit;
 165                 }
 166                 goto ReadAgain;
 167             default:
 168                 fprintf(stderr, "unknown dictionary header\n");
 169                 free(dict);
 170                 goto Exit;
 171             }
 172         }
 173         isEmpty = false;
 174         bIn.size = n;
 175         bIn.pos = 0;
 176         while (bIn.pos < bIn.size) {
 177             bOut.size = bufOutSize;
 178             bOut.pos = 0;
 179             zCode = ZSTD_decompressStream(ctx, &bOut, &bIn);
 180             if ((zCode != 0) && (ZSTD_isError(zCode))) {
 181                 fprintf(stderr, "can not decompress: %s\n", ZSTD_getErrorName(zCode));
 182                 goto Exit;
 183             }
 184             n = fwrite(bufOut, 1, bOut.pos, stdout);
 185             if (n != bOut.pos) {
 186                 perror("can not fwrite(stdout)");
 187                 goto Exit;
 188             }
 189             written += n;
 190             if (zCode == 0) {
 191                 offset += bIn.pos;
 192                 if (fdOff != NULL) {
 193                     fprintf(fdOff, "%zu\t%zu\n", offset - offsetPrev, written);
 194                 }
 195                 offsetPrev = offset + bIn.pos;
 196                 written = 0;
 197             }
 198         }
 199         if (lastBlock) {
 200             break;
 201         }
 202         offset += bIn.pos;
 203     }
 204
 205     if (isEmpty) {
 206         fputs("empty input\n", stderr);
 207         goto Exit;
 208     }
 209     if (zCode != 0) {
 210         fprintf(stderr, "unfinished decompression: %s\n", ZSTD_getErrorName(zCode));
 211         goto Exit;
 212     }
 213     rc = EXIT_SUCCESS;
 214
 215 Exit:
 216     if (bufOut != NULL) {
 217         free(bufOut);
 218     }
 219     if (bufIn != NULL) {
 220         free(bufIn);
 221     }
 222     ZSTD_freeDCtx(ctx);
 223     if ((fdOff != NULL) && (fclose(fdOff) != 0)) {
 224         perror("can not fclose(4)");
 225         return EXIT_FAILURE;
 226     }
 227     return rc;
 228 }